PaperlayerView exact revision on arXiv

Source-linked AI summary

Understanding Self-Predictive Learning for Reinforcement Learning

Yunhao Tang, Zhaohan Daniel Guo, Pierre Harvey Richemond, Bernardo Ávila Pires, Yash Chandak, Rémi Munos, Mark Rowland, Mohammad Gheshlaghi Azar, Charline Le Lan, Clare Lyle, András György, Shantanu Thakoor, Will Dabney, Bilal Piot, Daniele Calandriello, Michal Valko

arXiv:2212.03319v1cs.LGcs.AIstat.ML

TL;DR

Self-predictive learning can minimize its prediction error with constant representations, leaving unclear why it learns useful representations. The paper analyzes its learning dynamics and shows that, under idealized assumptions, they connect to spectral decomposition of the state transition matrix.

  • Problem

    Self-predictive learning’s prediction error can be minimized by constant representations, and there is no obvious theoretical justification that minimizing it yields meaningful representations.

    Source: §1 · p. 1

  • Method

    The paper analyzes self-predictive learning dynamics in an idealized setting to identify conditions for avoiding collapse and characterize the information captured by latent predictions.

    Source: §1 · pp. 1–2, §2.1 · p. 2

  • Results

    Under idealized assumptions, the learning dynamics locally improves a trace objective, whose maximization corresponds to spectral decomposition of the state transition matrix.

    Source: §1 · p. 2

  • Takeaways & Limitations

    The spectral-decomposition connection provides a partial theoretical understanding of why self-predictive learning can be useful in practice.

    Source: §1 · p. 2

    The dynamics are not guaranteed to converge to the top k eigenvectors and can instead converge to a non-eigenvector critical point.

    Source: §3.2 · p. 5

Abstract

from arXiv · show

We study the learning dynamics of self-predictive learning for reinforcement learning, a family of algorithms that learn representations by minimizing the prediction error of their own future latent representations. Despite its recent empirical success, such algorithms have an apparent defect: trivial representations (such as constants) minimize the prediction error, yet it is obviously undesirable to converge to such solutions. Our central insight is that careful designs of the optimization dynamics are critical to learning meaningful representations. We identify that a faster paced optimization of the predictor and semi-gradient updates on the representation, are crucial to preventing the representation collapse. Then in an idealized setup, we show self-predictive learning dynamics carries out spectral decomposition on the state transition matrix, effectively capturing information of the transition dynamics. Building on the theoretical insights, we propose bidirectional self-predictive learning, a novel self-predictive algorithm that learns two representations simultaneously. We examine the robustness of our theoretical insights with a number of small-scale experiments and showcase the promise of the novel representation learning algorithm with large-scale experiments.

1. Introduction

Self-predictive representation learning has empirical success, but its prediction objective also admits constant representations that retain no information. The paper analyzes how optimization dynamics avoid collapse, connects the learning process to spectral decomposition, and develops bidirectional self-predictive learning.

1. Introduction

  • Self-prediction extends the bootstrapping idea used in value-based temporal-difference learning to representation learning, where a learned transition function predicts future latent representations.In value-based RL, the next-step value estimate is the target for the current value estimate.Source: §1 · p. 1
  • Constant representations minimize the prediction error while retaining no information, leaving the objective without an obvious theoretical basis for learning meaningful representations.The paper asks how this apparent defect can be reconciled with practical algorithms’ empirical success and whether their learned representations help downstream RL.Source: §1 · p. 1
  • Two-time-scale optimization of the predictor and representation, together with semi-gradient updates on the representation, prevents convergence to trivial solutions from random initializations.The semi-gradient update maintains the representation’s capacity during learning.Source: §1 · pp. 1–2
  • In an idealized setting, the learning dynamics locally improve a trace objective for transition information, whose maximization corresponds to spectral decomposition of the state transition matrix.This offers a partial theoretical account of self-predictive learning’s practical usefulness.Source: §1 · p. 2
  • Bidirectional self-predictive learning simultaneously learns two representations from forward and backward predictions and has more general theoretical guarantees than the one-directional method.Tabular and deep RL experiments report more consistent and stable performance.Source: §1 · p. 2

