Learning Invariant Representations for Reinforcement Learning without Reconstruction

Amy Zhang; Rowan McAllister; Roberto Calandra; Yarin Gal; Sergey Levine

Source-linked AI summary

Learning Invariant Representations for Reinforcement Learning without Reconstruction

Amy Zhang, Rowan McAllister, Roberto Calandra, Yarin Gal, Sergey Levine

arXiv:2006.10742v2Jun 18, 2020cs.LGcs.AIstat.ML

TL;DR

The paper addresses how to learn data-efficient control from rich images without pixel reconstruction or manually specified visual knowledge. It trains bisimulation-aware latent representations for downstream reinforcement learning, reporting robustness to visual distractions and performance gains, while leaving partial observability and uncertainty handling as future work.

Problem
Representation learning for image-based reinforcement learning must support downstream control while discarding task-irrelevant visual details without relying on reconstruction or domain knowledge.
Method
DBC trains encoders so latent distances correspond to bisimulation distances, using reward differences and Wasserstein distances between predicted transition distributions, and integrates them with SAC.
Results
DBC is substantially more robust to task-irrelevant distractors than reconstruction- and contrastive-loss approaches, with performance gains when representation distances match bisimulation distances.
Takeaways & Limitations
Bisimulation-based representations can encode task-relevant information while remaining invariant to distracting visual details in outdoor and realistic-image control settings.
Takeaways & Limitations
The paper does not explicitly address partial observability and identifies uncertainty estimation for distributional shifts as future work.

Abstract

from arXiv · show

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 both 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.

1 Introduction

The paper targets data-efficient control from complex visual observations by learning representations that retain task-relevant information while ignoring visual distractions. It proposes bisimulation-based representations whose latent distances reflect behavioral differences between observations.

Motivation: Learning control from real images is challenging because outdoor scenes contain complex, unstructured visual information and real trials are costly.Representation learning is used to compress images into vectors better suited for reinforcement learning.
Motivation: Robust representations should ignore irrelevant objects and details while treating observations with equal task-relevant information equivalently.Examples of relevant information include road direction and the locations of other cars.
Motivation: A task-aware representation should predict both rewards and its own future representations, capturing the recursive relevance of state information.Future state elements matter when they influence later rewards.
Approach: DBC learns invariant representations by matching distances between observation encodings to bisimulation distances measuring behavioral differences.The method is designed for downstream control rather than pixel reconstruction.

2 Related Work

Prior representation-learning methods use reconstruction or contrastive objectives, while bisimulation provides a task-oriented notion of behavioral similarity for state abstraction. The paper positions DBC as a scalable, data-driven alternative for grouping visually different observations by downstream relevance.

Reconstruction-based Representations: Reconstruction-based approaches can learn low-dimensional representations, but their features are not guaranteed to be useful for control and often require expert tuning.Model-based variants jointly train encoders and dynamics, yet long-term prediction remains difficult.
Contrastive-based Representations: Contrastive methods enforce similarity constraints using augmentations, patches, or other data relationships, often relying on domain knowledge.Without suitable task-aware similarity, prediction-based representations may retain all predictable visual features.
Contrastive-based Representations: A task-aware similarity function can group pixel-wise different images, such as those differing in lighting or texture, when their downstream objectives are similar.This motivates learning similarity from the control task rather than manual visual engineering.
Bisimulation: Bisimulation groups states with identical reward-sequence behavior under action sequences, while bisimulation metrics measure graded behavioral similarity.The metric extends state abstraction to continuous settings.
Bisimulation: Exact bisimulation relations preserve behavioral properties but can be impractical for continuous state spaces, motivating metric-based soft state abstractions.The related-work discussion also notes scalability limitations in prior Wasserstein-based computation.

3 Preliminaries

The paper models the environment as a discounted Markov decision process and introduces bisimulation as a recursive behavioral equivalence. Because exact equivalence is brittle in continuous spaces, bisimulation metrics soften it using reward and transition-distribution distances.

