Title: Dynamics-Aligned Latent Imagination in Contextual World Models for Zero-Shot Generalization
Venue: NeurIPS 2025
Accepted: yes
Keywords: world models

contextual reinforcement learning

representation learning

counterfactual consistency
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

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 contextunaware baselines, often surpassing context-aware baselines in extrapolation tasks, enabling zero-shot generalization to unseen contextual variations.

1 Introduction

The ability to generalize across diverse environmental conditions is a fundamental challenge in reinforcement learning (RL). Contextual Markov Decision Processes (cMDPs) formalize this challenge by modeling task variations through latent parameters, such as friction, gravity, or object mass that govern dynamics [Hallak et al., 2015]. However, real-world agents rarely have direct access to these parameters. Consider a robotic control task where a 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 to contextual RL often rely on explicit context conditioning, where models are trained with domain-specific instrumentation or carefully crafted architectures [Seyed Ghasemipour et al., 2019, Ball et al., 2021, Eghbal-zadeh et al., 2021, Mu et al., 2022, Beukman et al., 2023, Benjamins 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 or difficult to measure. The challenge is to develop RL agents that infer and adapt to hidden contexts in a self-supervised manner, enabling robust generalization without direct supervision.

To address this challenge, we propose Dynamics-Aligned Latent Imagination (DALI), a framework integrated within the DreamerV3 architecture [Hafner et al., 2025] that enables zero-shot generalization by inferring latent contexts directly from interaction histories. Unlike DreamerV3, which struggles

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 environmental variations, overcoming information bottlenecks and sample inefficiencies in learning diverse contexts [Prasanna et al., 2024, Hafner et al., 2025].

By inferring latent contexts from interactions, DALI 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 nearoptimal 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 baselines in extrapolation tasks.

Physically consistent counterfactuals : We show that the learned latent space exhibits counterfactual 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., 2015, Modi et al., 2018]. A recent survey [Kirk et al., 2023] highlights its relevance for zero-shot generalization, emphasizing how clear context sets for training and evaluation enable systematic analysis. One research direction assumes context is explicitly observed as privileged information and integrates it into learning [Chen et al., 2018, Seyed Ghasemipour et al., 2019, Ball et al., 2021, Eghbal-zadeh et al., 2021, Sodhani et al., 2021, Mu et al., 2022, Benjamins et al., 2023, Prasanna et 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 et 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-purpose model-free 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 [Prasanna et al., 2024], we build on it with a dynamics-aligned inference mechanism to enhance generalization without context supervision.

Meta-RL. Meta-RL trains agents to adapt rapidly to new tasks with minimal experience [Beck et al., 2023], typically by learning adaptive policies that infer task-specific information from past interactions. However, most meta-RL methods require fine-tuning on new tasks [Duan et al., 2016, Finn et al., 2017, Nagabandi et al., 2018, Rakelly et al., 2019, Zintgraf et al., 2019, Melo, 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 in partially observable environments, we adopt a contextual Markov Decision Process (cMDP) framework, following Hallak et al. [2015], Kirk et al. [2023].

A cMDP is defined as a tuple M = ( S , A , O , C , R , P , E , µ, p C ) , where S , A , and O are resp. the state, action, and observation spaces, C ⊆ R d is the context space with contexts c ∈ C drawn i.i.d. from a distribution p C , R : S × A × C → R is the reward function, P : S × A × C → ∆( S ) is the stochastic transition function specifying the probability distribution over next states given the current state, action, and context, E : S × C → ∆( O ) is the observation function specifying the distribution over observations given the state and context, µ : C → ∆( S ) is the initial state distribution where µ ( s 0 | c ) specifies the probability of initial state s 0 given context c . In this partially observable setting, the agent does not observe the state s t or context c directly but receives observations o t ∈ O . The objective is to learn a policy π ( a t | o 1: t , a 1: t -1 ) that maximizes the expected return E [ ∑ T t =0 r t ] , where T is the episode length.

In our framework, each context c 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 p C . The agent infers c from interaction histories to adapt its policy. To formalize zero-shot generalization, we define two distributions over contexts: the training distribution p train ( c ) , from which contexts are sampled during training, and the evaluation distribution p eval ( c ) , representing unseen test contexts. The goal is to learn a policy using samples from p train ( c ) that maximizes the expected return for contexts drawn from p eval ( c ) without further adaptation or retraining.

4 Dynamics-Aligned Latent Imagination (DALI)

4.1 Background: DreamerV3

DreamerV3 [Hafner et al., 2025] is a model-based RL algorithm that enables agents to learn and plan in the latent space of a 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 o t (e.g., images or sensor data) and actions a t . The world model, structured as a Recurrent State-Space Model (RSSM), encodes these observations into a compact latent state s t = { h t , z t } , where h t = f θ ( h t -1 , z t -1 , a t -1 ) is a deterministic recurrent state capturing temporal dependencies, and z t is a stochastic state encoding uncertainty about the current observation. The separation of deterministic ( h t ) and stochastic ( z t ) states enables long-horizon 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 o t to infer z t , ensuring alignment with real-world dynamics (posterior inference); and (ii) during imagination, the model predicts future states without access to observations, sampling ˆ z t to generate hypothetical trajectories (prior prediction). The world model reconstructs observations ˆ o , predicts rewards ˆ r t , and estimates episode continuations ˆ n t from the latent state s t . These components are learned via the following structured models: posterior state representations (encoder) z t ∼ q θ ( z t | h t , o t ) , prior state representations ˆ z t ∼ p θ (ˆ z t | h t ) , reward predictor ˆ r t ∼ p θ (ˆ r t | h t , z t ) , continue predictor ˆ n t ∼ p θ (ˆ n t | h t , z t ) , and decoder ˆ o t ∼ p θ (ˆ o t | h t , z t ) .

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 τ ∼ π ϕ 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 in fixed environments, its reliance on static latent states limits adaptability to contextual variations, such as changes in gravity or friction. To address this, we introduce Dynamics-Aligned Latent Imagination (DALI), which extends DreamerV3 with a self-supervised context encoder that learns explicit context representations from interaction histories.

4.2.1 Forward Dynamics Alignment

The context encoder g φ maps a history of observations and actions ( o t -K : t , a t -K : t -1 ) to a context representation z t = g φ ( o t -K : t , a t -K : t -1 ) , where K defines the temporal window. To align z t with environmental transition dynamics, we optimize a forward dynamics loss:

where f w φ ( o t , a t , z t ) predicts the next observation ˆ o t +1 given the current observation o t , action a t , and context z t . Minimizing L FD ensures that z t encodes contextual factors critical for accurate dynamics prediction, such as variations in physical parameters. The parameters φ of g φ and f w φ are trained jointly, enabling the encoder to capture essential dynamics 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 z t = g φ ( o t -K : t , a t -K : t -1 ) with the DreamerV3 world model's posterior state z t ∼ q θ ( z t | h t , o t ) . The cross-modal loss enforces bidirectional alignment between z t and z t :

where W z and W z are linear maps between the context and state spaces. This bidirectional reconstruction leverages z t 's observation-informed representation, which captures instantaneous dynamics relevant to the current context. By aligning with z t instead of the full latent state s t = { h t , z t } , z t avoids encoding redundant trajectory-specific information from the deterministic h t , which could impair generalization. The bidirectional constraints prevent degenerate solutions (e.g., z t collapsing to a constant) by enforcing invertibility, ensuring z t remains a rich, context-specific representation consistent with the z t 's latent dynamics. The total loss combines both objectives:

where λ cross ∈ { 0 , 1 } 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 z t .

Shallow Integration. This approach appends z t to the observation embedding in the world model's encoder, modifying it to z t ∼ q θ ( z t | h t , o t , z t ) . All other world model components, including the sequence model h t = f θ ( h t -1 , z t -1 , a t -1 ) , predictors ( ˆ z t , ˆ r t , ˆ n t , ˆ o t ), and actor-critic networks ( π ϕ ( a τ | s τ ) , v ψ ( s t ) ), remain unchanged. This lightweight modification enriches latent states with dynamics-aware context, enabling the world model to adapt implicitly to contextual variations.

Deep Integration. For deeper context awareness, z t is integrated into the world model and policy as follows: h t = f θ ( h t -1 , z t -1 , a t -1 , z t ) , ˆ r t ∼ p θ (ˆ r t | h t , z t , z t ) , ˆ n t ∼ p θ (ˆ n t | h t , z t , z t ) , ˆ o t ∼ p θ (ˆ o t | h t , z t , z t ) . The posterior z t ∼ q θ ( z t | h t , o t ) and prior ˆ z t ∼ p θ (ˆ z t | h t ) access z t indirectly through h t . The actor and critic explicitly condition on z t to optimize imagined trajectories over horizon H : a τ ∼ π ϕ ( a τ | s τ , z t ) , v ψ ( s t , z t ) ≈ E π ( ·| s τ , z t ) ∑ t + H τ = t γ τ -t r τ . This ensures that policy optimization explicitly accounts for contextual variations, enabling context-conditioned imagination.

Training. Training the context encoder, the loss (3) aligns z t = g φ ( o t -K : t , a t -K : t -1 ) with environmental dynamics and requires careful handling of the DreamerV3 world model's recurrence. Both integration strategies unroll the world model over a K -length window, initializing h t -K = 0 (or a learned initial state), generating z τ = g φ ( o τ -K : τ , a τ -K : τ -1 ) for τ = t -K to t . In Shallow Integration, z τ ∼ q θ ( z τ | h τ , o τ , z τ ) and h τ +1 = f θ ( h τ , z τ , a τ ) ; in Deep Integration, z τ ∼ q θ ( z τ | h τ , o τ ) and h τ +1 = f θ ( h τ , z τ , a τ , z τ ) . For both, gradients through h τ and z τ are stopped in the recurrent dynamics, and through h τ in the world model's encoder, preventing updates to θ . In Shallow Integration, z τ gradients are preserved in the world model's encoder and losses L FD and L cross , allowing updates to φ , W z , and W z . In Deep Integration, z τ gradients are stopped in the recurrent dynamics and preserved only in L FD and L cross . This decouples context learning from the world model's recurrence, ensuring z t captures relevant dynamics for both strategies. For detailed pseudocode, see Algorithms 1 and 2 in Appendix B for Shallow Integration, and Algorithms 3 and 4 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 a cMDP framework with continuous contexts c (e.g., physical parameters such as gravity or actuator strength), sampled i.i.d. per episode from a distribution p C . Observations are noisy (e.g., o t = s t + η t , η t ∼ N (0 , σ 2 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 β -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. β -mixing captures realistic decorrelation in RL tasks, enabling short interaction windows to capture c . An exploratory policy further ensures that distinct contexts yield distinguishable trajectories, a common feature in RL training.

Domain randomization trains DreamerV3 across diverse contexts, with c ∼ p C , using h RSSM t = f θ ( h t -1 , z t -1 , a t -1 ) and z t ∼ q θ ( z t | h t , o t ) . The recurrent state h RSSM t implicitly accumulates information about c 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 h RSSM t can gradually encode after observing enough transitions. However, this strategy has limitations. The recurrent state h RSSM t 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 c 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 h RSSM t needs a long temporal window to reliably disambiguate c .

DALI's context encoder decouples context inference from dynamics modeling. It functions as a specialized module that learns a representation z t capturing context information from short, local histories, allowing its recurrent state h DALI t to focus on dynamics and reducing the burden on the GRU. The β -mixing property ensures that K transitions suffice to encode c , as the dynamics' dependence on past states fades rapidly. This enables faster adaptation compared to DreamerV3's h RSSM t . We formalize this intuition for the case of Deep Integration, where the context encoder enhances the RSSM's context awareness. The functions h ( · ) and I ( · ; · ) denote the differential entropy and mutual information, respectively [Thomas and Cover, 1999].

Theorem 1. In a cMDP with β -mixing and Lipschitz dynamics, DALI's context encoder z t captures near-optimal context information, I ( c ; z t ) ≥ (1 -δ ) h ( c ) for δ ∈ (0 , 1) , using N = O (1 /δ 2 ) windows of K = Ω(log(1 /δ ) /λ ) transitions, where λ is the mixing rate. Moreover, DALI's RSSM retains more context information than DreamerV3's, I ( c ; h DALI t ) ≥ I ( c ; h RSSM t ) -ϵ ( K ) , with ϵ ( K ) = O ( e -λK/ 2 ) . Compared to DreamerV3's RSSM processing full episodes of T ≫ K transitions, DALI achieves a sample complexity gain of O ( T/K ) .

For the formal statement and proof of Theorem 1, see Appendix A. The β -mixing property ensures that a short window of K = Ω(log(1 /δ ) /λ ) transitions is sufficient to achieve high context fidelity. For instance, when δ = 0 . 01 and the typical mixing rate in DMC tasks is λ ≈ 0 . 1 , a window of K ≈ 64 steps results in a conditional entropy error bounded by δ ′ ≈ 0 . 1 . DALI requires N = O (1 /δ 2 ) such windows, totaling O ( K/δ 2 ) transitions. In contrast, DreamerV3's h RSSM t , constrained by the GRU's finite capacity and observation noise, loses context information and requires O ( T/δ 2 ) transitions per episode (e.g., T = 1000 ), where T ≫ K . The sample complexity gain of O ( T/K ) reflects DALI's efficiency in exploiting local histories. Furthermore, z t achieves superior fidelity, with I ( c ; z t ) approaching h ( c ) within δ , while DreamerV3's RSSM incurs a larger information loss, ϵ ( K ) = O ( e -λK/ 2 ) , due to its bottleneck. While the theorem assumes Deep Integration, this sample complexity gain remains consistent across Deep and Shallow Integration, as z t 's efficiency derives from its training on K -length windows, which is identical in both configurations. These results highlight the necessity and efficiency of DALI's context encoder in enabling robust generalization.

6 Experiments and Analysis

In Sections 6.1 and 6.2, we evaluate DALI's ability to generalize in a zero-shot manner across unseen context variations. In Section 6.3, we show that the dynamics-aligned context encoder learns a structured latent representation, where perturbations to individual dimensions produce physically plausible counterfactuals (e.g., shorter ball swings for higher imagined gravity). Our code is available at https://github.com/frankroeder/DALI .

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-unaware methods, denoted DALI-S/Dχ /blank, employ either Shallow (S) or Deep (D) integration with the inferred context z t (see Section 4.3). These models are trained using either the forward dynamics loss alone (1) (denoted as 'blank', e.g., DALI-D), or the cross-modal regularized objective (2) (denoted as χ , e.g., DALI-Sχ ). We compare our models to the contextunaware 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 c [Prasanna et al., 2024]. For further details, see Appendix A.1.

Setup. We train our methods using the small variant of DreamerV3 [Hafner et al., 2025], with a transformer-based context encoder [Vaswani et al., 2017], for 200K timesteps (Ball-in-Cup) or 500K timesteps (Walker) across 10 random seeds, following the setup of Prasanna et al. [2024]. Hyperparameters and architectural details are provided in Appendix C.

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 et 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-Cup, the context parameters are gravity (training: [4 . 9 , 14 . 7] , evaluation: [0 . 98 , 4 . 9) ∪ (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: [0 . 1 , 0 . 5) ∪ (1 . 5 , 2 . 0] , default: 1.0). These OOD ranges, particularly in Ball-in-Cup, pose major generalization challenges, while Walker's 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 (PoI) metrics 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. PoI quantifies the probability that a 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 PoI across 10 seeds in Figure 1, comparing DALI-S, DALI-Sχ , Dreamer-DR, cRSSM-S, and cRSSM-D for Featurized and Pixel observations on the Ball-in-Cup and Walker Walk tasks.

In Ball-in-Cup's Interpolation, all methods achieve high IQM ( 0 . 92 -0 . 95 ), with DALI-Sχ competitive ( 0 . 9490 Featurized, 0 . 9440 Pixel). In Extrapolation, DALI-Sχ excels, outperforming Dreamer-DR by 87 . 9% (Featurized, IQM 0 . 3720 ) to 96 . 4% (Pixel, IQM 0 . 2730 ), surpassing cRSSMS by 63 . 9% (Featurized) and 45 . 9% (Pixel), and cRSSM-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

0

.

928 0

.

936 0

.

944 0

.

952

Normalized Score

Dreamer-DR

cRSSM-S

cRSSM-D

DALI-S

DALI-S-

c

Interpolation

0

.

18

0

.

24

0

.

30

0

.

36

Normalized Score

Extrapolation

0

.

48

0

.

56

0

.

64

0

.

72

Normalized Score

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Probability

P(DALI-S-

c

>

Y) (Extrap)

(a) Featurized

0

.

935 0

.

940 0

.

945

Normalized Score

Dreamer-DR

cRSSM-S

cRSSM-D

DALI-S

DALI-S-

c

Interpolation

0

.

12

0

.

18

0

.

24

0

.

30

Normalized Score

Extrapolation

0

.

40

0

.

48

0

.

56

0

.

64

Normalized Score

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Probability

P(DALI-S-

c

>

Y) (Extrap)

Walker Walk

0

.

945

0

.

960

0

.

975

Normalized Score

Dreamer-DR

cRSSM-S

cRSSM-D

DALI-S

DALI-S-

c

Interpolation

0

.

69

0

.

72

0

.

75

0

.

78

Normalized Score

Extrapolation

0

.

800 0

.

825 0

.

850 0

.

875

Normalized Score

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Probability

P(DALI-S

>

Y) (Extrap)

(a) Featurized

0

.

95

0

.

96

0

.

97

Normalized Score

Dreamer-DR

cRSSM-S

cRSSM-D

DALI-S

DALI-S-

c

Interpolation

0

.

69

0

.

72

0

.

75

0

.

78

Normalized Score

Extrapolation

0

.

80

0

.

82

0

.

84

0

.

86

Normalized Score

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

Probability

P(DALI-S

>

Y) (Extrap)

( 0 . 14 -0 . 37 ) reflect the extreme OOD context ranges, such as gravity values of 0 . 98 or 19 . 6 , far from the training range of [4 . 9 , 14 . 7] . In the Mixed regime, DALI-Sχ 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 . 6250 ), leveraging ground-truth context for visual dynamics. DALI-Sχ shows moderate PoI ( > 0 . 6 ) against Dreamer-DR in both modalities for Extrapolation (see Figure 1).

In Walker's Interpolation, methods score high ( 0 . 94 -0 . 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 ([ 0 . 1 , 2 . 0 ] vs. [ 0 . 5 , 1 . 5 ]), reducing generalization demands. In Pixel modality, context-aware cRSSM-S leads (IQM 0 . 7770 ), leveraging ground-truth context for robust visual feature extraction. DALI-S attains an IQM of 0 . 7580 , compared to DALI-Sχ at 0 . 7330 . In 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 PoI ( ∼ 0 . 6 ) against Dreamer-DR and cRSSM-D, while cRSSM-S leads, leveraging ground-truth context for visual dynamics (see Figure 1).

6.2 Generalization Trends Across Environments and Modalities

We provide additional interpretations of the IQM results in Figure 1, focusing on the generalization patterns of DALI-S and DALI-Sχ . We emphasize modality-specific effects, environmental nuances, 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χ consistently outperforms DALI-S in Ball-inCup across all regimes and modalities, indicating that cross-modal regularization ( L cross , see 2) enhances the context encoder's ability to capture pendulum-like dynamics. Aligning z t with the world model's posterior z t 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χ in Walker across most regimes and modalities, indicating that the forward dynamics loss ( L FD, see 1) 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 L cross 's complexity in DALI-Sχ . This underscores the need for task-specific regularization. Ball-in-Cup exhibits larger performance drops from Featurized to Pixel modalities (e.g., DALI-Sχ : 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χ 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 z t 's inferred context generalizes better than ground-truth c , which tends to overfit to training contexts. In Walker Featurized Extrapolation, both DALI-S and DALI-Sχ surpass cRSSM-S and cRSSM-D ( 0 . 7810 , 0 . 7770 vs. 0 . 7020 , 0 . 7490 ), benefiting from z t 's inference of actuator torque scaling. However, in Ball-in-Cup Pixel Mixed, cRSSM-D ( 0 . 6250 ) outperforms DALI-Sχ ( 0 . 6030 ), indicating that Deep Integration of ground-truth c 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χ , consistently performs well, leading to its exclusive use in the results reported here. Shallow Integration incorporates the inferred context z t solely in the world model's encoder, z t ∼ q θ ( z t | h t , o t , z t ) , allowing context information to propagate indirectly to the recurrent state h t = f θ ( h t -1 , z t -1 , a t -1 ) through recurrence. This design regularizes the world model, potentially mitigating overfitting to noisy z t 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 β -mixing property in practice.

6.3 Validating Physically Consistent Counterfactuals

We assess whether structured perturbations to the latent space z induce counterfactual trajectories that adhere to Newtonian physics in the Ball-in-Cup task, validating that z 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 our setup, the learned context representation z is an 8 -dimensional vector, i.e., z ∈ R 8 , produced by the dynamics-aligned encoder. Our goal is to identify which latent dimensions in z dominantly encode these mechanistic factors.

For each latent dimension z j , j = 1 , . . . , 8 , we sample a fixed observation o t from a test episode and use the learned context encoder g φ to infer the representation z t = g φ ( o t -K : t , a t -K : t -1 ) with K = 50 . Encoding o t into the latent state s t = { h t , z t } , we then roll out actions a t : t + H under both the original z and a perturbed z ′ = z + ∆ · e j , where e j is a one-hot vector, and ∆ is the standard deviation of z 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 z , T (0) = { ˆ o t , ˆ o t +1 , ..., ˆ o t + H } , and counterfactual trajectory under the perturbed z ′ , T ′ ( j ) = { ˆ o ′ t , ˆ o ′ t +1 , ..., ˆ o ′ t + H } with H = 50 . We repeat this process N = 2500 times to generate 2500

baseline trajectories and 2500 perturbed trajectories for each z j , yielding a total of 5000 trajectory samples: D ( j ) cf = {T (0) 1 , T ′ ( j ) 1 , . . . , T (0) 2500 , T ′ ( j ) 2500 } for z j . Wetrain a binary classifier c ( j ) ν to distinguish between perturbed trajectories T ′ ( j ) (label 1) and baseline trajectories T (0) (label 0). The classifier outputs the probability p n = P ( label = 1 |T n ) that a given trajectory T n ∈ D ( j ) cf is perturbed. High classifier accuracy for a 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 z j and rank the 8 dimensions based on the AUCs, and select the top dimension of z for validating physically consistent counterfactuals. See Appendix D for 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 counterfactual trajectories generated through perturbations to the top-ranked latent dimension ( z 6 ) in DALI-Sχ (Shallow Integration with cross-modal regularization). We show how z 6 encodes coupled gravity-string dynamics in the Ball-in-Cup 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 counterfactual counterparts generated by perturbing the latent context representation in a 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 a 64-step imagination horizon (excluding the final frame for visual clarity).

Leveraging our AUC-based ranking analysis, we identify z 6 as the dominant latent dimension across modalities. To isolate the influence of this dimension on passive dynamics, we generate counterfactual trajectories by perturbing z 6 and enforce zero-action rollouts. This suppresses policy-driven behavior, exposing the system'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 2b and 2c quantify these effects. The counterfactual trajectory (orange) in Figure 2b exhibits reduced amplitude (lower peak Z-position) compared to the original (blue), consistent with the pixel-based evidence of a shorter string. This alignment 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 z 6 . Both modalities reveal that z 6 encodes a coupled relationship between gravity and string length, as a perturbation simultaneously alters both parameters in a physically plausible manner.

7 Discussion and Outlook

We introduced Dynamics-Aligned 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χ achieve substantial performance gains over the context-unaware Dreamer-DR baseline, ranging from +4 . 0% to +96 . 4% across Featurized and Pixel modalities in the DMC Ball-in-Cup and Walker Walk environments. Notably, DALI-Sχ surpasses the ground-truth context-aware 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 to unseen contexts.

Limitations for DALI's Context Inference. Theorem 2 demonstrates that DALI's context encoder achieves near-optimal context inference, I ( c ; z t ) ≥ (1 -δ ) h ( c ) using short windows of length

K = Ω(log(1 /δ ) /λ ) , yielding a sample complexity gain of O ( T/K ) over DreamerV3, assuming β -mixing and Lipschitz dynamics. As an information-theoretic result, it does not address the downstream impact of z t on 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 β -mixing assumption may not hold in 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 of overfitting in high-dimensional observation spaces, are also beyond its scope. Future theoretical work should model how context representations influence policy learning and robustness in complex RL environments.

Integration Strategies and Practical Implications. Empirically, DALI's Shallow Integration strategy, where z t is appended only to the encoder input, consistently delivers strong performance in the tasks we evaluated, particularly under OOD conditions. Its simplicity acts as a regularizer, implicitly propagating z t 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 to policy performance. These directions would deepen our understanding of DALI's integration mechanisms and their practical impact on generalization.

0

5

10

15

20

25

30

35

40

45

50

55

60

(a) Pixel Reconstructions

0

10

20

30

40

50

60

Time

0.1

0.0

0.1

0.2

0.3

Ball Z Position

Original

Counterfactual

(b)

0

10

20

30

40

50

60

Time

1.5

1.0

0.5

0.0

0.5

Ball Z Velocity

Original

Counterfactual

(c)

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 Marc G Bellemare. Deep reinforcement learning at the edge of the statistical precipice. Advances in Neural Information Processing Systems , 2021.

Philip J Ball, Cong Lu, Jack Parker-Holder, and Stephen Roberts. Augmented world models facilitate zero-shot dynamics generalization from a single offline environment. In 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-reinforcement learning. arXiv preprint arXiv:2301.08028 , 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 L Bartlett, Ilya Sutskever, and Pieter Abbeel. RL 2 : Fast reinforcement learning via slow reinforcement learning. arXiv preprint arXiv:1611.02779 , 2016.

Hamid Eghbal-zadeh, Florian Henkel, and Gerhard Widmer. 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 on 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 1126-1135, 2017.

Scott Fujimoto, Pierluca D'Oro, Amy Zhang, Yuandong Tian, and Michael Rabbat. Towards general-purpose model-free reinforcement learning. In International Conference on Learning Representations , 2025.

Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak 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 to control: Learning behaviors by latent imagination. In International Conference on Learning Representations , 2020.

Danijar Hafner, Timothy P 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.02259 , 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 Research , 76:201-264, 2023.

Kimin Lee, Younggyo Seo, Seunghyun Lee, Honglak Lee, and Jinwoo Shin. Context-aware dynamics model for generalization in model-based reinforcement learning. In International Conference on Machine Learning , pages 5757-5766, 2020.

Luckeciano C Melo. Transformers are meta-reinforcement learners. In 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 597-618, 2018.

Mehryar Mohri, Afshin Rostamizadeh, and Ameet 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-reinforcement learning. Advances in Neural Information Processing Systems , 35:27563-27575, 2022.

Anusha Nagabandi, Ignasi Clavera, Simin Liu, Ronald S Fearing, Pieter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt in dynamic, real-world environments 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 Learning Journal , 1, 2024.

Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-policy metareinforcement learning via probabilistic context variables. In International Conference on Machine Learning , pages 5331-5340, 2019.

Emmanuel Rio. Asymptotic Theory of Weakly Dependent Random Processes . Springer, 2017.

Younggyo Seo, Kimin Lee, Ignasi Clavera Gilaberte, Thanard Kurutach, Jinwoo Shin, and Pieter Abbeel. Trajectory-wise multiple choice learning for dynamics generalization in reinforcement learning. Advances in Neural Information Processing Systems , 33:12968-12979, 2020.

Seyed 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 representations. In International Conference on Machine Learning , pages 9767-9779, 2021.

Shagun Sodhani, Franziska Meier, Joelle Pineau, and Amy Zhang. Block contextual MDPs for continual learning. 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 arXiv:1801.00690 , 2018.

Joy A. Thomas and Thomas M. Cover. Elements of Information Theory . John Wiley & Sons, 1999.

Josh Tobin, Rachel Fong, Alex Ray, Jonas Schneider, Wojciech Zaremba, and Pieter Abbeel. Domain randomization 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, Łukasz 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. Hyperdynamics: 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.

Tianhe Yu, Deirdre Quillen, Zhanpeng He, Ryan Julian, Karol Hausman, Chelsea Finn, and Sergey Levine. Meta-world: A benchmark and evaluation for multi-task and meta reinforcement learning. In Conference on Robot Learning , pages 1094-1100, 2020.

Luisa Zintgraf, Kyriacos Shiarli, Vitaly Kurin, Katja Hofmann, and Shimon Whiteson. Fast context adaptation via meta-learning. In International Conference on Machine Learning , pages 7693-7702, 2019.

A Formal Results and Proofs

Theorems 2 and 5 provide formal underpinnings for Theorem 1 in the main text, establishing the necessity and efficiency of DALI's context encoder with Deep Integration (see Section 4.3) compared to 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 z t , against DreamerV3 with domain randomization, using recurrent state h t , in a cMDP. The theorem establishes: (1) DALI's z t captures at least as much mutual information with context c as DreamerV3's h t , up to an error ϵ ( K ) = O ( e -λK/ 2 ) h ( c ) that decays exponentially with window size K ; (2) DALI achieves nearoptimal I ( c ; z t ) ≥ (1 -δ ) h ( c ) with N = O (1 /δ 2 ) samples of K = Ω(log(1 /δ ) /λ ) transitions, leveraging shorter windows compared to DreamerV3'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 = ( S , A , O , C , R , P , E , µ, p C ) with:

Continuous context c ∈ C ⊆ R d , static within episodes, drawn i.i.d. as c ∼ p C , with h ( c ) < ∞ .

Lipschitz dynamics: ∃ L > 0 such that ∀ s, s ′ ∈ S , a ∈ A , c, c ′ ∈ C ,

Observation noise: The observation function E satisfies o t ∼ N ( s t , σ 2 I ) , with fixed variance σ 2 > 0 .

β -mixing: The process τ t = ( o t , a t ) is β -mixing with

where σ ( τ s , s ≤ t ) denotes the sigma-algebra generated by { τ s : s ≤ t } , and similarly for σ ( τ s , s ≥ t + K ) , for constants C, λ > 0 .

̸

Sufficient exploration: The policy ensures that for any c = c ′ ∈ C , there exists a constant α > 0 such that KL( p ( o t -K : t , a t -K : t -1 | c ) ∥ p ( o t -K : t , a t -K : t -1 | c ′ )) ≥ αK .

Bounded log-likelihood ratios: ∣ ∣ ∣ log p ( o i ,a i | c ′ ,τ t -K : i -1 ) p ( o i ,a i | c,τ t -K : i -1 ) ∣ ∣ ∣ ≤ M , ensured by restricting observations and actions to compact sets or imposing suitable policy constraints.

Universal approximator: DALI's g φ (neural network) approximates any continuous function on compact domains.

Training data: Both models are trained on trajectories from M , with contexts c ∼ p C i.i.d. per episode, using an exploratory policy. DreamerV3 with domain randomization trains its RSSM on episodes ( o 1: T , a 1: T -1 ) of length T . DALI trains its context encoder g φ on K -length windows ( o t -K : t , a t -K : t -1 ) with loss L FD = E [ ∥ o t +1 -f w φ ( o t , a t , z t ) ∥ 2 2 ] , and its RSSM on full episodes, conditioning on z t .

Both models use a GRU f θ with finite parameters for their recurrent state updates.

Let h t = f θ ( h t -1 , z t -1 , a t -1 ) be DreamerV3's recurrent state, with z t ∼ q θ ( z t | h t , o t ) , and z t = g φ ( o t -K : t , a t -K : t -1 ) DALI's context encoding, with h t = f θ ( h t -1 , z t -1 , a t -1 , z t ) . Then:

Context identifiability : For any K , there exists g φ such that, within an episode:

Sample complexity : For K = Ω ( log(1 /δ ) λ ) and δ ∈ (0 , 1) , DALI achieves

with N = O ( 1 δ 2 ) samples of K transitions, vs. a hypothetical DreamerV3 context estimator g ψ ( o 1: T , a 1: T -1 ) requiring N = O ( 1 δ 2 ) samples of T transitions.

Remark 3 (Non-trivial sample complexity gain of DALI's Context encoder) . The sample complexity gain of O ( T/K ) for DALI over DreamerV3 with domain randomization, as shown in Theorem 2, arises from DALI's use of K -length windows versus DreamerV3's full episodes of length T . While this gain appears to scale with input lengths, it is non-trivial: Both models achieve N = O (1 /δ 2 ) samples, but DALI's window size K = Ω(log(1 /δ ) /λ ) leverages the β -mixing property (Lemma 6) to ensure h ( c | τ t -K : t ) ≤ C ′ e -λ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 λ and desired accuracy δ .

Remark 4 (Consistency of DALI's sample complexity across integration strategies) . The sample complexity advantage of DALI established in Theorem 2 for Deep Integration extends to Shallow Integration (Section 4.3), since the context encoder z t is trained identically on K -length windows in both configurations. Incorporating the cross-modal loss L cross in (2) aligns z t with z t , which may allow the context representation z t to leverage structured priors from the world model, potentially enhancing its robustness. This addition does not affect sample complexity, as L cross operates on the same K -length windows (Algorithm 2). Recurrent unrolls are limited to K steps, and φ updates depend solely on the current window, with detached states blocking backpropagation beyond it. The world model parameters ( θ ) are frozen during context learning (inputs to f θ and q θ are detached), so gradient stopping isolates φ and preserves the O ( T/K ) efficiency.

Theorem 5 proves that DALI's sequence model, incorporating a context encoder z t , reduces the information bottleneck in cMDPs compared to DreamerV3's RSSM. Specifically, it shows that DALI's recurrent state retains more context information, satisfying I ( c ; h DALI t ) ≥ I ( c ; h RSSM t ) -ϵ ( K ) , where ϵ ( K ) = O ( e -λK/ 2 ) h ( c ) , leveraging the encoder's ability to efficiently capture context over short windows.

Theorem 5 (Context encoder reduces information bottleneck) . Consider a cMDP M = ( S , A , O , C , R , P , E , µ, p C ) with:

Continuous context c ∈ C ⊆ R d , static within episodes, drawn i.i.d. as c ∼ p C , with h ( c ) < ∞ .

Lipschitz dynamics: ∃ L > 0 such that ∀ s, s ′ ∈ S , a ∈ A , c, c ′ ∈ C ,

Observations o t = s t + η t , η t ∼ N (0 , σ 2 I ) , σ 2 > 0 .

β -mixing with β ( K ) ≤ Ce -λK , exploratory policy with KL( p ( τ t -K : t | c ) ∥ p ( τ t -K : t | c ′ )) ≥ αK , α > 0 .

Context encoder g φ is a universal approximator, trained with forward dynamics loss L FD = E [ ∥ o t +1 -f w φ ( o t , a t , z t ) ∥ 2 2 ] , satisfying h ( c | z t ) ≤ C ′ e -λK h ( c ) + δ ′ , δ ′ = O ( e -λK/ 2 ) (Lemma 7).

Both models use a GRU f θ with finite parameters.

Let h RSSM t = f θ ( h t -1 , z t -1 , a t -1 ) be DreamerV3's recurrent state, with z t ∼ q θ ( z t | h t , o t ) , and h DALI t = f θ ( h t -1 , z t -1 , a t -1 , z t ) be DALI's recurrent state. Then, for any window size K , DALI's sequence model satisfies:

where ϵ ( K ) = O ( e -λK/ 2 ) h ( c ) , implying that the context encoder z t reduces the information bottleneck compared to DreamerV3's RSSM.

For β -mixing sequences with β ( K ) ≤ Ce -λK , blocks of K -separated transitions are approximately independent. Lemma 6 establishes that, in a β -mixing cMDP, the conditional entropy of the context c given a history τ t -K : t decays exponentially as h ( c | τ t -K : t ) ≤ C ′ e -λK h ( c ) . This result is crucial for Theorems 2 and 5, 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 a cMDP with context c ∼ p C , observation noise o t = s t + η t , η t ∼ N (0 , σ 2 I ) and history τ t -K : t = ( o t -K : t , a t -K : t -1 ) , where the process τ t = ( o t , a t ) is β -mixing with β ( K ) ≤ Ce -λK for constants C, λ > 0 , bounding dependence errors in tail probabilities by β ( K ) , log-likelihood ratios are bounded as ∣ ∣ ∣ log p ( o i ,a i | c ′ ,τ t -K : i -1 ) p ( o i ,a i | c,τ t -K : i -1 ) ∣ ∣ ∣ ≤ M , the policy ensures exploration such that for any c = c ′ ∈ C , there exists a constant α ≥ M √ 2 λ satisfying KL( p ( τ t -K : t | c ) | p ( τ t -K : t | c ′ )) ≥ αK , and c has bounded density on C ⊆ R d with differential entropy h ( c ) < ∞ . Then, there exists a constant C ′ > 0 such that:

Proof of Lemma 6. We show that uncertainty about c decreases exponentially with K , enabling DALI's z t to infer c .

By assumption, the process τ t = ( o t , a t ) is β -mixing, with β ( K ) ≤ Ce -λK . The condition KL( p ( τ t -K : t | c ) ∥ p ( τ t -K : t | c ′ )) ≥ αK , with α ≥ M √ 2 λ , ensures distinct histories. We need h ( c | τ t -K : t ) ≤ C ′ e -λK h ( c ) , so we show:

̸

Discretize C into N points, log N ≈ h ( c ) . Let ˆ c = arg max c ′ p ( τ t -K : t | c ′ ) and let P e = P (ˆ c = c | τ t -K : t ) be the probability of incorrectly estimating c . Fano's inequality gives [Thomas and Cover, 1999]:

̸

Error P e = P (ˆ c = c | c ) occurs if there exists any c ′ = c such that p ( τ t -K : t | c ′ ) > p ( τ t -K : t | c ) . By the union bound:

̸

̸

̸

̸

For c ′ = c , define: Z = log p ( τ t -K : t | c ′ ) p ( τ t -K : t | c ) . Then Z > 0 corresponds to the maximum likelihood estimator ˆ c incorrectly favoring c ′ over c . The probability of this error under context c is P ( Z > 0 | c ) . By the exploration assumption:

By the chain rule in the cMDP, p ( τ t -K : t | c ) = ∏ t -1 i = t -K p ( o i , a i | τ t -K : i -1 , c ) · p ( o t | τ t -K : t -1 , c ) , where τ t -K : i -1 = ( o t -K : i -1 , a t -K : i -1 ) .

Thus, Z = ∑ t -1 i = t -K log p ( o i ,a i | τ t -K : i -1 ,c ′ ) p ( o i ,a i | τ t -K : i -1 ,c ) +log p ( o t | τ t -K : t -1 ,c ′ ) p ( o t | τ t -K : t -1 ,c ) .

Define Z i = log p ( o i ,a i | τ t -K : i -1 ,c ′ ) p ( o i ,a i | τ t -K : i -1 ,c ) for i = t -K to t -1 , and Z t = log p ( o t | τ t -K : t -1 ,c ′ ) p ( o t | τ t -K : t -1 ,c ) . So Z = ∑ t i = t -K Z i . Z i measures how much more (or less) likely ( o i , a i ) is under context c ′ compared to c , given past observations and actions. By the bounded log-likelihood assumption, | Z i | ≤ M . For independent Z i , Hoeffding's inequality [Boucheron et al., 2013] bounds the tail probability for bounded random variables:

Since τ t is β -mixing with β ( K ) ≤ Ce -λK , the dependence adds an error β ( K ) , as β ( K ) bounds the deviation from independence between events separated by K time steps:

Set κ = α 2 2 M 2 . Since α ≥ M √ 2 λ , we have κ ≥ λ , so exp( -κK ) ≤ e -λK . Thus:

Sum over N -1 contexts:

Apply to (4):

For small β ( K ) , h ( C 1 β ( K )) ≈ 0 , so:

With β ( K ) ≤ Ce -λK ,

Lemma 7 provides a generalization bound for DALI's context encoder, showing that the conditional entropy h ( c | z t ) is bounded by the entropy given the history plus a small error, h ( c | z t ) ≤ h ( c | τ t -K : t )+ δ ′ , with δ ′ = O ( √ δ ) and δ = Ce -λK . This lemma is pivotal for Theorems 2 and 5, as it ensures that the encoder z t effectively captures context information with a sample complexity of N = O (1 /δ 2 ) , facilitating DALI's efficiency and reduced bottleneck.

Lemma 7 (DALI's generalization bound for context inference) . Consider a cMDP satisfying Theorem 2's assumptions: context c ∼ p C with bounded density on C ⊆ R d and finite entropy h ( c ) < ∞ , observation noise o t = s t + η t , η t ∼ N (0 , σ 2 I ) , β -mixing with β ( K ) ≤ Ce -λK for constants C, λ > 0 , exploratory policy ensuring KL( p ( τ t -K : t | c ) ∥ p ( τ t -K : t | c ′ )) ≥ αK for some α > 0 , Lipschitz dynamics with ∥P c ( s ′ | s, a ) -P c ′ ( s ′ | s, a ) ∥ 1 ≤ L ∥ c -c ′ ∥ 2 , and universal approximator g φ . DALI's context encoding z t = g φ ( o t -K : t , a t -K : t -1 ) , trained with the forward dynamics loss L FD = E [ ∥ o t +1 -f w φ ( o t , a t , z t ) ∥ 2 2 ] , satisfies:

with N = O ( 1 δ 2 ) samples, where δ = Ce -λK .

Proof of Lemma 7. The context encoder g φ minimizes the forward dynamics loss:

where τ t -K : t = ( o t -K : t , a t -K : t -1 ) , o t = s t + η t , and η t ∼ N (0 , σ 2 I ) . The empirical loss over N samples is:

To benchmark L FD, consider an ideal predictor with access to the true context c . Since o t +1 = s t +1 + η t +1 , η t +1 ∼ N (0 , σ 2 I ) , and o t = s t + η t , the predictor f ( o t , a t , c ) uses noisy observations. The loss is:

For d s -dimensional η t +1 , E [ ∥ η t +1 ∥ 2 2 ] = d s σ 2 . The ideal predictor f ∗ ( o t , a t , c ) = E [ s t +1 | o t , a t , c ] , the expected next state given o t , a t , c , minimizes L FD by setting f ∗ to the best estimate of s t +1 , yielding E [ ∥ o t +1 -f ∗ ( o t , a t , c ) ∥ 2 2 ] = E [ ∥ s t +1 -E [ s t +1 | o t , a t , c ] ∥ 2 2 ] + d s σ 2 , where the first term is the variance of s t +1 and the second is the noise variance. By the universal approximation property, f w φ can approximate f ∗ if z t encodes c .

Consider the loss function for fixed φ, w :

with | l | ≤ M (by bounded S , A ) and L l -Lipschitz in φ, w . The function class F = { ( o i , a i , o i -K : i , a i -K : i -1 ) ↦→ f w φ ( o i , a i , g φ ( o i -K : i , a i -K : i -1 )) } has Rademacher complexity R N ( F ) ≤ C F √ N , so the loss class L has:

For i.i.d. samples and a fixed φ, w , the generalization error is bounded by:

By McDiarmid's inequality, since changing one sample affects ˆ L FD by at most 2 M N , with probability at least 1 -δ [Mohri et al., 2018]:

For β -mixing data with β ( K ) ≤ Ce -λK , the deviation from i.i.d. adds an error bounded by Ce -λK to the expected loss for sequences spaced by K steps:

Set the generalization error to ϵ 2 = O ( δ ) , where δ = Ce -λK , to match the β -mixing decay. Solve:

The first term requires: 2 C F L l √ N ≤ C ′′ δ = ⇒ N ≥ ( 2 C F L l C ′′ δ ) 2 = O ( 1 δ 2 ) .

The second term requires: M √ 2 log(1 /δ ) N ≤ C ′′ δ = ⇒ N ≥ 2 M 2 log(1 /δ ) ( C ′′ δ ) 2 = O ( log(1 /δ ) δ 2 ) .

The third term Ce -λK = δ ≤ C ′′ δ holds for C ′′ ≥ 1 .

The first term dominates, so N = O ( 1 δ 2 ) suffices.

If E [ ˆ L FD ] ≤ E [ ∥ o t +1 -f ∗ ( o t , a t , c ) ∥ 2 2 ] + ϵ 2 , then:

Given the Markov chain c → τ t -K : t → z t , the entropy difference is:

The cMDP's Lipschitz dynamics, ∥P c ( s ′ | s, a ) -P c ′ ( s ′ | s, a ) ∥ 1 ≤ L ∥ c -c ′ ∥ 2 , and β -mixing ensure p ( c | τ t -K : t ) is L p -Lipschitz in Wasserstein distance, as small changes in τ t -K : t reflect proportional changes in c . Assume there exists a mapping ψ : Z → C such that the context encoder satisfies:

where Z is the encoder's output space. Then:

Thus:

From Lemma 6, h ( c | τ t -K : t ) ≤ Ce -λK h ( c ) = O ( δ ) h ( c ) . Since h ( c ) < ∞ and √ δ ≥ δ for small δ , the bound becomes:

yielding δ ′ = O ( √ δ ) with N = O ( 1 δ 2 ) samples.

Lemma 8 demonstrates that DreamerV3's RSSM suffers from a persistent information bottleneck in cMDPs, with the conditional entropy h ( c | h t ) ≥ κh ( c ) for a constant κ ∈ (0 , 1) . This result is essential for Theorems 2 and 5, as it establishes a baseline limitation in DreamerV3's ability to retain context information, against which DALI's superior performance in context identifiability and bottleneck reduction is compared.

Lemma8 (Information bottleneck in DreamerV3's RSSM) . Consider DreamerV3's RSSM in a cMDP with context c ∼ p C with bounded density on C ⊆ R d and finite entropy h ( c ) < ∞ , observation noise o t = s t + η t , η t ∼ N (0 , σ 2 I ) , Lipschitz dynamics with ∥P c ( s ′ | s, a ) -P c ′ ( s ′ | s, a ) ∥ 1 ≤ L ∥ c -c ′ ∥ 2 , and recurrent state h t = f θ ( h t -1 , z t -1 , a t -1 ) , where z t ∼ q θ ( z t | h t , o t ) , and f θ is a GRU with finite parameters. Then, there exists κ ∈ (0 , 1) such that:

Proof of Lemma 8. The context c determines the cMDP's dynamics P c ( s ′ | s, a ) . Since h t depends on the trajectory τ 1: t = ( o 1: t , a 1: t -1 ) and latent variables z 1: t -1 , we have a Markov chain c → τ 1: t → h t , so p ( h t | τ 1: t , c ) = p ( h t | τ 1: t ) . By the data processing inequality (DPI) [Thomas and Cover, 1999]:

Since o 1: t = s 1: t + η 1: t , and a 1: t -1 depend on o 1: t -1 , we have:

The GRU f θ , with finite parameters and Lipschitz continuity (due to bounded weights and standard GRU operations), produces h t ∈ R d , where d is a fixed dimension. The GRU is trained to maximize the ELBO:

which optimizes h t to predict observations o t , rewards r t , and termination signals n t . The input τ 1: t includes:

Context information: c affects s 1: t through dynamics.

Noise information: η 1: t ∼ N (0 , σ 2 I t · d s ) , with d s being the dimension of the state space S . η 1: t contributes significant entropy, h ( η 1: t ) = t · d s 2 ln(2 πeσ 2 ) , which grows linearly with t .

State transients: s 1: t includes dynamic behaviors not directly tied to c .

The GRU's finite parameters and fixed dimension d limit its capacity, forcing h t to compress τ 1: t , potentially discarding context information.

Assume h t is Gaussian with covariance Σ h , where Var ( h i t ) ≤ V for i = 1 , . . . , 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 h t ∼ N ( µ c , Σ h ) , where µ c depends on c , and noise variance at most σ 2 , the mutual information is bounded by [Thomas and Cover, 1999]:

The GRU's compression limits I ( c ; h t ) relative to h ( c ) for all h ( c ) ≥ 0 , as follows. The ELBO's KL-term KL( q θ ( z t | h t , o t ) ∥ p θ ( z t | h t )) regularizes q θ ( z t | h t , o t ) , reducing the information about o t (and thus c ) in z t . Since h t depends on z 1: t -1 , this limits context information in h t . The GRU's finite parameters and dimension d further constrain capacity. Let the channel capacity be reduced by a factor α ∈ (0 , 1) , reflecting this compression:

where α < 1 (e.g., α = 1 -ϵ , with ϵ > 0 depending on the GRU's architecture).

To satisfy the lemma, we seek a constant κ ∈ (0 , 1) such that h ( c | h t ) ≥ κh ( c ) for all p C . Since h ( c | h t ) = h ( c ) -I ( c ; h t ) and I ( c ; h t ) ≤ α · d 2 ln ( 1 + dV σ 2 ) , define:

For h ( c ) > 0 , we want h ( c ) -I ( c ; h t ) ≥ κh ( c ) . Using I ( c ; h t ) ≤ αB , this holds if h ( c ) -αB ≥ κh ( c ) , or κ ≤ 1 -αB h ( c ) . Thus, define:

If h ( c ) > 0 , then I ( c ; h t ) ≤ αB < B (since α < 1 ), and with κ = 1 -αB h ( c ) , we have h ( c | h t ) = h ( c ) - I ( c ; h t ) ≥ h ( c ) -αB = ( 1 -αB h ( c ) ) h ( c ) = κh ( c ) . Since αB > 0 and h ( c ) > 0 , then κ = 1 -αB h ( c ) < 1 . If h ( c ) ≥ αB , then κ ≥ 0 . If h ( c ) < αB , then κ < 0 , but h ( c | h t ) ≥ 0 ≥ κh ( c ) , so the inequality holds.

If h ( c ) = 0 , then I ( c ; h t ) = 0 , so h ( c | h t ) = 0 . Choosing κ = 1 2 ∈ (0 , 1) , we have h ( c | h t ) = 0 ≥ 1 2 · 0 = κh ( c ) , and any κ ∈ (0 , 1) would suffice.

Thus, for all p C , there exists κ ∈ (0 , 1) such that h ( c | h t ) ≥ κh ( c ) .

Lemma 9 proves that DALI's sequence model, incorporating the context encoder z t , reduces the information bottleneck compared to DreamerV3, with h ( c | h t ) ≤ κ ′ h ( c ) , where κ ′ = Ce -λK + O ( e -λK/ 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 9 (Information bottleneck in DALI's sequence model) . Consider DALI's sequence model in a cMDPwith context c ∼ p C with bounded density on C ⊆ R d and finite entropy h ( c ) < ∞ , observation noise o t = s t + η t , η t ∼ N (0 , σ 2 I ) , Lipschitz dynamics with ∥P c ( s ′ | s, a ) -P c ′ ( s ′ | s, a ) ∥ 1 ≤ L ∥ c -c ′ ∥ 2 , context encoding z t = g φ ( o t -K : t , a t -K : t -1 ) , and recurrent state h t = f θ ( h t -1 , z t -1 , a t -1 , z t ) , where z t ∼ q θ ( z t | h t , o t ) , and f θ is a GRU with finite parameters. Suppose that the context encoding satisfies h ( c | z t ) ≤ Ce -λK h ( c ) + δ ′ , where δ ′ = O ( e -λK/ 2 ) and δ = Ce -λK (Lemma 7). Then, there exists κ ′ ∈ (0 , 1) such that:

where κ ′ = Ce -λK + O ( e -λK/ 2 ) , implying a reduced information bottleneck compared to DreamerV3's κ (Lemma 8).

Proof of Lemma 9. The context c determines the cMDP's dynamics P c ( s ′ | s, a ) . DALI's sequence model computes the context encoding z t = g φ ( o t -K : t , a t -K : t -1 ) from trajectory τ t -K : t = ( o t -K : t , a t -K : t -1 ) , and the recurrent state h t = f θ ( h t -1 , z t -1 , a t -1 , z t ) , where z t -1 ∼ q θ ( z t -1 | h t -1 , o t -1 ) and f θ is a Lipschitz continuous GRU.

Since h t = f θ ( h t -1 , z t -1 , a t -1 , z t ) , we have a Markov chain:

conditional on h t -1 , z t -1 , a t -1 , as c affects h t through z t given these variables. By the DPI, we have

By Lemma 7, the context encoding satisfies:

where δ = Ce -λK , so δ ′ = O ( √ δ ) = O ( e -λK/ 2 ) . From Lemma 6, β -mixing implies:

Thus,

Applying the DPI bound, we have

Define

To satisfy h ( c | h t ) ≤ κ ′ h ( c ) , consider two cases. If h ( c ) > 0 , let δ ′ = C ′′ e -λK/ 2 , so:

Since δ ′ h ( c ) = O ( e -λK/ 2 ) , κ ′ = Ce -λK + O ( e -λK/ 2 ) when the first term is selected, and κ ′ ≤ 1 2 always. If h ( c ) = 0 , then I ( c ; h t ) = 0 , so h ( c | h t ) = 0 . Set κ ′ = 1 2 , satisfying h ( c | h t ) = 0 ≤ κ ′ h ( c ) .

Compared to DreamerV3's h ( c | h RSSM t ) ≥ κh ( c ) , κ ∈ (0 , 1) (Lemma 8), DALI's κ ′ → 0 as K →∞ , indicating a reduced information bottleneck.

Proof of Theorem 5. We prove that DALI's recurrent state h DALI t achieves higher mutual information with the context c than DreamerV3's h RSSM t , up to an error ϵ ( K ) = O ( e -λK/ 2 ) h ( c ) .

By Lemma 8, DreamerV3's recurrent state satisfies:

Thus,

By Lemma 9, DALI's recurrent state satisfies:

Thus,

Compare:

Set

since h ( c ) < ∞ . Thus,

As K →∞ , κ ′ → 0 , so I ( c ; h DALI t ) → h ( c ) , while I ( c ; h RSSM t ) ≤ (1 -κ ) h ( c ) < h ( c ) , demonstrating that the context encoder reduces the bottleneck.

Proof of Theorem 2. We prove both parts of the theorem comparing DALI's context encoding z t = g φ ( o t -K : t , a t -K : t -1 ) with DreamerV3's recurrent state h t = f θ ( h t -1 , z t -1 , a t -1 ) .

1. Context identifiability

We show that DALI's context encoding satisfies:

where τ t -K : t = ( o t -K : t , a t -K : t -1 ) is the history, and h t is DreamerV3's recurrent state.

By Lemma 7, DALI's context encoder, trained with the forward dynamics loss:

satisfies

By Lemma 6, the history's conditional entropy is:

Combining,

Thus,

The mutual information for DALI is

By Lemma 8, DreamerV3's recurrent state h t has so that

Compare:

I -I ≥ ----Since κh ( c ) > 0 , set

as the e -λK/ 2 term dominates. Thus,

Additionally, Lemma 9 shows that DALI's recurrent state h t = f θ ( h t -1 , z t -1 , a t -1 , z t ) has a reduced bottleneck:

implying,

This suggests that even DALI's recurrent state retains more context information than DreamerV3's h t , reinforcing the necessity of DALI's context encoder z t , which achieves near-optimal I ( c ; z t ) .

2. Sample complexity

We prove that DALI achieves I ( c ; z t ) ≥ (1 -δ ) h ( c ) with K = Ω ( log(1 /δ ) λ ) and N = O ( 1 δ 2 ) samples of K transitions, while a hypothetical DreamerV3 context estimator g ψ ( o 1: T , a 1: T -1 ) requires N = O ( 1 δ 2 ) samples of T transitions.

DALI's sample complexity From Lemma 7, we have

Combining, we have

To achieve I ( c ; z t ) ≥ (1 -δ ) h ( c ) , we need

Set C ′ e -λK h ( c ) + C ′′ e -λK/ 2 ≤ δh ( c ) . Since e -λK/ 2 dominates, approximate

Solve:

From Lemma 7, achieving δ ′ = O ( e -λK/ 2 ) requires

Solve:

DreamerV3's sample complexity

From Lemma 8, DreamerV3's recurrent state h t has h ( c | h t ) ≥ κh ( c ) , so I ( c ; h t ) ≤ (1 -κ ) h ( c ) , where κ is constant, preventing I ( c ; h t ) ≥ (1 -δ ) h ( c ) for small δ . To enable comparison, assume DreamerV3 with domain randomization trains a hypothetical context estimator g ψ ( o 1: T , a 1: T -1 ) to estimate c over N episodes of length T . By Lemma 6 for K = T :

Training g ψ to minimize E [ g ψ ( o 1: T , a 1: T -1 ) E [ c τ 1: T ] 2 ] requires

Solve:

The total number of transitions is N · T = O ( T δ 2 ) .

Note that DreamerV3's standard RSSM cannot achieve I ( c ; h t ) ≥ (1 -δ ) h ( c ) due to h ( c | h t ) ≥ κh ( c ) (Lemma 8). The estimator g ψ 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 9 indicates that DALI's h t achieves I ( c ; h t ) ≥ (1 -κ ′ ) h ( c ) , with κ ′ = O ( e -λK/ 2 ) , suggesting that DALI's sequence model could approach the desired information level with sufficient K , unlike DreamerV3's h t .

A.1 DALI and Dreamer-based Baselines

A.1.1 World Models

DreamerV3 with Domain Randomization (context-unaware baseline) [Tobin et al., 2017, Hafner et al., 2025]:

DALI-S (DALI with S hallow Integration):

Sequence model:

h t = f θ ( h t - 1 , z t - 1 , a t - 1 )

Encoder:

z t ∼ q θ ( z t | h t , o t )

Dynamics predictor:

ˆ z t ∼ p θ (ˆ z t | h t )

Reward predictor:

ˆ r t ∼ p θ (ˆ r t | h t , z t )

Continue predictor:

ˆ n t ∼ p θ (ˆ n t | h t , z t )

Decoder:

ˆ o t ∼ p θ (ˆ o t | h t , z t ) .

Context encoder:

Sequence model:

h t = f θ ( h t - 1 , z t - 1 , a t - 1 ) ,

Encoder:

z t ∼ q θ ( z t | h t , o t , z t ) ,

Dynamics predictor:

ˆ z t ∼ p θ (ˆ z t | h t ) ,

Reward predictor:

ˆ r t ∼ p θ (ˆ r t | h t , z t ) ,

Continue predictor:

ˆ n t ∼ p θ (ˆ n t | h t , z t ) ,

Decoder:

ˆ o t ∼ p θ (ˆ o t | h t , z t ) .

DALI-D (DALI with D eep Integration):

cRSSM-S (ground-truth context-aware baseline) [Prasanna et al., 2024]:

cRSSM-D (ground-truth context-aware baseline) [Prasanna et al., 2024]:

The baselines cRSSM-S and cRSSM-D correspond to concat-context and cRSSM , respectively, in [Prasanna et al., 2024].

A.1.2 Actor-Critic Models

DreamerV3 with Domain Randomization , DALI-S , and cRSSM-S :

DALI-D :

cRSSM-D :

B Algorithms

Algorithm 2: DALI-Sχ (Shallow Integration with Cross-modal Regularization)

Algorithm 4: DALI-Dχ (Deep Integration with Cross-modal Regularization)

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 (out-of-distribution, OOD) capabilities, we define two distinct context ranges: one for training and interpolation evaluation, and another for OOD evaluation. See Table 1. Additionally, we consider both Featurizedand Pixel-based observation modalities for each environment.

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 is held fixed at its 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 on 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 a 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 influence on oscillation frequency and string length's effect on swing amplitude. Extreme gravity values, such as 0 . 98 or 19 . 6 , significantly alter the ball's trajectory, while short or 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, as the context encoder must adapt to substantial shifts in motion patterns, particularly in the OOD regime.

DMC Walker Walk. This task involves a 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 multijoint coordination and contact interactions. Generalization is less challenging than in Ball-in-Cup, as the actuator strength evaluation range ( 0 . 1 to 0 . 5 and 1 . 5 to 2 . 0 ) is 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.

∪

C.3 Architectural details

Figure 3 shows a high-level illustration of the DALI architecture.

Encoder

Decoder

Encoder

Decoder

Encoder

Decoder

gradient flow

History

Forward Dynamics Loss

Context Encoder. The context encoder g ϕ employs a standard transformer encoder block [Vaswani et al., 2017] to process a sequence of K transitions, ( o t -K : t , a t -K : t -1 ) , and produce the context representation z t ∈ R 8 . The input sequence is fed directly into a dense layer with 256 units, bypassing a 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 output. The attention output is combined with the skip connection via a residual connection, followed by a 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 256-dimensional intermediate representation to z t ∈ R 8 . All hidden layers, including the MLP and attention head, use 256 units and SiLU activation functions, consistent with DreamerV3 [Hafner et al., 2025]. No masking is applied in the self-attention mechanism, as the input sequence is complete, enabling full contextual processing to generate z t .

Forward Model. The forward model f w φ in (1) for dynamics alignment is a two-layer MLP with 128 units and SiLU activations.

W z ∈ R 32 × 8 and W z t ∈ R 8 × 32 in (2) are 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 2 GPUs on 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 5 algorithm variants, 2 environments, and 2 modalities, evaluated over 3 context settings ( single_0 , single_1 , double ), yielding a total of 600 individual training runs. Table 2 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 is approximately 1,051 GPU hours ( ≈ 44 days). In the more typical worst-case scenario with 2 GPUs, wall-clock time increases to approximately 2,101 GPU hours ( ≈ 88 days).

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-small while scaling effectively to high-dimensional tasks. We evaluate it on DMCQuadruped Walk (56D observations, 12D actions), a larger environment than Walker Walk (24D, 6D). The 56D observation space emphasizes the challenge of torque-sensitive locomotion, where increased coordination and a higher number of actuators (12 vs. 6) demand more robust context inference. Trained for 600K timesteps and contextualized with gravity and actuator strength (same ranges as Walker Walk), DALI attains Featurized Extrapolation IQMs of 0 . 326 ± 0 . 043 (DALI-S), and 0 . 389 ± 0 . 027 (DALI-Dχ ), yielding up to 76 . 8% improvement over the context-unaware baseline 0 . 220 ± 0 . 023 (Dreamer-DR), and up to 50 . 8% improvement over the context-aware baselines ( 0 . 258 ± 0 . 028 , cRSSM-D; 0 . 317 ± 0 . 031 , cRSSM-S).

Ball-in-Cup

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

1

.

50

1

.

75

2

.

00

Environment Steps

×

10 5

0

50

100

150

200

250

300

Training Return

DMCBall in Cup - Featurized

DALI-D

DALI-D-

c

DALI-S

DALI-S-

c

Dreamer-DR

cRSSM-D

cRSSM-S

(a) Featurized

0

.

00

0

.

25

0

.

50

0

.

75

1

.

00

1

.

25

1

.

50

1

.

75

2

.

00

Environment Steps

×

10 5

0

50

100

150

200

250

300

Training Return

DMCBall in Cup - Pixel

DALI-D

DALI-D-

c

DALI-S

DALI-S-

c

Dreamer-DR

cRSSM-D

cRSSM-S

Walker Walk

0

1

2

3

4

5

Environment Steps

×

10 5

0

50

100

150

200

250

300

Training Return

DMCWalker Walk - Featurized

DALI-D

DALI-D-

c

DALI-S

DALI-S-

c

Dreamer-DR

cRSSM-D

cRSSM-S

(c) Featurized

0

1

2

3

4

5

Environment Steps

×

10 5

0

50

100

150

200

250

300

Training Return

DMCWalker Walk - Pixel

DALI-D

DALI-D-

c

DALI-S

DALI-S-

c

Dreamer-DR

cRSSM-D

cRSSM-S

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.1). Next, we analyze actuator-driven recovery dynamics in the DMC Walker Walk task, analogous to the counterfactual analysis conducted for the Ball-in-Cup task in Section 6.3.1 (Section D.2).

D.1 AUC Rankings and Significance of Counterfactual Perturbations

To identify latent dimensions that systematically alter imagined dynamics, we train a binary classifier to distinguish trajectories generated under perturbed ( T ′ ( j ) ) and original ( T (0) ) context representations. For each dimension z j , we generate N = 2500 trajectory pairs, perturbing z j by ∆ = σ ( z j ) to induce counterfactuals (see Section 6.3).

We implement an ensemble classifier comprising a support vector machine, an MLP, and an AdaBoost classifier, trained via stratified 5-fold cross-validation. Inputs are trajectory sequences T (0) (label 0) and T ′ ( j ) (label 1). Predictions are aggregated across folds to compute AUC. To ensure robustness, we compute 95% bootstrap confidence intervals for each z j '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 5, 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.

z

1

z

2

z

3

z

4

z

5

z

6

z

7

z

8

0

.

0

0

.

2

0

.

4

0

.

6

0

.

8

AUC

Mean AUC

D.2 Actuator-Driven Recovery and Locomotion in Latent Imagination

Analogous to the Ball-in-Cup analysis (Section 6.3.1), we evaluate DALI's ability to generate physically consistent counterfactuals in the DMC Walker Walk task contextualized with gravity and actuator strength. We show that perturbations to the top-ranked latent dimension z 3 in DALI-S, identified via AUC-based rankings (see Section D.1), influence the Walker's capacity to recover from destabilizing falls and sustain locomotion, indicating that z 3 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 recovery and locomotion behavior. For the Featurized modality, we 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 over a 64-step imagination horizon.

Pixel Modality. Figure 6a contrasts the original (top) and counterfactual (middle) trajectories. Both agents initially fall (frames 0 -20 ), but the perturbed z 3 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 6b 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 z 3 's role in encoding torque dynamics critical for stability.

0

5

10

15

20

25

30

35

40

45

50

55

60

(a) Pixel Reconstructions

0

10

20

30

40

50

60

Time

0.2

0.4

0.6

0.8

1.0

1.2

1.4

Torso Z

Original

Counterfactual

(b) Torso Z Position

人

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 in the abstract and introduction are fully supported by our theoretical and empirical findings, including a comprehensive analysis.

2. Limitations

Question: Does the paper discuss the limitations of the work performed by the authors?

Answer: [Yes]

Justification: Limitations are discussed in the corresponding section of this paper.

3. Theory assumptions and proofs

Question: For each theoretical result, does the paper provide the full set of assumptions and a 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 experimental 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 experiments in a 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 to the data and code, with sufficient instructions 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, hyperparameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?

Answer: [Yes]

Justification: Yes, we provide all details for training the agent in-distribution and testing zero-shot generalization.

7. Experiment statistical significance

Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?

Answer: [Yes]

Justification: For this purpose, we follow established methods for measuring the agent's performance [Agarwal et al., 2021].

8. Experiments compute resources

Question: For each experiment, does the paper provide sufficient information on the computer resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?

Answer: [Yes]

Justification: DALI adds only a small context encoder module to DreamerV3. As reported by Hafner et al. [2025], DreamerV3 can be trained on a single GPU, and our extension has the same requirement. Each trial with a single random seed takes roughly 1-2 days, depending on the observation modality.

9. Code of ethics

Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https://neurips.cc/public/EthicsGuidelines ?

Answer: [Yes]

Justification: No disagreement amoung the team.

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

Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?

Answer: [NA]

Justification: [NA]

12. Licenses for 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]

Justification: This work builds upon and extends the codebase of Prasanna et al. [2024], which is properly credited. All associated licenses and terms of use from the original work are respected.

13. New assets

Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?

Answer: [Yes]

14. Crowdsourcing and research with human subjects

Question: For 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 as details about compensation (if any)?

Answer: [NA]

Justification: [NA]

15. Institutional review board (IRB) approvals or 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

Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in 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]

Peer reviews:
comments: 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.

STRENGTHS AND WEAKNESSES:
**Strengths：**

- DALI outperforms context-unaware baselines on the  contextualized DMC Ball-in-Cup and Walker Walk tasks.
- This paper provides both empirical and theoretical analyses.

**Weaknesses：**

- The readability could be improved, as there is no figure provided to illustrate the pipeline.
- Experiments were conducted on only two tasks, and more experimental validation is needed.
- 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, such as [1], [2], and [3].

References:

[1] Wu et al. "Pre-training Contextualized World Models with In-the-wild Videos for Reinforcement Learning", NeurIPS, 2023.

[2] Li et al. "Open-World Reinforcement Learning over Long Short-Term Imagination", ICLR, 2025.

[3] Fujimoto et al. "MR.Q: Towards General-Purpose Model-Free Reinforcement Learning", ICLR, 2025.

QUESTIONS:
- See the weaknesses.
- I suggest adding the corresponding numerical comparisons for the results in Figures 1 and 2.

LIMITATIONS:
Yes.
TITLE: Official Review
CLARITY: 2
IMPACT: 3
ORIGINALITY: 2
DATE: 2025-06-20

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\nI 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.\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 1, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-25"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis 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.\n\nSTRENGTHS AND WEAKNESSES:\n## 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.\n\nQUESTIONS:\nWould techniques introduced in this work be applicable to improving sim2real transfer of robot locomotion policies, like quadruped running or bipedal walking?\n\nLIMITATIONS:\nLimitations of the proposed method are not sufficiently addressed. Section 7, “Discussion and Outlook”, should include a comprehensive discussion of the limitations of DALI.", "TITLE": "Official Review", "RECOMMENDATION": 5, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 3, "DATE": "2025-06-30"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nReinforcement 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.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\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.\n\nQUESTIONS:\n1. 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?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 3, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-02"}