# Generalized Animal Imitator:
Agile Locomotion with Versatile Motion Prior

Ruihan Yang1∗, Zhuoqun Chen1∗, Jianhan Ma1∗, Chongyi Zheng2∗, Yiyu Chen3, Quan
Nguyen3, Xiaolong Wang1

###### Abstract

The agility of animals, particularly in complex activities such as running,
turning, jumping, and backflipping, stands as an exemplar for robotic system
design. Transferring this suite of behaviors to legged robotic systems
introduces essential inquiries: How can a robot be trained to learn multiple
locomotion behaviors simultaneously? How can the robot execute these tasks
with a smooth transition? And what strategies allow for the integrated
application of these skills? This paper introduces the Versatile Instructable
Motion prior (_VIM_) – a Reinforcement Learning framework designed to
incorporate a range of agile locomotion tasks suitable for advanced robotic
applications. Our framework enables legged robots to learn diverse agile low-
level skills by imitating animal motions and manually designed motions with
_Functionality_ reward and _Stylization_ reward. While the _Functionality_
reward guides the robot’s ability to adopt varied skills, the _Stylization_
reward ensures performance alignment with reference motions. Our evaluations
of the VIM framework span both simulation environments and real-world
deployment. To our understanding, this is the first work that allows a robot
to concurrently learn diverse agile locomotion tasks using a singular
controller. Further details and supportive media can be found at our project
site: https://rchalyang.github.io/VIM

Figure 1:  Real-Robot Trajectory. Our robot exhibits back-flipping skill in
the second row by imitating the reference motion in the first row.

†† *Equal contribution1UC San Diego, CA, USA.2Carnegie Mellon University, PA,
USA. 3University of Southern California, CA, USA

##  I INTRODUCTION

Research efforts have been invested for years in equipping legged robots with
agility comparable to that of natural quadrupeds. Picture a golden retriever
gracefully maneuvering in a park: darting, leaping over obstacles, and
pursuing a thrown ball. These tasks, effortlessly performed by many animals,
remain challenging for contemporary legged robots. To accomplish such tasks,
robots need not only master individual agile locomotion skills like running
and jumping but also the capacity to adaptively select and configure these
skills based on sensory inputs. We regard this kind of complicated task
requiring highly agile locomotion skills as advanced parkour for legged
robots. The inherent ability of quadrupeds to smoothly execute diverse
locomotion skills across varied tasks inspires our pursuit of a control system
with a general locomotion motion prior that includes these skills. In this
direction, we introduce a novel RL framework, Versatile Instructable Motion
prior (_VIM_) aiming to endow legged robots with a spectrum of reusable agile
locomotion skills by integrating existing agile locomotion knowledge.

Figure 2: Learning Agile Locomotion Skills from Reference Motions: Our system
learns a single instructable motion prior from a diverse reference motion
dataset.

Historically, agile gaits[5, 33, 34] for legged robots have been sculpted
using model-based or optimization methods. While promising, these methods
demand significant engineering input and precise state estimation. Learning-
based controllers enable robots to walk or run while addressing these
limitations, although they still fall short of agility. Imitation-based
controllers are also proposed to learn from motion sequences from animals [37]
or optimization methods [16]. Research on incorporating sensory information,
such as visual observations [2, 53, 54, 24, 57, 27, 59] or elevation maps [30,
41] further enables legged robots to traverse complex terrain like stones. In
spite of the encouraging results, most of these works focus on building a
single controller from scratch, even though much of the learned locomotion
skills could be shared across tasks. Recent works start building a
reconfigurable low-level motion prior [43, 23, 7, 38, 26, 22] for downstream
applications. However, the previous methods failed to make the best use of
existing skills to learn diverse locomotion skills with high agility.

In this work, we focus on building low-level motion prior to utilize existing
locomotion skills in nature and previous optimization methods, and learn
multiple highly agile locomotion skills simultaneously, as shown in Figure 2.
Even though we cannot fully comprehend the agility of animals and lack a
unified framework for model-based controls, we recognize that motion sequences
offer a consistent representation of diverse agile locomotion skills. Our
motion prior extracts and assimilates a range of locomotion skills from
reference motions, effectively mirroring their dynamics. These references
comprise motion capture (mocap) sequences from quadrupeds, augmented
generative model sequences complementing mocap data, and optimized motion
trajectories. Throughout the training phase, we translate varied reference
motion clips into a unified latent command space, guiding the motion prior to
recreate locomotion dynamics based on these latent commands and the robot’s
inherent state.

For legged robots, we define a locomotion skill as the ability of the robot to
produce a specific trajectory. To break down the intricacies of movement, we
classify it into two primary facets: _Functionality_ and _Style_.
_Functionality_ pertains to the fundamental movement objectives a robot aims
to achieve, such as advancing forward at a predefined speed. _Style_ , in
contrast, delves into the specific mechanics of how a robot accomplishes a
task. To illustrate, two robots might be programmed to progress at an
identical speed, but the intricacies of their movement—like step size or
frequency—might differ considerably. Simultaneously instructing a robot in
both these domains is nontrivial[52]. Drawing inspiration from how humans
learn complicated tasks, especially in fields demanding physical prowess like
athletics, we identify three core feedback modalities: objective performance
metrics, qualitative assessments, and granular kinematic guidance. Adopting
this structured feedback approach, our robot starts with mastering the basic
functional objective and subsequently turns into refining the detailed
locomotion gaits.

By incorporating diverse reference motions and our reward design, our
Versatile Instructable Motion prior (_VIM_) learns diverse agile locomotion
skills and makes them available for intricate downstream tasks. With our VIM,
we enable legged robots to perform advanced robotics parkour in the real
world. We also evaluate our method in the simulation and real world, as Figure
1. Our method significantly outperforms baselines in terms of final
performance and sample efficiency.

##  II Related Work

Blind Legged Locomotion: Classical legged locomotion controllers [19, 56, 49,
31, 5, 40] based on model-based methods [9, 17, 11, 12, 4, 20, 46] and
trajectory optimization [8, 34] have shown promising results in diverse tasks
with high levels of agility. Nonetheless, these methods normally come with
considerable engineering tuning for the specific task, high computation
requirements during deployment, or fragility to complex dynamics. Learning-
based methods [2, 29, 30, 44, 10, 15] controllers are proposed to offer robust
and lightweight controllers for deployment at the cost of offline computation.
Peng et al. [14] developed a controller producing non-agile life-like gaits by
imitating animal. Though previous works offer robust or agile locomotion
controllers across complex environments, these works focus on finishing a
single task at a time without reusing previous experience. Smith et al. [45]
utilize existing locomotion skills to solve specific downstream tasks.
Vollenweider et al. [51] utilize multiple AMP [39] to develop a controller to
solve a fixed task set. In this paper, our motion prior captures diverse agile
locomotion skills from reference motions generated by trajectory optimization
and provides them for intricate future downstream tasks.

Motion Priors: Due to the notorious low sample efficiency and considerable
effort required for reward engineering of RL, low-level skill pretraining has
drawn growing attention in recent years. Singh et. al [43] utilize the flow-
based model to build an actionable motion prior with motion sequences
generated by scripts. More recent works [23, 7, 38, 26, 13, 22] focus on
building low-level motion prior for downstream tasks but fail to include
diverse highly agile locomotion skills. In this work, we build motion prior
with reference motions consisting of mocap sequences, synthesized motion
sequences, and trajectories from optimization methods and learn multiple
highly agile locomotion skills with a single controller.

##  III Learn Versatile Instructable Motion Prior

We present the Versatile Instructable Motion prior (_VIM_), depicted in Figure
3, designed to acquire a wide range of agile locomotion skills concurrently
from multiple reference motions. The development of our motion prior involves
three stages: assembling a comprehensive dataset of reference motions sourced
from diverse origins, crafting a motion prior that processes varying reference
motions and the robot’s proprioceptive feedback to generate motor commands,
and finally, utilizing an imitation-based reward mechanism to effectively
train this motion prior.

Figure 3: Versatile Instructable Motion prior (_VIM_): Reference motion
encoder maps reference motions into latent skill space indicating target robot
pose and low-level policy output motor command. Reward Design: Our includes
_Functionality_ reward and _Style_ reward.

###  III-A Reference motion dataset.

Our primary objective was to curate a skill set for the robot that covers
diverse functions and agility levels, equipping it to handle complex
downstream tasks. Our dataset encompasses reference motions for locomotion
skills, including but not limited to canter, pace, walk, trot, turns,
backflips, and various jumps. These reference motions are derived from:_(a)_
mocap data of quadrupeds, specifically a subset from previous work [58],
despite its inherent challenges like noise due to the unpredictability of
animal behavior; _(b)_ synthesized (Syn) motions generated using a generative
model [58], aimed at enhancing dataset diversity by capturing challenging
locomotion actions;_(c)_ motions crafted through trajectory optimization
methods (Opt).

To address the methodology disparities between quadrupeds and our robot, we
retarget both mocap and synthesized sequences to our robot as Peng et al.
[37]. While mocap and synthesized motions offer extensive data, not all
sequences may be practically achievable by the robot. Thus, our dataset is
supplemented with motion sequences from trajectory optimization, emphasizing
intricate moves like jumps and backflips. The comprehensive reference motion
list can be found in Table II. Each trajectory in our dataset, represented as
$(s^{\text{ref}}_{0},s^{\text{ref}}_{1},\cdots,s^{\text{ref}}_{T})$, focuses
on the robot’s trunk and joint movements, excluding specific motor commands
which are absent in the captured and synthesized data. We denote the dataset
as
$\mathcal{D}=\\{(s^{\text{ref}}_{0},s^{\text{ref}}_{1},\cdots,s^{\text{ref}}_{T})_{i}\\}_{i=1}^{N}$.

###  III-B Motion Prior Structure

Our motion prior consists of a reference motion encoder, and a low-level
policy. Reference motion encoder maps varying reference motions into a
condensed latent skill space, and low-level policy utilizes our imitation
reward, reproduces the robot motion given a latent command.

Reference motion encoder: Our reference motion encoder
$\mathbb{E}_{\text{ref}}(\cdot)$ maps segments of reference motion to latent
commands in a latent skill space that outline the robot’s prospective
movement. These segments span both imminent and distant future states,
expressed as
$\hat{s}^{\text{ref}}_{t}=\\{s^{\text{ref}}_{t+1},s^{\text{ref}}_{t+2},s^{\text{ref}}_{t+10},s^{\text{ref}}_{t+30}\\}$.
We model the latent command as a Gaussian distribution
$\mathcal{N}(\mathbb{E}_{\text{ref}}^{\mu}(\hat{s}^{\text{ref}}_{t}),\mathbb{E}_{\text{ref}}^{\sigma}(\hat{s}^{\text{ref}}_{t}))$
from which we draw a sample at each interval to guide the low-level policy.

To maintain a _temporal-consistent_ latent skill space, our training
integrates an information bottleneck [48, 3] objective $L_{\text{AR}}$, where
the prior follows an auto-regressive model [6]. Specifically, given the
sampled latent command for the previous time step $z_{t-1}$, we minimize the
KL divergence between the current latent Gaussian distribution and a Gaussian
prior parameterized by $z_{t-1}$,

$\displaystyle
L_{\text{AR}}(\hat{s}^{\text{ref}}_{t},z_{t-1})=\beta\text{KL}\left(\mathcal{N}(\mu_{t},\sigma_{t}^{2})\parallel\mathcal{N}(\alpha
z_{t-1},(1-\alpha^{2})I)\right),$

