Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization
Authors: Frank Röder
Jan Benad
Manfred Eppe
Pradeep Kr. Banerjee
Venue: NeurIPS 2025 poster ()
Keywords: world models
contextual reinforcement learning
representation learning
counterfactual consistency
Primary area: reinforcement_learning
TL;DR: We show that dynamics-aligned representations improve zero-shot generalization for contextual world models.

Abstract: Real-world reinforcement learning demands adaptation to unseen environmental conditions without costly retraining. Contextual Markov Decision Processes (cMDP) model this challenge, but existing methods often require explicit context variables (e.g., friction, gravity), limiting their use when contexts are latent or hard to measure. We introduce Dynamics-Aligned Latent Imagination (DALI), a framework integrated within the Dreamer architecture that infers latent context representations from agent-environment interactions. By training a self-supervised encoder to predict forward dynamics, DALI generates actionable representations conditioning the world model and policy, bridging perception and control. We theoretically prove this encoder is essential for efficient context inference and robust generalization. DALI’s latent space enables counterfactual consistency: Perturbing a gravity-encoding dimension alters imagined rollouts in physically plausible ways. On challenging cMDP benchmarks, DALI achieves significant gains over context-unaware baselines, often surpassing context-aware baselines in extrapolation tasks, enabling zero-shot generalization to unseen contextual variations.

Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization
Frank Röder* Jan Benad
Pradeep Kr. Banerjee" Institute for Data Science Foundations, Blohmstraße 15, 21079 Hamburg, Germany
Manfred Eppe
Abstract
Real-world reinforcement learning demands adaptation to unseen environmen- tal conditions without costly retraining. Contextual Markov Decision Processes (cMDP) model this challenge, but existing methods often require explicit context variables (e.g., friction, gravity), limiting their use when contexts are latent or hard to measure. We introduce Dynamics-Aligned Latent Imagination (DALI), framework integrated within the Dreamer architecture that infers latent context representations from egent-cnvironment interactions. By training self-supervised encoder to predict forward dynamics, DALI generates actionable representations conditioning the world model and policy, bridging perception and control. We the- oretically prove this encoder is essential for efficient context inference and robust generalization. DALI's latent space enables counterfactual consistency: Perturbing a gravity-encoding dimension alters imagined rollouts in physically plausible ways. On challenging cMDP benchmarks, DALI achieves significant gains context- unaware baselines, often surpassing context-aware baselines extrapolation tasks, enabling zero-shot generalization to unseen contextual variations.
Introduction
The ability generalize across diverse environmental conditions is fundamental challenge in reinforcement learning (RL). Contextual Markov Decision Processes (cMDPs) formalize this chal- lenge by modeling task variations through latent parameters, such friction, gravity, object mass that govern dynamics [Hallak et al., (215). However, real-world agents rarely have direct access these parameters. Consider robotic control task where legged agent must adapt to different surfaces, such as smooth tiles, rough gravel, or slippery ice. If the agent were explicitly provided with surface friction coefficients, adaptation would be straightforward. In practice, though, such ground- truth annotations are often unavailable or prohibitively expensive to obtain. Instead, the agent must infer these variations implicitly by observing how its actions influence the environment, using proprioceptive feedback and gait dynamics.
Existing approaches contextual RL often rely explicit context conditioning, where models trained with domain sppeiifc instrumentation or carefully crafted architectures [Seyed Ghasemipour al., 2019 [Bll et al. 2021 Eghbal-zzdeh al.] 202 Mu al.) 2022 Beukman et al. 2023 Benjaminns et al. 2023] Prasanna et al. 2024]. While effective in controlled settings, such approaches scale poorly and break down in unstructured environments where contexts are latent difficult to measure. The challenge to develop RI. agents that infer and adapt to hidden contexts a self-supervised manner, enabling robust generalization without direct supervision.
To address this challenge, we propose Dynamics- Aligned Latent Imagination (DALD, framework in- tegrated within the DreamerV3 architecture (Hafner et al. 2025 that enables zero-shot generalization by inferring latent contexts directly from interaction histories. Unlike Dreamer V3, which struggles
Equal contribution
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

to generalize across diverse latent contexts due to its limited ability to retain critical environmental information [Prasanna et al., 2024]. DALI overcomes these limitations through a self-supervised context encoder. This encoder learns to predict forward dynamics, capturing the relationship between actions, states, and their consequences. For example, by analyzing how applied forces influence motion trajectories, the encoder distills contextual factors (e.g., gravitational pull or surface friction) into compact, actionable embeddings. These embeddings condition the world model and policy, enabling the agent to reason about hidden dynamics (e.g., how increased inertia affects motion stability) and adapt its control strategies accordingly. Our theoretical analysis demonstrates that DALI's context encoder enables robust zero-shot generalization by efficiently capturing latent vironmental variations, overcoming information bottlenecks and sample inefficiencies in learning diverse contexts [Prasanna al. 2024 Hafner et al., 2025].
By inferring latent contexts from interactions, DALJ enhances DreamerV3's learning framework, achieving robust zero-shot generalization in partially observable environments. Our work makes the following key contributions:
Theoretical foundations: We establish that a dedicated context encoder is essential for robust generalization, proving that it infers latent contexts from short interaction windows with near- optimal sample complexity and mitigates information bottlenecks in partially observable settings.
Strong zero-shot generalization: On challenging cMDP benchmarks, DALI achieves up to +96.4% gains over context-unaware baselines, often surpassing ground-truth context-aware base- lines in extrapolation tasks.
Physically consistent counterfactuals: We show that the learned latent space exhibits counterfac- tual consistency: perturbing a dimension encoding gravity, for instance, results in imagined rollouts where objects fall faster or slower, faithfully mirroring real-world physics.
2
Related Work
Contextual RL for Zero-Shot Generalization. Contextual RL has been studied in various forms, from cMDPs to domain randomization and meta-RL (Hallak et al. 2013 Modi et al. 20188] A recent survey Kirk et al. 2023] highlights its relevance for zero-shot generalization, emphasizing how clear context sets for traming and evaluation enable systematic analysis. One research direction assumes context is explicitly observed as privileged information and integrates into learning [Chen et al.)
2018
Seyed Ghasemipour et al. 2019 Ball et al. 2021 Eghbal-zadeh et al., [021] Sodhani et al.
2021
Mu et aal.|2022 Benjamins et al.
2023 Prasanna et al. al.| 2024]. In contrast, we follow recent
work that assumes context is latent and must be inferred [Chen et al. 2018, Xu et al.] 2019 Lee
cct al. 2020 Seo et al. 2020 Xian et al. 2021 Sodhani et al., 2022 Melo 2022 Evans et al. 2022
focusing on self-supervised context inference through forward dynamics alignment.
Model-Based RL. Model-based RL improves sample efficiency by learning predictive environment
models for planning and imagination. Approaches like Dreamer [Hafner et al.) 2019 2020 2021
2025) and TD-MPC [Hansen et al.) (2024) achieve strong performance by learning compact latent representations of environment dynamics. Recent work has also explored general general-purpose model- -fee RL. that leverages model-based representations for broader applicability [Fujimoto et al. 2025 While Dreamer has been explored for zero-shot generalization with explicit context conditioning [Prasannal al. 2024), , we build on it with a dynamics-aligned inference mechanism to enhance generalization without context supervision.
Meta-RL. Meta-RL trains agents adapt rapidly to new tasks with minimal experience [Beck
ct al. 2023 typically by learning adaptive policies that infer task-specific information from past
interactions. However, most meta-RL RL methods require fine-tuning on new tasks [Duan et al. 2016
Finn et al. [017 Nagabandi et al.) 2018 Rakelly et al. 2019 Zintgraf et al.) 2019 Mclo 2022 Our
approach. in contrast, aims for zero-shot generalization by leveraging structured latent representations that transfer across environments.
3 Preliminaries
To investigate zero-shot generalization partially observable environments, we adopt a contextual Markov Decision Process (cMDP) framework, following Hallak et al. (2015), Kirk et al. (0033

---

A cMDP is defined as a tuple M = (S, -S..A......R.,,,E,,... where S, A, and O are resp. the state, action, and observation spaces, C C Rd is the context space with contexts E C drawn i.i.d. from distribution Pc. R: S x A C R the reward function, P 8 x A C A(S) the stochastic transition function specifying the probability distribution over next states given the current state, action, and context, & S + A(O) is the observation function specifying the distribution over observations given the state and context, μ C (SS) is the initial state distribution where ussule) /e) specifies the probability of initial state so given context c. In this partially observable setting. the agent does not observe the state St or context e directly but receives observations E O. The
objective is to learn a policy x aaa 01:1, Q1:--1) that maximizes the expected return Σ 1 T0
where T is the episode length.
In our framework, each context e induces a distinct variation in the transition dynamics P, such as changes in physical properties (e.g., gravity, friction), while the reward function R remains fixed across contexts. The context is latent and remains fixed within an episode but varies across episodes according to pe. The agent infers e from interaction histories to adapt its policy. To formalize zero-shot generalization, we define two distributions over contexts: the training distribution Puain (c), from which contexts are sampled during training, and the evaluation distribution Peval (c), representing unseen test contexts. The goal is to learn a policy using samples from Purin (c) that maximizes the expected return for contexts drawn from (eva(() without further adaptation or retraining.
4 Dynamics-Aligned Latent Imagination (DALI)
4.1 Background: DreamerV3
DreamerV3 Hafner et al. 2025 is a model- RL algorithm that enables agents to learn and plan in the latent space of learned world model. It follows the general structure of the Dreamer family of algorithms [Hafner et al.) 2019 2020 2021), incorporating three key components: (i) learning a generative world model to simulate environment dynamics, (ii) optimizing policies entirely within the latent space of the world model using imagined rollouts, and (iii) refining the world model and policy through interaction with the real environment.
World model. The agent interacts with the environment through observations on (e.g., images or sensor data) and actions a,, The world model, structured as Recurrent State-Space Model (RSSM), encodes these observations into a compact latent state 8, = (h,, zx where h, = [h( žt- ot-1) is a deterministic recurrent state capturing temporal dependencies, and is a stochastic state encoding uncertainty about the current observation. The separation of deterministic (h₂) and stochastic (2t) states enables long-horize
temporal reasoning while maintaining diversity in imagined rollouts. The RSSM operates in two distinct modes: (i) during training. the model conditions on the current observation of to infer 2,, ensuring alignment with real-world dynamics (posterior inference); and (ii) during imagination, the model predicts future states without access to observations, sampling & to generate hypothetical trajectories (prior prediction). The world model reconstructs observations ô, predicts rewards fr. and estimates episode continuations îi from the latent state 8t. These components learned via the following structured models: posterior state representations (encoder) 2, qo (zx h4, 0t), prior state representations å ~ po (x), reward predictor rt po(fx) he, 2t), continue predictor fūІ ~ (o, 2t), and decoder 6, po (h, 21).
Learning in imagination. Once the world model is trained, behavior is learned by optimizing a policy entirely within the latent space. An actor-critic architecture guides this process, with the critic estimating the cumulative future reward (value v.)) and the actor selecting actions a, ~ nhe to maximize this value. By decoupling policy learning from real-world interaction, DreamerV3 achieves strong sample efficiency, iteratively refining its world model and behaviors through cycles of imagination and execution.
4.2 Context Encoder for Dynamics-Aligned Representations
While DreamerV3 is effective fixed environments, its reliance on static latent states limits adapt- ability to contextual variations, such as changes in gravity or friction. To address this, we introduce Dynamicx-Alignen Latent Imagination (DAL.I, which extends DreamerV3 with a self-supervised context encoder that learns explicit context representations from interaction histories.
3

---

4.2.1 Forward Dynamics Alignment
The context encoder 9e maps history of observations and actions (oi-Kæt (():t-1) to context representation åt = (pt Ktt Ot-K=-1), where defines the temporal window. To align ås with environmental transition dynamics, we optimize a forward dynamics loss:
(+p = E 0++1 f (ot. 04, 3t)
where (or, (ot. 3t) predicts the next observation 0++1 given the current observation Ot, action at, and context šc- Minimizing LTD ensures that ši encodes contextual factors eritical for accurate dynamics prediction, such as variations in physical parameters. The parameters yp of 9₂ and f are trained jointly, enabling the encoder to capture essential dynamies from short interaction histories.
4.2.2 Cross-Modal Regularization
To enhance the context encoder's robustness, we introduce a cross-modal regularization that aligns the context representation åt 9p(0r-K (ott (t-K:t1) with the Dreamer V3 world model's posterior state 2 ~ q0(2) (h,, or). The cross tross-modal loss enforces bidirectional alignment between åe and žy:
Leross (4) E 2| W W 112 + E 1|sc W
where W, and W are linear maps between the context and state spaces. This bidirectional recon- struction leverages 2, 's ceservation-informer representation, which captures instantaneous dynamics relevant to the current context. By aligning with 2, instead of the full latent state 8, = (hh, z+ ). åt avoids encoding redundant trajectory-specific information from the deterministic h. which could impair generalization. The bidirectional constraints prevent degenerate solutions (e.g- ås collapsing to constant) by enforcing invertibility, ensuring åe remains a rich, context-specifet representation consistent with the 2t's latent dynamics. The total loss combines both objectives:
Loual (4) = LFD(ş) + Across Lenoss (pp);
where Agross E {0, 13 balances dynamics prediction and cross-modal alignment.
4.3
Integrating Context into DreamerV3
To enable context-conditioned imagination and policy learning in DreamerV3 without access to ground-truth context variables, we propose two integration strategies for incorporating the context representation 3r-
This approach appends 3: the observation embedding in the world model's encoder, modifying it to 24 go(2) /he, or, 3t). All other world model components, including the sequence model h fo (14-11 2;-1, 1, predictors (ε,, ft. fu. o₂), and actor-critic networks (ay](,,), (st)), remain unchanged. This lightweight modification enriches latent states with dynamics- aware context, enabling the world model adapt implicitly to contextual variations.
Shallow Integration.
Deep Integration. For deeper context awareness, is integrated into the world model and policy as follows: he = fo (14-1, 1, (,))
~- po [ht, 21, åt); nq ~ po (1, àt);
po(or [4, 2t, 3t). The posterior zt ~ (r, 0:) and prior & ~ po (p( access indirectly through h4. The actor and critic explicitly condition on åt optimize imagined trajectories over horizon H: a, ~ (a-(ssrr)))) 0gg (st, Ex((),)))) D y -r-"r, This ensures that policy optimization explicitly accounts for contextual variations, enabling context-condditoned imagination.
Training. Training the context encoder, the loss (3) aligns 3: go (ottoKKtt 4t- k2--1) with environ- mental dynamics and requires careful handling of the DreamerV3 world model's recurrence. Both integration strategies unroll the world model over K Kllength window, initializing h,-K learned initial state), generating 3r = g (or a -xr-=-1) for = t K to t. In Shallow Inte- gration, 2, qo (2+ (t, 0r, 37) and h fo Z,, a,); Deep Integration, 2, ~ qo (2 o,) and h++1 = (h,,, 2r, år) For both, gradients through h, and 2+ are stopped in the recurrent dy- In Shallow Integration,
(or
namics, and through h- in the world model's 's encoder, preventing updates 3r gradients are preserved in the world model's 's encoder and losses LHD and Leross allowing updates to ip Wa, and W,. Deep Integration, år gradients are stopped in the recurrent dynamics and preserved only in LTD and Lemss- This decouples context learning from the world model's recurrence, ensuring 3e captures relevant dynamics for both strategies. For detailed pseudocode, see Algorithms and in Appendix B for Shallow Integration, and Algorithms 3 and for Deep Integration.

---

5 Theoretical Insights into DALI's Context Encoder
In this section, we provide an exposition of our key theoretical results, elucidating why DALI's context encoder is essential for robust generalization. We emphasize conceptual insights, with formal statements and proofs deferred to Appendix A
Our analysis is grounded in cMDP framework with continuous contexts (e.g., physical parameters such as gravity or actuator strength), sampled i.i.d. per episode from distribution Pc. Observations are noisy (e.g., Or = St T/. 1)+ ~ AN(0, I)), and dynamics are assumed to be Lipschitz continuous, so that small changes in context (e.g., slight shifts in gravity) lead to smoothly varying behaviors (e.g., comparable object trajectories or locomotion patterns). The observation-action process is 3-mixing [Rio| 2017], meaning distant observations become nearly independent over time (e.g.. periodic motions or gait cycles decorrelate under exploratory actions). These structural assumptions are characteristic of many continuous control environments such as DMC [Tassa et al.) 2018).
CARL [Benjamins et al. 2023], and MetaWorld (Yu et al. 2020], and they underpin DALI's efficient context inference. f-mmxing captures realistic decorrelation in RI. tasks, enabling short interaction windows to capture c. An exploratory policy further ensures that distinet contexts yield distinguishable trajectories, a common feature in RI. training.
Domain randomization trains DreamerV3 across diverse contexts, with e PC, using hRSSM fe(hz-1, -111 (x-1) and 2 ~ qo ( (te 0z). The recurrent state A.SSMM implicitly accumulates information about e over time through its interaction history. For instance, in environments with pendulum-like dynamics, increased gravity may result in faster oscillatory behavior, a pattern that hr &RSSM can gradually encode after observing enough transitions. However, this strategy has limitations. The recurrent state h.SSSM RSSM compresses all episode information, including context, states, and noise into a fixed-size GRU, introducing an information bottleneck. As dynamic state information (e.g.. object motion) and transient noise compete for limited capacity, essential cues about the underlying context (e.g., gravity) may be lost or delayed. Identifying e often requires accumulating evidence across many transitions, potentially spanning an entire episode of length T, since early actions may not sufficiently excite the dynamics to reveal contextual differences. This delay hinders rapid adaptation to novel contexts, as RRRS needs a long temporal window to reliably disambiguate
DALI's context encoder decouples context inference from dynamics modeling. It functions as a specialized module that learns a representation år capturing context information from short, local histories, allowing its recurrent state h. h.PAL to focus on dynamies and reducing the burden on the GRU, The B-mixing property ensures that K transitions suffice to encode as the dynamics' dependence on past states fades rapidly. This enables faster adaptation compared to DreamerV3's h 1.RSSM, We formalize this intuition for the case of Deep Integration, where the context encoder enhances the RSSM's context awareness. The functions h(-) and Z(:+-) denote the differential entropy and mutual information, respectively (Thomas and Cover 1999].
Theorem 1. In a cMDP with 9-mixing and Lipschitz dynamics, DALI's context encoder åt captures near-optimal context information, I(c; 2 (1-6))(c)for E 10, using N = (115²²) windows of K = 21 / transitions, where X is the mixing rule. Moreover, DALI's RSSM relains more context information than DreamerV3's, I(c; hDALl) 2 I(c; hySSM) <(K), with <(K) = (x) Compared to DreamerV3's 's RSSM processing full episodes T >>>>>>> K transitions, DALI achieves sample complexity gain of (TT/K).
For the formal statement and proof of Theorem I see Appendix The 3-mixing property ensures that a short window of K = 2[log(1/ (1/5)/) transitions is sufficient to achieve high context fidelity. For instance, when 8 = 0.01 and the typical mixing rate in DMC tasks steps results in a conditional entropy error bounded by 8' x 0.1. DALI requires N = ((1/²²) such windows, totaling ((/88)) transitions. In contrast, DreamerV3's h. RRSSM conntrainee by the GRU's finite capacity and observation noise, loses context information and requires ( transitions per episode (e.g.. = 1000), where T >>>>> K. The sample complexity gain of O(T/K) reflects DALI's efficiency in exploiting local histories. Furthermore, åt achieves superior fidelity, with I(c;st) approaching h(c) within 6, while DreamerV3's RSSM incurs a larger information loss, <(K) ( )x/22), due to its bottleneck. While the theorem assumes Deep Integration, this sample complexity gain remains consistent across Deep and Shallow Integration, as 3, 's efficiency derives from its training on K-length windows, which is identical both configurations. These results highlight the necessity and efficiency of DALI's context encoder in enabling robust generalization.
0.1, a window of K zs 64
5

---

Experiments and Analysis
In Sections 6.1 and 6.2 we evaluate DALI's ability to generalize in zero-shot manner across unseen
context variations. In Section 6..] we show that the dynamies-aligned context encoder learns a structured latent representation, where perturbations individual dimensions produce physically plausible counterfactuals shorter http:/hwwww..ccm//aaaattrrerrDad
shorter ball swings for higher imagined gravity). Our code is available
6.1 Zero-Shot Generalization
We evaluate DALI's zero-shot generalization on contextualized DMC Ball-in-Cup and Walker Walk tasks from the CARL benchmark [Benjamins et al. 2023).
Methods. Our context-unawenn methods, denoted DALI-S/D X/blank, employ either Shallow (S) or Deep (D) integration with the inferred context še (see Section 4.3 These models are trained using either the forward dynamics loss alone T (denoted as "blank" e.g., DALI-D), or the cross-modal regularized objective (2) (denoted as X, e.g., DALI-S-X). We compare our models to the context- unaware baseline Dreamer-DR, which corresponds to DreamerV3 with domain randomization [Tobin) et al. 2017]. We also evaluate against the context-aware baselines cRSSM-S and cRSSM-D, which use the same world model and actor-critic architecture but directly incorporate the ground-truth context [Prasanna al. 2024]. For further details, Appendix A.. A.1
Setup. We train our methods using the small variant of DreamerV3 [Hafner al. 2025 with a transformer-bar context encoder (Vaswani et al. [017]. for 200K timesteps (Ball-in -in-CinCi or 500K timesteps (Walker) across 10 random seeds, following the setup of Prasanna et al. 2024 Hyperparameters and architectural details are provided in Appendix
We conduct evaluation on two observation modalities: Featurized, which provides structured state inputs with low partial observability, and Pixel, which requires latent state estimation from raw images. We assess performance under three generalization regimes (Kirk al. 2023 : Interpolation (contexts within the training range), Extrapolation (OOD contexts beyond the training range), and Mixed (one context OOD, one within training range). The Extrapolation regime tests generalization beyond seen contexts, requiring agents to capture underlying physical principles. The Mixed regime probes the ability to interpolate and recombine learned representations. Together, these settings reveal whether the agent learns meaningful abstractions or simply memorizes specific contexts.
Context ranges. For Ball-in-Cur inp, the context parameters are gravity (training: [4.9, 14.7], evalu- ation: [0.98, 4.9) U (14.7, 19.6), default: 9.81) and string length (training: [0.15, 0.45). evaluation: [0.03, 0.15) (0.45, 0.6). default: 0.3). For Walker, the parameters are gravity (same ranges as Ball-in-Cup) and actuator strength (training: [0.5, 1.5), evaluation: 10.1, 0.5) U (1.5, 2.0]. default: 1.0). These OOD ranges, particularly in Ball-in-Cup. pose major generalization challenges, while Walker's is actuator strength remains closer to training. Full training settings, including single and dual context variations, are in Appendix C
Evaluation Metrics. To evaluate zero-shot generalization, we use the Interquartile Mean (IQM) and Probability of Improvement (Pol) metries from the rliable framework (Agarwal et al. 2021 IQM averages the central 50% of performance scores, reducing the impact of outliers prevalent in RL due to high variance in training dynamics, particularly in OOD cMDP contexts (e.g., extreme gravity in Ball-in-Cup), We ensure statistical reliability by computing stratified bootstrap 95% confidence intervals over seeds and aggregated contexts. Pol quantifies the probability that randomly selected run of one algorithm outperforms a randomly selected run of another, offering a robust comparative metric across methods and modalities.
Results. We report IQM and Pol across 10 seeds in Figure comparing DALI-S, DALI-S-X, Dreamer-DR, cRSSM-S, and cRSSM-D for Featurized and Pixel observations the Ball-in-Cal -u- and Walker Walk tasks.
In Ball-in-Cup in-Cup's Interpolation, all methods achieve high IQM (0.92-0.95), with DALI-S-X com- petitive (0.9490 Featurized, 0.9440 Pixel). In Extrapolation. DALI-S-X excels, outperforming Dreamer-DR by 87.9% (Featurized, IQM 0.3720) to 96.4% (Pixel, IQM 0. 0.2730), surpassing cRSSM- by 63.9% (Featurized) and 45.9% (Pixel), and eRSSM-D by 33.8% (Featurized) and 12.8% (Pixel). This suggests ground-truth context may overfit, limiting OOD adaptability. Low absolute IQM scores

---

Ball-in-Cup
Interpolation
Extrapolation
Mixed Regime
PIDALI-S-Z Y) (Extrap)
DALI-S-Z
DALI-S
cRSSM-D
cRSSM-S
Dreamer-DR
0.928 0.936 0,944 0.952 0.18 0.24 0,30 0.36
0.48 0.56
0.64
0.00
0.25
Normalized Score
Normalized Score
Normalized Score
Probability
(a) Featurized
Interpolation
Extrapolation
Mixed Regime
PIDAL- > Y) (Extrap)
DALI-S-X DALI-S
cRSSM-D
cRSSM-S
Dreamer-DR
0.935 0.940 0.945 Normalized Score
0.72 0.18
0.24
0.30
0.40 0.48
0.56 0.64
0.00 0.25 0.50 0.75 1.00
Normalized Score
Normalized Score
Probability
(b) Pixel
Walker Walk
Interpolation
Extrapolation
Mixed Regime
P(DALI-S Y) (Extrap)
DALI-S-X DALI-S
CRSSM-D
cRSSM-S
Dreamer-DR
0.945
0.960
0.975
0.69
0.72
0.75 0.78
0,800 0.825 0.850 0,875 Normalized Score
0.00 0.25 0.50 Probability
Normalized Score
Nomalized Score
(a) Featurized
Interpolation
Extrapolation
Mixed Regime
P(DALI-S Y) (Extrap)
DALI-S-X
DALI-S
cRSSM-D D cRSSM-S
Dreamer-DR
0.95 0.96 0.97
0.69 0.72 0.75 0.78
0.80 0.82 0.84 0.86
0.00 0.25 0.50 0.75 1.00
Normalized Score
Normalized
Normalized Score
Prohahility
(b) Pixel
Figure I: Interquartile Mean (IQM) scores and Probability of Improvement (Pol) for the DMC Ball- in-Cup and Walker Walk tasks under contextual variations: gravity and string length for Ball-in-Cup. and gravity and actuator strength for Walker Walk. Results are shown for Featurized and Pixel observations in each environment. Scores aggregate across single and combined contexts. Shaded intervals represent 95% stratified bootstrap confidence intervals over seeds and aggregated contexts. The rightmost panel each plot displays Pol in the Extrapolation regime for DALI-S-X (Ball-in-Cup) and DALI-S (Walker Walk), relative to baseline methods.
(0. 14-07737 reflect the extreme OOD context ranges, such as gravity values 0.98 19.6, far from the training range of [4.9. 14.7). In the Mixed regime, DALI-S-X leads in Featurized modality (51.1% over Dreamer- DR, 20.3% over cRSSM-S, 1.9% over cRSSM-D, IQM 0.6830), while cRSSM-D excels in Pixel (IQM 0. 0.6250), leveraging ground-truth context for visual dynamics. DALI-S-X shows moderate Pol (> 0.6) against Dreamer-DR DR both modalities for Extrapolation (see Figure IT.
In Walker's Interpolation, methods score high (0997...97). with DALI-S leading (0.9710 Featurized). In Extrapolation, DALI-S achieves 4.0% over Dreamer-DR (Featurized, IQM 0.7810), 11.3% over cRSSM-S, and 4.3% over cRSSM-D. Higher IQM scores (0.7-0.78) across methods reflect less extreme OOD actuator strength (10.1, 2.0] vs. [0.5, 1.5]), reducing generalization demands. Pixel modality, context-aware cRSSM-S leads (IQM 0.7770), leveraging ground-truth context for robust visual feature extraction. DALI-S attains an IQM 7580,, compared DALI-S-X 0.7300. the Mixed regime, context-aware cRSSM-D leads (Featurized, IQM 0.8720), with DALI methods gaining

---

1.1%-1.8% over Dreamer-DR; in Pixel. cRSSM-S leads (0.8610), with DALI gains 1.1%-5.5%. In Pixel Extrapolation, DALI-S shows moderate Pol (~~ 0.6) against Dreamer-DR and cRSSM-D, while cRSSM-S leads, leveraging ground-truth context for visual dynamics (see Figure
6.2 Generalization Trends Across Environments and Modalities
We provide additional interpretations of the IQM results in Figure focusing on the generalization patterns of DALI-S and DALI-S-X. We emphasize modality-speciff effects, environmental nu- ances, integration strategies, and comparisons with context-aware baselines (cRSSM-S, cRSSM-D), highlighting the role of cross-modal regularization.
Cross-Modal Regularization Effects. DALI-S-X consistently outperforms DALI-S in Ball-in- Cup across all regimes and modalities, indicating that cross-modal regularization (Leross see 2 enhances the context encoder's ability to capture pendulum dynamics. Aligning de with the world model's 's posterior 2 enhances context inference for nonlinear dynamics, benefiting from low partial observability in Featurized inputs and stabilizing inference in Pixel modality's noisy visual inputs. In contrast, DALI-S outperforms DALI-S-X in Walker across most regimes and modalities, indicating that the forward dynamics loss (LrD. see alone suffices for contexts where actuator strength linearly scales joint torques. By simply optimizing for next-step predictions, DALI-S effectively handles predictable torque scaling, particularly in the Pixel modality where visual noise may amplify Leross complexity in DALI-S-X. This underscores the need for task-specific regularization. Ball-in- exhibits larger performance drops from Featurized to Pixel modalities (c.g., DALI-S-X: 0,3720 to 0.2730 in Extrapolation) compared to Walker, reflecting higher partial observability in visual settings with nonlinear dynamics.
Comparison with Context-Aware Baselines. DALI-S-X outperforms cRSSM-S and cRSSM-D in Ball-in-Cup Featurized Extrapolation (0.3720 vs. 0.2270, 0.2780) and Pixel Extrapolation (0.2730 vs. 0.1870, 0.2420), indicating that 3x 's inferred context generalizes better than ground-truth c, which tends to overfit to training contexts. Walker Featurized Extrapolation, both DALI-S and DALI-S-X surpass cRSSM-S and cRSSM-D (0. 7810, 0.7770 vs. 0.7020, 0. 497),, benefiting from 's inference of actuator torque scaling. However, in Ball-in-Cup Pixel Mixed, cRSSM-D (0.6250) outperforms DALI-S-X (0.6030), indicating that Deep Integration of ground truth e better supports modeling of complex visual dynamics. Similarly, in Walker Pixel Extrapolation, cRSSM-S (0.7770) surpasses DALI-S (0.7580), leveraging ground-truth context for precise visual modeling.
Shallow Context Propagation as Regularization. In our experiments, Shallow Integration, as implemented in DALI-S and DALI-S-X, consistently performs well, leading to its exclusive use in the results reported here. Shallow Integration incorporates the inferred context 3r solely in the world model's encoder, 2, go(22 [ir, 3t), allowing context information to propagate indirectly to the recurrent state h = fo 2, at-1) through recurrence. This design regularizes the world model, potentially mitigating overfitting to noisy 3e estimates, which can be particularly beneficial in OOD settings. The empirical effectiveness of Shallow Integration in our experiments suggests that implicit context propagation can effectively leverage the 3-mixing property in practice.
6.3 Validating Physically Consistent Counterfactuals
We assess whether structured perturbations to the latent space induce counterfactual trajectories that adhere to Newtonian physics in the Ball-in-Cup task, validating that encodes mechanistic factors (e.g., gravity, string length) critical for generalization. For this task, the agent must swing a ball into a cup under variable gravity and string length. In setup, the learned context representation 3 is an S-dimensional vector, i.e., E IR produced by the dynamics-aligned encoder. Our goal is to identify which latent dimensions in dominantly encode these mechanistic factors.
......., 8, we sample a fixed observation or from a test episode and use the learned context encoder 9, infer the representation åt (or Qt-K:t-1) with K = 50. Encoding or into the latent state st = {hr, 3. we then roll out actions 0r:t+ under both the original and perturbed j' 3 + A ej. where e, is one-hot vector, and A is the standard deviation of ăj across the dataset. Using the frozen world model and policy, we decode the predicted observation sequences, yielding two diverging trajectories: baseline trajectory under the original 3. (v) = (ōr, ōttt, 0z J. and counterfactual trajectory under the perturbed 3'. T (ō), 1',1, .( with H 50. We repeat this process N 2500 times to generate 2500
For each latent dimension J I,....8,
8

---

baseline trajectories and 2500 perturbed trajectories for each ăj. yielding a total of 5000 trajectory samples: D.j) (71 T T2500: T2500 J for 3y. We train a binary classifier c to distinguish between perturbed trajectories T'U) (label 1) and baseline trajectories T (label 0). The classifier outputs the probability Pn, = P(label = T.) that given trajectory Tn E DD.))
of perturbed. High classifier accuracy for given dimension suggests it captures mechanistic relationships rather than superficial correlations, with perturbations inducing systematic and physically consistent changes.
We compute 95% bootstrap confidence intervals (CIs) for the classifier AUC for each latent dimension ăj and rank the 8 dimensions based on the AUCs, and select the top dimension of j for validating physically consistent counterfactuals. See Appendix D details on the classifier implementation and additional experiments on physically consistent counterfactuals.
6.3.1 Mechanistic Alignment in Latent Imagination: Gravity-String Dynamics
We demonstrate DALI's capacity for physically consistent latent imagination by analyzing coun- terfactual trajectories generated through perturbations to the top-ranked latent dimension (36) in DALI-S-X (Shallow Integration with cross-modal regularization). We show how åe encodes coupled gravity-string dynamics in the Ball-in-Cuu task.
Our analysis reveals consistent physical behavior across both Featurized and Pixel-based modalities. We demonstrate systematic alignment between the original imagined trajectories and their counterfac- tual counterparts generated by perturbing the latent context representation in trained world model. For the Featurized experiments, we tracked predefined state variables (e.g., ball position and velocity), while in the Pixel experiments, we captured rendered environment frames at 5-step intervals over 64-step imagination horizon (excluding the final frame for visual clarity).
Leveraging our AUC-based ranking analysis, we identify åe as the dominant latent dimension across modalities. To isolate the influence of this dimension on passive dynamics, we generate counterfactual trajectories by perturbing 36 and enforce zero-action rollouts. This suppresses policy-drive behavior, exposing the system's 's intrinsic swinging dynamics.
Pixel Modality. Figure 2a contrasts the original (top) and perturbed (middle) imagined trajectories. The perturbed rollout (initialized from the same observation) reveals two key effects: (1) Shorter string length: The blue ball (counterfactual) hovers higher than the original (yellow) in frame 40, indicating a shorter string, and (2) Higher acceleration: The counterfactual ball overtakes the original in frame 15 and 45, demonstrating faster swing cycles due to increased gravitational pull. These observations align with Newtonian mechanics: shorter strings reduce string length, increasing oscillation frequency, while higher gravity amplifies acceleration.
Featurized Modality. Figures Saauressbbanne 2b]and [2 quantify these effects. The counterfactual trajectory (orange) in Figure 25 exhibits reduced amplitude (lower peak Z-position) compared to the original (blue), consistent with the pixel-based evidence of a shorter string. This aligmment confirms that perturbing shortens the string length, physically constraining the ball's swing. In Figure 2c the counterfactual's velocity peaks earlier and higher, confirming faster acceleration under perturbed 36. Both modalities reveal that 36 encodes a coupled relationship between gravity and string length, as perturbation simultaneously alters both parameters in a physically plausible manner.
Discussion and Outlook
We introduced Dynamics-Aligned I Latent Imagination (DALI), a framework built on DreamerV3 that enables zero-shot generalization by inferring latent contexts from interaction histories. In the Extrapolation regime, our methods DALI-S and DALI-S-X achieve substantial performance gains over the context-unawan Dreamer-DR bascline, ranging from +4.0% to +96.4% across Featurized and Pixel modalities in the DMC Ball-in- and Walker Walk environments. Notably, DALI-S-X surpasses the ground-trum -trththhhhhhhh context-award baselines cRSSM-S by +45.9% to +63.9% and cRSSM-D by +12.8% to + 33.8% in Ball-in-Cup across both modalities. Additionally, DALI-S outperforms cRSSM-S by + +11.3% and cRSSM-D by +4.3% in Walker's Featurized modality, demonstrating effective zero-shot generalization unseen contexts.
Limitations for DALI's Context Inference. Theorem Cheorem@demonee demonstrates that DALI's context encoder achieves near-optimal context inference. I(c) 2 (1 8)h(c) using short windows of length

---

= ( yielding a sample complexity gain of O(T/K) over DreamerV3, assuming S- mixing and Lipsehitz dynamics. As an cotommmiion-thhoretti result, it does not address the downstream impact of å policy performance, which depends on joint optimization of the context encoder, world model, and actor-critic components. Reliance on an exploratory policy may falter in sparse-reward or high-dimensional settings, producing noisy estimates. Furthermore, the f-mixing assumption may not hold environments with slow-mixing dynamics, such as highly correlated trajectories or restricted exploration, potentially limiting generalization. Practical challenges, including training instability, sensitivity to hyperparameters, and the risk overfitting high-dimensional observation spaces, also beyond its scope. Future theoretical work should model how context representations influence policy learning and robustness in complex RL enviromments.
Integration Strategies and Practical Implications. Empirically, DALI's Shallow Integration strategy, where de appended only to the encoder input, consistently delivers strong performance in the tasks we evaluated, particularly under OOD conditions. Its simplicity acts as regularizer, implicitly propagating åe through recurrence and mitigating overfitting, yielding gains of up to +87.9% (Featurized) and +96.4% (Pixel) over Dreamer-DR. Future work could investigate the regularization effects of Shallow Integration, explore hybrid strategies interpolating between Shallow and Deep, and develop theoretical models connecting context inference policy performance. These directions would deepen our understanding DALI's integration mechanisms and their practical impact generalization.
0
5
10
15
20
25
30
35
40
45
50
55
60
(a) Pixel Reconstructions
Original Counterfactual
0.3
Bail Posttion
Z Velocity
0.0
0.2
-0.5
-1.0
0.0
Counterfactusl
52
20
30
Time
Time
(b)
(c)
Figure 2: (a) (Pixel Modality) Counterfactual Trajectories in Pixel Space: Top: Original imagined trajectory the Ball-in-Cup under default gravity and string length. Middle: Perturbed trajectory after adding noise A to the top-ranked latent dimension Bottom: Pixel-wise differences (6 = jo, 6(1). The perturbed trajectory (blue) exhibits shorter string (frame 40) and faster acceleration (overtaking the original trajectory in frames 15 and 45), aligning with increased gravitational effects. Rollouts use zero actions to isolate passive dynamics. (b) (Featurized Modality) Ball Z-Posiiion Under Latent Comparison of original (blue) and counterfactual (orange) ball height trajectories. The perturbed åe reduces oscillation amplitude (lower peak Z- position) and accelerates descent, consistent with shorter string length and higher gravity. (c) (Featurized Modality) Ball Z- Z-Velocity Under Latent Perturbation: Velocity profiles for original (blue) and counterfactual (orange) trajectories. The perturbed induces earlier and higher velocity peaks, confirming faster swing dynamics. This mirrors the pixel-beee based evidence increased gravitational acceleration.
Perturbation:
10

---

Acknowledgments and Disclosure of Funding
The authors thank the anonymous reviewers for their valuable feedback and the NeurIPS community for fostering open and rigorous scientific exchange. JB and ME gratefully acknowledge funding by the German Research Foundation DFG through the MoReSpace (402776968) project.
References
Rishabh Agarwal, Max Schwarzer, Pablo Samuel Castro, Aaron Courville, and Mare G Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in Neural Information Processing Systems, 2021.
Philip Ball, Cong Lu, Jack Parker-Holder and Stephen Roberts, Augmented world models facilitate zero-shot dynamics generalization from single offline environment. International Conference on Machine Learning. pages 619 629, 2021.
Jacob Beck, Risto Vuorio, Evan Zheran Liu, Zheng Xiong, Luisa Zintgraf, Chelsea Finn, and Shimon Whiteson. A survey of meta-reinforcemmntt learning. arXiv preprint arXiv:2301 2023.
Carolin Benjamins, Theresa Eimer, Frederik Schubert. Aditya Mohan, Sebastian Döhler, André Biedenkapp. Bodo Rosenhahn, Frank Hutter, and Marius Lindauer. Contextualize me - the case for context in reinforcement learning. Transactions on Machine Learning Research, 2023.
Michael Beukman, Devon Jarvis, Richard Klein, Steven James, and Benjamin Rosman. Dynamics generalisation in reinforcement learning via adaptive context-aware policies. Advances in Neural Information Processing Systems, 36:40167-40203, 2023.
Stéphane Boucheron, Gábor Lugosi, and Pascal Massart. Concentration Inequalities A Nonasymptotic Theory of Independence. Oxford University Press, 2013.
Tao Chen, Adithyavairavan Murali, and Abhinav Gupta. Hardware conditioned policies for multi-robot transfer learning. Advances in Neural Information Processing Systems, 31, 2018.
Yan Duan, John Schulman, Xi Chen, Peter I. Bartlett, Ilya Sutskever, and Pieter Abbeel. RL3: Fast reinforcement learning via slow reinforcement learning. arXiv preprint 16..v:::11110.2777 2016.
Hamid Eghbal-zadee zadeh, Florian Henkel, and Gerhard Widmer. Context- Context-adaptive reinforcement learning using unsupervised learning of context variables. In NeurIPS 2020 Workshop on Pre-registration in Machine Learning, pages 236-254, 2021.
Ben Evans, Abitha Thankaraj, and Lerrel Pinto. Context is everything: Implicit identification for dynamics adaptation. In 2022 International Conference Robotics and Automation (ICRA), pages 2642-2648, 2022.
Chelsea Finn, Pieter Abbeel, and Sergey Levine. Model-agnostic meta-learning for fast adaptation of deep networks. In International conference on machine learning, pages 1 1126-1 135, 2017.
Scott Fujimoto, Pierluca D'Oro, Amy Zhang, Yuandong Tian, and Michael Rabbat. Towards general-purpose model-free reinforcement learning. International Conference on Learning Representations, 2025.
Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak I Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Conference on Machine Learning, pages 2555-2565, 2019.
Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream control: Leamnnng behaviors by latent imagination. In International Conference on Learning Representations, 2020.
Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba. Mastering Atari with discrete world models. In International Conference on Learning Representations, 2021.
Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap. Mastering diverse control tasks through world models. Nature, pages 1-7, 2025.
Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual Markov decision processes. arXiv preprint arXiv: 1502..2259009 2015.
Nicklas Hansen. Hao Su, and Xiaolong Wang. TD-MPC2: Scalable, robust world models for continuous control. In International Conference on Learning Representations, 2024.

---

Robert Kirk, Amy Zhang, Edward Grefenstette. and Tim Rocktäschel. A survey of zero-shot generalisation in deep reinforcement learning. Journal of Artificial Intelligence Rexearch, 76:201-264, 2023.
Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context- dynamics model for generalization model-based reinforcement learning. In International Conference on Machine Learning. pages 5757-5766, 2020.
Luckeciano C Melo. Transformers are meta-reinforeement learners, International Conference on Machine Learning, pages 15340-15359, 2022.
Aditya Modi, Nan Jiang. Satinder Singh, and Ambuj Tewari. Markov decision processes with continuous side information. In Algorithmic Learning Theory, pages 397-618,20 2018.
Mehryar Mohri, Afshin Rostamizadch, and Amect Talwalkar. Foundations of Machine Learning. MIT press, 2018.
Yao Mu, Yuzheng Zhuang. Fei Ni, Bin Wang. Jianyu Chen, Jianye Hao, and Ping Luo. Domino: Decomposed mutual information optimization for generalized context in meta-reinforceemnt learning. Advances in Neural Information Processing Systems, 35:27563-27577 2022.
Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald S Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt in dynamic, real-world envirouments through meta-reinforcement learning. In International Conference on Learning Representations, 2018.
Sai Prasanna, Karim Farid, Raghu Rajan, and André Biedenkapp. Dreaming of many worlds: Learning contextual world models aids zero-shot generalization. Reinforcement Leurning Journal, 1, 2024.
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off ooflcc meta- reinforcement learning via probabilistic context variables. In International Conference Machine Learming. pages 5331-5340, 2019.
Emmanuel Rio. Asymptotic Theory Weakly Dependent Random Processes. Springer, 2017.
Younggyo Seo, Kimin Lec, Ignasi Clavera Gilaberte, Thanard Kurutach, Jinwoo Shin, and Picter Abbeel. Trajectory-wise multiple choice learning for dynamics generalization reinforcement learning. Advances in Neural Information Processing Systems, 33:12968- 11112668-2299,, 2020.
Scyed Kamyar Seyed Ghasemipour, Shixiang Shane Gu, and Richard Zemel. Smile: Scalable meta inverse reinforcement learning through context-conditional policies. Advances in Neural Information Processing Systems, 32, 2019.
Shagun Sodhani. Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context-based represen- tations. In International Conference Machine Learning, pages 9767-9779, 2021.
Shagun Sodhani, Franziska Mcier, Joelle Pineau, and Amy Zhang. Block contextual MDPs for continual Icarning. In Learning for Dynamics and Control Conference, pages 608 623, 2022.
Yuval Tassa, Yotam Doron, Alistair Muldal, Tom Erez, Yazhe Li, Diego de Las Casas, David Budden, Abbas Abdolmaleki, Josh Merel, Andrew Lefrancq, et al. DeepMind control suite. arXiv preprint 1801:18001.0000 2018.
Joy A. Thomas and Thomas M. Cover. Elements of Information Theory. John Wilcy & Sons, 1999.
Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojeiech Zaremba, and Pieter Abbeel. Domain random- ization for transferring deep neural networks from simulation to the real world. In International Conference on Intelligent Robots and Systems, pages 23-30, 2017.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention is all you need. Advances in Neural Information Processing Systems, 30, 2017.
Zhou Xian, Shamit Lal, Hsiao- Yu Tung, Emmanouil Antonios Platanios, and Katerina Fragkiadaki. Hyper- dynamies: Meta-learning object and agent dynamics with hypernetworks. In International Conference on Learning Representations, 2021.
Zhenjia Xu, Jiajun Wu, Andy Zeng, Joshua B Tenenbaum, and Shuran Song. Densephysnet: Learning dense physical object representations via multi-step dynamic interactions. In Robotics: Science and Systems (RSS), 2019.
12

---

Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman. Chelsea Finn, and Sergey Levine. Mela-world: A benchmark and evaluation for multi-task and meta reinforcement learning, In Conference on Robot Learning, pages 1094-1100, 2020.
Luisa Zinigraf, Kyriacos Shiarli, Vitaly Kurin, Katja Hofmann, and Shimon Whiteson. Fast context adaptation via meta- vecaalea In International Conference Machine Learning, pages 7693-7702, 2019.
A Formal Results and Proofs
Theorems 2 and provide formal underpinnings for Theorem the main text, establishing the necessity and efficiency of DALI's context encoder with Deep Integration (see Section 4.3) compared DreamerV3 with domain randomization in cMDPs, focusing on context identifiability, sample complexity, and information bottleneck reduction.
Theorem 2 compares DALI with Deep Integration, using context encoder ăe, against DreamerV3 with domain randomization, using recurrent state h4, a cMDP. The theorem establishes: (1) DALI's captures at least much mutual information with context
DreamerV3's he, up to error <(K) O O(e-A//2)M(C that decays exponentially with window size K; (2) DALI achieves near- optimal Z(e; 3c) 2 (1 5))(c) with N = (1/88²) samples of K = (llog((/0)
transitions, leveraging shorter windows compared DreamerV3's s full episodes of length T. This demonstrates the necessity and efficiency of DALI's context encoder for context identifiability and learning.
Theorem 2 (Necessity and efficiency of Context encoder). Consider a cMDP M R, & p. Pe) with:
Continuous context C C Rd, static within episodes, drawn i.i.d. as ~ pc, with h(c) oc.
Lipschitz dynamics: 3L 0 such that Vs, e a. A,d
[Pa('"|s, (e s Elle- d112.
Observation noise: The observation function & satisfies ~ N(s,, with fixed variance
d² O.
8-mixing: The process T. = (0r, 0z) 3-mixing with
((K) sup
sup
P(AnB)- AK
where aTs,, S t) denotes the sigma-algebra generated by {Ts: t), and similarly for aTT,, 2 t for constants C, 0.
Sufficient exploration: The policy ensures that for any + d' E C, there exists a constant >> 0 such that KL(p(ot- Kit, at K [e)]pp(pp at (0-Ka-11 ((t-)(d)) 2 aK.
Moo Bounded log-likelihood ratios: log plozseleet Ki and actions to compact sets or imposing suitable policy constraints.
< M, ensured by restricting observations
Universal upproximator: DALI's 9e (neural network) approximates any continuous function on compact domains.
Training data: Both models trained trajectories from M, with contexts Pe i..d. per episode, using an exploratory policy. DreamerV3 with domain randomization trains its RSSM on episodes (ov.T,ay: a1:T-1) of length T. DALI trains its context encoder gip on K K-length windows (ot-Ka, Kt-Ka-1) with loss LFD [[[|o0+1 f (oc, Oe 11)) and its RSSM on full episodes, conditioning on ât-
Both modely use a GRU So with finite parameters for their recurrent state updates.
h4 = So(hr- 2t-1, 10-11) DreamerV3's recurrent state, with z ~ qo(ze(h4, ot), and åt g. (oro0-Ki ar ka-1) DALI's context encoding, with h, fo (hi-1,, 2t-1 0c-1.5t). Then:
13

---

1. Context identifiability: For any K, there exists 9, such that, within an episode:
3(c;3e) 2 I(c;hs) ((), c(K) -<(e-AA/2)M((e)
Sample complexity: For K = Ω [oR and ó E (0, DALI achieves
I(G3)) (1 8)h(c)
with N O (+)) samples of K transitions. vs. hypothetical DreamerV3 context estimator
Ie (olT, a1:T-1) requiring N = (+²) samples of' transitions.
Remark 3 (Non-trivial sample complexity gain of DALI's Context encoder). The sample complexity
gain of O(T K) for DALL over DreamerV3 with domain randomization, as shown in Theorem 2 arises from DALI's use of K K-length windows versus DreamerV3's 3's full episodes of length T. While
this gain appears to scale with input lengths, it is non-trivial: Both models achieve
(1)
samples, but DALI's window size K = (2loo(111))) (16) / X] leverages the ß-mixing property (Lemma 0 to ensure h(c(T=Kka) S Cee -K h(c). enabling efficient context identification. In contrast, DreamerV3 relies on longer trajectories (T > K) to achieve comparable context identification. The gain reflects DALI's context encoder exploiting short, local histories determined by the mixing rate X and desired accuracy 6.
Remark 4 (Consistency of DALI's sample complexity actoss integration strategies). The sample complexity advantage of DALI established Theorem for Deep Integration extends to Shallow Integration (Section 4.3, since the context encoder àe is trained identically
K K-length windows in both configurations. Incorporating the cross-modal loss Lenus in (2) aligns ăt with 24, which may allow the context representation àt leverage structured priors from the world model, potentially enhancing its robusiness. This addition does not affect sample complexity, as Lenss operates on the same K Kllength windows (Algorithm Z Recurrent unrolls are limited to K steps, and y updates depend solely on the current window, with detached states blocking backpropagation beyond it. The world model parameters (0) are frozen during context learning (inputs fe and qo are detached), so gradient stopping isolates vp and preserves the O(T/K) K) efficiency.
Theorem 5 proves that DALI's sequence model, incorporating context encoder åt, reduces the information bottleneck in cMDPs compared to DreamerV3's RSSM. Specifically, it shows that DALI's recurrent state retains context information, satisfying I(c; N.PAL)) I(c; hRSSM) <(K), where <(K) (()))) leveraging the encoder's ability to efficiently capture context over short windows.
Theorem 5 (Context encoder reduces information bottleneck). Consider a cMDP M (S, O, E,p,p))) pe) with:
Continuous context E C C Rd, static within episodes, drawn i..dd as c ~ pe, with h(c) < co.
Observations og = st + Th, To ~ N'O, o20, 00
Lipschitz dynamics: 3L > such that Vs,
[|Pa('lls,)) as1l, & Lllc-éll é'12-
8-mixing with B(K) < CK-AK, exploratory policy with KTop(r a <>0 O.
aK,
Context encoder 9o is a universal approximator, trained with forward dynamics loss LFD Elllo++1 - f (or, 01. 31)), satisfying hi(laa)) S - + 8', 8' = Oee-(((ee))) (Lemma
Both models use a GRU So with finite parameters.
Let RSSS = fa 1, 21-1 11-1) be DreamerV3's recurrent state, with 2, ~ qa(zelhe, oz), and lipALi = [ohhh Zt 1, be DALI's recurrent state. Then, for any window size K, DALI's
sequence model satisfies:
I(e; hpAAt) 2 -((K),
where <(K) O(e- implying that the context encoder St reduces the information bottleneck compared to DreamerV3's RSSM.
14

---

For B-mixing sequences with ((K) S Ce blocks of K -seeaarteed transitions are approximately independent. Lesaalbestablishes that, a 3-mixing cMDP, the conditional entropy of the context e given a history 7 a decays exponentially as h(c[Tr-K:)) K.)) I C h(c). This result is crucial for Theorems 2 and | as it quantifies how historical trajectories reduce uncertainty about the context, enabling DALI's context encoder to achieve near-optimal information capture with short windows, supporting both the identifiability and sample complexity claims.
Lemma 6 (Entropy decay from mixing), Consider cMDP with context
Pc, observation (Ot-K:t, Ot-K:t1),, where the process noise Ot 8t + The Ttt ~ NO, o² T) and history Tt-K:t = (Ot- K:t,
To (or, ae) 3-mixing with (KK S Ce-AK for constants C, C O, bounding dependence errors in tail probabilities by B(K), log-likkeliieeee likelihood ratios are bounded as logg plai,a; c'm -ii--]
plo,,aele;TL Rit 11 SM, the policy ensures exploration such that for any c + E EC, there exists a constant a 2 MV2X satisfying KL(p(Te Tre-o--e-llll kxlc))p(TT T(T (())kk)))))) Z aK, and e has bounded density on C C IR with differential entropy h(c) < cc. Then, there exists a constant C' > such that:
M((TT-Ka)) S C'e AK h(cc).
Proof of Lemma 6 We show that uncertainty about decreases exponentially with K, enabling DALI's dr to infer c.
By assumption, the process Tr (or, ae) is B-mixing, with ((K) Ce-AK, The condition (L(p(n-- KL(p(T-R:elellleg (etrr-::eellllpp_-all kall')) 2 aK, with M MV2X, ensures distinct histories. We need (())TKK)) S C"e (ee-Ah((), so we show:
(x) & ((K)(()
and let P. P(è of C/Ty Ka) be the probability of incorrectly estimating c. Fano's 's inequality gives Thomas and Cover 1999|
Discretize C into N points, log N h(c). Let è arg max,
MclT--Kal ka) < ((P.)) + Pah(c).
(4)
Error Pa = P(é # ele) occurs if there exists any # e such that p(Fr-Kald) p(T)-. By the union bound:
P Utp(t-kale))> poolxol I Σ
P(p(T-Kuud))> > p(t)ek e'
c'
For é' * define: 2 log kall- Then Z > 0 corresponds to the maximum likelihood
xull)
estimator incorrectly favoring c' over e. The probability of this error under context is P(Z > 0fc).
By the exploration assumption:
E, [2] = KKL(L(T--ETETT ( & -aK.
By the chain rule in the cMDP, -KTTk) = IT-= a.l (lo---)-- p(o1)T----11.c)) where T = (o--K:- at a
Thus, Z = Dist log pfo Ki 1.d)
plo /T- k.- log p(oi/T p 1.d')
p(o: (xt
K to 1, and Z, log pla.l Ra 1.o) 7 = -= - 7, Z; measures how much more (or less) likely (o,,a;) is under context d' compared to c, given past observations and actions. By the bounded log likelihodoo assumption, |Z,| < M. For independent Z,, Hoeffding's inequality [Boucheron et al. [013 bounds the tail probability for bounded random variables:
Define Z; log p(a,a; 7t K4-
for
p(ollte-x-nie)) So
2(aK)² K(2M)
a²K
P(Z [.[2] 2 aK) < exp (
exp
2M2
Since 7i is 3-mixing with 3(K) Ce-Ae the dependence adds error ((K), 3(K) bounds the deviation from independence between events separated by K time steps:
a K
P(z - E. 2 aK) < exp
+ (()).
2M
15

---

Set
o² 2M
Since 2 M Mv2X, we have 2 so exp((xx]] S e Thus:
P(Z 0/c) S exp(-kK) + ((K) exp(-AK) Ce- AK (1 Cke
Sum over N
contexts:
Pe s (N 1)(1 ()(1+))++++++111 +Je Ci3(K), C (N 1)(1 +
Apply to
b(dT-Kk)) S + For small 3(K), h(C] = 0, So:
(((T-Kk)) S CG,(KK)M())
With B(K) - Ce-Ak
M(TT-Ka) S Ci Ce -ekk-Kp CC'e
Lemma 7 provides a generalization bound for DALI's context encoder, showing that the conditional
entropy h( clas) is bounded by the entropy given the history plus small error, h((las) < h(clT=-K=)- 8', with 8' O v)) and 8 = Ce-AK This lemma is pivotal for Theorems and 5 as it ensures that the encoder 3r effectively captures context information with a sample complexity of N (1188)) facilitating DALI's efficiency and reduced bottleneck.
Lemma (DALL's generalization bound for context inference). Consider a cMDP satisfying Theo- rem 2 assumptions: context Pe with bounded density on C C Rdd and finite entropy h(c) co, observation noise ot 8t Th, Th (N0 I² B-mixing with 3(K) < Ce-AK for constants 2 aK for some a > O, (s'l as S Llfc c' d22,, and universal approximator 9 (or,,K K-tK-(-)., trained with the forward dynamics loss T.DD = F[[ J (ot, at de 113), satisfies:
C, > O, exploratory policy ensuring KL(p(Ti-Ka) (ca)(r--aa)))l(((TTT [t-
Lipschitz dynamics with [(Pe(s'|s.a)
Jr. DALI's context encoding 3e
h(cl3t) S h(cTT-Ka) +8', 8' O(vo),
with N o (++) samples, where d Ce-AK
Proof of Lemmal 7
The context encoder 94 minimizes the forward dynamics loss:
LTD Elllo-1 (o,,
(o,, au 11) åe = 9p(0t- (:0 at-K:t-1),
where Ti-Kit samples is:
(or (Ot-Kit: at Ot-Kit1); or st + ne and ne N(0,021) The empirical loss over N
N N
LED
10041 - a: 9 a
To benchmark LED. consider an ideal predictor with access to the true context e. Since 0++1 8141 + 17+1, M1+1 NN o² I), and o, 8t + me. the predictor f(o,, f(o,,a,,e) uses noisy observations.
The loss is:
Billo++1 f(o,, at, 1)3 B[[|s++1 f(o, at, c) 1)2) + B 112)-
For d. dd-dimensional 7441, Bil,, B[] [ d_,²2. The ideal predictor f (Or, az, E[st+ 10:, dz, c]. the expected next state given 02, a, e, minimizes LED by setting f to the best estimate of 8141. yielding E||lo0+1 - f" a,, 1| Billst+1 E(s loo, de, c]]c]]1]] d. where the first term is the variance of $111 and the second is the noise variance. By the universal approximation property, for approximate fr if encodes
Consider the loss function for fixed yp, w:
((ip, W. O Ki ai- K.)) = ||o.-1 1 (oi, ais (o
Kri, 0ix
16

---

with 14| I M (by bounded A, and Læ-Lipschitz in 9: w. The function class F
{(o,, Cj. 02- Kui 0;-K: + f (ot, a, Ip Kui O2- k==-))) has Rademacher complexity
RN (F) v VN
Cz so the loss class L has:
CrL Rx(C) S VN
For i.i.d. samples and a fixed p, w, the generalization error is bounded by:
Bliml LeDl s 2Rr(C) 2C
VN
By McDiarmid's inequality, since changing one sample affects LTD by most 2M at least 1 - [Mohri et al. 2018):
N with probability
B[LFol & LTD + 2C
llo((1/8)
+ M
VN
For fi-mixing data with B(K) Ce- A the deviation from i.d. adds error bounded by Ce AK to the expected loss for sequences spaced by K steps:
2CIL, + MV VN
1lg((1/6)
B[LAD] & LAD
+ AK
N
Set the generalization error to c² =2=000 0()), where ā = =Ce -e-AK to match the f-mixing decay. Solve:
2CFI1
22log(1/8) -Ce < 8C'
+MV
VN
N
The first term requires: 20rL
S C"8 = N2 (CzLz)² (
VN
The second term requires: M / 2108(1/8] N
S C = 2
2M 1/g(1/A) (C"8)²
log(1/8)
62
The third term Ce-AK = 6 S C holds for C" 2
The first term dominates, so N = O (++² suffices.
If E(Lo] < B[[|o0+1 (or, az [ + e², then:
(or, (ot, 01, 3t) f (or, a. c)l22 c=0(v)).
Given the Markov chain e - T- + ăts the entropy difference is:
h(clax) (cc k <
The cMDP's Lipschitz dynamics, [ a) - (s'ls' a)]l < L|lc e/ and 3-mixing ensure p(c[Te-K:e) K:e) is Ly-Lipschitz in Wasserstein distance, small changes in T.- reflect proportional changes in c. Assume there exists mapping 2b Z + C such that the context encoder satisfies:
[T= - cll2 s Loe = O(Vo),
where the encoder's output space. Then:
xx)p(c/7==xl[((cb))l) S LpLac O(V8).
Thus:
h(clst) S h(c)T +c
From Lemma h(c[TTKK)) 5 Ce h()))(((e-------- -())((())((() Since h(()) and vo 2824 for small 8, the bound becomes:
h(cl3r)) < O(Vo),
yielding 8' = =(VV) with N = (++) samples.
17

---

Lemma 8 demonstrates that DreamerV3's RSSM suffers from persistent information bottleneck in cMDPs, with the conditional entropy h(c/h, 2 kh(c) for a constant (0, 1). This result is essential for Theorems 2 and 5 as it establishes baseline limitation in DreamerV3's ability to
retain context information, against which DALI's superior performance context identifiability and bottleneck reduction is compared.
Lemma 8 (Information bottleneck in DreamerV3's RSSM), Consider DreamerV3's RSSM in a cMDP with context ~ PC with bounded density on _ IK and finite entropy h(c) < oc, observation noise + Tht ~ (0,,0 o² Lipschitz dynamics with [[P (iPc(s'ss, a) Per (s'|ss Is, a)llı S Lllc c dl and recurrent state h4 fo fo 2t-1, at-1), where 2t ~ (o(2) hhg 0t), and fo a GRU with finite parameters. Then, there exists K E (0, such that:
h(c[b))) 2 xh(c).
Proof of Lemma g The context determines the cMDP's dynamics Po(s'(s,
ss. a) Since h4 depends on the trajectory T1:t (01:1, 01:1-1) and latent variables 21:--1, we have a Markov chain c - Tie h so p(h,lTr. c) p(h, Ti.e). By the data processing inequality (DPI) (Thomas and Cover 1999
I(c; s I(c; Tht) h(c).
Since 01:1 = $13 + M1.e. and 01:1-1 depend on 01:1-1, we have:
M(cTTa)= = h(clora, 414-1) 1 = h(c|81.. ma) = h(c(s1:)) as 7ha 1 c.
The GRU So. with finite parameters and Lipschitz continuity (due bounded weights and standard GRU operations), produces h, E IRd, where d a fixed dimension. The GRU is trained to maximize the ELBO:
T
Σ
oo) [log (o Tt, ne = KL (2) (p)
which optimizes h4 to predict observations o,, rewards rt, and termination signals ny- The input Ti: includes:
Context information: e affects $1:t through dynamics.
Noise information: Th1 ~ NO, o² I,,,), with de being the dimension the state space S. The
contributes significant entropy, 1((m):)) = d/ (n(2rea7), which grows linearly with t.
State transients: $1z includes dynamic behaviors not directly tied c.
The GRU's finite parameters and fixed dimension limit its capacity, forcing he to compress T1: potentially discarding context information.
Assume h is Gaussian with covariance Σn. where Var(h) S V for = I....,d The variance bound V follows from compact C, S. A, Lipschitz dynamics, and the Lipschitz GRU, ensuring bounded outputs. For a d-dimensional Gaussian channel with output he ~ ~N(,, E,), where 44, depends on c, and noise variance most 2², the mutual information is bounded by (homas and Cover 1999]
I(c; ho) S d In (1 (1+
tr tr(En) o²
S d In
dV
1
o²
The GRU's compression limits I(c; h2) relative to h(c) for all h(() 2 O, as follows. The ELBO's KL-term (gg((z)) oe) (z regularizes ( (hz, or), reducing the information about o (and thus e) in 2, Since he depends on 21.t-1, this limits context information in h₁ The GRU's finite parameters and dimension d further constrain capacity. Let the channel capacity be reduced by a factor E (0, 1), reflecting this compression:
d
2 In (1
(1+ dV
I(c; ha) sa.
0²
where < (e.g.. - with > 0 depending the GRU's architecture).
18

---

To satisfy the lemma, we seek constant a E (0.1) such that h(clhs) 2 kh(c) for all Pe. Since
h(clh+) = h(c) - I(c; he) and I(c; he) I a. in (1 + dV define:
d In B
1+ )z
dV
I(c:h) S aB.
For h(c)) 0, we want h(c) I[c; 2 kh(c). Using I(c; h4) & aB, this holds if h(c) o. 2
aB kh(c), or S 1 h(c))
Thus, define:
oB h(()
if > 0, if)(c) = 0.
K
If h(c)) 0, then hc) & aB < (since < 1), and with & aB h(o) M() kh(c). Since aB 0 and h(c) 0, then
have h(clfh))
(()) I(c; hx) 2 h(c) aB (1
aB h(c)
aB < 1. If h(()) 2 aB, then & 2 0. If h(c) aB, then but hicche) 2 2 xh(c).
h())
so the inequality holds.
h(c) O, then I(e; h4) = 0, so h(c[h+)= = 0. Choosing & = E (0, 1), we have h(clh+) = 0 2 0 - =-h0))((,, and any E (0, 1) would suffice.
Thus, for all pc, there exists E (0, 1) such that h(c(hh)) 2 sh(c)-
Lemma 9 proves that DALI's sequence model, incorporating the context encoder 3t, reduces the information bottleneck compared to DreamerV3, with h(clhe) S k'h(c), where k' = -e-AK AA
D(e-AK/2). This lemma directly supports Theorem 5 by quantifying the improved context retention in DALI's recurrent state, and aids Theorem 2 by reinforcing the necessity of the context encoder for achieving near-optimal information capture.
Lemma (Information bottleneck DALI's sequence model). Consider DALI's sequence model in cMDP with context e Pe with bounded density on C C R and finite entropy h(c) < oo, observation noise o = St + TJ, πe NO I²I), Lipschitz dynamics with [|Pe(s')
(s') (ss a) Per (s) S L|lc c'Il2, context encoding = 9 (0p(0k Ka, (e-Kat1). and recurrent state h, fa ( 2t-1, at-1 113), where 2n ~ qo(zz [hg ot), and Ja is GRU with finite parameters. Suppose that the context encoding satisfies in(cla)) S Ce-AKK ((e)) 8', where 8' =Oee- and Ce- (Lemma) 7 Then,
there exists w 10, such that:
h(c(tt)) S k'h(c)).
where K² = Ce-AK + O(e ((e--K/22) implying reduced information bottleneck compared to Dream-
erV3's K. (Lemma gy.
Proof of Lemma 9 The context determines the cMDP's dynamics Po(s'Is, a). DALI's se- quence model computes the context encoding år (o: (x=k=11) from trajectory Tt-Kat K0 Kx-K; and the recurrent state h4 fo 2t-1, at-1.3t), where
2t-1 ~
1-1) and fo is Lipschitz continuous GRU.
Since h4 f0 21 (1,3 have Markov chain:
3t + he,
conditional on h 21-1, 11-1, affects h; through given these variables. By the DPI, we have I(e; hx) < I(())) M(clhe)) S h(clar).
By Lemma 7 the context encoding satisfies:
h(clāx) h(c(T.-Ka) +
where = Ce-AK so 0(V5) = From Lemma 8-mixing implies:
h(clTT-Ka) S Ce AK h(c).
Ce
19

---

Thus,
Thus, h(clat) S Ce CC-AX((c)) +
Applying the DPI bound, have
h(clhs) S SCe <Ce-Ak +
To satisfy h(c/hr) S K h(e)), consider two cases. If h(c) > O, let 8' = =C"e e-xi2 s0:
h(clhe)) S Ce (Ce-AK C-Ke-
h(c)).
h(c)
Define
+ e-k-2212
= min Ce- min
h(c) 2
Since h(c) = (e(e-A//)), Ce-A + O(e-AK//) when the first term is selected, and S 2 always. hc)) 0, then [[;; h1) 0, h(clht) = 0. Set k² satisfying h(c(he) k² (c))
Compared DreamerV3's (csM))) 2 xh(c), (0, (Lemma 8p. DALI's k' as K xc
indicating a reduced information bottleneck.
Proof of Theorem 5 We prove that DALI's recurrent state 6,Pal achieves higher mutual information with the context than DreamerV3's fiSSMM up an error e(K) (Xk/))))))(().
By Lemma DreamerV3's recurrent state satisfies:
h(c/h, RSSM) 2 sh(c),
K KE (0, 1).
I(e; 11 RSSM) h(()) M[c/hRSSM S (1 k))(c).
By Lemma DALI's recurrent state satisfies:
MclhDALL) & k(h((), ='e + Ole -XK/2).
Thus,
I(c (pALd) = h(c) h(c/hpaLI 2 x())(c).
Compare:
I(c; (ALL) - I(c; ARSSM) 2 (1 1. r()b(c) - (1 - -<))((c) = (K - -x()((c).
Set
<(K) = k'hlc) = (C'E-Ak + +le (Ax13) M(()) = =0 -xk/2))))()))
since h(c) < 00. Thus,
I(c; AAAAA 2 I(c; NSSSM) ((K).
As K + 00, n' + 0, so I(c; hDALL)- + h(c). while I(c; hRSSM) S (1 k))(c) < h(c), demonstrat- ing that the context encoder reduces the bottleneck.
We prove both parts of the theorem comparing DALI's context encoding åt. Dep (Or- Kit, at Ka-1) with DreamerV3's recurrent state h fo 2t-1, 41-1).
Proof of Theorem 2
1. Context identifiability
We show that DALI's context encoding satisfies:
I(c; 36) 2 I(c; h4) e(K), <(K) = O -(e-))(e-))))
where Ti-Kit (Or-Kit, at k:t-1) is the history, and h; is DreamerV3's recurrent state.
By Lemma
DALI's context encoder, trained with the forward dynamics loss:
LED EE[lor-1 S (or, at, 11)
satisfies
h(cl3)) h(clTT-Ka) + +8, 8' =(e -XK13).
20

---

By Lemma 6 the history's s conditional entropy is:
M(cTT-Ka) S C'e h((c).
Combining,
Mc(b))) S C' <C'EeA +8, CCE
Thus,
6 h(clat) & AO'
The mutual information for DALI is
I(c = h(c) h(cl3r) 2 h(c) - (C'e + C²e -AAI1)
By Lemma DreamerV3's recurrent state h has
h(cl/s) 2 kh(c), WE (0,
so that
I(c = h(c) h(c(he) & (1 k))(c).
Compare:
((caa)I((ch))) 2 AK M(cttt" (1-x))(() =
Since xh(c) > 0, set
(K) CCe-AK((c)) + -KK22 10- -KK//)()(()).
the -XK/2 term dominates. Thus,
I(c: 2 I(c;h,) - -((K).
Additionally, Lemma Lemma shows that DALI's recurrent state he (0o- 20- has a reduced
bottleneck:
h(clh,)) S <kh(c), K' C'e -AK O O(eAK(2))
implying,
I(c;h:) 2 (1- x')((c) = h(c) for large K.
This suggests that even DALI's recurrent state retains more context information than DreamerV3's ht reinforcing the necessity of DALI's context encoder år, which achieves near-optimal 3(c;
Sample complexity
We prove that DALI achieves 3(c; (1 8)h(c) with K Ω (lng(16)) and N O (1²) sam- ples of K transitions, while hypothetical DreamerV3 context estimator gy (01.T, 01:T-1) requires
N O (+)) samples of I transitions.
DALI's sample complexity From Lemma] 7
we have
h(cl3))) h(clTe- +8', 8' =0=0(e- AKK/22
From Lemma
6 we have
h(clT-Ka)) S C'ke-AK((c).
Combining, we have
h(clae) s C'e C'e-AK((c)) + c- -X//2_
To achieve I(GG3t) 2 (1 -5)h(c), need
h(cl3)) S öt((c).
Set C'e ('ke-Ak((c) + -AK/2
2 < 8h(c). Since XK/2 dominates, approximate
8 C -2K/2 -X12 Ö O"
h(c)
h(())
Solve:
8
AK
2ln(C/8)
(log(1/5)
6 -A//2
In
K
S
C
2
A
A
21

---

From Lemma 7 achieving 8' = requires
E Bl - f$ (0t, az, Do (p( =²==0 VN
= Ole -X112).
Solve:
(2 0 (0) = N 82
82 (12)
1
0
VN
DreamerV3's sample complexity
From Lemma 8 DreamerV3's recurrent state he has h(c/he) 2 xh(c), so I(c: he) S (1-x)b(c))
where k is constant, preventing Z(c; ht) 2 (1 - 6)h(c) for small 8. To enable comparison, assume DreamerV3 with domain randomization trains hypothetical context estimator (01.T, @1.7-1) to estimate с over N episodes of length T. By Lemma for =T
h(dm)))) S C'e MTM(M)...
Training Iy to minimize Blllg (01.7, 01.r-1) E/ ]12] requires
[[]lg(TTTTT) - B[cmm.r]] v 6 = V/
= (Oe
Solve:
1 o (2) N= VN
1 82
0
0
8²
The total number of transitions is N (T)
Note that DreamerV3's 's standard RSSM cannot achieve I(c; h,) 2 (1-5)h(c) due h(c[h4) 2 kh(c) (Lemma 8). The estimator ge is introduced to enable a fair comparison, acknowledging that this is not standard in DreamerV3. Since K < T typically. DALI is more efficient. Moreover, Lemma indicates that DALI's hú achieves I(c; h2) 2 (1 xc)),,) with k' O suggesting that DALI's sequence model could approach the desired information level with sufficient K, unlike DreamerV3's ht.
A.1 DALI and Dreamer-based Baselines
A.1.1 World Models
DreamerV3 with Domain Randomization (context-unaware baseline)
Tobbin al. 2017 Hafer
et al. 2025 :
Sequence model: h fo 11, 21-1 ax-1)
Encoder:
~ g he, ot)
Dynamics predictor:
2. ~ ( h2)
Reward predictor:
F. ~ (f: he, 2t)
Continue predictor:
âr po (fee h4 21)
Decoder:
6e ~ po h4, zt).
DALI-S (DALI with Shallow Integration):
Context encoder:
3r = 9o (or- Kit, ae- Kit 1).
Sequence model:
h, fa 21 a1-1);
Encoder:
~ qo (zx he, Or, 3c),
Dynamics predictor:
2. ~ po ht);
Reward predictor:
ir ~ po he, zx);
Continue predictor:
ter ~ ( he, 2z),
Decoder:
o, ~ Po
h4g 2,)
22

---

DALI-D (DALI with Deep Integration):
Context encoder:
åt = 9 kit, Ot- ktt-1).
Sequence model:
hr = fo(ho -,,, 11-30)
Encoder:
21 90 (2 Ot),
Dynamics predictor:
& ~ pa p he):
Reward predictor:
Pe (fe h4, 2t, at);
Continue predictor:
fg ~ po
huy 2ts 3t);
Decoder:
or po
heg (1,
cRSSM-S (ground-trutt -ruth context-aware baseline)
Prasanna et al. 2024
Sequence model:
h so(h4-1, 21 1,00 1-11)
Encoder:
2t (0(2) h4, oz, c),
Dynamics predictor:
Se po(Ee hhe,
Reward predictor:
P. pf( h2 zz),
Continue predictor:
ùe ( he: 2t),
Decoder:
ōR ~ po(oe he, 22).
cRSSM-D (ground-truth context-aware baseline)
Prasanna et 2024
Sequence model: h, fo - 21-1,
Encoder:
2t ~ (0( (zt
hr, 0z)
Dynamics predictor:
2. (2t h₂)
(2t
hes 2t, c) Continue predictor: tn po
Reward predictor: Ft ~ (f
h4 Zt, c) Decoder: 6, ~ po(ô,
he, 24, c).
The baselines cRSSM-S and cRSSM-D correspond to cencat-conter and CRSSM, respectively, in [Prasaana et al. 2024
A.1.2 Actor-Critic Models
DreamerV3 with Domain Randomization, DALI-S, and eRSSM-S:
Actor:
a, (x)))) x Bur(-10=) E H-! reter
Critic:
DALI-D:
Actor:
a (o(ar(srr3)
Critic:
(w(s),)) Ex(lered E H
retr
CRSSM-D:
Actor:
To(arls,,c) up c) Ex((,,,)) H-t ret.
Critic:
23

---

B
Algorithms
Algorithm 1: DALI-S (Shallow Integration with Forward Dynamics Loss)
Model components
Hyper parameters
Context encoder
9. (o ae- (o-K:t-1)
Seed episodes
S
One-step model
1 (0t, Uts utsāt)
Collect interval
C B
Deterministic State
fo(h, 2 1,00-111 Qt-1)
Batch size
Stochastic state
Po he))
Sequence length
L
Encoder
90 h4, 0t:
Imagination horizon
H
Reward
po (f: he, 2t)
Learning rate
a
Continue
ftt
he, Zt)
Learning rate g
Og
Decoder
pa (ō
hes 2x)
Episode length
T
Actor
Ta (te st))
Context history
K
Critic
(5))
Initialize dataset with S random seed episodes.
Initialize neural network parameters 0, d, v. randomly.
while not converged do for update step = I....,Cdd C'do
Dynamics learning (World Model 0) Draw B sequences {os, ae, rt)) D.
Initialize deterministic state he 0. for all steps t from sequences batch D do
if t < K then
Pad and da.c-1 with
length K.
Compute context j. 9x (ot- Kit, ac K de--kch
12
Encode observation 2; ~ ga(2: he, 0z.
13
Compute deterministic state h, = fo ac-1)
Set latent state 8, 4 [=t, he]
Update 8 using representation learning (decoder, reward, continue).
// Context representation learning (s)
16
Draw B data sequences {(oi, m)) ] for each sample (01.k-1, (1:k) ~ do
D
18
Compute 31.K = (Ork, (1:K-1)
Update & f 4 0g v ΣK 0+ fo (o,, (or, 31
// Behavior learning (Actor/Critic v)) Infer context: jt + g Kit Ot- (t-aa))) // 8, frozen
20
Imagine trajectories [(*+, a)) (++ from each oe (with fixed ás once it is provided to the initial observation encoding)
21
Predict rewards Erpo(r> sy)] and values (sr)).
23
Compute value estimates (s(sr).
Update aV * Vx(*7)-
Update +- 3 V * (p((s) - va(s+))??
25
// Environment interaction Or env. (nv.rese
26
27
for time step 1= 1....,'' do
28
if < K then
Pad and 01:-1 with zeros length K.
Infer context online: jt De (tt
31
Compute 8, ~ (s: (t-1,00-1000).
32
Compute a ~ (a st).
33
0141 <- env. (..)
Add experience to dataset D +- D U ((o,, ae, Te)
34
24

---

Algorithm 2: DALI-S-X (Shallow Integration with Cross-modal Regularization)
Initialize dataset D with S random seed episodes
Initialize parameters 0, $, v,, se, Wa, W, randomly
while not converged do for update step
1 to C' do
// Dynamics learning (World Model 0)
Draw B sequences (o Initialize hu f- for all steps in batch do if < K then
Pad 0max(1.1-- K).t and 0max(1,1- K)=-- with
Compute context jr f- ge (or Kit, ac- -.t
11
Encode observation 2 ~ ga(2; h4 0t,
Compute deterministic state h. 1, ar-1)
Set latent state 8 - [2x,
Update 9 using representation learning (decoder, reward, continue).
// Context representation learning (p, W., W2) Draw trajectory segments {(ot-Kee, (ttt-KK--0KKK:t000+ Initialize losses LED + 0, Lexes f for each segment (0t at- KitKit, 0c+1) do Initialize h for 1 K to do
15
16
17
if then
Pad 0max(1,7- and Omax(1 with zeros.
= 90 Kir, 0r K:r-1)
(zr hy detach(), 0r))))
h fe(h- detachh),,
25
& = W2(3:)
3. d2ttcch)
27
LND = | f (or, at)
28
Lasss = detach)) - & lla: - 3,
LAD B TTD. Tecruss 11 Tecrss
30
Update 6,, W.,, W, (gradient descent) using Luua = LAD + devosa
Behavior learning (Actor/Critic v))
Draw B sequences {(o,, ae rt)liin Initialize ho +- for all steps in batch do if < K then
32
33
Pad ()) and (1 with
Infer Je + go(0t-Kit, Ot-K - detach))
// 8, ip frozen
37
Encode initial observation ~ [0( he, 3t,
Set initial latent state <- [z+, hie]
for imagination step t
do
Sample action a, ~ (a; 8r)
41
Sample prior state & (a, h4)
Compute next deterministic state 4 fo
Set latent state $r+1 +- Ar+1)
44
Compute value estimates Vx(8,) using TD(A)
Update Update +- + v Σ
V Σ VA
46
& -ll Vx
// Environment interaction env.reset()
47
for time step 1,.... T do
49
< then
59
Pad 0max(1,t-Ky. and 0max(1,t-K)r-1 with zeros.
51
Infer context online: jt ge (Oi-Ka,,
Compute &; (a 5t-1,0c-1,00
Compute a, (a 8t)
54
Ft, 0++1 f- env. sttp(a;)
Add experience dataset + DU DUt( as, re)
25

---

Algorithm 3: DALI-D (Deep Integration with Forward Dynamics Loss)
Model components
Hyper parameters
Context encoder
9e (ot- Kit; (tt-Kit-1
Seed episodes
S
One-step model
[ (0t, 0z,30
Collect interval
C
Deterministic State
Jo(1 (4-11 1, 2t-1, 1,
Batch size
B
Stochastic State
Po he)
Sequence length
L
Encoder
(0(2t
124 Oc)
Imagination horizon
H
Reward
Pf hes Zo: 3t)
Learning rate
Continue
Po f4
her 3t)
Learning rate !
Og
Decoder
po Ô
hes 2t,
Episode length
T
Actor
[a (5,3t)
Context history
K
Critic
Initialize dataset ID' with S random seed episodes.
Initialize neural network parameters 0, Φ, f, randomly.
while not converged do for update step c == l....., Codo
// Dynamics learning (World Model 0)
Draw B sequences {o,, az, 7=) D.
Initialize deterministic state ha 0. for all steps from sequences batch do
if < K then
Pad and di:t-1 with zeros length K. Compute context jr 9y Kit, Ottt, R -
Encode observation 2, (2 h4,, 01)
13
Compute deterministic state hu fo (o(hr-1, 21 1;
Set latent state s, + [ [z]
Update 9 using representation learning (decoder, reward, continue).
// Context representation learning (p)
Draw B data sequences {(or, 44) ( K+1 ~ 'D
16
for each sample (01:K-1, 01:K) do
Compute 31:K = 90 01:K-1)
Update ip 6 6- ----on 0g v Ein 0c+1 (or, at, (o., st)
// Behavior learning (Actor/Critic v) Infer context: jt <- 9h (4 // 0, frozen
20
21
Imagine trajectories a7 from each 8 (with fixed āt)
22
Predict rewards E[po(r= [r,3(]] and values Ue 3t).
23
Compute value estimates (8(5..5.).
Update & a. E++/ (x(8+,3)).
Update + 1 o E (p((5)3t) - Vx(o+,3t)|
// Environment interaction env. reweteee(()
26
for time step 1,..., T do
28
< then
29
Pad 01.: and d=a-1 with zeros to length K.
30
Infer context online: jt gu (or- Kit, Ct- k(--)..
Compute 9; pa(st 8t-1,
Compute a, [o(a; St.3e)-
33
Ft, or+1 <- env. step(aa).
Add experience dataset +- DU f(or, r.)
26

---

Algorithm 4: DALI-D-X (Deep Integration with Cross-modal Regularization)
Initialize dataset DD with S random seed episodes
Initialize parameters 0, $, sp, Wa, W, randomly
while not converged do for update step
1 to C' do
// Dynamics learning (World Model 0)
Draw B sequences (o Initialize hu f- for all steps in batch do if < K then
Pad 0max K).t and "1ax((111-):1---1 with
Compute context jr f- ge (or Kit, Rt- dee
11
Encode observation 2 ~ ga(2; h2, 0i)
Compute deterministic state h. 1, 1:-1;
Set latent state 8; - [2x,
Update 9 using representation learning (decoder, reward, continue).
// Context representation learning (p, W., W2) Draw trajectory segments Initialize losses LEp + 0, Lesss f- for each segment (0t-Kit: Qt-Kit; 0c+1) do Initialize hi-k
15
16
17
for 7 =t K to do
if K then
Pad 0max((17--Xx Kyr and dmox(1 K)r-1 with zeros.
21
= go(0r- Kir, (r:k:-1))
~ g0(2r hy d),
23
h So (hr (o(h..den detach (), Zr dww.dettahhhheem. ar, dr...dd
(a+) 3. ( detacc(t)))))
Lm ||01 1 (or, ai, )
Taros = || detach () - 2,112 llae - 3.112
28
LED +- B LAD, Lecoas t- + Lenos
Update 4, W., W, (gradient descent) using Loval LoD A Lowos
// Behavior learning (Actor/Critic 0,²))
30
Draw B sequences {(o,, ae, re) Initialize ho +0 for all steps in batch do t < K then
31
0max((1:--]). and 4max (1,1 KJat1 with zeros.
35
Infer jr < ge Kit, at K detachh)
// 0, frozen
(a(z) 0:) Set initial latent state St [2t, he] for imagination step T = to t + do
Encode initial observation
37
3s
Sample action a, To (a 8r)
Sample prior state & ~ (e- h-)
Compute next deterministic state h++1 So a7.3)) Set latent state 8++1 6- [1+,/
Compute value estimates (*(5+. using TD(A)
44
Update & f Ø + V. Σ Vx
Update +- 10 V. Σ fll Vx
// Environment interaction on <- env, rever() for time step t 1,..., do
46
47
48
if < K then
Pad 0mar(1,1--)) K): and 0maa(1, K):-- with zeros.
Infer context online: jt Te
51
Compute s, ~ Po (5t (5-1,4(-1,0,)
52
Compute a ~ (a st,))).
53
Ft,
+ env w.v. sttp(aa).
54
Add experience to datasel D <- DU To)
27

---

C Experimental Setup and Implementation Details
We evaluate DALI's zero-shot generalization on contextualized DMC Ball-in-Cup (gravity, string
length) and Walker Walk (gravity, actuator strength) tasks [Tassa et al. 2018 These environments feature context parameters with distinct training and evaluation ranges, presenting unique dynamics and generalization challenges.
C.1 Training and Evaluation
Our setup follows the CARL benchmark (Benjamins et al. 2023]. which defines default values for the two context dimensions per environment. To assess both interpolation and extrapolation Nout-of-dittribution, OOD) capabilities, we define two distinct context ranges: one for training and interpolation evaluation, and another for OOD evaluation. See Table [ both Featurized- and Pixel-based observation modalities for each environment.
Additionally, we consider
We consider two training scenarios. In the single context variation setting, each context dimension is varied independently by uniformly sampling 100 values within its respective training range, while the other dimension held fixed at default value. In the dual context variation setting, both context dimensions are varied jointly by uniformly sampling 100 context pairs from the Cartesian product of their respective training ranges.
Following Kirk et al. 2023), we evaluate our agents under three generalization regimes. The Interpolate regime tests performance unseen context values sampled within the training range, while the Extrapolate regime assesses generalization to OOD contexts beyond the training range. The Mixed regime evaluates generalization when both context dimensions vary: one is sampled from the extrapolation range and the other from the interpolation range. Results in the Interpolate and Extrapolate regimes are aggregated over cases where one or both context dimensions vary.
C.2 Environments
DMC Ball-in-Cup. This task involves an agent controlling a cup attached to a ball by string, aiming to swing the ball into the cup, with context parameters gravity and string length. The ball's pendulum-like motion exhibits complex dynamics due to gravity's 's influence on oscillation frequency and string length's s effect on swing amplitude. Extreme gravity values, such as 0.98 or 19.6, significantly alter the ball's trajectory, while short long strings (e.g., 0.03 or 0.6) change the swing's frequency and reach, creating complex interactions between these parameters. The interaction amplifies nonlinear effects, and poses significant generalization challenges, the context encoder must adapt to substantial shifts in motion patterns, particularly in the OOD regime.
DMC Walker Walk. This task involves planar bipedal robot walking forward, with context parameters gravity and actuator strength. Actuator strength scales torque applied to joint motors (e.g., hips, knees), modulating walking speed and stability, while gravity affects balance and ground reaction forces. Despite simple torque scaling. locomotion dynamics are intricate, driven by multi- joint coordination and contact interactions. Generalization is less challenging than in Ball-in-Cup, as actuator strength evaluation range (0.1 to 0.5 and 1.5 to 2.0) closer to training (0.5 to 1.5), and its torque effects are predictable. However, extreme gravity values test the context encoder's ability to maintain stable, coordinated walking in the OOD regime.
Table 1: Context ranges for considered environments.
Environment
Context
Training/ Interpolation
Extrapolation
DMC ball in tn_ccp-catch-vv
gravity string length
[4.9, 14.7] [0.15, 0.45]
0.98, 4.9) U (14.7, 19.6 0..03 0.15) U (0.45, 0.6
DMC walker-walk-v0
gravity
[4.9, 14.7] [0.5, 1.5]
0.98, 4.9) U (14.7, 19.6 0..1 (1.5, 2.0]
actuator strength
28

---

C.3
Architectural details
Figure shows high-level illustration of the DALI architecture.
P K
at
Gt
ft
9 K
81-
8t
Enzodar
Docoder
Encoder
Docoder
Encoder
Decoder
(0--1,3-11
δt-1
0b, 3t
Ôt
a+1 1
0t-1
3t-1
f
LFD(F) = EE f (o--1, 41-1, 0,
9
Forward Dynamics Loss
(0i-K; 0/-K3 ................................ 2/-2, 04-1)
gradient flow
History
Figure 3: DALI architecture overview.
Context Encoder. The context encoder de employs standard transformer encoder block [Vaswani et al. 2017) to process sequence of K transitions, (ot- Kit, dz- Ka-1), and produce the context representation 3r E $ The input sequence is fed directly into a dense layer with 256 units, bypassing traditional embedding layer. This is followed by layer normalization and a single-head self-attention block with a skip connection from the dense layer's s output. The attention output is combined with the skip connection via residual connection. followed by second layer normalization. A two-layer MLP, each layer with 256 units, processes this output, with another residual connection combining the MLP output with the second layer norm's output. A final dense layer projects the 356-dimensional intermediate representation to je E R All hidden layers, including the MLP and attention head, 256 units and SiLU activation functions, consistent with DreamerV3 (Hafner et al.) (025). No masking is applied in the self-attention mechanism, as the input sequence is complete, enabling full contextual processing generate ăt-
Forward Model. The forward model f for dynamics alignment a two-layer MI.P with 128 units and SiLU activations.
We E 32xx and W,, 128x32 in learnable parameter matrices that map between the context and state spaces.
DreamerV3. We adopt the small DreamerV3 variant with hyperparameters from Hafner et al. 2025 following the setup of Prasanna et al. 2024 to ensure fair and reproducible comparison with their cRSSM-S/D baselines.
DALI adds only about 4% parameter overhead (e.g., Dreamer-DR: 15.73M vs. DALI-S: 16.45M) while consistently improving performance with minimal additional complexity.

---

C.4
Computational Setup and Resource Requirements
Training and evaluation of the baselines and our DALI approaches were conducted on NVIDIA A100 GPUs with 80GB of VRAM and Intel Xeon Platinum 8352V CPUs. Typically, our setup provides access to GPUs average, with up to 4 GPUs available in the best-case scenario. Parallelization differs between modalities: featurized experiments can run 10 seeds in parallel, while pixel-based experiments are limited to 4 seeds due to higher memory requirements.
Our complete experimental setup includes 10 seeds across algorithm variants, environments, and 2 modalities, evaluated over context settings (single_0) single_1 1, double), yielding total of 600 individual training runs. Table summarizes the computational requirements, where reported GPU hours account for 10 parallel processes in the featurized case and 4 in the pixel-based case. Note that training times per trial are similar for both variants, as the bottleneck lies in environment interaction rather than network size.
In a purely sequential execution, the total computational cost would reach approximately 24,000 GPU hours. Leveraging parallel execution significantly reduces wall-clock time: in the best-case scenario with 4 GPUs, total wall-clock time approximately 1,051 GPU hours (=:44 days). In the more typical worst-case scenario with 2 GPUs, wall-clock time increases to approximately 2,1 2,101 GPU hours (=:88 days).
GPU Hrs
Runs
Wall-Clock (GPU Hrs)
Task
Modality
per Trial
Total
2 GPUs
4 GPUs
Featurized
20-26
150
195
98
Ball Cup
Pixel-based
20-26
150
488
244
Featurized
46-54
150
405
203
Walker
Pixel-based
46-54
150
1,013
506
Total Featurized
300
600
301
Total Pixel-based
300
1,501
750
Grand Total
600
2,101
1,051
Table 2: Computational requirements for the full experimental evaluation. Each task is run across 5 algorithm variants with 10 seeds each in both featurized and pixel-based modalities. Wall-clock times assume optimal GPU utilization given the stated parallel capacities.
C.5 Additional Experiments
Figure 4 shows the learning curves for DMC Ball-in-Cup and Walker Walk tasks under Featurized and Pixel-based observation modalities.
DALI's context encoder adds only about 4% more parameters to DreamerV3-smal while scaling effectively to nigh-dimensional tasks. We evaluate it on DMC Quadruped Walk (56D observations, 12D actions), larger environment than Walker Walk (24D, 6D). The 56D observation space em- phasizes the challenge of torque-sent senuiessensiive locomotion, where increased coordination and higher number of actuators (12 vs. demand more robust context inference. Trained for 600K timesteps and contextualized with gravity and actuator strength (same ranges Walker Walk), DALI attains Featurized Extrapolation IQMs of 0.326 + 0.043 (DALI-S), and 0.389 + 0.027 (DALI-D-X), yielding up to 76.8% improvement over the context-unaware baseline 0.220 + 0.023 (Dreamer-DR), and up 50.8 50.8% improvement over the context-awar baselines (0.258 0.028, cRSSM-D; 0.317 0.031, cRSSM-S).
30

---

Ball-in-Cup
DMC Ball Cup- Featurized
DMC Ball Pixel
100
Thaining Roturn
250
Thaining Rotum
250
300
300
150
DALES
DALLS
DALL--22
DALL---
50
Steps
Steps
(a) Featurized
(b) Pixel
Walker Walk
DMC Walker Walk Featurized
DMC Walker Walk Pixel
300
300
250
Training Return
Training Return
200
1.00
100
DALL-S-2
DALLS
100
DALL-S-Z
DALLS
JRSSSS
aRSSM55
Environment Steps
Steps
(c) Featurized
(d) Pixel
Figure 4: Learning curves for DMC Ball-in-Cup and Walker Walk tasks under Featurized and Pixel-based observation modalitics. Results show mean episode returns with 25th-75th percentile confidence intervals.
D Supplementary Experiments on Physically Consistent Counterfactuals
We provide additional details supporting Section 6.3 First, we validate the statistical significance
of the influence of latent dimensions on counterfactuals using AUC-based rankings (Section D.I). Next, we analyze actuator-driven recovery dynamics in the DMC Walker Walk task, analogous the counterfactual analysis conducted for the Ball-in-Cup task in Section 6.3.
(Section D.2)
D.1 AUC Rankings and Significance of Counterfactual Perturbations
To identify latent dimensions that systematically alter imagined dynamics, train binary classifier to distinguish trajectories generated under perturbed (7"0)) and original (7) context representa- tions. For each dimension āj. we generate N 2500 trajectory pairs, perturbing 3; by = o(3)) to induce counterfactuals (see Section 6.3)
We implement an ensemble classifier comprising a support vector machine, an MLP, and AdaBoost classifier, trained via stratified 5-fold voss-validation. Inputs trajectory sequences ( (label 0) and ( (label 1). Predictions aggregated across folds to compute AUC. To ensure robustness, we compute 95% bootstrap confidence intervals for each 3n 's AUC using 500 resamples. Permutation tests (1000 iterations) validate statistical significance between top-ranked dimensions by comparing observed AUC gaps to a shuffled null distribution.
For the Ball-in-Cup task, results are shown in Figure where higher AUCs indicate dimensions whose perturbations induce physically distinguishable dynamics (e.g., gravity or string length changes). Confidence intervals highlight stability across bootstrap samples.
31

---

Mear
0.4
0.0
-9-
49
Figure 5: AUC + 95% CI per context dimension šj for the Ball-in-Cup task.
D.2 Actuator-Driven Recovery and Locomotion in Latent Imagination
Analogous to the Ball-in-Cup analysis (Section 6.3.1p we evaluate DALI's ability to generate physically consistent counterfactuals the DMC Walker Walk task contextualized with gravity and actuator strength. We show that perturbations to the top-ranked latent dimension às in DALI-S, identified via AUC-based rankings (see Section D.I), influence the Walker's capacity to recover from destabilizing falls and sustain locomotion, indicating that 33 is mechanistically aligned with actuator strength. Imagined rollouts are initialized with a real observation and then rolled out in latent imagination. Unlike in the Ball-in-Cup analysis, we retain the original policy actions during imagination to examine how actuator strength modulates the Walker's 's recovery and locomotion behavior. For the Featurized modality, tracked low-level state variables (e.g., torso height) at each timestep. while in the Pixel modality, we captured rendered environment frames at 5-step intervals 64-step imagination horizon.
Pixel Modality. Figure ba contrasts the original (top) and counterfactual (middle) trajectories. Both agents initially fall (frames 0-20), but the perturbed trajectory (counterfactual) exhibits enhanced actuator strength, enabling the Walker to stand from a challenging pose (left leg extended, right leg bent in frames 30-45) and locomote forward (frames 50-64). The original trajectory fails to recover, collapsing after the fall. Pixel-wise differences (bottom) highlight kinematic deviations, particularly in leg articulation and torso alignment.
Featurized Modality. Figure (bb shows torso height dynamics. While both trajectories share similar initial descent (time steps 0-20), the counterfactual (orange) achieves sustained elevation (after time step 20), reflecting successful recovery and locomotion. This aligns with the pixel- based evidence of increased actuator strength, confirming às 's role in encoding torque dynamics critical for stability.
32

---

5
10
20
25
30
35
40
45
50
55
60
(a) Pixel Reconstructions
Original
1.2
Counterfactual
0.8
0.4
(b) Torso Z Position
Figure 6: (a) (Pixel Modality) Counterfactual Trajectories in Pixel Space: Top: Original imagined trajectory of the Walker Walk under default gravity and actuator strength. Middle: Perturbed trajectory after adding noise to the top-ranked latent dimension 33- Original trajectory shows the Walker failing to recover after fall. The perturbed trajectory with enhanced actuator strength enables recovery and forward locomotion. Bottom: Pixel-wise differences = lo. - o%)). The counterfactual agent refines leg kinematics (frames 30-45) and stabilizes upright (frames 50-64). (b) (Featurized Modality) Torso Z--Position Under Latent Perturbation: Comparison of original (blue) and counterfactual (orange) torso heights. The perturbed 33 trajectory maintains higher elevation after time step 20, reflecting successful stabilization and locomotion.
33

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: The claims presented the abstract and introduction are fully supported by theoretical and empirical findings, including comprehensive analysis.
2. Limitations
Question:
Does the paper discuss the limitations the work performed by the authors?
Answer: [Yes]
Justification: Limitations are discussed in the corresponding section of this paper.
3. Theory assumptions and proof's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
Justification: We provide a comprehensive proof of the necessity and benefit of our approach in Section 5
and Appendix A
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Answer: [Yes]
Justification:
This paper provides the complete code and instructions for running the experi- ments reproducible manner. Python dependencies, their versions, and all random seeds used are clearly specified.
5. Open access to data and code
Question: Does the paper provide open access the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: The full source code is provided with this submission along with all relevant details required to reproduce the presented experiments.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Answer: [Yes]
Justification: Yes, provide all details for training the agent in-distribution and testing zero-shot generalization.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: For this purpose, we follow established methods for measuring the agent's performance [Agarwal et al. 2021].
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
34

---

Answer: [Yes]
Justification: DALI adds only a small context encoder module to DreamerV3. As reported by [Hainer al. 2025]. DreamerV3 can be trained on single GPU, and our extension has the same requirement. Each trial with single random seed takes roughly 1-2 days, depending on the observation modality.
9. Code of ethies
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics
Answer: [Yes]
Justification: No disagreement amoung the team.
Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
[NA]
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: [NA]
Licenses existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
This work builds upon and extends the codebase of Prasanna et [2024] which is properly credited. All associated licenses and terms of use from the original work are respected.
Justification:
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Crowdsourcing and research with human subjects
Question: crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: [NA]
15. Institutional review board (IRB) approvals equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
16. Declaration of LLM usage
35

---

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology, scientific rigorousness, or originality of the research, declaration is not required.
Answer: [NA]
Justification: [NA]
36


Reviews:
[{"review_id": "ExuMeOjhlw", "reviewer": "Reviewer_i6n9", "summary": "Reinforcement learning algorithms like DreamerV3 learn a world model to train an agent entirely in (latent) imagination, eliminating the need for direct environment interaction during training. However, these methods typically suffer from either a lack of imagination accuracy, exploration capabilities, or efficient context inference and robust generalization. The paper proposes DALI, a framework based on the Dreamer architecture that that seeks to infer latent context representations from agent-environment interactions by training a self-supervised encoder for prediction. Some numerical experiments are performed for cMDP benchmarks against existing context-aware and context-unaware baseline solutions.", "questions": "1. What are the main insights? Under which conditions is the proposed solution better than existing context-aware baselines and when it is worse? What are the reasons for the performance differences?\n2. Does it scale to larger problems (dimensions of the state space or action space)? How many training will be required? Please discuss.\n3. What are the required runtimes for the performed experiments?\n4. How problematic is the tuning of the hyperparameters used?\n5. The discussion of related work appears a short. Regarding the aim of zero-shot generalization when context is latent and must be inferred: Which components of the model already exist and which ones are novel? Please clarify.\n6. Can you provide more benchmark evaluations? How do results relate to cRSSM-S?", "limitations": "yes", "rating": 3, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 3, "significance": 3, "originality": 2, "strengths_and_weaknesses": "Strengths:\n * The investigated problem is timely and relevant. The proposed model addressing contextual MDPs and effective zero-shot learning is interesting.\n * The paper provides some analytical results for the context encoder, cf. Theorem 1, page 5.\n\nWeaknesses:\n * Only few benchmark experiments (Ball-in-Cup and Walker) against existing state-of-the-art models (Dreamer, cRSSM) are provided. Due to the lack of extensive evaluations the results of the experiments are not yet fully convincing. The general applicability of the framework and its performance is not demonstrated and hence, remains unclear for other settings.\n * Limitations of the approach could be better discussed. It is not clear when or for which tasks the framework performs particularly good / not good.", "was_revised": false, "final_justification": "", "rebuttal": "Dear Reviewer i6n9,\n\nWe appreciate your thoughtful engagement with our work and your recognition of its timeliness and theoretical contributions. Below, we address each of your concerns comprehensively, providing new insights and empirical evidence to demonstrate DALI's significance and broader applicability.\n\n### (C1)  Main insights and performance conditions\n\n(a) **Main insights:**\n1. **Learned context inference is more broadly applicable.** \nDALI's self-supervised context encoder infers contexts directly from observation-action sequences, enabling robust zero-shot generalization when ground-truth context is unavailable, unlike cRSSM’s reliance on known contexts. This makes our method particularly valuable in realistic settings where such ground-truth context is unavailable or costly to measure.\n2. **Relying on ground-truth context may reduce OOD adaptability.**\nContext-aware methods may overfit to training context distributions and fail in extreme OOD settings (e.g., gravity values\nof 0.98 or 19.6 far from the training range of [4.9, 14.7]). For instance, in Ball-in-Cup Featurized Extrap., our DALI‑S‑$\\chi$ achieves IQM 0.3720, significantly outperforming context-aware baselines, cRSSM‑S (0.2270) and cRSSM‑D (0.2780). This suggests that when the context is inferred from data (as in DALI) rather than given (as in cRSSM), the model can generalize better to unseen contexts when it learns the underlying physical dynamics.\n3. **Cross-modal regularization can boost performance in nonlinear, partially observable tasks.**\nAligning the context representation with the world-model's latent $z_t$ via $L_{\\text{cross}}$ (Eq. 2) can enhance inference under noisy visual input and nonlinear dynamics. For example, in Ball-in-Cup Pixel Extrap., DALI‑S‑$\\chi$ surpasses Dreamer‑DR by up to 96.4%, showing the strength of this alignment.\n4. **Task-specific regularization benefits.**\nWhile cross-modal regularization can enhance context inference in extreme OOD settings with complex physical dynamics (e.g., Ball-in-Cup), DALI-S outperforms DALI-S-$\\chi$ in Walker Walk across most regimes and modalities: By simply optimizing for next-step predictions, DALI-S effectively handles Walker’s torque-dependent locomotion, particularly in the Pixel modality where visual noise may amplify $L_{\\text{cross}}$’s complexity in DALI-S-$\\chi$. This highlights the need for tailored regularization.\n5. **Integration strategy matters: Shallow context propagation is more robust than Deep conditioning.** \nShallow integration, as implemented in DALI-S and DALI-S-$\\chi$, mostly outperforms deep integration (DALI-D, DALI-D-$\\chi$), leading to its exclusive use in our reported results. Shallow Integration incorporates the inferred context solely in the world model’s encoder, allowing context information to propagate indirectly to the recurrent state $h_t$ through recurrence. This avoids overfitting to potentially noisy $\\mathfrak{z}_t$ estimates, which is critical in OOD settings.\n6. **Performance varies with environment complexity and modality; No single method dominates everywhere.**\nDALI-S-$\\chi$ excels in Ball-in-Cup, particularly under OOD generalization, while context-aware baselines (e.g., cRSSM-D, cRSSM-S) can outperform in Pixel settings of Walker, where visual complexity is high and ground-truth context might aid robust feature extraction. \n\n(b) **When is the proposed solution better (or worse) than context-aware baselines?**\n\nFrom our empirical findings, we observe the following broad trends: DALI is potentially **better** than context-aware baselines when (i) the ground-truth context is unavailable, uncertain, or overfits, esp. under Extrap., (ii) the task exhibits complex nonlinear dynamics (e.g., coupled gravity-string interactions in Ball-in-Cup). Conversely, it can be **worse** when (i) ground-truth context is reliable and directly encodes task variation, and (ii) visual dynamics dominate and precise reconstruction of features benefits from explicit context conditioning.\n\n\n(c) **What explains the performance differences?**\n\n1. **Overfitting of ground-truth context** reduces adaptability in OOD settings. This is true for context-aware models (cRSSM-S/D), especially in Extrap..\n2. **Cross-modal regularization (DALI-S-$\\chi$)** improves generalization in visually noisy, nonlinear environments, by stabilizing context inference.\n3. **Shallow integration regularizes latent context use**, reducing over-reliance on $\\mathfrak{z}_t$ and making models more robust to context inference errors.\n4. **Task structure and modality influence what works best**: Walker’s actuator scaling is quasi-linear and less sensitive to context modeling, while Ball-in-Cup’s nonlinear pendulum dynamics demand better learned generalization.\n\n\n\n### (C2) Scalability and training requirements\n\nDALI's context encoder adds minimal overhead (< 5\\% parameters to DreamerV3-small) and scales to high-dimensional tasks. We tested DALI on **Quadruped Walk** (56D observation, 12D action), a larger environment than Walker (24D observation, 6D action). Trained for 500K timesteps and contextualized with gravity and actuator strength (same range as DMC Walker), preliminary IQMs for the Featurized modality are 0.139 (Dreamer-DR) and 0.22 (DALI-S-$\\chi$), a **58.3%** improvement. The 56D observation space highlights the challenge of torque-sensitive locomotion with increased coordination complexity and higher degrees of freedom (12 vs. 6 in Walker), demanding more robust context inference. In general, training time scales approximately linearly with the dimensionality of the state and action spaces. Pixel tasks take 1.5–2x longer due to CNN encoding. We will add a summary in the App. emphasizing DALI’s scalability to complex cMDPs.\n\n### (C3) Runtime efficiency\n\nDALI's evaluation used NVIDIA A100 GPUs (80GB) and Intel Xeon CPUs, with 2–4 GPUs available. It comprised 600 runs: 10 seeds × 5 algorithm variants × 2 environments (Ball-in-Cup, Walker) × 2 modalities (featurized, pixel-based) × 3 context settings. Featurized runs paralleled 10 seeds; pixel-based, 4 seeds. Per-trial GPU hours: 20–26 (Ball-in-Cup), 46–54 (Walker). **Total cost: ~24,000 GPU hours** sequentially, reduced to 2,101 hours (88 days) with 2 GPUs or 1,051 hours (44 days) with 4 GPUs, highlighting the extensive computational complexity. We will add a summary of the runtimes in the App..\n\n\n### (C4) Hyperparameter tuning\n\nPlease refer to Response (C3) to Reviewer ogTR, where we address the same question.\n\n### (C5) Related work\n\nOur related work centers on contextual RL for zero-shot generalization and model-based methods (e.g., DreamerV3, TD-MPC2, and variants). We will revise Section 2 to incorporate the additional citations (suggested by Reviewer mG63).\n\n### (C6) Existing and Novel Components\n\n- **DALI’s Novelty**:\n  - **Self-supervised context encoder**: Infers $\\mathfrak{z}_t$ from $K$-step windows without ground-truth context (unlike cRSSM) (Sec. 4.2).\n  - **Dynamics-aligned training**: Forward dynamics loss (Eq. 1) and cross-modal alignment (Eq. 2).\n  - **Theoretical foundation**: Sample efficiency of context inference (Theorem 1).\n  - **Integration strategies**: Shallow/deep conditioning (Sec. 4.3).\n\n- **Prior Work**:\n  - **DreamerV3** [Hafner et al. 2023]: World model backbone (RSSM, actor-critic).\n  - **cRSSM** [Prasanna et al. 2024]: Requires explicit context; no latent inference.\n\n\n\n\n### (C7) Limited evaluation; Additional benchmark evaluations needed\n\nWe agree that broader benchmarks strengthen our claims. Alongside Ball-in-Cup and Walker, we include three new environments (Featurized Extrapolation setting) despite tight compute limits. Preliminary IQMs show DALI improvements in the last column relative to Dreamer-DR.\n\n| Environment       | Dreamer-DR    | DALI-S        | DALI-S-$\\chi$   | DALI-D        | DALI-D-$\\chi$   | Best Method    | Extrap. Improvement |\n|-------------------|-------------:|-------------:|---------------:|-------------:|---------------:|---------------|-------------------|\n| **Acrobot**       | 0.970±0.011  | 0.985±0.002  | **0.986**±0.001    | 0.978±0.002  | 0.974±0.002    | DALI-S-$\\chi$ | +1.65%             |\n| **Pendulum**      | 0.843±0.020  | 0.869±0.023  | **0.942**±0.015    | 0.857±0.015  | 0.884±0.012    | DALI-S-$\\chi$ | +11.74%            |\n| **Quadruped Walk** | 0.139±0.022  | 0.136±0.025  | 0.148±0.029    | 0.137±0.024  | **0.220**±0.027    | DALI-D-$\\chi$ | +58.27%            |\n\n**Note**: Full results (10 seeds, all baselines) are infeasible within rebuttal. We focus on the most critical comparison (DALI vs. context-unaware Dreamer-DR) for latent-context settings.\n\n### (C8)  What is “near-optimal sample complexity”?\n\nBy “near-optimal sample complexity,” we mean that $\\mathfrak{z}_t$ captures context information close to the differential entropy $h(c)$ of the context distribution, with an error bounded by $\\delta \\in (0, 1)$. This is optimal up to a constant factor, as $\\mathcal{I}(c; \\mathfrak{z}_t)$ approaches the information-theoretic limit $h(c)$, and the sample complexity $\\mathcal{O}(K/\\delta^2)$ scales logarithmically with $1/\\delta$ due to the $\\beta$-mixing property. The choice of $K \\approx 64$ (for $\\delta = 0.01$, $\\lambda \\approx 0.1$ in DMC tasks) is theoretically justified and empirically validated through ablations.\n\n### (C9)  Formatting: \n- cMDP not yet introduced: Thanks for pointing this out. We will revise lines 2-3: \"Contextual Markov Decision Processes (cMDP) model this challenge ... hard to measure.\"\n- All appendices (A–D) are included in the Supplementary Material, per NeurIPS guidelines. The main paper adheres to the 9-page limit.\n\nWe have addressed all your concerns to the best of our ability, incorporating detailed clarifications and preliminary results to strengthen our manuscript. Thank you for your valuable feedback; we hope our response merits raising your score to an Accept for DALI’s impactful contributions."}, {"review_id": "XOWaMTcWPz", "reviewer": "Reviewer_serC", "summary": "This submission proposes a method for learning to detect the world context (like the mass of an object) from trajectories while learning a world model and a policy. The proposed method, DALI, is built on top of DreamerV3, using a separate transition model, loss function and conditioning mechanism. DALI is shown to correctly identify the context variables in two DeepMind Control environments even when the ground-truth context is out of distribution. The authors also perform a counterfactual analysis, showing that the trajectories generated by DALI respond to perturbations of the inferred context variable.", "questions": "Would techniques introduced in this work be applicable to improving sim2real transfer of robot locomotion policies, like quadruped running or bipedal walking?", "limitations": "Limitations of the proposed method are not sufficiently addressed. Section 7, “Discussion and Outlook”, should include a comprehensive discussion of the limitations of DALI.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n\n**General-purpose context inference method that could be applied to many domains.** DALI learns a module that infers the latent context of the environment from a trajectory of states and actions, this module is trained using a pair of forward modeling and latent space alignment losses. Since the formulation of this problem is very general, this approach could be used to, e.g., infer the ground dynamics during locomotion or object dynamics during robotic manipulation. The submission also proposes two approaches to integrating the context into the predictions of the world model and the policy learning algorithm. Effectively steering world models and policies with context variables is another crucial research direction.\n\n**Strong empirical analysis in two domains.** The experimentation performed in this submission is rigorous. Experiments are performed in three different regimes (interpolation, extrapolation and mixed regime) that test in-domain and out-of-domain generalization of the context inference module. Results are reported as Interquantile Means and statistical testing is performed to highlight the significance of the results. The three baselines and one ablation included in the experiments are sufficient. The submission further contains a qualitative example of the difference between generated trajectories with different context variables.\n\n## Weaknesses\n\n**Testing on only two tasks.** The submission uses a modified version of the CARL benchmark for evaluation. The original benchmark seems to include 19 tasks, but only two are used in this paper. Demonstrating that DALI can learn many tasks with the same architecture and hyper-parameters would strengthen the claims made in this paper.\n\n**Contribution tied to DreamerV3.** The contribution is currently framed as a modification of DreamerV3. The paper could have higher impact if the main contribution was agnostic to the specific instantiation of the world modeling and policy learning method. For example, showing a combination of DALI and TD-MPC2 could be interesting.\n\n**Only two ground-truth context axes.** The experiments vary only two environmental variables (gravity and string length in Ball-in-Cup and gravity and actuator strength in Walker). In practice, a robot may need to infer many variables. For example, domain randomization in sim2real training for locomotion usually involves 10+ domain variables.", "was_revised": true, "final_justification": "The rebuttal contains useful clarification as well as supporting experiments. I remain in favor of accepting this paper.", "rebuttal": "Dear Reviewer serC,\n\nWe are deeply grateful for your insightful review and recognition of DALI's general-purpose context inference, rigorous evaluation, and potential for real-world impact. Your questions highlight opportunities to clarify DALI’s broader impact and strengthen its claims, which we address below.\n\n### (C1) Testing on only two tasks\n\nWe sincerely thank you for noting the evaluation scope, referencing the CARL benchmark’s 19 tasks. Your suggestion to test DALI across more tasks with consistent architecture and hyperparameters is valuable, and we address it with new preliminary evaluations.\n\nAlongside Ball-in-Cup and Walker, we include three new environments (Featurized Extrapolation setting). Preliminary IQMs show DALI improvements in the last column relative to Dreamer-DR.\n\n| Environment       | Dreamer-DR    | DALI-S        | DALI-S-$\\chi$   | DALI-D        | DALI-D-$\\chi$   | Best Method    | Extrap. Improvement |\n|-------------------|-------------:|-------------:|---------------:|-------------:|---------------:|---------------|-------------------|\n| **Acrobot**       | 0.970±0.011  | 0.985±0.002  | **0.986**±0.001    | 0.978±0.002  | 0.974±0.002    | DALI-S-$\\chi$ | +1.65%             |\n| **Pendulum**      | 0.843±0.020  | 0.869±0.023  | **0.942**±0.015    | 0.857±0.015  | 0.884±0.012    | DALI-S-$\\chi$ | +11.74%            |\n| **Quadruped Walk** | 0.139±0.022  | 0.136±0.025  | 0.148±0.029    | 0.137±0.024  | **0.220**±0.027    | DALI-D-$\\chi$ | +58.27%            |\n\n**Note**: Full results (10 seeds, all baselines) are infeasible within rebuttal. We focus on the most critical comparison (DALI vs. context-unaware Dreamer-DR) for latent-context settings.\n\n### (C2)  Contribution tied to DreamerV3\n\nWe greatly value your suggestion to frame DALI’s contribution as agnostic to the world modeling and policy learning method, such as integrating with TD-MPC2. Your question about compatibility is critical, and we provide a detailed, honest assessment of DALI’s integration mechanics and theoretical implications.\n\n- **DALI’s core contribution**: DALI’s context encoder $g_\\varphi(o_{t-K:t}, a_{t-K:t-1})$ infers latent context $\\mathfrak{z_\\mathrm{t}}$ using $L_{\\text{FD}}$ and $L_{\\text{cross}}$ (Eqs. 1-3) from $K$-length histories, a modular component that can enhance any model-based RL framework with a latent state representation (Section 4.3).\n\n- **RSSM’s role in Theorem 1**: The guarantee of $\\mathcal{I}(c; \\mathfrak{z_\\mathrm{t}}) \\geq (1 - \\delta) h(c)$ with $\\mathcal{O}(T/K)$ sample efficiency relies on the RSSM’s recurrence ($h_t = f_\\theta(h_{t-1}, z_{t-1}, a_{t-1})$), which propagates context across time. This recurrence is specific to DreamerV3 but common in recurrent world models.\n\n- **TD-MPC2 integration**: TD-MPC2 is a model-based RL algorithm that does not use a RSSM like Dreamer. This means it does not maintain a deterministic recurrent state (like Dreamer's\n$h_t$) over time. Instead, it uses a Markovian latent dynamics model ($z_{t+1} = f_\\theta(z_t, a_t)$) without a persistent $h_t$. DALI’s $g_\\varphi$ can integrate via:\n\n   - **Shallow Integration**: Append $\\mathfrak{z_\\mathrm{t}}$ to observations, encoding $z_t = \\text{enc}_\\theta(o_t, \\mathfrak{z}_t)$, modifying TD-MPC2’s encoder to incorporate context.\n\n    - **Deep Integration**: Condition dynamics on $\\mathfrak{z_\\mathrm{t}}$, yielding $z_{t+1} = f_\\theta(z_t, a_t, \\mathfrak{z_\\mathrm{t}})$, embedding context in transitions.\n\n- **Training**: $L_{\\text{FD}}$ can be applied (predict $o_{t+1}$ from $o_t, a_t, \\mathfrak{z_\\mathrm{t}}$) as it is a self-supervised loss that does not depend on the RL algorithm's internals.\n\n- **Sample efficiency implications**:  Since TD-MPC2 does not have a recurrent state that accumulates information over time, the theoretical guarantee of Theorem 1, which is derived for a recurrent model (like DreamerV3), may not directly transfer to TD-MPC2.\n\n- We will include a comment in Section 7 to discuss DALI’s potential integration with non-recurrent models like TD-MPC2.\n\n### (C3) Only two ground-truth context axes; Domain randomization in sim2real training for locomotion usually involves 10+ domain variables\n\nWe appreciate your concern about evaluating only two context variables (gravity/string length for Ball-in-Cup, gravity/actuator strength for Walker) when real-world tasks involve 10+ variables.\n\nWe prioritized in-depth analysis along two axes for rigor and interpretability. In principle, DALI can handle higher-dimensional contexts, but practical compute constraints arise.\n\nTo see this, we highlight the computational complexity for our existing setup for context inference along two axes. DALI's evaluation used NVIDIA A100 GPUs (80GB) and Intel Xeon CPUs, with 2–4 GPUs available. It comprised 600 runs: 10 seeds × 5 algorithm variants × 2 environments (Ball-in-Cup, Walker) × 2 modalities (featurized, pixel-based) × 3 context settings. Featurized runs paralleled 10 seeds; pixel-based, 4 seeds. Per-trial GPU hours: 20–26 (Ball-in-Cup), 46–54 (Walker). **Total cost: ~24,000 GPU hours** sequentially, reduced to 2,101 hours (88 days) with 2 GPUs or 1,051 hours (44 days) with 4 GPUs, highlighting the extensive computational complexity.\n\nScaling to 10+ variables would require ~1,000 GPU-hours per task for combinatorial sampling, increasing the compute budget to ~**25,000 GPU-hours**.\n\nWe will add a discussion in Section 6 on scalability to higher-dimensional contexts.\n\n### (C4) Can techniques introduced here be applicable to improving sim2real transfer of robot locomotion policies e.g. quadruped/bipedal walking?\n\nWe value your question on sim2real transfer for quadruped/bipedal walking, which highlights DALI’s real-world potential. We assess applicability using existing results.\n\n- **Sim2Real Relevance**: DALI’s $g_\\varphi$ infers latent contexts (e.g., gravity, actuator strength) from trajectories, addressing sim2real challenges where robots face real-world variations (e.g., friction, mass) not seen in simulation. DALI’s $K$-step inference aligns with real-time adaptation windows (e.g., 1–2 gait cycles). The context encoder adds minimal overhead (< 5\\% parameters to DreamerV3-small) and scales to high-dimensional tasks. In general, training time scales approximately linearly with the dimensionality of the state and action spaces. Pixel tasks take 1.5–2x longer due to CNN encoding.\n\n- **Bipedal Walking**: Walker tests torque-dependent locomotion (24D state, 6D action), with DALI-S achieving IQM 0.7810 (Featurized Extrapolation), inferring gravity and actuator strength. This supports sim2real robustness for bipedal walking by adapting to OOD dynamics (e.g., actuator variations).\n- **Quadruped Walking**: Walker’s locomotion dynamics are analogous to quadruped tasks but more complex. We tested DALI on Quadruped Walk (56D observation, 12D action), a larger environment than Walker (24D observation, 6D action). Trained for 500K timesteps and contextualized with gravity and actuator strength (same range as DMC Walker), preliminary IQMs for the Featurized modality are 0.139 (Dreamer-DR) and 0.22 (DALI-S-$\\chi$), a **58.3%** improvement. The 56D observation space highlights the challenge of torque-sensitive locomotion with increased coordination complexity and higher degrees of freedom (12 vs. 6 in Walker), demanding more robust context inference.\n- We will add a sim2real discussion in Section 6, referencing Walker and Quadruped results, and note future sim2real experiments with real robot data.\n\n\nWe thank you for your invaluable feedback. Your vision for DALI’s broader impact inspires our future work, and we hope you now consider it a Strong Accept."}, {"review_id": "FmvzDs5wah", "reviewer": "Reviewer_ogTR", "summary": "The authors propose to build an approach to train agents for contextual MDPs where the context (variable parts of an environment, e.g. gravity, friction) is latent and fixed for each episode and the objective is to generalize across contexts. The context is identified automatically using additional context encoders on top of DreamerV3 which processes information over shorter windows of size K=50 and output 8-dimensional continuous context embeddings. They use transformers for this encoder. Their approach can have two forms of integration into the DreamerV3 arhcitecture - \"shallow\" (DALI-S), where the context is used only as an additional input to the original world model encoder and \"deep\" (DALI-D), where it is inetgrated into the RSSM, decoder, reward and continue predictors instead. They also, I believe, introduce an additional network that predicts the next observations using the context to ensure the context learnt is useful (see eqn. 1). Additionally, an optional cross-modal loss is introduced to regularize the context to be more aligned with the hidden stochastic latent state z of DreamerV3.\n\nThey prove in Theorem 1 that, assuming beta-mixing and Lipschitz dynamics, their approach is more sample efficient than Dreamer with a gain of O(T/K) where T is the episode length (I did not check the proof).\n\nThey evaluate the *small* variant of DreamerV3 on pixel and feature-based versions of Ball-in-Cup and Walker variants in the CARL benchmark on interpolation, extrapolation and mixed regime settings to evaluate generalization. Across many experiments, one of their two evaluated approaches performs best in extrapolation scenarios, with the experiment on DMC Walker-Pixel being an exception to this. The results are somewhat mixed in terms of which agent performs best accross the different tasks and the settings of interpolation vs extrapolation vs mixed regimes.\n\nThey also perform counterfactual rollouts with perturbing elements of the learnt context vector and show some experiments which show that gravity and string length may be captured by one of the dimensions of this context vector.", "questions": "I have raised many question in the above. And also provided ideas on what can be done to improve the paper in my opinion. They are listed somewhat in decreasiong order of importance for me, with the parts up to and including \"Other Technical issues\" being the most important.", "limitations": "yes", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 1, "significance": 2, "originality": 3, "strengths_and_weaknesses": "## Strengths:\nThe idea of the paper to extend previous approaches to a setting where the context is latent and is identified automatically is a very important setting with great potential. The experimental setup seems good in general, with 10 random seeds evaluated and a fair selection of environments and their variations evaluated. The counterfactual evaluation also seems promising.\n\n## Weaknesses:\n### Presentation:\nThe paper was hard for me to follow, with me having to put together the many components to try and understand the approach.\n@authors, Please provide an overview block diagram of the different components of your approach, the gradient flow, etc. There are so many moving parts, that it is hard to understand everything without an overview. For example, I may be wrong about claiming that you introduced an additional network above. This depends on what $f^w_{phi}$ is in equation 1? Is it another MLP or GRU or what?\n\n### Lack of evaluating their own baseline DALI-D:\nThe paper does not seem to evaluate their own method, DALI-D, that is introduced as one of the main methods sec 4.3. At least it is missing in figs. 1 and 2.\n\n### No hyperparameter optimisation:\nNo hyperparameter optimisation was performed if I understood the Appendix correctly and default settings of previous work were taken even though the experimental setting is different in the paper here.\ncRSSM which is proposed as a baseline does not seem to have been used as in its original form. Instead -S and -D variants of it conforming to the authors shallow and deep DALI variants are introduced. When changes to a baseline, such as cRSSM, are introduced without tuning for the new experimental setup, it is hard to make claims about the results in this paper. It seems to me that no hyperparameter optimisation was performed even for the authors' own approaches. That could explain the mixed results and why different agents perform best for different settings.\n\n\n### Other Technical issues:\nCross modality loss (see eqn. 2):\n\nThe cross modality loss only aligns the context closer to the stochastic state z which is something learned over long windows while the context is supposed to be learned from K-sized windows. I do not see the reason the context needs alignment with z when the original reason to introduce the context was that it is captured over shorter windows. Does not the cross-modal loss (regularization) break the claim from Theorem 1: *Compared to DreamerV3’s RSSM processing full episodes of T >> K transitions, DALI achieves a sample complexity gain of O(T /K).* Do you now not need to process full espisodes to make the context and z more similar?\n\nIn equation 2, are the W's just matrices or actual MLPs? I would expect MLPs to be represented as functions (and not W) since they include non-linearities. But since invertibility is being talked about later, I am assuming they are actually matrices?\n\n\nI do not see any learning curves, it is fine to provide aggregate results in the paper, but without individual learning curves, I cannot say if the agents converged and whether, e.g., 200K timesteps for ball-in-cup is a fair training length. I can be fairly confident about the 500K for Walker because I am familiar with that environment, but in general one would need learning curves.\n\n### Lack of ablations:\nI would like to see ablations on the setting of K. Without this, it is hard to see in practice how this value is influencing the agent.\n\nI would be interested in seeing other latent dimensions and their effects as well, not just the \"best dimension\" $z_6$. So, I would also propose to show the effect of the worst context dimension on counterfactual settings.\n\nA couple of side remarks here about something I found confusing:\n1) Does $z_6$ capture both gravity and string length at the same time? Do we not want it to capture a single factor at one time? Adding some remarks about this in the limitations section would be nice.\n2) Why is it that element number 6 of the context captures the same gravity + string length perturbation across pixel and feature modalities? These 2 modalities mean 2 separate training processes for the agents. Is this just a coincidence that element number 6 of the context was not only the *best* dimension in both feature and pixel modalities but also captured the same 2-factor perturbation across the 2 settings?\n\n\n### Too strong claims:\n\n*essential for robust generalization, proving that it infers latent contexts from short interaction windows with near-optimal sample complexity*\n\nI do not agree with the wording here. While I did not check the proofs, from what I understand, DALI improves generalization but is not *essential* because DreamerV3 can do so but only at a poorer sample complexity. What do the authors mean by \"near-optimal sample complexity\"? If the theorem does prove that DALI has improved sample complexity by O(T/K), it does not mean it has near-optimal complexity, I would expect that at least K needs tuning / ablations to show any form of near optimality.\n\n\n*Physically consistent counterfactuals: ...  faithfully mirroring real-world physics.*\n\nTo show that the counterfactuals are faithful, we need to see how the relationship changes for different values of that dimension. For example, acceleration for different values of g would change linearly wrt g for a ball falling down (if we ignore other forces such as air resistance). Can the authors show that for, say, 3 different values of context dimension 6 for the feature-based Ball in cup environment, there is a largely linear relationship for acceleration experienced?\n\n\nIn lines 115 and 152: *is a deterministic recurrent state capturing temporal dependencies, and $z_t$ is a stochastic state encoding uncertainty about the current observation*\n\nThe context *avoids encoding redundant trajectory-specific information from the deterministic ht, which could impair generalization*\n\nThese are again strong assumptions that h_t only captures deterministic temporal dependencies and $z_t$ is a stochastic state encoding uncertainty. As far as I know, DreamerV3 authors did not prove this. In my opinion, $h_t$ needs to capture uncertainty information as well given that $z_t$ is coming from the encoder which takes $h_t$ and $o_t$ as inputs and since $o_t$ itself has no information on uncertainty by itself, $h_t$ needs to capture it.\n\nIn line 366, they say the DALI-S setting outperforms DALI-D and claim this is because it acts as a regularizer. It is a strong claim again that feels like it was something fit to their observations from the experiments.\n\n\n### Less important points:\nSection 6.2:\n  Which of the 4 DALIs that were mentioned in sec 6.1 is used for the counterfactual experiments in 6.2? This is only mentioned later in 6.2.1.\n\nLine 310: *For a trajectory pair ... classifier outputs*? Should it not be operating on a single trajectory? Or is the classifier providing 2 outputs, one for each trajectory?\n\nProvide one line intuitive explanations for differential entropy and mutual information and why MI being >= 0.99 times DE is a good thing to capture near-optimal context information.\n\nThe paragraph beginning in line 157 already discusses the deep integration stuff from Section 4.3 without having introduced it.\n\nProviding too many exact numbers in textual discussion (page 6) hampers reading. These can be provided in a table in the Appendix and the discussion can be compressed and made easier to read.\n\nMake \"Deep integration\" italic. It is too generic a name and I thought at some point that it just meant *Deep Learning* being integrated into the setup.", "was_revised": true, "final_justification": "I am still a bit concerned about different variants performing best in different experiments and it feels to me like the hyperparameter tuning is lacking.\n\nAdditionally, they introduce the cross-modality loss leading to 4 variants and the loss itself is not completely convincing to me.", "rebuttal": "Dear Reviewer ogTR,\n\nThank you for your thoughtful review and valuable feedback. We appreciate your recognition of our contributions and welcome the chance to clarify and improve clarity in our rebuttal.\n\n### (C1) Presentation and overview diagram. Clarify $f^w_\\varphi$ in Eq.1.\n\nWe acknowledge that DreamerV3’s complexity can make our contributions seem intricate. We clarify: \n\n* **DreamerV3 (Prior Work)**: RSSM backbone with deterministic ($h_t$) and stochastic ($z_t$) states from observations $o_t$, actions $a_t$.\n* **DALI additions**:\n\n  * Context encoder ($g_\\varphi$): transformer on $K$-step window $(o_{t-K:t}, a_{t-K:t-1})$, outputs $\\mathfrak{z}_t \\in \\mathbb{R}^8$.\n  * Forward model ($f^w_\\varphi$): lightweight 2-layer MLP predicting $o_{t+1}$ from $(o_t, a_t, \\mathfrak{z}_t)$, ensuring $\\mathfrak{z}_t$ captures dynamics (Eq. 1).\n  * $g_\\varphi$ and $f^w_\\varphi$ trained jointly.\n* Integration: *Shallow* ($\\mathfrak{z}_t$ conditions world-model encoder) and *Deep* (conditions entire model and policy).\n\nAs NeurIPS disallows figures in rebuttals, we’ll add this overview text and include a gradient-flow diagram in the final version.\n\n### (C2) Results for DALI-D are missing.   \n\nDALI-D is a proposed method, not baseline, as you noted. DALI-D underperforms DALI-S in our experiments (e.g., Ball-in-Cup Pixel Extrap.: 0.225 vs. 0.273; Walker Featurized Extrap.: 0.771 vs. 0.781). We’ll add a footnote noting this and refer to full results in the Appendix.\n\n### (C3) No hyperparameter tuning? Baselines (cRSSM-S/D) were modified without retuning? \n\nWe followed [Prasanna et al. (2024)]’s protocol:\n- **DreamerV3**: Used the *small* variant with same hyperparameters from [Hafner et al. (2023)].  \n- **Context Encoder**: Explored transformer- and GRU- based designs; chose transformer (App. C.3) via validation.\n- **$K=64$**: Chosen via Theorem 1 ($K = \\Omega(\\log(1/\\delta)/\\lambda) \\approx 64$ for $\\delta=0.01$, $\\lambda=0.1$). Verified also via ablations (Response C7 below).\n- **Confusion about Baselines**: We clarify that our baselines cRSSM-S, cRSSM-D, and Dreamer-DR are directly adopted without any modifications from Prasanna et al. (2024) (see their arXiv version, Sec 5, P.7), a recent work that extends DreamerV3 for cMDPs. Specifically, our naming maps as follows (App. A.1 DALI and Dreamer-based Baselines):\n  - `cRSSM-S` ≡ `concat-context` in [Prasanna et al.].  \n  - `cRSSM-D` ≡ `cRSSM` in [Prasanna et al.].  \n  - `Dreamer-DR` ≡ `hidden-context` in [Prasanna et al.].  \nWe used their hyperparameters unchanged for fair, reproducible comparison, leveraging their proven cMDP performance. We’ll add this mapping in Sec. 6.\n\n**Mixed results** (e.g., DALI-S-$\\chi$ leads Ball-in-Cup Extrap., IQM 0.3720 Featurized, while cRSSM-S leads Walker Pixel Extrap., IQM 0.7770) reflect cMDP challenges: nonlinear dynamics favor DALI in Ball-in-Cup, more linear dynamics favor cRSSM in Walker Pixel (see App. D.2: Generalization Trends Across Environments and Modalities).\n\n\n### (C4) Cross-Modal Loss vs. Theorem 1: Aligning $\\mathfrak{z}_t$ with $z_t$ (long-window state) breaks $O(T/K)$ sample efficiency claim.  \n\n**No conflict: Cross-modal loss preserves $O(T/K)$ efficiency:**\nTheorem 1 (Sec. 5; formal versions in App. A Theorems 2 & 5) covers Deep integration’s forward-dynamics loss $L_{\\text{FD}}$ only. The $O(T/K)$ sample efficiency also holds for Shallow integration since both train the context encoder on $K$-step windows (App. A, Remark 4). Adding $L_{\\text{cross}}$ does not change this, as it uses the same $K$-length data (App. B, Algorithm 3), aligning $\\mathfrak{z}_t$ and $z_t$ computed from identical windows with detached history:\n- **Gradient isolation**:  \n  - The recurrent unroll is limited to $K$ steps; $\\varphi$ updates depend only on the current window with detached states blocking backpropagation beyond it. The world model ($\\theta$) is **not updated** during context learning (inputs to $f_\\theta$, $q_\\theta$ are detached, App. B Algorithm 3, lines 16–29). Gradient stopping isolates $\\varphi$, preserving the $\\mathcal{O}(T/K)$ efficiency; $L_{\\text{cross}}$ needs no full episodes.\n- **Empirically: Why align $\\mathfrak{z}_t$ with $z_t$?**:  \n  $z_t$ encodes instantaneous dynamics (e.g., ball velocity in Ball-in-Cup). Aligning $\\mathfrak{z}_t$ (context) with $z_t$ can avoid overfitting to transient noise and leverage structured priors from the world model, enhancing $\\mathfrak{z}_t$’s robustness (see App. Sec. D.2 for a discussion).\n\nWe’ll revise Sec. 4.3 to clarify: “$L_{\\text{cross}}$ uses $K$-length windows, preserving the $\\mathcal{O}(T/K)$ sample complexity gain.”\n\n### (C5) $W_z$ and $W_{\\mathfrak{z}}$ matrices or MLPs?\n\nYou are correct that $W_z \\in \\mathbb{R}^{32 \\times 8}$ and $W_{\\mathfrak{z}_t} \\in \\mathbb{R}^{8 \\times 32}$ in Eq.2 are matrices, not MLPs. We’ll fix this typo in the revision.\n\n\n### (C6) Show learning curves\n\nWe confirm convergence:  \n- **Ball-in-Cup**: Returns plateaued by 150K steps (200K allocated).  \n- **Walker**: Returns stabilized by 400K steps (500K allocated).  \nWe will add convergence curves to the camera-ready App..  \n\n\n### (C7) Ablation on $K$  \n\nWe validated $K$ by training a small MLP $f_\\varrho: \\mathbb{R}^8 \\rightarrow \\mathbb{R}^2$ to predict ground-truth context from $\\mathfrak{z}_t$ in Ball-in-Cup Featurized. $K = 64$ achieved the lowest MSE, outperforming $K = 16, 32$ (insufficient context) and $K = 128, 256$ (diminishing returns due to noise accumulation in longer windows). Will add a summary table in the revision.\n\n\n### (C8) Counterfactual (CF) analysis\n\n**a) Effect of lowest-ranked dimension**:  \nWe analyzed lowest-ranked $\\mathfrak{z}_3$ (AUC 0.56, Featurized); perturbations caused random ball position changes (<0.01 Z-amplitude), indicating noise or weak context. Full AUCs to be added in App. D.\n\n**b) Coupled or disentangled representations:**\n$\\mathfrak{z}_6$ encodes coupled gravity-string dynamics in Ball-in-Cup, as perturbations induce trajectories with shorter string lengths and higher accelerations. In RL, where the goal is to maximize cumulative rewards, coupled representations can outperform disentangled ones when dynamics involve nonlinear context interactions (e.g., $T_p \\approx 2\\pi \\sqrt{L/g}$). $\\mathfrak{z}_6$’s joint encoding aids policy adaptation to unseen gravity-string combos. Our losses don’t enforce disentanglement, benefiting from this coupling. Contrast this with Walker (App. D.3.2), where the top-ranked $\\mathfrak{z}_3$ aligns with a single factor (actuator strength). We’ll add a note: \"DALI learns task-relevant, not disentangled, representations.\"\n\n**c) Coincidence of $\\mathfrak{z}_6$ across modalities:**\n$\\mathfrak{z}_6$ as the top latent in DALI-S-$\\chi$ for both Featurized and Pixel Ball-in-Cup (AUC 0.82, 0.80), encoding coupled gravity-string dynamics. Though trained separately, shared seeds and identical zero-action conditions ensure its consistency. In Featurized, $\\mathfrak{z}_6$ causes interpretable changes (e.g., shorter string, higher acceleration); in Pixel, we iterated experiments to intentionally synchronize visuals and select the same dimension for comparable trajectories. Effect sizes vary, but shared dynamics yield consistent latent structure. Sec. 6.2 will clarify that this consistency stems from our controlled setup for clear exposition, not coincidence, enabling Pixel and Featurized plots to jointly show a full oscillation cycle governed by the same physics.\n\n\n### (C9) Theoretical claims \n\n**a) \"Essential\" and \"Near-Optimal\"**:  \n- “Essential” highlights the encoder’s key role in reducing DreamerV3 RSSM’s information bottleneck (Theorem 3, App. A): $\\mathcal{I}(c; h_t^{\\text{DALI}}) \\geq \\mathcal{I}(c; h_t^{\\text{RSSM}}) - \\epsilon(K)$. Without it, DreamerV3’s GRU needs full episodes to capture context, slowing OOD adaptation (e.g., Ball-in-Cup Extrap. IQM 0.1980 vs. 0.3720). We’ll revise “essential” to “highly beneficial.”\n- “Near-optimal sample complexity” means $\\mathfrak{z}_t$ captures context info close to the context entropy $h(c)$ within error $\\delta \\in (0,1)$. This is optimal up to constants, as $\\mathcal{I}(c; \\mathfrak{z}_t)$ approaches $h(c)$, and sample complexity $\\mathcal{O}(K/\\delta^2)$ scales logarithmically with $1/\\delta$ via the $\\beta$-mixing property.\n\n\n**b) Linearity in Counterfactuals**:  \nWe probed $\\Delta \\in {-\\sigma_6, 0, +\\sigma_6}$ in the Featurized modality and observed monotonic but nonlinear changes in peak Z-velocity: -10%, 0%, +12%. This reflects the pendulum-like dynamics of Ball-in-Cup, where $\\mathfrak{z}_6$ jointly modulates gravity ($g$) and string length ($L$), affecting the oscillation period $T_p \\approx 2\\pi\\sqrt{L/g}$. A linear relationship with $g$ would apply to free fall ($a=g$). But when the ball is *tethered by a string*, the effect of $g$ on velocity is **not linear** due to the coupling with string length $L$.\n\n**c) $h_t$ vs. $z_t$ for uncertainty**:  \nWe agree that $h_t$ influences $z_t$’s uncertainty via the encoder, but $h_t$’s deterministic nature means it captures historical trajectory information (e.g., past ball positions in Ball-in-Cup) without directly modeling stochasticity, while $z_t$ *explicitly* represents observation uncertainty. Aligning with $z_t$ alone yields better results empirically.\n\n**d) Shallow Integration as regularizer**:  \nWe agree that this is an interpretation grounded in evidence from our experiments. Shallow integration has fewer parameters (no context conditioning in RSSM), and it propagates $\\mathfrak{z}_t$ implicitly via recurrence. \n\n\n### (C10) Counterfactual Classifier  \n\nClassifier $c_\\nu^{(j)}$ takes one trajectory, outputs one context probability. For pair $(\\mathcal{T}^{(0)}, \\mathcal{T}'^{(j)})$, it evaluates each separately.\n\nWe appreciate your incisive feedback, which has greatly improved our work. We will address your formatting-related points and hope our responses resolve all your concerns, encouraging you to raise your score to Accept."}, {"review_id": "Pc4q5fH5aT", "reviewer": "Reviewer_mG63", "summary": "This paper introduces Dynamics-Aligned Latent Imagination (DALI) for visual reinforcement learning, which infers latent contexts from interactions.  Experiments demonstrate the effectiveness of the proposed method.", "questions": "- See the weaknesses.\n- I suggest adding the corresponding numerical comparisons for the results in Figures 1 and 2.", "limitations": "Yes.", "rating": 4, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 2, "strengths_and_weaknesses": "**Strengths：**\n\n- DALI outperforms context-unaware baselines on the  contextualized DMC Ball-in-Cup and Walker Walk tasks.\n- This paper provides both empirical and theoretical analyses.\n\n**Weaknesses：**\n\n- The readability could be improved, as there is no figure provided to illustrate the pipeline.\n- Experiments were conducted on only two tasks, and more experimental validation is needed.\n- The paper does not compare ContextWM [1], which is a strong baseline that explicitly separates context and dynamics modeling.\n- The paper lacks references related to visual RL, such as [1], [2], and [3].\n\nReferences:\n\n[1] Wu et al. \"Pre-training Contextualized World Models with In-the-wild Videos for Reinforcement Learning\", NeurIPS, 2023.\n\n[2] Li et al. \"Open-World Reinforcement Learning over Long Short-Term Imagination\", ICLR, 2025.\n\n[3] Fujimoto et al. \"MR.Q: Towards General-Purpose Model-Free Reinforcement Learning\", ICLR, 2025.", "was_revised": true, "final_justification": "The rebuttal addressed most of my concerns. However, additional experiments are still needed to demonstrate the effectiveness of the proposed method, and the readability of the paper could be further improved.", "rebuttal": "Dear Reviewer mG63,\n\nWe appreciate your engagement with our work and your suggestion to compare with ContextWM. Below, we address your concerns while clarifying the scope, contributions, and positioning of DALI. We hope this response alleviates your concerns and demonstrates why DALI merits a higher score.  \n\n### (C1) No overview figure provided to illustrate the pipeline.\n\nWe acknowledge that DreamerV3’s complexity can make our contributions seem intricate. We clarify: \n\n* **DreamerV3 (Prior Work)**: RSSM backbone with deterministic ($h_t$) and stochastic ($z_t$) states from observations $o_t$, actions $a_t$.\n* **DALI additions**:\n\n  * Context encoder ($g_\\varphi$): transformer on $K$-step window $(o_{t-K:t}, a_{t-K:t-1})$, outputs $\\mathfrak{z}_t \\in \\mathbb{R}^8$.\n  * Forward model ($f^w_\\varphi$): lightweight 2-layer MLP predicting $o_{t+1}$ from $(o_t, a_t, \\mathfrak{z}_t)$, ensuring $\\mathfrak{z}_t$ captures dynamics (Eq. 1).\n  * $g_\\varphi$ and $f^w_\\varphi$ trained jointly.\n* Integration: *Shallow* ($\\mathfrak{z}_t$ conditions world-model encoder) and *Deep* (conditions entire model and policy).\n\nAs NeurIPS disallows figures in rebuttals, we’ll add this overview text and include a overview figure in the final version.\n\n### (C2) Experiments were conducted on only two tasks, and more experimental validation is needed.\n\nWe agree that broader benchmarks strengthen our claims. Alongside Ball-in-Cup and Walker, we include three new environments (Featurized Extrapolation setting). Preliminary IQMs show DALI improvements in the last column relative to Dreamer-DR.\n\n| Environment       | Dreamer-DR    | DALI-S        | DALI-S-$\\chi$   | DALI-D        | DALI-D-$\\chi$   | Best Method    | Extrap. Improvement |\n|-------------------|-------------:|-------------:|---------------:|-------------:|---------------:|---------------|-------------------|\n| **Acrobot**       | 0.970±0.011  | 0.985±0.002  | **0.986**±0.001    | 0.978±0.002  | 0.974±0.002    | DALI-S-$\\chi$ | +1.65%             |\n| **Pendulum**      | 0.843±0.020  | 0.869±0.023  | **0.942**±0.015    | 0.857±0.015  | 0.884±0.012    | DALI-S-$\\chi$ | +11.74%            |\n| **Quadruped** | 0.139±0.022  | 0.136±0.025  | 0.148±0.029    | 0.137±0.024  | **0.220**±0.027    | DALI-D-$\\chi$ | +58.27%            |\n\n**Note**: Full results (10 seeds, all baselines) are infeasible within rebuttal. We focus on the most critical comparison (DALI vs. context-unaware Dreamer-DR) for latent-context settings. \n\n**Runtimes:** For an accurate assessment of **computational cost**, please see Response (C3): Runtime Efficiency to Reviewer i6n9.\n\n### (C3) The paper does not compare ContextWM [1], which is a strong baseline that explicitly separates context and dynamics modeling. The paper lacks references related to visual RL.\n\n### We respectfully clarify that **ContextWM Is not a valid baseline for DALI**:\n\n- ContextWM focuses on **visual RL pre-training** from passive videos to separate static visual features (e.g., texture, background) from dynamics. In contrast, DALI addresses **contextual MDPs (cMDPs)**, inferring **latent dynamics parameters** (e.g., gravity, friction) for **zero-shot generalization** to unseen environments. ContextWM does not operate under the cMDP formulation and cannot serve as a meaningful baseline for DALI. \n- **DALI solves contextual MDPs, not Visual RL**: DALI is designed to infer **latent context vectors** that modulate **transition dynamics**. We test pixel-based tasks to assess robustness to partial observability. All model components, viz., context encoder, RSSM, loss functions, are tailored for **dynamics inference**, not visual scene understanding. \n- **Key difference in “Context” definition**: ContextWM [1] defines context as a single observation frame $c := o_t$, $t \\sim \\text{Uniform}\\{1, \\ldots, T\\}$, capturing static visual attributes (e.g., texture, color) to simplify dynamics modeling in visual RL pre-training. In contrast, DALI defines context as a vector $c$ (e.g., gravity, string length) parameterizing cMDP dynamics.\nDALI’s context modulates how the world behaves (e.g., how gravity affects motion); ContextWM’s context describes how the world looks. This fundamental difference in **definition, sampling, and purpose** makes ContextWM inappropriate for comparison.\n\n## Emphasizing DALI’s Merits\n\nWe respectfully highlight DALI’s key contributions to justify a higher score:\n1. **Strong zero-shot generalization**:  \n   - Up to **+96.4%** gains over context-unaware baselines, often surpassing ground-truth context-aware baselines in extrapolation tasks.\n2. **Theoretical foundations**:  \n   - Theorem 1 (formal version: Appendix A, Theorem 2) provides sample efficiency guarantees for context inference (Section 5), while Theorem 3 (Appendix A) introduces a novel information bottleneck result for DreamerV3’s RSSM. Together, these represent new theoretical contributions to contextual MDPs.\n3. **Physical consistency**:  \n   - Perturbations in DALI's latent space induces physically plausible counterfactuals, showing actionable context encoding (Sec. 6). Our counterfactual analysis is, to our knowledge, the first in the Dreamer framework to probe latent imagination, offering new interpretability insights into learned context representations.\n4. **Architecture-agnostic design**:  \n   - DALI’s components (context encoder, dynamics losses) can augment *any* model-based algorithm (e.g., TD-MPC2) (see Response (C2) to Reviewer serC).  \n\n**Request**: We respectfully urge you to reassess DALI based on its core contributions and novel insights, and to consider raising your score in light of its potential impact."}]

Author rebuttal:

Meta-review: The authors propose a method to jointly encode forward dynamics, using it to condition both control and perception problems downstream. The overall method is an extension of the Dreamer architecture and while reviewers were initially concerned about the lack of more rigorous evaluation in more settings to prove the usefulness of the method, it appears to have been addressed by the authors during rebuttals. In fact, multiple reviewers have raised their scores and if the authors update their main paper with everything they have stated in the discussion period, then this paper warrants an acceptance.
Decision: Accept (poster)