MDP assumptions: An MDP is represented by M = (S, A, P, R, γ), where states, actions, transitions, rewards, and discounting define the control problem.The agent selects actions through a policy and seeks to maximize expected cumulative discounted rewards.
MDP assumptions: The paper approximates stacked pixel observations as the fully observed state and does not explicitly address partial observability.This is an assumption for the image-based experiments.
Bisimulation: Bisimilar states share the same immediate rewards and equivalent transition distributions over groups of bisimilar successor states.Equivalently, they produce identical probabilistic reward sequences for any action sequence.
Bisimulation metrics: Bisimulation metrics replace exact state partitions with a pseudometric whose distance quantifies behavioral similarity between states.Different states may have zero distance because the metric is a pseudometric.
Bisimulation metrics: The metric combines reward distance with Wasserstein distance between transition distributions, treating distribution matching as an earth-mover transport cost.The Wasserstein metric is defined through couplings between probability distributions.

4 Learning Representations for Control with Bisimulation Metrics

DBC learns control-oriented encoders whose latent ℓ1 distances correspond to bisimulation metrics, using reward and predicted transition-distribution distances. The representation is trained with a dynamics model and integrated with SAC so it can support online policy improvement.

DBC representation learning: The representation-learning architecture reuses one encoder for both states and trains with reward and discounted transition-distribution distances measured using Wasserstein distance.The shared architecture is analogous to a Siamese network.
Training procedure: DBC iteratively trains a policy, encoder, and latent dynamics model, using permuted state batches for encoder training.The dynamics model predicts the next latent representation from the current representation and action.
DBC representation learning: DBC learns representations Z rather than only state distances, making latent ℓ1 distances correspond to bisimulation metrics for reinforcement learning.The learned encoders map states into a representation space used to improve control.
Training objective: Because the dynamics model outputs Gaussian distributions, DBC uses the closed-form 2-Wasserstein distance for its transition-distribution term.The method continues using ℓ1 distances for other distance terms.
Incorporating control: DBC combines its representation learning with SAC and allows the value function to backpropagate to the encoder.The paper describes this modification as a way to improve performance further.

5 Generalization Bounds and Links to Causal Inference

The paper derives value-function bounds for bisimulation-based representations and connects them to causal feature sets, task generalization, and robustness to distractor interventions.

DBC's theoretical analysis bounds the suboptimality of value functions trained on representations learned via DBC.
A continuously improving policy induces a bisimulation operator whose least fixed point is a π∗-bisimulation metric.
Bisimulation distances are independent of pixels and recursively depend on future transition probabilities and rewards, while the optimal value function is Lipschitz with respect to the metric.
Larger ϵ produces a more compact bisimulation partition but permits a looser optimal-value bound.
With encoder learning error L, the value discrepancy is bounded by 2ϵ+2L over (1−γ)(1−c).
Bisimulation partitions correspond to causal feature sets: representations retain the causal ancestors of current and future reward, forming a minimal sufficient statistic.
Representations remain robust to interventions on distractors that alter rendering without affecting reward and can generalize to new rewards sharing causal ancestors.
The fixed representation learned from walker_walk was evaluated with new walker_stand and walker_run rewards, providing empirical evidence for this causal generalization claim.

6 Experiments

Across visual-control experiments, DBC remains robust to increasingly complex distractors, learns representations organized by task-relevant behavior, generalizes across backgrounds and reward functions, and performs well in autonomous driving.

6.1 Control with Background Distraction: DBC remains consistent with simple background distractors while other methods’ performance decreases.The distractors are colored circles following easy-to-predict ideal-gas dynamics.
6.1 Control with Background Distraction: With natural-video backgrounds, DBC attains substantially better results than prior methods that can work without distractors.The natural videos come from the Kinetics dataset.
6.1 Control with Background Distraction: DBC maps observations with similar robot configurations near each other even when their backgrounds differ drastically.The learned bisimulation representation ignores irrelevant visual information while preserving task-relevant configuration.
6.2 Generalization Experiments: DBC representations generalize from simple distractors to natural videos and from walker_walk training to walker_stand and walker_run reward functions.The reported experiments use fixed learned representations for the new settings and tasks.
6.3 Comparison with other Bisimulation Encoders: Compared with Castro (2020), DBC learns control faster while using an ℓ1 encoding distance instead of an additional network ψ.The Castro representation can also learn control, despite not being designed for it.
6.4 Autonomous Driving with Visual Redundancy: In CARLA, DBC is evaluated for highway driving with visual observations containing moving traffic and task-irrelevant environmental variation.The task requires driving safely along a highway without collisions.

