PaperlayerView exact revision on arXiv

Source-linked AI summary

Deep Reinforcement and InfoMax Learning

Bogdan Mazoure, Remi Tachet des Combes, Thang Doan, Philip Bachman, R Devon Hjelm

arXiv:2006.07217v3cs.LGstat.ML

TL;DR

Model-free agents can struggle to generalize when environments or rewards change, motivating representations predictive of future-state properties beyond expected rewards. DRIML adds a Deep InfoMax-based contrastive objective to RL agents to learn such representations. The paper reports that it implicitly learns a transition model, improves continual-learning performance, and boosts Procgen training performance.

  • Problem

    Model-free agents often generalize poorly when environments or reward functions change, motivating the study of representations predictive of future states beyond expected rewards.

    Source: §1 · p. 1

  • Method

    DRIML augments RL with an auxiliary objective that maximizes concordance between state-action representations and future-state representations.

    Source: §1 · p. 2, §4 · p. 4

  • Results

    DRIML implicitly learns a transition model by boosting state similarity, improves deep RL performance in continual learning, and boosts training performance on Procgen.

    Source: §7 · p. 9

  • Takeaways & Limitations

    The reported results show that a contrastive objective over future states can support continual-learning performance and training performance in complex domains such as Procgen.

    Source: §7 · p. 9

    Under nonlinear function approximation, the authors know of no reliable way to estimate the transition operator’s second eigenvalue.

    Source: §8.4 · p. 15

Abstract

from arXiv · show

We begin with the hypothesis that a model-free agent whose representations are predictive of properties of future states (beyond expected rewards) will be more capable of solving and adapting to new RL problems. To test that hypothesis, we introduce an objective based on Deep InfoMax (DIM) which trains the agent to predict the future by maximizing the mutual information between its internal representation of successive timesteps. We test our approach in several synthetic settings, where it successfully learns representations that are predictive of the future. Finally, we augment C51, a strong RL baseline, with our temporal DIM objective and demonstrate improved performance on a continual learning task and on the recently introduced Procgen environment.

1 Introduction

Model-free agents can excel in high-dimensional settings but often generalize poorly when environments or rewards change, motivating representations that capture predictive dynamics beyond reward. The paper proposes a Deep InfoMax-based auxiliary objective and reports faster continual adaptation and improved Procgen performance when it augments an RL agent.

1 Introduction

  • Model-free methods can generalize poorly when environments or reward functions change and can have high sample complexity.They often perform well when reward functions are fixed.Source: §1 · p. 1
  • Predictive representations of future-state properties beyond expected rewards are hypothesized to help model-free agents solve and adapt to new RL problems.The hypothesis is motivated by the richer learning signal available from predictive models compared with reward alone.Source: §1 · p. 1
  • The paper uses a self-supervised Deep InfoMax-derived objective and examines learned representations in synthetic domains and complex environments, including Procgen.The experiments include disjoint and glued Markov chains, a 2D Ising model, and sequential Ms. PacMan.Source: §1 · p. 2
  • The authors propose an auxiliary objective that maximizes concordance between successive-state representations given the action, with an adaptive mechanism for adjusting contrastive-task timescales.They report faster adaptation in continual learning and improved overall performance on Procgen.Source: §1 · p. 2

2 Background

The background motivates learning transferable dynamics information because future task rewards can be difficult to anticipate and model-free agents can struggle to adapt. DRIML builds on predictive and contrastive representation-learning approaches, focusing on settings where rewards change more rapidly than environment dynamics.

2 Background

  • The motivation for transferable representations is that downstream tasks can be difficult to anticipate, and adaptation remains challenging even for state-of-the-art systems.This is especially relevant to complex environments involving other intelligent agents such as humans.Source: §2 · p. 2
  • Continual learning benchmarks adaptation to new tasks using auxiliary information about relatedness across tasks and timescales.Meta-learning and multi-task learning have shown good continual-learning performance by explicitly training transfer.Source: §2 · p. 2
  • Prior RL work uses predictive-state tasks, empowerment, latent transition models, and contrastive objectives, including CPC, EMI, CURL, and HOMER.CPC augments A2C with a temporal contrastive task; CURL uses augmented views of the same frame rather than future frames.Source: §2 · pp. 2–3
  • Contrastive representation learning trains encoders to capture information shared across data views and commonly maximizes a scoring-function lower bound on mutual information.Views may come from spatial or temporal locations, or from data augmentation.Source: §2 · p. 2
  • The paper focuses on settings where reward functions change more rapidly than environment dynamics, while noting this assumption does not hold in every generalization setting.Auxiliary losses are used to encourage learning about environmental generative factors and dynamics.Source: §2 · p. 2
  • DRIML predicts future states conditioned on the current state-action pair at multiple scales, using global features, local features, or both.Adapting contrastive-task timescales to the predictability of subsequent actions under the current policy is reported to improve performance substantially.Source: §2 · p. 3