2. Background

The background sets up a reward-free, finite-state MDP and a tabular state representation. Self-predictive learning models transitions in latent space with a prediction matrix, although naive optimization can yield a zero representation.

2. Background

  • The analysis considers a reward-free MDP with finite state and action spaces, a transition kernel, a discount factor, and a fixed policy.Rewards are omitted because they are not needed for the subsequent discussion.Source: §2 · p. 2
  • States are represented as one-hot vectors, while a representation matrix embeds each state into a k-dimensional latent vector, typically with k much smaller than the state-space size.Useful representations should share information across states and support downstream tasks such as policy evaluation or control.Source: §2 · p. 2

2.1. Self-predictive learning

  • Self-predictive learning models transitions between latent states with a prediction matrix P and minimizes latent-space reconstruction loss.The predicted next latent is P^T Φ^T x, and the target is Φ^T y.Source: §2.1 · p. 2
  • Naively optimizing the reconstruction objective can yield the zero representation, which achieves zero loss despite being trivial.The framework analyzes sequentially sampled state pairs under the policy; its prediction target uses a stop-gradient in practice.Source: §2.1 · p. 2, Figure 1 · p. 2

3. Understanding learning dynamics of self-predictive learning

The section analyzes how self-predictive learning avoids collapse and how its representation dynamics relate to spectral information in the state-transition matrix. It identifies optimization conditions for non-collapse and establishes a trace-objective result under symmetric transitions, while noting limits on convergence and non-symmetric cases.

3.1. Non-collapse property of self-predictive learning

  • When either safeguard is removed, representation columns start to collapse; Figure 2 tracks normalized-column cosine similarity across 100 random-MDP runs.The columns start orthogonal, so higher cosine similarity indicates stronger evidence of collapse.Source: Figure 2 · p. 3, §3.1 · p. 4
  • Faster optimization of predictor P and semi-gradient updates to Φ preserve representation diversity, preventing differently initialized vectors from converging to the same vector under the analyzed dynamics.The covariance matrix ΦᵀΦ is conserved; the non-collapse result also extends to a broader loss class and a modified cosine-similarity loss.Source: §3.1 · pp. 3–4

3. Understanding learning dynamics of self-predictive learning

  • The analysis assumes orthonormal initialization and a uniform first-state distribution for the idealized learning dynamics.The orthonormal-initialization assumption can be approximately satisfied when the state space is large relative to the representation dimension.Source: §3.1 · p. 4, §3.2 · p. 4

3.2. Self-predictive learning as eigenvector decomposition

  • Even in the two-state, k=1 example, non-eigenvector critical points coexist with eigenvector points, and initialization near the bottom eigenvector can converge away from the top.The analysis leaves a more comprehensive study of critical points in the general case to future work, especially when Pπ has complex eigenvectors.Source: §3.2 · pp. 4–5, Figure 3 · p. 5
  • For symmetric Pπ, the trace objective is non-decreasing and strictly improves away from critical points; its constrained maximizers span the top-k eigenvectors by absolute eigenvalue.The dynamics therefore tends toward representations with high trace objective, interpretable as principal-subspace PCA on Pπ.Source: §3.2 · p. 5
  • For non-symmetric Pπ, trace-objective improvement need not be monotone; simulations on 100 generated tabular MDPs found overall gains but also slow or marginal-improvement cases.The dynamics is not guaranteed to reach the maximum objective value in these cases.Source: §3.2 · p. 5

4. Extensions of the theoretical analysis

The extensions examine how self-predictive learning relates to extra prediction functions, multi-step and action-conditional transitions, and partially observable settings. They also note that finite learning rates can yield partial collapse despite retaining spectral information.

Scope

  • The analysis uses a deliberately simple model setup to expose essential features while leaving extensions open.The authors frame this simplicity as a scope choice, not as a general characterization of practical systems.Source: §4 · p. 6

Additional prediction function

  • With an additional prediction matrix Q, composing P and Q during optimization extends the non-collapse result to the extra prediction function used in practice.P models latent transitions, while Q adds flexibility for matching the predicted latent to the next representation.Source: §4 · p. 6