where $\alpha=0.95$ is the scalar controlling the effect of correlation,
$\beta$ is the coefficient balancing regularization.

Low-level policy training: Our low-level policy $\pi_{\text{low}}$ takes
latent command $z_{t}$ representing the desired robot pose and robot’s current
proprioceptive state $s_{t}$ as input, and outputs actual motor commands
$a_{t}$ for the robot, where $s_{t}$ is encoded with a proprioception encoder
$\mathbb{E}_{\text{prop}}$. We train low-level policy and reference motion
encoder using PPO [42] in an end-to-end manner. Additionally, we introduce a
motion embedding for the critic to distinguish diverse reference motions.
Episodes initiate with randomized starting time steps from the dataset to
avert overfitting and conclude when the root pose tracking error escalates
beyond an acceptable range.

###  III-C Imitation Reward for Functionality and Style

Given the formulation of our motion prior, the robot learns diverse agile
locomotion skills with our imitation reward and reward scheduling mechanics.
Our reward offers consistent guidance, ensuring the robot captures both the
functionality and style inherent to the reference motion.

Learning Skill Functionality: To mirror the functionality of the reference
motion, we translate the root pose discrepancy between agent trajectories and
reference motion into a reward. The functionality reward $r_{\text{func}}$ is
subdivided into tracking rewards for robot root position
$r_{\text{func}}^{\text{pos}}$ and orientation $r_{\text{func}}^{\text{ori}}$.
Recognizing the distinct importance of vertical movement in agile tasks, the
root position tracking is further split into rewards for vertical
$r_{\text{func}}^{\text{pos-z}}$ and horizontal movements
$r_{\text{func}}^{\text{pos-xy}}$.

$\displaystyle
r_{\text{func}}(s_{t},\hat{s}_{t}^{\text{ref}})=w_{{\text{func}}}^{\text{ori}}*r_{\text{func}}^{\text{ori}}+w_{\text{func}}^{\text{pos-
xy}}*r_{\text{func}}^{\text{pos-
xy}}+w_{\text{func}}^{\text{pos-z}}*r_{\text{func}}^{\text{pos-z}}$

The specific formulation of our functionality rewards is provided as follows,
which is similar to previous work[37].

$\displaystyle r_{\text{func}}^{\text{ori}}(s_{t},\hat{s}^{\text{ref}}_{t})$
$\displaystyle=\exp\left(-10\left\lVert\hat{\mathbf{q}}_{t}^{\text{root}}-\mathbf{q}_{t}^{\text{root}}\right\rVert^{2}\right)$
$\displaystyle r_{\text{func}}^{\text{pos-
xy}}(s_{t},\hat{s}^{\text{ref}}_{t})$
$\displaystyle=\exp\left(-20\left\lVert\hat{\mathbf{x}}_{t}^{\text{root-
xy}}-\mathbf{x}_{t}^{\text{root-xy}}\right\rVert^{2}\right)$ $\displaystyle
r_{\text{func}}^{\text{pos-z}}(s_{t},\hat{s}^{\text{ref}}_{t})$
$\displaystyle=\exp\left(-80\left\lVert\hat{\mathbf{x}}_{t}^{\text{root-z}}-\mathbf{x}_{t}^{\text{root-z}}\right\rVert^{2}\right)$

where $\mathbf{q},\hat{\mathbf{q}}$ and $\mathbf{x},\hat{\mathbf{x}}$ denote
the root orientation and position from both the robot and reference motion,
respectively. Notably, in contrast to previous work [37], we allocate a
greater emphasis on root height in our reward, crucial for mastering agile
locomotion skills such as backflips and jumps.

Learning Skill Style: Capturing the style of a reference motion, in addition
to its functionality, enriches the application by meeting criteria such as
energy efficiency, robot safety, and facilitating human-robot interaction.
Drawing inspiration from how humans learn - starting by emulating the broader
style before focusing on intricate joint movements - our robot first mimics
the broader locomotion style with an adversarial style reward and later
refines its technique with a joint angle tracking reward.

Adversarial Stylization Reward: To swiftly encapsulate the style of the
locomotion skill, we train distinct discriminators $D_{i},\ i=1..n$ for all
$n$ reference motions separately to distinguish robot transitions from the
transition of the specific reference motion[39, 51] and use the output to
provide high-level feedback to the agent. Specifically, our discriminator is
trained with the following objective:

$\displaystyle\mathop{\mathrm{argmin}}_{D_{i}}\mathop{\mathbb{E}}_{d_{i}^{\mathcal{M}}(s,s^{\prime})}\left(D_{i}(s,s^{\prime})-1\right)^{2}+\mathop{\mathbb{E}}_{d_{i}^{\pi}(s,s^{\prime})}\left(D_{i}(s,s^{\prime})+1\right)^{2}$

where $d_{i}^{\mathcal{M}}(s,s^{\prime})$, $d_{i}^{\pi}(s,s^{\prime})$ denote
the transition distribution of the dataset and policy for $i$th reference
motion respectively.

For each reference motion, the likelihood from the discriminator is then
converted to a reward with:

$r_{\text{style}}^{\text{adv}}(s_{t},s_{t}^{\prime})=1-\frac{1}{4}*\left(1-D(s_{t},s_{t}^{\prime})\right)^{2}.$

Initially, our adversarial stylization reward provides dense reward and
enables the robot to learn a credible gait, but it can not provide more
detailed instructions as the training proceeds, which leads to mode collapse
and unstable training.

Joint Angle tracking Reward: On the other end, joint angle tracking reward
[36, 7] provides sparse but stable instruction for the robot to mimic the gait
of reference motion. Similar to our root pose tracking reward, our joint angle
tracking reward has the following formulation:

$\displaystyle
r_{\text{style}}^{\text{joint}}(s_{t},\hat{s}^{\text{ref}}_{t})$
$\displaystyle=\exp\left(-5\sum_{j\in\mathrm{joints}}\left\lVert\hat{\mathbf{q}}_{t}^{j}-\mathbf{q}_{t}^{j}\right\rVert^{2}\right)$
$\displaystyle+\exp\left(-20\sum_{f\in\mathrm{feet}}\left\lVert\hat{\mathbf{x}}_{t}^{f}-\mathbf{x}_{t}^{f}\right\rVert^{2}\right)$

where $\mathbf{q}_{t}^{j},\hat{\mathbf{q}}_{t}^{j}$ are the joint angle of
robot and reference motion and $\mathbf{e}_{t}^{f},\hat{\mathbf{e}}_{t}^{f}$
are the end-effector positions of robot and reference motion.

When learning diverse agile locomotion skills, only combining the joint angle
tracking reward and functionality reward leads to the failure of tracking
functionality or tracking the style of reference motion. Since different
locomotion skills are sensitive to different rewards.

Stylization Reward Scheduling: To take the best of both worlds, we propose to
use both adversarial stylization reward and joint angle tracking reward with a
balanced scheduling mechanism. Considering the discriminator as a "coach", we
utilize the mean adversarial reward as an indication of how the coach is
satisfied with the current performance. When it’s not satisfied with the
current performance of the robot, it provides detailed instruction for the
robot to learn. Specifically our stylization reward follows:

$\displaystyle r_{\text{style}}(s_{t},\hat{s}_{t}^{\text{ref}})$
$\displaystyle=w_{{\text{style}}}^{\text{adv}}*r_{\text{style}}^{\text{adv}}+w_{\text{style}}^{\text{joint}}*r_{\text{style}}^{\text{joint}}$
$\displaystyle+w_{{\text{style}}}^{\text{adv}}*(1-\displaystyle\mathop{\mathbb{E}}_{s_{t}\in
S}(r_{\text{style}}^{\text{adv}}(s,s^{\prime})))*r_{\text{style}}^{\text{joint}}$

With the given formulation, our stylization reward provided dense rewards
during the beginning of training, enabling the robot to quickly catch the
essence of different agile locomotion skills. Our stylization reward also
provides detailed instruction as the training proceeds, enabling the robot to
refine its gait and lead to more stable training.

###  III-D Solving Downstream Tasks with Motion Prior:

For intricate tasks like jumping over gaps, addressing them from scratch
presents challenges including acquiring necessary agile locomotion skills,
such as jumping and rapid running within limited interactions, and the
intensive engineering needed to harmonize the reward for top-tier tasks while
regularizing the robot’s motion. With a low-level motion prior, robots can
instantly harness existing skills encapsulated within the prior and channel
their efforts into high-level strategizing. For each distinct downstream task,
we train a high-level policy $\pi_{\text{high}}$ takes the high-level
observation $\mathbf{o}_{\text{high}}$ and outputs latent command for low-
level motion prior to utilize the existing agile locomotion skills:
$a_{t}=\pi_{\text{low}}(\pi_{\text{high}(\mathbf{o}_{\text{high}},\mathbf{E}_{\text{prop}}(s_{t}))})$.

##  IV Experiments

We evaluate our system in simulation and real-world, benchmarking it against
prior work for low-level locomotion skill learning and various downstream
tasks. Our robot demonstrates life-like agility in the real world.

Figure 4: Real World Jump Forward Trajectory Comparison: Each row represents a single trajectory (From top to bottom: Reference Motion, VIM, GAIL, Motion Imitation). TABLE I: Evaluation of Motion Prior in Simulation: We compare Horizontal and Vertical Root Position (Root Pos (XY), Root Pos (Height)), Root Orientation (Root Ori), Joint Angle, and End Effector Position (EE Pos) tracking errors and RL objectives of all methods. Our methods outperform all baselines in terms of smaller tracking errors, higher episodic returns, and longer episode lengths. GAIL baseline shows a smaller root position tracking error since it can’t follow the reference motion leading to early termination of the episode. | Tracking Error $\downarrow$ | RL Objectives $\uparrow$
---|---|---
Method | Root Pos (XY) | Root Pos (Height) | Root Ori | Joint Angle | EE Pos | Episode Return | Episode Length
VIM | $\mathbf{1.24\scriptstyle{\pm 0.62}}$ | $0.01\scriptstyle{\pm 0.02}$ | $0.11\scriptstyle{\pm 0.06}$ | $\mathbf{0.08\scriptstyle{\pm 0.06}}$ | $\mathbf{0.03\scriptstyle{\pm 0.03}}$ | $13.313\scriptstyle{\pm 11.48}$ | $166.783\scriptstyle{\pm 120.217}$
VIM (w/o Scheduling) | $1.28\scriptstyle{\pm 0.67}$ | $0.009\scriptstyle{\pm 0.0123}$ | $\mathbf{0.1\scriptstyle{\pm 0.06}}$ | $0.1\scriptstyle{\pm 0.08}$ | $0.05\scriptstyle{\pm 0.04}$ | $\mathbf{13.963\scriptstyle{\pm 11.395}}$ | $\mathbf{179.047\scriptstyle{\pm 121.788}}$
Motion Imitation | $1.39\scriptstyle{\pm 0.66}$ | $\mathbf{0.0077\scriptstyle{\pm 0.0114}}$ | $0.11\scriptstyle{\pm 0.05}$ | $0.25\scriptstyle{\pm 0.14}$ | $0.14\scriptstyle{\pm 0.08}$ | $9.536\scriptstyle{\pm 9.049}$ | $143.393\scriptstyle{\pm 114.514}$
GAIL | $1.04\scriptstyle{\pm 0.86}$ | $0.03\scriptstyle{\pm 0.03}$ | $0.13\scriptstyle{\pm 0.05}$ | $0.17\scriptstyle{\pm 0.1}$ | $0.09\scriptstyle{\pm 0.05}$ | $3.586\scriptstyle{\pm 6.166}$ | $54.723\scriptstyle{\pm 75.984}$
TABLE II: Evaluation of Motion Prior in Real (Left): We collect representative
metrics for different locomotion skills with corresponding metrics from
reference motion. $N/A$ denotes completely failed skills in real. Full
Reference Motion List (Right)