7 Discussion

DBC is presented as a control-oriented representation method that improves performance on driving from outdoor images while leaving several extensions for future work.

46.8% better final performance than the next-best baseline was achieved by DBC on the driving task.The comparison used 3 seeds; DBC outperformed DeepMDP, SAC, and contrastive methods.
DBC learns representations invariant to task-irrelevant details in outdoor or otherwise distracted images.The discussion frames this invariance as important for learning control from such observations.
The latent dynamics model was used only to train the encoder and could additionally support multi-step planning in latent space.
Future work includes estimating uncertainty to detect and adapt to distributional shifts between training and test observations.
Partially observed settings remain an open issue because the approach approximately assumes full observability through stacked images.Possible directions include explicit or implicit memory, such as an LSTM.

A Additional Theorems and Proofs

The appendix establishes convergence and value guarantees for bisimulation metrics, then connects latent bisimulation representations to causal ancestors of reward.

The operator F has a least fixed point that is a π*-bisimulation metric.
Policy improvement and fixed-point arguments show that policy iteration finds a fixed-point bisimulation metric under the stated assumptions.The proof uses monotonic value improvement across policies.
Theorem 2 bounds differences between optimal value functions in the original and aggregated MDPs when states are clustered within an ε-neighborhood.The bound is expressed using cluster distances and the discount factor.
For an encoder whose L1 latent distances equal the latent bisimulation metric, the representation encodes all causal ancestors of the reward.
A state factor outside the reward’s causal ancestors has bisimulation distance zero when changed, whereas an ancestor can have positive bisimulation distance.Figure 10 illustrates reward depending directly on s1 while s1 depends on s2, making both causal ancestors.

B Definition of State

The paper models image observations with a latent-state Block MDP, whose block structure ensures that each observation identifies its generating state.

The Block MDP separates the unknown state space S from the image observation space O through a rendering function q.
Each observation uniquely determines its generating state, so O is partitioned into disjoint blocks associated with states.
This block assumption supplies the Markov property in observation space; stacked pixels from one camera angle can represent rendered underlying actuator states.

C Additional DMC Results

Additional DMC experiments evaluate DBC against reconstruction, contrastive, DeepMDP, and SLAC baselines across default and distracted visual settings.

Default-setting results cover 9 DMC environments and compare DBC with reconstruction-loss, contrastive-loss, and SLAC baselines.The experiments use 10 seeds with one standard-error shading in Figure 11.
Simple-distractor results cover all 9 environments and compare DBC with reconstruction, contrastive, DeepMDP, and SLAC baselines.Figure 12 reports 10 seeds with one standard-error shading.
Natural-video results cover all 9 environments and compare the bisimulation metric method with reconstruction, contrastive, DeepMDP, and SLAC baselines.Figure 13 reports 10 seeds with one standard-error shading.

D Additional Visualizations

Figure 14 visualizes latent spaces learned with a bisimulation metric and shows that nearby embeddings preserve robot configuration while varying in background appearance.

The visualization averages observations from 10 nearby points in the t-SNE embedding.
The averaged robot agent remains crisp, indicating that neighboring points encode similar robot positions.
The averaged backgrounds become blurry because nearby points contain substantially different backgrounds.

E Implementation Details

The implementation combines a visual encoder and forward models with Soft Actor-Critic, using specified network architectures and stabilization procedures whose hyperparameters are summarized in Table 2.

The visual encoder uses 3 × 3 convolutional kernels with 32 channels, stride 2 only initially, ReLU activations, and a 50-dimensional tanh output.
The reconstruction baseline uses a fully connected layer followed by four deconvolutional layers, with pixel output from the final layer.
The dynamics and reward models are MLPs with two hidden layers of 200 neurons and ReLU activations.
SAC evaluates soft Q-functions from replay-buffer transitions and improves policies through a soft actor-critic maximum-entropy procedure.
The target value uses a Monte-Carlo expectation, while the target soft Q-function is updated by exponential moving average and policy improvement minimizes KL divergence to a Q-induced Boltzmann distribution.
The RL implementation shares an encoder between actor and critic, uses forward models with two 200-neuron hidden layers, and lists experiment settings in Table 2.