Multi-step and action-conditional latent prediction

  • For multi-step prediction, the spectral view applies to the n-step transition (P π)^n; action-conditional matrices Pa and representations Φa correspond to the transition matrix P πa for each action.The cited passage notes that multi-step prediction improves practical performance.Source: §4 · p. 6

Partial observability

  • In partially observable settings, recurrent representations can avoid enumerating histories, and under certain conditions the analysis extends to the history-transition matrix.The authors call for more refined analysis to better bridge theory and empirical results in POMDPs.Source: §4 · p. 6

Bidirectional self-predictive learning

  • Algorithm 2 alternates optimal forward and backward prediction-matrix computations with representation updates, then outputs the final pair of representations.The pseudocode specifies representations initialized for k ≤ |X| and iterates through the updates.Source: §4 · p. 7

Finite learning rate and other factors

  • At finite learning rates, representations can partially collapse while still capturing spectral information; the analysis relies on continuous-time dynamics.The authors identify formalizing finite-rate behavior as a future direction.Source: §4 · pp. 6–7

5. Bidirectional self-predictive learning with left and right representations

The section introduces bidirectional self-predictive learning, which learns paired representations through forward and backward prediction. Under stated assumptions, its dynamics improve an SVD trace objective, with experiments illustrating an advantage over single-representation learning when left and right singular vectors differ.

5. Bidirectional self-predictive learning with left and right representations

  • Bidirectional self-predictive learning uses left and right representations, each predicting across time in opposite directions with targets from the other representation.Both representations use semi-gradient updates.Source: §5 · p. 7
  • Non-collapse for both representations follows when the forward and backward predictors are optimally adapted by solving their least-squares problems.The section states this optimal adaptation is a key requirement for non-collapse in the new dynamics.Source: §5 · p. 7

5.1. Bidirectional self-predictive learning as singular value decomposition

  • Backward prediction requires the reverse process to be Markov, meaning (P π)T is also a transition matrix, a weaker condition than symmetry of P π.The symmetry condition is required for the single-representation dynamics discussed in Theorem 6.Source: §5.1 · p. 7
  • The SVD trace objective increases strictly away from critical points, and its constrained maximizers span the top k singular-vector pairs of P π.Any k singular-vector pairs can be critical points, but the objective measures how informative those pairs are.Source: §5.1 · p. 8
  • Because left and right singular vectors can differ, paired representations can learn both, whereas one representation must interpolate between them and may have non-monotonic trace-objective behavior.The distinction motivates using two representations rather than compromising between the two singular-vector sets.Source: §5.1 · p. 8
  • In an MDP designed with very different left and right singular vectors, bidirectional learning steadily improves its trace objective while single-representation learning mostly stalls near initialization.The comparison uses 100 random orthonormal initializations; the plot reports objective ratios over training iterations.Source: §5.1 · p. 8
  • Bidirectional learning captures both top left and right singular vectors, providing a more comprehensive account of transition dynamics and richer representation-learning signals with nonlinear function approximation.The passage connects left singular vectors to approximating value functions, while noting that learning them as standalone objects is challenging.Source: §5.1 · pp. 8–9

6. Prior Work

The section situates the paper among work on non-collapse mechanisms, spectral methods for RL representations, forward-backward representations, and gradient-based PCA and SVD algorithms. It positions self-predictive learning as a connection between empirically oriented RL algorithms and spectral decomposition.

Non-collapse mechanisms

  • Prior non-collapse analyses for non-contrastive learning emphasize semi-gradients and regularization, prediction heads, or data augmentation; this paper describes its RL mechanism as qualitatively different.The authors suggest the result may also inform understanding of unsupervised learning.Source: §6 · p. 9

RL representation learning via spectral decomposition

  • Transition- and successor-matrix spectral methods include eigenvector approaches with symmetry or real-diagonalizability assumptions and more generally applicable SVD approaches.The cited SVD approaches have shown theoretical benefits for downstream RL tasks.Source: §6 · p. 9

