Learning Invariant Representations for Reinforcement Learning without Reconstruction
Authors: Amy Zhang
Rowan Thomas McAllister
Roberto Calandra
Yarin Gal
Sergey Levine
Venue: ICLR 2021 Oral (2021)
Keywords: rich observations
bisimulation metrics
representation learning
state abstractions

Abstract: We study how representation learning can accelerate reinforcement learning from rich observations, such as images, without relying either on domain knowledge or pixel-reconstruction. Our goal is to learn representations that provide for effective downstream control and invariance to task-irrelevant details. Bisimulation metrics quantify behavioral similarity between states in continuous MDPs, which we propose using to learn robust latent representations which encode only the task-relevant information from observations. Our method trains encoders such that distances in latent space equal bisimulation distances in state space. We demonstrate the effectiveness of our method at disregarding task-irrelevant information using modified visual MuJoCo tasks, where the background is replaced with moving distractors and natural videos, while achieving SOTA performance. We also test a first-person highway driving task where our method learns invariance to clouds, weather, and time of day. Finally, we provide generalization results drawn from properties of bisimulation metrics, and links to causal inference.

Published as a conference paper at ICLR 2021
LEARNING INVARIANT REPRESENTATIONS FOR REIN- FORCEMENT LEARNING WITHOUT RECONSTRUCTION
Amy Zhang112 Rowan McAllister
Roberto Calandra²
Yarin Gal
Sergey Levine
1 McGill University
2 Facebook Al Research
University of California, Berkeley
OATML group, University of Oxford
ABSTRACT
We study how representation learning can accelerate reinforcement learning from rich observations, such images, without relying either domain knowledge pixel-reconstruction. Our goal is to learn representations that provide for effective downstream control and invariance to task-irrelevant details. Bisimulation metries quantify behavioral similarity between states in continuous MDPs, which we pro- pose using to learn robust latent representations which encode only the task-relevant information from observations. Our method trains encoders such that distances in latent space equal bisimulation distances in state space. We demonstrate the effec- tiveness of our method at disregarding task-irrelevant information using modified visual MuJoCo tasks, where the background is replaced with moving distractors and natural videos, while achieving SOTA performance. We also test first-person highway driving task where our method learns invariance to clouds, weather, and time of day. Finally, provide generalization results drawn from properties of bisimulation metrics, and links to causal inference.
Introduction
Inrelevant
Relevant
Learning control from images is important for many real world applications. While deep reinforcement learning (RL) has enjoyed many successes in simulated tasks, learning control from real vision more complex, especially outdoors, where images reveal detailed scenes of complex and unstructured world. Furthermore, while many RL algorithms can eventually learn control from real images given unlimited data, data-efliciency is often a necessity in real trials which are expensive and constrained to real-time. Prior methods for fata-efficient learning of simulated visual tasks typically use representation learning. Representation learning summarizes images by encoding them into smaller vectored representations better suited for RI. For example, sequential autoencoders aim to learn lossless representations of streaming observations- sufficient to reconstruct current observations and predict future abservations-fivee various RL algorithms can be trained (Hainer et al., 2018 Lee al. 2019; Yarats et al.] 2019). However, such methods are task- agnostic: the models represent all dynamic elements they observe the world, whether they are relevant the task not. We argue such representations can easily "distract" RL algorithms with irrelevant information in the case of real images. The issues of distraction less evident in popular simulation MuJoCo and Atari tasks, since any change in observation space is likely task-relevant, and thus, worth representing. By contrast, visual images that autonomous cars observe contain predominately task-irrelevant information, like cloud shapes and architectural details, illustrated in Figure
Labent Space
Figure 1: tions of the visual scene should be insensitive to irrelevant objects (e.g., clouds) details (e.g.. car types), and encode observa- tions equivalently if their relevant details are equal (c.g., road direc- tion and locations of other cars).
Robust representa-
which
"qqual contribution. Corresponding author: amy xx zhang@mai magi 11

---

Published as a conference paper at ICLR 2021
Rather than learning control-agnostic representations that focus accurate reconstruction of clouds and buildings, we would rather achieve a more compressed representation from a lossy encoder, which only retains state information relevant to task. If we would like to learn representations that capture only task-relevant elements of the state and are invariant to task-irrelevant information, intuitively we can utilize the reward signal to help determine task-relevance, as shown by Donschkowski & Brock (2015). As cumulative rewards are our objective, state elements are relevant not only if they influence the current reward, but also if they influence state elements in the future that in turn influence future rewards. This recursive relationship can be distilled into a recursive tasre notion of state abstraction: an ideal representation is one that is predictive of reward, and also predictive of itself in the future.
We propose learning such an invariant representation using the bisimulation metrie, where the dis- tance between two observation encodings correspond to how "behaviourally different" (Ferns & Precup 2014) both observations are. Our main contribution a practical representation learning method based on the bisimulation metric suitable for downstream control, which we call deep bisimulation for control (DBC). We additionally provide theoretical analysis that proves value bounds between the optimal value function of the true MDP and the optimal value function of the MDP constructed by the learned representation. Empirical evaluations demonstrate our non- reconstructive approach using bisimulation is substantially more robust to task-irrelevant distractors when compared to prior approaches that use reconstruction losses or contrastive losses. Our initial experiments insert natural vidcos into the background of MoJoCo control task as complex distrac- tion. Our second setup is a high-fidelity highway driving task using CARLA (Dosovitskiy et al. 2017). showing that our representations can be trained effectively even on highly realistic images with many distractions, such as trees, clouds, buildings, and shadows. For example videos see https://sites. http:///esteessogoeee.o.oooo woamm/vieww httpg//e...oo/eee////eeep../.ee.
2 Related Work
Our work builds on the extensive prior research on bisimulation in MDP state aggregation.
Reconstruction-basel Representations. Early works deep reinforcement learning from im- ages (Lange & Riedmiller, 2010; Lange et al., 2012) used a two-step learning process where first an auto-encoder was trained using reconstruction loss to learn a low-dimensional representation, and subsequently a controller was learned using this representation. This allows effective leveraging of large. unlabeled datasets for learning representations for control. In practice, there no guarantee that the learned representation will capture useful information for the control task, and significant expert knowledge and tricks are often necessary for these approaches to work. In model-based RL, one solution to this problem has been to jointly train the encoder and the dynamics model end-to-end (Watter et al. 2015) Wahlström et al.] (015)) this proved effective in learning useful task-oriented representations, Hafner et al. (0018) and Lee et al. (2019) learn latent state models using a reconstruction loss, but these approaches suffer from the difficulty of learning accurate long-term predictions and often still require significant manual tuning. Gelada et al. (2019) also propose latent dynamics model-based method and connect their approach to bisimulation metrics, using a reconstruction loss in Atari. They show that lg distance in the DeepMDP representation upper bounds the bisimulation distance, whereas our objective directly learns a representation where distance in latent space the bisimulation metric. Further, their results rely on the assumption that the learned representation is Lipschitz, whereas we show that, by directly learning a bisimilarity-basee represen- tation, we guarantee a representation that generates a Lipschitz MDP. We show experimentally that our non-reconttrrettve DBC method is substantially more robust to complex distractors.
Lee
Contrastive- Representations. Contrastive losses are a self-supervised approach learn useful representations by enforcing similarity constraints between data (van den Oord et al.) 2018 Chen et al., 2020). Similarity functions can be provided as domain knowledge in the form of heuristic data augmentation, where we maximize similarity between augmentations of the same data point (Laskin et al. 2020) or nearby image patches (Hénaff et al. 2019), and minimize similarity between different data points. In the absence of this domain knowledge, contrastive representations can be trained by predicting the future (yan den Oord et al., 2018). We compare to such an approach in our experiments, and show that DBC is substantially more robust. While contrastive losses do not require reconstruction. they do not inherently have a mechanism to determine downstream task relevance without manual engineering, and when trained only for prediction, they aim to capture all
2

