#  Braxlines: Fast and Interactive Toolkit for
RL-driven Behavior Engineering
beyond Reward Maximization

Shixiang Shane Gu
<EMAIL_ADDRESS>
&Manfred Diaz
<EMAIL_ADDRESS>
Daniel C. Freeman
<EMAIL_ADDRESS>
&Hiroki Furuta
<EMAIL_ADDRESS>
&Seyed Kamyar Seyed Ghasemipour
<EMAIL_ADDRESS>
&Anton Raichuk
<EMAIL_ADDRESS>
&Byron David
<EMAIL_ADDRESS>
Erik Frey
<EMAIL_ADDRESS>
&Erwin Coumans
<EMAIL_ADDRESS>
&Olivier Bachem
<EMAIL_ADDRESS>

###### Abstract

The goal of continuous control is to synthesize desired behaviors. In
reinforcement learning (RL)-driven approaches, this is often accomplished
through careful task reward engineering for efficient exploration and running
an off-the-shelf RL algorithm. While reward maximization is at the core of RL,
reward engineering is not the only – sometimes nor the easiest – way for
specifying complex behaviors. In this paper, we introduce Braxlines, a toolkit
for fast and interactive RL-driven behavior generation beyond simple reward
maximization that includes Composer, a programmatic API for generating
continuous control environments, and set of stable and well-tested baselines
for two families of algorithms –mutual information maximization (MI-max) and
divergence minimization (D-min)– supporting unsupervised skill learning and
distribution sketching as other modes of behavior specification. In addition,
we discuss how to standardize metrics for evaluating these algorithms, which
can no longer rely on simple reward maximization. Our implementations build on
a hardware-accelerated Brax simulator in Jax with minimal modifications,
enabling behavior synthesis within minutes of training. We hope Braxlines can
serve as an interactive toolkit for rapid creation and testing of environments
and behaviors, empowering explosions of future benchmark designs and new modes
of RL-driven behavior generation and their algorithmic research.

##  1 Introduction

In the current era of neural networks and high-performance computing, large
and well-engineered benchmarks [26] have enabled key architectural and
algorithmic breakthroughs, in computer vision [82, 78, 26, 79, 58, 59, 114],
natural language processing [137, 12] and molecular biology [67]. When scaled
to the extreme levels of data and computing, these innovations can train
models to master an impressive range of capabilities [27, 108, 16, 109, 111,
18, 13] exhibiting developer-aware generalization [21], where the trained
models can function even on datasets outside the developer’s imagination.
Undoubtedly, the massive success of supervised learning (SL) has been largely
due to large, diverse, and high-quality datasets, that are relatively simple
to collect. In this paradigm, data generation essentially comes down to
finding paired data: speech with texts [49], images with texts [26, 138, 6,
111, 109], words with words [90, 137, 27, 108, 16, 18], pixels with pixels
[73, 48, 116, 134, 74]. Moreover, a significant amount of this data can be
mined automatically through the web [27, 16, 109, 18] or collected relatively
effortlessly with minor bottlenecks like manual human labeling [24] or slow
simulation [67].

Figure 1:  Braxlines vision: from behavior _engineering_ to _distillation_.
Our _behavior engineering_ toolkit starts from Braxlines Composer, a
programmatic API that facilitates environment engineering and will accelerate
the proliferation of novel environments for continuous control. Components in
green represent multiple modes of RL-driven behavior engineering. Red
components require human inputs and, except for environment engineering,
require different modes of specifications from the designers (see Table 2).
Importantly, all behaviors from these modes can be aggregated and fed to
downstream behavior distillation modules. We believe that maximizing
simulation throughput and minimizing human inputs as the keys to scalable
behavior generation, where in this work we focus on enabling MI-max and D-min
as new modes of behavior engineering.

To mirror this success in RL, we may need efficient strategies for data
generation. Data generation in RL [130, 95, 89, 52] is significantly more
costly (than in SL), the tasks design space is infinite, and in real-world
environments like robotics, healthcare, or dialog modeling [52, 70, 98, 65]
data is scarse. Even simulated environments [10, 14, 131, 143, 23, 139, 121,
19, 31] are arguably far from having sufficient diversity, learnability,
regularity, and downstream applicability to come close to enabling the levels
of architectural breakthrough we have witnessed in SL.

Today, there exist data-centric techniques in RL such as policy distillation
[117, 62, 118, 84], hindsight imitation learning or RL [85, 92, 4, 105, 70,
66, 91, 17] and offline RL [86, 42, 80, 40, 47, 128, 93] that aim to learn
faster by better utilizing prior policies or experiences. Behavior
distillation techniques offer an important alternative to popular approaches
[10, 14, 131, 95, 124, 125, 51, 54, 41] that often require to start _tabula
rasa_ and rediscover meaningful behaviors over and over again. However, we
argue that the question of how to obtain interesting and useful data for
distillation methods is under-explored. Prior work [39] largely focuses on
collecting human trajectories or obtaining trajectories from policies trained
with engineered rewards. Both approaches are, practically and computationally,
prohibitively expensive. In this paper, we focus on alternative methods to
generate diverse and useful behaviors.

We introduce Braxlines, an interactive toolkit for principled behavior
creation for continuous control tasks (see Figure 1) and claim the following
key contributions:

  * •

Braxlines introduces a novel programmatic API for environment composition,
Braxlines Composer, that enables RL researchers to create continuous control
environments from scratch or pre-defined components.

  * •

Braxlines provides clean, minimalistic implementations for two broad families
of algorithms suitable for reward-free behavior engineering: mutual
information maximization (MI-max) and divergence minimization (D-min). They
provide a rich set of alternative, sometimes complementary, algorithms for
behavior engineering, including goal-conditioned RL [69, 122, 4, 105],
empowerment maximization [75, 34, 20], state marginal matching and adversarial
inverse RL [63, 38, 46, 83].

  * •

The accompanying toolkit provides absolute and stationary metrics for
evaluating MI-max and D-min algorithms, whose reward functions dynamically
varies over training and across different implementation choices [75, 34, 20].
Some of these metrics essentially measure characteristics of behaviors
themselves, such as how much control of degrees of freedom of the world an
agent has, in a task-agnostic sense, providing more intrinsic insights into
agent-environment interactions [136, 15].

  * •

The open-source implementation based on the Brax [37] physics library is
extremely fast, allowing interactive creation of novel behaviors in a few
minutes through Google Colab on a freely available 2$\times$2 with 8 cores TPU
(see Table 1). This effectively enables real-time interactive RL-driven
environment and behavior designs.

Library | Run-time | Steps | Resource | Simulator | Backends
---|---|---|---|---|---
Baselines [29, 144] | 3.5 hours | $4.0\times 10^{7}$ | 16 CPU + 1 GPU | MuJoCo | TensorFlow, OpenMPI
Braxlines | 3.5 min | $2.0\times 10^{8}$ | 2$\times$2 TPU (free) | Brax | Jax
Table 1: Run-time speeds for DIAYN [34] on _Ant_ with different library and
simulators. We adopt a recent open-source DIAYN implementation released by
Zhao et al. [144] built on Baselines [29].

##  2 Related Work

Benchmarks and Baselines. Numerous benchmarks proposed over the years have
spanned a wide range of tasks including simulated games of varying
complexities [10, 139], continuous control and robotics [133, 131, 143, 121,
31], or procedurally generated environments [23, 132]. Braxlines differentiate
by focusing on improving the roundtrip time from conception to experimental
results, increasing the interactivity of RL research. We leverage Brax [37]
highly-paralellizable simulations to significantly reduce both time and
computational requirements for experimentation in RL. Furthermore, using
Braxlines researchers can fix well-tested algorithmic baselines and iterate
quickly over environment designs. Also, while many existing libraries [32,
105, 100, 29, 110, 88] provide fine-tuned collections of common deep RL [95]
algorithms, these works focus mostly on optimizing painstakingly hand-designed
rewards (see Table 5 in Appendix A) using algorithms of varying complexity,
from basic [141] to more advanced techniques [89, 124, 125, 41, 54]. In
contrast, Braxlines provides generalized implementations for mutual
information maximization (MI-max) [20] and divergence minimization (D-min)
algorithms [46]. While these approaches have been studied extensively [97, 50,
3, 34, 127], few benchmark implementations are available; often in isolated
and not well-maintained repositories. To the best of our knowledge, we are the
first to provide fast, testable, reproducible and minimalist baselines in this
area (see Table 6 in Appendix A).

Reward-free Behavior Engineering. While arguably rewards are enough to produce
intelligent behavior [129], reward engineering [28] bottlenecks the
application of RL. Prior work have sought to address this issue by introducing
other objectives such as human preferences elicitation [55, 22] or task-
agnostic reward functions [123, 103, 9, 104]. On the latter, some recent works
have emphasized empowerment [75, 119, 68, 97] and skills/options discovery
[50, 36, 34, 140, 56, 126] through MI-max objectives. Choi et al. [20]
connects these two areas with a _variational_ approach to goal-conditioned RL
[69, 4, 105, 106, 140]. Section 3.1 discusses how Braxlines leverages this
variational approach to implement MI-max algorithms. On the other hand, an
alternative to generating reward-free behaviours is to learn from human
demonstrations [7]. The range of techniques has spanned from feature
expectation matching [101, 1], maximum entropy formulations [145, 142], to
adversarial approaches to generative modelling [63, 38, 35]. Recently,
unifying frameworks [46, 71] emerged under a probabilistic view, based on
algorithms the minimization of _f_ -divergences [115, 25] between state-action
and state marginal distributions. Braxlines leverages, as we discuss in
Section 3.2, Ghasemipour et al. [46] to implement algorithms of the D-min
family.

Metrics for Reward-Free Behavior. One limitation of reward-free approaches to
RL has been the lack of quantitative evaluation metrics. Previous work in both
MI-max [34, 3, 127, 126] and D-min [46, 71, 63, 38, 35] families have relied
on either qualitative analysis of the learned behavior, the value of the
objective (which usually saturates), and the performance on existing reward
functions for downstream tasks. In the case of MI-max approaches, as mutual
information is difficult to scale to high-dimensional spaces [107, 8, 135], it
seems particularly hard to come up with quantitative metrics. Braxlines
toolkit provides a non-parametric, particle-based approximation to mutual
information similar to PIC metrics [45]. We also leverage this approach to
compute metrics for D-min algorithms. Finally, we provide an implementation of
Laten Goal Reaching (LGR) [20] that leverages the connection between goal-
conditioned RL and MI-max objectives. We discuss this metrics further in
Section 4 and contrast with prior evaluation techniques in Appendix C and
Table 7.

Comparison to Other Benchmarks We surveyed the support of common RL algorithms
in well-maintained _open-source_ repositories of baselines: Baselines [29],
Stable Baselines [110], RLLab [32], Tf-Agents [53], RLLib [88], and PFRL /
ChainerRL [43]. For each, we verified the support for continuous control,
reward maximization algorithms like DDPG [89], A3C [96], PPO [125], TRPO
[124], TD3 [41], SAC [54] and ES [120]. Similarly, we investigated whether
existing baselines provide support for algorithms of the MI-max and D-min
families. With the exception of RLLab and OpenAI Baselines, there were no
other benchmark implementation for MI-max GCRL [69], DIAYN [34], DISCERN
[140], VISR [56] or D-min algorithms [63, 46, 71]. Table 5 and 6 in Appendix A
show the findings of our survey. Braxlines is, to the best of our knowledge,
the first reproducible and reusable open-source library of baseline
implementations for reward-free RL algorithms.

##  3 Background

Notations We consider a Markov Decision Process (MDP) to be a tuple
characterized by $(\mathcal{S},\mathcal{A},p,r,\rho_{0},\gamma)$; state space
$\mathcal{S}$, action space $\mathcal{A}$, transition probability
$p:\mathcal{S}\times\mathcal{A}\rightarrow\mathcal{S}$, reward function
$r:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$, initial state
distribution $\rho_{0}$, and discount factor $\gamma\in(0,1]$. A policy is a
function $\pi:\mathcal{S}\rightarrow\mathcal{A}$ that defines behaviors. In
this work, we denote the state-action and marginal state distribution of a
policy as $\rho^{\pi}(s,a)$ and $\rho^{\pi}(s)$, respectively.

###  3.1 Mutual Information Maximization (MI-max)