Metrics | VIM | Motion Imitation | GAIL | Reference Motion
---|---|---|---|---
Height (Jump While Running) $(m)$ | $\mathbf{0.50\scriptstyle{\pm 0.003}}$ | $0.42\scriptstyle{\pm 0.01}$ | $0.41\scriptstyle{\pm 0.04}$ | $0.53\scriptstyle{\pm 0.005}$
Height (Jump Forward) $(m)$ | $\mathbf{0.44\scriptstyle{\pm 0.01}}$ | $0.42\scriptstyle{\pm 0.01}$ | $0.27\scriptstyle{\pm 0.006}$ | $0.59\scriptstyle{\pm 0.006}$
Height (Jump Forward (Syn)) $(m)$ | $\mathbf{0.52\scriptstyle{\pm 0.01}}$ | $N/A$ | $N/A$ | $0.55\scriptstyle{\pm 0.007}$
Height (Backflip) $(m)$ | $\mathbf{0.62\scriptstyle{\pm 0.01}}$ | $0.49\scriptstyle{\pm 0.01}$ | $N/A$ | $0.60\scriptstyle{\pm 0.005}$
Distance (Jump While Running) $(m)$ | $\mathbf{0.48\scriptstyle{\pm 0.08}}$ | $0.35\scriptstyle{\pm 0.02}$ | $0.40\scriptstyle{\pm 0.003}$ | $0.56\scriptstyle{\pm 0.008}$
Distance (Jump Forward) $(m)$ | $\mathbf{0.76\scriptstyle{\pm 0.05}}$ | $0.40\scriptstyle{\pm 0.01}$ | $0.10\scriptstyle{\pm 0.002}$ | $0.82\scriptstyle{\pm 0.003}$
Distance (Jump Forward (Syn)) $(m)$ | $\mathbf{0.49\scriptstyle{\pm 0.04}}$ | $N/A$ | $N/A$ | $0.54\scriptstyle{\pm 0.007}$
Linear Velocity (Pace) $(m/s)$ | $\mathbf{0.76\scriptstyle{\pm 0.01}}$ | $0.97\scriptstyle{\pm 0.07}$ | $0.50\scriptstyle{\pm 0.02}$ | $0.72\scriptstyle{\pm 0.05}$
Linear Velocity (Canter) $(m/s)$ | $\mathbf{1.49\scriptstyle{\pm 0.15}}$ | $N/A$ | $N/A$ | $3.87\scriptstyle{\pm 0.17}$
Linear Velocity (Walk) $(m/s)$ | $0.90\scriptstyle{\pm 0.04}$ | $\mathbf{0.96\scriptstyle{\pm 0.06}}$ | $0.53\scriptstyle{\pm 0.58}$ | $0.97\scriptstyle{\pm 0.42}$
Linear Velocity (Trot) $(m/s)$ | $1.33\scriptstyle{\pm 0.17}$ | $\mathbf{1.05\scriptstyle{\pm 0.02}}$ | $0.93\scriptstyle{\pm 0.01}$ | $1.16\scriptstyle{\pm 0.12}$
Angular Velocity (Left Turn) $(rad/s)$ | $1.71\scriptstyle{\pm 0.04}$ | $0.00\scriptstyle{\pm 0.00}$ | $\mathbf{0.91\scriptstyle{\pm 0.04}}$ | $1.01\scriptstyle{\pm 0.05}$
Angular Velocity (Right Turn) $(rad/s)$ | $0.81\scriptstyle{\pm 0.02}$ | $\mathbf{0.62\scriptstyle{\pm 0.02}}$ | $0.63\scriptstyle{\pm 0.05}$ | $0.41\scriptstyle{\pm 0.09}$
Joint Angle Tracking Error ($rad^{2}/\textit{joint}$) | $\mathbf{0.10\scriptstyle{\pm 0.08}}$ | $0.27\scriptstyle{\pm 0.16}$ | $0.22\scriptstyle{\pm 0.10}$ | -

Skill Name | Source
---|---
Walk | Mocap
Trot | Mocap
Jump while Running | Mocap
Right Turn | Mocap
Left Turn | Mocap
Backflip | Opt
Jump Forward (Syn) | Syn
Left Turn (Syn) | Syn
Jump Forward | Opt
Canter | Mocap
Pace | Mocap

###  IV-A Evaluation of Learned Motion Priors

Our system’s proficiency in learning a range of agile locomotion skills from
the reference motion dataset (discussed in Sec III-A) is initially assessed.

Baselines: We benchmark our method against two representative baselines:
Motion Imitation [37, 7, 22] baseline represents a thread of recent works
whose imitation rewards are defined solely with errors between current robot
states and the corresponding reference states. Generative Adversarial
Imitation Learning (GAIL) baseline represents a thread of recent work [38],
whose imitation reward is solely provided by the discriminator trained to
distinguish trajectories generated by the policy from the ground truth
reference motions. Given that our reference motions consist only of state
sequences, they offer less supervision compared to expert action sequences,
rendering motion prior learning more challenging. Each method trains for
$2\times 10^{9}$ iterations across 3 random seeds. Both our technique and the
Motion Imitation baseline adopt identical reward scales for all motion error-
tracking rewards. Likewise, our approach and GAIL maintain the same scale for
adversarial stylization rewards.

Simulation Evaluation: In the simulation, we measure average imitation
tracking errors for various agile locomotion skills, episode returns, and
trajectory lengths across random seeds. Specifically as listed in Table I,
where the tracking error of root pose represents the ability of the robot to
reproduce the locomotion skill, and the tracking error of joint angle and end
effector position represents the ability of the robot to mimic the style of
reference motion. Our method achieves a similar root pose tracking error as
the motion imitation baseline with a much smaller joint angle tracking error.
This shows our method striking a balance between functionality and style,
superior to the motion imitation baseline that focuses solely on
functionality. Meanwhile, the GAIL baseline failed to learn the functionality
of the reference motions which leads to short episode length and the least
episode return. We surmise that the GAIL baseline’s inadequacy arises for two
main reasons: First, exclusive reliance on adversarial stylization reward does
not offer temporally consistent guidance throughout skill learning due to
misaligned rewards across timesteps. Second, the mode collapse issue inherent
in adversarial training hinders the robot from mastering highly agile skills,
such as backflipping. The shortcomings of the Motion Imitation baseline may
stem from the challenges of balancing different terms and selecting suitable
hyperparameters when concurrently learning multiple agile locomotion skills.
Comparing our VIM with and without stylization reward scheduling, we find the
former exhibits enhanced style tracking performance, underscoring the value of
stylization reward scheduling in refining robot gait tracking.

Real World Evaluation: We gauge learned agile locomotion skills in real-world
scenarios. Due to the lack of precise robot pose, we resort to specific
metrics tailored to different locomotion skills, detailed in Table II. For
Jump While Running & Jump Forward & Jump Forward (Syn) & Backflip, we measure
the jumping height and jumping distance. For Pace & Canter & Walk & Trot and
Left Turn & Right Turn, we measure the linear and angular velocity,
respectively. Results reveal that our method retains most of the reference
motion functionality. The only significant deviation, observed in the Canter
motion, arises from inherent differences between animal movement (its source)
and our robot’s capabilities. Even with comparable root pose tracking errors
in simulations, our method outshines the Motion Imitation baseline in real-
world metrics like jumping height, distance, and velocity tracking error. This
suggests that mirroring the style of the reference motion improves sim2real
transfer for natural gaits. GAIL baseline struggled to reproduce most real-
world locomotion skills. A visual comparison of real-world trajectories is
available in Figure 4, showing our method’s superiority in capturing both
motion functionality and style.

Figure 5: Latent Skill Space t-SNE. We visualize the latent skill space by
visualizing the latent embedding for varying motion segments.

Latent Skill Space Visualization: To better understand our low-level motion
prior, we visualize the learned latent skill space in Figure 5 by visualizing
the latent embedding corresponding to motion segments in our reference motion
dataset via t-SNE [50]. We find that different skills are separated into
different regions with clear boundaries. Our reference motion encoder also
clusters the skills with similar semantic meaning together: the embeddings
from Left Turn sequence is close to the embeddings from Right Turn sequence,
which enables the smooth transition between different skills, the embeddings
from Jump While Running & Jump Forward & Jump Forward (Syn) sequence are
clustered together. These observations suggest that our system learned a
smooth and semantically meaningful latent skill space for solving downstream
tasks.

###  IV-B Evaluation on Downstream Tasks

To further evaluate how our method leverages learned agile locomotion skills
for downstream tasks, we designed a set of tasks and tested our method against
standard baselines in both simulation and real-world scenarios.

Downstream Tasks: Our task suite comprises: Following Command: This involves
directing the robot to move with specific linear and angular velocities,
sampled uniformly between $0\sim 2$ m/s and $-2\sim 2$ rads/s. In our motion
prior, the robot is trained to move and turn at the reference motion’s speed;
hence, to follow a command precisely, the high-level policy should smoothly
interpolate between different speeds. Jump Forward: This task requires the
robot to execute a jump during a forward run. We have adapted a subset of
jumping rewards from CAJun [55] to evaluate policy interpolation between
jumping and running motions within a fixed timeframe. Following Command + Jump
Forward: Here, the robot must either jump forward or adjust to changing
commanded speeds. To optimize episode return, the robot should not only use
the agile locomotion skills from the reference motion dataset but also develop
unobserved skills like executing sharp turns.

Baselines: Considering the baseline’s subpar performance in low-level motion
prior training, we compare our system with three representative baselines
without pre-trained low-level controller: PPO [42]: Demonstrates controllers
trained exclusively on downstream task rewards. AMP [39] utilize existing
reference motion to provide styling reward in an adversarial imitation
learning manner and learn the policy for the downstream task while mimicking
the behavior of reference motions. Jain et al.Hierarchical Reinforcement
Learning (HRL) adapts from [25] which learns a high-level policy sending
latent commands to a low-level motor controller. HRL resembles a broad
category of prior works that decompose temporally extended reasoning into sub-
problems [47, 1, 35, 32, 18, 28, 21]. For a fair comparison, we made
modifications like removing the trajectory generator in [25], using PPO for
AMP and HRL, and supplying full reference motion data to AMP and HRL
integrated with AMP.

Figure 6: Downstream Tasks Evaluation in Simulation: Solid line and shaded
area denote the mean and std across random seeds. Our system outperforms all
baselines on all tasks.

Evaluation in Simulation & Real World: We train all methods on each downstream
task for $4\times 10^{8}$ environment samples with 3 random seeds. The
simulation results are detailed in Figure 6, and real-world results are
provided in Table III. For the Following Command task, while all methods
mastered basic locomotion, ours excelled in efficiency and smoothly
transitioned between diverse linear and angular velocities. The other tasks,
Jump Forward and Following Command + Jump Forward, demanded advanced jumping
abilities, which baselines couldn’t emerge. These baseline methods either
continuously moved forward, remained grounded when prompted to jump, or
toppled to evade energy consumption penalties. In contrast, our system
seamlessly bridged jumping and running actions, securing the highest episode
return. Despite providing with a comprehensive reference motion dataset,
baselines couldn’t harness the skills. This shortcoming possibly stems from
the challenge of deriving agile locomotion skills from the dataset using only
adversarial stylization rewards, mirroring the GAIL baseline’s poor
performance in low-level motion prior training.

