Regularized Latent Dynamics Prediction is a Strong Baseline For Behavioral Foundation Models
Authors: Pranaya Jajoo, Harshit Sikchi, Siddhant Agarwal, Amy Zhang, Scott Niekum, Martha White
Organizations: Department of Computing Science, University of Alberta, Canada · Alberta Machine Intelligence Institute (Amii) · The University of Texas at Austin · University of Massachusetts Amherst · Canada CIFAR AI Chair
Behavioral Foundation Models (BFMs) produce agents with the capability to adapt to any unknown reward or task. These methods, however, are only able to produce near-optimal policies for the reward functions that are in the span of some pre-existing state features, making the choice of state features crucial to the expressivity of the BFM. As a result, BFMs are trained using a variety of complex objectives and require sufficient dataset coverage, to train task-useful spanning features. In this work, we examine the question: are these complex representation learning objectives necessary for zero-shot RL? Specifically, we revisit the objective of self-supervised next-state prediction in latent space for state feature learning, but observe that such an objective alone is prone to increasing state-feature similarity, and subsequently reducing span. We propose an approach, Regularized Latent Dynamics Prediction (RLDP), that adds a simple orthogonality regularization to maintain feature diversity and can match or surpass state-of-the-art complex representation learning methods for zero-shot RL. Furthermore, we empirically show that prior approaches perform poorly in low-coverage scenarios where RLDP still succeeds.
Figures & tables
Figure 1: Average Cosine similarity between state-representations sampled uniformly from the training dataset. Feature similarity increases over the course of training; once adding our orthogonality regularizer (with λ=1 ), we obtain more diverse representations. Shaded region shows standard deviation over 4 seeds.
Figure 2: RLDP combines latent next state prediction + regularization for diversity (an orthogonality regularizer) to learn representations for BFMs.
Task
Random Features
FB
PSM
RLDP
Walker
Stand
392.40 ± 58.03
918.29 ± 28.83
899.54 ± 30.73
890.40 ± 27.33
Run
75.39 ± 20.97
381.31 ± 17.32
450.57 ± 28.95
334.26 ± 49.69
Walk
193.84 ± 112.98
779.29 ± 63.60
875.61 ± 33.44
779.77 ± 137.16
Flip
132.02 ± 67.85
977.08 ± 2.76
621.36 ± 75.62
492.94 ± 22.79
Cheetah
Run
31.82 ± 36.88
129.39 ± 37.63
181.85 ± 54.17
157.12 ± 29.92
Run Backward
60.08 ± 12.82
142.41 ± 36.77
158.64 ± 18.56
170.52 ± 15.30
Table 1: Comparison (over 4 seeds) of zero‐shot RL performance between using an untrained initialized encoder, FB, PSM, and RLDP with representation size d=512 . Bold indicates the best mean and any method whose mean plus one standard deviation overlaps with the best mean.
Figure 3: Pair-wise comparison of RLDP against prior offline representation learning methods using per-task oracle normalized performance differences ( Δ = RLDP – Baseline) in SMPL Humanoid environment. The gray diamond represents the IQM (Interquartile Mean).
Figure 4: Pair-wise comparison of RLDP against baseline representation learning methods in low-coverage D4RL dataset. Each point represents Δ=RRLDP−Rbaseline for a single {task,seed} pair. The gray diamond represents the IQM (Interquartile Mean).
Figure 5: Evaluating the impact of Orthogonality Regularization: We ran one-sided Mann–Whitney U tests on the per-seed returns over 4 seeds to compare different values of the orthogonality regularization, and we observe that adding small orthogonality regularization coefficient λ=0.01 gives a statistically significant improvement over coefficient λ=0.0 .
Task
RLDP
RLDP w/o SN
Quadruped
Stand
794.94 ± 43.25
661.73 ± 95.75
Run
457.41 ± 74.70
378.97 ± 148.47
Walk
465.40 ± 185.29
519.39 ± 251.11
Jump
733.32 ± 55.30
495.98 ± 133.81
Average(*)
612.77 ± 83.17
514.02 ± 86.94
Pointmass
Top Left
890.41 ± 60.79
892.13 ± 41.74
Table 2: Study of encoder architecture (subset). Table shows mean ± std; RLDP significantly outperforms RLDP w/o SN on Pointmass, Quadruped, and pooled. SN: Spherical Normalization on g .
Appendix figures & tables14 assets
Supplementary material from the paper’s appendix.
Appendix
Domain
Description
Type
Observation/Action Dimension
Tasks
Reward
Walker
two-legged robot
Locomotion
24/6
stand walk run flip
Dense
Quadruped
four-legged robot
Locomotion
78/12
jump walk run stand
Dense
Cheetah
planar, 2D robot
Locomotion
17/6
walk run walk backward run backward
Sparse
Pointmass
navigation in 2D plane
Goal-reaching
4/2
reach top left reach top right reach bottom right reach bottom left
Sparse
Appendix
Table 3: ExoRL dataset summary . Domain is the environment name in the ExoRL benchmark. Description is a natural language description of the agent embodiment and environment. Type refers to the broad task category. Observation/Action Dimension refers to the size of observation and action vectors from the environment. Tasks refers to the suite of evaluation tasks provided in the ExoRL benchmark. Reward refers to the density of non-zero reward signals from the environment.
Domain
Task Name
# Samples
Gym-MuJoCo
hopper-medium
106
hopper-medium-expert
2×106
halfcheetah-medium
106
halfcheetah-medium-expert
2×106
walker2d-medium
106
walker2d-medium-expert
2×106
Appendix
Table 4: Gym-MuJoCo tasks from D4RL.
Figure 6: Architecture of latent next-state prediction network in RLDP .
Figure 7: Evaluating the impact of Orthogonality Regularization on representations learned across four environments: Cheetah (top left), Pointmass (top right), Quadruped (bottom left), and Walker (bottom right).
Figure 8: Evaluating the impact of Encoding Horizon
Task
RLDP
RLDP w/o SN
RLDP w/o A
RLDP w/o SN & A
Walker
Stand
890.40 ± 27.33
860.74 ± 62.47
810.79 ± 100.90
881.69 ± 6.92
Run
334.26 ± 49.69
324.04 ± 6.73
290.78 ± 26.52
276.30 ± 47.21
Walk
779.77 ± 137.16
728.29 ± 43.09
715.83 ± 92.43
583.60 ± 28.26
Flip
492.94 ± 22.79
501.59 ± 45.04
477.95 ± 37.88
447.73 ± 33.59
Average(*)
624.34
603.66
573.84
547.33
Cheetah
Run
157.12 ± 29.92
84.99 ± 67.31
115.25 ± 14.13
118.67 ± 32.67
Appendix
Table 5: Study of encoder architecture. Cells show mean ± std over 4 seeds; boldface indicates the highest mean per task.
Figure 9: t-SNE visualizations of state features in Pointmass. Each panel shows the 2D projection of 10,000 uniformly sampled states.
Task
Laplace
FB
HILP
PSM
RLDP
Walker
Stand
243.70 ± 151.40
902.63 ± 38.94
607.07 ± 165.28
872.61 ± 38.81
877.69 ± 45.03
Run
63.65 ± 31.02
392.76 ± 31.29
107.84 ± 34.24
351.50 ± 19.46
324.85 ± 54.57
Walk
190.53 ± 168.45
877.10 ± 81.05
399.67 ± 39.31
891.44 ± 46.81
790.94 ± 67.55
Flip
48.73 ± 17.66
206.22 ± 162.27
277.95 ± 59.63
640.75 ± 31.88
491.64 ± 37.30
Cheetah
Run
96.32 ± 35.69
257.59 ± 58.51
68.22 ± 47.08
244.38 ± 80.00
236.31 ± 20.75
Run Backward
106.38 ± 29.40
307.07 ± 14.91
37.99 ± 25.16
296.44 ± 20.14
322.08 ± 39.28
Appendix
Table 6: Comparison of zero‐shot offline RL performance between different methods. Entries in bold are within one standard deviation of the per-task best mean (i.e., μi≥μ∗−σ∗ ) aggregated over 4 seeds.
Task
RLDP
Learning USFAs on RLDP representations
Walker
Stand
890.40 ± 27.33
854.84 ± 64.50
Run
334.26 ± 49.69
304.89 ± 63.15
Walk
779.77 ± 137.16
665.98 ± 117.64
Flip
492.94 ± 22.79
497.77 ± 70.80
Cheetah
Run
157.12 ± 29.92
139.00 ± 6.72
Run Backward
170.52 ± 15.30
172.52 ± 25.72
Appendix
Table 7: Comparison (over 4 seeds) of zero‐shot RL performance when equation 10 is used to train the critic and when Universal Successor Features are trained on top of the state features. For fair comparison, we set RLDP representation dimension d=512 for both methods.
metric
FB-CPR (Oracle)
FB
PSM
HILP
RLDP
crawl-0.4-0-d
191.75 ± 43.60
26.06 ± 39.43
38.38 ± 14.73
52.31 ± 23.15
86.48 ± 45.87
crawl-0.4-0-u
101.76 ± 15.90
8.38 ± 9.01
4.52 ± 7.47
18.59 ± 20.42
25.00 ± 27.32
crawl-0.4-2-d
19.00 ± 4.00
3.05 ± 3.65
6.52 ± 1.27
8.97 ± 5.62
11.21 ± 4.72
crawl-0.4-2-u
15.02 ± 6.03
0.79 ± 1.07
0.64 ± 0.82
2.95 ± 1.37
2.76 ± 3.73
crawl-0.5-0-d
131.13 ± 64.97
43.27 ± 34.66
46.17 ± 13.56
52.41 ± 27.82
55.82 ± 18.40
crawl-0.5-0-u
101.92 ± 16.39
4.04 ± 5.87
4.18 ± 4.83
21.14 ± 24.26
20.22 ± 30.91
Appendix
Table 8: Comparing (over 4 seeds) FB, PSM, HILP, RLDP performance on SMPL Humanoid. FB-CPR (online oracle baseline) results are from Tirinzoni et al. (2025) . Bold indicates the best mean across methods.
Figure 10: Evaluating offline representation learning methods using an online oracle policy in high dimensional 3D humanoid. Solid lines shows mean performance across tasks for each method.
Task
FB
PSM
HILP
RLDP
halfcheetah-medium-expert-v2
52.17 ± 11.37
49.92 ± 26.89
68.47 ± 8.20
86.03 ± 8.36
halfcheetah-medium-v2
39.27 ± 8.71
42.64 ± 0.64
43.85 ± 1.49
49.08 ± 1.93
hopper-medium-expert-v2
54.64 ± 19.47
14.59 ± 26.35
68.18 ± 18.19
77.21 ± 16.77
hopper-medium-v2
43.75 ± 6.65
33.49 ± 26.71
52.19 ± 4.02
44.93 ± 13.08
walker2d-medium-expert-v2
60.17 ± 28.94
79.32 ± 41.16
93.72 ± 21.12
103.87 ± 3.31
walker2d-medium-v2
43.72 ± 24.95
55.70 ± 22.01
56.34 ± 15.08
83.83 ± 2.66
Appendix
Table 9: Normalized returns comparing FB, PSM, HILP, and RLDP in low-coverage setting. RLDP shows significant gains over approaches that rely on explicit Bellman backups for representation learning. Table shows mean ± std over 10 seeds; Boldface indicates the highest mean return per environment. Statistical comparisons use per-seed returns with Welch’s t-test and Holm correction; cases where the bolded method is not significantly better than the runner-up are discussed in the text..
Figure 11: Visualization of successor measures Mπz(s0,a0,s+) for randomly sampled z (a) and (b); and goal-conditioned z (c) and (d).
Environment
Returns (mean ± std)
halfcheetah-medium-expert
LOrtho
51.36±6.93
LBellman
64.42±7.59
LPrediction
49.00±4.34
LBellman+LPrediction+LOrtho
89.45±3.81
LBellman+LOrtho (FB)
55.46±7.75
Appendix
Table 10: Results in D4RL expert environments across different loss combinations aggregated over 4 seeds.
Behavioral Foundation Models (BFMs) aim to solve a wide range of downstream tasks without test-time policy learning by inferring a task vector from the reward function. While efficient, the retrieved policies are often suboptimal because of how this task vector is inferred, typically with ordinary least squares (OLS). OLS minimizes reward reconstruction error but leaves the ordering of rewards unconstrained, which can bias the successor measure of the retrieved zero-shot policy away from that of the optimal policy. In this work, we propose BLS, an efficient test-time inference method that balances minimizing reward reconstruction error with reducing successor-measure mismatch. Theoretically, we provide a suboptimality gap upper bound characterized by both successor-measure and reward-function residuals. Empirically, we evaluate BLS on top of state-of-the-art BFMs across benchmarks for locomotion, manipulation, and humanoid control. BLS outperforms existing task inference baselines with negligible computational overhead. Project page: https://embodiedai-ntu.github.io/BLS
Building Behavioral Foundation Models (BFMs) for humanoid robots has the potential to unify diverse control tasks under a single, promptable generalist policy. However, existing approaches are either exclusively deployed on simulated humanoid characters, or specialized to specific tasks such as tracking. We propose BFM-Zero, a framework that learns an effective shared latent representation that embeds motions, goals, and rewards into a common space, enabling a single policy to be prompted for multiple downstream tasks without retraining. This well-structured latent space in BFM-Zero enables versatile and robust whole-body skills on a Unitree G1 humanoid in the real world, via diverse inference methods, including zero-shot motion tracking, goal reaching, and reward optimization, and few-shot optimization-based adaptation. Unlike prior on-policy reinforcement learning (RL) frameworks, BFM-Zero builds upon recent advancements in unsupervised RL and Forward-Backward (FB) models, which offer an objective-centric, explainable, and smooth latent representation of whole-body motions. We further extend BFM-Zero with critical reward shaping, domain randomization, and history-dependent asymmetric learning to bridge the sim-to-real gap. Those key design choices are quantitatively ablated in simulation. A first-of-its-kind model, BFM-Zero establishes a step toward scalable, promptable behavioral foundation models for whole-body humanoid control.
As unsupervised pretraining becomes increasingly ubiquitous in reinforcement learning, a more thorough theoretical understanding of these methods becomes of equal importance to their empirical success. We focus on the setting of unsupervised learning via interaction, where the forward-backward (FB) representation learning serves as a prototypical and popular example. In this paper, we shed light on FB by formally contextualizing the method within a broader class of recent methods that use regression to obtain a low-rank approximation of a successor measure ratio. Our analysis clarifies when FB representations can exist and how the low-rank approximation converges in practice. Building upon the theory, we propose a variant of FB that is both more amenable to theoretical understanding and simpler to optimize in practice. Experiments in didactic settings, as well as in 10 state-based and image-based continuous control domains, demonstrate that our method converges to desired representations with 105× smaller errors than FB, achieving +24% improved zero-shot performance on average. We also demonstrate that zero-shot policies inferred by our algorithm provide an efficient initialization if the user prefers further fine-tuning on downstream tasks. Our project website is available at https://chongyi-zheng.github.io/onestep-fb.
Chongyi Zheng, Royina Karegoudra Jayanth, Benjamin Eysenbach