PaperlayerView exact revision on arXiv

Source-linked AI summary

RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability

Chuning Zhu, Max Simchowitz, Siri Gadipudi, Abhishek Gupta

arXiv:2309.00082v2cs.LGcs.AIcs.RO

TL;DR

Visual model-based RL representations can retain task-irrelevant image information, leaving agents vulnerable to spurious variation. RePo learns reward-predictive, predictable latents under an information constraint and uses test-time encoder adaptation for novel visual environments. The paper reports effectiveness in dynamic, diverse simulation and real-world settings, while noting vulnerability to systematic distribution shift.

  • Problem

    Visual model-based RL representations may retain redundant, task-irrelevant image information, making agents susceptible to spurious variations such as distractors and lighting changes.

    Source: §1 · pp. 1–2

  • Method

    RePo maximizes reward predictability while constraining information from observations into latent representations, then adapts visual encoders at test time using mostly unlabeled data.

    Source: §1 · p. 2

  • Results

    RePo representations and policies succeed in environments with significant dynamic distractors and variability, with efficacy demonstrated across simulation and real-world domains.

    Source: §1 · p. 2

  • Takeaways & Limitations

    RePo is effective for learning in dynamic, distracted environments, though its learned dynamics model captures only task-relevant information.

    Source: §6 · p. 10

    Learned latents and dependent policies may not generalize under systematic distribution shift, because unseen visual inputs can make latents behave erratically even when system dynamics are unchanged.

    Source: §4.1 · p. 6

Abstract

from arXiv · show

Visual model-based RL methods typically encode image observations into low-dimensional representations in a manner that does not eliminate redundant information. This leaves them susceptible to spurious variations -- changes in task-irrelevant components such as background distractors or lighting conditions. In this paper, we propose a visual model-based RL method that learns a latent representation resilient to such spurious variations. Our training objective encourages the representation to be maximally predictive of dynamics and reward, while constraining the information flow from the observation to the latent representation. We demonstrate that this objective significantly bolsters the resilience of visual model-based RL methods to visual distractors, allowing them to operate in dynamic environments. We then show that while the learned encoder is resilient to spirious variations, it is not invariant under significant distribution shift. To address this, we propose a simple reward-free alignment procedure that enables test time adaptation of the encoder. This allows for quick adaptation to widely differing environments without having to relearn the dynamics and policy. Our effort is a step towards making model-based RL a practical and useful tool for dynamic, diverse domains. We show its effectiveness in simulation benchmarks with significant spurious variations as well as a real-world egocentric navigation task with noisy TVs in the background. Videos and code at https://zchuning.github.io/repo-website/.

1 Introduction

Visual model-based RL must handle task-irrelevant changes that can confuse vision-based agents, especially in dynamic and diverse scenes. RePo learns reward-predictive, predictable latents and adds test-time encoder adaptation for novel visual environments.

1 Introduction

  • Prior compression methods target lossless encoding of image evolution without downstream-task demands, and can struggle when observations contain unpredictable but irrelevant content.Standard end-to-end visual RL methods also struggle with spurious variations in part because they fail to discard task-irrelevant elements.Source: §1 · pp. 1–2
  • RePo learns latents that predict their own dynamics and reward while encouraging predictability to discard uncontrollable variations such as background, lighting, and visual traffic.Its objective constrains information flow from observations while retaining task-relevant dynamics for policy learning.Source: §1 · p. 2
  • RePo representations and policies succeed amid dynamic, uncontrollable distractors and environments with substantial variability and complexity.The paper reports efficacy across simulation and real-world domains.Source: §1 · p. 2
  • RePo's visual encoders can suffer under distribution shift, so its test-time procedure adapts only those encoders using mostly unlabeled test data while keeping other model components fixed.The paper says matching test-time encodings to training-time representations can recover optimal performance with minor adaptation.Source: §1 · p. 2

2 Related Work

Related work spans visual model-based RL, task-oriented representation learning, and unsupervised domain adaptation. RePo differs by targeting spurious variation without reconstruction and by matching test-distribution support rather than forcing distributions to be identical.

2 Related Work

  • Visual model-based RL learns latent representations and dynamics from images; RePo forgoes DREAMER's reconstruction loss, while TD-MPC predicts value but may retain irrelevant and policy-related information.The paper identifies DREAMER as especially relevant and discusses the differences in Section 4.Source: §2 · p. 3
  • Unsupervised visual pretraining does not account for downstream tasks, while reward-guided control methods often pursue value-equivalence but do little to counteract spurious variation.Other cited approaches analyze exogenous distractors or task relevance and controllability; RePo does not impose a specific form on spurious variations.Source: §2 · p. 3
  • Domain adaptation often matches source and target distributions, but online RL exploration may produce mismatched marginals; RePo instead uses support matching without making the distributions identical.The stated objective is to keep the test distribution within the support of the training distribution.Source: §2 · p. 3

