🤖 AI

Regularized Latent Dynamics Prediction is a Strong Baseline For Behavioral Foundation Models

This paper proposes Regularized Latent Dynamics Prediction (RLDP), a method that enhances zero-shot reinforcement learning in Behavioral Foundation Models by adding simple orthogonality regularization to latent dynamics prediction, thereby maintaining feature diversity and outperforming complex state-of-the-art approaches, particularly in low-coverage scenarios.

Original authors: Pranaya Jajoo, Harshit Sikchi, Siddhant Agarwal, Amy Zhang, Scott Niekum, Martha White

Published 2026-03-18

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Pranaya Jajoo, Harshit Sikchi, Siddhant Agarwal, Amy Zhang, Scott Niekum, Martha White

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a robot to do anything a human can do: walk, run, dance, or carry groceries. The problem is, you can't write a specific instruction manual for every single possible task the robot might face in the future. That's where Behavioral Foundation Models (BFMs) come in.

Think of a BFM as a universal "muscle memory" library. Instead of learning one specific skill at a time, the robot learns a general understanding of how the world works during a "pre-training" phase. Later, when you give it a new goal (like "go to the kitchen"), it can instantly figure out how to move without needing to relearn everything from scratch.

However, there's a catch. To build this library, the robot needs a good way to "see" and understand the world. If it sees the world through a blurry, distorted lens, it will fail at new tasks.

The Problem: The "Blurry Lens" and the "Sticky Note"

Current methods for teaching robots this general understanding are like trying to solve a massive, complex math puzzle while blindfolded. They try to predict the future by simulating every possible path a robot could take.

The Issue: These methods are computationally heavy and prone to errors. If the robot tries to imagine a path it hasn't seen before, the simulation gets messy, and the "lens" (the representation) gets distorted.
The Collapse: A simpler method involves just predicting "if I do action X, what happens next?" But the paper found that if you only do this, the robot's brain starts to get lazy. It begins to think that every state looks the same. It's like a student who stops studying and decides that "all history is just a blur of dates." The robot loses the ability to tell the difference between walking and running because its internal map has collapsed into a single, useless point.

The Solution: RLDP (Regularized Latent Dynamics Prediction)

The authors propose a new, simpler method called RLDP. They use a clever analogy to fix the "lazy brain" problem.

The Analogy: The Gymnast and the Stretching Mat

Imagine the robot's internal map of the world is a gymnast trying to learn a routine.

The Old Way (Complex Methods): The coach tries to teach the gymnast by showing them a video of every possible routine they could ever do, calculating the perfect score for each one. It's exhausting, confusing, and if the video glitches, the gymnast gets confused.
The Simple Way (Latent Dynamics): The coach just says, "If you jump here, you land there." This is easy to understand.
The Problem with the Simple Way: The gymnast gets bored. They start doing the same move over and over because it's the easiest way to get a "good" score. They stop stretching, and their muscles (the features) all look the same. They can't do a split anymore because they've forgotten how to stretch.
The RLDP Fix (The Orthogonality Regularizer): The coach adds a simple rule: "You must stay stretched out."
- They force the gymnast to keep their limbs in different directions (orthogonality).
- They say, "Don't let your 'jump' look like your 'run'. Keep them distinct."
- This simple rule prevents the gymnast from collapsing into a single, useless pose. It forces the robot to keep its internal map diverse and detailed.

Why This Matters

The paper shows that this simple "stretching rule" (orthogonality regularization) works better than the complex, heavy-duty methods.

It's Robust: Even if the robot only has a small, messy dataset (like a low-coverage scenario where it hasn't seen many moves), RLDP still works. The complex methods often fail here because they get confused by the missing data.
It's Simple: You don't need a supercomputer to run the complex math. You just need the prediction rule and the "stretching" rule.
It's General: The robot learns a map that is so good, it can handle new tasks (like a new reward function) instantly, just like a human who knows how to walk can instantly adapt to walking on sand, ice, or stairs without a new lesson.

The Takeaway

The authors discovered that you don't need a complicated, over-engineered system to teach a robot general skills. You just need a simple way to predict the future, combined with a simple rule to ensure the robot doesn't get "lazy" and forget the differences between things.

In short: Instead of trying to predict every possible future outcome (which is hard and error-prone), just teach the robot to predict what happens next, but force it to keep its mental map diverse and distinct. It's the difference between a confused student cramming for a test and a disciplined athlete who knows exactly how to stretch for any sport.

1. Problem Statement

Behavioral Foundation Models (BFMs) aim to train agents capable of solving any task in an environment (zero-shot reinforcement learning) after a single pre-training phase on reward-free data. The core mechanism involves learning a state representation $\phi(s)$ such that any reward function $r(s)$ can be approximated as a linear projection $r(s) \approx \phi(s)^\top z$ , where $z$ is a task embedding.

The Challenge:

Representation Dependency: The performance of BFMs relies heavily on the expressivity of the state representation $\phi$ . If $\phi$ cannot span the space of relevant reward functions, the agent cannot learn optimal policies for unseen tasks.
Complexity of Current Methods: State-of-the-art methods (e.g., Forward-Backward, PSM, HILP) learn representations by predicting Successor Measures (the discounted state visitation distribution under a policy). This requires iterative Bellman backups, which are prone to:
- Bias and Instability: Due to bootstrapping in function approximation.
- Out-of-Distribution (OOD) Issues: In low-coverage datasets, policies may select actions not present in the data, leading to incorrect generalization and degenerate representations.
- Complexity: These methods require defining specific classes of policies or rewards a priori.
The Gap: It is unclear if these complex objectives are necessary. Can a simpler objective, such as Latent Dynamics Prediction (predicting the next latent state given current state and action), suffice?