Reward-free RL typically considers the maximization of mutual information
between actions and future states. Objectives such as empowerment [75, 97] or
unsupervised skills discovery [50, 34, 3, 140] propose the maximization of
mutual information between states and abstracted actions $z\sim p(z)$
(commonly known as _skills_ [112]). Algorithms in this family learn policies
$\pi(a|s,z)$ conditioned on this latent actions. Recently, Choi et al. [20]
presented Variational Goal-Conditioned RL (V-GCRL), a framework that unifies
Goal-Conditioned RL (GCRL) [69] and a variational approach to MI-max, derived
from the mutual information maximization objective:

$\displaystyle\max_{\pi}\text{MI}(z,s;\pi)$
$\displaystyle=\max_{\pi}H(s;\pi)-H(s|z;\pi)$ (1)
$\displaystyle=\max_{\pi}H(z)-H(z|s;\pi)$ $\displaystyle=\max_{\pi}\iint
p(z)\rho^{\pi}(s|z)\log p^{\pi}(z|s)-\int p(z)\log p(z)$ (2)

where $MI(z,s;\pi)$ is the mutual information between the distribution of
latent actions or skills $z\sim p(z)$ and the state marginal distribution
$s\sim\rho^{\pi}(s)$, under $z$-conditioned policy $\pi(a|s,z)$. The objective
in Eq. 1 is a generative formulation while Eq. 2 defines a discriminative of
the MI-max objective. Due to the intractability of the true posterior
$p^{\pi}(z|s)$, a variational bound of Eq. 2 is used instead:

$\displaystyle\max_{\pi}\text{MI}(z,s;\pi)$ $\displaystyle\geq\max_{\pi}\iint
p(z)\rho^{\pi}(s|z)\log q(z|s)-\int p(z)\log p(z)$ (3)

that enables most of the algorithms in the MI-max family [69, 87, 34, 140, 56,
127].

###  3.2 Divergence Minimization (D-min)

Inverse RL (IRL) [101, 1] aims to infer the reward function and subsequently
obtain an optimal policy from an expert behavior. Seminal work [63] has
demonstrated that the problem of MaxEnt IRL is equivalent to matching the
state-action marginal of a policy, $\rho^{\pi}(s,a)$, to that of the expert’s,
$\rho^{\textnormal{exp}}(s,a)$. Thus, distribution-matching GAN [48]
techniques can be used for IRL with little expert data, leading to Adversarial
Imitation Learning (AIL) of methods [63, 38]. Recently, the work in [46, 71]
have shown that to correspond to different choice of $f$-divergences for
state-action distribution matching.

An interesting modification studied in [46, 83, 57] matches $\rho^{\pi}(s)$ to
a desired distribution $\rho^{\textnormal{target}}(s)$ and allows AIL to match
behavior not necessarily generated by an expert policy, and can even be
designed without any access to the underlying MDP. For a given choice of
$f$-divergence, minimizing $D_{f}$ of state marginals can be accomplished
through an adversarial optimization setup [46].

$\displaystyle\min_{\pi}D_{f}(\rho^{\pi}(s),\rho^{\textnormal{target}}(s))=\max_{\pi}\mathds{E}_{\rho^{\pi}(s)}\Big{[}f^{*}(T_{\omega}(s))\Big{]}=\max_{\pi}\mathds{E}_{\rho^{\pi}(s)}\Big{[}r^{*}(s)\Big{]}$
(4)

recovering optimization objectives similar to GAIL[63] (Jensen-Shannon
divergence), AIRL[38] (reverse KL divergence), and F-AIRL[46] (forward KL
divergence).

In these settings, the discriminator is a binary classifier distinguishing
between target states and states visited by the policy, trained using cross
entropy. The reward functions for training the policy are obtained from:

$\displaystyle r^{\text{{GAIL}}}(s)$ $\displaystyle:=\text{log
}D(s)=\log\rho^{\textnormal{target}}(s)-\log\left(\rho^{\textnormal{target}}(s)+\rho^{\pi}(s)\right)$
(5) $\displaystyle h(s)$ $\displaystyle:=\text{log }D(s)-\text{log }(1-D(s))$
$\displaystyle r^{\text{{AIRL}}}(s)$
$\displaystyle:=h(s)=\log\rho^{\textnormal{target}}(s)-\log\rho^{\pi}(s)$ (6)
$\displaystyle r^{\text{{MLE}}}(s)$
$\displaystyle:=\log\rho^{\textnormal{target}}(s).$ (7)

where $D(s)$ is the probability the discriminator assigns for $s$ being from
the target distribution. For a more in depth discussion, we refer the reader
to [46, 71]111The last reward Eq. 7 is technically not an instance of D-min
algorithms and is a plain stationary reward function, but it is effectively
AIRL(Eq. 6) without $H^{\pi}(s)$ term. In Figure 4, we show that sometimes
this presents competitive results for matching multi-modal distribution as a
full D-min algorithm..

##  4 Tools for Behavior Engineering

Braxlines for behavior engineering is built over three pillars: (1) an
environments composer, (2) a set of stable baselines, and (3) a collection of
metrics for reward-free behavior evaluation.

###  4.1 Braxlines Composer

Composer is designed for modularity and reusability. Figure 2 illustrates
examples of environments that can be composed easily using Composer. See
Appendix D for an Ant Push task example constructed under 50 lines. Since all
observations, reward functions, and scene components can be combined and
reused, Composer allows programmatic procedural generation of parameterized
environments. Combined with Brax’s PPO/ES methods and Braxlines’ MI-max/D-min
implementations that allow training within minutes, these enable designers to
quickly design, debug, and tune tasks.

Figure 2: Braxlines Composer API is a flexible and fast environment
engineering tool for continuous control. The space of environments designs
spans from (left) goal-based locomotion + manipulation task, (middle)
parameterized morphologies, to (right) multi-agent systems. See Appendix D for
an example on how to construct Ant Push (left) in under 50 lines of code.

###  4.2 Stable Baselines for Reward-free Behavior Generation

Braxlines is designed to optimize for speed and minimalism of algorithms in
the MI-max and D-min families. Prior implementations [34, 63, 38, 46, 83, 99]
rely on heavily modifying different RL optimizers (e.g. SAC [54], TD3 [41],
PPO [125], MPO [2]), making comparative analyses and code reuses difficult.
The benchmark implementations we provide are unified under Brax’s PPO [37],
enabling stable and speedy training on medium-size problems within minutes.
Table 1 shows Braxlines speed gains of orders of magnitude compared to
existing DIAYN [34] baselines. Furtheremore, we leverage recent unification
papers [46, 20] to keep Braxlines design as clean and _minimal_ as possible. A
minimalistic approach to RL [40] ensures that codebases does not rely on
extensive hyperparameter designs [60, 5, 33, 44], and can be debugged,
reproduced, extended, and mixed easily downstream.

Algorithm | Family | Implementation | User Specification
---|---|---|---
GCRL [69, 4, 105, 20] | MI-max | Fix $q(z|s)=\mathcal{N}(o(s),\sigma I)$, $p(z)$ \+ Offset |  $p(z)$, $o(\cdot)$
V-GCRL [20, 34, 140, 56] | MI-max | Param $q_{\theta}(z|o(s))$\+ Fix $p(z)$ \+ SN + Offset | $o(\cdot)$
GAIL [63] | D-min | Eq. 5 \+ Offset |  $o(\cdot)$, $\rho^{\textnormal{target}}(o(s))$
AIRL [38] | D-min | Eq. 6 \+ Offset |  $o(\cdot)$, $\rho^{\textnormal{target}}(o(s))$
Table 2: A summary of example Braxlines algorithms explained in terms of the
Algorithm Family/Implementation decomposition in [44]. Braxlines emphasizes on
minimal implementations with little to no code-level optimization (only
spectral normalization (“SN”) [94, 20] on discriminator $q_{\theta}(z|\cdot)$)
and adding positive constant offsets (“Offset”), equivalent survival bonus, to
ensure that rewards almost always positive (see Appendix A.1). User
Specification lists what Braxlines users are required to define to perform
feature engineering $o(\cdot)$, and marginal state distribution matching
$\rho^{\textnormal{target}}(\cdot)$. Importantly, since V-GCRL has
representation learning capacity, $p(z)$ can often be simply a fixed
uninformative prior, e.g. zero-mean Gaussian or a uniform Categorial [34, 20].

We present in Table 2 an analysis of the flexibility of Braxlines using an
_algorithm-implementation_ decomposition similar to Furuta et al. [44] that
separates algorithm (i.e., mathematical and algorithmic choices) from
implementation details (i.e., implementation and code-level optimizations). In
addition, we define what are users required to specify to turn a family of
algorithms into a concrete instance (e.g., MI-max to DIAYN). In the sections
that follow, we present more detailed explanations of both MI-max and D-min
families, and how are they supported in Braxlines. We support flexible
parametrizations of $q(z|s)$ that could take the form of an isotropic Gaussian
distribution $q(z|s)=\mathcal{N}(O;\mu,\sigma^{2}I)$,or a parametric tractable
posterior $q_{\theta}(z|s)$ learned using function approximation (i.e., neural
networks). Similarly, Braxlines provides baseline implementations for
variations of the AIL family of algorithms. In the case of D-min algorithms,
we provide the flexibility for users to specify feature engineering functions
as well as behavior distributions over this feature space. Together,
algorithms in both families are a key ingredient for generating diverse
behaviors through little human intervention.

###  4.3 Metrics

Both MI-max and D-min algorithms infer dynamic reward functions during
learning (except in few special cases like GCRL), which can change arbitrarily
across training iterations, as well as across different implementation and
hyperparameter choices (e.g. discrete or continuous $z$-space in MI-max, or
different choices of reward transformation in D-min as in Table 2). This makes
episodic rewards inapplicable for quantitatively evaluating optimization
convergence or final performances. Common alternatives are also limited and
hard to automate: MI-max algorithms often use qualitative visual inspections
or down-stream task performances [34, 127, 126], while D-min algorithms mostly
focus on imitation learning and frequently use task rewards directly to
compare how close to the experts they got [63, 38, 46].

Inspired by prior works [20], we provide objective metrics for MI-max and
D-min, that (1) provide an intuitive and interpretable absolute units of
performance measure applicable to any algorithm instance within the same
family, and (2) are stationary throughout learning, meaning that they do not
depend of quantities that change over the course of training or across runs of
an experiment (such as rewards given by the discriminator).

####  4.3.1 MI-max Metrics

We implemented two metrics for evaluating the behaviors generated by
algorithms in the MI-max family that leverage approximations to mutual
information. These metrics compute, in general, how much an agent learned to
control each dimension of the environment.

##### Particle-based Mutual Information Approximation.

We can directly adopt similar techniques in Furuta et al. [45] to estimate
empowerment of a trained policy. It uses straight-forward discretization for
tractable non-parametric estimation in 1D and 2D. Intuitively this metric
quantifies how much predictive control over each dimension(s) of the
environment an agent has learned, a direct measure of the original MI-max
objective in Eq. 1 (without variational approximation). Note that reliable
mutual information estimation in high dimensions is still an actively
researched problem, but some scalable estimation techniques [107, 8, 135] will
be added later.

Input: agent $\pi(a|s,z)$; intent $p(z)$; feature function $o(s)$; bin size
$B$; bin range $(a,b)$

Output: Mutual information estimate $\hat{\text{MI}}$

for  _$n=1,...,N$_ do

Sample $z_{n}\sim p(z)$ // Sample an agent intent

Sample $s_{n,m,t}\sim\pi(\cdot|\cdot,z_{n})$ // Collect $M$ episodes of
horizon $T$ for $z_{n}$

// Estimate conditional entropy in $o$-space for $z_{n}$ using
np.histogram(2d)($B,(a,b)$)

Estimate $h_{n}$ with
$p(o|z_{n})\approx\frac{1}{MT}\sum_{m,t}\delta(o=o(s_{n,m,t}))$

end for

// Estimate marginal entropy in $o$-space for $z_{n}$ using
np.histogram(2d)($B,(a,b)$)

Estimate $h$ with $p(o)\approx\frac{1}{NMT}\sum_{n,m,t}\delta(o=o(s_{n,m,t}))$

$\hat{\text{MI}}=h-\frac{1}{N}(\sum_{n}h_{n})$ // Compute MI estimate

Algorithm 1 Particle-based Mutual Information Approximation

Latent Goal Reaching. In our experiments, we fixed $N=1$ for Algorithm 3 and
use the deterministic sampling (for Gaussians, take the mean; for
Categoricals, take the argmax). LGR for GCRL in Section 3.1 corresponds
exactly to the standard goal-reaching reward evaluation. While the LGR metric
in Algorithm 3 (see Appendix B) is motivated intuitively as a task-oriented
metric from GCRL [20] and appears unrelated to the MI-max objective (Eq. 2) or
its approximation (Algorithm 1), it can also be viewed as a crude
approximation to MI, as we show in Appendix B.1.

