Action-Conditional Self-Prediction (BYOL-AC)

• The paper introduces BYOL-AC, which conditions on future actions to improve state representations in reinforcement learning.
• It provides an exact ODE analysis and a variance decomposition that link policy-marginalized and action-conditional objectives with convergence guarantees.
• Empirical results across linear and deep RL settings demonstrate BYOL-AC’s superior performance in Q-function fitting and latent dynamics modeling.

Action-Conditional Self-Prediction (BYOL-AC) is a framework for representation learning in reinforcement learning (RL), situated within self-predictive learning paradigms. BYOL-AC extends bootstrap-based objectives by explicitly conditioning on future actions, rather than marginalizing over the agent’s action distribution. This modification yields improved state representations optimized for both model-based and model-free RL, enabling more precise dynamics modeling and value estimation in environments with complex, action-dependent transitions (Khetarpal et al., 2024).

1. Objective Function and Formalism

Let $X$ denote a finite state space with states represented one-hot as $x \in \mathbb{R}^{|X|}$, and $A$ the finite action set. For a fixed sampling distribution $d_X$ over $X$, a policy $\pi(\cdot|x)$, and per-action transition matrices $T_a \in \mathbb{R}^{|X| \times |X|}$, the goal is to learn a $k$-dimensional embedding matrix $\Phi \in \mathbb{R}^{|X| \times k}$ and per-action predictors $P_a \in \mathbb{R}^{k \times k}$.

The BYOL-AC loss is the expected squared prediction error in embedding space:

$$\min_{\Phi,\,\{P_a\}} \mathbb{E}_{x\sim d_X,\,a\sim\pi(\cdot|x),\,y\sim T_a(\cdot|x)} \left\| P_a^T \Phi^T x - \mathrm{sg}(\Phi^T y) \right\|_2^2,$$

where $\mathrm{sg}(\cdot)$ denotes the stop-gradient operator. The loss enforces that, for each $(x, a)$, the transformed embedding $P_a^T \Phi^T x$ predicts the embedding of the successor state $\Phi^T y$.

2. Continuous-Time ODE and Solution Characterization

BYOL-AC admits an exact continuous-time ODE characterization under simplifying assumptions: (i) orthonormal initialization and maintenance of $\Phi$, (ii) uniform $d_X$, (iii) symmetric $T_a$, (iv) uniform $\pi$, and (v) shared eigenvectors among $T_a$. The learning proceeds on two timescales:

• Inner loop: For fixed $\Phi$, each $P_a$ is set to its least-squares optimal solution, $P_a^* = \Phi^T T_a \Phi$.
• Outer loop: $\Phi$ evolves via the semi-gradient:

$$\dot{\Phi} = (I - \Phi\Phi^T)\frac{2}{|A||X|} \sum_a T_a\Phi\Phi^T T_a\Phi.$$

This ODE admits the following properties:

• Non-collapse: $\Phi^T \Phi = I$ is maintained for all $t$.
• Lyapunov function: $$f_{\mathrm{BYOL-AC}}(\Phi) = \frac{1}{|A|}\sum_a \mathrm{Tr}(\Phi^T T_a \Phi \Phi^T T_a \Phi)$$ is strictly increasing along ODE trajectories except at critical points, guaranteeing convergence.
• Stationary points: Local maximizers $\Phi^*$ span the top-$k$ eigenspace of $\frac{1}{|A|} \sum_a T_a^2$, aligning the learned representation with high-variance, action-sensitive dynamical modes.

3. Variance Identity and Connections to Policy-Marginalized Objectives

The action-conditional BYOL-AC formulation contrasts with its policy-marginalized counterpart, BYOL-$\Pi$, which replaces explicit $a$-conditioning with an expected transition operator $T^\pi = \mathbb{E}_a [T_a]$. The two are connected by a variance decomposition:

$$\frac{1}{|A|}\sum_a T_a^2 = (T^\pi)^2 + \mathrm{Var}_{ac},$$

