Source-linked AI summary
Unsupervised State Representation Learning in Atari
Ankesh Anand, Evan Racah, Sherjil Ozair, Yoshua Bengio, Marc-Alexandre Côté, R Devon Hjelm
TL;DR
The paper tackles the challenge of learning state representations without reward supervision while capturing latent environmental factors rather than pixel-level details. It introduces ST-DIM, which maximizes mutual information across spatially and temporally distinct features, and evaluates it with an Atari benchmark. ST-DIM excels at capturing small objects and remains effective when easy-to-exploit features are present, while the benchmark exposes qualitative and quantitative differences among representation-learning methods.
Problem
Unsupervised state representation learning must capture latent environmental factors, whereas reconstruction-based methods often emphasize pixel-level details.
Method
ST-DIM maximizes mutual information across spatial and temporal axes, and the Atari benchmark evaluates representations by linearly recovering annotated state variables.
Results
ST-DIM excels at capturing underlying latent factors for small objects and when many objects are present, outperforming generative and other contrastive techniques in these difficult cases.
Takeaways & Limitations
The benchmark supports studying qualitative and quantitative differences between representation-learning techniques across multiple generative factors.
Takeaways & Limitations
The evaluations use a fixed CNN encoder architecture and 256-dimensional representation for every method to ensure fair comparison.
Abstract
from arXiv · showhide
State representation learning, or the ability to capture latent generative factors of an environment, is crucial for building intelligent agents that can perform a wide variety of tasks. Learning such representations without supervision from rewards is a challenging open problem. We introduce a method that learns state representations by maximizing mutual information across spatially and temporally distinct features of a neural encoder of the observations. We also introduce a new benchmark based on Atari 2600 games where we evaluate representations based on how well they capture the ground truth state variables. We believe this new framework for evaluating representation learning models will be crucial for future representation learning research. Finally, we compare our technique with other state-of-the-art generative and contrastive representation learning methods. The code associated with this work is available at https://github.com/mila-iqia/atari-representation-learning
1 Introduction
The paper addresses unsupervised state representation learning by proposing ST-DIM and an Atari benchmark that evaluates whether representations capture latent state variables.
- State representations should capture underlying environmental factors while discarding low-level visual details, but reconstruction-based methods often emphasize pixel-level information.
- ST-DIM learns representations by maximizing mutual information across spatial and temporal axes of visual observations.
- Representations are evaluated by how accurately linear classifiers recover annotated state variables from them.
- The Atari benchmark uses 22 games and source-code annotations of variables such as object locations to evaluate learned representations.
- The work compares ST-DIM with existing unsupervised representation-learning techniques using this benchmark.
2 Spatiotemporal Deep Infomax
SpatioTemporal DeepInfoMax learns abstract, semantically focused representations by maximizing multiple mutual-information objectives across spatially and temporally distinct encoder features. It uses InfoNCE with consecutive observations as positives and non-consecutive minibatch observations as negatives, combining global-local and local-local objectives.
- The goal is to capture latent generative factors and high-level semantics while ignoring low-level observation details such as background texture.The intended semantics include agents, enemies, objects, and score.
- ST-DIM combines global-local and local-local mutual-information objectives across spatially and temporally distinct features.The global-local objective relates a full observation at time t to small patches at time t + 1, while local features come from an intermediate convolutional layer.
- Multiple small mutual-information objectives make the lower-bound estimates easier than a single large objective, addressing looseness and incomplete feature capture at high mutual information.
- InfoNCE learns scores that are large for positive joint samples and small for negative product-of-marginals samples.In this setting, positives are consecutive observation pairs, whereas negatives pair an observation with a randomly sampled non-consecutive observation from the same minibatch.
- A bilinear score function encourages the encoder to learn linearly predictable representations, which the authors associate with semantic-level representation learning.The global-local and local-local scores use learned bilinear mappings between global or local feature vectors.
3 The Atari Annotated RAM Interface (AtariARI)
AtariARI addresses the need for broader representation evaluation by exposing ground-truth game-state factors from Atari RAM and probing whether representations recover them linearly.
- AtariARI uses Atari 2600 games because their temporal structure and accessible ground-truth states support visual representation evaluation.The benchmark uses the Arcade Learning Environment and extracts underlying factors from game source code.
- The benchmark verifies important state variables for 22 games by interpreting RAM through commented disassemblies or source code.ALE exposes 128 bytes of RAM per timestep, which programmers use for information such as sprite locations, clocks, and rooms.
- State variables are organized into agent localization, small-object localization, other localization, score/clock/lives/display, and miscellaneous categories.Examples include controllable sprites, balls or missiles, enemies, game scores, and game-specific variables.
- The benchmark evaluates explicitness by training linear classifiers to predict each ground-truth state variable and reporting mean F1 score.Explicitness measures whether underlying generative factors can be recovered using a linear transformation from the representation.
4 Related Work
Prior work learns representations through mutual-information objectives, temporal self-supervision, controllability, and domain-specific evaluation; ST-DIM combines temporal information with local spatial structure.
- Mutual-information methods estimate dependencies between representations and inputs, images, or patches using variational bounds or noise-contrastive estimation.The cited approaches include Jensen–Shannon divergence and multi-sample noise-contrastive estimation.
- Earlier state-representation methods use VAEs, robotic priors, or temporal objectives, whereas ST-DIM additionally exploits local spatial structure.ST-DIM is described as an extension of TDC and TCN.
- Controllability-focused methods capture factors under the agent’s control, while ST-DIM targets every temporally evolving environmental factor.Examples of targeted factors include enemies, scores, balls, missiles, and moving obstacles.
- Representation evaluation is domain-specific, with vision, NLP, and other fields using benchmarks that probe different kinds of learned information.Vision commonly tests linearly separable label-relevant information, while SentEval and GLUE provide linguistic evaluations.
5 Experimental Setup
The experiments freeze each learned encoder and evaluate representations with linear probes on Atari ground-truth factors, comparing random, generative, predictive, contrastive, and supervised methods.
- Each encoder is trained first, frozen, and then evaluated by linear probes predicting every identified generative factor in each game.Probe gradients are not backpropagated through the encoder.
- Data are collected using either a random agent or a PPO agent trained for 50M timesteps, with eight differently initialized workers.PPO data receive additional stochasticity through an epsilon-greedy mechanism.
- The comparison includes RANDOM-CNN, VAE, PIXEL-PRED, CPC, and a supervised model providing a best-case performance bound.CPC maximizes mutual information between current and future latents; the supervised model backpropagates probe gradients through its encoder.
- All methods use the same base CNN architecture and 256-dimensional representations to ensure a fair comparison.The CNN is adapted to full 160x210 Atari frames.
- Each probe uses 35,000 training frames, 5,000 validation frames, and 10,000 test frames after pruning low-entropy variables and preventing train-test duplicates.The probes are 256-way linear classifiers with early stopping and a plateau-based learning-rate scheduler.
6 Results
Across Atari probing tasks, ST-DIM generally achieves the strongest unsupervised performance, with contrastive methods ahead of generative methods but below supervised representations.
- ST-DIM largely outperforms other methods in mean F1 score across games using random-agent data.Similar trends hold when data are collected by a pretrained PPO agent.
- Contrastive methods, including ST-DIM and CPC, generally perform better than generative VAE and PIXEL-PRED methods on probing tasks.RandomCNN is also reported as a strong Atari prior.
- A sizable gap remains between ST-DIM and the fully supervised approach, leaving room for improvement in unsupervised representation learning.The supervised model serves as a best-case performance bound in the experimental comparison.
7 Discussion
ST-DIM’s spatial and temporal contrastive design improves capture of small objects and resists representations dominated by easy-to-exploit features. Its patch-level mutual-information objectives help explain these advantages over generative methods and ablations.
- Ablations: InfoNCE outperforms JSD, and contrasting spatiotemporally rather than only temporally is important for capturing representations.These are the two reported ablation findings for ST-DIM.
- Interpretation: The reported advantages of ST-DIM are attributed to maximizing mutual information across patch representations.This explanation links its small-object performance and robustness to easy-to-exploit features.
- Capturing small objects: ST-DIM captures small objects better than generative methods and Global-T-DIM.The authors attribute this to patch representations, which retain information from small image regions.
- Robustness to easy-to-exploit features: ST-DIM remains effective when easy-to-exploit features could otherwise saturate a mutual-information or contrastive objective.In Boxing, ST-DIM predicts the clock well while also encoding score and boxer positions.
- Robustness to easy-to-exploit features: Generative and contrastive methods show complementary preferences: PIXEL-PRED favors large, low-entropy features, whereas contrastive methods favor high-entropy features.PIXEL-PRED performs better on some low-entropy location features, while ST-DIM excels on clock and score variables.
8 Conclusion
The paper concludes that ST-DIM learns representations across spatial and temporal axes and introduces an Atari benchmark targeting multiple latent factors. It reports strong performance on small objects and scenes with many objects, while documenting the common experimental setup.
- 8 Conclusion: ST-DIM maximizes mutual information across spatial and temporal axes.This is the paper’s proposed representation-learning technique.
- 8 Conclusion: The Atari 2600 benchmark evaluates whether representations capture multiple underlying generative factors.The benchmark is intended to support qualitative and quantitative comparison of representation-learning techniques.
- 8 Conclusion: ST-DIM excels at capturing latent factors involving small objects and states with many objects.These cases are described as difficult for generative and other contrastive techniques, respectively.
- 8 Conclusion: ST-DIM and its ablations add a 1x256x256 bilinear layer to the base encoder.The listed methods use the same architecture for this comparison.
- 8 Conclusion: Experiments use full-resolution 210x160 images and validation-loss plateau learning-rate schedules.The schedule is used for both contrastive training and probing.
C Results with Probes Trained on Data Collected By a Pretrained RL agent
The appendix evaluates representations on observations collected by a pretrained PPO agent. The agent was trained for 50M steps per game, sampled actions stochastically, and used additional ε-greedy stochasticity.
- C Results with Probes Trained on Data Collected By a Pretrained RL agent: Representations are evaluated on data collected by pretrained PPO agents trained for 50M steps on each game.Actions are sampled from each PPO policy’s distribution at every time step.
- C Results with Probes Trained on Data Collected By a Pretrained RL agent: Additional exploration uses an ϵ-greedy mechanism with ϵ = 0.2.This stochasticity is added on top of sampling from the PPO action distribution.
D More Detailed Ablation Results
The appendix expands ablation results for ST-DIM and adds Static-DIM, while also reporting preliminary probing results for pretrained RL-agent features. These RL features perform poorly relative to baselines, but the authors describe the investigation as preliminary.
- D More Detailed Ablation Results: The appendix provides game-by-game and category-wise breakdowns for JSD-ST-DIM and Global-T-DIM ablations.It also includes Static-DIM, which removes temporal context by sampling negatives from a different game.
- E Probing Pretrained RL Agents: Pretrained PPO-agent features perform poorly on probing tasks compared with the baselines.The agents were trained for 50 million frames, and the authors characterize these results as preliminary.
- E Probing Pretrained RL Agents: The RL-agent probing findings require investigation across different policies and models.The authors explicitly limit the scope of these preliminary results.
F Accuracy Metric
Tables 11 and 12 report probe results using percent accuracy, whereas tables 2 and 3 used F1 scores. The appendix also presents accuracy and F1 results across games, categories, methods, and ST-DIM ablations.
- Tables 11 and 12 use standard percent accuracy instead of the F1 score used in tables 2 and 3.
- Table 7 reports probe F1 scores across methods and categories using data collected by a pretrained PPO agent trained for 50M steps.
- Tables 8 and 9 report F1 scores for ST-DIM ablations across categories, with data collected by random agents.
- Table 10 compares probe results from features of a PPO agent trained for 50 million timesteps with majority-classifier and random-CNN baselines.
- Table 11 reports probe accuracy scores averaged across categories for each game using data collected by random agents.
- Table 12 reports probe accuracy scores for different methods averaged across all games for each category using data collected by random agents.