---
title: Action-Conditional Self-Prediction (BYOL-AC)
url: https://www.emergentmind.com/topics/action-conditional-self-prediction-byol-ac
type: topic
---

# Action-Conditional Self-Prediction (BYOL-AC)

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 [2406.02035].

## 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 [2406.02035].

Source: https://www.emergentmind.com/topics/action-conditional-self-prediction-byol-ac