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

Pranaya Jajoo; Harshit Sikchi; Siddhant Agarwal; Amy Zhang; Scott Niekum; Martha White

cs.AIMar 16, 2026

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 ·

Abstract

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 $\lambda=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 $\pm$ 58.03	918.29 $\pm$ 28.83	899.54 $\pm$ 30.73	890.40 $\pm$ 27.33
	Run	75.39 $\pm$ 20.97	381.31 $\pm$ 17.32	450.57 $\pm$ 28.95	334.26 $\pm$ 49.69
	Walk	193.84 $\pm$ 112.98	779.29 $\pm$ 63.60	875.61 $\pm$ 33.44	779.77 $\pm$ 137.16
	Flip	132.02 $\pm$ 67.85	977.08 $\pm$ 2.76	621.36 $\pm$ 75.62	492.94 $\pm$ 22.79
Cheetah	Run	31.82 $\pm$ 36.88	129.39 $\pm$ 37.63	181.85 $\pm$ 54.17	157.12 $\pm$ 29.92
	Run Backward	60.08 $\pm$ 12.82	142.41 $\pm$ 36.77	158.64 $\pm$ 18.56	170.52 $\pm$ 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 ( $\Delta$ = 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 $\Delta=R_{\text{{RLDP}}}-R_{\text{baseline}}$ for a single $\{\text{task},\text{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 $\lambda=0.01$ gives a statistically significant improvement over coefficient $\lambda=0.0$ .

	Task	RLDP	RLDP w/o SN
Quadruped	Stand	794.94 $\pm$ 43.25	661.73 $\pm$ 95.75
	Run	457.41 $\pm$ 74.70	378.97 $\pm$ 148.47
	Walk	465.40 $\pm$ 185.29	519.39 $\pm$ 251.11
	Jump	733.32 $\pm$ 55.30	495.98 $\pm$ 133.81
	Average(*)	612.77 $\pm$ 83.17	514.02 $\pm$ 86.94
Pointmass	Top Left	890.41 $\pm$ 60.79	892.13 $\pm$ 41.74

Table 2: Study of encoder architecture (subset). Table shows mean

\pm

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	$10^{6}$
	hopper-medium-expert	$2\times 10^{6}$
	halfcheetah-medium	$10^{6}$
	halfcheetah-medium-expert	$2\times 10^{6}$
	walker2d-medium	$10^{6}$
	walker2d-medium-expert	$2\times 10^{6}$

Appendix

Table 4: Gym-MuJoCo tasks from D4RL.

Figure 6: Architecture of latent next-state prediction network in RLDP .

Panel of Figure 7 — 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 $\pm$ 27.33	860.74 $\pm$ 62.47	810.79 $\pm$ 100.90	881.69 $\pm$ 6.92
	Run	334.26 $\pm$ 49.69	324.04 $\pm$ 6.73	290.78 $\pm$ 26.52	276.30 $\pm$ 47.21
	Walk	779.77 $\pm$ 137.16	728.29 $\pm$ 43.09	715.83 $\pm$ 92.43	583.60 $\pm$ 28.26
	Flip	492.94 $\pm$ 22.79	501.59 $\pm$ 45.04	477.95 $\pm$ 37.88	447.73 $\pm$ 33.59
	Average(*)	624.34	603.66	573.84	547.33
Cheetah	Run	157.12 $\pm$ 29.92	84.99 $\pm$ 67.31	115.25 $\pm$ 14.13	118.67 $\pm$ 32.67

Appendix

Table 5: Study of encoder architecture. Cells show mean

\pm

std over 4 seeds; boldface indicates the highest mean per task.

Panel of Figure 9 — 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 $\pm$ 151.40	902.63 $\pm$ 38.94	607.07 $\pm$ 165.28	872.61 $\pm$ 38.81	877.69 $\pm$ 45.03
	Run	63.65 $\pm$ 31.02	392.76 $\pm$ 31.29	107.84 $\pm$ 34.24	351.50 $\pm$ 19.46	324.85 $\pm$ 54.57
	Walk	190.53 $\pm$ 168.45	877.10 $\pm$ 81.05	399.67 $\pm$ 39.31	891.44 $\pm$ 46.81	790.94 $\pm$ 67.55
	Flip	48.73 $\pm$ 17.66	206.22 $\pm$ 162.27	277.95 $\pm$ 59.63	640.75 $\pm$ 31.88	491.64 $\pm$ 37.30
Cheetah	Run	96.32 $\pm$ 35.69	257.59 $\pm$ 58.51	68.22 $\pm$ 47.08	244.38 $\pm$ 80.00	236.31 $\pm$ 20.75
	Run Backward	106.38 $\pm$ 29.40	307.07 $\pm$ 14.91	37.99 $\pm$ 25.16	296.44 $\pm$ 20.14	322.08 $\pm$ 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.,

\mu_{i}\geq\mu^{*}-\sigma^{*}

) aggregated over 4 seeds.

	Task	RLDP	Learning USFAs on RLDP representations
Walker	Stand	890.40 $\pm$ 27.33	854.84 $\pm$ 64.50
	Run	334.26 $\pm$ 49.69	304.89 $\pm$ 63.15
	Walk	779.77 $\pm$ 137.16	665.98 $\pm$ 117.64
	Flip	492.94 $\pm$ 22.79	497.77 $\pm$ 70.80
Cheetah	Run	157.12 $\pm$ 29.92	139.00 $\pm$ 6.72
	Run Backward	170.52 $\pm$ 15.30	172.52 $\pm$ 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 $\pm$ 43.60	26.06 $\pm$ 39.43	38.38 $\pm$ 14.73	52.31 $\pm$ 23.15	86.48 $\pm$ 45.87
crawl-0.4-0-u	101.76 $\pm$ 15.90	8.38 $\pm$ 9.01	4.52 $\pm$ 7.47	18.59 $\pm$ 20.42	25.00 $\pm$ 27.32
crawl-0.4-2-d	19.00 $\pm$ 4.00	3.05 $\pm$ 3.65	6.52 $\pm$ 1.27	8.97 $\pm$ 5.62	11.21 $\pm$ 4.72
crawl-0.4-2-u	15.02 $\pm$ 6.03	0.79 $\pm$ 1.07	0.64 $\pm$ 0.82	2.95 $\pm$ 1.37	2.76 $\pm$ 3.73
crawl-0.5-0-d	131.13 $\pm$ 64.97	43.27 $\pm$ 34.66	46.17 $\pm$ 13.56	52.41 $\pm$ 27.82	55.82 $\pm$ 18.40
crawl-0.5-0-u	101.92 $\pm$ 16.39	4.04 $\pm$ 5.87	4.18 $\pm$ 4.83	21.14 $\pm$ 24.26	20.22 $\pm$ 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 $\pm$ 11.37	49.92 $\pm$ 26.89	68.47 $\pm$ 8.20	86.03 $\pm$ 8.36
halfcheetah-medium-v2	39.27 $\pm$ 8.71	42.64 $\pm$ 0.64	43.85 $\pm$ 1.49	49.08 $\pm$ 1.93
hopper-medium-expert-v2	54.64 $\pm$ 19.47	14.59 $\pm$ 26.35	68.18 $\pm$ 18.19	77.21 $\pm$ 16.77
hopper-medium-v2	43.75 $\pm$ 6.65	33.49 $\pm$ 26.71	52.19 $\pm$ 4.02	44.93 $\pm$ 13.08
walker2d-medium-expert-v2	60.17 $\pm$ 28.94	79.32 $\pm$ 41.16	93.72 $\pm$ 21.12	103.87 $\pm$ 3.31
walker2d-medium-v2	43.72 $\pm$ 24.95	55.70 $\pm$ 22.01	56.34 $\pm$ 15.08	83.83 $\pm$ 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

\pm

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..

Panel of Figure 11 — Figure 11: Visualization of successor measures $M^{\pi_{z}}(s_{0},a_{0},s^{+})$ for randomly sampled $z$ (a) and (b); and goal-conditioned $z$ (c) and (d).

Environment	Returns (mean $\pm$ std)
halfcheetah-medium-expert
$\mathcal{L}_{\text{Ortho}}$	$51.36\pm 6.93$
$\mathcal{L}_{\text{Bellman}}$	$64.42\pm 7.59$
$\mathcal{L}_{\text{Prediction}}$	$49.00\pm 4.34$
$\mathcal{L}_{\text{Bellman}}+\mathcal{L}_{\text{Prediction}}+\mathcal{L}_{\text{Ortho}}$	$89.45\pm 3.81$
$\mathcal{L}_{\text{Bellman}}+\mathcal{L}_{\text{Ortho}}$ (FB)	$55.46\pm 7.75$

Appendix

Table 10: Results in D4RL expert environments across different loss combinations aggregated over 4 seeds.

Behavioral Foundation Models

See all 15 papers

Task Inference Beyond Least Squares in Behavioral Foundation Models
Oct 4, 2026Kuan-Hsun Tu, Chien-Sheng Chiang, Hsin-Wei Chen +2Reinforcement Learning Multi-Task RL
National Taiwan University
BFM-Zero: A Promptable Behavioral Foundation Model for Humanoid Control Using Unsupervised Reinforcement Learning
Nov 6, 2025Yitang Li, Zhengyi Luo, Tonghe Zhang +10Humanoid Robot Control Behavioral Foundation Models
1Carnegie Mellon University · 2Meta
Can We Really Learn One Representation to Optimize All Rewards?
Feb 11, 2026Chongyi Zheng, Royina Karegoudra Jayanth, Benjamin EysenbachZero-Shot RL Representation Learning
Princeton University

Oct 4, 2026J/K move · Enter open · S save

Task Inference Beyond Least Squares in Behavioral Foundation Models

Kuan-Hsun Tu, Chien-Sheng Chiang, Hsin-Wei Chen, Ping-Chun Hsieh+1

National Taiwan University

Reinforcement Learning Multi-Task RL Behavioral Foundation Models Zero-Shot RL

Open paper AskPDF arXiv

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

Abstract

Behavioral Foundation Models

Explore similar work

Task Inference Beyond Least Squares in Behavioral Foundation Models

BFM-Zero: A Promptable Behavioral Foundation Model for Humanoid Control Using Unsupervised Reinforcement Learning

Can We Really Learn One Representation to Optimize All Rewards?