TABLE III: Downstream Tasks Evaluation in Real: We compare Following Command +
Jump Forward policies of all methods in real, and $N/A$ denotes completely
failed skills in real. Our methods outperform all baselines in real for most
metrics.

Metrics (Vel for Velocity) | Ours | AMP | PPO | HRL
---|---|---|---|---
Max Linear Vel $(m/s)$ | $\mathbf{1.78\scriptstyle{\pm 0.13}}$ | $1.74\scriptstyle{\pm 0.21}$ | $1.75\scriptstyle{\pm 0.26}$ | $1.70\scriptstyle{\pm 0.08}$
Max Angular Vel (Left) $(rad/s)$ | $1.78\scriptstyle{\pm 0.004}$ | $1.07\scriptstyle{\pm 0.09}$ | $\mathbf{2.24\scriptstyle{\pm 0.05}}$ | $0.00\scriptstyle{\pm 0.00}$
Max Angular Vel (Right) $(rad/s)$ | $\mathbf{2.05\scriptstyle{\pm 0.02}}$ | $0.83\scriptstyle{\pm 0.09}$ | $1.75\scriptstyle{\pm 0.19}$ | $0.95\scriptstyle{\pm 0.37}$
Jump Distance $(m)$ | $\mathbf{0.50\scriptstyle{\pm 0.07}}$ | $0.00\scriptstyle{\pm 0.00}$ | $N/A$ | $N/A$
Jump Height $(m)$ | $\mathbf{0.50\scriptstyle{\pm 0.02}}$ | $0.38\scriptstyle{\pm 0.01}$ | $N/A$ | $N/A$

##  V Conclusion

In this paper, we propose Versatile Instructable Motion prior (_VIM_) which
learns agile locomotion skills from diverse reference motions with a single
motion prior. Our results in simulation and in the real world show that our
VIM captures both the functionality and the style of locomotion skills from
reference motions. Our VIM also provides a temporally consistent and compact
latent skill space representing different locomotion skills for different
downstream tasks. With agile locomotion skills in our VIM, complex downstream
tasks can be solved efficiently with minimum human effort.

## References

  * [1] “Hierarchical reinforcement learning with the maxq value function decomposition,” _Journal of artificial intelligence research_ , vol. 13, pp. 227–303, 2000.
  * [2] A. Agarwal, A. Kumar, J. Malik, and D. Pathak, “Legged locomotion in challenging terrains using egocentric vision,” in _6th Annual Conference on Robot Learning_ , 2022.
  * [3] A. A. Alemi, I. Fischer, J. V. Dillon, and K. Murphy, “Deep variational information bottleneck,” in _International Conference on Learning Representations_.
  * [4] G. Bledt and S. Kim, “Extracting legged locomotion heuristics with regularized predictive control,” in _2020 IEEE International Conference on Robotics and Automation (ICRA)_. IEEE, 2020, pp. 406–412.
  * [5] G. Bledt, M. J. Powell, B. Katz, J. Di Carlo, P. M. Wensing, and S. Kim, “Mit cheetah 3: Design and control of a robust, dynamic quadruped robot,” in _2018 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_. IEEE, 2018, pp. 2245–2252.
  * [6] S. Bohez, S. Tunyasuvunakool, P. Brakel, F. Sadeghi, L. Hasenclever, Y. Tassa, E. Parisotto, J. Humplik, T. Haarnoja, R. Hafner _et al._ , “Imitate and repurpose: Learning reusable robot movement skills from human and animal behaviors,” _arXiv preprint arXiv:2203.17138_ , 2022.
  * [7] S. Bohez, S. Tunyasuvunakool, P. Brakel, F. Sadeghi, L. Hasenclever, Y. Tassa, E. Parisotto, J. Humplik, T. Haarnoja, R. Hafner, M. Wulfmeier, M. Neunert, B. Moran, N. Siegel, A. Huber, F. Romano, N. Batchelor, F. Casarini, J. Merel, R. Hadsell, and N. Heess, “Imitate and repurpose: Learning reusable robot movement skills from human and animal behaviors,” 2022. [Online]. Available: https://arxiv.org/abs/2203.17138
  * [8] J. Carius, R. Ranftl, V. Koltun, and M. Hutter, “Trajectory optimization for legged robots with slipping motions,” _IEEE Robotics and Automation Letters_ , vol. 4, no. 3, pp. 3013–3020, 2019.
  * [9] J. D. Carlo, P. M. Wensing, B. Katz, G. Bledt, and S. Kim, “Dynamic locomotion in the MIT cheetah 3 through convex model-predictive control,” in _2018 IEEE/RSJ International Conference on Intelligent Robots and Systems, IROS 2018, Madrid, Spain, October 1-5, 2018_. IEEE, 2018, pp. 1–9. [Online]. Available: https://doi.org/10.1109/IROS.2018.8594448
  * [10] X. Cheng, A. Kumar, and D. Pathak, “Legs as manipulator: Pushing quadrupedal agility beyond locomotion,” in _2023 IEEE International Conference on Robotics and Automation (ICRA)_ , 2023.
  * [11] J. Di Carlo, P. M. Wensing, B. Katz, G. Bledt, and S. Kim, “Dynamic locomotion in the mit cheetah 3 through convex model-predictive control,” in _2018 IEEE/RSJ International Conference on Intelligent Robots and Systems (IROS)_. IEEE, 2018, pp. 1–9.
  * [12] Y. Ding, A. Pandala, and H.-W. Park, “Real-time model predictive control for versatile dynamic motions in quadrupedal robots,” in _2019 International Conference on Robotics and Automation (ICRA)_. IEEE, 2019, pp. 8484–8490.
  * [13] B. Eysenbach, A. Gupta, J. Ibarz, and S. Levine, “Diversity is all you need: Learning skills without a reward function,” in _International Conference on Learning Representations_ , 2019. [Online]. Available: https://openreview.net/forum?id=SJx63jRqFm
  * [14] P. Fankhauser, M. Bloesch, C. Gehring, M. Hutter, and R. Siegwart, “Robot-centric elevation mapping with uncertainty estimates,” in _International Conference on Climbing and Walking Robots (CLAWAR)_ , 2014.
  * [15] Z. Fu, X. Cheng, and D. Pathak, “Deep whole-body control: Learning a unified policy for manipulation and locomotion,” _Conference on Robot Learning (CoRL)_ , 2022.
  * [16] Y. Fuchioka, Z. Xie, and M. van de Panne, “Opt-mimic: Imitation of optimized trajectories for dynamic quadruped behaviors,” 2022. [Online]. Available: https://arxiv.org/abs/2210.01247
  * [17] C. Gehring, S. Coros, M. Hutter, M. Blösch, M. A. Hoepflinger, and R. Siegwart, “Control of dynamic gaits for a quadrupedal robot,” in _2013 IEEE International Conference on Robotics and Automation, Karlsruhe, Germany, May 6-10, 2013_. IEEE, 2013, pp. 3287–3292. [Online]. Available: https://doi.org/10.1109/ICRA.2013.6631035
  * [18] J. Gehring, G. Synnaeve, A. Krause, and N. Usunier, “Hierarchical skills for efficient exploration,” _Advances in Neural Information Processing Systems_ , vol. 34, pp. 11 553–11 564, 2021.
  * [19] H. Geyer, A. Seyfarth, and R. Blickhan, “Positive force feedback in bouncing gaits?” _Proceedings of the Royal Society of London. Series B: Biological Sciences_ , vol. 270, no. 1529, pp. 2173–2183, 2003.
  * [20] R. Grandia, F. Farshidian, A. Dosovitskiy, R. Ranftl, and M. Hutter, “Frequency-aware model predictive control,” _IEEE Robotics and Automation Letters_ , vol. 4, no. 2, pp. 1517–1524, 2019.
  * [21] A. Gupta, V. Kumar, C. Lynch, S. Levine, and K. Hausman, “Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning,” in _Conference on Robot Learning_. PMLR, 2020, pp. 1025–1037.
  * [22] L. Han, Q. Zhu, J. Sheng, C. Zhang, T. Li, Y. Zhang, H. Zhang, Y. Liu, C. Zhou, R. Zhao, J. Li, Y. Zhang, R. Wang, W. Chi, X. Li, Y. Zhu, L. Xiang, X. Teng, and Z. Zhang, “Lifelike agility and play on quadrupedal robots using reinforcement learning and generative pre-trained models,” 2023.
  * [23] L. Hasenclever, F. Pardo, R. Hadsell, N. Heess, and J. Merel, “CoMic: Complementary task learning & mimicry for reusable skills,” in _Proceedings of the 37th International Conference on Machine Learning_ , ser. Proceedings of Machine Learning Research, H. D. III and A. Singh, Eds., vol. 119. PMLR, 13–18 Jul 2020, pp. 4105–4115. [Online]. Available: https://proceedings.mlr.press/v119/hasenclever20a.html
  * [24] C. Imai, M. Zhang, Y. Zhang, M. Kierebinski, R. Yang, Y. Qin, and X. Wang, “Vision-guided quadrupedal locomotion in the wild with multi-modal delay randomization,” in _2022 IEEE/RSJ international conference on intelligent robots and systems (IROS)_.
  * [25] D. Jain, K. Caluwaerts, and A. Iscen, “From pixels to legs: Hierarchical learning of quadruped locomotion,” in _Conference on Robot Learning_. PMLR, 2021, pp. 91–102.
  * [26] J. Juravsky, Y. Guo, S. Fidler, and X. B. Peng, “Padl: Language-directed physics-based character control,” in _SIGGRAPH Asia 2022 Conference Papers_ , ser. SA ’22. New York, NY, USA: Association for Computing Machinery, 2022. [Online]. Available: https://doi.org/10.1145/3550469.3555391
  * [27] S. Kareer, N. Yokoyama, D. Batra, S. Ha, and J. Truong, “Vinl: Visual navigation and locomotion over obstacles,” _arXiv preprint arXiv:2210.14791_ , 2022.
  * [28] M. Klissarov and D. Precup, “Flexible option learning,” _Advances in Neural Information Processing Systems_ , vol. 34, pp. 4632–4646, 2021.
  * [29] A. Kumar, Z. Fu, D. Pathak, and J. Malik, “Rma: Rapid motor adaptation for legged robot,” _Robotics: Science and Systems_ , 2021.
  * [30] T. Miki, J. Lee, J. Hwangbo, L. Wellhausen, V. Koltun, and M. Hutter, “Learning robust perceptive locomotion for quadrupedal robots in the wild,” _Sci Robot_ , vol. 7, no. 62, p. eabk2822, Jan. 2022.
  * [31] H. Miura and I. Shimoyama, “Dynamic walk of a biped,” _The International Journal of Robotics Research_ , vol. 3, no. 2, pp. 60–74, 1984.
  * [32] O. Nachum, S. S. Gu, H. Lee, and S. Levine, “Data-efficient hierarchical reinforcement learning,” _Advances in neural information processing systems_ , vol. 31, 2018.
  * [33] C. Nguyen, L. Bao, and Q. Nguyen, “Continuous jumping for legged robots on stepping stones via trajectory optimization and model predictive control,” in _2022 IEEE 61st Conference on Decision and Control (CDC)_. IEEE, 2022, pp. 93–99.
  * [34] Q. Nguyen, M. J. Powell, B. Katz, J. Di Carlo, and S. Kim, “Optimized jumping on the mit cheetah 3 robot,” in _2019 International Conference on Robotics and Automation (ICRA)_. IEEE, 2019, pp. 7448–7454.
  * [35] R. Parr and S. Russell, “Reinforcement learning with hierarchies of machines,” _Advances in neural information processing systems_ , vol. 10, 1997.
  * [36] X. B. Peng, P. Abbeel, S. Levine, and M. van de Panne, “Deepmimic: Example-guided deep reinforcement learning of physics-based character skills,” _ACM Trans. Graph._ , vol. 37, no. 4, pp. 143:1–143:14, Jul. 2018. [Online]. Available: http://doi.acm.org/10.1145/3197517.3201311
  * [37] X. B. Peng, E. Coumans, T. Zhang, T.-W. E. Lee, J. Tan, and S. Levine, “Learning agile robotic locomotion skills by imitating animals,” in _Robotics: Science and Systems_ , 07 2020.
  * [38] X. B. Peng, Y. Guo, L. Halper, S. Levine, and S. Fidler, “Ase: Large-scale reusable adversarial skill embeddings for physically simulated characters,” _ACM Trans. Graph._ , vol. 41, no. 4, Jul. 2022.
  * [39] X. B. Peng, Z. Ma, P. Abbeel, S. Levine, and A. Kanazawa, “Amp: Adversarial motion priors for stylized physics-based character control,” _ACM Trans. Graph._ , vol. 40, no. 4, Jul. 2021. [Online]. Available: http://doi.acm.org/10.1145/3450626.3459670
  * [40] M. H. Raibert, “Hopping in legged systems—modeling and simulation for the two-dimensional one-legged case,” _IEEE Transactions on Systems, Man, and Cybernetics_ , vol. SMC-14, no. 3, pp. 451–463, 1984.
  * [41] N. Rudin, D. Hoeller, P. Reist, and M. Hutter, “Learning to walk in minutes using massively parallel deep reinforcement learning,” in _Conference on Robot Learning_. PMLR, 2022, pp. 91–100.
  * [42] J. Schulman, F. Wolski, P. Dhariwal, A. Radford, and O. Klimov, “Proximal policy optimization algorithms,” _arXiv preprint arXiv:1707.06347_ , 2017.
  * [43] A. Singh, H. Liu, G. Zhou, A. Yu, N. Rhinehart, and S. Levine, “Parrot: Data-driven behavioral priors for reinforcement learning,” 2020.
  * [44] J.-P. Sleiman, F. Farshidian, and M. Hutter, “Versatile multicontact planning and control for legged loco-manipulation,” _Science Robotics_ , vol. 8, no. 81, p. eadg5014, 2023. [Online]. Available: https://www.science.org/doi/abs/10.1126/scirobotics.adg5014
  * [45] L. Smith, J. C. Kew, T. Li, L. Luu, X. B. Peng, S. Ha, J. Tan, and S. Levine, “Learning and adapting agile locomotion skills by transferring experience,” _arXiv preprint arXiv:2304.09834_ , 2023.
  * [46] Y. Sun, W. L. Ubellacker, W.-L. Ma, X. Zhang, C. Wang, N. V. Csomay-Shanklin, M. Tomizuka, K. Sreenath, and A. D. Ames, “Online learning of unknown dynamics for model-based controllers in legged locomotion,” _IEEE Robotics and Automation Letters (RA-L)_ , 2021.
  * [47] R. S. Suttona, D. Precup, and S. Singha, “Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning,” _Artificial Intelligence_ , vol. 112, pp. 181–211, 1999.
  * [48] N. TISHBY, “The information bottleneck method,” _Computing Research Repository (CoRR)_ , 2000.
  * [49] N. Torkos and M. van de Panne, “Footprint–based quadruped motion synthesis,” in _Proceedings of the Graphics Interface 1998 Conference, June 18-20, 1998, Vancouver, BC, Canada_ , June 1998, pp. 151–160. [Online]. Available: http://graphicsinterface.org/wp-content/uploads/gi1998-19.pdf
  * [50] L. Van der Maaten and G. Hinton, “Visualizing data using t-sne.” _Journal of machine learning research_ , vol. 9, no. 11, 2008.
  * [51] E. Vollenweider, M. Bjelonic, V. Klemm, N. Rudin, J. Lee, and M. Hutter, “Advanced skills through multiple adversarial motion priors in reinforcement learning,” 2022.
  * [52] J. Xu, Y. Tian, P. Ma, D. Rus, S. Sueda, and W. Matusik, “Prediction-guided multi-objective reinforcement learning for continuous robot control,” in _Proceedings of the 37th International Conference on Machine Learning_ , 2020.
  * [53] R. Yang, G. Yang, and X. Wang, “Neural volumetric memory for visual locomotion control,” in _Conference on Computer Vision and Pattern Recognition 2023_ , 2023. [Online]. Available: https://openreview.net/forum?id=JYyWCcmwDS
  * [54] R. Yang, M. Zhang, N. Hansen, H. Xu, and X. Wang, “Learning vision-guided quadrupedal locomotion end-to-end with cross-modal transformers,” in _International Conference on Learning Representations_ , 2022. [Online]. Available: https://openreview.net/forum?id=nhnJ3oo6AB
  * [55] Y. Yang, G. Shi, X. Meng, W. Yu, T. Zhang, J. Tan, and B. Boots, “Cajun: Continuous adaptive jumping using a learned centroidal controller,” _arXiv preprint arXiv:2306.09557_ , 2023.
  * [56] K. Yin, K. Loken, and M. Van de Panne, “Simbicon: Simple biped locomotion control,” _ACM Transactions on Graphics (TOG)_ , vol. 26, no. 3, pp. 105–es, 2007.
  * [57] W. Yu, D. Jain, A. Escontrela, A. Iscen, P. Xu, E. Coumans, S. Ha, J. Tan, and T. Zhang, “Visual-locomotion: Learning to walk on complex terrains with vision,” in _Proceedings of the 5th Conference on Robot Learning_ , ser. Proceedings of Machine Learning Research, A. Faust, D. Hsu, and G. Neumann, Eds., vol. 164. PMLR, 08–11 Nov 2022, pp. 1291–1302. [Online]. Available: https://proceedings.mlr.press/v164/yu22a.html
  * [58] H. Zhang, S. Starke, T. Komura, and J. Saito, “Mode-adaptive neural networks for quadruped motion control,” _ACM Transactions on Graphics (TOG)_ , vol. 37, no. 4, pp. 1–11, 2018.
  * [59] Z. Zhuang, Z. Fu, J. Wang, C. G. Atkeson, S. Schwertfeger, C. Finn, and H. Zhao, “Robot parkour learning,” in _7th Annual Conference on Robot Learning_ , 2023. [Online]. Available: https://openreview.net/forum?id=uo937r5eTE