where the diagonal term captures mean dynamics and the variance term quantifies diversity across actions in the same eigenbasis. Consequently, BYOL-AC representations are optimized for directions critical under the policy-averaged operator and those with large across-action variance. This identity underlies the introduction of the "variance-like" action-conditional objective (BYOL-VAR), whose fixed points relate to principal "advantage" directions, as further formalized in model-free interpretations.

4. Model-Based and Model-Free Interpretations

The BYOL-AC trace objective admits both model-based and model-free interpretations:

• Model-based (low-rank transition fitting): Maximizing $f_{\mathrm{BYOL-AC}}$ is equivalent to minimizing

$$\min_{(P_a)} \frac{1}{|A|} \sum_a \| T_a - \Phi P_a \Phi^T \|_F^2 + C$$

for constant $C$, i.e., fitting each $T_a$ with a rank-$k$ factorization. This situates BYOL-AC as learning a $k$-dimensional linear model of action-conditional dynamics, potentially advantageous for latent-space planning.

• Model-free (value, Q, and advantage fitting): For states equipped with i.i.d. isotropic Gaussian rewards $R$ ($\mathbb{E}[RR^T]=I/|X|$), minimizing the BYOL objectives yields minimum MSE projections of value-like functions onto the $\Phi$-span:
  • BYOL-$\Pi$ $f_{\mathrm{BYOL-\Pi}}$ fits one-step values: $$\min_{\theta,\omega}\mathbb{E}\left[ \| T^\pi R-\Phi\theta \|^2 + \| T^\pi \Phi\Phi^T R - \Phi\omega \|^2 \right]$$
  • BYOL-AC $f_{\mathrm{BYOL-AC}}$ fits one-step Q-functions: average over $a$ of $$\min_{\theta_a,\omega_a}\left(\|T_a R-\Phi\theta_a\|^2 + \|T_a\Phi\Phi^T R-\Phi\omega_a\|^2\right)$$
  • Variance objective fits one-step advantage functions.

This duality clarifies the relationship between action awareness in the objective and the learned latent features' suitability for different RL signal decompositions.

5. Empirical Evaluation

Empirical comparisons cover both linear function approximation and deep RL environments:

• Linear setting: Over 100 random symmetric MDPs ($|X|=10, |A|=4, k=4$), BYOL-AC, BYOL-$\Pi$, and BYOL-VAR representations minimize their respective trace objectives and best fit Q, value, and advantage functions in MSE.
• Deep RL benchmarks: On four Minigrid tasks with an on-policy V-MPO agent (augmented by each auxiliary loss) and on CartPole, MountainCar, Acrobot with DQN plus BYOL-$\Pi$ or BYOL-AC auxiliary losses, the action-conditional (BYOL-AC) variant matches or outperforms BYOL-$\Pi$. In contrast, the variance objective performs poorly for control tasks.
• Ablations: Variation of prediction horizon and policy noise indicates that action-conditional prediction with BYOL-AC yields more robust and task-relevant features compared to policy-marginalized objectives.

A summary of the correspondence between objectives and RL signals:

Objective	Optimizes	Best fits
BYOL-$\Pi$	$(T^\pi)^2$	Value function
BYOL-AC	$\mathbb{E}_a[T_a^2]$	Q-function
BYOL-VAR	$\text{Var}_{ac}$	Advantage function

6. Significance and Theoretical Implications

BYOL-AC constitutes an extension of self-predictive representation learning to account for action-conditional dynamics, bridging the gap between theoretical analyses grounded in policy-marginalized assumptions and practical implementations in deep RL. The ODE theory specifies the geometry of fixed points (as principal subspaces of expected $T_a^2$), supports convergence guarantees, and links the approach to well-known decompositions in RL. The variance decomposition clarifies the benefit of action conditioning: features can capture task-relevant, action-specific structure overlooked by value-centric, policy-averaged approaches. Empirically, BYOL-AC outperforms policy-marginalized and variance-focused analogues in both low- and high-dimensional control environments (Khetarpal et al., 2024).