3 Preliminaries

The preliminaries define the discounted MDP setting and explain how visual observations are encoded into latent states. They characterize spurious variation as changing state content that does not affect task reward or relevant transitions.

3 Preliminaries

  • A discounted MDP specifies states, actions, discounting, transition and initial-state distributions, and rewards; the learning goal is a policy maximizing expected returns.The paper assumes deterministic rewards for simplicity.Source: §3 · p. 3
  • Visual RL treats observations as states, maps images through an encoder h to encodings x, and compresses those encodings into latent states z.The paper notes that images or sequences may be assumed sufficient to determine rewards and transitions, or a belief state may be introduced.Source: §3 · p. 3
  • Spurious variation includes changing state components that do not affect reward or relevant transitions, such as backgrounds, dynamic television content, or lighting.In the canonical factorization, reward and the relevant next-state distribution depend only on s(1), making s(2) spuriously variable.Source: §3 · pp. 3–4

4 RePo: Parsimonious Representation Learning without Reconstruction

RePo learns compact latents that retain information about reward and dynamics while discarding task-irrelevant observations. For test-time distribution shifts, it adapts the visual encoder using support constraints while keeping the rest of the model fixed.

4 RePo: Parsimonious Representation Learning without Reconstruction

  • RePo encourages representations to retain reward- and dynamics-relevant information while discarding information that is neither pertinent to reward nor dynamics.Its objective maximizes latent–reward mutual information while bottlenecking latent–observation information.Source: §4 · pp. 4–5
  • Matching the observation-dependent latent posterior to a prior that omits the latest observation encoding encourages latents to discard unpredictable spurious variations.The paper illustrates this with background TV images, whose removal reduces next-step environmental stochasticity.Source: §4 · p. 5
  • RePo implements its objective with an RSSM and alternates representation learning with actor-critic policy learning in latent space.The RSSM uses an encoder, latent prior, latent posterior, and reward predictor; KL balancing is applied during optimization.Source: §4 · p. 5
  • Unlike DeepMDP and bisimulation, RePo compresses image encodings into latents by matching the observation-conditioned posterior to a prior, ruling out identity-like representations.The paper contrasts this with DeepMDP's potentially non-parsimonious encodings and bisimulation's lack of a guarantee that irrelevant information is discarded.Source: §4 · p. 6

4.1 Transferring Invariant Latent Representations via Test-Time Adaptation

  • RePo's latents may fail to generalize under systematic visual distribution shifts, so test-time adaptation adjusts the encoder while leaving the latent model fixed.The encoder can encounter novel visual features even when system dynamics remain unchanged.Source: §4.1 · pp. 6–7
  • Because initial adaptation can visit unseen state regions, the method matches reweighted training-feature support to test features rather than requiring marginal distributions to match.Calibration pairs regularize the support constraint to prevent degenerate encodings from collapsing to one point.Source: Figure 3 · p. 6, §4.1 · p. 7

5 Experimental Evaluation

The evaluation tests RePo on distracted control tasks, adaptation to new visual domains, realistic robotic manipulation, and real-world navigation. Results cover learning performance, test-time adaptation, representation properties, and objective ablations.

Dynamic visual distractors

  • Across six distracted DeepMind Control environments, RePo outperforms baselines in learning speed and asymptotic performance.The reported comparison indicates resilience to spurious visual variations.Source: Figure 5 · p. 8, §5 · p. 8

Evaluation domains

  • Evaluation spans distracted DeepMind Control, realistic-background Maniskill manipulation, and real-world Lazy TurtleBot navigation.The benchmarks include dynamic video backgrounds, varied static backgrounds, and egocentric navigation in a furnished room.Source: §5 · p. 8

Test-time adaptation

  • Within 50k steps, RePo nearly recovered full training performance after adapting across three DeepMind Control domains.Replacing the support constraint with distribution matching degraded performance; removing calibration examples also made both matching variants perform worse.Source: Figure 6 · p. 9, §5 · p. 9

Diverse static environments

  • On LiftCube, PushCube, and TurnFaucet, RePo saturates performance; Dreamer misses optimal task performance, and several baselines lag in sample efficiency.The tasks use realistic backgrounds in the Maniskill benchmark.Source: Figure 7 · p. 9, §5 · p. 10

Real-world navigation

  • At 15K environment steps, RePo reached nontrivial success on Lazy TurtleBot while Dreamer failed to reach the goal.The real-world navigation setup includes two TVs playing random videos along critical paths.Source: Table 1 · p. 9, §5 · p. 9