# Striking the Right Balance: Recall Loss for Semantic Segmentation

Junjiao Tian*, Niluthpol Mithun$\dagger$, Zachary Seymour$\dagger$, Han-Pang
Chiu$\dagger$, Zsolt Kira*
<EMAIL_ADDRESS>
<EMAIL_ADDRESS>
###### Abstract

Class imbalance is a fundamental problem in computer vision applications such
as semantic segmentation. Specifically, uneven class distributions in a
training dataset often result in unsatisfactory performance on under-
represented classes. Many works have proposed to weight the standard cross
entropy loss function with pre-computed weights based on class statistics,
such as the number of samples and class margins. There are two major drawbacks
to these methods: 1) constantly up-weighting minority classes can introduce
excessive false positives in semantic segmentation; 2) a minority class is not
necessarily a hard class. The consequence is low precision due to excessive
false positives. In this regard, we propose a hard-class mining loss by
reshaping the vanilla cross entropy loss such that it weights the loss for
each class dynamically based on instantaneous recall performance. We show that
the novel recall loss changes gradually between the standard cross entropy
loss and the inverse frequency weighted loss. Recall loss also leads to
improved mean accuracy while offering competitive mean Intersection over Union
(IoU) performance. On Synthia dataset111Synthia-Sequence Summer split, recall
loss achieves $9\%$ relative improvement on mean accuracy with competitive
mean IoU using DeepLab-ResNet18 compared to the cross entropy loss. Code
available at https://github.com/PotatoTian/recall-semseg.

##  I Introduction

Dataset imbalance [1, 2] is an important problem for many computer vision
tasks, such as semantic segmentation and image classification. In semantic
segmentation, imbalance occurs as a result of natural occurrence and varying
sizes of different classes. For example, in an outdoor driving segmentation
dataset, light poles and pedestrians are considered minority classes compared
to large classes such as building, sky, and road. These minority classes are
often more important than large classes for safety reasons. Imbalance also
affects the metrics used by segmentation tasks. Recent works [3] [4] [5] use
both mean accuracy and mean Intersection over Union (IoU) to measure
segmentation performance so that large classes do not dominate the evaluation.
While mean IoU provides a more holistic view of the performance of a model
than accuracy because it considers false positive predictions (see Table I),
mean accuracy (also as mean recall) does not consider the effects of false
positives [6]. One might argue that high recall for each class regardless of
imbalance is important for safety. It can degrade segmentation quality, as
shown in Fig. 3, where small classes become indistinguishable from each other.
Therefore, it is important to improve recall while maintaining a competitive
mean IoU without introducing excessive false positives.

Researchers have studied the imbalance problem for classification, detection,
and segmentation extensively. Most prior research has been on designing
balanced loss functions. We classify existing loss functions under three
categories: region-based losses, statistics-balanced losses and performance-
balanced losses. Region-based losses directly optimize region metrics (e.g.,
Jaccard index [7]) and are mainly popular in medical segmentation
applications; Statistics-balanced losses (e.g., LDAM [8], Class-Balanced (CB)
loss [9]) up/down weights the contribution of a class based on its class
margin or class size; however, they tend to encourage excessive false
positives in minority classes to improve mean accuracy especially in
segmentation. A recent study in [10] also shows that the weighting undermines
the generic representation learning capability of the feature extractors;
Performance-balanced losses (e.g., Focal loss [11]) use a certain performance
indicator to weight the loss of each class. As an example, Focal loss assumes
that the difficulty of a class is correlated with imbalance and can be
reflected by the predicted confidence. However, it has not been very
successful in other applications (such as image classification and
segmentation) for dealing with imbalance, as reported by [9]. We investigate
the reasons of failure in Sec. IV-D.

Figure 1: Mean accuracy (ACC: $\frac{TP}{TP+FN}$) vs. mean intersection over
union (IOU: $\frac{TP}{TP+FN+FP}$) on Synthia dataset with Deeplab (Resnet18).
Many losses improve accuracy (recall) at the expense of IoU. Only losses in
the upper-right quadrant improve accuracy without introducing excessive false
positives. Note: TP: true positive, FN: false negative, FP: false positive.

We propose a novel performance-balanced loss, called Recall loss, using the
recall metric to address the imbalance problem. Recall loss down/up weights a
class based on the training recall performance of that class. It is an example
of hard class mining as supposed to the hard example mining strategy in Focal
loss. Unlike the statistics-balanced losses, Recall loss dynamically changes
its weights with training based on per-class recall performance. The dynamism
is the key to overcome many drawbacks of the statistics-balanced losses. In
our experiments, the CB loss improves accuracy at the expense of Intersection
over Union (IOU) which considers false positives in semantic segmentation.
However, our Recall loss can effectively balance between precision and recall
of each class, and hence, it improves accuracy but maintains a competitive
IOU. Fig. 1 shows a mean accuracy VS. mean IoU plot for different losses on
the Synthia segmentation dataset [12]. We use Focal loss as a reference in the
figure because it does not introduce excessive false positives. While some
losses provide better mean accuracy, the improvement comes at the expense of
lower mean IOU. Our proposed Recall loss improves accuracy considerably while
still maintaining a competitive IoU. Further, experiments on two semantic
segmentation datasets demonstrate that Recall loss shows significantly better
performance than state-of-the-art loss functions in balancing recall and
precision.