3 Preliminaries

The preliminaries describe C51 as the value-based RL agent and define the contrastive objective as mutual information between state-action pairs and later states. InfoNCE estimates a lower bound by distinguishing joint samples from samples formed using negative states.

3 Preliminaries

  • C51 extends DQN to predict the full distribution of potential future rewards, and its loss uses a projected distributional Bellman target.The projection maps the return distribution onto a fixed support of 51 atoms.Source: §3 · p. 3

3.1 State-action mutual information maximization

  • Mutual information quantifies information shared between variables and is used in representation learning to align features from different views.Examples include image patches, sequence timesteps, and augmented versions of an image.Source: §3.1 · p. 3
  • The method seeks state-action representations with high mutual information with representations of later trajectory states at a fixed temporal offset.The tuple (s_t, a_t, s_t+k) is generated by following policy π in the MDP.Source: §3.1 · p. 3
  • InfoNCE maximizes a lower bound on mutual information by discriminating joint-distribution samples from product-of-marginals samples.Negative samples replace the positive future state with a state drawn from another trajectory.Source: §3.1 · p. 3
  • The encoders map state-action pairs and states to features, and minimizing the loss maximizes mutual information between those features.The minibatch supplies N − 1 negative samples for each tuple of N samples.Source: §3.1 · p. 4

4 Architecture and Algorithm

DRIML adds a contrastive auxiliary objective to a C51-style encoder, using local and global state features and actions to score present state-action representations against future-state representations. The algorithm jointly optimizes the RL and auxiliary objectives, with the added compute cost described as minimal.

4 Architecture and Algorithm

  • The encoder Θ maps input states to progressively more global, or less local, feature spaces, with the final features serving as C51 value heads.The auxiliary objective uses local f3 and global f4 features as views of state-action-future tuples.Source: §4 · p. 4
  • DRIML conditions local and global representations on an embedded action and scores them against future-state representations.The functions process local and global features using 1 × 1 convolutions.Source: §4 · p. 4
  • The contrastive loss uses state-action-future tuples sampled from trajectories in the experience replay buffer.Experiments can use the local head, global head, or both, as illustrated by Figure 1’s architecture and batch-sample panels.Source: §4 · p. 4, Figure 1 · p. 5
  • DRIML combines the RL update with the DIM objective, and the additional compute cost is described as minimal.The added cost comes from extra passes through small state/action embedding functions followed by an outer product.Source: §4 · pp. 4–5
  • The auxiliary loss improves value-based agents’ predictive capabilities by boosting similarity between representations close in time.The algorithm jointly optimizes the encoder and auxiliary functions together with the RL objective.Source: §4 · pp. 4–5

5 Finding the Best Task Timescale

DRIML adapts its contrastive look-ahead horizon to action predictability, avoiding a single fixed timescale across games. The learned action-transition statistics drive non-homogeneous geometric sampling of the horizon.

5 Finding the Best Task Timescale

  • A fixed temporal offset must be chosen in advance, although games differ in the timescales of their dynamics and therefore in predictive-task difficulty.The adaptive offset is intended to make DRIML more generally useful across environments without the added complexity of simultaneous prediction at multiple timescales.Source: §5 · p. 5
  • A learned qπ(ai, aj) estimates the log-odds of taking action aj after ai under policy π, and its softmax rows form action-transition matrix A.The method learns qπ off-policy from replay-buffer data and updates it at the main DRIML update frequency.Source: §5 · pp. 5–6
  • DRIML samples k from a non-homogeneous geometric distribution using A, selecting an adaptive, game-specific look-ahead value between 1 and H.The authors motivate this distribution because it can represent any discrete positive random variable and bounds its expectation using the minimum and maximum q values.Source: §5 · pp. 5–6
  • DRIML-ada repeats movement actions in navigation games and firing actions in shooting games, while k tends toward 1 in Bigfish and Plunder as training progresses.The authors interpret the latter trend as hinting at an exploration-exploitation-like trade-off.Source: §5 · p. 6, Figure 2 · p. 6
  • Procgen look-ahead selection considers actual actions alongside visible actions, which can number up to 15.This accommodates games that lack special actions such as firing or diagonal movement.Source: §5 · p. 6