####  4.3.2 D-min Metrics

Particle-based Divergence Approximation. For lower dimension cases, we use the
energy distance for non-parametric estimation. Given two distributions $p$ and
$q$, the energy distance is defined as,

$\displaystyle d(p,q):=2\mathds{E}_{x\sim p,y\sim
q}||x-y||-\mathds{E}_{x,x^{\prime}\sim
p}||x-x^{\prime}||-\mathds{E}_{y,y^{\prime}\sim q}||y-y^{\prime}||$

where $||\cdot||$ is the Euclidean norm but may be any desired metric. For
higher dimensional and more complex distributions, adhoc notions of distance
(such as Frechet Inception Distance [61]) have been proposed.

Input: agent $\pi(a|s)$; $o(s)$; target samples
$o^{\textnormal{target}}_{1:L}=o(s^{\textnormal{target}}_{1:L})\sim
p^{\textnormal{target}}(s)$

Output: Energy distance estimate $\hat{\text{D}}$

Sample $o_{m,t}=o(s_{m,t}),s_{m,t}\sim\pi$ in $\mu$ // Collect $M$ episodes of
horizon $T$

// Compute energy distance

$\hat{\text{D}}=\frac{2}{MTL}\sum_{l,m,t}||o^{\textnormal{target}}_{l}-o_{m,t}||-\frac{1}{M^{2}T^{2}}\sum_{m,t,m^{\prime},t^{\prime}}||o_{m,t}-o_{m^{\prime},t^{\prime}}||-\frac{1}{L^{2}}\sum_{l,l^{\prime}}||o^{\textnormal{target}}_{l}-o^{\textnormal{target}}_{l^{\prime}}||$

Algorithm 2 Energy Distance Figure 3:  MI-max results on Ant averaged over 10
seeds: (from left to right) (a) episodic discriminator reward, (b, c) MI and
$H(s)$ estimates from Algorithm 1222XY velocities correspond to dimensions
$(13,14)$ of Ant env. Since each dimension had similar MI estimates, we only
show for MI and entropies for dim$=(13,)$. See Appendix A.1 for more
discussions., and (d) LGR from Algorithm 3. Since the discriminator reward
scale (a) can vary arbitrarily among different instances of MI-max algorithms,
it is not useful for comparisons. Conversely, (b-d) metrics have the same
units for any MI-max algorithm, allowing comparisons.

##  5 Experiments

We provide both illustrative Google Colab examples
https://github.com/google/brax/tree/main/notebooks/braxlines for interactive
environment composition and training, and a set of quantitative benchmark
results that could accelerate future algorithmic research in MI-max and D-min
RL. Importantly, due to two orders of training speedups as measured in Table
1, a medium-size environment can be composed and trained in a few minutes333If
you enable dense evaluations of the quantitative metrics in Section 4, it
could slow down to 10-20 minutes, but we mainly recommend these during initial
debugging or full quantitative benchmarking. on Google Colab with 2$\times$2,
8 cores TPU, enabling interactive environment designs and algorithm
deployments. All Braxlines codes, documentation, result videos, and Colab
examples are accessible at
https://github.com/google/brax/tree/main/brax/experimental/braxlines.
Experimental details and hyperparameters, additional benchmark results, and
ablation studies are available in Appendix A.

Environment |  MI-max algo. | MI$(s,z)$ | $H(s)$ | -LGR
---|---|---|---|---
HalfCheetah | DIAYN | $\mathbf{1.815\pm 0.201}$ | $5.164\pm 0.455$ | $-0.492\pm 0.154$
HalfCheetah | DIAYN_FULL | $0.490\pm 0.156$ | $\mathbf{5.519\pm 0.178}$ | $-1.397\pm 0.287$
HalfCheetah | GCRL | $1.626\pm 0.157$ | $5.307\pm 0.173$ | $\mathbf{-0.293\pm 0.356}$
HalfCheetah | cDIAYN | $1.551\pm 0.098$ | $5.104\pm 0.353$ | $-0.741\pm 0.365$
Ant | DIAYN | $1.115\pm 0.268$ | $5.558\pm 0.286$ | $-2.260\pm 0.539$
Ant | DIAYN_FULL | $0.034\pm 0.006$ | $5.044\pm 0.039$ | $-3.010\pm 0.010$
Ant | GCRL | $\mathbf{1.218\pm 0.192}$ | $\mathbf{6.066\pm 0.051}$ | $\mathbf{-2.074\pm 0.232}$
Ant | cDIAYN | $0.750\pm 0.185$ | $5.253\pm 0.270$ | $-2.885\pm 0.291$
Humanoid | DIAYN | $\mathbf{0.927\pm 0.128}$ | $\mathbf{6.074\pm 0.088}$ | $\mathbf{-2.313\pm 0.266}$
Humanoid | DIAYN_FULL | $0.071\pm 0.007$ | $5.676\pm 0.081$ | $-3.051\pm 0.028$
Humanoid | GCRL | $0.769\pm 0.146$ | $6.058\pm 0.063$ | $-2.755\pm 0.115$
Humanoid | cDIAYN | $0.297\pm 0.147$ | $5.677\pm 0.060$ | $-2.738\pm 0.112$
Table 3: Benchmark for MI-max algorithms averaged over 10 seeds. Environment |  D-min algo. | -Energy Distance
---|---|---
HalfCheetah | AIRL | $-1.448\pm 0.843$
HalfCheetah | GAIL | $\mathbf{-0.516\pm 0.640}$
HalfCheetah |  GAIL2 | $-2.488\pm 0.355$
HalfCheetah | MLE | $-2.614\pm 0.415$
Ant | AIRL | $-2.588\pm 1.194$
Ant | GAIL | $\mathbf{-0.802\pm 0.460}$
Ant |  GAIL2 | $-2.898\pm 0.946$
Ant | MLE | $-1.829\pm 0.573$
Humanoid | AIRL | $\mathbf{-1.122\pm 0.210}$
Humanoid | GAIL | $-1.790\pm 0.413$
Humanoid |  GAIL2 | $-1.471\pm 0.528$
Humanoid | MLE | $-2.037\pm 0.973$
Table 4: Benchmark for D-min algorithms averaged over 10 seeds.

MI-max footnote 2 shows the training curves with respect to (a) inferred
reward, (b, c) mutual information MI and state entropy $H(s)$ estimates, and
(d) LGR estimate. The scales of reward between cDIAYN (continuous DIAYN)/GCRL
and DIAYN are very different due to the continuous and discrete $z$
respectively, and they also vary non-monotonically during training due to the
interplay between RL and discriminator fitting. In contrast, both MI and LGR
estimates have the same comparable absolute units for any algorithm in the
family, making them possible for estimating the effectiveness of learning and
compare algorithmic performances. Additional results are presented in Table 3
and Appendix A.2.

Figure 4: D-min results for Ant averaged over 10 seeds: (from left to right)
(a, b) episode reward and energy distance metric across training iterations;
(c, d) visualizations of learned policy state marginal distribution and the
target. Similarly to footnote 2 of MI-max, (a) shows that episode rewards
cannot be used as a metric for evaluating learning progress within each run or
across runs of different choices. Our proposed metric (b), however, can
successfully measure the convergence and algorithm performances (GAIL has
lowest energy distance and performs the best). (c) shows the result of
matching a bi-modal distribution in XY velocity space, where the Ant acquired
a hopping behavior.

D-min Similarly to footnote 2, Figure 4 (left) presents the training curves
for reward and divergence estimate. Unsurprisingly, while the reward
discriminator loss curves are incomparable, the energy distance estimate curve
can be reliably used to measure convergence and compare which algorithmic
variant performs the best in terms of state-marginal matching. Additional
results are presented in Table 4 and Appendix A.2.

Limitations

  * •

Feature Engineering and Metric Space Assumption: Both MI-max and D-min
objectives indirectly involve density estimation, which is sensitive to high
dimensionality and poor conditioning. An implicit assumption in all the above
metrics is that $o(s)$ by the user specifies a low-dimensional, good metric
space where Euclidean distance is a good measure for proximity. Since our
emphasis is on behavior generation in simulation with access to simulator
states (which could be used by downstream behavior distillation procedures
[84, 92]), this remains a reasonable assumption in many cases.

  * •

Sample Efficiency: Sample efficiencies are not the primary concern for
Braxlines, as the code is designed to maximize the benefits of hardware-
accelerated simulations and serve as an interactive toolkit for behavior
designers and RL researchers exploring high-level concepts. While these are
not tuned with respect to sample efficiency compared to many prior works
[126], they exhibit orders of magnitudes speedups in training and often better
final performances (e.g. DIAYN on Humanoid has not been successful [34, 127,
20]).

##  6 Conclusion

In this paper, we introduced Braxlines, a fast and interactive toolkit for
behavior synthesis beyond reward engineering in continuous control that
unlocks four key bottlenecks in RL research: (i) fast environment generation,
(ii) difficulty and limitations of reward engineering, (iii) slow iteration
speeds, and (iv) lack of metrics on data or tasks properties. Our experimental
analysis showcased how two families of algorithms – mutual information
maximization (MI-max) and divergence minimization (D-min) – can generate
interesting behaviors in continuous control environments, and provided
concrete evaluation metrics for interactive debugging. We expect that these
approaches, complementary with classical behavior generation techniques
through reward engineering would be useful for the long-term goal of creating
large data sets of interesting behaviors, and enable algorithmic and
architectural breakthrough in RL. In turn, more diverse data sets of behaviors
have the potential to enable more efficient algorithms (e.g. via distillation)
for few-shot learning in reinforcement learning, as recently observed in
supervised learning.

## Acknowledgments and Disclosure of Funding

We appreciate feedback and advice from Sergey Levine, Nicolas Heess, Marc
Bellemare, Yutaka Matsuo, Pierre Sermanet, Douglas Eck, and Vincent Vanhoucke.

