PaperlayerView exact revision on arXiv

Source-linked AI summary

Reinforcement Learning with Prototypical Representations

Denis Yarats, Rob Fergus, Alessandro Lazaric, Lerrel Pinto

arXiv:2102.11271v2cs.LGcs.AI

TL;DR

Image-based RL must learn representations from diverse exploratory data, but exploration itself requires coherent representations and task-specific rewards may be absent. Proto-RL jointly pre-trains visual representations and prototypes through task-agnostic exploration, then reuses them for downstream tasks. The method enables state-of-the-art downstream policy learning, exploration, and multi-task generalization, particularly in sparse-reward settings.

  • Problem

    Representation learning in image-based RL is constrained by coupled exploration and data-coverage requirements, while task-specific rewards may be unavailable and task-tied representations generalize poorly.

  • Method

    Proto-RL jointly learns a visual encoder and prototypical embeddings during task-agnostic exploration, using prototype-based intrinsic entropy rewards, then reuses them for downstream RL.

  • Results

    Proto-RL enables state-of-the-art downstream policy learning, effective exploration, and generalization across multiple tasks, especially in sparse-reward settings.

  • Takeaways & Limitations

    Task-agnostic representations and prototypes can support efficient downstream exploration and learning across unseen tasks without access to downstream task information during pre-training.

  • Takeaways & Limitations

    Comparisons are restricted to environments reported in the original paper because the publicly released code underperforms the reported numbers.

Abstract

from arXiv · show

Learning effective representations in image-based environments is crucial for sample efficient Reinforcement Learning (RL). Unfortunately, in RL, representation learning is confounded with the exploratory experience of the agent -- learning a useful representation requires diverse data, while effective exploration is only possible with coherent representations. Furthermore, we would like to learn representations that not only generalize across tasks but also accelerate downstream exploration for efficient task-specific training. To address these challenges we propose Proto-RL, a self-supervised framework that ties representation learning with exploration through prototypical representations. These prototypes simultaneously serve as a summarization of the exploratory experience of an agent as well as a basis for representing observations. We pre-train these task-agnostic representations and prototypes on environments without downstream task information. This enables state-of-the-art downstream policy learning on a set of difficult continuous control tasks.

1. Introduction

Proto-RL addresses coupled representation-learning and exploration challenges by pre-training visual embeddings and prototypes without task-specific rewards. These representations support rapid downstream exploration, task generalization, and state-of-the-art performance.

  • Image-based RL needs compact encoders, but reward-only supervision is sample inefficient and reward-free settings lack task-specific exploration signals.
  • Representation learning and exploration are mutually dependent because exploration determines the observed data, while exploration requires a coherent latent space.
  • Proto-RL pre-trains an encoder and prototypical embeddings through task-agnostic exploration, using prototype-based intrinsic entropy rewards without downstream task rewards.
  • The pre-trained encoder and prototypes are reused for downstream RL, where they accelerate learning of new tasks with few environment interactions.
  • Proto-RL achieves state-of-the-art performance by combining accurate latent-state representations with prototype-based exploration.
  • Proto-RL generalizes learned representations and prototypes across unseen DeepMind Control Suite tasks and improves downstream exploration, especially with sparse rewards.

2. Related Work

Proto-RL builds on self-supervised visual representation learning, representation learning for pixel-based RL, and task-agnostic intrinsic-motivation methods. Its distinction is using prototypes to estimate state-visitation entropy for downstream exploration.

  • Self-supervised learning in Computer Vision (CV): Self-supervised computer-vision methods learn representations from unlabeled data and provide initialization for downstream fine-tuning.
  • Representation learning in RL: Representation-learning methods for pixel-based RL learn coherent latent representations alongside reinforcement learning to improve sample efficiency.
  • Exploration and Intrinsic Motivation in RL: Intrinsic-motivation approaches address exploration in sparse- or reward-free settings using visitation counts, curiosity, and related inductive biases.
  • Exploration and Intrinsic Motivation in RL: Proto-RL maximizes state-visitation entropy and uses prototypical representations to estimate entropy more effectively than prior work.