##  II Related Work

Region-balanced Loss. In image segmentation, Dice and Jaccard indices (IoU)
are commonly used as the evaluation metrics. However, the most common training
criterion, cross entropy, does not directly optimize these metrics. In medical
imaging, researchers propose to optimize soft/surrogate version of these
indices. SoftIOU [7] proposes to optimize a soft version of the Jaccard index;
Lovasz Softmax [13] also optimizes the Jaccard index based on the Lovasz
convex extension; SoftDice [14] and softTversky [15] optimize a soft version
of the Dice index and Tversky index (see Table I). However, concerns have been
raised in [16] on whether these losses consider the trade-off between false
positives and false negatives. We show that they also tend to yield high mean
accuracy at the expense of lower mean IoU.

Statistics-balanced Loss. Various losses have been proposed to deal with
imbalance or long-tail distributions using weighted losses. The most popular
loss is the inverse-frequency loss (Weighted Cross-Entropy). It weights the
cross entropy loss of each class by its inverse frequency. Class-Balanced Loss
[9] motivates a weighted cross entropy loss with the concept of effective
number of samples in each class. LDAM [8] also derives a weighted cross
entropy loss based on margins between classes. Some losses require changes to
the network and require an iterative optimization process, so they are not
easily adaptable to semantic segmentation. We compare to the representative
Class-Balanced loss in our experiments.

Performance-balanced Loss. Imbalance is also a problem in object detection,
where the foreground-background imbalance is extreme and undermines learning.
Online Hard Example Mining (OHEM) [17] proposes to find hard examples by
ranking the losses and only keeping those with the highest losses. Focal loss
[11] chooses to down weight easy samples and emphasizes hard samples by
weighting each sample by $1-p$ where $p$ is the predicted probability for the
sample. The weight for each sample dynamically changes with training, and the
method never completely discards any samples. Focal loss is especially
successful because it is easy to implement and proves effective in the binary
foreground-background imbalance setting. However, it has not been very
successful in other applications for dealing with imbalance, as reported by
[9].

##  III Recall Loss

###  III-A Motivation: From Inverse Frequency Loss to Recall Loss

To motivate our proposed loss, we first analyze the standard cross entropy
loss. Let $\\{x_{n},y_{n}\\}\forall n\in\\{1,...,N\\}$, where $x_{n}\in
R^{d},y_{n}\in\\{1,...,C\\}$ denote the set of training data and corresponding
labels. Let $P_{n}$ denotes the predictive softmax-distribution over all
classes for input $x_{n}$ and $P_{n}^{i}$ denotes the probability of the
$i$-th class. The cross entropy loss used in multiclass classification is
defined as:

$\displaystyle CE$ $\displaystyle=-\sum_{n=1}^{N}\log(P_{n}^{y_{n}})$ (1)
$\displaystyle=-\sum_{c=1}^{C}\sum_{n:y_{n}=c}\log(P_{n}^{y_{n}})=-\sum_{c=1}^{C}N_{c}\log(P^{c})$

where $\\{n:y_{i}=c\\}$ denotes the set of samples whose label is $c$ and
$P^{c}=(\prod_{n:y_{n}=c}P_{n}^{y_{n}})^{1/N_{c}}$ denotes the geometric mean
confidence of a class $c$ and $N_{c}$ denotes the number of samples in that
class. As shown in Eq. 1, the conventional cross entropy optimizes the
geometric mean confidence of each class, weighted by the number of pixels in
each class. When there is a significant class imbalance in the dataset, the
loss function biases towards large classes as a result of larger $N_{c}$.

One commonly used loss for imbalanced datasets is inverse frequency cross
entropy loss [18, 19] which assigns more weight to the loss of minority
classes. Let $N$ denote the total number of pixels in the training set and
$N_{c}$ denotes the number of pixels belonging to class $c\in\\{1,..,C\\}$.
The frequency of a class is calculated as $freq(c)=N_{c}/N$. We show that
while the unweighted cross entropy loss optimizes the overall confidence, the
loss weighted by inverse frequency optimizes mean confidence. If we use an
inverse frequency weighting, the loss is rebalanced. Note, we leave out the
$N$ in $freq(c)$ as it is shared by all classes.

$\displaystyle InvCE$
$\displaystyle=-\sum_{c=1}^{C}\frac{1}{freq(c)}N_{c}\log(P^{c})$ (2)
$\displaystyle=-\sum_{c=1}^{C}\frac{1}{N_{c}}N_{c}\log(P^{c})=-\sum_{c=1}^{C}\log(P^{c})$

As shown in Eq. 2, the weighted loss optimizes the geometric mean of accuracy
directly. However, the inverse frequency loss might not be optimal in practice
because it over-weights the minority classes and introduces excessive false
positives, i.e., it sacrifices precision for recall. This problem is
especially severe in semantic segmentation [20]. Applying the inverse
frequency loss to segmentation increases recall for each class. However, the
improvement comes at the cost of excessive false positives, especially for
small classes (See Fig. 3).

While the inverse frequency loss solves the problem of imbalance, it focuses
only on improving one aspect of the problem in classification, i.e. the recall
of each class. To solve this issue, we propose to weight the inverse frequency
loss in Eq. 2 with the false negative ($FN_{c}$) counts for each class. The
first insight is that the $FN_{c}$ is bounded by the total number of samples
in a class and zero, i.e.

$\displaystyle N_{c}\geq FN_{c}\geq 0$ (3)

By weighting the inverse frequency cross entropy loss in Eq. 2 by the false
negative counts for each class, we obtain a moderate loss function which sits
between the regular cross entropy loss and inverse frequency loss.

$\displaystyle RecallCE$
$\displaystyle=-\sum_{c=1}^{C}FN_{c}\log(P^{c})=-\sum_{c=1}^{C}\frac{FN_{c}}{N_{c}}N_{c}\log(P^{c})$
(4)
$\displaystyle=-\sum_{c=1}^{C}\frac{FN_{c}}{FN_{c}+TP_{c}}N_{c}\log(P^{c})$

As Eq. 4 shows, the loss can be implemented as the regular cross entropy loss
weighted by class-wise false negative rate (FNR). The second insight is that
minority classes are most likely more difficult to classify with higher FNR
and large classes with smaller FNR. Therefore, similar to inverse frequency
loss, gradients of minority classes will be boosted and gradients of majority
classes will be suppressed. However, unlike frequency weighting, the weighting
will not be as extreme as motivated in Eq. 3.

Unlike frequency and decision margin [8] which are characteristics of the
dataset, FNR is a metric of a model’s performance. As we continue to update
the model’s parameters, FNR changes. Therefore, the weights for each class
change dynamically to reflect a model’s instantaneous performance. We rewrite
Eq. 4 using recall ($1-FNR$) with a subscript $t$ to denote the time
dependency.

$\displaystyle RecallCE$
$\displaystyle=-\sum_{c=1}^{C}(1-\frac{TP_{c,t}}{FN_{c,t}+TP_{c,t}})N_{c}\log(p^{c,t})$
(5)
$\displaystyle=-\sum_{c=1}^{C}\sum_{n:y_{i}=c}(1-\mathcal{R}_{c,t})\log(p_{n,t})$

where $\mathcal{R}_{c,t}$ is the recall for class $c$ at optimization step
$t$.

###  III-B Gradient Analysis: Recall Loss Balances Recall and Precision

To validate the claim that recall loss balances recall and precision, we
conduct a gradient analysis of the recall loss. Let’s assume a binary
classification task for clarity. Let $[z_{1},z_{2}]$ and $[P_{1},P_{2}]$
denote the pre-softmax logits and post-softmax probabilities of a classifier.
We first obtain the gradient for the standard Cross-Entropy loss with respect
to the logits for a single input.

$\displaystyle\nabla_{z_{i}}CE$
$\displaystyle=-\log(P^{y})=-\log(\frac{e^{z_{y}}}{\sum_{i}e^{z_{i}}})=P_{i}-\mathcal{I}(y=i)$
(6)

where $\mathcal{I}(y=i)=1$ if $y=1$ and $0$ otherwise. Assume we have a
dataset with $N_{1}$ and $N_{2}$ number of samples for each class. The
gradients of the recall loss w.r.t the logit of the first class, $z_{1}$, is:

$\displaystyle-\nabla_{z_{1}}\left[(1-\mathcal{R}_{1})\sum_{n=1}^{N_{1}}\log(p_{n}^{1})+(1-\mathcal{R}_{2})\sum_{n=1}^{N_{2}}\log(p_{n}^{2})\right]$
(7)
$\displaystyle=(1-\mathcal{R}_{1})\sum_{n=1}^{N_{1}}(p_{1,n}^{(1)}-1)+(1-\mathcal{R}_{2})\sum_{n=1}^{N_{2}}(p_{1,n}^{(2)})$
$\displaystyle=FN_{1}(P_{1}^{(1)}-1)+FN_{2}P_{1}^{(2)}$

where the superscript $(j)$ in $P^{(j)}_{i}$ denotes the ground truth class
and the subscript denotes the class w.r.t which the gradient is calculated.
$P_{1}^{(1)}=\frac{1}{N_{1}}\sum_{n=1}^{N_{1}}p_{1,n}^{(1)}$ denotes the
average confidence of class 1 when it is the ground truth class, and
$P_{1}^{(2)}=\frac{1}{N_{2}}\sum_{n=1}^{N_{2}}p_{1,n}^{(2)}$ denotes the
average confidence of class 1 when the ground truth class is 2.

To see how the recall loss affects gradients backprobagation to the logits, we
can use the fact that in a binary classification problem, a false negative
(FN) in one class is a false positive (FP) in the other to replace $FN_{2}$ by
$FP_{1}$ in Eq. 7.

$\displaystyle\nabla_{z_{1}}RecallCE$
$\displaystyle=FN_{1}(P_{1}^{(1)}-1)+FN_{2}P_{1}^{(2)}$ (8)
$\displaystyle=FN_{1}(P_{1}^{(1)}-1)+FP_{1}P_{1}^{(2)}$

The gradient is a sum of two terms: the first term directly encourages recall
and the second term regularizes precision, i.e., reduces false positive.

  * •

The first term $FN_{1}(P_{1}^{(1)}-1)$ consists of gradients from samples with
the ground truth class $1$. It incurs a larger negative gradient to the class
with larger false negatives ($FN$). The logit of $z_{1}$ will increase because
we subtract the gradient in gradient descent.

  * •

The second term $FP_{1}P_{1}^{(2)}$ consists of the non-ground truth gradient
contribution. It incurs a larger positive gradient to the class with larger
false positive ($FP$). Therefore, the logit of $z_{1}$ will decrease.

###  III-C Necessity Analysis: Recall, Precision, Dice, Jaccard and Tvesky