---

Published as a conference paper at ICLR 2021
predictable features in the observation, which performs poorly real images for the same reasons world models do. A better method would be to incorporate knowledge of the downstream task into the similarity function in a data-driven way, that images that are very different pixel-wise (e.g. lighting or texture changes), can also be grouped as similar w.r.t. downstream objectives.
Bisimulation. Various forms of state abstractions have been defined in Markov decision processes (MDPs) to group states into clusters whilst preserving some property (e.g. the optimal value, or all values, or all action values from each state) (Li et al.. 2006). The strictest form, which generally preserves the most properties, is bisimulation (Larsen & Skou 1989). Bisimulation only groups states that are indistinguishable w.r.t. reward sequences output given any action sequence tested. A related concept is bisimulation metrics (Femn & Precup 2014), which measure how "behaviorally similar" states are. Ferns et al. (2011) defines the bisimulation metric with respect to continuous MDPs, and propose a Monte Carlo algorithm for learning it using an exact computation of the Wasserstein distance between empirically measured transition distributions. However, this method does not scale well to large state spaces. Taylor et al. (2009) relate MDP homomorphisms to lax probabilistic bisimulation, and define a lax bisimulation metric. They then compute a value bound based on this metric for MDP homomorphisms, where approximately equivalent stae-actiion pairs aggregated. Most recently, CCattro (2020) propose an algorithm for computing on on-policy bisimulation metrics, but does directly, without learning a representation. They focus on deterministic settings and the policy evaluation problem. We believe our work is the first to propose a gradient-baser method for directly learning a representation space with the properties of bisimulation metrics and show that works in the policy optimization setting.
3 Preliminaries
We start by introducing notation and outlining realistic assumptions about underlying structure in the environment. Then, we review state abstractions and metrics for state similarity.
We assume the underlying environment is Markov decision process (MDP), described by the tuple M = (S, P, R, 7). where S is the state space, A the action space, (s''|s, a) the probability of transitioning from state E S to state s' E S, and y E [o, 1) a discount factor, An "agent" chooses actions a E AA according to policy function a ~ π (s), which updates the system state s' ~ P(ss, a). yielding a reward rewards by learning good policy: max, Ep (R:)s While our primary concern is learning from images, we do not address the aattial-bbservablii problem explicitly: we instead approximate stacked pixel observations the fully-observed system state s (explained further in Anppeeddix B).
(s) R. The agent's ' goal is to maximize the expected cumulative discounted
Bisimulation is form of state abstraction that groups states s; and S, that are "behaviorally equiv- alent" (Li et al. 2006). For any action sequence a0:00, the probabilistic sequence of rewards from s; and S, are identical. A more compact definition has a recursive form: two states are bisimilar if they share both the same immediate reward and equivalent distributions over the next bisimilar states (Larsen& & Skou 1989 Givan et al. 2003).
Definition 1 (Bisimulation Relations GGian et al. 2003)). Given an MDP M, an equivalence
relation B between states is a bisimulation relation if. for all states S, E S that are equivalent under B (denoted s; = s;) the following conditions hold:
R(s; a) = R(s,, a) Va E A, P(G|s,,
P(G|s,, a)
Va E A, VG E SB
(2)
where Sn is the partition of under the relation (the set of all groups of equivalent states), and P(G [ -Eveg P(s' a).
Exact partitioning with bisimulation relations is generally impractical in continuous state spaces, as the relation is highly sensitive to infinitesimal changes in the reward function or dynamics. For this reason, Bisimulation Metries (Fems al. 2011 Ferns & Precup 2014 Castro (020) softens the concept of state partitions, and instead defines a pseudometric space (S, d). where a distance function d: S x S + R20 measures the "behavioral similarity" between two states")
Defining distance d between states requires defining both a distance between rewards (to soften Equation (I). and distance between state distributions (to soften Equation (2).. Prior works use the Wasserstein metric for the latter, originally used in the context of bisimulation metrics by van Breugel
Note that d is pseudometric, meaning the distance between two different states can be zero, corresponding behavioral equivalence.
3

---

Published as a conference paper at ICLR 2021
& Worrell (2001). The pdh Wasserstein metric is defined between two probability distributions Pa and P, as WP,,, py; d) = ( Ssxs ((s,,,)) dy' (s, 8,)) where I(P;, P,) is the set of all couplings of P. and P,. This is known the "earth mover" distance, denoting the cost of transporting mass from one distribution to another (Villani 2003). Finally, the bisimulation metric is the reward difference added the Wasserstein distance between transition distributions:
Definition 2 (Bisimulation Metric), From Theorem 2.6 in Ferns et al. (2017) with E (0,1): d(s;, Sy) max (1 ()-c) c) [Rg. Rg + (I(Ps; 8: P dd
(3)
aeA
4 Learning Representations for Control with Bisimulation Metries
Raward Model
Enooder
Algorithm 1 Deep Bisimulation for Control (DBC)
Dynanics
Model
for Time / 0 co
2: Encode state 2, = o(s))
Replay
Reword Modia
3: Execute action a, ~ IT (Z+)
Encoder
4: Record data: D 4- D U {s:, ar, Sc+1, Fo41)
Model
Sample batch B, D
Permute batch: B3 = permute (B:)
Figure 2: Learning bisimulation metric represen-
7: Train policy: Bg, (Jx))
b Algorithm
tation: shaded in blue is the main model architecture,
8: Train encoder: Eg, B, [J()]]
Equation (4)
is reused for both states, like a Siamese network.
Train dynamics: J(P,)) (P(p(s) a Ze
The loss is the reward and discounted transition dis- tribution distances (using Wasserstein metric W).
We propose Deep Bisimulation for Control (DBC), a data-efficient approach to learn control policies from unstructured, high-dimensional states. In contrast prior work bisimulation, which typically aims to learn a distance function of the form d: S x S + R: between states, our aim is instead to learn representations Z under which & distances correspond to bisimulation metries, and then use these representations to improve reinforcement learning. Our goal is to learn encoders á - Z that capture representations of states that are suitable to control, while discarding any information that is irrelevant for control. Any representation that relies on reconstruction of the state cannot do this, as these irrelevant details are still important for reconstruction. We hypothesize that bisimulation metrics can acquire this type of representation, without any reconstruction.
Bisimulation metrics are useful form of state abstraction, but prior methods to train distance functions either do not scale to pixel observations (Ferns et al.) (001]) (due to the max operator in Equation (3). or were only designed for the (fixed) policy evaluation setting (Castro) 2020). By contrast, we learn improved representations for policy inputs, as the policy improves online. Our x -bisimulatiinn metric is learned with gradient decent, and we prove it converges to a fixed point in Theorem
I under some assumptions. To train encoder é towards our desired relation ((8,)]l1,, we draw batches of state pairs, and minimise the mean square error between the on-policy bisimulation metric and & distance in the latent space:
d(s; 8,) = [<x))
((b) = (11-- - -2111 - Ira - ryl - WW2 az), P(-12,, ay))
where z; ((s:)) 2, = ((s)), are rewards, and z denotes (ss) with stop gradients. Equation (4) also uses probabilistic dynamics model P which outputs Gaussian distribution. For this reason, we use the 2-Wasserstein metric W2 in Equation (4) as opposed to the I-Wasserstein in [qqation (3) since the W2 metric has convenient closed form: W2(V E,), NW,,, Σ,))² =1m My
- - 113, where || ||ا is the Frobenius norm. For all other distances we continue using the & norm. Our model architecture and training is illustrated by Figure 2 and Algorithm
Incorporating control. We combine our rep- resentation learning approach (Algorithm with the soft actor-critic (SAC) algorithm
Algorithm 2 Train Policy (changes to SAC in blue)
Get value: V V/=mine min..11 Q.(((s)) a log =og=((a)(s))
(Haarnoja et al., 2018) to devise practical
2: Train critics: JQQ,, () (o)() ~V)
3: Train actor: J(T) = a logp(alo(s)
-in.. Q:
reinforcement learning method. We modified
4: Train alpha: J(a) -a log p(a ((s))
SAC slightly in Algoriihh 2/to allow the value function to backprop to our encoder, which improve performance further (Yarats et al.) 2019; Rakelly et al.) 2019). Although, in principle, our method could combined with any RL
5: Update target critics: Q; TqQ, (1 -
6: Update target encoder: Φ + T6 + To) lb
algorithm, including the model-free DQN (Mnih 2015), model-based PETS (Chua et al.
4

---

Published as a conference paper at ICLR 2021
2018). Implementation details and hyperparameter values of DBC are summarized in the appendix, We train DBC by iteratively updating three components in turn: a policy π (in this case SAC), an encoder 6, and dynamics model P (lines 7-9, [Aloorith I). We found single loss function was less stable to train. The inputs of each loss function J(.) in Algorithm 1 rrepresents which components updated. After each training step, the policy x is used step in the environment, the data is collected in a replay buffer D, and a batch is randomly selected to repeat training.
Table 2
5
Generalization Bounds and Links to Causal Inference
While DBC enables representation learning without pixel reconstruction, leaves open the question of how good the resulting representations really are. In this section, we present theoretical analysis that bounds the suboptimality of a value function trained on the representation Iearned via DBC. First, we show that our x¹ -bisimulation metric converges to fixed point, starting from the initialized policy To and converging to an optimal policy π².
Theorem 1.
Let met be the space of bounded pseudometrics on S and a policy that is continuously improving. Define F mmet + met by
F(d, T (S;, 8j) = (1 -1-c))1a-rr]1++ - k.o PT,))
Then F has a least fixed point d which is a n:" -bisimulation metric.
Proof in appendix. As evidenced by DDeinition 2] the bisimulation metric has no direct dependence on the state space. Pixels can change, but bisimilarity will stay the same. Instead, bisimilarity is grounded in recursion future transition probabilities and rewards, which closely related to the optimal value function. In fact, the bisimulation metric gives tight bounds on the optimal value function with discount factor yy We show this using the property that the optimal value function is Lipschitz with respect to the bisimulation metric, see Theorem 5 Appendix (Ferns al. (004). This result also implies that the closer two states are in terms of d, the more likely they are to share the same optimal actions. This leads us generalization bound on the optimal value function of an MDP constructed from a representation space using bisimulation metrics, (lo(s.) - (s;)]] == d(s., 8,)- We construct partition of this space for some c 0, giving
partitions where 1 < (1 e)c. We denote & as the encoder that maps from the original state space S to each e-cluster. This denotes the amount of approximation allowed, where large e leads to a more compact bisimulation partition at the expense of looser bound on the optimal value function.
Theorem (Value bound based on bisimulation metries). Given an MDP M constructed by aggre- gating states in an c-neighborhood, and an encoder & that maps from states in the original MDP M to these clusters, the optimal value functions for the two MDPs are bounded as
2c
[V) - V* <
(6)
(1 (1-7) (1-7)(1-) c) c)
Proof in appendix. As 0 the optimal value function of the aggregated MDP converges to the original. Further, by defining a learning error for Φ, L == sups, [lo(s.) o(s,) - s,)|,
2c+2c we can update the bound in Theorem 2 to incorporate L: (s) - (d(s)) S (I->X(-a)'
MDP dynamics have a strong connection to causal inference and causal graphs. which are directed acyclic graphs (Jonsson & Barto| 2006) Schölkopf 2019 Zhang et al. 2020). Specifically, the state and action at time & causally affect the next state at time + 1. In this work, care about the components of the state space that causally affect current and future reward. Deep bisimulation for control representations connect to causal feature sets, the minimal feature needed to predict target variable (Zhang et al. 2020).
Theorem 3 (Connections causal feature sets (Thm 1 in Zhang al. (0200)) If we partition observations using the bisimulation metric, those clusters (a bisimulation partition) correspond to the causal feature set of the observation space with respect to current and future reward.
This connection tells us that these features are the minimal sufficient statistic of the current and future reward, and therefore consist of (and only consist of) the causal ancestors of the reward variable r.
Definition 3 (Causal Ancestors). In causal graph where nodes correspond to variables and directed edges between a parent node and child node C are causal relationships, the causal ancestors (C) of a node are all nodes in the path from C' to root node.
If there are interventions on distractor variables, or variables that control the rendering function and fherefore the rendered observation but do not affect the reward, the causal feature set will be
5

---

Published as a conference paper at ICLR 2021
robust to these interventions, and correctly predict current and future reward in the linear function approximation setting (Zhang et al. (2020). As an example, in autonomous driving, an intervention can be a change from day night which affects the observation space but not the dynamics reward. Finally, we show that representation based on the bisimulation metric generalizes to other reward functions with the same causal ancestors.
Theorem (Task Generalization). Given an encoder Ø +. Z that maps observations to a latent bisimulation metric representation where ([)(s;) o(s;)1l2 == d(s,.s)) Z encodes information about all the causal ancestors of the reward (R...
Proof in appendix. This result shows that the learned representation will generalize to unseen reward functions, as long the new reward function has a subset of the same causal ancestors. As an example, a representation learned for a robot to walk will likely generalize to learning to run, because the reward function depends forward velocity and all the factors that contribute forward velocity. However, that representation will not generalize to picking up objects, as those objects will be ignored by the learned representation, since they are not likely be causal ancestors of a reward function designed for walking. [Thorem shows that the learned representation will be robust to spurious correlations, or changes in factors that are not in (N) This complements Theorem 5| that the representation is a minimal sufficient statistic of the optimal value function, improving generalization over non-minimal representations.
Theorem (V* is Lipschitz with respect to d). Let V be the optimal value function for given discount factor y. Ifc 7; then V is Lipschitz confinuous with respect d with Lipschitz, constant I-c where d a n² -bisimullatoon metric.
1 T-c = dss 8z).
(s: VV
See Theorem 5.1 in Ferns et al. (2004) for proof. We show empirical validation of these findings in Section 6.2
6 Experiments
Our central hypothesis is that our non-reconstrucctiv bisimulation based representation learning approach should be substantially more robust to task-irrelevant distractors. To that end, we evaluate our method in a clean setting without distractors, as well as a much more difficult setting with distractors. We compare against several baselines. The first is Stochastic Latent Actor-Critic (SLAC, Lee et al. (2019)). state-of-tle tfate-of-the-are art method for pixel observations on DeepMind Control that learns dynamics model with a reconstruction loss. The second is DeepMDP (Gelada al., 2019), a recent method that also learns latent representation space using a latent dynamics model, reward model, and distributional Q learning, but for which they needed reconstruction loss to scale up to Atari. Finally, compare against two methods using the same architecture as ours but exchange our bisimulation loss with (1) reconstruction loss ("Reconstruction") and (2) contrastive predictive coding (Oord) et al. 2018) ("Contrastive") to ground the dynamics model and learn a latent representation.
6.1 Control with Background Distraction
In this section, we benchmark DBC and the previously described baselines on the DeepMind Control (DMC) suite (Tassa al. 2018) in two settings and nine environments (Figure 3). finger_spp spin. cheetah run, and wa ker wa wakk and additional environments in the appendix.
Default Setting. Here, the pixel observations have simple backgrounds as shown in Figure (top row) with training curves for our DBC and baselines. We see SLAC, recent state-of-the-an model-based representation learning method that uses reconstruction, generally performs best.
Simple Distractors Setting, Next, we include simple background distractors, shown in Figure 3 (middle row), with easy-to-predict motions. We use a fixed number of colored circles that obey the dynamics of an ideal gas (no attraction repulsion between objects) with no collisions. Note the performance of DBC remains consistent, as other methods start decreasing.
Natural Video Setting, Then, we incorporate natural video from the Kinetics dataset (Kay et al. 2017) as background (Zhang et al. 2018). shown in Figure 3 (bottom row). The results confirm our hypothesis: although number of prior methods learn effectively in the absence of distractors, when complex distractions are introduced, our non-reconstructive bisimulation based method attains substantially better results.

---

Published as a conference paper at ICLR 2021
DBC (ours)
Reconstruction
Contrastive
DeepMDP
SLAC
Figure 3: observations: Pixel observations DMC in default setting (top row) finger spin (left
column), cheetah (middle column), and walker (right column), with simple distractors (middle row), and natural video distractors (bottom row). Right training curves: Results comparing out DBC method to baselines on 10 seeds with standard error shaded in the default setting. The grid- grid-location of each graph corresponds to the grid-location each observation.
80
40
20
70
0
25.
-40
-50
60
-60
50
-80
100
-75
-50
-25
Figure t-SNE of latent spaces learned with bisimulation metric (left 1-SNE) and VAE (right t-SNE) after training has completed, color-coded with predicted state values (higher value yellow, lower value purple). Neighboring points in the embedding space learned with bisimulation metric have similar states and correspond observations with the same task-related information (depicted pairs of images with their corresponding embeddings), whereas no such structure is seen the embedding space learned by VAE, where the same image pairs are mapped away from each other.
To visualize the representation learned with bisimulation metric loss function in Equation (4) use a t-SNE plot (Figure 4p. We see that even when the background looks drastically different, our en- coder learns ignore irrelevant information and maps observations with similar robot configurations near each other. See AAppendix D for another visualization.
6.2 Generalization Experiments
We test generalization of our learned representation in two ways. First, we show that the learned representation space can generalize to different types of distractors, by training with simple distractors and testing on the natural video setting. Second, we show that our learned representation can be useful reward functions other than those it was trained for.
Generalizing over backgrounds. We first train on the simple distractors setting and eval- uate natural video. Figure 5 shows an example of the simple distractors setting and performance during training time of two experiments, blue being the zero-shot transfer to the natural video setting, and orange the baseline which trains natural video. This result empirically validates that the representations learned by DBC are able to effectively learn to ignore the background, regardless of what the background contains or how dynamic is.

---

Published as a conference paper at ICLR 2021
Generalizing over reward functions. We evaluate (Figure 5) the generalization capabilities of the learned representation by training SAC with new reward functions walker_ _sd and walker run using the fixed representation learned from walkers _aak. This is empirical evidence that confirms [Theormm 4] if the new reward functions are causally dependent on a subset of the same factors that determine the original reward function, then representation is sufficient.
S4C tranned on ebierration
SBC trained
3RC Trakedce dietic
BAC trtee with boom oec crcnder usted with trgaa Complite ancada
Trapper deal gasto linaka
Figure 5: Generalization of model trained simple distractors environment and evaluated on kinetics (left). Generalization of an encoder trained on walker_wald enviroument and evaluated on walker_se berkkrrstadd (center) and wa 1 ker_rur (right), all the simple distractors setting. seeds, standard error shaded.
6.3
Comparison with other Bisimulation Encoders
Even though the purpose of bisimulation metries by Castro (020) is learning distances d, not representation spaces 2, it nevertheless implements d with function approximation: ((s),,)) a (ó(s;), o(s,)) by encoding observations with & before computing distances with v), trained as: J(, v)- = (d(s), ((s)) - Ir rryl -r (P((ss,, (s:))) (P((s), T ((s))))))
where à and s are target networks. natural question is: how does the encoder & above perform in control tasks? We com- bine ó above with our policy Algooriim 2and use the same network v (single hidden layer 729 wide). Figure 6jshows rep- resentations from Castro (2020) can learn control (surprisingly well given it not designed to), but our method learns faster. Further, our method simpler: by comparing Equation (8) to Equuaiion (4) our method uses the & distance between the encoding instead of introducing an addition network eh.
we her wa with returs dite
Figure 6: Bisim. results. Blue is DBC
6.4 Autonomous Driving with Visual Redundancy
and orange is Castro (2020).
Real-world control systems such robotics and autonomous vehicles must contend with huge variety of task-irrelevant information, such as irrelevant objects (e.g. clouds) and ir- relevant details (e.g. obstacle color). To eval- uate DBC tasks with more realistic obser- vations, we construct a highway driving sce-
Figure 7: The driving task is to drive the red ego (left) safely in traffic (middle) along highway (right). nario with photo-realistic visual observations using the CARLA simulator (Dosovitskiy al.) 2017) shown in Figure 7| The agent's goal is to drive as far as possible along CARLA'S Town04's figure-8 8 the highway in 1000 time-steps without colliding into the 20 other moving vehicles or barriers. Our objective function rewards highway progression and penalises collisions: rs ego Uhighway At- X; impulse >, |steer|, where Vego is the velocity vector of the ego vehi- elc, projected onto the highway's 's unit vector ünighway. and multiplied by time discretization At = 0.05 measure highway progression in meters. Collisions result impulses Rt, measured in Newton- seconds. We found steering penalty steer E [-1, 1, 1] helped, and used weights A; 10- and A 1. While more specialized objectives exist like lane-ng this experiment's purpose is only to compare representations with observations more characteristic of real robotic tasks. We use five cameras on the vehicle's 's roof, each with 60 degree views. By concatenating the images together, our vehicle has a 300 degree view, observed as 84 x 420 pixels. Code and install instructions in appendix.
Results in Figure compare the same baselines as before, except for SLAC which is easily distracted used SAC, which does not explicitly learn representation, but performs
(Figure 3). Instead surprisingly well from raw images. DeepMDP performs well too, perhaps given its similarly to bisimulation. But, Reconstruction and Contrastive methods again perform poorly with complex
8

---

Published as a conference paper at ICLR 2021
images. More intuitive metrics are in Table and Figure 8 depicts the representation space a t-SNE
with corresponding observations. Each run took 12 hours on GTX 1080 GPU.
Figure 8: A t-SNE diagram of encoded first-person driving observations after 10k training steps of Algorithm color coded by value (V in Algorithm 2). Top: the learned representation identifies obstacle the right side. Whether that obstacle is dark wall, bright car, or truck is lask-irrelevant: these states are behaviourally equivalent. Left: the ego vehicle has flipped onto its left side. The different wall colors, due to a setting sun, irrelevant: all states are equally stuck and low-value (purple t tNS color). Right: clear highway driving. Clouds and sun position are irrelevant.
Figure 9: Performance comparison with 3 seeds the driving task. Our DBC method (red) performs better than DeepMDP (purple) or learning direct from pixels without representation (SAC, green). and much better than contrastive methods (blue). Our method's ' final performance 46.8% hetter than the next hest baseline.
Contrastive
175
Reconstruction
150
SAC
apsede reward
125
DeepMDP
DBC (ours)
100
Table 1: Driving metries, averaged over 100 episodes. after 100k training steps, with standard error. Arrow direction indicates if metric desired larger or smaller.
SAC
DeepDDF
DBC (eurs)
distance
123.2
intensity
4004
1958 15.6
2677 38.5
20000
40000
60000
80000
100000
average soces
16.0% #: 0.19%
10.4% = 0.013%
7.8% = 0.12%
step
average beake
1.33%
+ 0.00%
4.3% + 0.033%
1.6% + 0.022%
Discussion
This paper presents Deep Bisimulation for Control: new representation learning method that considers downstream control. Observations are encoded into representations that are invariant to different task-irrelevant details in the observation. We show this is important when learning control from outdoor images, or otherwise images with background "distractions". In contrast to other bisimulation methods, we show performance gains when distances in representation space match the bisimulation distance between observations.
Future work: Several options exist for future work. First, our latent dynamics model P was only used for training our encoder in [Equation (4) but could also be used for multi-step planning in latent space. Second, estimating uncertainty could also be important to produce agents that can work in the real world, perhaps via ensemble of models { to detect-and and adapt to -o-iitriibutional shifts between training and test observations. Third, an undressed issue is that of partially observed settings (that assumed approximately full observability by using stacked images), possibly using explicit memory or implicit memory such as an LSTM. Finally, investigating which metrics (L. or 1.2) and dynamics distributions (Gaussians not) would be beneficial.

---

Published as a conference paper at ICLR 2021
References
Pablo Samuel Castro.
Scalable methods for computing state similarity in deterministic Markov decision processes. In Association for the Advancement of Artificial Intelligence (AAAI), 2020.
Ting Chen, Simon Kornblith, Mohammad Norouzi, and Geoffrey Hinton. A Simple Framework for Contrastive Learning Visual Representations. arXiv:2002.0577 05000 (cs, stat], February 2020. URL
007009 arXiv: 2002.05709.
Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine. Deep reinforcement learning in handful of trials using probabilistic dynamics models. In Neural Information Processing Systems (NeurIPS), pp. 4754-4765, 2018.
Alexey Dosovitskiy, German Ros, Felipe Codevilla, Antonio Lopez, and Vladlen Koltun. CARLA: An open urban driving simulator. arXiv preprint arXiv: 2017.
Simon S. Du, Akshay Krishnamurthy, Nan Jiang, Alekh Agarwal, Miroslav Dudík, and John Langford. Provably efficient RL with rich observations via latent state decoding. Computing Research Repository (CoRR),
2019. TRLLhtpp:22 pttttp//arxiv> wrg/ htt/aa///190..001111
Norm Ferns, Prakash Panangaden, and Doina Precup. Metrics for finite Markov decision processes. In Uncertainty in Artificial Intelligence (UAI), pp. 1696669,96266 2004. ISBN 0-9749039-0-6. URL
http: /ttt:///d... acm.
c wtt.:ioo....mm.m1133333333333 1036886
Norm Ferns, Prakash Panangaden, and Doina Precup. Bisimulation metrics for continuous Markov
decision processes. Society for Industrial and Applied Mathematics, 4006): December 2011. ISSN 0097-5397. doi: 10. 10.1137/10080484X, URL https:
10080484X
Norman Ferns and Doina Precup. Bisimulation metrics optimal value functions. In Uncertainty in Artificial Intelligence (UAI), pp. 210-219, 2014.
Carles Gelada, Saurabh Kumar, Jacob Buckman, Ofir Nachum, and Marc G. Bellemare. DeepMDP: Learning continuous latent space models for representation learning. In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), International Conference Machine Learning (ICML), volume 97, pp. 2170-2179, Jun 2019.
Robert Givan, Thomas L. Dean, and Matthew Greig. Equivalence notions and model minimization in Markov decision processes. Artificial Intelligence, 147:163-22 223,, 2003.
Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critic: Off-policy maxi- mum entropy deep reinforcement learning with a stochastic actor. arXiv preprint wxXiv:801.0019%0 2018.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. arXiv preprint 1rXiv:1811.064 1.04511, 2018.
Olivier J Hénaff, Aravind Srinivas, Jeffrey De Fauw, Ali Razavi, Carl Doersch, SM Eslami, and Aaron van den Oord. Data-efficient image recognition with contrastive predictive coding. arXiv preprint 1
2019.
Rico Jonschkowski and Oliver Brock. Learning state representations with robotic priors. Autonomous Robots, 39(3))4077-44 2015.
Anders Jonsson and Andrew Barto. Causal graph based decomposition of factored MDPs. J. Mach. Learn. Res., 7:2259 -2001, December 2006. ISSN 1532-4435.
Will Kay, João Carreira, Karen Simonyan, Brian Zhang, Chloe Hillier, Sudheendra Vijayanarasimhan, Fabio Viola, Tim Green, Trevor Back, Paul Natsev, Mustafa Suleyman, and Andrew Zisserman. The kinetics human action video dataset. Computing Research Repository (CoRR), 2017. URL
http:://aret http://wwww..ooog///////////// www.:/orggg./oooeoo 06950
10

---

Published as a conference paper at ICLR 2021
Sascha Lange and Martin Riedmiller. Deep auto-encoder neural networks in reinforcement learning. In International Joint Conference on Neural Networks (IJCNN), pp. 1-8. IEEE, 2010.
Sascha Lange, Martin Riedmiller, and Amn Voigtländer. Autonomous reinforcement learning on raw visual input data in real world application. In International Joint Conference Neural Networks (IJCNN), pp. 1-8, 2012. doi: 10.
K. G. Larsen and A. Skou. Bisimulation through probabilistic testing (preliminary report). In Symposium on Principles of Programming Languages, pp. 344 352. Association for Computing Machinery, 1989. ISBN 0897912942. doi: 10.1 10.1145/75277.753000
httpsts::://do org/
10.1145/7522 75307
Michael Laskin, Aravind Srinivas, and Pieter Abbeel. CURL: Contrastive unsupervised representa- tions for reinforcement learning. arXiv:2003.06417, 2020.
Alex X Lee, Anusha Nagahandi, Pieter Abbeel, and Sergey Levine. Stochastic latent actor-critic: Deep reinforcement learning with a latent variable model. arXiv preprint arXiv:1907.009:: 2019.
Lihong Li, Thomas J Walsh, and Michael L Littman. Towards unified theory of state abstraction for MDPs. In ISAIM, 2006.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A. Rusu, Joel Veness, Mare G. Bellemare, Alex Graves, Martin Riedmiller, Andreas K. Fidjeland, Georg Ostrovski, Stig Pe- tersen, Charles Beattie, Amir Sadik, Ioannis Antonoglou, Helen King, Dharshan Kumaran, Daan Wierstra, Shane Legg, and Demis Hassabis. Human-level control through deep rein- forcement learning. Nature. 518(7540):529 533, February 2015. ISSN 00280836. URL httt::/dx.do.: org.ioog//11 10. 1038 www.1///////ttre/eeee 42366
Aaron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. arXiv preprint arXiv: 1807.0378, 2018.
Kate Rakelly, Aurick Zhou, Deirdre Quillen, Chelsea Finn, and Sergey Levine. Efficient off- off-co meetareeinforcement learning via probabilistic context variables, arXiv preprint arXiv: 1903.0900.0880 082544 2019.
Bernhard Schölkopf. Causality for machine learning, 2019.
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. Technical report, DeepMind, January 2018. URL https
00690
Jonathan Taylor, Doina Precup, and Prakash Panagaden. Bounding performance loss in approximate MDP homomorphisms. In Neural Information Processing (NeurIPS), pp. 1649-1656, 2009.
Franck van Breugel and James Worrell. Towards quantitative verification of probabilistic transition systems. In Fernando Orejas, Paul G. Spirakis, and Jan van Leeuwen (eds.). Automata, Languages and Programming, pp. 422-432. Springer, 2001. ISBN 978-3-540-48224-6. doi: 10.1007/
3-540-48224-5
Aäron den Oord, Yazhe Li, and Oriol Vinyals. Representation learning with contrastive predictive coding. ArXiv, abs/1807 0///1807.03774 03748, 2018.
Cédric Villani. Topics in optimal transportation. American Mathematical Society, 01 2003.
Niklas Wahlström, Thomas Schön, and Marc Deisenroth. From pixels to torques: Policy learning with deep dynamical models. arXiv preprint arXiv: 550.00251, 2015.
Manuel Watter, Jost Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed control: A locally linear latent dynamics model for control from images. In Neural Information Processing Systems (NeurIPS), pp. 2728-2736, 2015.
Denis Yarats and Ilya Kostrikov. Soft actor-critic (SAC) implementation in PyTorch. https 2020.
7/gi //github.r www.wo/w./eeeeee http:///eeeeiisaaaaastpptttorooa

---

Published as a conference paper at ICLR 2021
Denis Yarats, Amy Zhang, Ilya Kostrikov, Brandon Amos, Joelle Pineau, and Rob Fergus. Im- proving sample efficiency in model-free reinforcement learning from images. arXiv preprint arXiv:1910.01741, 2019.
Amy Zhang, Yuxin Wu, and Joelle Pineau. Natural environment benchmarks for reinforcement learning. Computing Research Repository (CoRR), abs/1811 1.00022 2018. URL http:////aaxx
axg/abs/1811. 06032
Amy Zhang, Clare Lyle, Shagun Sodhani, Angelos Filos, Marta Kwiatkowska, Joelle Pineau, Yarin Gal, and Doina Precup. Invariant causal prediction for block mdps. In International Conference on Machine Learning (ICML), 2020.
12

---

Published as a conference paper at ICLR 2021
A Additional Theorems and Proofs
Theorem I Let met be the space of bounded pseudometries on S and æ a policy that is
continuously improving in the space of policies IL Define F : met x 11 + met by
F(d, (s,, 8,) = (1 c)(r=. - cd) 8, P
Then F has a least fixed point d which is nt -isimmllation metric.
Proof. Ideally, to prove this theorem we show that F is monotonically increasing and continuous, and apply Fixed Point Theorem to show the existence of a fixed point that F converges to. Unfortunately, we can show that F under π as monotonically converges to π is not also monotonic, unlike the original bisimulation metric setting (Ferns al. 2004) and the policy evaluation setting (Castro) 2020). We start the iterates F from bottom I, denoted as F (L). In Ferns et al. (2004) the maXac can be thought of as learning policy between every two pairs of states to maximize their distance,
and therefore this distance can only stay the same or grow over iterations of F. In Castro (2020), is fixed, and under deterministic MDP it can also be shown that distance between states d,, (8., will only expand, not contract increases. In the policy iteration setting, however, with æ starting from initialization To and getting updated:
= arg max E'ES
Σ Tras + +7V*E-1(s')
(10)
there is no guarantee that the distance between two states d (s., 8,) < d (s;, sy) under policy iterations Tk-1, and distance metric iterations da-1, dn for k, n E N, which is required for monotonicity.
Instead, we show that using the policy improvement theorem which gives us V** (s) 2 V (s), Vs E S, will converge to fixed point using the Fixed Point Theorem, and taking the result by Castro (2020) that J has a fixed point for every æ E II, we can show that fixed point bisimulation metric will be found with policy iteration.
(11)
Theorem Given a new aggregated MDP M constructed by aggregating states an e- neighborhood, and an encoder that maps from states in the original MDP M to these clusters, the optimal value functions for the two MDPs are bounded as
2c
V S (1->)( (1-7)(1-))
-
(12)
Proof. From Theorem 5.1 in Ferns al. (2004) we have:
max g(u, d) (1 (1-c)lV v)lV' V (ø(s)) S g(s, d) 1- wES
where g the average distance between state and all other states in its equivalence class under the bisimulation metric d. By specifying e-neighborhood for each cluster of states can replace g:
7 2c
(1 - v) (s) - V (o(s)) < 2c +
1-7
VV*(s) ()(s)) <
1
(2c +
7
2c)
1 c
2c
(1 - (1-7)(1-) cc)
Theorem Given an encoder á S Z that mups observations to latent bisimulation metric representation where [[ő(s:) ((s;)]l d(s.. s,), Z encodes information about all the causal ancestors of the reward (N(R).
Proof. We assume a MDP with a state space S (S .... SKy that can be factorized into K variables with 1-step causal transition dynamics described by a causal graph & (example in Figurr 10). We break the proof up into two parts: 1) show that if factor S' & AN(R) changes, the bisimulation distance between the original state s and the new state s' is O. and 2) show that factor S E AN(R) changes, the bisimulation distance be > 0.
13

---

Published as a conference paper at ICLR 2021
S!
Figure 10: Causal graph of transition dynamics, Reward depends only s causal parent, but s! causally depends on s², so AN(R) is the set , s²).
Sf & (R),,, an intervention on that factor does not affect current or future reward. s,) max(1 c)lr r. cW (d) Pa
max cW (d) P.,) and sy have the same reward. acAA
S does not affect future reward, then states s; and s, will have the same future reward conditioned on all future actions. This gives us
d(s.s') =
2) there is an intervention on S AN (R) then current and/or future reward can change. current reward changes, then we already have maXaed - c)Ir2, - > > 0, giving us ((s,,s)) > only future reward changes, then those future states will have nonzero bisimilarity, and max W (d) po > 0, giving us d(ss; s,) >
B Definition of State
Since we are concerned primarily with learning from image observations, we could explicitly distinguish the image observation space O from an unknown state space S. However, since not tackling the general POMDP problem, we consider the Block MDP (Du et al. (019), which assumes the state space is latent, and that and rendering function S + O. The crucial assumption that distinguishes the Block MDP from partially observable MDPs is the following:
instead given access to observation space O
Assumption 1 (Block structure (Du et al. 2019)). Each observation uniquely determines its generating state s. That is, the observation space can be partitioned into disjoint blocks O., each containing the support of the conditional distribution q(oss).
This assumption gives the Markov property in the observation space E O. As an example, one can think of the proprioceptive state consisting of positions and velocities of actuators as underlying state, and stacked pixel observations from a specific camera angle as a particular rendering function and corresponding observation space.
C Additional DMC Results
In Figure 11 show performance on the default setting 9 different environments from DMC. Figures 12 and 13 give performance on the simple distractors and natural video settings for all environments.
14

---

Published as a conference paper at ICLR 2021
Togerap3
Figure 11: Results for DBC in the default setting, in comparison baselines with reconstruction loss, contrastive loss, and SLAC on 10 seeds with standard error shaded.
Figure 12: Results for DBC in the simple distractors setting, in comparison to baselines with
reconstruction loss, contrastive loss, DeepMDP, and SLAC on 10 seeds with standard error shaded.
15

---

Published as a conference paper at ICLR 2021
1RU5T 1980
1800 with rali WiRl
SUA
Figure 13: Results for our bisimulation metric method in the natural video setting, in comparison to baselines with reconstruction loss, contrastive loss, DeepMDP, and SLAC on 10 seeds with standard error shaded.
D Additional Visualizations
In addition to Figure 4 shown the far left of Figure 14 Note the robot agent is quite crisp, which means neighboring points encode the agent similar positions, but the backgrounds are very different, and so are blurry when averaged.
also took 10 nearby points the t-SNE plot and average the observations,
80
60
20
0
-20
440
-60
-80
-75
-50
-25
25
50
75
100
Figure 14: t-SNE of latent spaces learned with bisimulation metric after training has completed, color-coded with predicted state values (higher value yellow, lower value purple). Neighboring points (right) in the embedding space leamned with a bisimulation metric have similar encodings (middle). When we sample from the same latent point, and average the images, agent similar positions, but the backgrounds are very different, and so blurry when averaged.
the robot configuration crisp, meaning neighboring points encode the
16

---

Published as a conference paper at ICLR 2021
E
Implementation Details
We use the same encoder architecture in Yarats et al. (2019), which is an almost identical encoder architecture as Tassa al. 2018), with two more convolutional layers to the convnet trunk. The encoder has kernels of size 3 x 3 with 32 channels for all the convolutional layers and set stride to everywhere, except of the first convolutional layer, which has stride 2, and interpolate with ReLU activations. Finally, add tanh nonlinearity to the 50 dimensional output of the fully-connected layer.
For the reconstruction method, the decoder consists of fully- connected layer followed by four deconvolutional layers. We use ReLU activations after each layer, except the final deconvolutional layer that produces pixels representation. Each deconvolutional layer has kernels of size x 3 with 32 channels and stride 1, except of the last layer, where stride is 2.
The dynamics and reward models are both MLPs with two hidden layers with 200 neurons each and ReLU activations.
Soft Actor Critic (SAC) (Haarnoja al. 2018) an off-policy actor-critic method that the maximum entropy framework for soft policy iteration. At each iteration, SAC performs soft policy evaluation and improvement steps. The policy evaluation step fits a parametric soft Q-function Qs,, ar) using transitions sampled from the replay buffer DD by minimizing the soft Bellman residual,
J(Q) (.p..ee. ( (olst. a.) >V (s.+1)
(13)
The target value function V is approximated via Monte-Carlo estimate of the following expectation,
((s+1)) Em., ($s++1 a(41)) - log (aa++1/st+1]]
(14)
where Q is the target soft Q-function parameterized by weight vector obtained from an exponentially moving average the Q-function weights to stabilize training. The policy improvement step then attempts project parametric policy (a))s; by minimizing KL divergence between the policy and Boltzmann distribution induced by the Q-function, producing the following objective,
((T) = ES
E Ea (ar log(lo (at((a))))) - Q(s:, [(ss,aa)]
(15)
We modify the Soft Actor- PyTorch implementation by Yarats & Kostrikov (2020) and augment with shared encoder between the actor and critic. the general model fa and task-specific models fg. The forward models are multi-layer perceptions with ReLU non-lincarities and two hidden layers of 200 neurons each. The encoder a linear layer that maps to a 50-dim hidden representation. The hyperparameters used for the RL experiments are in Table 2
Parameter name
Value
Replay buffer capacity
Batch size
128
Discount
0.99
Optimizer
Adam
Critic learning rate
10-5
Critic target update frequency
2
Critic Q-function soft-update rate TQ
0.005
Critic encoder soft-update rate T
0.005
Actor learning rate
10-5
Actor update frequency
2
Actor log stddev bounds
|-5, 2]
Encoder learning rate
10-5
Decoder learning rate
10-5
Decoder weight decay Temperature learning rate Temperature Adam's Br
10-7
10 4
0.9
Init temperature
0.1
Table A complete overview used hyper parameters.
17


Reviews:
[{"review_id": "RbVptpSUFSy", "reviewer": "AnonReviewer1", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "The authors propose an approach to robust representation learning of observations for reinforcement learning by training a model to align the euclidean distance between two observations with bisimulation metrics that quantify how similar the states that generated the observations are in terms of the control problem.  This reduces the effect of irrelevant features in the observations on the representations.\n\nThe paper is well written, the problem is clearly motivated and the approach and technical contribution is easy to follow.\n\nThe approach to use the state bisimulation metric to supervise observation representation is intuitive and clearly motivated.  Theoretical analysis is provided with generalization guarantees.\n\n\"As an example, in the context of autonomous driving, an intervention can be a change in weather, or a\nchange from day to night which affects the observation space but not the dynamics or reward.\"  I do not agree with the example as weather can directly alter the dynamics and desired behavior of an AV system.  The point in this paragraph is still clear but I would suggest a different example.\n\nThe evaluations are strong and run on a number of different experiment settings against multiple strong SOTA models.\n\nIn Figure 4 the proposed approach is outperformed by contrastive learning in the default setting.  I understand that the goal is to learn robust representations for the natural setting, but can the authors comment on why it fails to beat the contrastive approach here and provide insight on how this may be addressed.  Some problems have only a few distractors and may fall between the natural and the default setting.\n\nrecommendation and reasoning\n\nThe paper is well written and presents a clear approach to a well motivated problem with strong evaluation results.  I recommend acceptance.", "extra_scores": {}, "extra_text": {}}, {"review_id": "6PJ0jwCtAZ2", "reviewer": "AnonReviewer2", "summary": "", "questions": "", "limitations": "", "rating": 9, "confidence": 5, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "# Paper Summary \n\nThe paper presents a new method for embedding visual images into a state space suitable for effective control by an actor-critic style RL algorithm. They show how a previously explored idea of using a bisimulation between state abstractions and reward sequences to group states that are similar from a decision theoretic perspective can be extended to a continuous deep embedded representation using twin network style learning through standard gradient descent optimization. They call the approach Deep Bisimulation for Control (DBC).  The paper also argues for the correctness of their approach using a contraction proof and a theoretical argument for generalization to new problem domains. The approach contrasts directly with algorithms that use an autoencoder to find a compact representation by reconstructing input frames or predicting future input frames. These approaches necessarily represent enough information to reconstruct both task relevant and incidental details. Experimentally, the paper shows that performance of reconstruction-based approaches degrades by a large and significant amount when extraneous background detail is present in image frames, while the proposed method is immune. The evaluation is done on widely respected benchmark of Deep Mind MuJoCo based articulated figure simulations and a CARLA realistic image simulation of car driving which shows that focusing on task specific detail is important on realistic tasks. \n\n# Pros and Cons \n\nThe work addresses a key problem in reinforcement learning, which is learning effective policies from unlabeled visual images.  The ability to find compact, task relevant representations of high-dimensional inputs is central to the ICLR community. \n\nThe paper covers relevant background in state abstraction in RL and it is clear how it relates to the paper's contribution.  \n\nThe paper clearly explains the background of bisimulation and how the idea of grouping states according to their probable reward sequences can be practically realized by grouping states by their immediate reward and future state distributions. It is also clear in principle that a twin network can be used to induce a latent space with these properties.  \n\nThe paper shows that the policy will converge and that the error will be bounded and that the representation will generalize to any problem domain in which causal dependencies are a subset of the problem domain the representation was trained on. Which is nice.  \n\nThere is a nice evaluation on both MuJoCo articulated figure problems from the Deep Mind Control suite as well as the CARLA simulated driving application that uses large realistically rendered images. Comparison against state-of-the-art RL algorthims makes results convincing. I think the CARLA example is nice as it shows that practical problems without artificial augmentations have the property that the input details can be distracting if fully reconstructed. The superiority of DBC in this case really makes this point well. \n\nThe paper does not report computational load associated with their approach. Presumably due to the closed form Wasserstein approximations, we are probably looking at only a fractional increase in time for the encoder network above the policy, dynamics and reward models. So roughly 30% more?? \n\nKey figures for central results are too small to get meaningful interpretations unless enlarged by a factor of 4 or 5. Somehow this seems to go against the spirit of page limits to me. \n\nFigure 6 and the description don't seem to be aligned, but I can imagine that it can be readily fixed.\n \n\n# Recommendation  \n\nI recommend acceptance. The paper shows how to extend bisimulation principle for grouping states to continuous deep actor critic methods and provides convincing evidence on standard benchmarks that it is effective in extracting a task-relevant abstraction that is robust to noise in observations and focuses on task specific detail. \n\n \n\n# Questions  \n\nFigure 6: it seems that the “simple_distractors” environment is different than “setting 2” natural video setting. This caption could use some reworking to get parallelism clear. The ‘ideal gas’ is the same as simple_distractors?  Also the graphs seems to be about different experimental types. Does the first experiment also use frozen encoder?  The graphs do not seem to relate to the text which talks about walker_stand, walker_run reward functions. I am confused. \n\nCan the paper have any insight into the relative performance of their algorithm versus benchmarks algorithms across different tasks (finger spin, cheetah, walker) given that they are all stick figures? \n\n# Feedback  \n\nI was able to work through Definition 1 and assure myself it made sense. In the end I made a small picture that helped. \n\nThere seems to be some ambiguity in notation between observations, underlying state and latent variable spaces.  For instance, in section 3, script S is defined as a state space. In section 4, the function d is defined on script S x script S which is described as an observation space.  Admittedly, the work seems to be situated in an approximately fully observable world in which states and observations are somewhat equivalent. I suspect that is why the 5-camera 300-degree suround view was necessary in the CARLA experiments. \n\nInteresting that the paper employ stop gradients on the latent representation terms when they appear in the reward and Wasserstein terms in the loss function. This is to enforce the separateness of the optmizations in algorithm 1? \n\nDefinition 2, the bisimulation metric contains a max … so this is worst case discrepancy between the futures between state space actions and empirical actions. Was average discrepancy considered? This might be relevant later in discussion about intractability in previous approaches (section 4 paragraph 3). \n\nEquation 4 specifies an L1 metric for the distance in abstraction space ||z_i – Z_j ||. Is there a motivation for this choice? Again trying to bound the worst error? \n\nTheorem 4 references Theorem 5 which is not in the main text. \n\nIn section 5, ideally epsilon would be briefly described before it appears in a bound. \n\nFigure 3 is very small... the labels – particularly the subscripts are unreadable. It may not be making a point important enough to include it. The causal variables section could be shorted in general. It is a good point but not completely surprising. \n\nFigure 4 is a key figure to support the paper's hypothesis that deep bisimulation is effective for state abstraction in noisy images. This figure really needs to be bigger. In particular, I had to strain to read the legends to understand if the axes were different between the uncluttered video of articulated figures and the cluttered video of articulated figures with a background movie in them.  It was also hard to make out which line was which.  In particular, the “cheetah” column does not show DBC improving on cluttered video scenario --- it tops out at around 250 in both top and bottom graphs, but the scales are very different. It confuses the message a bit.  \n\nThere is a statement “a single loss function would be less stable and require balancing the components”. Is this speculation or based on experience? Separate optimizations implicitly define a balance between these terms wouldn’t they? Namely equal balance? \n\nTo some degree, figure 9 is making the same argument as figure 5. Could leave this out if you were tight for space. \n\nFigure 5: What is the first column of figures to the left? It doesn’t seem to be relevant? Otherwise, I think the figure is very effective in conveying the structured embedding space does a better job of grouping similar states. It is also very small. Drawing a white border between the figure pairs would underscore visually that there are two states in each “image”  \n\nFigure 9: I could not make sense of the images on the sides of the figure. Particularly, the images on the left side.  They seemed to be abstract geometric shapes and I could not get an intuition about what the driving scenario was. \n\nSection 6.4 “reward highway progression an penalizes collisions”  an => and \n\nFuture work – another likely avenue for future work would be to introduce some sort of memory to handle partially observable worlds. For instance, can the agent drive a car with only a forward view if given memory? Does this break down if it does not have memory? This could either be an explicit neural memory or implicit memory such as an LSTM …  Estimating uncertainty could also be important to produce agents that can work in the real world and assess when they know what they are doing and when they do not. Another future work area is in modeling of transition distributions as something more complex than Gaussians … what if there are distinct possible futures that are equally valid: it is ok for the robot to turn left or right as long as it avoids the object straight ahead? ", "extra_scores": {}, "extra_text": {}}, {"review_id": "5slSBViBBY4", "reviewer": "AnonReviewer3", "summary": "", "questions": "", "limitations": "", "rating": 7, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "main_review": "", "summary_of_the_paper": "", "summary_of_the_review": "", "strength_and_weaknesses": "", "limitations_and_societal_impact": "", "review_text": "**Summary**\nThe paper focuses on how learning state-representations that encode information relevant to the task can improve reinforcement learning from pixels. Often, observations in an MDP can contain information that are irrelevant (“distractors”) to the task at hand and can likely “distract” the downstream RL algorithm used. Unlike existing reconstruction based approaches (which don’t explicitly incentivize ignoring task-irrelevant information), the authors propose Deep Bisimulation Control (DBC) that relies on bi-simulation metrics (as the task-aware criterion) that encode behavioral similarity b/w states with respect to the reward structure. Instead of explicitly learning a bi-similarity distance function, authors enforce the representations which under L1 distances correspond to bi-simulation metrics. DBC demonstrates learning these representations in conjunction with the control policy, reward and a dynamics model. Furthermore, the authors highlight connections to causal inference which can hopefully further provide insights into which “new” reward structures can the learned representations generalize to (since bi-similarity metrics themselves are heavily dependent on the reward structure). Results obtained by the authors demonstrate that DBC can learn task specific representations and a control policy in a robust manner in the presence of distractors on the Deepmind Control Suite and the CARLA simulator. Additionally, the authors also demonstrate how DBC can generalize to new reward functions on Mujoco.\n\n**Strengths**\n\n- The paper is generally well-written and easy to follow for the most part. The authors do a good job of walking the reader through the preliminaries, highlighting distinctions with prior usage of bisimilarity metrics and how DBC ties in with slight modifications to the Soft Actor-Critic algorithm.\n- The results obtained on DM control suite demonstrate that while DBC is competitive or slightly worse compared to other approaches in the absence of distractors, it performs better compared the set of baselines when distractors are introduced — significantly outperforming the reconstruction and the contrastive approaches. Furthermore, compared to a VAE, DBC places observations that are similar in terms of task information closer in the learnt embedding space.\n- The authors also show that compared to an adaptation of prior usage of bisimulation metrics, DBC is more sample efficient. Additionally, within the considered family of reward functions to generalize to “walker_run” and “walker_stand”, DBC shows stronger generalization compared to DeepMDP. Furthermore, learned embeddings on CARLA demonstrate qualitatively that DBC learns to group behaviorally similar distractor states (manifesting as obstacles).\n\n**Weaknesses**\nI don’t have any major weaknesses to point out. I will highlight minor comments / weaknesses which, I think if addressed, would definitely make the paper stronger.\n- While the authors state that DBC in practice can be combined with any model-free of model-based algorithm, the paper would definitely benefit if this claim was backed up with results demonstrating DBC combined with another model-free or model-based algorithm since it might be unclear off-the-shelf if an algorithm requires major or minor changes to work in conjunction with DBC.\n- Task generalization (or generalizing over different reward functions) relies on the assumption that the new reward function depends on contributing factors that likely influence the earlier “seen” reward function as well. While empirically validating this claim from walker_walk -> walker_run / walker_stand is a reasonable starting point, demonstrating another instance where inferring the source and target reward functions is not immediately obvious would definitely benefit the paper.", "extra_scores": {}, "extra_text": {}}]

Author rebuttal:

Meta-review: This paper proposed using the state bisimulation metric to learn invariant representations for reinforcement learning.  The method is generic, effective, and is supported by both theoretical and experimental results.  All reviewers and I think this is a strong contribution to the area.
Decision: Accept (Oral)