Forward-backward representations

  • Forward-backward representations learn value functions and visitation distributions, connecting to left and right singular vectors; bidirectional self-prediction seeks those vectors with a simpler design based purely on self-prediction.The passage describes the connection as high-level rather than an identical design.Source: §6 · p. 9

Algorithms for PCA and SVD with gradient-based update

  • The paper connects self-predictive dynamics to incremental gradient-based PCA and SVD algorithms that compute leading eigenvectors or singular vectors.Examples include Oja’s subspace algorithm and its SVD extension.Source: §6 · p. 9

7. Deep RL implementation

The deep bidirectional method adapts BYOL-RL with backward prediction and is evaluated alongside BYOL-RL on DMLab-30. Across tasks, the methods perform comparably overall, while backward prediction yields sizable gains on some tasks.

7. Deep RL implementation

  • The deep method builds on BYOL-Explore’s representation learning but omits its exploration bonuses to focus on representation learning.The paper presents the adaptation as inspired by its theoretical discussion.Source: §7 · p. 9
  • In BYOL-RL, a convolutional network and LSTM encode observation histories, and an open-loop LSTM predicts a multi-step future latent target.The target is computed from the observation at the n-step future time.Source: §7 · p. 9
  • The backward objective mirrors forward prediction with the time axis reversed; PBL’s reverse prediction instead matches an observation-derived target and does not predict backward over time.The passage distinguishes PBL’s motivation in partial observability from the deep bidirectional method.Source: §7 · pp. 9–10

7.1. Experiments

  • The comparison uses BYOL-RL, which is built on V-MPO and shapes representations through policy gradients without an explicit representation-learning objective.The testbed is the multi-task setting across 30 partially observable DMLab-30 tasks.Source: §7.1 · p. 10
  • Across 30 DMLab tasks, deep bidirectional self-predictive learning and BYOL-RL perform comparably overall, while backward prediction improves some tasks by as much as 0.4 human normalized score.A normalized score of 1 corresponds to human-level performance on a task.Source: §7.1 · p. 10

8. Conclusion

The paper analyzes self-predictive learning dynamics, relates them to spectral decomposition of transition dynamics, and derives a bidirectional representation-learning algorithm. Its deep RL instantiation is reported as empirically comparable to state-of-the-art performance, while per-game comparisons with BYOL-RL are mixed.

8. Conclusion

  • The analysis connects self-prediction to spectral decomposition of transition dynamics and identifies algorithmic components for non-collapse.The authors describe this as a first justification for the practical efficacy of some empirically motivated algorithms.Source: §8 · p. 10
  • Against BYOL-RL, bidirectional learning gains on some games but degrades performance on others.Figure 6 reports end-of-training mean human normalized scores averaged across 3 seeds.Source: Figure 6 · p. 10
  • A deep RL instantiation of the bidirectional algorithm is reported as empirically comparable to state-of-the-art performance.The algorithm mirrors SVD-based methods and offers more flexibility than single-representation self-prediction.Source: §8 · p. 10
  • The authors identify limitations of linear models as a reason to study nonlinear latent predictions and representations.They also point to studying interactions with RL algorithms and unsupervised learning as future directions.Source: §8 · p. 10

APPENDICES: Understanding Self-Predictive Learning for Reinforcement Learning

The appendix derives the continuous-time ODE systems for self-predictive and bidirectional learning, then relates assumptions used in the bidirectional derivation. The derivations rely on stated distributional and transition-matrix conditions.

A. Detailed derivations of ODE systems

  • The appendix derives the ODE systems by expanding expectations over one-hot state encodings and a diagonal matrix of state probabilities.The derivations cover both the self-predictive and bidirectional dynamics.Source: A. Detailed derivations of ODE systems · p. 13

A.1. Equation (5) for self-predictive learning

  • For Equation (5), the derivation makes the predictor’s optimality condition explicit and simplifies the resulting dynamics under Assumptions 3 and 4.The appendix identifies the simplified result as the ODE system in Equation (5).Source: §A.1 · p. 13