## References

  * Abbeel and Ng [2004] Pieter Abbeel and Andrew Y Ng.  Apprenticeship learning via inverse reinforcement learning.  In _International Conference on Machine Learning_ , 2004.
  * Abdolmaleki et al. [2018] Abbas Abdolmaleki, Jost Tobias Springenberg, Yuval Tassa, Remi Munos, Nicolas Heess, and Martin Riedmiller.  Maximum a posteriori policy optimisation.  In _International Conference on Learning Representations_ , 2018.
  * Achiam et al. [2018] Joshua Achiam, Harrison Edwards, Dario Amodei, and Pieter Abbeel.  Variational option discovery algorithms.  _arXiv preprint arXiv:1807.10299_ , 2018.
  * Andrychowicz et al. [2017] Marcin Andrychowicz, Filip Wolski, Alex Ray, Jonas Schneider, Rachel Fong, Peter Welinder, Bob McGrew, Josh Tobin, Pieter Abbeel, and Wojciech Zaremba.  Hindsight experience replay.  In _Advances in neural information processing systems_ , 2017.
  * Andrychowicz et al. [2021] Marcin Andrychowicz, Anton Raichuk, Piotr Stańczyk, Manu Orsini, Sertan Girgin, Raphaël Marinier, Leonard Hussenot, Matthieu Geist, Olivier Pietquin, Marcin Michalski, Sylvain Gelly, and Olivier Bachem.  What matters for on-policy deep actor-critic methods? a large-scale study.  In _International Conference on Learning Representations_ , 2021.
  * Antol et al. [2015] Stanislaw Antol, Aishwarya Agrawal, Jiasen Lu, Margaret Mitchell, Dhruv Batra, C. Lawrence Zitnick, and Devi Parikh.  VQA: Visual Question Answering.  In _International Conference on Computer Vision_ , 2015.
  * Argall et al. [2009] Brenna D Argall, Sonia Chernova, Manuela Veloso, and Brett Browning.  A survey of robot learning from demonstration.  _Robotics and autonomous systems_ , 2009.
  * Belghazi et al. [2018] Mohamed Ishmael Belghazi, Aristide Baratin, Sai Rajeshwar, Sherjil Ozair, Yoshua Bengio, Aaron Courville, and Devon Hjelm.  Mutual information neural estimation.  In _International Conference on Machine Learning_ , 2018.
  * Bellemare et al. [2016] Marc Bellemare, Sriram Srinivasan, Georg Ostrovski, Tom Schaul, David Saxton, and Remi Munos.  Unifying count-based exploration and intrinsic motivation.  In _Advances in neural information processing systems_ , 2016.
  * Bellemare et al. [2013] Marc G. Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling.  The arcade learning environment: An evaluation platform for general agents.  _Journal of Artificial Intelligence Research_ , 2013.
  * Berseth et al. [2021] Glen Berseth, Daniel Geng, Coline Devin, Nicholas Rhinehart, Chelsea Finn, Dinesh Jayaraman, and Sergey Levine.  Smirl: Surprise minimizing reinforcement learning in unstable environments.  In _International Conference on Learning Representations_ , 2021.
  * Bojar et al. [2014] Ondrej Bojar, Christian Buck, Christian Federmann, Barry Haddow, Philipp Koehn, Johannes Leveling, Christof Monz, Pavel Pecina, Matt Post, Herve Saint-Amand, Radu Soricut, Lucia Specia, and Ale s Tamchyna.  Findings of the 2014 workshop on statistical machine translation.  In _Workshop on Statistical Machine Translation_ , 2014.
  * Bommasani et al. [2021] Rishi Bommasani, Drew A. Hudson, Ehsan Adeli, Russ Altman, Simran Arora, Sydney von Arx, Michael S. Bernstein, Jeannette Bohg, Antoine Bosselut, Emma Brunskill, Erik Brynjolfsson, Shyamal Buch, Dallas Card, Rodrigo Castellon, Niladri Chatterji, Annie Chen, Kathleen Creel, Jared Quincy Davis, Dora Demszky, Chris Donahue, Moussa Doumbouya, Esin Durmus, Stefano Ermon, John Etchemendy, Kawin Ethayarajh, Li Fei-Fei, Chelsea Finn, Trevor Gale, Lauren Gillespie, Karan Goel, Noah Goodman, Shelby Grossman, Neel Guha, Tatsunori Hashimoto, Peter Henderson, John Hewitt, Daniel E. Ho, Jenny Hong, Kyle Hsu, Jing Huang, Thomas Icard, Saahil Jain, Dan Jurafsky, Pratyusha Kalluri, Siddharth Karamcheti, Geoff Keeling, Fereshte Khani, Omar Khattab, Pang Wei Kohd, Mark Krass, Ranjay Krishna, Rohith Kuditipudi, Ananya Kumar, Faisal Ladhak, Mina Lee, Tony Lee, Jure Leskovec, Isabelle Levent, Xiang Lisa Li, Xuechen Li, Tengyu Ma, Ali Malik, Christopher D. Manning, Suvir Mirchandani, Eric Mitchell, Zanele Munyikwa, Suraj Nair, Avanika Narayan, Deepak Narayanan, Ben Newman, Allen Nie, Juan Carlos Niebles, Hamed Nilforoshan, Julian Nyarko, Giray Ogut, Laurel Orr, Isabel Papadimitriou, Joon Sung Park, Chris Piech, Eva Portelance, Christopher Potts, Aditi Raghunathan, Rob Reich, Hongyu Ren, Frieda Rong, Yusuf Roohani, Camilo Ruiz, Jack Ryan, Christopher Ré, Dorsa Sadigh, Shiori Sagawa, Keshav Santhanam, Andy Shih, Krishnan Srinivasan, Alex Tamkin, Rohan Taori, Armin W. Thomas, Florian Tramèr, Rose E. Wang, William Wang, Bohan Wu, Jiajun Wu, Yuhuai Wu, Sang Michael Xie, Michihiro Yasunaga, Jiaxuan You, Matei Zaharia, Michael Zhang, Tianyi Zhang, Xikun Zhang, Yuhui Zhang, Lucia Zheng, Kaitlyn Zhou, and Percy Liang.  On the opportunities and risks of foundation models.  _arXiv preprint arXiv:2108.07258_ , 2021.
  * Brockman et al. [2016] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba.  Openai gym.  _arXiv preprint arXiv:1606.01540_ , 2016.
  * Brooks [1991] Rodney A. Brooks.  Intelligence without representation.  _Artificial Intelligence_ , 1991.
  * Brown et al. [2020] Tom B. Brown, Benjamin Mann, Nick Ryder, Melanie Subbiah, Jared Kaplan, Prafulla Dhariwal, Arvind Neelakantan, Pranav Shyam, Girish Sastry, Amanda Askell, Sandhini Agarwal, Ariel Herbert-Voss, Gretchen Krueger, Tom Henighan, Rewon Child, Aditya Ramesh, Daniel M. Ziegler, Jeffrey Wu, Clemens Winter, Christopher Hesse, Mark Chen, Eric Sigler, Mateusz Litwin, Scott Gray, Benjamin Chess, Jack Clark, Christopher Berner, Sam McCandlish, Alec Radford, Ilya Sutskever, and Dario Amodei.  Language models are few-shot learners.  _arXiv preprint arXiv:2005.14165_ , 2020.
  * Chebotar et al. [2021] Yevgen Chebotar, Karol Hausman, Yao Lu, Ted Xiao, Dmitry Kalashnikov, Jake Varley, Alex Irpan, Benjamin Eysenbach, Ryan Julian, Chelsea Finn, and Sergey Levine.  Actionable models: Unsupervised offline reinforcement learning of robotic skills.  _arXiv preprint arXiv:2104.07749_ , 2021.
  * Chen et al. [2021] Mark Chen, Jerry Tworek, Heewoo Jun, Qiming Yuan, Henrique Ponde de Oliveira Pinto, Jared Kaplan, Harri Edwards, Yuri Burda, Nicholas Joseph, Greg Brockman, Alex Ray, Raul Puri, Gretchen Krueger, Michael Petrov, Heidy Khlaaf, Girish Sastry, Pamela Mishkin, Brooke Chan, Scott Gray, Nick Ryder, Mikhail Pavlov, Alethea Power, Lukasz Kaiser, Mohammad Bavarian, Clemens Winter, Philippe Tillet, Felipe Petroski Such, Dave Cummings, Matthias Plappert, Fotios Chantzis, Elizabeth Barnes, Ariel Herbert-Voss, William Hebgen Guss, Alex Nichol, Alex Paino, Nikolas Tezak, Jie Tang, Igor Babuschkin, Suchir Balaji, Shantanu Jain, William Saunders, Christopher Hesse, Andrew N. Carr, Jan Leike, Josh Achiam, Vedant Misra, Evan Morikawa, Alec Radford, Matthew Knight, Miles Brundage, Mira Murati, Katie Mayer, Peter Welinder, Bob McGrew, Dario Amodei, Sam McCandlish, Ilya Sutskever, and Wojciech Zaremba.  Evaluating large language models trained on code.  _arXiv preprint arXiv:2107.03374_ , 2021.
  * Chevalier-Boisvert et al. [2018] Maxime Chevalier-Boisvert, Lucas Willems, and Suman Pal.  Minimalistic gridworld environment for openai gym.  https://github.com/maximecb/gym-minigrid, 2018.
  * Choi et al. [2021] Jongwook Choi, Archit Sharma, Honglak Lee, Sergey Levine, and Shixiang Shane Gu.  Variational empowerment as representation learning for goal-based reinforcement learning.  In _International Conference on Machine Learning_ , 2021.
  * Chollet [2019] François Chollet.  On the measure of intelligence.  _arXiv preprint arXiv:1911.01547_ , 2019.
  * Christiano et al. [2017] Paul F Christiano, Jan Leike, Tom B Brown, Miljan Martic, Shane Legg, and Dario Amodei.  Deep reinforcement learning from human preferences.  In _Advances in neural information processing systems_ , 2017.
  * Cobbe et al. [2020] Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman.  Leveraging procedural generation to benchmark reinforcement learning.  _arXiv preprint arXiv:1912.01588_ , 2020.
  * Crowston [2012] Kevin Crowston.  _Amazon Mechanical Turk: A Research Tool for Organizations and Information Systems Scholars_ , pages 210–221.  Springer, 2012.
  * Csiszar [1967] Imre Csiszar.  Information-type measures of difference of probability distributions and indirect observation.  _Studia Scientiarum Mathematicarum Hungarica_ , 1967.
  * Deng et al. [2009] Jia Deng, Wei Dong, Richard Socher, Li-Jia Li, Kai Li, and Li Fei-Fei.  Imagenet: A large-scale hierarchical image database.  In _Conference on Computer Vision and Pattern Recognition_ , 2009\.
  * Devlin et al. [2019] Jacob Devlin, Ming-Wei Chang, Kenton Lee, and Kristina Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2019.
  * Dewey [2014] Daniel Dewey.  Reinforcement learning and the reward engineering principle, 2014.
  * Dhariwal et al. [2017] Prafulla Dhariwal, Christopher Hesse, Oleg Klimov, Alex Nichol, Matthias Plappert, Alec Radford, John Schulman, Szymon Sidor, Yuhuai Wu, and Peter Zhokhov.  Openai baselines.  https://github.com/openai/baselines, 2017.
  * Do and Tran [2020] Kien Do and Truyen Tran.  Theory and evaluation metrics for learning disentangled representations.  In _International Conference on Learning Representations_ , 2020.
  * Dosovitskiy et al. [2017] Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun.  Carla: An open urban driving simulator.  In _Conference on Robot Learning_ , 2017.
  * Duan et al. [2016] Yan Duan, Xi Chen, Rein Houthooft, John Schulman, and Pieter Abbeel.  Benchmarking deep reinforcement learning for continuous control.  In _International conference on machine learning_ , 2016.
  * Engstrom et al. [2019] Logan Engstrom, Andrew Ilyas, Shibani Santurkar, Dimitris Tsipras, Firdaus Janoos, Larry Rudolph, and Aleksander Madry.  Implementation matters in deep rl: A case study on ppo and trpo.  In _International Conference on Learning Representations_ , 2019.
  * Eysenbach et al. [2019] Benjamin Eysenbach, Abhishek Gupta, Julian Ibarz, and Sergey Levine.  Diversity is all you need: Learning diverse skills without a reward function.  In _international conference on learning representations_ , 2019.
  * Finn et al. [2016] Chelsea Finn, Paul F. Christiano, Pieter Abbeel, and Sergey Levine.  A connection between generative adversarial networks, inverse reinforcement learning, and energy-based models.  _arXiv preprint arXiv:1611.03852_ , 2016.
  * Florensa et al. [2017] Carlos Florensa, Yan Duan, and Pieter Abbeel.  Stochastic neural networks for hierarchical reinforcement learning.  In _International Conference on Learning Representations_ , 2017.
  * Freeman et al. [2021] C. Daniel Freeman, Erik Frey, Anton Raichuk, Sertan Girgin, Igor Mordatch, and Olivier Bachem.  Brax - a differentiable physics engine for large scale rigid body simulation, 2021.  URL http://github.com/google/brax.
  * Fu et al. [2018] Justin Fu, Katie Luo, and Sergey Levine.  Learning robust rewards with adverserial inverse reinforcement learning.  In _International Conference on Learning Representations_ , 2018.
  * Fu et al. [2020] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine.  D4RL: datasets for deep data-driven reinforcement learning.  _arXiv preprint arXiv:2004.07219_ , 2020.
  * Fujimoto and Gu [2021] Scott Fujimoto and Shixiang Shane Gu.  A minimalist approach to offline reinforcement learning.  _arXiv preprint arXiv:2106.06860_ , 2021.
  * Fujimoto et al. [2018] Scott Fujimoto, Herke van Hoof, and David Meger.  Addressing function approximation error in actor-critic methods.  In _International Conference on Machine Learning_ , 2018.
  * Fujimoto et al. [2019] Scott Fujimoto, David Meger, and Doina Precup.  Off-policy deep reinforcement learning without exploration.  In _International Conference on Machine Learning_ , 2019.
  * Fujita et al. [2021] Yasuhiro Fujita, Prabhat Nagarajan, Toshiki Kataoka, and Takahiro Ishikawa.  Chainerrl: A deep reinforcement learning library.  _Journal of Machine Learning Research_ , 2021.
  * Furuta et al. [2021a] Hiroki Furuta, Tadashi Kozuno, Tatsuya Matsushima, Yutaka Matsuo, and Shixiang Shane Gu.  Co-adaptation of algorithmic and implementational innovations in inference-based deep reinforcement learning.  _arXiv preprint arXiv:2103.17258_ , 2021a.
  * Furuta et al. [2021b] Hiroki Furuta, Tatsuya Matsushima, Tadashi Kozuno, Yutaka Matsuo, Sergey Levine, Ofir Nachum, and Shixiang Shane Gu.  Policy information capacity: Information-theoretic measure for task complexity in deep reinforcement learning.  In _International Conference on Machine Learning_ , 2021b.
  * Ghasemipour et al. [2020] Seyed Kamyar Seyed Ghasemipour, Richard Zemel, and Shixiang Gu.  A divergence minimization perspective on imitation learning methods.  In _Conference on Robot Learning_ , 2020.
  * Ghasemipour et al. [2021] Seyed Kamyar Seyed Ghasemipour, Dale Schuurmans, and Shixiang Shane Gu.  Emaq: Expected-max q-learning operator for simple yet effective offline and online rl.  In _International Conference on Machine Learning_ , 2021.
  * Goodfellow et al. [2014] Ian J. Goodfellow, Jean Pouget-Abadie, Mehdi Mirza, Bing Xu, David Warde-Farley, Sherjil Ozair, Aaron Courville, and Yoshua Bengio.  Generative adversarial networks.  In _Advances in neural information processing systems_ , 2014.
  * Graves et al. [2013] Alex Graves, Abdel rahman Mohamed, and Geoffrey Hinton.  Speech recognition with deep recurrent neural networks.  In _International Conference on Acoustics, Speech, and Signal Processing_ , 2013.
  * Gregor et al. [2016] Karol Gregor, Danilo Jimenez Rezende, and Daan Wierstra.  Variational intrinsic control.  _arXiv preprint arXiv:1611.07507_ , 2016.
  * Gu et al. [2016] Shixiang Gu, Timothy Lillicrap, Ilya Sutskever, and Sergey Levine.  Continuous deep q-learning with model-based acceleration.  In _International Conference on Machine Learning_ , 2016.
  * Gu et al. [2017] Shixiang Gu, Ethan Holly, Timothy Lillicrap, and Sergey Levine.  Deep reinforcement learning for robotic manipulation with asynchronous off-policy updates.  In _International Conference on Robotics and Automation_ , 2017.
  * Guadarrama et al. [2018] Sergio Guadarrama, Anoop Korattikara, Oscar Ramirez, Pablo Castro, Ethan Holly, Sam Fishman, Ke Wang, Ekaterina Gonina, Neal Wu, Efi Kokiopoulou, Luciano Sbaiz, Jamie Smith, Gábor Bartók, Jesse Berent, Chris Harris, Vincent Vanhoucke, and Eugene Brevdo.  TF-Agents: A library for reinforcement learning in tensorflow.  https://github.com/tensorflow/agents, 2018.
  * Haarnoja et al. [2018] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine.  Soft actor-critic: Off-policy maximum entropy deep reinforcement learning with a stochastic actor.  In _International Conference on Machine Learning_ , 2018.
  * Hadfield-Menell et al. [2017] Dylan Hadfield-Menell, Smitha Milli, Pieter Abbeel, Stuart Russell, and Anca D Dragan.  Inverse reward design.  In _Advances in neural information processing systems_ , 2017.
  * Hansen et al. [2020] Steven Hansen, Will Dabney, Andre Barreto, Tom Van de Wiele, David Warde-Farley, and Volodymyr Mnih.  Fast Task Inference with Variational Intrinsic Successor Features.  In _International Conference on Learning Representations_ , 2020.
  * Hazan et al. [2019] Elad Hazan, Sham M. Kakade, Karan Singh, and Abby Van Soest.  Provably efficient maximum entropy exploration.  In _International Conference on Machine Learning_ , 2019.
  * He et al. [2016] Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun.  Deep residual learning for image recognition.  In _Conference on Computer Vision and Pattern Recognition_ , 2016\.
  * He et al. [2017] Kaiming He, Georgia Gkioxari, Piotr Dollar, and Ross Girshick.  Mask r-cnn.  In _International Conference on Computer Vision_ , 2017.
  * Henderson et al. [2018] Peter Henderson, Riashat Islam, Philip Bachman, Joelle Pineau, Doina Precup, and David Meger.  Deep reinforcement learning that matters.  In _AAAI conference on artificial intelligence_ , 2018.
  * Heusel et al. [2017] Martin Heusel, Hubert Ramsauer, Thomas Unterthiner, Bernhard Nessler, and Sepp Hochreiter.  Gans trained by a two time-scale update rule converge to a local nash equilibrium.  _Advances in neural information processing systems_ , 30, 2017.
  * Hinton et al. [2015] Geoffrey Hinton, Oriol Vinyals, and Jeff Dean.  Distilling the knowledge in a neural network.  _arXiv preprint arXiv:1503.025315_ , 2015.
  * Ho and Ermon [2016] Jonathan Ho and Stefano Ermon.  Generative adversarial imitation learning.  In _Advances in neural information processing systems_ , 2016.
  * Huang et al. [2020] Wenlong Huang, Igor Mordatch, and Deepak Pathak.  One policy to control them all: Shared modular policies for agent-agnostic control, 2020.
  * Jaques et al. [2020] Natasha Jaques, Judy Hanwen Shen, Asma Ghandeharioun, Craig Ferguson, Agata Lapedriza, Noah Jones, Shixiang Gu, and Rosalind Picard.  Human-centric dialog training via offline reinforcement learning.  In _Conference on Empirical Methods in Natural Language Processing_ , 2020.
  * Jiang et al. [2019] Yiding Jiang, Shixiang Gu, Kevin Murphy, and Chelsea Finn.  Language as an abstraction for hierarchical deep reinforcement learning.  In _Advances in neural information processing systems_ , 2019.
  * Jumper et al. [2021] John Jumper, Richard Evans, Alexander Pritzel, Tim Green, Michael Figurnov, Olaf Ronneberger, Kathryn Tunyasuvunakool, Russ Bates, Augustin Žídek, Anna Potapenko, Alex Bridgland, Clemens Meyer, Simon A A Kohl, Andrew J Ballard, Andrew Cowie, Bernardino Romera-Paredes, Stanislav Nikolov, Rishub Jain, Jonas Adler, Trevor Back, Stig Petersen, David Reiman, Ellen Clancy, Michal Zielinski, Martin Steinegger, Michalina Pacholska, Tamas Berghammer, Sebastian Bodenstein, David Silver, Oriol Vinyals, Andrew W Senior, Koray Kavukcuoglu, Pushmeet Kohli, and Demis Hassabis.  Highly accurate protein structure prediction with AlphaFold.  _Nature_ , 2021.
  * Jung et al. [2011] Tobias Jung, Daniel Polani, and Peter Stone.  Empowerment for continuous agent—environment systems.  _Adaptive behavior_ , 2011.
  * Kaelbling [1993] Leslie Pack Kaelbling.  Learning to achieve goals.  In _International Joint Conference on Artificial Intelligence_ , 1993\.
  * Kalashnikov et al. [2018] Dmitry Kalashnikov, Alex Irpan, Peter Pastor, Julian Ibarz, Alexander Herzog, Eric Jang, Deirdre Quillen, Ethan Holly, Mrinal Kalakrishnan, Vincent Vanhoucke, and Sergey Levine.  QT-Opt: Scalable deep reinforcement learning for vision-based robotic manipulation.  In _Conference on Robot Learning_ , 2018.
  * Ke et al. [2021] Liyiming Ke, Sanjiban Choudhury, Matt Barnes, Wen Sun, Gilwoo Lee, and Siddhartha Srinivasa.  Imitation learning as f-divergence minimization.  In Steven M. LaValle, Ming Lin, Timo Ojala, Dylan Shell, and Jingjin Yu, editors, _Algorithmic Foundations of Robotics_. Springer, 2021.
  * Kim et al. [2021] Jaekyeom Kim, Seohong Park, and Gunhee Kim.  Unsupervised skill discovery with bottleneck option learning.  In _International Conference on Machine Learning_ , 2021.
  * Kingma and Welling [2014] Diederik P Kingma and Max Welling.  Auto-encoding variational bayes.  In _International Conference on Learning Representations_ , 2014.
  * Kingma et al. [2016] Diederik P. Kingma, Tim Salimans, Rafal Jozefowicz, Xi Chen, Ilya Sutskever, and Max Welling.  Improving variational inference with inverse autoregressive flow.  In _Advances in neural information processing systems_ , 2016.
  * Klyubin et al. [2005] A.S. Klyubin, D. Polani, and C.L. Nehaniv.  Empowerment: a universal agent-centric measure of control.  In _IEEE Congress on Evolutionary Computation_ , 2005.
  * Kostrikov [2018] Ilya Kostrikov.  Pytorch implementations of reinforcement learning algorithms.  https://github.com/ikostrikov/pytorch-a2c-ppo-acktr-gail, 2018.
  * Kostrikov et al. [2019] Ilya Kostrikov, Kumar Krishna Agrawal, Debidatta Dwibedi, Sergey Levine, and Jonathan Tompson.  Discriminator-actor-critic: Addressing sample inefficiency and reward bias in adversarial imitation learning.  In _International Conference on Learning Representations_ , 2019.
  * Krizhevsky et al. [2010] Alex Krizhevsky, Vinod Nair, and Geoffrey Hinton.  Cifar-10 (canadian institute for advanced research), 2010.  URL http://www.cs.toronto.edu/~kriz/cifar.html.
  * Krizhevsky et al. [2012] Alex Krizhevsky, Ilya Sutskever, and Geoffrey E Hinton.  Imagenet classification with deep convolutional neural networks.  In _Advances in Neural Information Processing Systems_ , 2012.
  * Kumar et al. [2020] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine.  Conservative q-learning for offline reinforcement learning.  In _Advances in Neural Information Processing Systems_ , 2020.
  * Kurin et al. [2021] Vitaly Kurin, Maximilian Igl, Tim Rocktäschel, Wendelin Boehmer, and Shimon Whiteson.  My body is a cage: the role of morphology in graph-based incompatible control, 2021.
  * LeCun et al. [2010] Yann LeCun, Corinna Cortes, and CJ Burges.  Mnist handwritten digit database.  _ATT Labs [Online]. Available: http://yann.lecun.com/exdb/mnist_ , 2010.
  * Lee et al. [2020] Lisa Lee, Benjamin Eysenbach, Emilio Parisotto, Eric Xing, Sergey Levine, and Ruslan Salakhutdinov.  Efficient exploration via state marginal matching.  _arXiv preprint arXiv:1906.05274_ , 2020.
  * Levine et al. [2016a] Sergey Levine, Chelsea Finn, Trevor Darrell, and Pieter Abbeel.  End-to-end training of deep visuomotor policies.  _arXiv preprint arXiv:1504.00702_ , 2016a.
  * Levine et al. [2016b] Sergey Levine, Peter Pastor, Alex Krizhevsky, and Deirdre Quillen.  Learning hand-eye coordination for robotic grasping with deep learning and large-scale data collection.  _arXiv preprint arXiv:1603.02199_ , 2016b.
  * Levine et al. [2020] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu.  Offline reinforcement learning: Tutorial, review, and perspectives on open problems.  _arXiv preprint arXiv:2005.01643_ , 2020.
  * Li et al. [2017] Yunzhu Li, Jiaming Song, and Stefano Ermon.  Infogail: Interpretable imitation learning from visual demonstrations.  In _Advances in neural information processing systems_ , 2017.
  * Liang et al. [2017] Eric Liang, Richard Liaw, Robert Nishihara, Philipp Moritz, Roy Fox, Joseph Gonzalez, Ken Goldberg, and Ion Stoica.  Ray rllib: A composable and scalable reinforcement learning library.  _arXiv preprint arXiv:1712.093810_ , 2017.
  * Lillicrap et al. [2016] Timothy P. Lillicrap, Jonathan J. Hunt, Alexander Pritzel, Nicolas Heess, Tom Erez, Yuval Tassa, David Silver, and Daan Wierstra.  Continuous control with deep reinforcement learning.  In _International Conference on Learning Representations_ , 2016.
  * Luong et al. [2015] Minh-Thang Luong, Hieu Pham, and Christopher D. Manning.  Effective approaches to attention-based neural machine translation.  In _Conference on Empirical Methods in Natural Language Processing_ , 2015.
  * Lynch and Sermanet [2021] Corey Lynch and Pierre Sermanet.  Language conditioned imitation learning over unstructured data.  In _Robotics: Science and Systems_ , 2021.
  * Lynch et al. [2019] Corey Lynch, Mohi Khansari, Ted Xiao, Vikash Kumar, Jonathan Tompson, Sergey Levine, and Pierre Sermanet.  Learning latent plans from play.  In _Conference on Robot Learning_ , 2019.
  * Matsushima et al. [2021] Tatsuya Matsushima, Hiroki Furuta, Yutaka Matsuo, Ofir Nachum, and Shixiang Shane Gu.  Deployment-efficient reinforcement learning via model-based offline optimization.  In _International Conference on Learning Representations_ , 2021.
  * Miyato et al. [2018] Takeru Miyato, Toshiki Kataoka, Masanori Koyama, and Yuichi Yoshida.  Spectral normalization for generative adversarial networks.  In _International Conference on Learning Representations_ , 2018.
  * Mnih et al. [2015] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al.  Human-level control through deep reinforcement learning.  _Nature_ , 2015.
  * Mnih et al. [2016] Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu.  Asynchronous methods for deep reinforcement learning.  _arXiv preprint arXiv:1602.01783_ , 2016.
  * Mohamed and Rezende [2015] Shakir Mohamed and Danilo J Rezende.  Variational information maximisation for intrinsically motivated reinforcement learning.  In _Advances in Neural Information Processing Systems_ , 2015.
  * Murphy et al. [2001] Susan A Murphy, Mark J van der Laan, James M Robins, and Conduct Problems Prevention Research Group.  Marginal mean models for dynamic regimes.  _Journal of the American Statistical Association_ , 2001.
  * Nachum et al. [2018] Ofir Nachum, Shixiang Gu, Honglak Lee, and Sergey Levine.  Data-Efficient Hierarchical Reinforcement Learning.  In _Advances in neural information processing systems_ , 2018.
  * Nair et al. [2018] Ashvin Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine.  Visual reinforcement learning with imagined goals.  _arXiv preprint arXiv:1807.04742_ , 2018.
  * Ng and Russell [2000] Andrew Y Ng and Stuart Russell.  Algorithms for inverse reinforcement learning.  In _International Conference on Machine Learning_ , 2000.
  * Orsini et al. [2021] Manu Orsini, Anton Raichuk, Léonard Hussenot, Damien Vincent, Robert Dadashi, Sertan Girgin, Matthieu Geist, Olivier Bachem, Olivier Pietquin, and Marcin Andrychowicz.  What matters for adversarial imitation learning?  _arXiv preprint arXiv:2106.00672_ , 2021.
  * Oudeyer and Kaplan [2009] Pierre-Yves Oudeyer and Frederic Kaplan.  What is intrinsic motivation? a typology of computational approaches.  _Frontiers in Neurorobotics_ , 2009.
  * Pathak et al. [2017] Deepak Pathak, Pulkit Agrawal, Alexei A. Efros, and Trevor Darrell.  Curiosity-driven exploration by self-supervised prediction.  In _International Conference on Machine Learning_ , 2017.
  * Pong et al. [2018] Vitchyr Pong, Shixiang Gu, Murtaza Dalal, and Sergey Levine.  Temporal difference models: Model-free deep rl for model-based control.  _International Conference on Learning Representations_ , 2018.
  * Pong et al. [2020] Vitchyr H Pong, Murtaza Dalal, Steven Lin, Ashvin Nair, Shikhar Bahl, and Sergey Levine.  Skew-fit: State-covering self-supervised reinforcement learning.  In _International Conference on Machine Learning_ , 2020.
  * Poole et al. [2019] Ben Poole, Sherjil Ozair, Aaron Van Den Oord, Alex Alemi, and George Tucker.  On variational bounds of mutual information.  In _International Conference on Machine Learning_ , 2019.
  * Radford et al. [2019] Alec Radford, Jeff Wu, Rewon Child, David Luan, Dario Amodei, and Ilya Sutskever.  Language models are unsupervised multitask learners, 2019.
  * Radford et al. [2021] Alec Radford, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, Amanda Askell, Pamela Mishkin, Jack Clark, Gretchen Krueger, and Ilya Sutskever.  Learning transferable visual models from natural language supervision.  _arXiv preprint arXiv:2103.00020_ , 2021.
  * Raffin et al. [2020] Antonin Raffin, Ashley Hill, Maximilian Enerstus, Adam Gleave, Anssi Kanervisto, and Noah Dormann.  Stable baselines3, 2020.  URL https://github.com/DLR-RM/stable-baselines3.
  * Ramesh et al. [2021] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever.  Zero-shot text-to-image generation.  _arXiv preprint arXiv:2102.12092_ , 2021.
  * Ravindran and Barto [2002] Balaraman Ravindran and Andrew G Barto.  Model minimization in hierarchical reinforcement learning.  In _Abstraction, Reformulation, and Approximation_ , 2002.
  * Reda et al. [2020] Daniele Reda, Tianxin Tao, and Michiel van de Panne.  Learning to locomote: Understanding how environment design matters for deep reinforcement learning.  In _Motion, Interaction and Games_ , 2020.
  * Redmon et al. [2016] Joseph Redmon, Santosh Divvala, Ross Girshick, and Ali Farhadi.  You only look once: Unified, real-time object detection.  In _Conference on Computer Vision and Pattern Recognition_ , 2016\.
  * Rényi [1961] Alfréd Rényi.  On measures of entropy and information.  In _Berkeley Symposium on Mathematical Statistics and Probability_ , 1961.
  * Rezende and Mohamed [2016] Danilo Jimenez Rezende and Shakir Mohamed.  Variational inference with normalizing flows.  In _International Conference on Machine Learning_ , 2016.
  * Rusu et al. [2016a] Andrei A. Rusu, Sergio Gomez Colmenarejo, Caglar Gulcehre, Guillaume Desjardins, James Kirkpatrick, Razvan Pascanu, Volodymyr Mnih, Koray Kavukcuoglu, and Raia Hadsell.  Policy distillation.  _arXiv preprint arXiv:1511.06295_ , 2016a.
  * Rusu et al. [2016b] Andrei A. Rusu, Neil C. Rabinowitz, Guillaume Desjardins, Hubert Soyer, James Kirkpatrick, Koray Kavukcuoglu, Razvan Pascanu, and Raia Hadsell.  Progressive neural networks.  _arXiv preprint arXiv:1606.04671_ , 2016b.
  * Salge et al. [2014] Christoph Salge, Cornelius Glackin, and Daniel Polani.  Empowerment–An introduction.  In _Guided Self-Organization: Inception_ , pages 67–114. Springer, 2014.
  * Salimans et al. [2017] Tim Salimans, Jonathan Ho, Xi Chen, Szymon Sidor, and Ilya Sutskever.  Evolution strategies as a scalable alternative to reinforcement learning.  _arXiv preprint arXiv:1703.03864_ , 2017.
  * Savva et al. [2019] Manolis Savva, Abhishek Kadian, Oleksandr Maksymets, Yili Zhao, Erik Wijmans, Bhavana Jain, Julian Straub, Jia Liu, Vladlen Koltun, Jitendra Malik, Devi Parikh, and Dhruv Batra.  Habitat: A platform for embodied ai research.  In _International Conference on Computer Vision_ , 2019.
  * Schaul et al. [2015] Tom Schaul, Dan Horgan, Karol Gregor, and David Silver.  Universal Value Function Approximators.  In _International Conference on Machine Learning_ , 2015.
  * Schmidhuber [2010] Jrgen Schmidhuber.  Formal theory of creativity, fun, and intrinsic motivation (1990-2010).  _IEEE transactions on autonomous mental development_ , 2010.
  * Schulman et al. [2015] John Schulman, Sergey Levine, Philipp Moritz, Michael Jordan, and Pieter Abbeel.  Trust region policy optimization.  In _International Conference on Machine Learning_ , 2015.
  * Schulman et al. [2017] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  _arXiv preprint arXiv:1707.06347_ , 2017.
  * Sharma et al. [2020a] Archit Sharma, Michael Ahn, Sergey Levine, Vikash Kumar, Karol Hausman, and Shixiang Gu.  Emergent real-world robotic skills via unsupervised off-policy reinforcement learning.  In _Robotics: Science and Systems_ , 2020a.
  * Sharma et al. [2020b] Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman.  Dynamics-aware unsupervised discovery of skills.  In _International conference on learning representations_ , 2020b.
  * Siegel et al. [2020] Noah Y. Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmaleki, Michael Neunert, Thomas Lampe, Roland Hafner, and Martin A. Riedmiller.  Keep doing what worked: Behavioral modelling priors for offline reinforcement learning.  In _International Conference on Learning Representations_ , 2020.
  * Silver et al. [2021] David Silver, Satinder Singh, Doina Precup, and Richard S Sutton.  Reward is enough.  _Artificial intelligence_ , 2021.
  * Sutton and Barto [2018] Richard S. Sutton and Andrew G. Barto.  _Reinforcement Learning: An Introduction_.  The MIT Press, 2018.
  * Tassa et al. [2018] Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, Timothy Lillicrap, and Martin Riedmiller.  Deepmind control suite.  _arXiv preprint arXiv:1801.00690_ , 2018.
  * Team et al. [2021] Open Ended Learning Team, Adam Stooke, Anuj Mahajan, Catarina Barros, Charlie Deck, Jakob Bauer, Jakub Sygnowski, Maja Trebacz, Max Jaderberg, Michael Mathieu, Nat McAleese, Nathalie Bradley-Schmieg, Nathaniel Wong, Nicolas Porcel, Roberta Raileanu, Steph Hughes-Fitt, Valentin Dalibard, and Wojciech Marian Czarnecki.  Open-ended learning leads to generally capable agents.  _arXiv preprint arXiv:2107.12808_ , 2021.
  * Todorov et al. [2012] Emanuel Todorov, Tom Erez, and Yuval Tassa.  Mujoco: A physics engine for model-based control.  In _International Conference on Intelligent Robots and Systems_ , 2012\.
  * van den Oord et al. [2016] Aaron van den Oord, Nal Kalchbrenner, and Koray Kavukcuoglu.  Pixel recurrent neural networks.  In _International Conference on Machine Learning_ , 2016.
  * van den Oord et al. [2019] Aaron van den Oord, Yazhe Li, and Oriol Vinyals.  Representation learning with contrastive predictive coding.  _arXiv preprint arXiv:1807.03748_ , 2019.
  * Varela et al. [1991] Francisco J. Varela, Eleanor Rosch, and Evan Thompson.  _The Embodied Mind: Cognitive Science and Human Experience_.  The MIT Press, 1991.
  * Vaswani et al. [2017] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Ł ukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  In _Advances in Neural Information Processing Systems_ , 2017.
  * Vinyals et al. [2015] Oriol Vinyals, Alexander Toshev, Samy Bengio, and Dumitru Erhan.  Show and tell: A neural image caption generator.  In _International Conference on Computer Vision_ , 2015.
  * Vinyals et al. [2017] Oriol Vinyals, Timo Ewalds, Sergey Bartunov, Petko Georgiev, Alexander Sasha Vezhnevets, Michelle Yeo, Alireza Makhzani, Heinrich Küttler, John Agapiou, Julian Schrittwieser, John Quan, Stephen Gaffney, Stig Petersen, Karen Simonyan, Tom Schaul, Hado van Hasselt, David Silver, Timothy Lillicrap, Kevin Calderone, Paul Keet, Anthony Brunasso, David Lawrence, Anders Ekermo, Jacob Repp, and Rodney Tsing.  Starcraft ii: A new challenge for reinforcement learning.  _arXiv preprint arXiv:1708.04782_ , 2017.
  * Warde-Farley et al. [2019] David Warde-Farley, Tom Van de Wiele, Tejas Kulkarni, Catalin Ionescu, Steven Hansen, and Volodymyr Mnih.  Unsupervised Control Through Non-Parametric Discriminative Rewards.  In _International Conference on Learning Representations_ , 2019.
  * Williams [1992] Ronald J. Williams.  Simple statistical gradient-following algorithms for connectionist reinforcement learning.  _Machine Learning_ , 1992.
  * Wulfmeier et al. [2015] Markus Wulfmeier, Peter Ondruska, and Ingmar Posner.  Maximum entropy deep inverse reinforcement learning.  _arXiv preprint arXiv:1507.04888_ , 2015.
  * Yu et al. [2019] Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine.  Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning.  In _Conference on Robot Learning_ , 2019.
  * Zhao et al. [2021] Rui Zhao, Yang Gao, Pieter Abbeel, Volker Tresp, and Wei Xu.  Mutual information state intrinsic control.  In _International Conference on Learning Representations_ , 2021.
  * Ziebart et al. [2008] Brian D. Ziebart, Andrew L. Maas, J. Andrew Bagnell, and Anind K. Dey.  Maximum entropy inverse reinforcement learning.  In _AAAI conference on artificial intelligence_ , 2008.