3. Background

The background formulates reward-free image control as a POMDP converted to an MDP over stacked observations, with exploration defined by visitation-distribution entropy. Nearest-neighbor estimators provide a tractable entropy proxy in continuous spaces.

  • 3.1. Task-Agnostic RL from Images: Task-agnostic image control is modeled as an infinite-horizon reward-free POMDP over image observations, actions, transitions, discounting, and an initial-state distribution.
  • 3.1. Task-Agnostic RL from Images: Stacking three consecutive image observations converts the task-agnostic POMDP into an MDP with state space X and transition dynamics over trajectory snippets.
  • 3.1. Task-Agnostic RL from Images: A policy induces a discounted state-visitation distribution, and task-agnostic exploration seeks to maximize its entropy.
  • Task-specific downstream RL adds a reward function and optimizes the expected discounted sum of rewards.
  • 3.3. Nearest Neighbor Entropy Estimation: Monte Carlo entropy estimation is difficult in high-dimensional continuous spaces because it requires pointwise density estimation.
  • Proto-RL jointly learns image embeddings and prototypes, uses clustering loss for representation learning, and derives intrinsic rewards from prototype-based entropy.
  • 3.3. Nearest Neighbor Entropy Estimation: Nearest-neighbor entropy estimation assigns each sample a contribution proportional to its distance from its kth nearest neighbor and is asymptotically unbiased and consistent.

4. Proto-RL Algorithm

Proto-RL jointly learns task-agnostic visual representations, prototypes, and entropy-based exploration from visual observations. The learned prototypes define the latent basis and support lower-variance exploration rewards by balancing candidate sampling across clusters.

  • Prototypical Representation Learning: Proto-RL learns an encoder, projector, predictor, and prototype vectors simultaneously using self-supervised comparisons between consecutive augmented observations.A target network produces prototype-assignment targets, while gradients update only the online network; the target network uses exponential moving averages to prevent collapse.
  • Prototypical Representation Learning: The framework maps augmented image frames through an encoder, projector, and predictor before assigning the resulting features to a softmax distribution over prototypes.The prototypes form a basis for the latent space, and the softmax temperature controls assignment probabilities.
  • Maximum Entropy Exploration: Proto-RL trains an exploration policy with intrinsic reward based on entropy estimated from nearest-neighbor distances in the learned low-dimensional latent space.This replaces similarity estimation in the original high-dimensional image space with distances between learned embeddings.
  • Maximum Entropy Exploration: Prototype-based clustering uniformly samples candidates for the nearest-neighbor queue, reducing the high variance caused by random replay-buffer batches and under-represented novel states.Candidates are sampled from prototype-defined distributions and stored in a fixed-size queue with a constant number per cluster.
  • Evaluation: In eight DeepMind Control Suite environments, Proto-RL consistently beats baselines and often exceeds the fully supervised DrQ approach after task-agnostic pretraining.Most methods receive 500k pretraining steps followed by 500k task-specific training steps; DrQ uses task reward from the outset.

5. Experiments