A.2. Equation (9) for bidirectional self-predictive learning

  • For Equation (9), identity sampling and a doubly-stochastic transition matrix make the next-state second moment expressible using the state-distribution matrix.The derivation also identifies the backward predictor as a Markov transition matrix for the reverse chain.Source: §A.2 · p. 13

A.3. Equivalence between assumptions in deriving Equation (9)

  • Given a uniform first-state distribution, a doubly-stochastic transition matrix and a uniform next-state marginal are equivalent assumptions.The appendix shows that either condition makes the uniform distribution stationary for the transition matrix.Source: §A.3 · p. 14

B. Proof of theoretical results

The proofs characterize non-collapse and objective behavior for self-predictive and bidirectional dynamics, then extend the non-collapse result to a broader class of losses. The spectral conclusions depend on the stated symmetry or Assumption 8 conditions.

B. Proof of theoretical results

  • Under Equation (4), the covariance matrix is constant, so differently initialized representation vectors cannot converge to the same vector.The covariance result is also extended to general losses whose minimizer satisfies the first-order optimality condition.Source: B. Proof of theoretical results · p. 14
  • The critical points include matrices with orthonormal columns spanning any set of k eigenvectors of P π.The proof verifies this using the eigenvector relation P πU = UΛ.Source: B. Proof of theoretical results · pp. 14–15
  • For symmetric P π, the trace objective is non-decreasing and is maximized by representations spanning the eigenvectors with the top absolute eigenvalues.The derivative is strictly positive away from critical points.Source: B. Proof of theoretical results · p. 15
  • In bidirectional dynamics, both representation covariance matrices remain constant, and critical points include orthonormal representation pairs spanning singular-vector pairs.The proof constructs these critical points from singular vectors of P π.Source: B. Proof of theoretical results · p. 16
  • Under Assumption 8, the bidirectional SVD trace objective is non-decreasing and is maximized by representation pairs spanning the top-k singular-vector pairs.The objective increases strictly away from its critical points.Source: B. Proof of theoretical results · pp. 16–17

C. Extension of non-collapse property to general loss function

  • The non-collapse result extends to general losses when the predictor minimizer satisfies the first-order optimality condition.Examples named include L1 loss and regularized cosine-similarity loss.Source: C. Extension of non-collapse property… · p. 18

D. Discussions on critical points to the self-predictive learning dynamics

Eigenvector subspaces are critical points of the self-predictive dynamics, but they need not be the only critical points. In one example, initialization near the bottom eigenvector leads to a non-eigenvector critical point rather than the top eigenvector.

D. Discussions on critical points to the self-predictive learning dynamics

  • Orthonormal representations spanning any k eigenvectors are critical points when the transition matrix is real diagonalizable.The paper denotes this set by C, which is a subset of all critical points CPπ.Source: D. Discussions on critical points to… · p. 18
  • For a symmetric transition matrix with eigenvalues 1 and 0.8, all critical points are eigenvector-based, so C = CPπ.The argument uses the nonzero predictor at a critical point to show the representation span is invariant and composed of eigenvector subspaces.Source: D. Discussions on critical points to… · pp. 18–19
  • With a symmetric transition matrix whose eigenvalues are 1 and −0.8, the dynamics also has non-eigenvector critical points formed from combinations of eigenvectors.These points belong to CPπ but not C.Source: D. Discussions on critical points to… · p. 19
  • Starting near the bottom eigenvector in this example, the dynamics converges to one of four non-eigenvector critical points, not the top eigenvector.The figure depicts four eigenvector critical points in blue, four non-eigenvector points in red, and local update directions as black arrows.Source: D. Discussions on critical points to… · p. 19, Figure 7 · p. 20

E. Algorithmic and implementation details on BYOL-RL

The appendix describes BYOL-RL and its bidirectional extension, then details the tabular experiments. The reported trace objective generally improves, while target-network updates retain improvement but can make it less monotonic.

E. Algorithmic and implementation details on BYOL-RL

  • BYOL-RL embeds observations and recurrent histories, predicts future history representations from action sequences, and uses semi-gradient descent with a stop-gradient target.Its observation embedding is implemented as a convolutional network, and its recurrent embedding as an LSTM core output.Source: E. Algorithmic and implementation… · p. 20