Key Observation: The authors found that while latent dynamics prediction is policy-independent and avoids Bellman backups, using it naively leads to a "mild feature collapse." Over training, different states become increasingly similar in the latent space (high cosine similarity), reducing the span of the representation and harming zero-shot performance.

2. Methodology: Regularized Latent Dynamics Prediction (RLDP)

The authors propose RLDP, a simple yet effective framework that combines latent dynamics prediction with a diversity-preserving regularization.

A. Core Objective

The method learns a state encoder $\phi: S \to \mathbb{R}^d$ and a latent dynamics model $g: \mathbb{R}^d \times A \to \mathbb{R}^d$ with weights $w$ . The goal is to predict the next latent state:
$\phi(s') \approx g(\phi(s), a)^\top w$

The primary loss is the Latent Dynamics Prediction Loss ( $L_d$ ):
$L_d(\phi, g, w) = \mathbb{E}_{\tau} \left[ \sum_{t=1}^H \| h_t - \bar{\phi}(s_t) \|^2 \right]$
Where $h_t$ is the unrolled latent state and $\bar{\phi}$ is a slowly moving target encoder.

B. Mitigating Feature Collapse

To prevent the representation from collapsing (where distinct states map to similar vectors), the authors introduce an Orthogonality Regularization ( $L_r$ ).

Hyperspherical Projection: Both the encoded states $\phi(s)$ and the predicted next states are projected onto a hypersphere $S^{d-1}$ (normalized to unit length).
Regularization Loss: The loss minimizes the cosine similarity between any two distinct state representations:
$L_r(\phi) = \mathbb{E}_{s, s' \sim \rho} [\phi(s)^\top \phi(s')]$
Note: Since vectors are normalized, the dot product equals the cosine similarity. Minimizing this encourages orthogonality/diversity.

C. Final Objective

The total loss is a weighted sum:
$L_{RLDP} = L_d + \lambda L_r$
where $\lambda$ is a small regularization coefficient (e.g., 0.01 or 1.0).

D. Zero-Shot RL Pipeline

Pre-training: Train $\phi$ using $L_{RLDP}$ on offline reward-free data.
Successor Measure Estimation: Freeze $\phi$ . Train the successor feature model $\psi$ and policy $\pi_z$ using standard BFM objectives (alternating between estimating successor measures and policy improvement).
Inference: For a new reward $r_{test}$ , solve for $z$ via linear regression ( $r_{test} \approx \phi^\top z$ ) and extract the optimal policy $\pi_z$ .

3. Key Contributions

Identification of Feature Collapse: The paper identifies that naive latent dynamics prediction leads to a reduction in feature diversity (increased cosine similarity), which degrades the span of learnable reward functions.
RLDP Algorithm: Proposes a simple, scalable method adding orthogonality regularization to latent dynamics prediction to maintain feature diversity.
Theoretical Justification: Provides a lemma showing that minimizing $L_{RLDP}$ upper-bounds the prediction error of successor measures for Lipschitz policies, theoretically linking dynamics prediction to BFM requirements.
Robustness in Low-Coverage: Demonstrates that RLDP outperforms complex Bellman-backup-based methods in low-coverage datasets where OOD actions cause other methods to fail.

4. Experimental Results

The authors evaluated RLDP across multiple domains: DeepMind Control Suite (ExoRL), SMPL Humanoid, and D4RL (low-coverage).

Offline Zero-Shot RL (DMC):
- RLDP matches or surpasses state-of-the-art baselines (FB, PSM, HILP) across tasks like Walker, Cheetah, Quadruped, and Pointmass.
- It achieves comparable performance to PSM and FB but with a significantly simpler training objective (no Bellman backups).
Online Zero-Shot RL (SMPL Humanoid):
- In the complex 3D Humanoid environment (358-dim obs, 69-dim action), RLDP representations enable competitive multi-task learning.
- RLDP outperforms baselines on specific tasks (e.g., raise arms, lie on ground) and remains competitive overall, even beating oracle performance on some tasks.
Low-Coverage Scenarios (D4RL):
- Critical Finding: In low-coverage datasets (e.g., halfcheetah-medium, hopper-medium), methods relying on Bellman backups (FB, PSM) struggle due to OOD action queries.
- RLDP significantly outperforms these baselines in 5 out of 6 tasks, proving its robustness when data coverage is sparse.
Ablation Studies:
- Orthogonality: Removing regularization ( $\lambda=0$ ) causes performance drops and increased feature similarity.
- Spherical Normalization: Projecting states to a hypersphere is crucial for stability and performance.
- Loss Components: Combining prediction loss and orthogonality yields the best results; adding Bellman backups to this combination often degrades performance in low-coverage settings.

5. Significance and Conclusion

Simplicity vs. Complexity: The paper challenges the necessity of complex successor measure estimation (which requires Bellman backups and policy classes) for training BFMs. A simple, policy-independent dynamics prediction with diversity regularization is sufficient.
Practicality: RLDP is more amenable to optimization and handles high-dimensional observation spaces better than reconstruction-based or Bellman-based methods.
Robustness: It offers a reliable solution for real-world robotics and control scenarios where data coverage is often limited (low-coverage), a regime where current SOTA methods fail.
Future Impact: By providing a stable, simple baseline, RLDP paves the way for more accessible and effective unsupervised learning approaches for generalist agents.

In summary, RLDP demonstrates that maintaining feature diversity through simple orthogonality constraints is the key to unlocking the potential of latent dynamics for zero-shot reinforcement learning, outperforming complex, Bellman-dependent baselines especially in data-scarce environments.