## Appendix

##  Appendix A Braxlines Benchmark, Hyperparameters, Ablations, and Analyses

We detail the additional information for experimental results and Braxlines
Benchmark. While Braxlines is designed more as a tool to enable rapid creation
of environments and behaviors, it also provides minimal quantitative
benchmarks that could be utilized as a reference for how to evaluate new
environments or new MI-max and D-min algorithms. Such baselines for reward-
free algorithms have been lacking, or non-existent, in open-source libraries
(see Tables 5 and 6).

Framework | DDPG | A3C | PPO | TRPO | TD3 | SAC | ES
---|---|---|---|---|---|---|---
Baselines [29] | $\checkmark$ | $\checkmark$ | $\checkmark$ | $\checkmark$ |  |  |
Stable Baselines [110] | $\checkmark$ | $\checkmark$ | $\checkmark$ |  | $\checkmark$ | $\checkmark$ |
RLLab [32] | $\checkmark$ |  |  | $\checkmark$ |  |  | $\checkmark$
Tf-Agents [29] | $\checkmark$ |  | $\checkmark$ |  | $\checkmark$ | $\checkmark$ |
RLLib [88] | $\checkmark$ | $\checkmark$ | $\checkmark$ |  | $\checkmark$ | $\checkmark$ | $\checkmark$
PFRL [43] | $\checkmark$ | $\checkmark$ | $\checkmark$ | $\checkmark$ | $\checkmark$ | $\checkmark$ |
Brax/Braxlines(ours) |  |  | $\checkmark$ |  |  | $\checkmark$ | $\checkmark$
Table 5: A survey on the support of the most common continuous control, reward maximization algorithms from a set of well-maintained, stable RL baselines. Framework | AIRL | GAIL | GCRL | DIAYN | cDIAYN |  Oth. MI-max
[56, 140]
---|---|---|---|---|---|---
Baselines [29] |  | $\checkmark$ |  |  |  |
Stable Baselines [110] |  |  |  |  |  |
RLLab [32] |  |  |  |  |  |
Tf-Agents [29] |  |  |  |  |  |
RLLib [88] |  |  |  |  |  |
PFRL [43] |  |  |  |  |  |
Braxlines(ours) | $\checkmark$ | $\checkmark$ | $\checkmark$ | $\checkmark$ | $\checkmark$ |  $\checkmark$
Kostrikov [76] |  | $\checkmark$ |  |  |  |
Eysenbach et al. [34] |  |  |  | $\checkmark$ |  |
Sharma et al. [127] |  |  |  | $\checkmark$ | $\checkmark$ |
Choi et al. [20] |  |  |  |  |  |
Table 6: A survey on the support of the most common algorithms in MI-max and
D-min families across a set of well-maintained, stable RL baselines. Since
standard RL baselines rarely support advanced algorithms, we added a few
additional codebases (some from original algorithmic papers). Note that Choi
et al. [20] is empty since no code is open-sourced.