| $Recall(\mathcal{G}_{c},\mathcal{P}_{c})$ | $Precision(\mathcal{G}_{c},\mathcal{P}_{c})$ | $Dice(\mathcal{G}_{c},\mathcal{P}_{c})$ | $Jaccard(\mathcal{G}_{c},\mathcal{P}_{c})$ | $F1(\mathcal{G}_{c},\mathcal{P}_{c})$ | $Tversky(\mathcal{G}_{c},\mathcal{P}_{c})$
---|---|---|---|---|---|---
Set Rep. | $\frac{|\mathcal{G}_{c}\cap\mathcal{P}_{c}|}{|\mathcal{G}_{c}|}$ | $\frac{|\mathcal{G}_{c}\cap\mathcal{P}_{c}|}{|\mathcal{P}_{c}|}$ | $\frac{2|\mathcal{G}_{i}\cap\mathcal{P}_{c}|}{|\mathcal{P}_{c}|+|\mathcal{G}_{c}|}$ | $\frac{|\mathcal{G}_{c}\cap\mathcal{P}_{c}|}{|\mathcal{G}_{c}\cup\mathcal{P}_{c}|}$ | $\frac{|\mathcal{G}_{c}\cap\mathcal{P}_{c}|}{|\mathcal{G}_{c}\cup\mathcal{P}_{c}|+\frac{1}{2}|\mathcal{P}_{c}|+\frac{1}{2}|\mathcal{G}_{c}|}$ | $\frac{|\mathcal{G}_{c}\cap\mathcal{P}_{c}|}{|\mathcal{G}_{c}\cup\mathcal{P}_{c}|+\alpha|\mathcal{P}_{c}|+\beta|\mathcal{G}_{c}|}$
Boolean Rep. | $\frac{TP_{c}}{TP_{c}+FN_{c}}$ | $\frac{TP_{c}}{TP_{c}+FP_{c}}$ | $\frac{2TP_{c}}{2TP_{c}+FP_{c}+FN_{c}}$ | $\frac{TP_{c}}{TP_{c}+FP_{c}+FN_{c}}$ | $\frac{TP_{c}}{TP_{c}+\frac{1}{2}FP_{c}+\frac{1}{2}FN_{c}}$ | $\frac{TP_{c}}{TP_{c}+\alpha FP_{c}+\beta FN_{c}}$

TABLE I: Region Metrics: We show both set representation and Boolean
representation. TP, FN and FP stand for true positive, false negative and
false positive respectively. The subscript $c$ means that the metrics are
calculated for each class.

Why do we not use other metrics such as F1, Dice, Jaccard and Tvesky as the
weights? Following previous convention, let $\mathcal{G}_{c}$ and
$\mathcal{P}_{c}$ denote the set of ground truth (positive) samples and
predicted samples for the class $c$. Let $FP_{c}$, $TN_{c}$ denote the set of
false positive and true negative samples respectively for class $c$, and other
terms are defined similarly. Recall is different from the other metrics in
that it does not have false positive in the denominator and this distinction
makes it ideal for weighting cross entropy loss (and others not) as shown in
table I. Referring back to Eq. 5, where recall loss is defined as weighted
cross entropy by $1-\mathcal{R}_{c}$, replacing recall by any other metrics
above would result in FP appearing in the numerator of the weights. For
example, a hypothetical precision loss can be defined as the following.

$\displaystyle PrecisionLoss$
$\displaystyle-\sum_{c=1}^{C}\sum_{n:y_{i}=c}\left(\frac{FP_{c,t}}{FP_{c,t}+TP_{c,t}}\right)\log(p_{n,t})$
(9)

This formulation will introduce unexpected behavior. A large false positive
count in a class will result in a large weight, which further encourages false
detection for that class. This will cause the number of false positives to
explode.

##  IV Experiments

###  IV-A Experimental Setting

Datasets. We evaluate our recall loss on two popular large-scale outdoor
semantic segmentation datasets, Synthia [12] and Cityscapes [21]. Synthia is a
photorealistic synthetic dataset with different seasons, weather, and lighting
conditions. Specifically, we use the Synthia-sequence Summer split for
training (4400), validation (624), and testing (1272). Cityscapes consists of
real photos of urban street scenes in several cities in Europe. It has 5000
annotated images for training and another 5000 for evaluation.

Evaluation Metrics. We use mean accuracy and mean IoU as evaluation metrics,
which are commonly used in prior works on semantic segmentation [3] [4] [5].
We note that both mean accuracy and mean IoU are important metrics for
semantic segmentation. While a good mean IOU indicates a balanced performance
of precision and recall, mean accuracy is an indicator of the detection rate
of each class, which is important for safety-critical applications such as
self-driving.

Compared Methods. We compare with two region-based losses, SoftIOU [7] and
Lovasz softamax [13]. We also compare with statistics-balanced loss including
Weighted Cross-Entropy loss, also known as the Inverse-frequency (Weighted
CE), Balanced Cross-Entropy loss (Balanced CE) [9], and Online Hard Example
Mining (OHEM) loss [17]. We keep top $70\%$ samples in OHEM. We also compare
to performance-based losses, Focal loss [11], in experiments.

Implementation. We use DeepLabV3 [4] with resnet-$\\{18,101\\}$ [22] backbones
for semantic segmentation experiments. On Synthia, images are resized to 768
by 384. The resnet models are trained for 100 thousand iterations. On
Cityscapes, images are resized to 769 by 769, and models are trained for 90
thousand iterations following [21]. We use the Adam optimizer [23] with a
learning rate of $10^{-3}$ and $10^{-4}$ without annealing, respectively.

Figure 2: Mean accuracy vs. mean intersection over union (all-class average).
Left: Cityscapes dataset with Deeplab (Resnet18), Middle: Cityscapes dataset
with Deeplab (Resnet101), Right: Synthia dataset with Deeplab (Resnet101).
Only losses in the upper-right quadrant improve accuracy without introducing
excessive false positives.

###  IV-B Synthia

Figure 3: Visualization of Segmentation results on Synthia with Resnet18.
Recall loss encourages models to predict more small classes such as poles,
lights and pedestrians. Compared to the cross-entropy trained model, the
recall loss trained model is able to output finer details especially for small
classes. In contrast, Weighted CE yields excessive false positives on small
classes and degrades segmentation quality significantly. (Zoom in the
highlighted region for details).

We first present overall results on the synthetic Synthia segmentation dataset
[12] using DeepLab-ResNet 101 in Fig. 2. On the plot, the x-axis represents
mean IoU and the y-axis is mean accuracy. Weighted CE and Balanced CE improve
the mean accuracy considerably compared to the baseline Cross-Entropy loss.
However, the improvement comes at a cost of lower mean IoU. Focal loss
performs similarly to the standard cross entropy Cross-Entropy loss without
noticeable improvement [8, 9]. OHEM performs worse on both metrics. We think
this is because OHEM completely discards $30\%$ training samples at each
iteration, and this negatively affects feature learning. SoftIoU performs
reasonably well and improves both IoU and accuracy compared to the baseline.
Nevertheless, Recall loss achieves the highest mean recall besides the
weighted CE and Balanced CE with competitive IoU performance. This validates
our analysis that Recall loss can balance precision and recall.

| Accuracy | IoU
---|---|---
| pole | sign${\dagger}$ | pedestrian${\dagger}$ | bike${\dagger}$ | car${\dagger}$ | light${\dagger}$ | ave.${\dagger}$ | pole | sign${\dagger}$ | pedestrian${\dagger}$ | bike${\dagger}$ | car${\dagger}$ | light${\dagger}$ | ave.${\dagger}$
CE | 98.87 | 77.42 | 76.75 | 5.35 | 77.86 | 80.34 | 63.54 | 97.62 | 70.76 | 67.53 | 5.00 | 68.98 | 70.33 | 56.5
SoftIOU | 98.59 | 82.22 | 82.37 | 50.38 | 79.62 | 86.29 | 76.18 | 97.20 | 75.70 | 70.79 | 7.07 | 69.38 | 75.94 | 59.77
Balanced | 98.74 | 94.68 | 94.28 | 97.80 | 92.92 | 98.31 | 95.60 | 96.98 | 50.75 | 54.44 | 9.30 | 51.93 | 50.66 | 43.41
Focal | 98.70 | 78.91 | 79.70 | 20.14 | 74.42 | 80.09 | 66.65 | 97.48 | 71.18 | 67.53 | 14.89 | 67.41 | 69.42 | 58.09
Recall | 98.28 | 83.94 | 87.37 | 72.69 | 89.19 | 87.24 | 84.09 | 97.37 | 71.40 | 67.30 | 24.80 | 66.81 | 71.10 | 60.28

TABLE II: DeepLab-ResNet18 per-class accuracy and IoU performance on Synthia-
Summer (minority class). ${\dagger}$ denotes performance on classes of IoU
lower than $80$ with CE loss. Recall loss achieves the best IoU and the second
best accuracy. We highlight the Best and Second Best baseline method.

In Table II, we show per-class accuracy and IoU on selected classes using
DeepLab-ResNet18. Specifically, we show the results for minority classes with
original IoU lower than 80 using Cross-Entropy loss. We observe that Weighted
CE, Balanced CE and Recall loss can improve accuracy on small classes
significantly. However, they deteriorate IoU for those classes significantly.
The bike class is the smallest class in the dataset. It only occupies
$2*10^{-3}\%$ of pixels. Recall loss improves its accuracy from $5.35\%$ to
$72.69\%$ while simultaneously improving IoU from $5\%$ to $24.8\%$ compared
to the baseline Cross-Entropy loss. The other noticeable class is pedestrain.
Recall loss improves accuracy from $76.75\%$ to $87.37\%$ without sacrificing
IoU. In contrast, Balanced loss and Weighted loss drops IoU in exchange for
high accuracy. Note that while the pole class only occupies $1\%$ of pixels,
it is an easy class, therefore weighting is not necessary and can have a
negative effect.

This observation supports our claim that Recall loss balances recall and
precision because of its dynamic adaptability to performance. We further
provide a visual comparison between a model trained with the Cross-Entropy
loss, Weighted CE loss and the proposed recall loss in Fig. 3. Our method
provides more fine details on small classes which are often suppressed in
traditional cross entropy training. Recall loss also provides more precise
segmentation than the Weighted CE loss. The Weighted CE loss enlarges all
small classes and make them indistinguishable at distance.

###  IV-C Cityscapes

| Accuracy | IoU
---|---|---
| person | wall${\dagger}$ | fence${\dagger}$ | pole${\dagger}$ | light${\dagger}$ | terrain${\dagger}$ | rider${\dagger}$ | mo.cycle${\dagger}$ | ave.${\dagger}$ | person | wall${\dagger}$ | fence${\dagger}$ | pole${\dagger}$ | light${\dagger}$ | terrain${\dagger}$ | rider${\dagger}$ | mo.cycle${\dagger}$ | ave.${\dagger}$
CE | 82.98 | 69.60 | 65.63 | 58.93 | 76.76 | 62.87 | 70.16 | 67.01 | 67.63 | 71.31 | 55.23 | 55.88 | 45.75 | 57.54 | 57.14 | 53.49 | 55.77 | 55.23
SoftIOU | 85.00 | 61.23 | 68.11 | 57.70 | 79.15 | 72.53 | 71.51 | 68.57 | 68.03 | 72.93 | 46.70 | 51.92 | 45.36 | 60.41 | 57.41 | 54.39 | 44.91 | 46.34
Balanced | 90.07 | 61.43 | 81.13 | 71.12 | 93.36 | 83.01 | 77.90 | 84.43 | 80.88 | 66.48 | 45.46 | 45.82 | 40.16 | 36.58 | 48.66 | 47.62 | 38.63 | 29.58
Focal | 85.24 | 55.30 | 71.04 | 56.26 | 69.02 | 67.33 | 66.34 | 71.50 | 63.38 | 70.58 | 47.91 | 55.89 | 45.33 | 56.18 | 60.02 | 50.66 | 55.17 | 51.63
Recall | 88.31 | 72.16 | 75.55 | 69.17 | 80.00 | 70.12 | 75.26 | 83.04 | 75.06 | 69.06 | 57.92 | 52.50 | 44.67 | 54.64 | 57.44 | 51.14 | 41.11 | 53.36