Representation analysis

  • RePo latents largely omit backgrounds while retaining agent information, and are more compact across backgrounds than Dreamer, supporting cross-background data sharing.The representation probes reconstruct the agent but little background; the latent comparison uses principal components across backgrounds.Source: Figure 9 · p. 10, §5 · p. 10, Figure 10 · p. 10

Ablations

  • Performance depends crucially on the information bottleneck ϵ and KL balancing.These components are examined in the objective ablations.Source: Figure 8 · p. 10, §5 · p. 10

6 Discussion

RePo learns parsimonious representations resilient to spurious variations and can adapt them to new domains. Its task-relevant dynamics model is not task-agnostic, motivating future multi-reward and continual-learning directions.

Discussion

  • RePo learns parsimonious representations resilient to spurious variations and is effective in dynamic, distracted environments.The discussion also reports that representations degraded under distribution shift but could be quickly adapted through semi-supervised test-time adaptation.Source: §6 · p. 10
  • The learned dynamics model captures task-relevant information only, so it is no longer task-agnostic.The authors suggest that simultaneously predicting multiple reward objectives could potentially address this limitation.Source: §6 · p. 10
  • The authors identify multi-task representations and continual learning as future directions for RePo.They express promise for these settings, particularly for real robots in dynamic, human-centric environments.Source: §6 · p. 10

Supplementary Materials for “RePo: Resilient Model-Based Reinforcement Learning by Regularizing Posterior Predictability”

The supplementary pseudocode outlines RePo training and its test-time visual-encoder adaptation procedure.

A Algorithm Pseudocode

  • Training initializes replay data and model parameters, then iterates through dynamics learning, behavior learning, and data collection.Dynamics learning encodes image sequences, infers latent states and rewards, and updates the objective; behavior learning updates actor and value networks using imagined trajectories.Source: A Algorithm Pseudocode · p. 16
  • Adaptation fixes the training encoder, model, policy, and replay buffer, then initializes a test-time encoder and collects paired calibration observations.The procedure also initializes adaptation networks and a test-time replay buffer.Source: A Algorithm Pseudocode · p. 16
  • Each adaptation update encodes training and test observations, computes the adaptation objective, and adds calibration-encoding error to the Lagrangian.The pseudocode then updates the test-time encoder and associated parameters using dual gradient descent.Source: A Algorithm Pseudocode · p. 16

B Derivations

The derivations formalize RePo as maximizing information between latent states and rewards while constraining information shared with observations. They derive variational bounds for the objective and a tractable support constraint.

B.1 RePo objective

  • RePo’s objective maximizes latent-reward mutual information subject to an upper bound on latent-observation mutual information.The objective is written in Lagrangian form, and its expectation is taken over experienced actions, observations, and rewards.Source: §B.1 · p. 17
  • A variational lower bound on reward prediction uses a reward model conditioned on the latent state; the sufficiency assumption is not required, though its failure may make the bound loose.The derivation uses a reward conditional on the current latent and action.Source: §B.1 · p. 17
  • The dynamic-compression upper bound assumes Markov latent transitions and uses a variational transition approximation; the objective is agnostic to compressing image encodings or observations.The Markov assumption can be discarded, but the bound is sharper when it holds.Source: §B.1 · pp. 17–18

B.2 Tractable objective for support constraint

  • The support-constraint derivation yields a tractable objective using a reweighting function and an f-divergence, with χ2 divergence used in the paper’s formulation.The reweighting constraint ensures the reweighted training distribution is valid, and a Lagrange multiplier incorporates the constraint.Source: §B.2 · pp. 18–19

C Implementation Details

The implementation builds on Dreamer’s recurrent state-space model and trains online with replay-buffer trajectories and imagined policy rollouts. The paper also describes baselines and a test-time encoder adaptation procedure.

Model architecture

  • RePo uses a Dreamer-based recurrent state-space model with a CNN image encoder and Gaussian latent prior and posterior.The image encoder maps 64 × 64 images to 1024-dimensional embeddings; the model uses a 200-unit belief and 30-dimensional state.Source: C Implementation Details · p. 19

Optimization

  • Online training alternates replay-buffer model updates with 15-step policy rollouts in the learned dynamics model.Each update samples 50 trajectories of length 50; training uses Adam, and the initial Lagrange multiplier, target KL, and KL balancing ratio are tuned on evaluation tasks.Source: C Implementation Details · p. 19

