Value-Consistent Representation Learning

• VCR is a representation learning framework that enforces alignment between imagined latent states and real Q-value estimates to bridge the gap between visual similarity and decision relevance.
• It leverages a latent transition model and an auxiliary self-predictive representation loss to preserve value-relevant features during state encoding.
• Empirical evaluations on Atari and DeepMind Control Suite demonstrate superior performance and improved sample efficiency compared to traditional visual representation methods.

Value-Consistent Representation Learning (VCR) is a framework for representation learning in deep reinforcement learning (RL) that explicitly aligns learned state representations with value predictions, addressing the semantic gap between visual similarity and value-relevant features. Unlike contrastive or reconstruction-based visual representation learning, VCR targets the ability to imagine future states in latent space such that their action-value predictions (Q-values) are consistent with those of real environment states. This approach is designed to boost data efficiency, particularly in regimes with limited environment interactions, by enforcing value-consistency across imagined and real transitions (Yue et al., 2022).

1. Motivation and Problem Definition

Conventional deep RL from pixels employs visual representation learning techniques like contrastive learning and data augmentation to train state encoders. While effective at distinguishing observations visually, these methods do not guarantee that encoded states preserve value-relevant information critical for decision making. Two visually similar observations may yield divergent returns, and standard transition models may not focus on preserving aspects necessary for accurate value estimation. VCR addresses this limitation by training the representation space such that state-action pairs propagated in latent space (“imagined states”) yield Q-value predictions that match those produced by the encoder when applied to real observations. This value-alignment target enforces that the latent space supports downstream control requirements, not just perceptual similarity (Yue et al., 2022).

2. Theoretical Formulation

Consider the following notation:

• $s_t \in \mathcal{S}$: raw observation at time $t$
• $a_t$: action at time $t$
• $f(\cdot;\theta_f)$: encoder mapping $s_t$ to latent $z_t = f(s_t)$
• $h(\cdot, \cdot; \theta_h)$: latent transition model producing $z_{t+1} = h(z_t, a_t)$
• $Q(z,a; \theta_Q)$: Q-value function in latent space
• Target networks $t$0: exponentially moving averages for stable learning
• $t$1: imagination horizon (steps for unrolling the transition model)
• $t$2: multi-step return length (off-policy target computation)

Imagined latent rollouts are obtained by recursively applying the transition model: $t$3

For each step $t$4, the value-consistency pairing is enforced by comparing the action-value outputs from $t$5 and $t$6, where $t$7 is encoded from the true next observation.

The per-step value-consistent distance is defined as

$t$8

where

$t$9

and $a_t$0 is the n-step return.

The total VCR loss over horizon $a_t$1 is

$a_t$2

An auxiliary self-predictive representation (SPR) cosine loss is included: $a_t$3

The combined training objective is: $a_t$4 with standard RL losses for $a_t$5 (e.g., Rainbow DQN for discrete, SAC for continuous control).

3. Implementation Details

Discrete-Action Setting (Atari 100K)

• Encoder: 3 conv layers (channels 32–64–64, kernels 8–4–3, strides 4–2–1) with linear projection, $a_t$6.
• Latent Transition Model: 2 conv layers, concatenating $a_t$7 and $a_t$8 in channel space.
• Q-Head: 2-layer MLP yielding categorical (C51) Q-value distribution; employs noisy-net, dueling architecture, double Q.
• Auxiliary Loss: SPR predictor head attached to Q-head’s first layer, used only for $a_t$9.
• Loss Formulation: $t$0 follows Rainbow DQN, $t$1 aligns imagined and target distributions with cross-entropy.

Continuous-Action Setting (DeepMind Control)

• Encoder: 4 conv layers plus linear, $t$2.
• Latent Transition: 2-layer MLP (1024 hidden units, layer norm).
• Critics: Twin Q-functions $t$3 and target critics.
• Actor: Policy $t$4 optimized as in SAC.
• Action Sampling: For each step, considers the real $t$5 plus $t$6 random actions (from $t$7 and/or uniform sampling).
• VCR Loss: Mean squared error between Q-values and pseudo-targets as described, leveraging n-step returns and target critics.
• Optimization: Independent Adam optimizers for critic, actor, $t$8, and VCR losses (batch sizes: 512 for SAC, 128 for VCR).

4. Algorithmic Workflow

The VCR learning process can be summarized as:

• Initialize all network components (encoder, transition, Q-head, actor where relevant), and target networks as EMA copies.
• Use a replay buffer $t$9 to accumulate environment transitions.
• For each iteration:
  • Sample a minibatch of $f(\cdot;\theta_f)$1 trajectories of length $f(\cdot;\theta_f)$2 from $f(\cdot;\theta_f)$3.
  • Apply random-shift augmentation to observations.
  • Encode initial states; unroll imaginations via the latent transition model.
  • Encode real subsequent observations with the target encoder.
  • For each step $f(\cdot;\theta_f)$4, compute $f(\cdot;\theta_f)$5 and $f(\cdot;\theta_f)$6 over all considered actions.
  • Sum losses over horizon, compute $f(\cdot;\theta_f)$7.
  • Update parameters by minimizing $f(\cdot;\theta_f)$8.
  • EMA update of all target networks.

This process robustly ties the utility of imagination in latent space to the anticipated return, enforcing utility-aligned representation learning.

5. Empirical Results and Benchmark Comparisons

Domains and Metrics

Experiments were conducted on:

• Atari 100K (26 games, 100K frames, 10 seeds)
• DeepMind Control Suite, 6 environments (100K, 500K frames, 10 seeds)
• Evaluation: Interquartile Mean of Human-Normalized Score (IQM HNS), Optimality Gap as in Agarwal et al. (NeurIPS 2021).

Discrete-Action Performance

• Baselines: SimPLe, DER, OTR, CURL, DrQ, SPR, PlayVirtual.
• VCR achieves IQM 39.7%, outperforming SPR (33.7%) and PlayVirtual (37.4%).
• Optimality Gap for VCR: 54.4% versus 57.7% for SPR.
• Performance-profile curves indicate VCR surpasses baselines over a range of score thresholds.

Continuous-Action Performance

• Baselines: Dreamer, SAC+AE, SLAC, CURL, DrQ, SPR, PlayVirtual.
• VCR achieves IQM 745 versus 700 for SPR; Optimality Gap 283 vs 335 for SPR.
• At 500K frames, VCR achieves a median score of 934 compared to DrQ's 929.

Ablation Studies

• Using mean squared error (MSE) distance only on the real action in $f(\cdot;\theta_f)$9 degrades performance (IQM: 30.9%, worse than SPR).
• This suggests the inclusion of all actions in the value-consistent pairing is important for performance.

6. Significance and Implications

VCR demonstrates that enforcing value consistency in representation learning substantially improves sample efficiency in RL, particularly in regimes with scarce observational data. By closing the gap between visually-aligned and value-aligned features, VCR achieves superior empirical performance across both discrete and continuous control tasks. The methodology indicates that shaping the learned latent space to reflect decision-centric, rather than perception-centric, similarity is critical for efficient policy learning. A plausible implication is that similar value-targeted representation learning approaches could benefit other domains where latent imagination and value estimation must be tightly coupled (Yue et al., 2022).