6 Experiments

The experiments examine DRIML’s predictive representations in random walks, partially deterministic visual dynamics, continual Ms. PacMan, and Procgen. Across these settings, they test predictability signals and adaptive performance, including a Procgen evaluation on 500 levels with fixed order.

6 Experiments

  • The experiments assess state similarity in Markov chains, local predictability in partially deterministic systems, continual adaptation in Ms. PacMan, and performance on Procgen’s fixed-order levels.The authors report that the DIM loss converges faster for more deterministic Ms. PacMan tasks.Source: §6 · p. 7

6.1 DRIML learns a transition ratio model

  • In the biased random walk, the learned successive-state prediction matrix aligns with the true transition matrix, while DRIML’s convergence rate is affected by the transition operator’s spectral gap.The lowest spectral gap occurs near α = 0.5, where the system is least predictable; the experiment uses α = 0.499 because the transition matrix is not invertible at 0.5.Source: §6.1 · p. 7, Figure 3 · p. 7

6.2 DRIML can capture complex partially deterministic dynamics

  • Local DIM assigns higher similarity scores to Ising-pattern regions than to noisy regions, identifying areas of greater temporal predictability.The experiment overlays a locally dependent Ising patch on independently evolving random pixels and trains a convolutional encoder on local views.Source: §6.2 · pp. 7–8, Figure 4 · p. 8

6.3 A continual learning experiment on Ms. PacMan

  • In Ms. PacMan, DRIML’s performance improves across successive lethal-ghost changes, while C51 reaches similar end-of-task performance and relearns when the lethal ghost changes.DRIML encodes predictable objects, including harmless ghosts; it also outperforms C51 after a lethal-ghost change when Ising dynamics add distractors in walled areas.Source: §6.3 · p. 8, Figure 5 · p. 9

6.4 Performance on Procgen Benchmark

  • DRIML is evaluated on the first 500 levels of all 16 Procgen tasks, with algorithms trained for 50M frames; the supplied passages report no comparative table values.Table 1 reports average training returns and one standard deviation; DRIML’s local-local loss also benefits a smaller set of games such as Starpilot.Source: §6.4 · pp. 8–9, Table 1 · p. 9

7 Discussion

DRIML maximizes representation-level concordance between state-action pairs and future states. The paper reports implicit transition-model learning, improved continual-learning performance, and boosted training performance on Procgen.

7 Discussion

  • DRIML implicitly learns a transition model by boosting state similarity.The paper also characterizes the objective as maximizing concordance between state-action pairs and future states at the representation level.Source: §7 · p. 9
  • DRIML can improve deep RL agents’ performance in a continual learning setting.Source: §7 · p. 9
  • DRIML boosts training performance in complex domains such as Procgen.Source: §7 · p. 9

Broader Impact

The objective improves continual-learning performance and average training rewards across complex video games, while the authors note limits on interpretability and performance guarantees. They recommend failsafes for real-world deployment.

Broader Impact

  • The objective improves performance in a continual learning setting and average training rewards across a suite of complex video games.Source: Broader Impact · p. 9
  • The authors identify healthcare, dialog systems, crop management, and robotics as possible deep RL application areas.They state that methods more robust to environmental changes or better in continual learning can lead to improvements in these applications.Source: Broader Impact · p. 10
  • Deep learning tools and architectures are hard to interpret and prone to failures, while deep RL agents lack formal performance guarantees.The authors say failsafes are essential when deploying such agents in the real world.Source: Broader Impact · p. 10

8 Appendix

The appendix defines Markov chains and decision processes, then relates state-distribution convergence under a fixed policy to an invariant distribution. It notes that absorbing or terminal states often prevent real-life MDPs from having one.

8.1 Markov Chains

  • A discrete-time homogeneous Markov chain is specified by a transition matrix whose entries give next-state probabilities, and its invariant distribution summarizes long-term behavior.Source: §8.1 · p. 13

8.2 Markov Decision Processes

  • A discounted MDP comprises states, actions, transitions, rewards, and a discount factor; its policy aims to maximize expected discounted returns.Source: §8.2 · p. 13