Baselines

  • Dreamer and TIA are implemented on RePo’s codebase, while DBC and DeepMDP use their official implementations and reported hyperparameters.DBC uses a deterministic transition model, which the authors found performed better than a stochastic one.Source: C Implementation Details · p. 19

Adaptation

  • Test-time adaptation updates the visual encoder using a support constraint objective and paired calibration observations while keeping the rest of the RePo model fixed.The procedure collects 10 source and target calibration trajectories, then minimizes encoding distance for paired observations alongside the support constraint objective.Source: C Implementation Details · pp. 19–20

Adaptation baselines

  • Adaptation baselines use GAN distribution matching, while ILA additionally minimizes dynamics violation measured by prior-posterior KL divergence.For ILA, gradients are backpropagated through the model to the image encodings.Source: C Implementation Details · p. 20

D Environment Details

The evaluation environments introduce visual variation in simulated control, manipulation, and real-world navigation settings. The section specifies the distractors, tasks, and evaluation protocols for each setting.

Distracted DeepMind Control

  • Distracted DeepMind Control replaces standard environments’ static backgrounds with grayscale videos from Kinectics-400.The evaluation uses a 1000-step time limit, action repeat of 2, and four random seeds.Source: D Environment Details · p. 20

Realistic Maniskill

  • Realistic ManiSkill evaluates Push Cube, Lift Cube, and Turn Faucet tasks against randomly selected realistic Matterport scenes.The authors curate 90 scenes and use four random seeds, a 100-step limit, and action repeat of 1.Source: D Environment Details · p. 20

Lazy TurtleBot

  • Lazy TurtleBot tests egocentric goal-reaching in a household-like room with two TVs playing random videos along critical paths.Training uses 5K offline transitions, 10K offline steps, and 10K online environment steps; evaluation uses two seeds due to time constraints.Source: D Environment Details · p. 20

E.1 Results on standard DMC environments

On standard DeepMind Control environments, RePo matches Dreamer’s asymptotic performance despite having no reconstruction objective. It slightly lags in sample efficiency, which the authors attribute to reward signals being sparser than pixel reconstruction signals.

E.1 Results on standard DMC environments

  • RePo matches Dreamer in asymptotic performance on standard DeepMind Control despite not using a reconstruction objective.It slightly lags Dreamer in sample efficiency, which the authors attribute to reward signals being inherently sparser than pixel reconstruction signals.Source: §E.1 · p. 20, Figure 11 · p. 21

E.2 Results on Distracted Walker Walk with scoreboard

On distracted Walker Walk with a scoreboard, RePo’s latent represents joint positions rather than the score. Additional ablations show trade-offs in bottleneck and KL-balancing settings, while support constraint behaves differently from distribution matching under limited calibration coverage.

E.2 Results on Distracted Walker Walk with scoreboard

  • 872.86 ± 43.71 with the scoreboard versus 868.72 ± 53.93 without; RePo’s latent reconstructs joint positions and ignores the score.The visualization attributes this to RePo predicting dynamics and future rewards as well as current reward.Source: §E.2 · p. 21, Figure 12 · p. 21

E.3 Additional ablations

  • A β learning rate that is too high enforces the constraint early and can impair exploration, while one that is too low makes dual gradient descent converge slowly.The reported sweet spot relaxes the information bottleneck early, then gradually tightens it to discard redundant information.Source: §E.3 · p. 21, Figure 13 · p. 22
  • Too small a KL target can make the representation degenerate and increase reward loss, while too large a target can make the dynamics model inaccurate.The figure caption describes the general trade-off as tighter bottlenecks yielding more parsimonious representations and more accurate dynamics, but overly tight ones thwarting reward learning.Source: §E.3 · p. 21, Figure 14 · p. 22
  • A larger KL-balancing parameter tends to improve performance because the dynamics model is harder to train, but training the prior too aggressively leaves the posterior poorly regularized.The figure caption likewise cautions that overly aggressive prior training leads to a poorly regularized posterior.Source: §E.3 · p. 21, Figure 15 · p. 22

E.4 An illustrative example of support constraint

  • With good calibration coverage, support constraint plus dynamics consistency recovers the ground-truth latent distribution, whereas distribution matching conflicts with dynamics consistency and calibration.The comparison visual describes the conflict under insufficient exploration; support constraint conforms to the dynamics-consistency and calibration objectives.Source: §E.4 · pp. 22–23, Figure 17 · p. 23
  • With skewed calibration coverage, support constraint is accurate only near examples; dynamics consistency propagates calibration, though RSSM consistency helps little under random-policy DMC calibration.The authors hypothesize that the recurrent architecture interferes with this optimization and leave potential solutions for future work.Source: §E.4 · p. 23, Figure 18 · p. 23
Loading 2309.00082v2…