TABLE III: DeepLab ResNet18 per-class accuracy and IoU performance on
Cityscapes (minority class). ${\dagger}$ denotes performance on classes of IoU
lower than $80$ with CE loss. Recall loss achieves the second best IoU and
accuracy. Note that balanced loss improves mean accuracy at the expense of the
worst IoU. We highlight the Best and Second Best method.

We further present results on the Cityscapes segmentation dataset [21] with
real images. Images in the dataset are real camera images and pixel-wise
labels are manually assigned. Therefore, there exists more data noise as
compared to the synthetic Synthia dataset. We show mean accuracy and mean IoU
in Fig. 2 for ResNet18 and Reset101, respectively. Weighted CE and Balanced CE
still yield the highest accuracy, but result in significantly worse IoU than
baseline Cross-Entropy. Surprisingly, SoftIoU also suffers in the presence of
data noise as its IoU does not match those of the Cross-Entropy loss, mostly
because of lower performance on minority classes. Compared to performance on
Synthia, this shows that some losses are not robust to label noise in a real
dataset. Recall loss again outperforms other losses by improving mean accuracy
and maintaining a good mean IOU. We further provide per-class performance on
selected losses in Table III using DeepLab-ResNet18. Recall loss improves mean
accuracy from $67.63\%$ to $75.06\%$ with minimum IoU drop on minority classes
compared to Cross-Entropy. In other words, recall loss improves the detection
rate of small classes such as pedestrians and light poles, while maintaining a
good precision. This demonstrates the effectiveness of Recall loss on both
synthetic and real outdoor segmentation datasets. No loss is perfect and
improves performance on all the datasets.

###  IV-D Analysis of Focal Loss

Figure 4: (a) Normalized focal weight over iterations. (b) Normalized recall
weight over iterations. (c) The ratio of recall weight and focal weight over
iterations. We present the three plots to argue that the focal loss can be
minimized by making the correct predictions more confident instead of
encouraging wrong predictions to become correct.

Focal loss[11] is another popular and important performance-balanced loss.
However, it’s not specifically successful in balancing for semantic
segmentation. We list Focal loss and Recall loss here for convenience.

$\displaystyle
RecallCE=-\sum_{c=1}^{C}\sum_{n:y_{i}=c}(1-\mathcal{R}_{c,t})\log(p_{n,t})$
$\displaystyle
FocalCE=-\sum_{c=1}^{C}\sum_{n:y_{n}=c}(1-p_{n,t})^{\gamma}\log(p_{n,t})$

where $y_{n}$ and $p_{n,t}$ denote the label of sample $n$ and predicted
probability (confidence) for the label at time $t$. ${R}_{c,t}$ denotes the
recall of class $c$ at time $t$. We use $\gamma=1$ in following discussion. We
will compare the average focal weight,
$\frac{1}{N_{c}}\sum_{n:y_{n}=c}(1-p_{n,t})$, with the recall
weight,$1-\mathcal{R}_{c,t}$, for each class. Because the weights are
independently calculated for each class in both losses, the weights across
class do not add up to one. We report normalized weights across class instead
because the relative percentage of each weight determines the focus of a loss
function. The absolute weight only scales the gradient. In fig 4, we report
the normalized focal weights, recall weights and recall-focal weight ratio
over time during training.

We have the following observations and follow-up discussion. 1) For Focal
loss, the average bike class weight increases relatively over time, whereas
the bike class weight percentage decreases for Recall loss. This means that
the focal weight $\frac{1}{N_{c}}\sum_{n:y_{n}=c}(1-p_{n,t})$ for the bike
class does not decrease as fast as other classes. In other words, the
confidence for the bike class does not increase relatively to others. The
recall-focal weight ratio plot tells the same story. 2) We observe that the
ratios of large classes are mostly larger than one and increasing, while
ratios of small classes are below one and constant. This means that Focal loss
tends to assign lower and decreasing weights, which means higher and
increasing confidence, to large classes. Colloquially, Focal loss finds it
easier to increase the confidence of large classes to reduce loss than
increase the confidence of small classes. This limits its ability to correct
wrong predictions from small classes. However, because Recall loss uses the
metric recall instead of predicted probability as the weights, there is no
benefit to continue to increase the confidence of a sample once it is already
a true positive. For Recall loss, the only way to further reduce loss is by
encouraging correct predictions from small classes. The decreasing recall
weight for the bike class in fig. 4 shows that the performance of the bike
class increases over time.

##  V conclusion

In this paper, we study the less-explored area of imbalance in semantic
segmentation and discuss the trade-off between precision and recall,
especially its asymmetric effect on large and small classes. We discovered
that minority classes are more sensitive to rebalancing losses and require
careful approaches to avoid excessive false positives. We propose a novel loss
function based on the metric recall. The loss function uses a hard-class
mining strategy to improve model performance on imbalanced datasets.
Specifically, Recall loss weights examples in a class based on its
instantaneous recall performance during training, and the weights change
dynamically to reflect relative change in performance among classes.
Experimentally, we demonstrate several advantages of the loss: 1) Recall loss
improves accuracy while maintaining a competitive IoU performance. Most
notably, Recall loss improves recall of minority classes without introducing
excessive false positives, striking the right balance between recall and
precision. 2) Recall loss avoids excessive weighting to minority but easy
classes.

##  VI Acknowledgement

This material is based on work supported by the Collaborative GPS-Denied
Navigation for Combat Vehicles Program under Contract W9132V19C0003. We would
like to thank Garry P. Glaspell, Delma B. Del Bosque, and Jean D. Nelson for
their valuable feedback on the project. This work is also partially supported
by ONR grant N00014-18-1-2829. Any opinions, findings, and conclusions or
recommendations expressed in this material are those of the author(s) and do
not necessarily reflect the views of the US government.

## References

  * [1] J. Tian, Y.-C. Liu, N. Glaser, Y.-C. Hsu, and Z. Kira, “Posterior re-calibration for imbalanced datasets,” _arXiv preprint arXiv:2010.11820_ , 2020.
  * [2] J. M. Johnson and T. M. Khoshgoftaar, “Survey on deep learning with class imbalance,” _Journal of Big Data_ , vol. 6, no. 1, p. 27, 2019.
  * [3] L.-C. Chen, G. Papandreou, I. Kokkinos, K. Murphy, and A. L. Yuille, “Semantic image segmentation with deep convolutional nets and fully connected crfs,” _arXiv preprint arXiv:1412.7062_ , 2014.
  * [4] L.-C. Chen, G. Papandreou, F. Schroff, and H. Adam, “Rethinking atrous convolution for semantic image segmentation,” _arXiv preprint arXiv:1706.05587_ , 2017.
  * [5] J. Jiang, Z. Zhang, Y. Huang, and L. Zheng, “Incorporating depth into both cnn and crf for indoor semantic segmentation,” in _2017 8th IEEE International Conference on Software Engineering and Service Science (ICSESS)_. IEEE, 2017, pp. 525–530.
  * [6] M. Grandini, E. Bagli, and G. Visani, “Metrics for multi-class classification: an overview,” _arXiv preprint arXiv:2008.05756_ , 2020.
  * [7] M. A. Rahman and Y. Wang, “Optimizing intersection-over-union in deep neural networks for image segmentation,” in _International symposium on visual computing_. Springer, 2016, pp. 234–244.
  * [8] K. Cao, C. Wei, A. Gaidon, N. Arechiga, and T. Ma, “Learning imbalanced datasets with label-distribution-aware margin loss,” in _Advances in Neural Information Processing Systems_ , 2019, pp. 1565–1576.
  * [9] Y. Cui, M. Jia, T.-Y. Lin, Y. Song, and S. Belongie, “Class-balanced loss based on effective number of samples,” in _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , 2019, pp. 9268–9277.
  * [10] B. Zhou, Q. Cui, X.-S. Wei, and Z.-M. Chen, “Bbn: Bilateral-branch network with cumulative learning for long-tailed visual recognition,” _CVPR_ , 2020\.
  * [11] T.-Y. Lin, P. Goyal, R. Girshick, K. He, and P. Dollár, “Focal loss for dense object detection,” in _Proceedings of the IEEE international conference on computer vision_ , 2017, pp. 2980–2988.
  * [12] G. Ros, L. Sellart, J. Materzynska, D. Vazquez, and A. M. Lopez, “The synthia dataset: A large collection of synthetic images for semantic segmentation of urban scenes,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2016, pp. 3234–3243.
  * [13] M. Berman, A. Rannen Triki, and M. B. Blaschko, “The lovász-softmax loss: A tractable surrogate for the optimization of the intersection-over-union measure in neural networks,” in _Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition_ , 2018, pp. 4413–4421.
  * [14] C. H. Sudre, W. Li, T. Vercauteren, S. Ourselin, and M. J. Cardoso, “Generalised dice overlap as a deep learning loss function for highly unbalanced segmentations,” in _Deep learning in medical image analysis and multimodal learning for clinical decision support_. Springer, 2017, pp. 240–248.
  * [15] S. S. M. Salehi, D. Erdogmus, and A. Gholipour, “Tversky loss function for image segmentation using 3d fully convolutional deep networks,” in _International Workshop on Machine Learning in Medical Imaging_. Springer, 2017, pp. 379–387.
  * [16] S. A. Taghanaki, Y. Zheng, S. K. Zhou, B. Georgescu, P. Sharma, D. Xu, D. Comaniciu, and G. Hamarneh, “Combo loss: Handling input and output imbalance in multi-organ segmentation,” _Computerized Medical Imaging and Graphics_ , vol. 75, pp. 24–33, 2019.
  * [17] A. Shrivastava, A. Gupta, and R. Girshick, “Training region-based object detectors with online hard example mining,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2016, pp. 761–769.
  * [18] D. Eigen and R. Fergus, “Predicting depth, surface normals and semantic labels with a common multi-scale convolutional architecture,” in _Proceedings of the IEEE international conference on computer vision_ , 2015, pp. 2650–2658.
  * [19] V. Badrinarayanan, A. Kendall, and R. Cipolla, “Segnet: A deep convolutional encoder-decoder architecture for image segmentation,” _IEEE transactions on pattern analysis and machine intelligence_ , vol. 39, no. 12, pp. 2481–2495, 2017.
  * [20] R. Chan, M. Rottmann, F. Hüger, P. Schlicht, and H. Gottschalk, “Application of decision rules for handling class imbalance in semantic segmentation,” _arXiv preprint arXiv:1901.08394_ , 2019.
  * [21] M. Cordts, M. Omran, S. Ramos, T. Rehfeld, M. Enzweiler, R. Benenson, U. Franke, S. Roth, and B. Schiele, “The cityscapes dataset for semantic urban scene understanding,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2016, pp. 3213–3223.
  * [22] K. He, X. Zhang, S. Ren, and J. Sun, “Deep residual learning for image recognition,” in _Proceedings of the IEEE conference on computer vision and pattern recognition_ , 2016, pp. 770–778.
  * [23] D. P. Kingma and J. Ba, “Adam: A method for stochastic optimization,” _arXiv preprint arXiv:1412.6980_ , 2014.