E.1. Details on deep bidirectional self-predictive learning with BYOL-RL

  • The bidirectional extension adds a backward prediction process that mirrors the forward process and shares its observation embedding to reduce computational cost.The backward process has its own observation and recurrent embeddings, latent dynamics, and projection function.Source: §E.1 · p. 21

E.2. Differences from PBL (Guo et al., 2019)

  • Unlike PBL, the backward prediction here seeks to predict the recurrent history embedding from the observation embedding, while sharing the forward observation encoder.The supplied passages describe PBL as using a projection function for its backward prediction.Source: §E.2 · p. 21

F. Transition matrix to illustrate the failure mode of self-predictive learning

  • The example transition matrix illustrates directional clustering: its top left singular vector groups the first two states, while the top right singular vector groups the last two.Left singular vectors correspond to similar outgoing distributions; right singular vectors correspond to similar incoming distributions.Source: F. Transition matrix to illustrate the… · p. 21

G. Experiment details; G.1. Tabular experiments

  • Tabular experiments use randomly generated transition matrices, typically with 20 states and k = 2 representation columns.The experiments include doubly stochastic matrices and symmetric matrices generated from them.Source: §G.1 · p. 22

G.1. Experiment details

  • Discretized tabular updates use a finite representation learning rate, with η = 10^-3 as the default so the non-collapse property is almost satisfied.The experiments also simulate exact ODE dynamics for other figures.Source: §G.1 · p. 22

G.1.1. COMPLETE RESULT TO FIG. 4

  • Over 100 iterations, the trace objective improves significantly overall, although nonsymmetric runs can show small oscillations and some instances improve slowly.The complete-result figure shows the longer evolution and individual runs across MDPs.Source: §G.1.1 · p. 22, Figure 8 · p. 23

G.1.2. EFFECT OF TARGET NETWORK

  • With a target representation updated by moving average, the trace objective still improves, but β = 0 yields more limited and potentially non-monotonic improvement.For β > 0, the dynamics behaves similarly to the dynamics without a target representation.Source: §G.1.2 · pp. 22–23

G.1.3. ABLATION ON HOW HYPER-PARAMETERS IMPACT NON-COLLAPSE DYNAMICS

The ablations test how finite learning rates and predictor imperfection affect non-collapse, alongside a deep-RL comparison. Higher learning rates and greater predictor noise correspond to greater collapse tendencies.

G.1.3. ABLATION ON HOW HYPER-PARAMETERS IMPACT NON-COLLAPSE DYNAMICS

  • As the finite learning rate increases, the cosine similarity between initially orthogonal representation vectors rises, indicating a more severe violation of non-collapse.The comparison uses η values 0.01, 0.1, 1, and 10, with a fixed total step size across rates.Source: §G.1.3 · p. 23, Figure 10 · p. 25
  • Increasing Gaussian noise in the predictor matrix increases the tendency for the representation vectors to collapse.The noise scale σ controls the corruption of the optimal prediction matrix.Source: §G.1.3 · p. 24

G.2. Deep RL experiments

  • The deep bidirectional algorithm adds backward prediction to BYOL-RL, weighting forward and backward predictions equally with α = 1.The method is compared with BYOL-RL, which combines a forward prediction loss with V-MPO's reinforcement-learning loss.Source: §G.2 · p. 24

G.1.3. ABLATION ON HOW HYPER-PARAMETERS IMPACT NON-COLLAPSE DYNAMICS

  • Figure 10 tracks cosine similarity, with values nearer 1 indicating that the initially orthogonal representations are aligning toward collapse.The ablation examines finite learning rates and non-optimal prediction matrices.Source: Figure 10 · p. 25

G.2. Deep RL experiments

  • BYOL-RL significantly outperforms the RL baseline in mean human-normalized scores across most games.Figure 11 reports scores averaged over three seeds at the end of training; the deep-RL test bed is DMLab-30's 30 partially observable tasks.Source: §G.2 · p. 25, Figure 11 · p. 26
Loading 2212.03319v1…