Proto-RL is evaluated in a two-phase unsupervised RL setting, testing task-agnostic representations, prototype-driven exploration, transfer across tasks, and pre-training efficiency. Across these evaluations, Proto-RL generally outperforms comparison methods and supports downstream learning with limited task-agnostic data.

  • Experimental Setup: Agents first explore without downstream-task information, then receive task rewards for a separate downstream RL phase.The experiments use 500k task-agnostic interactions followed by 500k task-specific interactions.
  • Experimental Setup: Proto-RL is compared with Random exploration, Curiosity, APT, Plan2Explore, and DrQ under the experimental evaluation protocols.Plan2Explore uses model-based optimization and receives an estimate of the downstream reward, while DrQ anchors task-specific performance ranges.
  • Task-Agnostic Pre-training: Proto-RL significantly improves upon Random exploration and APT across all environments, and outperforms Curiosity in 7/8 environments.The evaluation reports results on eight environments, with extended results on sixteen environments in Appendix F.
  • Task-Agnostic Pre-training: Proto-RL with 500k downstream RL steps outperforms DrQ trained on 1M steps in 6/8 environments.The result is reported on difficult image-based environments from the DeepMind Control Suite.
  • Multi-Task Generalization: Proto-RL significantly outperforms the compared baselines after 500k downstream steps across multiple downstream tasks in the multi-task evaluation.The frozen representation and prototypes are applied separately to four tasks in each of two domains.
  • Efficiency of Task-Agnostic Pre-training: Proto-RL consistently outperforms baselines across 200k, 500k, and 1M task-agnostic training steps.On Walker Run, baseline performance drops with increased task-agnostic training, highlighting overfitting to explored data.
  • Downstream Exploration: Using prototype-driven exploration during downstream RL outperforms setting α = 0, especially on sparse-reward tasks.Across tasks, every experiment with α ≥0.1 outperforms α = 0; the trade-off parameter balances task reward and entropy-based intrinsic reward.
  • Downstream Exploration: Uniformly spread prototypes provide candidates both near and far from an embedding, stabilizing entropy estimation and intrinsic-reward learning.The reported effect leads to a more stable exploration strategy.

6. Conclusion

Proto-RL combines unsupervised representation learning, prototypical embeddings, and task-agnostic exploration for downstream control. Its implementation uses SAC-based agents and fixed pre-trained representations during downstream training.

  • Conclusion: Proto-RL simultaneously learns visual representations and prototypes while exploring environments without task-specific rewards.The learned representations and prototypes are intended to support downstream exploration, learning, and generalization across multiple tasks.
  • Conclusion: The benchmark uses DeepMind Control Suite image-based continuous-control tasks with 1000-step episodes, except Reach Duplo at 250 steps.Actions repeat for two environment steps, and observations stack three 3×84×84 RGB frames into 9×84×84 tensors.
  • Conclusion: The encoder processes 9×84×84 observations through four convolutional layers and produces 39,200-dimensional flattened representations.A linear projector maps these representations to 128 dimensions, while the predictor uses a 128→512→128 MLP.
  • Conclusion: Proto-RL learns 512 prototypes represented as 128-dimensional continuous vectors, using temperature τ = 0.1 and three Sinkhorn-Knopp relaxation iterations.The prototype and assignment settings define the discrete basis used by the representation-learning procedure.

B.4. Soft-Actor Critic Architecture

The architecture jointly trains Proto-RL’s representation and exploration components, then freezes the learned visual components for downstream task learning. Its intrinsic reward is computed from nearest-neighbor distances in the prototype-based latent space.

  • B.4. Soft-Actor Critic Architecture: The SAC implementation adds a 39,200→50 fully connected layer with LayerNorm to both actor and critic networks.It uses learning rate 10^-4, minibatch size 512, actor update frequency 1, and critic target momentum 0.01.
  • B.4. Soft-Actor Critic Architecture: Proto-RL jointly optimizes self-supervised and reinforcement-learning losses during task-agnostic pre-training.Representation parameters receive gradients only from the self-supervised loss, while the exploration agent is trained with the reinforcement-learning loss.
  • B.4. Soft-Actor Critic Architecture: After pre-training, Proto-RL freezes the encoder, projector, and prototypes during downstream training.The downstream agent uses the fixed encoder representations and combines task reward with prototype-based intrinsic reward.
  • B.4. Soft-Actor Critic Architecture: The downstream task actor is initialized from the exploration actor’s weights to support initial exploration.This initialization is used alongside the pre-trained representations and intrinsic reward.
  • B.4. Soft-Actor Critic Architecture: Representation learning predicts prototype assignments from online observations and matches them to Sinkhorn assignments from target-network observations.The loss is a cross-entropy-style assignment objective, followed by online updates and exponential-moving-average target updates.
  • B.4. Soft-Actor Critic Architecture: Proto-RL’s intrinsic reward uses 128-dimensional normalized representations and a queue containing 2048 candidates, with three nearest neighbors used for distance estimation.The reward is the distance to the third-nearest candidate in the queue.