8.3 Link to invariant distribution

  • For ergodic MDPs under a fixed policy, the state distribution is assumed to converge to an invariant distribution; off-policy replay supplies gradients from varied timesteps, encouraging features predictive through time.Source: §8.3 · pp. 13–14
  • Most real-life MDPs have absorbing or terminal states and therefore do not have an invariant distribution; in that case, the DIM estimate of mutual information yields a convergence rate estimable from the transition spectrum.Source: §8.3 · p. 14

8.4 Predictability and Contrastive Learning

The appendix connects mutual information and contrastive learning to predictability and transition structure in Markov systems. Its convergence analysis assumes a fixed policy and near-converged InfoNCE training, while reliable second-eigenvalue estimation under nonlinear function approximation remains unavailable to the authors.

8.4 Predictability and Contrastive Learning

  • Mutual information between state-action pairs and future states reflects dependence, and its connection to transition-kernel eigenvalues links predictability to properties such as mixing time and cluster structure.Source: §8.4 · p. 14
  • The convergence argument considers a fixed policy and an InfoNCE scoring function trained close to convergence on a state pair.The learned ratio is compared with its stationary counterpart; the estimate converges to the time-specific ratio in the infinite-sample limit.Source: §8.4 · p. 14
  • Faster convergence of the learned ratio to its stationary value is suggested when the transition matrix has a larger spectral gap, but reliable estimation of the second eigenvalue under nonlinear function approximation is unavailable.Source: §8.4 · pp. 14–15

8.5 Code snippet for DIM objective scores

The code produces pointwise scores for a batch, then reduces them to a scalar by summing across locations and averaging across batch and representation dimensions.

8.5 Code snippet for DIM objective scores

  • The snippet yields pointwise, uncontracted scores for a batch of data.It does not initially reduce the scores to a single value.Source: §8.5 · p. 15
  • The positive-score tensor has dimensions n_batch × n_rkhs × n_locs.Source: §8.5 · p. 15
  • To obtain a scalar, sum over locations and average over the batch and representation dimensions.Source: §8.5 · p. 15

8.6 Experiment details

The experiments examine temporal predictability in a biased random walk, an Ising-pattern lattice, Ms.PacMan, and Procgen, alongside implementation settings. In Procgen, augmentation helps on some games but harms others, and the best prediction horizon varies by game.

8.6 Experiment details

  • RGB experiments use the settings in Table 2; global DIM heads use a 512-unit hidden layer with ReLU and a skip connection, while local heads use a 1 × 1 convolution.State and action representations are concatenated for DIM heads; the Ising model has no action component.Source: §8.6 · p. 15, Table 2 · p. 16

8.6.1 AMI of a biased random walk

  • For the biased random walk, the stationary state distribution depends on r = α/(1 − α), and the objective’s mutual information depends on the ratio P/ρ.The average mutual information can be estimated by maximizing a bound such as InfoNCE and computed as a function of α.Source: §8.6.1 · pp. 15–16

8.6.2 Ising model

  • At each timestep, a uniformly sampled 43 × 43 region evolves according to an Ising model at temperature β−1 = 0.4, while remaining nodes independently take −1 or 1 with equal probability.A subset of the region at timestep t allows mutual-information-maximizing models to predict the system configuration at t + 1, since the region has high mutual information across time.Source: §8.6.2 · p. 16

8.6.3 Ms.PacMan

  • The Ms.PacMan experiments vary ghost predictability and test adaptation when the lethal ghost changes among four levels after 5,000 episodes.A further experiment overlays an evolving Ising distractor on the game’s walls; four frames are stacked as model input.Source: §8.6.3 · p. 17

8.6.4 Procgen

  • Experiments use the first 500 levels of each Procgen game in a fixed order and the Nature DQN architecture, so results may differ from other Procgen baselines.CURL and CPC use the same similarity networks as the proposed models; No Action is tuned across k = 1, 5 and augmentation settings.Source: §8.6.4 · pp. 17–18
  • Input augmentation helps on bigfish, bossfight, chaser, coinrun, jumper, leaper, and ninja, but has detrimental effects on other tasks.The procedure uses a random crop of 0.8 of the original size with color jitter parameters 0.4.Source: §8.6.4 · p. 18
  • A lookahead k of 5 performs better on bigfish, chaser, climber, fruitbot, jumper, miner, maze, and plunder; k = 1 performs better on the remaining games.Source: §8.6.4 · p. 18

Related papers

Loading 2006.07217v3…