###  A.1 Hyperparameters

Hyperparameters for reproducing all the benchmark and experiment results can
be found in
https://github.com/google/brax/blob/main/brax/experimental/braxlines/experiments.
These experiments can be easily run serially on Colab
https://github.com/google/brax/blob/main/notebooks/braxlines/experiment_sweep.ipynb
based on a configuration list of dictionaries such as
https://github.com/google/brax/blob/main/brax/experimental/braxlines/experiments/mimax_sweep.py.
Else, users can modify run_experiment() from
https://github.com/google/brax/blob/main/brax/experimental/braxlines/experiments/__init__.py
to launch parallel sweeps on their custom computing clusters.

##### Experiment Parameters

We used 10 random seeds for each experiment result to compute its mean and
variance. Similarly to [34, 127, 20], we assume prior knowledge on dimensions
of interest, e.g. XY-velocities: obs_indices$=(11,)$ for HalfCheetah,
$(13,14)$ for Ant, and $(22,23)$ (see
https://github.com/google/brax/blob/main/brax/experimental/composer/obs_descs.py),
for learning (except DIAYN_FULL) and evaluation metrics.

##### PPO Parameters

We used the same exact hyperparameters for PPO as in task-reward RL examples
in Brax (see https://github.com/google/brax/blob/main/notebooks/training.ipynb
for their original Colab and
https://github.com/google/brax/tree/main/brax/experimental/braxlines/experiments/defaults.py
for these PPO hyperparameters), except multiplying the num_timesteps by 2 to
allow longer training than in the single-task setting. Fully-connected MLPs
with linen.swish activation function and hidden sizes of $[32,32,32,32]$ are
used for policy functions, and hidden sizes of $[256,256,256,256,256]$ for
value functions (see
https://github.com/google/brax/blob/main/brax/training/networks.py).

#####  MI-max Parameters

The hyperparameters are specified in
https://github.com/google/brax/blob/main/brax/experimental/braxlines/experiments/mimax_sweep.py.
For discrete latent experiments (DIAYNand DIAYN_FULL), number of skills for
$z$ diayn_num_skills is set to $8$. For continuous latent experiments
(cDIAYN), the $z$ dimension is set to $2$. Fully-connected MLPs with
linen.swish activation function and hidden sizes of $[32,32]$ are used for
discriminator functions $q(z|o(s))$. See Section A.3 for ablations on key
hyperparameters such as spectral_norm and diayn_num_skills.

#####  D-min Parameters

The hyperparameters are specified in
https://github.com/google/brax/blob/main/brax/experimental/braxlines/experiments/dmin_sweep.py.
Fully-connected MLPs with linen.swish activation function and hidden sizes of
$[32,32]$ are used for binary discriminator functions. See Section A.3 for
ablations on key hyperparameters such as gradient_penalty_weight.

#####  Composer Parameters

Braxlines Composer is designed as a toolkit for efficiently reusing
environment components, observation definitions, and reward functions for
composing new tasks or reward-free environments. Example “hyperparameters” of
these environments are listed in
https://github.com/google/brax/tree/main/brax/experimental/composer/env_descs.py.
Similar to the rest of experiments, hyperparameter sweeps are supported. If
the comparison is about the optimizability by PPO with respect to a task
score, it is important to ensure that score_fns are the same across all swept
variants. Examples include:

  * •

Sweeping Reward Function Parameters: When you have a sum of multiple reward
terms, it’s helpful to sweep over the scaling parameters to find the optimal
trade-off among these to have the best task score of interest
https://github.com/google/brax/blob/main/brax/experimental/braxlines/experiments/ant_push_sweep.py

  * •

Sweeping Morphology Parameters: You could parameterize a component to have
varying morphologies (e.g. different number of leggs for Ant)
https://github.com/google/brax/blob/main/brax/experimental/braxlines/experiments/ant_run_morphology_sweep.py

###  A.2 Additional Benchmark Results

Figure 5:  MI-max results on Humanoid averaged over 10 seeds: (from left to
right) (a) episodic discriminator reward, (b, c) MI(z,s)and H(s)estimates from
Algorithm 1, and (d) LGR from Algorithm 3. Unlike Ant results in footnote 2,
DIAYN performs the best both in terms of MI(z,s)and LGR metrics, better than
GCRL. Since H(s)is the same for DIAYN and GCRL, this means that DIAYN acquired
better on average controllability/consistency with respect to each given
target goal (reasonable since DIAYN is only learning 8 targets, while GCRL is
learning a infinite set of targets).  Figure 6:  MI-max DIAYN skill
visualization on Humanoid. Humanoid successfully learns to move in multiple
directions. See behavior videos in
https://github.com/google/brax/blob/main/brax/experimental/braxlines.

##### MI-max

Figure 5 and Figure 6 show learning curves and skill visualizations for
Humanoid respectively. To the best of our knowledge, DADS [127] is one of the
few unsupervised RL algorithms that scaled to Humanoid, and we are the first
to report successful DIAYN results on it.

Figure 7: D-min results for Humanoid averaged over 10 seeds: (from left to
right) (a, b) episode reward and energy distance metric across training
iterations; (c, d) visualizations of learned policy state marginal
distribution and the target. Unlike Ant results in Figure 4, AIRL and GAIL2
significantly outperforms GAIL. See behavior videos in
https://github.com/google/brax/blob/main/brax/experimental/braxlines.

##### D-min

Figure 7 shows learning curves and result visualization for Humanoid. Humanoid
successfully learns to match the given bi-modal target distribution.

Figure 8: Task parameter sweeps on (left and middle) Ant Push and (right)
morphological Ant Run in Figure 2. See behavior videos in
https://github.com/google/brax/blob/main/brax/experimental/braxlines.

#####  Braxlines Composer

Figure 8 show training results for environment examples in Figure 2. For Ant
Push, we observe that “scale” which balances between object-to-target-velocity
and ant-to-object reward terms has significant effect on final scores (since
“episode_reward” is a variable here, we defined “episode_score” as a
consistent metric for evaluation). For the morphologically-varying Ant Run, we
observe that the running performance is largely consistent, with an exception
in the case of two legs. Importantly, as described by simple examples in
Appendix A.1 and D, Composer allows simple programming of environment
variations and hyperparameter sweeps.

###  A.3 Ablation Studies and Analysis

This section mainly lists a few design parameters that are critical to getting
interesting emergent behaviors. For more examples, check out our codebase or
directly try out our interactive Google Colabs in
https://github.com/google/brax/tree/main/notebooks/braxlines.

Figure 9:  MI-max results on Ant with and without Spectral Normalization
(“sn_F” means no spectral norm) averaged over 10 seeds: (from left to right)
(a) episodic discriminator reward, (b, c) MI(z,s)and H(s)estimates from
Algorithm 1, and (d) LGR from Algorithm 3.  Figure 10:  MI-max results on
Humanoid with and without Spectral Normalization (“sn_F” means no spectral
norm) averaged over 10 seeds: (from left to right) (a) episodic discriminator
reward, (b, c) MI(z,s)and H(s)estimates from Algorithm 1, and (d) LGR from
Algorithm 3.  Figure 11: Skills visualization for (from left to right) (a)
DIAYN without spectral normalization (SN), (b) DIAYN with SN, (c) continuous
DIAYN (cDIAYN) without SN, and (d) cDIAYN with SN. DIAYN with continuous $z$
has been notoriously difficult even with feature engineering, where Eysenbach
et al. [34] and Sharma et al. [127] have not showed successful skill learning.
A simple combination of PPO [125] and spectral normalization (SN) [94] has
proven to be surprisingly effective and enabled training a continuous-latent
DIAYN (while results are slightly worse than those of discrete-latent DIAYN
according to footnote 2.

##### Spectral Normalization

MI-max algorithms exhibit complex learning dynamics through continuously
changing reward functions. Choi et al. [20] showed that the use of spectral
normalization (SN) can significantly stabilize learning and quality of
discovered skills. Figure 9 and Figure 10 show that SN does improve DIAYN
results substantially in MI(z,s)and H(s), in both Ant and Humanoid. We also
show in Figure 11 that even cDIAYN can discover meaningful skills, which has
been difficult in prior works [34, 127, 20]. This result, along with our
success with Humanoid, is likely because most prior results built on SAC which
is more sample-efficient but less stable than PPO.

Figure 12: D-min results for Ant with or without Gradient Penalty (“gpw”
stands for gradient penalty weight) averaged over 10 seeds: (from left to
right) (a, b) episode reward and energy distance metric across training
iterations for AIRL; (c, d) episode reward and energy distance metric across
training iterations for GAIL. Both Braxlines AIRL and GAIL are stable with or
without gradient penalties, and too much penalty can degrade performances.

##### Gradient Penalty

D-min algorithms also exhibit difficult learning dynamics, and many prior
works utilized numerous additional implementations to stabilize learning, such
as separately learning terminal reward for bias correction [77, 102] and
fictitious play [83]. Contrary to prior wisdom, Figure 12 shows that gradient
penalty does not significantly improve the learning performances for these bi-
modal target matching tasks in Ant or Humanoid.

(a) DIAYN_FULL (full state)

(b) DIAYN (xy)

(c) DIAYN (xy) + Run Reward

Figure 13: The importance of feature engineering (o(s)) and reward
augmentation (env_reward_multiplier) to discover diverse behaviors when
running (a) DIAYN [34] with full observation (DIAYN_FULL), (b) DIAYN with XY-
velocities (dim=(13,14)), (c) DIAYN with XY + adding original reward on Brax’s
Ant [37]. Naively using full states does not lead to interesting behaviors
like walking or hopping (it just learns to pose differently). With simple
specification of dimensions in observation, however, its learns very
effectively. Simple feature engineering and MI-max is therefore a powerful
tool for behavior engineering without much human efforts. Additionally, if
there is some task reward available, it can be added to automatically discover
multiple distinguishable ways of accomplishing the same task.

##### Feature Engineering $o(s)$

Figure 13 shows qualitative results of MI-max with or without feature
engineering $o(s)$ (DIAYN or DIAYN_FULL). As different dimensions have
different sensitivity to the agent’s actions, often the easiest dimension is
diversified first to maximize the objective in discriminative MI-max, i.e. in
DIAYN_FULL Ant just learns to make slightly different poses. By removing
uninteresting and trivial dimensions for control, the agent can effectively
focus on controlling the key dimensions such as XY velocities (dim$=(13,14)$).
Our fast interactive tool allows efficiently iterating different feature
choices.

Additionally, Figure 14 and Figure 15 provide a deeper look into MI and $H(s)$
estimates across first 30 dimensions Ant, DIAYN or DIAYN_FULL. This
essentially evaluates intrinsic controllability that the agent has over the
environment, from which we can gain quantitative intuitions about the agent’s
behaviors as well as its affinity with the environment. For example,
DIAYNlearns to run in different directions (dim=(13,14)), but as it runs it
does not change its Z-axis orientation (dim=1) and therefore indirectly gained
controllability over it. If we can measure such indirect relationships among
dimensions, we can more efficiently design minimal features $o(s)$.
Furthermore, combining with the insights from policy information capacity
(PIC) [45], we can choose to design environments or random initial policies
such that the initial MI is already high with respect to the dimensions of
interest, i.e. designing environments and agents for optimizability [45, 113].

##### Task-Reward Reward Augmentation

When a task reward is combined with MI-max (env_reward_multiplier $>0$), it
essentially diversifies policies while accomplishing a given task, i.e.
learning multiple independent ways of accomplishing the task. Such auxiliary
task reward can sometimes be as simple as a positive survival reward (which is
used in all our MI-max and D-min experiments as listed in Table 2), if the
environment itself is non-trivial and has termination like the Humanoid
environment [127]. In Figure 13, the standard Ant’s forward running reward is
combined, and it learns different gait behaviors.

Figure 14: Mutual information evaluations across first 30 dimensions on Ant.
We analyzed two configurations: DIAYN uses $o(s[13,14])$ (i.e., the XY
velocities of the root), while DIAYN_FULL uses the full state. The red
background indicates when DIAYN is higher than DIAYN_FULL more than 80% of the
time and blue if vice versa. While it is unsurprising that DIAYN excels at
gaining controllability over specified dimensions (13, 14), it is interesting
that DIAYN additionally gained controllability over dim=1 (i.e., the first
rotation angle of the root) and lost controllability over dim=(0, 2) (the
height and other rotation angles of the root) as an indirect a consequence of
empowering dim=(13,14). Figure 15: Entropy evaluations across first 30
dimensions on Ant. DIAYN uses $o(s[13,14])$ (corresponding to xy velocities),
while DIAYN_FULL uses identity $o$ (full observation state). It is noticeably
to note that in all dimensions DIAYN (XY) gains higher state entropies than
DIAYN (full state). Contrasting this result with Figure 14, we can conclude
that while DIAYN (XY) can thoroughly increase marginal entropies, it can only
gain controllabilities (low conditional entropies) over a limited set of the
dimensions. In cases where you only care about marginal entropies such as for
risk-seeking exploration, this may be desired over controllability.

##  Appendix B Latent Goal Reaching

Algorithm 3 estimates the LGR metric [20] of a policy $\pi(a|s,z)$. This has
exact correspondence with the standard goal-reaching objective in GCRL.

Input: agent $\pi(a|s,z)$; $o(s)$; goals $o^{*}_{1:L}=o(s^{*}_{1:L})\sim
p^{*}(s)$ ; inference $q(z|o(s))$

Output: LGR estimate $\hat{\text{LGR}}$

Sample $z_{l,n}\sim q(z|s^{*}_{l})$ // Infer $N$ latent intents per goal
$s^{*}_{l}$

Sample $s_{l,n,m,t}\sim\pi(\cdot|\cdot,z_{l,n})$ in $\mu$ // Collect $M$
episodes of horizon $T$ per $z_{l,n}$

// Compute average goal-reaching performance per $z_{l,n}$ for recovering
$s^{*}_{l}$

$\hat{\text{LGR}}=\sum_{l,n,m,t}-(o^{*}_{l}-o(s_{l,n,m,t}))^{2}/\sigma^{2}$

Algorithm 3 Latent Goal Reaching [20]

###  B.1 Justification for Negative Latent Goal Reaching as An Empowerment
Approximation

Given $p^{*}(z)=\arg\max_{p(z)}I_{\pi}(z,o({S^{\mu}}))$, the empowerment per
environment $\mu$ is given by:

$\displaystyle p^{*}(z)p(o({S^{\mu}})|z)$
$\displaystyle=p^{*}(o({S^{\mu}}))p^{*}(z|o({S^{\mu}}))$ (8) $\displaystyle
I_{\pi}^{*}(z,o({S^{\mu}}))$
$\displaystyle=H^{*}(o({S^{\mu}}))-H^{*}(o({S^{\mu}})|z)$
$\displaystyle=H^{*}(o({S^{\mu}}))+\mathbb{E}_{p^{*}(o({S^{\mu}}))p^{*}(z|o({S^{\mu}}))}\left[\log
p(o({S^{\mu}})|z)\right]$ (9)

Importantly, in joint distributions in Eq. 8, $p(o({S^{\mu}})|z)$ does not
depend on the choice of $p(z)$. We then make two important assumptions: (1) we
are given with samples of $p^{*}(o({S^{\mu}}))$ and $p^{*}(o({S^{\mu}}))$ is
fixed regardless of the choice of agent, and (2) we are given with (or could
train models to approximate) $q(z|o({S^{\mu}}))\approx p^{*}(z|o({S^{\mu}}))$.
While both assumptions are difficult in practice to satisfy accurately, since
each depends on $p^{*}(z)$ which clearly depends on the agent policy
$\pi(a|o,z)$, they do allow deriving a connection between generalized goal-
reaching objective and empowerment estimation. In particular, Eq. 9 can now be
estimated by:

$\displaystyle I_{\pi}^{*}(z,S^{\mu})$
$\displaystyle\approx\mathbb{E}_{p^{*}(o({S^{\mu}}))q(z|o({S^{\mu}}))}\left[\log
p(o({S^{\mu}})|z)\right]+C_{1},$ (10)

where $C_{1}=H^{*}(o({S^{\mu}}))$ is now just a constant offset due to our
assumption (1) that does not depend on the evaluation agent $\pi$. Eq. 10 is
equivalently the variational autoencoder objective [73], except in RL we need
to somehow estimate $\log p(o({S^{\mu}})|z)$. One way is to fit a density
estimation model [83, 11], but to make a direct connection to goal-reaching
performances, we instead use a non-parametric density estimation based on a
mixture of Gaussians. Following the notations in Algorithm 3,

$\displaystyle\log p(o^{*}_{l}|z_{l,n})$
$\displaystyle\approx\log\frac{1}{MT}\sum_{m,t}\mathcal{N}(o^{*}_{l}|o_{l,n,m,t},\sigma^{2}I)$
(11)
$\displaystyle\geq\sum_{m,t}-(o^{*}_{l}-o_{l,n,m,t})^{2}/\sigma^{2}+C_{2},$
(12)

where Eq. 12 follows from a simple application of Jensen’s inequality, and
$C_{2}$ is again a constant offset that does not depend on $\pi$.

##  Appendix C Prior Evaluation Methodologies for MI-max and D-min Algorithms

In this section, we briefly summarize prior evaluation methodologies for MI-
max and D-min families, and relation to the metrics in the Braxlines. Since
the inferred reward function dynamically changes during training, could have
an arbitrary range of values, and largely depends on the different
implementation and hyperparameter choices such as $z$-space in MI-max, or
reward transformation in D-min, it is difficult to apply episodic reward for
quantitatively evaluating convergence or final performances. As in Table 7,
most of evaluations for MI-max and D-min have relied on qualitative
comparisons, visualization of rollout videos or density plots [34, 127, 46,
83], or down-stream task performances [34, 127, 126]. Qualitative evaluation
might be hard to ensure objectiveness and automate the whole process, and the
evaluation with down-stream tasks may not correspond with the quality of
generated behaviors themselves. While previous D-min algorithms mostly focus
on imitation learning and frequently use task rewards directly to compare how
close to the experts they got [63, 38, 46], such evaluation is not suitable
for the distribution sketching case in our paper.

Recently, Kim et al. [72] employ SEPIN@$k$ and W SEPIN, which are proposed in
the disentangled representation learning literature [30], for the quantitative
evaluation of MI-max algorithms. Both SEPIN@$k$ and W SEPIN consider the
mutual information between the skills $Z$ and the last states of the
trajectories with specified dimensions $o(S_{T})$. SEPIN@$k$ is the top-$k$
average of $\text{MI}(o(S_{T}),Z_{i};Z_{\neq i})$ over skills $i=1,2,\dots$,
and W SEPIN is the weighted sum of $\text{MI}(o(S_{T}),Z_{i};Z_{\neq i})$ over
skills. They also propose to simply use $\text{MI}(o(S_{T}),Z)$ as a
information theoretic metrics. See Kim et al. [72] for the detailed
discussion. Because they consider the last states the agents reached and
measure how they are different depending on $z$, these might be similar
metrics to LGR.

Family | Metrics | Type | Hyperparameter-Agnostic | References
---|---|---|---|---
MI-max | Episodic Discriminator Reward | Quantitative | No | –
MI-max | Diversity in Rollout Videos | Qualitative | Yes | [34, 127, 36, 72]
MI-max | Downstream Task Performance | Quantitative | Yes | [34, 127, 36, 56, 50]
MI-max | SEPIN@$k$, W SEPIN [30] | Quantitative | Yes | [72]
D-min | Episodic Inferred Reward | Quantitative | No | –
D-min | Density Visualization | Qualitative | Yes | [46, 83]
D-min | Imitation Task Performance | Quantitative | Yes | [63, 38, 46]
MI-max | Particle-based MI Estimation | Quantitative | Yes | Ours, [45]
MI-max | Latent Goal Reaching | Quantitative | Yes | Ours, [20]
D-min | Energy Distance | Quantitative | Yes | Ours

Table 7: A summary of prior evaluation methods for MI-max and D-min
algorithms.

##  Appendix D Braxlines Composer API Example

We made possible in Braxlines Composer to design continuous control tasks with
under 50 lines of code (including comments). The code snippet below shows how
to construct a manipulation, goal-conditioned task that we coined _ant-push_.
We made available several off-the-shelf components444Available at
https://github.com/google/brax/tree/main/brax/experimental/composer/components
for rapid prototyping. These components were designed with flexibility in
mind, making them amenable for procedurally-generated continuous control tasks
or for learning modular policies [64, 81] that can control multiple
morphologies.

For a complete running example and a more in-depth introduction to Braxlines
Composer, check out the _Composer Basics_ 555Available at
https://colab.research.google.com/github/google/brax/blob/main/notebooks/composer/composer.ipynb
notebook.

⬇

from brax.experimental.composer import composer

\parenv_desc = dict(

components=dict( # component information

agent1=dict(

# pro_ant is a provided off-the-shelf component

# defined at components/pro_ant.py

component=’pro_ant’,

# among other things, designers can configure

# the ant morphology

component_params=dict(num_legs=6) # ant with 6 legs

),

cap1=dict(

# the singleton is the ball the ant controls

# defined in components/singleton.py

component=’singleton’,

component_params=dict(size=0.5), # a ball with radius 0.5

pos=(1, 0, 0), # where to place a capsule object

reward_fns=dict( # reward1: a GCRL task

goal=dict(

reward_type=’root_goal’,

sdcomp=’vel’,

target_goal=(4, 0, 0) # a target velocity for the object

)

)

)

),

edges=dict( # edge information

agent1__cap1=dict( # edge names use sorted component names

# observers extract relevant information for control

extra_observers=[ # add agent-object position diff as an extra obs

dict(observer_type=’root_vec’)

],

reward_fns=dict( # reward2: make the agent close to the object

dist=dict(reward_type=’root_dist’)

),

),

)

)

env = composer.create(env_desc=env_desc) # Composer returns a functional
environment