D. The PointMass Maze Experiment Details

The PointMass maze modifies the DMC PointMass Easy environment with walls, concentrated initial states, and separate task-agnostic and task-specific phases. Exploration before the task uses entropy-based intrinsic reward.

  • D. The PointMass Maze Experiment Details: The U-maze adds three walls to the PointMass Easy MuJoCo environment.The walls are specified as box geometries positioned along the x and y directions.
  • D. The PointMass Maze Experiment Details: Initial states are sampled uniformly from the smaller top-left region [−0.3, −0.15] × [0.15, 0.3] instead of the full [−0.3, 0.3]×[−0.3, 0.3] grid.This changes the starting-state distribution before task-agnostic exploration.
  • D. The PointMass Maze Experiment Details: During task-agnostic training there is no target location, and the agent explores using entropy-based intrinsic reward; during task-specific training, the target is centered at [0, 0] with radius 0.07.The agent receives reward 1 upon reaching the target and otherwise receives the task’s non-success outcome.

E. The Multitask Experiment Details

The multitask experiments extend DMC environments with additional locomotion and reaching objectives. Walker receives four directional or flipping tasks, while Reach Duplo varies the target location.

  • E. The Multitask Experiment Details: Walker is extended with Run Forward, Run Backward, Flip Forward, and Flip Backward tasks.These tasks require corresponding forward, backward, or flipping behavior and resemble the Cheetah tasks used in Plan2Explore.
  • E. The Multitask Experiment Details: Reach Duplo requires the agent to reach a Lego block placed at four fixed locations: top-left, top-right, bottom-left, and bottom-right.The locations are [−0.09, 0.09], [0.09, 0.09], [−0.09, −0.09], and [0.09, −0.09].

F. Full Results for the Task-Agnostic Pre-training Experiment

Figure 10 evaluates task-agnostic pre-training across 16 challenging DeepMind Control Suite environments. Proto-RL consistently outperforms the baselines and often exceeds fully supervised performance.

  • 16 challenging DeepMind Control Suite environments are evaluated after task-agnostic and downstream training stages.Most methods receive 500k environment steps of pre-training followed by 500k task-specific steps; DrQ uses task reward from the outset.
  • Proto-RL consistently beats the baselines across the evaluated environments.
  • Proto-RL exceeds the fully supervised result in many cases.

G. Full Results for the Efficiency of Task-Agnostic Pre-training Experiment

The experiments vary both the duration of task-agnostic pre-training and the balance between exploration and downstream reward. Proto-RL supports downstream tasks after 200k pre-training steps, while α > 0 is especially useful for sparse-reward tasks.

  • Varying task-agnostic pre-training: 200k task-agnostic steps suffice for Proto-RL to explore state space and learn representations supporting downstream tasks.Subsequent task-specific training uses 500k steps on top of the frozen representation.
  • Varying task-agnostic pre-training: The efficiency experiment uses 500k steps of downstream training on top of the frozen representation.
  • Balancing exploration and downstream reward: α > 0 facilitates learning when Proto-RL balances exploration with downstream reward.
  • Balancing exploration and downstream reward: The benefit of α > 0 is especially pronounced in Cartpole Swingup Sparse and Reach Duplo.

I. Full Results on the Downstream Performance without Pre-training Experiment

Without prior task-agnostic pre-training, Proto-RL optimizes the downstream objective from the first training step but does not match its pretrained variant. The comparison shows that Proto-RL’s performance depends substantially on pre-training.

  • Proto-RL’s superior performance significantly depends on its task-agnostic pre-training scheme.
  • Without prior task-agnostic pre-training, Proto-RL and DrQ optimize their downstream objectives from the first training step.
  • Proto-RL outperforms DrQ without pre-training but remains inferior to Proto-RL (Pretrained).
Loading 2102.11271v2…