Ada-Diffuser: Latent-Aware Adaptive Diffusion for Decision-Making
Authors: Fan Feng
Selena Ge
Minghao Fu
Zijian Li
Yujia Zheng
Zeyu Tang
Yingyao Hu
Biwei Huang
Kun Zhang
Venue: ICLR 2026 Poster (2026)
Keywords: POMDP; Latent Variable Models; RL generalization
Primary area: reinforcement learning
TL;DR: We present a unified framework that provides both theoretical identification guarantees and a practical generative modeling approach for identifying and learning with latent factors in decision-making

Abstract: Recent work has framed decision-making as a sequence modeling problem using generative models such as diffusion models. Although promising, these approaches often overlook latent factors that exhibit evolving dynamics, elements that are fundamental to environment transitions, reward structures, and high-level agent behavior. Explicitly modeling these hidden processes is essential for both precise dynamics modeling and effective decision-making. In this paper, we propose a unified framework that explicitly incorporates latent dynamic inference into generative decision-making from minimal yet sufficient observations. We theoretically show that under mild conditions, the latent process can be identified from small temporal blocks of observations. Building on this insight, we introduce Ada-Diffuser, a causal diffusion model that learns the temporal structure of observed interactions and the underlying latent dynamics simultaneously, and furthermore, leverages them for planning and control. With a modular design, Ada-Diffuser supports both planning and policy learning tasks, enabling adaptation to latent variations in dynamics, rewards, and latent actions. Experiments on locomotion and robotic manipulation benchmarks demonstrate its effectiveness in accurate latent inference, long-horizon planning, and adaptive policy learning

Published as a conference paper at ICLR 2026
ADA- LATENT-AWARE
ADAPTIVE DIFFUSION FOR DECISION-MAKING
Fan Feng!.3 Selena Ge', Minghao Fu', Zijian Li Yujia Zheng?, Zeyu Tang?*,
Yingyao Hu Biwei Huang" Kun Zhang 2,3g
1 University of California San Diego 2 Carnegie Mellon University
3 MBZUAI, Stanford University, Johns Hopkins University
f Equal Senior Authorship
ABSTRACT
Recent work has framed decision-making as a sequence modeling problem using generative models such diffusion models. Although promising, these approaches often overlook latent factors that exhibit evolving dynamics, elements that are fundamental to environment transitions, reward structures, and high-level agent behavior. Explicitly modeling these hidden processes is essential for both precise dynamics modeling and effective pecision-making In this paper, propose uni- fied framework that explicitly incorporates latent dynamic inference into generative decision-making from minimal yet sufficient observations. We theoretically show that under mild conditions, the latent process can be identified from small temporal blocks of observations. Building on this insight, we introduce Ada-Di fuser, a causal diffusion model that learns the temporal structure of observed interactions and the underlying latent dynamics simultaneously, and furthermore, leverages them for planning and control. With modular design, Ada-Di ffuse supports both planning and policy learning tasks, enabling adaptation to latent variations in dynamics, rewards, and latent actions. Experiments on locomotion and robotic manipulation benchmarks demonstrate its effectiveness in accurate latent inference, long-ho horizon planning, and adaptive policy learning'.
INTRODUCTION
Learning and planning in partially observable environments is fundamental challenge in building intelligent agents (Kaelbling et al., 1998). Recent work on casting decision-making as generative modeling problem, taking advantage of powerful models such transformers (Chen et al., 2021; Zheng et al., 2022; Kong et al., 2024) and diffusion models (Janner al., 2022; Chi et al., 2023; Ren al., 2025). has achieved impressive results in a wide range of tasks. However, these methods often fail to account for hidden latent variables and their temporal dynamics, factors that are prevalent in real-world settings such robotics (Lauri et al., 2022), autonomous driving (Huang et al., 2024), healthcare (Hauskrecht & Fraser, 2000; Ehrmann et al., 2023), and economics (Brero et al., 2022). Ignoring such latent processes can result in suboptimal decision-making particularly when the observational data does not provide full coverage of the latent factors underlying the environment's dynamics (Zintgraf et al., 2021; Xie al., 2021; Swamy et al., 2022; Belkhale al., 2023).
Early works address partial observability in reinforcement learning (RL) and imitation learning (IL) by encoding historical observations and actions into belief states or latent embeddings. which represent distribution over the underlying latent state (Kaclbling et al., 1998; Hauskrecht, 2000; Guo et al., 2018; Igl et al., 2018; Liang et al., 2024a; Xie et al., 2021). Policy optimization or planning then carried out based on these inferred belief states. However, learning such representations often requires access to the historical trajectories or data from a diverse set of environments. This can be prohibitively expensive, particularly in high-dimensional state or action spaces, posing challenges for integrating these methods into modern generative decision maaiing models, which typically prioritize scalability. Can identify the latent factors that govern environment dynamics and rewards, and
Project Page: http::/siiegggggl..oo///wwwsa/as

---

Published as a conference paper at ICLR 2026
integrate them into scalable generative decision-man making models to enable adaptive planning and policy learning, using only minimal observations, while preserving theoretical guarantees?
In this paper, pursue this goal by addressing two fundamental questions. First, what is the minimum set of observations required, in principle, to reliably identify the latent factors that govern the environment? Second, how can latent identification be effectively incorporated into generative models (e.g., diffusion models) to enable adaptive planning and policy learning? To answer the first question, we theoretically show that, under mild conditions, the latent factors the time step can be block-wise identified using only four surrounding observable measurements (i.e., state-action trajectories) within a small temporal window. This identification result implies that a small temporal block is sufficient to infer the latent factors in observational RL trajectories in an online manner.
Guided by the theoretical findings, we propose Ada -Di fuser, a novel causal diffusion framework with latent identification from temporal blocks, designed to model the data generation process of RL trajectories influenced by latent factors. To reflect the autoregressive nature of sequential decision making, we introduce causal denoising schedule that aligns the denoising steps with the underlying causal structure, drawing inspiration from recent advances in autoregressive diffusion models (Ho et al., 2022; Chen et al., 2024; Xie et al., 2024b; Sand-AI, 2025). For saapprralblook-wisee latent identification, during training, we propose a between denoising the observations and refining latent estimates. This enables Ada-Di ffuser to jointly learn a structured representation of latent variables and the corresponding observational distribution. At inference time, Ada-Di f fuser generates actions and states while estimating latent variables in an online fashion. Since states and actions are conditioned on the latent factors, we employ zig-zag sampling scheme that alternates between sampling state-action pairs and updating latent variables, ensuring consistency between generated sequences and their underlying latent dynamics.
teeeese--he---fifin procedure that iteratively alternates
Ada-Diffuser provides a unified generative framework for sequential decision-making It is applicable to both planning and policy learning tasks by conditioning on different types of obser- vations and adapting the conditional generative process accordingly. The framework is flexible and can accommodate various forms of latent, including ones that influence dynamics, rewards, or even represent high-level latent actions. Importantly, even in environments without explicitly designed latent variables, the block-wise latent identification mechanism improves generative modeling by implicitly capturing structured temporal dependencies.
Contributions: (1) We establish sufficient conditions under which latent factors influencing environ- ment dynamics and rewards can be identified from short temporal windows of RL trajectories, without requiring full trajectory access
data. (2) We develop Ada-Dif Difffurr,,ee a causal diffusion model that performs block-wise latent inference to jointly model latent contexts and observ- able trajectories. Unlike prior latent-augmented diffusion approaches, Ada-Di f fuser introduces minimal sufffciet block with backward refinement for identifiable latents and uses fully autoregres- sive denoising with zig-zag sampling couple inference and generation. (3) Ada- Diffuser can be adapted to a wide range of decision-making tasks by conditioning on different types of observation. We empirically show the improved performance on a wide range of planning and control tasks, including 8 environments under 23 different settings.
BACKGROUND AND RELATED WORK
In this section, we provide background and related work on diffusion-based decision-making, Addi- tional discussions are provided in Appendix E, including related work on (1) learning latent belief states in POMDPs (Kaelbling et al., 1998; Hauskrecht, 2000; Igl et al., 2018; Gregor et al., 2018; Goyal et al., 2021), particularly in the context of transfer, meta, and nonstationary RI/III (Zint- graf et al., 2021; Xie et al., 2021; Feng et al., 2022; Ni et al., 2023; Liang et al., 2024a), and (2) autoregressive diffusion models (Chen et al., 2024; Xie et al., 2024b; Sand-AI, 2025; Wu et al., 2023).
Recent advances use diffusion models as planners and policies for both RL and IL. Diffusson Plan- ner: ciitusioo-bbsedd planning leverages generative models to sample future state-action trajectories from a given state, using guidance techniques (Dhariwal& & Nichol, 2021: Ho & Salimans, 2022) to en- courage desirable properties such high expected rewards. Taking Denoising Diffusion Probabilistic Models (DDPM (Ho et al., based approaches as an example, these methods learn a generative model over expert trajectories ((So, ao).-.., (ST, aT)) by modeling forward-noising noising process:
2

---

Published as a conference paper at ICLR 2026
x') to reverse the process. Here, the superscript t denotes diffusion steps, denotes the planning horizon, x" is a clean subsequence sampled from the expert trajectory т, and a, controls the variance schedule at diffusion step t. During inference, trajectories are generated by starting from Gaussian noise and iteratively denoising through the learned reverse process. This generation conditioned on the initial state or other guidance signals (e.g., goals, rewards): f ~ pe(r So, So:y) y). II. Diffusion Policy: In contrast to diffusion planners, Diffusion Policy methods directly parameterize the policy (a s) using diffusion models. For example, Diffusion Policy (Chi et al., 2023) uses diffusion model generate multi-step actions with expressive multimodal distributions. DPPO (Ren et al., 2025) extends this idea by modeling two-layer MDP structure, which enables fine-tuning of diffusion-basee policies in RL. settings. Another line of work uses diffusion models to parameterize the policy networks for only the single current step (Wang et al., 2022; Wannen-Fsstruch et al., 2023; Chen et al., 2023; Lu et al., 2023). Ada-Di fffseer can generally accommodate both diffusion planner and policies within the same framework.
q(''
oI) and parameterized denoising model po(x'=1 -
be optionally
3 LATENT IDENTIFICATION IN POMDP
In this section, we seek to formally model the structure of the decision-makinn system by answering the following questions. First, where do the latent factors reside, and how do they influence the observable variables such as states, actions, and rewards? Second, can they be identified from demonstration data alone? We model the system that extends the standard MDP include un- observable, time-varying latent variables that affect both the transition dynamics and the reward function. This model generalizes the contextual MDP by allowing the context evolve stochastically over time. We then formalize the data generation process under this model using structural causal models (SCMs) (Pearl, 2010). Finally, we present theoretical results that characterize the minimal observational requirements for identifying the latent variables.
3.1 LATENT CONTEXTUAL POMDP WITH TIME-DEPENDENT CONTEXT
We model the latent factors using general contextual MDP framework, where the context it- self evolves over time. Formally, we define a latent time-varying contextual MDP as a tuple M = (S, A. C, T, R),), where S is the state space, A is the action space, C is the latent con- text space, T (s:
a,-1, Cr) is the transition distribution, R(s,, a;, C2) is the reward function, and y E 10, is the discount factor. The latent context C₂ C follows time-dependent (possibly stochastic) process: C₁ ~ p(c, 1, and is unobserved during training and inference. The agent only observes trajectories T {S ag)...., (ST, ar)), } and infers the latent context C, from the observational data. This is naturally relevant to several MDP models, including (dynamic) hidden parameter MDPs (Doshi-Velee & Konidaris, 2016; Perez al., 2020: Xie et al., 2021), Bayes-adaptive MDPs (Martin, 1965; Duff, 2002: Zintgraf et al., 2021). and factored MDPs (Guestrin et al.. 2003). full comparison and analysis is given in App. C.
Given trajectories generated under this model, we can describe the data generation process using the SCMs. Without the loss of generality, we consider the setting where an expert policy is assumed to generate the actions, as is standard in learning from demonstration data. The data generation process can therefore be expressed as (l.h.s. Fig. 1):
(a)
(b)
Latent Dynamics: Ct = (c(+,_1,7p) Tp)),
State Transitions: S; 1(5--1,8--11 Ct, Et)
Action Generation: a; = T(s;
Reward Function: F2 g(st, as, Ct2, 51),
Figure 1: (a) SCM of the Latent Contextual POMDP. Gray/white nodes are observed/latent variables; green/red edges represent transitions driven by latents/expert policies, respectively. (b) Examples where latents influence either dynamics rewards (affecting optimal actions).
where T/. Eg. and & denote i.i.d. exogenous noise variables. Fig. 1(a) shows the graphical model. Fig. 1(b) illustrates examples where latent factors on dynamics (c.g., external wind locomotion) and rewards (e.g., varying target objects in robot control) influence optimal decisions.
3

---

Published as a conference paper at ICLR 2026
3.2 IDENTIHIABILITY OF LATENT FACTORS WITH MINIMAL MEASUREMENTS
To learn accurate dynamics and make reliable decisions, is essential that the underlying latent factors influencing the environment are identifiable with observational data. We present theoretical results that characterize the minimal number of consecutive observations required for the identifiability of the latent variables, under a of mild and natural assumptions.
Assumption 1 (First-order MDP). We consider the following conditions:
P (St: az, T, C2 S a,-1, Ct-1 W wer-1) = P (s,, a,, tt, Ce st-1, a(-,
where wet-1 = ( .,,,,,,11 $1, 2 a1. c1- c1
This is naturally satisfied under our setting described in Section 3.1.
Assumption (Distributional Variability). There exist observed state and action variables x such that for any x; E Xt. there exists a corresponding Xt-1 E Xt-1 and neighborhood N around (xt, X(-1) satisfying that, for all X,-2 E 2,, Xt-1 E X4-1, X; X2, and X++1 E NNt11, the Sollowing conditional distribution operators are injective: (i) Lxx, 2|x,111 (ū) Lx e/x and (ii) Txxa 1 where the conditional operator I, represents transformations at the distribution level, that is, how one probability distribution pushed forward to another (Dunford Schwartz, 1971).
Assumption justification. Conceptually, the injectivity of these operator implies that different inputs induce different output distributions, thus imposing minimal condition on distributional variability. In RL systems, this condition is naturally satisfied in most stochastic environments where transitions produce sufficient diversity across different states and actions. The assumption also aligns with the conditions in identifiability theory, particularly in works using spectral decomposition and latent variable models (Hu Schennach, 2008; Hu & Shum, 2012; Fu et al., 2025). We further verify this empirically using MuJoCo RL trajectories with the context instantiated as time-varying wind (App. B.5.
Assumption 3 (Uniqueness of Spectral Decomposition). For any Xt E X4 and any è₁ + E. E C.. there exists X X1-1 and corresponding neighborhood N" satisfying Assumption 2 such that, for some (R,, x;-1) E N with Xs # Xx, X # X;-1;
0 < k(x,, X1, Xx- x1 c.) < C < 0a for any c₁ E C2 and some constant C:
it. k(x,, X1, X 1, 22) * k(x,, X1, x 8, E2), where
k(X,, X4. X, 1, c,) = Px. Px.1x.11.0 (xe
Xx-1 Ce) (x-- (Xx 81-1, Ct) Px(xx 1.04 (Xe Xe-1, Co Px, 11,, (Xi xx-11. c:)
(1)
Assumption justification. Conceptually, Assumption 3 requires that k, which captures second-order variations in transition dynamics at time - 1 and t under the latent variable c, yields distinct values for different c's. This requirement is typically met in RL, as varied latent dynamics or rewards often cause significant, observable shifts in behavior. Crucially, this variability is precisely what motivates the need for the identification of the latent variable ct as it governs meaningful differences in learning underlying decision decision-making process. We further verify this empirically using MuJoCo RL trajectories with the context instantiated as time-varying wind (App. B.5.2).
These assumptions are mild and natural. While Assumption 1 standard in RL, it can be relaxed without violating our theory (App. B.3.4). Assumptions 2-3 are naturally satisfied in practice, as they simply formalize that latent variables influence the dynamic, motivating why we need the identification of them. Further validation and discussion are provided in App. A.4. Importantly, the more strongly the context influences the dynamics (and thus the more critical it becomes to account for in decision-making! the more strongly these two assumptions are satisfied: the transition operator becomes more injective as required in Assumption 2, and the spectral ratio l becomes more separable across contexts required in Assumption 3 (See empirical validation in App. B.5.3 ). Under these assumptions, we establish an identifiability theory that characterizes the conditions under which the latent factors can be recovered, and specifies the level of identifiability that be achieved.
4

---

Published as a conference paper at ICLR 2026
Latent Factor Identification
Ada- DiffDiiiff-Fiiiie Planning
Ada-Diffuuse-PPolicy
...
C
Sequential Encoder
Causal Diffusion Model
Sequential Decoder
Diftusion VO
Mask O Inverse dynamics output (
Figure 2: Overview of the Ada-Di ffuser framework. The modular design consists of two main stages: latent context identification (Stage 1, Section 4.2), followed by causal diffusion model (Stage 2, Section 4.3) that models the generative structure of the trajectories. The learned model is then used for planning or policy learning conditioned on the inferred latent context.
Theorem (Identifiability on Latent Factors). Under Assumptions 1-3, the posterior distribution of latent factor with consecutive observations p(c, Xt-2:t+1) can be identifiable up to an invertible transformation the latents é h(cz), where é, estimated latents and h is invertible function.
The proof is in App. B.2. Theorem indicates that a short temporal window of observations (with future frame at & 1) contains sufficient information recover the posterior distribution over the true latent factors (up an invertible transformation) in an online manner, without requiring access the full trajectory. This form of identifiability standard in representation learning and is sufficient for downstream tasks such dynamics modeling, planning, and control. Any policy or dynamics model that conditions on & can implicitly compose with /- without loss expressiveness. We further discuss the implications of this finding in greater detail in App. A.4.
4 LATENT-AWARE ADAPTIVE DIFFUSION PLANNER AND POLICY
Building on Theorem 1, we introduce the Ada-Diffuse framework for learning and planning with latent identification. As illustrated in Fig. 2, Ada-Di f fuser models the trajectory generation process via two modules: (1) latent factor identification block, which estimates the sequence of latent variables from the observable trajectories; and (2) causal diffusion model, which learns the causal generative process of RL trajectories and explicitly infers latent context. Guided by the theoretical findings Theorem and the generative process (Sec C), we use autoregressive denoising step
for temporal dependencies and backward-refieeme via a tailored noise schedule and zig zag sampling, to recover the latent posterior an online manner.
minimal-sufficien sufficient block, designed
In this section, first present general formulation of conditional diffusion modeling with latent variables. We then describe the two modules of Ada-Di f fuser detail (Fig. 2). The complete algorithmic pseudocode of the training and inference procedures given in App. D.1.
4.1 LATENT-AUGMENTE DIFFUSION MODEL FOR PLANNING AND POLICY LEARNING
Without loss of generality, we denote the observable trajectory as which may correspond to state-action sequence Tea or a state-only sequence T., depending on the task setting. To incorporate latent structure, we augment the observable trajectory with the estimated latent context, yielding the full trajectory representation r [Tz. Tz). where т. denotes the inferred sequence of latent variables.
We train conditional diffusion model to generate trajectories conditioned on desired attributes (T) (e.g., reward or goal specification) and the identified c. The denoising model Eq is trained to predict the noise added during the forward diffusion process via the objective: Lai
E, Eyo.y vile [|ee(TT, (rooT",,,y(11 (,, - -el]] where -0 is clean trajectory sample, ~ NN I), and

---

Published as a conference paper at ICLR 2026
the noisy trajectory at diffusion step t is constructed as: -'= =Vi v1- ā.c, where & denotes the cumulative product of the forward noise schedule. Here, the superscript indexes diffusion steps, and should not be confused with the environment time step indices within the trajectory.
Ada--iis Dfffeerr can flexibly adapt to generate different components of the trajectory depending on the task. In the planning setting, the model generates full trajectories
{xx,, Xt+1,. Xx,TT, where Tμ denotes the planning horizon. Here, X; may have two cases: (i) x; = {S,, az), when both states and actions are generated, (ii) x, = {se), when only states are generated. In the latter case, we train an inverse dynamics model (IDM) (Ajay et al., 2023a) to infer the corresponding actions from state transitions. In the policy learning setting, the model generates only actions, i... T = {a++1, a(+2..... at+T. J. where Tg is the action generation horizon. While multi-step action generation methods (e.g., DP (Chi et al., 2023)) can also be viewed as form of plan- ning (Zhu et al., 2023), for generality, we categorize such settings under the policy framework. Ada-Di f--Diffuurr-Pool iy accommodates both variants: multi-step action generation (T, > 1), as in DP, and single-step decision-making (Ta 1), as in IDQL Clansen-Estruchh al., 2023).
4.2 STAGE 1: OFFLINE LATENT FACTOR IDENTIFIATION
Based on Theorem 1, we structure the latent inference process around temporal blocks, using short segments of trajectories identify the latent context each time step. We adopt variational inference framework (Kingma & Welling, 2014) in which the latent variable C; is inferred block-wise. That is, the prior distribution is conditioned on the latent variable from the previous step and the in- block history, while the posterior additionally incorporates future observations. Specifically, given a trajectory block T2 + 1, where T, is the block size, we have prior Po ct-1). and posterior 2o (Cr Xt-T,=*+1] 1), where denotes the observed variables and may correspond to {s}, (s. a), {s, r). We then optimize the evidence lower bound (ELBO) of the observed trajectories:
LELBO.: (e. |- log (xx x,-1, + DKL (ge (c: Xt-T,a+1) || Pa(cc cc-1)).
Here, the reconstruction term, log Px X;-1; Ct) is instantiated based the available observa- tion modalities. Specifically, (i) when only states are observed, the model reconstructs S, conditioned on (s--1, ce): and (ii) when rewards are available, the model also reconstructs r2 from (s,, a, Ct).
The stage is learned through sequential encoder and decoder (l.h.s,, Fig. 2).
4.3 STAGE 2: CAUSAL DIFFUSION MODEL
We propose a causal diffusion model for learning the generative process described in Sec.3. By "causal," we refer to the modeling of the true underlying data generation process, which incorporates two key desiderata: (1) the autoregressive process inherent in temporal sequential RL trajectories; and (2) the latent factor process, capturing the causal influence of the unobserved context variables C, on the observations (e.g., Xt [s,, ar, rz)). Thus, unlike prior diffusion- RL methods and latent-augmented variants (Sec. see Table A1l for comparison), our approach incorporates the following design choices.
Autoregressive Denoising To model the autoregressive structure of trajectory generation, and following the recent advances in autoregressive diffusion (Chen et al., 2024; Xie et al., 2024b; Wu et al., 2023), we introduce causal denoising schedule. Under this mechanism, each time step within a local temporal block is assigned denoising schedule that depends both on its temporal distance from the conditioning anchor and on the inferred latent variables. This reflects the intuition that later time steps exhibit higher uncertainty. Specifically, for a trajectory of length T,, we assign monotonically increasing noise levels (kg [Aj,..., kry. sampled linearly as k; = + K where i (I,..., denotes the maximum diffusion step.
.....,TT T} and K
Given the inferred latent context Čo:r, the model performs autoregressive denoising over the block in T' steps. The overall denoising process is defined as:
Pø
x!-1
X
Xi Co.T).
where x denotes the noisy observation at time step i, and xf is the clean, denoised output. Specifically. the first denoising step is: Po (xf. first observation Xy has been fully denoised and other observations are partially denoised, followed by
& x y.... x x ...., x..,,,,__ Co:r). where the

---

Published as a conference paper at ICLR 2026
the second step: pe(x (xx: X2 ................................ are denoised: Po Po(x x xg 2: X1 --1 Co:r).
ky-2 xg. x²" T-1
Cu:r), and finally until all observations T-
kr-1
Denoise-and-refine Mechanism Theorem indicates that both historical and future observations are required for recovering the latents. However, these future observations are not accessible during online inference, which results in mismatch between identifiability requirements and available information. Hence, guided by this insight with preserving the causal structure of the generative process, we propose a novel denoise-and-refin mechanism that alternates between denoising the observable sequences and refining the latent estimates, and is applied consistently during both training and inference to ensure sigh-quaity latent context recovery in an online manner. We introduce how we implement this during training and inference.
Training: Given a noisy input x with noise level kt, we first sample an initial latent context from the prior: cqrior Po(ce | Ct-1), and use it to denoise the observation: x = co (x kr, eprior). Then we infer the latent using the broader temporal window including future observations (accessible in offline data): cieat qy(cc, x(-k:t+1); and obtain a refined denoised prediction: x (00) = =x((x; ki, c.
eprior [-1
aprio
posterior network, conditioned
é
epost
eprior
We have two reconstruction losses: one from the prior-sampled latent, Cprir = ||x (0) x2112, and one from the posterior-sampled latent, Lpost ||x x To encourage the posterior latent produce better re-
post
Figure 3: zig-zag
constructions, we introduce a contrastive improvement loss: Crel
sampling (2 steps).
softplus (log Cpest - log sg (Eprior) + +m), where sg(-) denotes stop-gradient, softplus(1) u) log(1 + e"), and m 2 is margin hyperparameter. The final objective for this denoise-and-refine step is: La. Lpost + Aprior Eprioo + Are where Aprice and Arel weighting coefficients. Lai updates only 0, postt updates only 10, Lprior updates only o, and Leet
updates both @ and v.
Inference: During inference, future observations are not available, which prevents direct use of the posterior network for latent inference. To address this, we adopt zig-zag sampling strategy² that combines autoregressive denoising with latent refinement. Specifically, we first sample the entire trajectory by applying the forward diffusion process with the maximum noise level K. We then perform autoregressive denoising across time.
For each time step f, we begin by denoising x to an intermediate noise level h using è, sampled from the prior: a!rior Po(c, c.-1). We then obtain updated è; from the posterior latent distribution ePos (c, x xi x x which conditioned on the denoised history, the intermediate step with noise level kn, and the next step with noise level k2. We then use 6pext as the input to further denoise An to x!_ An illustration of the zig-zag inference process is provided in Fig. 33.
In summary, Ada-Di -----fffuee fuserr leverages autoregressive noise scheduling to reflect temporal structure, integrates latent context identification by the denoise-and-refee mechanism, and employs zig-zag sampling for online latent inference. This framework accommodates wide range of scenarios, including latent dynamics/rewards, learning from action-free data with latent actions, and both state- and image- based environments. All variants share the same core, with task-specific modifications to the input/output only. Details of these architectural and variations are in App. H.
5 EXPERIMENTS
We aim to answer the following questions in the evaluation: (1) Latent Identification: How well can Ada-Dif fuseer capture latent factors in the environment? (2) Learning with Latent Factors: How effective is Ada-Di f fuser r in planning and control when learning with the latent context on dy- namics and reward? And can AdafDiffuser infer latent actions from action-free demonstrations? (3) Learning with Environments w/o Explicit Latents: In environments without explicit latent factors,
2 Note on terminology: our use of "vig-rag" is purely descriptive, and there is no connection between the proposed sampling and Bai et al. (2024).
'A larger illustration with steps given in App. Fig. A2.
7

---

Published as a conference paper at ICLR 2026
(a]
MSE
Reward (b)
Lift
Can
Square
Success Rate
TDP (ki))
Lift
Our
Valun
Can
Oses (AP)
6as
0.44
Square
(MP)
8.54
Blck
Figure 4: (a). Identification Results (i.e., Linear Probing MSE, R²) and normalized rewards on the Cheetah environment with varevng wind as the latent factor, evaluated across different block sizes. (b). Results (i.e., average success rate) on planning with action-free demonstrations on Robomimic benchmark. "AF" denotes Action-free.
can modeling latent processes still bring performance gains? (4) Ablation Studies: What is the impact of key design choices in the framework?
5.1 SETTINGS
Benchmarks We consider a diverse set of benchmarks, including Mujoco-based locomotion tasks (Cheetah, Ant, Walker), a robot navigation task (Maze2D), and robot arm control task (Franka-Kitchen) (Gupta et al., 2020), all from the D4RL benchmark suite (Fu et al., 2020). We also consider robotic manipulation tasks from RobotMimic (Mandlckar et al., 2021) and LTBERO-10 (Liu et al., 2023). detailed description and illustration of these environments is provided in App. F. We introduce latent factors affecting both dynamics (c*) and reward functions (c") in the Cheetah and Ant environments, considering two types of variations: episodic changes (E) and fine-grained, time-varying step-wise changes (S). The specific change functions for each setting are detailed in App. F.1. For evaluating latent action modeling, we follow the setup from LDP (Xie al., 2025), using action-free, pixel-based demonstrations from the LIBERO benchmark (Liu et al., 2023). We use our framework to learn the inverse dynamics model to infer the latent actions (details are App. G. total, we evaluate on 8 environments with 23 settings.
Baselines We compare Ada-Di E fuser with diverse set of baselines for fair and comprehensive evaluation. (I) Vanilla diffusson models: For planning, we consider Diffuser (Janner et al., 2022) and DD (Ajay et al., 2022). For policy learning, we include DP and IDQL (Hansen-Estruch et al., 2023). also evaluate LDCQ (Venkatraman et al., 2024), which learns a latent skill space and optimizes a value function conditioned on both states and latent skills. (2) Latent context modeling: We include MetaDiffuser (Ni et al., 2023) that learns contextual representations from multiple environments. We also consider using LILAC (Xie et al., 2021) and DynaMITE (Liang et al., 2024a) which models nonstationarity in RI. through latent context learning using belief states. For a fair comparison, we integrate their context modules into diffusion planners and policies plug-in components (detailed analysis in App. H) (3) Latent action modeling: We compare with LDP (Xie et al., 2025) with action-free demonstrations for planning. total, we compare with 9 baselines across these settings.
Architecture Choices (Details are in App. D.2) For latent factor identification, we use GRU (Cho al.. 2014) embedding with MLP layers as both prior and posterior encoders to produce Gaussian distribution over latents. For decoders, we use MLP layers. For planning and policy learning, we use UNet (Ronneberger et al., 2015) or Transformers (Vaswani et al., 2017) as denoising networks and use MIPs to learn the IDM. We use VAE (Kingma & Welling, 2014) for the visual encoders.
5.2 RESULTS AND ANALYSIS
Results on Latent Identification To verify our identification theory, we evaluate model performance under different block sizes that contain varying amounts of temporal context. We include settings where all blocks have sufficient observations, as well as challenging case with insufficient obser- vations (i.e., without access to future observations). To quantify the quality of the learned latent representations, we adopt linear probing and the coefficient of determination R²2 as the evaluation metric. The results, together with normalized results, are shown in Fig. 4(a). Similarly, we also provide the clustering result in App. Fig. A6. The yellow region indicates settings with insufficient observations, resulting in lower identification results. The purple region corresponds to sufficient observations and yields relatively strong performance, and the green region reflects larger block sizes,
8

---

Published as a conference paper at ICLR 2026
DD
DF
LDCQ
Ours
Diffuser
DP
120
100
65
93
140
110
55
120
100
Kitchen (partial)
Maze2D (large)
Walker (medium expert)
LIBERO (object)
LIBERO
(long)
Figure 5: Results on environments without explicitly designed latent factors. Complete results are provided in App. Table A6-A9.
DF DynaMITE
DF LILAC
Environment
Diffuser
DF
MetaDiffuser
Ours
Cheetah- Hiinhh-Win--EEE Bind-B(n (c")
-120,4 = 12.7
-105.8
9.6
-82.3 8.2
-91.5 = 1.8
-95.3 + 7.4
-68.9 7.6
Cheetah-Wines Wind-S (c")
-148.5
9.8
-102.0 +: 10.2
-87.2 4: 10.4
-96.7 9.5
-105.6
14.5
-73.5 :: 8.7
Cheetah-Ves Vel--E (c")
-102.4
-85.6
-60.2
10.8
-67.8
-62.6
11.1
-45.8 9.5
Ant-Dir-E (c")
188.6
39.2
195.4
+ 47.0
266.7 28.1
233.6
31.9
229.4
32.6
285.3 24.5
Table I:
Results (5 seeds) Ada-Di wfwuefuullaaanne and rewards. c and c" indicate the changes on dynamics and reward, and S represent the episodic and time-step changes. All results are averaged over random seeds.
with latent factors that affects dynamics
which lead degraded results due redundant information inherent difficulty for optimization. Notably, the reward is positively associated with the accuracy of latent identification, validating the importance of identifying latent factors in RL trajectories.
Results on Decision-making We consider three groups based on the kind of latent factors.
Group I: Latent factors dynamics and reward: Table presents the results of learning under latent factors that affect dynamics and rewards in locomotion tasks. To ensure fair comparison, we implement autoregressive variants of DynaMITE and LILAC using the DF backbone. Results for the DP- UP-backbone counterparts are provided in App. Tables A4, which are consistently worse than DE. Additional results, including using DP as backbones and meta-learned versions of Ada-Di-i ffuser that use ground-truth latents input, are provided in App. Table A4-A5. From the results, we observe that Ada-Di fuser consistently achieves the best performance, with significant margin over all baselines. In particular, it outperforms Diffusion planners and policies even when those models are enhanced with latent context modules such DynaMITE and LILAC (pink area), which are most comparable our setting. Furthermore, Ada-Di Dffferer outperforms DF, showing the effectiveness our framework.
oracle variants
Group II: Latent Actions: Following Xie et al. (2025), consider learning from action-free demonstration data, where actions are treated as latent factors to be inferred. We adopt the same setup in (Xie et al., 2025), using pre-trained visual encoder obtained via VAE learn the latent space from pixel observations. We then train a latent planner and an IDM using diffusion- approach. Unlike prior work, our diffusion-based latent planner additionally incorporates latent factors с to model latent context. Importantly, we train only the planner using additional action- free demonstrations. Detailed training procedures are provided in App. G.I. Results on several tasks in Roboming inm i с benchmark show that we can bring improvements on all tasks via modeling the latent process supplementary to the latent planner (Xie al., 2025). Here, the IDM is trained solely on expert demonstrations. Complete results are provided in App. Table A3.
Group III: Environments w/o Explicitly Designed Latents: variable effectively serves as form of Bayesian filtering over the observed trajectories, capturing the inherent stochasticity in the data (a more detailed discussion in App. D.3). Such variability commonly arises from system noise, expert action noise, or high-level unobserved factors. The results, shown in Fig. 5 (full results provided in App. Table A6-A9), support this interpretation. Even in environments without explicitly designed latent contexts, incorporating latent modeling allows Ada-De D fuser to achieve performance that is comparable to or better than these baselines. By recovering the latent variables that capture stochasticity, nonstationarity, unobserved structure in
Crucially, in this scenario, the latent

---

Published as a conference paper at ICLR 2026
Latent Design Orig.
w/o latents
Freeze
0.5x
4x
Cheetah (c")
-73.5
-103.5 89.3
-10.44 90.2
-85,2
-77.6
-89.5
-102.4
LIBERO
93.4
90.9
89.4
87.6
85.0
Diffusion Design
Orig.
w/o refine
w/o zigzag
same NS
random NS
Cheetah (c")
-73.5 93.4
-82.0
99.6
-89.7
84.6
LIBERO
83.9
91.4
85.2
88.5
Table 2: Ablations on Cheetah- (planner) and LIBERO (DP-policy).
the offline trajectoriess, the model can produce rollouts that better match the underlying dynamics, even when the demonstrations are imperfect. These findings suggest that our framework can consistently capture implicit latent process in the data, improving both trajectory modeling and planning.
5.3 ABLATION STUDIES
We conduct ablation studies evaluate the contributions of key components in our framework. For latent factor identification, Fig. 4(a) shows the effect of different temporal block sizes, illustrating the benefit of incorporating future observations during inference. Here, we also consider ablations where (i) the entire latent identification module is removed, (ii) the latent identification network is frozen after the first 10% of training steps, and (iii) different numbers of latent updates are used. For the causal diffusion model, we examine the impact of the following design choices: (i) removing the refinement step (wo refine); (ii) removing zig-zag zag sampling (wo zig-zag): (iii) replacing the causal noise schedule with fixed noise level across time steps (vanilla diffusion)
with random noise scaling as in DF (Chen et al., 2024) (same NS, random NS). The results in Table 2 demonstrate the effectiveness of these modules in our framework in both settings: with and without explicit latent factors. Specifically, For the latent identification ablations, we find that the latent variables play critical role. In particular, freezing the latent module makes the model perform poorly, because the latent context follows a temporal process and must continue adapting during training. Varying the latent dimensionality within a moderate range (about 0.5x-2x does not significantly change performance, but using overly large latent dimensions (e.g., 4x-6x) degrades results. likely due to redundant capacity and harder optimization.
In terms of causal diffusion, for refinement and zig-zag, we hypothesize the gains come from reducing posterior mismatch. We therefore run latent probing test on Cheetah with changing wind and report linear- linear-probe MSE across variants; Ada-Di Difi fuser with both refinement and zig aag- attains the lowest error (Table A16; Details are in App. 1.2.4). Removing backward refinement yields the largest degradation (0.18 0.28), consistent with the role of refinement in letting future evidence within block update the latent posterior and reduce temporal lag. Disabling zig-zag also harms accuracy (0.18 + 0.23), suggesting that alternating conditioning helps align the denoising trajectory with the latent dynamics rather than purely following the forward temporal pass. Moreover, the gap between our full model (0.18) and the oracle that has access true futures (0.12) is small, verifying that the predicted future is already sufficiently informative for reliable latent inference in practice.
Additional ablations are provided in App. 1.2, including full results, comparisons of alternative noise schedules beyond linear (App. 1.2.2). sweeps over temporal block length (App. 1.2.3), and analyses of long-horizzn planning (App. I.2.5). Notably, we show that our method introduces no significant computational overhead in terms of training runtime and inference latency (App. 1.1, Table A12-A13).
6 CONCLUSIONS
We demonstrate that identifying latent factors from sequential observations is critical for effective decision-making. We provide theoretical results that establish conditions under which latent variables can be identified using small temporal blocks of observations. This insight enables principled integration of latent identification into a diffusion- generative framework, allowing us to capture the underlying causal process while maintaining scalability. Ada-Di fuser is broadly applicable to a variety of settings, including planning and control tasks with or without explicit latent structure, and even action-free demonstrations. Results across diverse benchmarks show substantial improvements, validating the effectiveness of our method not only in environments with designed latent factors but also in general settings where latent structure is implicit but in Influential.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENT
We would like to acknowledge the support from NSF Award No. 2229881, AI Institute for Societal Decision Making (AI-SDM), the National Institutes of Health (NIH) under Contract R and grants from Quris AI, Florin Court Capital, MBZUAI-WIS Joint Program, and the Al Deira Causal Education project.
REFERENCES
Anurag Ajay, Yilun Du, Abhi Gupta, Joshua Tenenbaum, Tommi Jaakkola, and Pulkit Agrawal. Is con- ditional generative modeling all you need for decision- decision-making? arXiv preprint arXiv:2211.15657, 2022.
Anurag Ajay. Yilun Du, Abhi Gupta, Joshua B. Tenenbaum, Tommi S. Jaakkola, and Pulkit Agrawal. Is conditional generative modeling all you need for decision making? In The Eleventh International Conference on Learning Representations, 2023a. URL https: /oppnreeeeeeeeeeeeeeeeeeeee net/ forum?
id=sP1 Fo2K9DFG.
Anurag Ajay, Seungwook Han, Yilun Du, Shaung Li, Abhi Gupta, Tommi Jaakkola, Josh Tenenbaum, Leslie Kaelbling, Akash Srivastava, and Pulkit Agrawal. Compositional foundation models for hierarchical planning. arXiv preprint arXiv:2309 08587, 2023b.
Lichen Bai, Shitong Shao, Zikai Zhou, Zipeng Qi, Zhiqiang Xu, Haoyi Xiong, and Zeke Xie. Zigzag diffuusion sampling: Diffusion models can self-improve via self-reflection. arXiv preprint 1088,1, 2024.
Suneel Belkhale, Yuchen Cui, and Dorsa Sadigh. Data quality in imitation learning. Advances in neural information processing systems, 36:80375-80395, 2023.
Andreas Blattmann, Tim Dockhorn, Sumith Kulal, Daniel Mendelevitch, Maciej Kilian, Dominik Lorenz, Yam Levi, Zion English, Vikram Voleti, Adam Letts, et al. Stable video diffusion: Scaling latent video diffusion models to large datasets. arXiv preprint arX1v:231..15127, 2023.
Gianluca Brero, Alon Eden, Darshan Chakrabarti, Matthias Gerstgrasser, Amy Greenwald, Vincent Li, and David C Parkes. Stackelberg pomdp: A reinforcement learning approach for economic design. arXiv preprint arXiv:2210.03855. 2022.
Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojeiech Zaremba. Openai gym. arXiv preprint arXiv: 2016.
Raymond Carroll, Xiaohong Chen, and Yingyao Hu. Identification and estimation of nonlinear
models using two samples with nonclassical measurement errors. Journal of nonparametric statistics, 22(4):379-33)) 2010.
Boyuan Chen, Diego Martí Monsó, Yilun Du, Max Simchowitz, Russ Tedrake, and Vincent Sitzmann. Diffusion forcing: Next-token prediction meets uull-sequence diffusion. Advances in Neural Information Processing Systems, 17:224081-24122 2024.
Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in newral information processing systems, 34: 14085444550777 2021.
Yuhui Chen, Haoran Li, and Dongbin Zhao. Boosting continuous control with consistency policy. arXiv preprint arXiv:2310.06333 2023.
Zhe Chen et al. Bayesian filtering: From kalman filters to particle filters, and beyond. Statistics, 182 (1):1-69, 2003.
Cheng Chi, Zhenjia Xu. Siyuan Feng, Eric Cousineau, Yilun Du. Benjamin Burchfiel. Russ Tedrake, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. The International Journal Robotics Research, pp. 02783649241173666 2023.
11

---

Published as a conference paper at ICLR 2026
Kyunghyun Cho, Bart Van Merriënboer, Caglar Gulcehre, Dzmitry Bahdanau, Fethi Bougares, Holger Schwenk, and Yoshua Bengio. Learning phrase representations using mn encoder-decodes decoder for statistical machine translation. In Proceedings of the 2014 Conference Empirical Methods in Natural Language Processing (EMNLP), pp. 1724-1734, 2014.
Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. Advances in neural information processing systems, 34:8780-877 -0944, 2021.
Zibin Dong. Yifu Yuan, Jianye HAO, Fei Ni, Yao Mu, YAN ZHENG, Yujing Hu, Tangjie Lv, Changjie Fan, and Zhipeng Hu. Aligndiff: Aligning diverse human preferences via eehavior-customiin diffusion model. The Twelfth International Conference on Learning Representations, 2024.
URL wwtpt:tt////// /tt:://opppe /pwpreveeeeeeeeeeeee www.wwwwweet...........m........ / forum? i id=bxfKIYfHyy
Finale Doshi-Veles and George Konidaris. Hidden parameter markov decision processes: A semi- parametric regression approach for discovering latent task parametrizations. In IJCAI: proceedings of the conference, volume 2016, pp. 1432, 2016.
Michael O'Gordon Duff. Optimal Learning: Computational procedures for Bayes-adaptive Markov decision processes. University of Massachusetts Amherst, 2002.
Nelson Dunford and Jacob T. Schwartz. Linear Operators. John Wiley & Sons, New York, 1971.
Daniel E Ehrmann, Shalmali Joshi, Sebastian D Goodfellow, Mjaye L Mazwi, and Danny Eytan. Making machine learning matter to clinicians: model actionability in medical decision-making. NPJ Digital Medicine, 6(1):7, 2023.
Fan Feng and Sara Magliacane. Learning dynamic attribute-factored world models for efficient multi-object reinforcement learning. Advances in Neural Information Processing Systems, 36: 19117-19144, 2023.
Fan Feng, Biwei Huang, Kun Zhang, and Sara Magliacane. Factored adaptation for non-stationary reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:200..02119, 2020.
Minghao Fu, Biwei Huang, Zijian Li, Yujia Zheng, Ignavier Ng, Guangyi Chen, Yingyao Hu, and Kun Zhang. Learning general causal structures with hidden dynamic process for climate analysis. arXiv preprint #xXiv:2501.22500 2025.
Tanmay Gangwani, Joel Lehman, Qiang Liu, and Jian Peng. Learning belief representations for imitation learning pomdps. In uncertainty in artificial intelligence, pp. 1061-1071, PMLR. 2020.
Jensen Gao, Annie Xie, Ted Xiao, Chelsea Finn, and Dorsa Sadigh. Efficient data collection for robotic manipulation via compositional generalization. arXiv preprint arXiv:2403.05/ 10, 2024a.
Kaifeng Gao, Jiaxin Shi, Hanwang Zhang, Chunping Wang, and Jun Xiao. Vid-gpt: Introducing gtt-style autoregressive generation in video diffusion models. arXiv preprint 240. 2024b.
Anirudh Goyal, Alex Lamb, Jordan Hoffmann, Shagun Sodhani, Sergey Levine, Yoshua Bengio, and Bernhard Schölkopf. Recurrent independent mechanisms. In International Conference on Learning Representations (ICLR), 2021. URL / /ttp///pppe/...ooooe/oo/e net forummide BiLemdlBUxy-
Karol Gregor, George Papamakarios, Frederic Besse, Lars Buesing, and Theophane Weber. Temporal 2018.
difference variational auto-encoder. arXiv preprint
Carlos Guestrin, Daphne Koller, Ronald Parr, and Shobha Venkataraman. Efficient solution algorithms for factored mdps. Journal of Artificial Intelligence Research, 19:399-468, 2003.
Zhaohan Daniel Guo, Mohammad Gheshlaghi Azar, Bilal Piot, Bernardo A Pires, and Rémi Munos. Neural predictive belief representations. arXiv preprint arXiv: / 1.06407, 2018.
12

---

Published as a conference paper at ICLR 2026
Abhishek Gupta, Vikash Kumar, Corey Lynch, Sergey Levine, and Karol Hausman. Relay policy learning: Solving
tasks via imitation and reinforcement learning. In Proceedings of the Conference Robot Learning (CoRL). PMLR, 2020.
Assaf Hallak, Dotan Di Castro, and Shie Mannor. Contextual markov decision processes. arXiv preprint arXiv:1502.022559 2015.
Nicklas Hansen, Xiaolong Wang, and Hao Su. Temporal difference learning for model predictive control. arXiv preprint arXiv:2203.04900 2022.
Philippe Hansen-Estruch, Hya Kostrikov, Michael Janner, Jakub Grudzien Kuba, and Sergey Levine. Idgl: Implicit q-learning as an actor-critic method with diffusion policies, 2023.
Milos Hauskrecht. Value-function approximations for partially observable markov decision processes. Journal of artificial intelligence research, 13:33- 94, 2000.
Milos Hauskrecht and Hamish Fraser. Planning treatment of ischemic heart discase with partially observable markov decision processes. Artificial intelligence in medicine, 18(3):221-244 2000.
Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang. Bin Zhao, and Xue- long Li. Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning. Advances in neural information processing systems, (1644966649111 2023.
Joey Hejna, Chethan Bhateja, Yichen Jiang, Karl Pertsch, and Dorsa Sadigh. Re-mix: Optimizing data mixtures for large scale imitation learning. arXiv preprint arXiv:2408. 14037, 2024.
Jonathan Ho and Tim Salimans. Classifier-free diffusion guidance. arXiv preprint arXiv:2207.12550 2022.
Jonathan Ho. Ajay Jain, and Pieter Abbeel. Denoising diffusion probabilistic models. Advances in neural information processing systems, 33:6840-6851, 2020.
Jonathan Ho. Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David Fleet. Video diffusion models. Advances in Neural Information Processing Systems, 35:8633 86666636666, 2022.
Yingyao Hu and Susanne M Schennach. Instrumental variable treatment of nonclassical measurement error models. Econometrica, 76(1):195-216, 2008
Yingyao Hu and Matthew Shum. Nonparametric identification of dynamic models with unobserved state variables. Journal of Econometrics, 171(1):32-44, 2012.
Biwei Huang, Fan Feng, Chaochao Lu, Sara Magliacane, and Kun Zhang. Adarl: What, where, and how to adapt in transfer reinforcement learning. arXiv preprint arXiv:2107.02729, 2021.
Zhiyu Huang, Chen Tang, Chen Lv, Masayoshi Tomizuka, and Wei Zhan. Learning online belief prediction for efficient pomdp planning in autonomous driving. IEEE Robotics and Automation Letters, 2024.
Ahmed Hussein, Mohamed Medhat Gaber, Eyad Elyan, and Chrisina Jayne. Imitation learning: survey of learning methods. ACM Computing Surveys (CSUR), 50(2):1-35, 2017.
Maximilian Igl, Luisa Zintgraf, Tuan Anh Le, Frank Wood, and Shimon Whiteson. Deep variational reinforcement learning for pomdps. In International conference on machine learning, pp. 2117- 2126. PMLR, 2018.
Michael Janner, Yilun Du, Joshua B Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. arXiv preprint arXiv:2205.0990 099911 2022.
Chiyu Jiang, Andre Commann, Cheolho Park, Benjamin Sapp. Yin Zhou, Dragomir Anguelov, et al. Motiondiffuser: Controllable multi-agent motion prediction using diffusion. In Proceedings of the JEEE/CVF conference on computer vision and pattern recognition, pp.
2023.
Leslie Pack Kaelbling, Michael Littman, and Anthony R Cassandra. Planning and acting in partially observable stochastic domains. Artificial intelligence, 101(1-2):99-134 1998.
13

---

Published as a conference paper at ICLR 2026
Diederik P Kingma. Adam: A method for stochastic optimization. arXiv preprint 2014.
Diederik P Kingma and Max Welling, Auto-encoding variational bayes. The International Conference on Learning Representations (ICLR), 2014.
Deqian Kong, Dehong Xu, Minglu Zhao, Bo Pang, Jianwen Xie, Andrew Lizarraga, Yuhao Huang, Sirui Xie, and Ying Nian Wu. Latent plan transformer for trajectory abstraction: Planning as latent space inference. Advances in Neural Information Processing Systems, 37:123379-123411, 2024,
Mikko Lauri, David Hsu, and Joni Pajarinen. Partially observable markov decision processes in robotics: A survey. IEEE Transactions on Robotics, 39(1):21 2022.
Shuang Li, Yihuai Gao, Dorsa Sadigh, and Shuran Song. Unified video action model. arXiv preprint arXiv:2503.00200 2025.
Wenhao Li. Efficient planning with latent diffusion. In The Twelfth International Conference on Learn- ing Representations, 2024. URL https: h/tpp:/s////eeeeiie...... wwt.iiewweeeeeeeeeemeeeeeeeeeeee forummrrmmmu pid-btgDDo j.
Anthony Liang, Guy Tennenholtz, Chih-wei Hsu, Yinlam Chow, Erdem BIyık, and Craig Boutilier. Dynamite-rl: A dynamic model for improved temporal meta-reinforcemment learning. arXiv preprint arXiv:2402. 15957, 2024a.
Zhixuan Liang, Yao Mu, Mingyu Ding, Fei Ni, Masayoshi Tomizuka, and Ping Luo. Adaptdiffuser: Diffusion models as adaptive self-evolving planners, arXiv preprint arXiv:230..0187. 2023.
Zhixuan Liang, Yao Mu, Hengbo Ma, Masayoshi Tomizuka, Mingyu Ding, and Ping Luo. Skilldif- fuser: Interpretable hierarchical planning via skill abstractions in diffusion-basee based task execution. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pp.
16467-16476, 2024b.
Bo Liu, Yifeng Zhu, Chongkai Gao, Yihao Feng, Qiang Liu, Yuke Zhu, and Peter Stone. Libero: Benchmarking knowledge transfer for lifelong robot learning. Advances in Neural Information Processing Systems (NeurIPS), 36, 2023.
Cheng Lu, Huayu Chen, Jianfei Chen, Hang Su, Chongxuan Li, and Jun Zhu. Contrastive energy prediction for exact energy-guided diffusion sampling in offline reinforcement learning. In International Conference on Machine Learning, pp. 20825222222222 PMLR, 2023.
Ajay Mandlekar, Danfei Xu, Josiah Wong, Soroush Nasiriany, Chen Wang, Rohun Kulkarni. Li Fei- Fei, Silvio Savarese, Yuke Zhu, and Roberto Martin-Martin What matters in learning from offline human demonstrations for robot manipulation. In Conference on Robot Learning (CoRL), 2021.
James John Martin. Some Bayesian decision problems in a Markov chain. PhD thesis, Massachusetts Institute of Technology, 1965.
Cuong C Nguyen, Thanh-Toan Do, and Gustavo Cameiro. Probabilistic task modelling for meta- learning. In Uncertainty in Artificial Intelligence, pp. 781-791, PMLR, 2021.
Fei Ni, Jianye Hao, Yao Mu, Yifu Yuan, Yan Zheng, Bin Wang, and Zhixuan Liang. Metadiffuser: Diffusion model as conditional planner for offline meta-rl. In International Conference on Machine Learning, pp. 26087-26105. PMLR, 2023.
Judea Pearl. Causal inference. Causality: objectives and assessment, pp. 39-58, 2010.
Christian Perez, Felipe Petroski Such, and Theofanis Karaletsos. Generalized hidden parameter mdps: Transferable model-based in handful of trials. In Proceedings of the AAAI Conference on Artificial Intelligence, volume 34, pp. 5403-5411, 2020.
Karl Pertsch, Youngwoon Lee, and Joseph Lim. Accelerating reinforcement learning with learned skill priors. In Conference on robot learning, pp. 188-204, PMLR, 2021.
Dean A Pomerleau. Efficient training of artificial neural networks for autonomous navigation. Neural computation, 3(1):88-97 9788-77,11 1991.
14

---

Published as a conference paper at ICLR 2026
Lawrence Rabiner and Biinghwang Juang. An introduction hidden markov models. ieee assp magazine, 3(1):4-16, 1986.
Kate Rakelly, Aurick Zhou, Chelsea Finn, Sergey Levine, and Deirdre Quillen. Efficient off-po meta-reinforcemmnt learning via probabilistic context variables. In International conference on machine learning, pp. 5331-5340. PMLR, 2019.
Allen Z. Ren, Justin Lidard, Lars Lien Ankile, Anthony Simeonov, Pulkit Agrawal, Anirudha Majumdar, Benjamin Burchfiel, Hongkai Dai, and Max Simchowitz. Diffusion policy policy optimization. In The Thirteenth International Conference on Learning Representations, 2025.
URL https: //pppneeeieeeee net/ forum? id=mEpgHvbD2h.
Olaf Ronneberger, Philipp Fischer, and Thomas Brox. U-net: Convolutional networks for biomedical image segmentation. In Medical image computing and compuler-asiisted intervention-MICCAI 2015: 18th international conference, Munich, Germany, October 5-9, 2015, proceedings, part IIl 18, pp. 234-241. Springer, 2015.
Sand-AL. Magi-1: Autoregressive video generation scale, 2025. URL URLhttps://static world/ss httssstatcccfllle 1 pdf.
Andy Shih, Suneel Belkhale, Stefano Ermon, Dorsa Sadigh, and Nima Anari. Parallel sampling of diffuusin models. Advances in Neural Information Processing Systems, 36:4263-4276, 2023.
Shagun Sodhani, Amy Zhang, and Joelle Pineau. Multi-task reinforcement learning with context- based representations. In International Conference on Machine Learning. pp. 9767-9779, PMLR, 2021.
Richard S Sutton, Andrew G Barto, et al. Reinforcement learning: An introduction, volume 1. MIT press Cambridge, 1998.
Gokul Swamy, Sanjiban Choudhury, J Bagnell, and Steven Z Wu. Sequence model imitation learning with unobserved contexts. Advances in Neural Information Processing Systems, 15:17665-77676, 2022.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Siddarth Venkatraman, Shivesh Khaitan, Ravi Tej Akella, John Dolan, Jeff Schneider, and Glen Berseth. Reasoning with latent diffusion in offline reinforcement learning. In The Twelfth International Conference on Learning Representations, 2024. URL bttps://openrevice. diiiggirijjdeeeee
Xinyue Wang and Biwei Huang. Modeling unseen environments with language-guided composable causal components in reinforcement learning. arXiv preprint arXiv:2505. 08361, 2025.
Zhendong Wang. Jonathan Hunt, and Mingyuan Zhou. Diffusson policies as an expressive policy class for offline reinforcement learning, arXiv preprint arXiv:2208.0671: 2022.
Tong WL, Zhihao Fan, Xiao Liu, Hai-Tao Zheng, Yeyun Gong, Jian Jiao, Juntao Li, Jian Guo, Nan Duan, and Weizhu Chen. Ar-diffusion: Auto-regressive diffusion model for text generation. Advances in Neural Information Processing Systems, 36:39957-39974, 2023.
Wei Xiao, Tsun-Hsuan Wang, Chuang Gan, Ramin Hasani, Mathias Lechner, and Daniela Rus. Safediffuser: Safe planning with diffusion probabilistic models. In The Thirteenth International Conference on Learning Representations, 2025. URL https: /wpp//eveeeee.ee/////e/ net/ forum? id=ig2wk7kK93
Amber Xie, Oleh Rybkin, Dorsa Sadigh, and Chelsea Finn. Latent diffusion planning for imitation learning. International Conference on Machine Learning (ICML), 2025.
Annie Xie, James Harrison, and Chelsea Finn. Deep reinforcement learning amidst continual structured non-stationarity. In Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. 1 11393-11444 11403. PMLR, 18-24 Jul 2021.
15

---

Published as a conference paper at ICLR 2026
Annie Xie, Lisa Lee, Ted Xiao, and Chelsea Finn. Decomposing the generalization gap in imitation learning for visual robotic manipulation. In 2024 IEEE International Conference on Robotics and Automation (ICRA), pp. 3153-3160. IEEE, 2024a.
Desai Xie, Zhan Xu, Yicong Hong, Hao Tan, Difan Liu, Feng Liu, Arie Kaufman, and Yang Zhou. Progressive autoregressive video diffusion models. arXiv preprint arXiv:2410.08/51, 2024b.
Zilai Zeng, Ce Zhang, Shijie Wang, and Chen Sun. Goal-conditioned predictive coding for offline reinforcement learning. Advances in Neural Information Processing Systems, 10-25528-225444 2023.
Qinqing Zheng, Amy Zhang, and Aditya Grover. Online decision transformer. In international conference on machine learning, pp. 27042-27059. PMLR, 2022.
Zangwei Zheng, Xiangyu Peng, Tianji Yang. Chenhui Shen, Shenggui Li, Hongxin Liu, Yukun Zhou, Tianyi Li, and Yang You. Open-sora: Democratizing efficient video production for all. arXiv preprint arXiv:2412.20404, 2024.
Chuning Zhu, Raymond Yu, Siyuan Feng, Benjamin Burchfiel, Paarth Shah, and Abhishek Gupta. Unified world models: Coupling video and action diffusion for pretraining on large robotic datasets. arXiv preprint arXiv:2504.02792, 2025.
Zhengbang Zhu, Hanye Zhao, Haoran He, Yichao Zhong, Shenyu Zhang, Haoquan Guo, Tingting Chen, and Weinan Zhang. Diffusion models for reinforcement learning: A survey. arXiv preprint arXiv:2311.01223. 2023.
Luisa Zintgraf. Sebastian Schulze, Cong Lu, Leo Feng, Maximilian Igl, Kyriacos Shiarlis, Yarin Gal, Katja Hofmann, and Shimon Whiteson. Varibad: Variational bayes- adaptive deep rl via meta-learning. Journal of Machine Learning Research, 22(289):1-39, 2021.
16

---

Published as a conference paper at ICLR 2026
Appendix of Latent-Aware Adaptive Diffusion for Decision-Making
A Discussions and Overview
18
A.I Broader Impact
18
A.2 Limitations and Future Work
19
A.3 Discussions on the Core Idea
19
A.4 Discussions the Theoretical Assumptions and Results
20
A.5 Discussions on the Model Design..
A.6 Overview
22
Theory
22
B.1 Notation List
22
B.2 Proof of Theorem
................................
23
B.3
Theory-Algorithm Alignment
26
B.3.1
Non-Parametrie Identifiability Theory
26
B.3.2
Model Design Guidance
27
B.3.3
Discussion on Assumptions
27
B.3.4
Relaxing Assumption (beyond first order Markov)
................................
27
B.3.5
Cases Assumption 2.
28
B.4
ELBO
................................
29
B.5
Assumption Verification
31
B.5.1
About Assumption 2
................................
B.5.2
About Assumption 3
32
B.5.3
Policy Learning under Different Separability
32
Summary on Different MDPs
32
C.1 Contexutal MDPs
32
C.2 Hidden-Parameter MDPs
33
C,3 Discussions and Comparisons
33
Details on Ada-Diffuser
34
D.I Full Algorithm and Results
34
D.2 Architecture Choices and ftyper-prrameten
35
D.2.1
Latent Factor Identification
35
D.2.2 Planner
37
D.2.3 Policy
37
17

---

Published as a conference paper at ICLR 2026
D.2.4 Hyperparameters of Contrastive Improvement Loss
38
D.3 Connection to Bayesian Filtering
38
Extended Related Works
38
E.1
Diffusion Model-based Decision-making
38
Latent Belief State Learning in POMDP
E,3 Autoregressive Diffusion Models
40
E,4
Summary
40
Benchmark Settings and Illustrations
40
F.1 Latent Change Factors Design
40
F.2 Overview on Other Benchmarks
G Other Details Ada-Diffuser
42
G.1 Latent Action Planner
42
G.2 Noise Scheduling
43
H Specific Design Choices for Baselines
43
H.1 Details on LILAC and DynaMITE
43
H.2 Details Diffusion Forcing
43
1
Ablation Analysis
44
L1
Training/Inference Time Analysis
44
1.2
Ablation Results
44
1.2.1
Full Results Supplement to Table
1.2.2
Noise Schedule: Linear vs. Logistic Sigmoid
44
I,2.3
Effect of Temporal Block Length on Latent Identification
45
I,2.4
Latent Probing: Effect of Backward Refinement and Zig-Za
45
1,2.5
On the Effect Planning and Execution Horizons: Long-horte horizon Planning
46
A DISCUSSIONS AND OVERVIEW
In this section, we expand on the design and motivation behind Ada Aida-Difee Diffuser, including the rationale for modeling latent factors in decision-making, key architectural choices, and additional analysis of the experimental results presented in Section 5. We then provide an overview of the remaining contents of this appendix.
A. BROADER IMPACT
Our work aims to identify and leverage latent processes in generative decision decision-making making, with applica- tions in real-world domains such as robotics and healthcare. While these tasks may entail potential
18

---

Published as a conference paper at ICLR 2026
societal risks, do not believe any specific concerns need to be highlighted here. Instead, by uncov- ering and modeling the underlying hidden processes, our approach promotes greater transparency in decision- decision-making, which can ultimately lead to more reliable and trustworthy outcomes.
A.2 LIMITATIONS AND FUTURE WORK
One current limitation is that this work focuses primarily theoretical formulation and algorithmic development. Although we evaluate on a variety established benchmarks, real-world deployment, such in self-driving, aerial drones, and physical robotics, remains important direction for future work.
A.3 DISCUSSIONS ON THE CORE IDEA
Q1: On Latent Modeling. Why is it necessary to model latent processes when we already have access to a large amount of demonstration data?
In many decision- makiin systems, there exist unobservable variables that influence both the dynamics and the reward structure. More generally, these latent variables often evolve over time. Such scenarios are common real-world settings, for example, in robotic control, system dynamics can be affected by external forces (e.g., wind, friction), or by varying user demands (e.g., different target positions). In these cases, learning an optimal policy requires conditioning on the latent factors, especially when they importance of latent variable modeling in both reinforcement learning (RL) and imitation learning (II.) (Zintgraf al., 2021; Liang al., 2024a; Nguyen al., 2021; Rakelly et al., 2019; Ni et al., 2023; Xie et al., 2021).
non-stationary when transferring new domains. Prior work has demonstrated the
Even with access to large demonstration datasets, it remains difficult to ensure sufficient coverage over the full space of environmental or task-specific latent factors relevant to decision-making, This limitation has been widely acknowledged in recent efforts focused on analyzing data quality and designing data collection protocols to promote generalization (Belkhale et al., 2023; Xie et al., 2024a; Hejna al., 2024; Gao al., 2024a). However, most of these works target fixed or task-specific latent variables. In contrast, we consider more general setting where latent factors evolve over time and are not predefined. Our framework provides theoretical guarantees for identifying such latent variables from partial observations and seamlessly integrates this identification process into diffusion models, enabling scalability across complex decision-making tasks.
Q2: On the Scenarios w/o Explicit Latents. What does latent modeling represent when no explicit latent factors defined, and why can still benefit decision-making!?
First, Latent stochasticity is always present (in real-world systems). Even in settings where all task-relevant observations are available, e.g., in locomotion tasks where full physical state information is provided, in robotic manipulation with access both proprioceptive and visual inputs, there may still exist underlying processes that are not directly observed. These include domain-specinn factors such as external forces (c.g., wind) or dynamically changing task goals (c.g.. target positions), which can be viewed as implicit latent variables. Hence, it is crucial to infer and condition on these latent factors
In the extreme case where such factors are also fully observed, latent modeling can still offer significant benefits. Specifically. can capture residual stochasticity present the environment or demonstration data, serving to explain variability not accounted for by observable features. As shown our formulation: 8, (st-1, a,-1, rz = 9 (St, 0z, be interpreted implicit latent variables (sometimes can be time-correlated) influencing transitions and rewards. The model can then identify meaningful structure from irrelevant or noisy variations, for instance, filtering visual background artifacts that are not predictive of dynamics optimal actions. In this sense, the learning framework is conceptually similar to Bayesian Filtering
0,,6,) 6,), the residual stochasticity (f, 5) can
Moreover, partial observability and attribution gaps exist even in clean data. Even in environ- ments with consistent near-deterministi demonstrations, the agent often lacks access to the full set of latent causal factors attributes that influence behavior. Specifically, many systems exhibit
19

---

Published as a conference paper at ICLR 2026
structured yet unobserved variability (e.g., task goals, preferences, intentions), and modeling this variability with latent variables improves generalization.
Q3: On the Identification Theory. What does the identification theory establish, and how does inform algorithm design?
The identification theory (Theorem 1) establishes that the distribution over latent variables can be provably recovered from observable trajectories using only a small temporal window, specifically, small temporal block of four time steps. This provides general non-parametric theoretical guarantee that latent factors can be identified without requiring strong inductive biases or restrictive assumptions on the model class or functional form.
This "four-step" result has direct implications for algorithm design. It suggests that latent identification can be effectively performed using short temporal block, which aligns naturally with block-wise generative modeling approaches such as diffusion models. These models operate over segments or chunks of data, and our theoretical results justify using local temporal blocks to infer latent variables in principled and scalable manner.
A.4 DISCUSSIONS ON THE THEORETICAL ASSUMPTIONS AND RESULTS
Q4: On the Assumptions. What do Assumption (Distributional Variability) and Assumption (Uniqueness of Spectral Decomposition) mean, and why they considered mild?
We expand on the intuition and practical relevance of these two assumptions below.
Distributional Variability (Assumption 2) refers to the requirement that the conditional distributions
p(x:-2 Xx+1), p(x++1 Xx, Ct); and p(x; X;-2; x,-1)
sufficiently sensitive to variations in their input. That is, for different input pairs within local neighborhood, the output distributions differ meaningfully, ensuring the system exhibits enough variability for identification. This assumption aligns with real-world decision- settings (e.g.. locomotion or robotic manipulation), where changes in inputs such as physical state, control policy, or reward function lead observable changes output distributions.
Uniqueness of Spectral Decomposition (A (Assumption 3) builds on this by ensuring that changes in the latent variable e, induce distinet influences on the transition dynamics, specifically on the mapping from - to x, To formalize this, we consider the operator l:
p(x, X,11.C) C.) p(x, 1,-1, c.)
k(x,, X1, X1- &.
(AI)
p(x; x:-1, Ct) -p(x &c,
which separates into two multiplicative components:
p(x, Xt-1, Ct)
h1 p(x,
(A2)
X 1, Ce)
p(&,
8 Ct)
h/
(A3)
p(x,
X;-1, Ct)
Here, hn and ly measure how changes in historical inputs affect the transition distribution at the current time step. The assumption requires that for any two distinct values of c,, the corresponding operator is different, indicating that the latent variable has a sufficiently strong influence on the system dynamics.
Since x is in the neighborhood of x, this formulation effectively captures second-order changes in the transition dynamics with respect the latent variable c, This reflects many real-world RL systems, where even unobservable latent factors (c.g., wind speed or goal target) cause noticeable and structured changes in transition behavior over time, for instance, by considering velocity states.
In summary, these two assumptions are not only theoretically necessary for identification, but also naturally hold in many RL and control systems. They justify the need explicitly model and identify latent variables, as such variables often induce meaningful and structured changes in both dynamics and optimal decision-makin behavior.
20

---

Published as a conference paper at ICLR 2026
Q5: On the Identification of Posterior Distribution and up the Invertible Function (Theorem 1). Why do we aim to identify the posterior distribution over latent variables, and what the role of the invertible function between the estimated and true latents?
Theorem establishes that the posterior distribution over latent factors given surrounding observations, p(c, xt-2:t+1 J, is identifiable up invertible transformation. That is, the estimated latent é, satisfies C2 = h((e)) for some invertible fuucction h.
This form of identifiability is sufficient for downstream tasks such dynamics modeling, planning, and control. Specifically, the learned dynamics or policy can be composed with /- without loss of expressiveness or utility. Since we only need to condition on the inferred latent é, perform these tasks, any invertible transformation of the latent space preserves the representational capacity required for decision-ne making.. other words, although we may not recover the true latent variable e₁ exactly, the recovered representation é, contains the same information and can be used equivalently practice.
Therefore, identifying the posterior distribution (up to an invertible transformation) is both theo- retically meaningful and practically sufficient for learning accurate dynamics models and optimal policies.
A.5 DISCUSSIONS THE MODEL. DESIGN
Q6: On Difffernt Settings (Planning and Policy). How Ada Ada- -Dfffffser applied to both planning and policy learning settings?
Ada-Diffuser designed unified and generic framework that accommodates different types of inputs x (e.g., states, state-action pairs) and outputs (e.g., actions, trajectories. or state sequences). This flexibility allows it to support a wide range of planning and policy learning paradigms. We summarize four representative settings below:
Planning with state-action generation: The model generates both states and actions, with latent variables influencing dynamics or rewards. This setting aligns with prior work such as Diffuser (Janner et al., 2022).
Planning with state-only generation: The model generates future states, and an inverse dynamics model is used to recover the corresponding actions. This setup follows Decision Diffuser (Ajay et al., 2023a).
Planning from action-free demonstrations: Only state sequences are available, and latent variables assumed to capture high-level behaviors skills. This setting extends latent diffusion planning (Xie et al., 2025).
Policy learning: The model generates actions conditioned on the current or recent history of states. This includes multi-step action generation (as in Diffusion Policy (Chi et al., 2023)) and one-step action generation (as Implicit Diffusion Q-Learning, IDQL (Hansen et al., 2022)). In both cases, latent factors may affect the underlying dynamics rewards.
These diverse settings demonstrate the universality of our framework and highlight that uncovering latent structure broadly applicable and critical problem in generative decision-making
Q7: On the Latent Identification. How is Stage (Latent Identification) trained, and does introduce additional computational overhead?
In Stage 1, we train the latent identification module using an offline dataset, as commonly done in offline RL and imitation learning tasks. Specifically, we employ lightweight variational autoencoder (VAE) to optimize the ELBO defined Section 4.2. Empirically, this stage introduces minimal computational overhead (Appendix I.1). We further provide ablation study in Appendix I.1 showing the impact the number of training samples on the effectiveness of the latent identification module.
21

---

Published as a conference paper at ICLR 2026
Q8: On the Temporal Block Design. How does this reflect Theorem 1, and why do we not use exactly four steps in practice?
Our approach reflects the theoretical result in Theorem by identifying latent variables using small temporal blocks in both Stage and Stage 2. In Stage 1, we segment trajectories into local blocks and optimize the ELBO to learn the posterior over latent variables. Stage 2, we apply block-wise refinement to improve the posterior estimates using both past and one-step future observations, making a more accurate identification than using the prior alone.
While Theorem I shows that four consecutive time steps sufficient for identifiability in principle, we do not strictly limit the block size to four in practice. Empirically, we find that using slightly larger blocks (typically between 6 and 20 steps) leads to more stable optimization and better performance. Our ablations in Appendix I.2 show that without access to future observations, identifiability degrades, aligning with the theory.
We treat the "four-step" condition not as a strict architectural constraint but as a theoretical justification (sufficient condition) for using small temporal blocks. The optimal number of steps in practice may vary depending on data properties, task complexity, and model capacity.
Q9: On the Refinement Step. Why is the refinement step necessary, how does it work, and does introduce additional computational overhead?
The refinement step is motivated by the identification theory, which suggests that incorporating the current and future observations (other than only using historical ones) allows the model to infer more informative posterior over latent variables than relying the prior alone. This posterior refinement helps the model better capture latent dynamics by leveraging richer temporal context.
During training, the refinement step encourages the model to extract meaningful information from the posterior. Since Stage optimizes the ELBO, the learned prior is already aligned with the posterior to some extent. This prevents the prior from collapsing into trivial solution. The refinement step builds on this by using the pre-trained prior while further improving inference through contrastive learning between prior and posterior samples.
Importantly, this procedure does not introduce significant computational overhead. As shown in Appendix I.2, the refinement uses the same denoising network with different latent inputs (c) and adds only lightweight contrastive loss, making it efficient in practice.
A.6 OVERVIEW
In this appendix, we first present the theoretical analysis in Section B. including the proof of Theorem and accompanying discussion, followed by the ELBO derivation for Ada-Di fuser. In Section C, we provide an in-depth analysis different types of MDPs and their interconnections. Section H details the full Ada-Di f fuser algorithm, model architectures, and its relation to Bayesian filtering. Section E expands on related work, covering diffusion-based decision-making, latent state estimation via belief learning, and autoregressive diffusion models. Finally, Sections F, G, H, and I provide additional details on benchmarks, baseline implementations, and complete experimental results.
B THEORY
B. NOTATION LIST
We summarize the key notations used throughout the paper in Table A1, including variables for observed and latent states, temporal indices, and relevant mappings. These notations are used consistently in our theoretical analysis and algorithmic framework.
Also, we formally define the operators used in the following.
Definition (Linear Operator (Dunford & Schwartz, 1971)). Let a and be random variables with supports A and B, respectively. The linear operator Taba defined as mapping from a probability
22

---

Published as a conference paper at ICLR 2026
Index
Explanation
Support
18, a.]. observed trajectories including state and action at time step dimension of observed variables
N, Rd.td.
dz
da + d, 8, C S a F A Y E E C C, [(50, a...... (s),ar))
8:
state variable at time t
B.
action variable time rewurd received time latent context variable at time t trajectory sequence of (sr, a.)
To
observable trajectory (states or statc-actions)
Tea or 7a
To
sequence of latent contexts
(Co...,, .r) TT: T.l
augmented trajectory with context
Function
T
transition dynamics conditioned
T(s; (..1
R
reward function conditioned state, action, and context expert policy used for generating demonstrations variational posterior latent inference
R as,
x
x (8,
Go
(p(c Xt-Taa11
Po
latent prior distribution
Po(C; Cz-1)
Pe
generative model for transitions
Pe(x; xx-1, (a()
denuising network in diffusion process
Symbol
fi 0z
exogenous noise latent dynamics, state transitions, and reward distribution operator from to ratio of joint probabilities used in uniqueness assumption cumulative nuise schedule diffusion maximum number diffusion steps
i.d. samples from noise distributions
Lulb
defined in Dunford & Schwartz (1971)
k(-))
defined in Eq.
product forward K N
factors
T Th
planning and action generation horizons temporal block size for latent inference
T,, E T. TEN
T,
Table A1: List of notations, explanations, and corresponding definitions.
function Pa E F(A) to probubility function ph E F(B), given by
F(A) F(B)
Pb = Lbja o Pa
Pbla(- (-
aa) Pa( da.
(A4)
Intuitively, this operator characterizes the transformation of probability distributions induced by the conditional distribution Pbja- provides a general representation of distributional change from a to b, without imposing any parametric assumptions on the underlying distributions.
Definition 2 (Diagonal Operator). Let and b be random variables with associated density functions pa and ph defined on supports A and B, respectively. For a fixed value b E B, the diagonal operator Duja is defined as linear operator that maps density function Pa F(A) function in F(A) via pointwise multiplication:
Dbja o Pa = Pbja (b) ()bb--).P..
(A5)
where Dbja = (bba[b] 1) acts as a multiplication operator indexed by b.
B.2 PROOF OF THEOREM 1
Proof. By the definition of data generation process (Fig. 1), the observed density is represented by:
Px:11.X130 1.Xt
[ f Jonn
Px Pri1iae 180 16x, dc,dc;-1
{ [ Presale IM 1.x.
Px. X 1 Pow.eeePPo----------- dc,dc,-1
c--1
Px 18% 1:wp111---PP------ 1 de,dc,-1
Jenn
[ (
Pxr - 21-2.01-11 do
=fo (
Px www.ne/.....i.............e dc,dd,,1-
23

---

Published as a conference paper at ICLR 2026
Then, the property of Markov process presents conditional independence, organized follows:
Nt [ ox
(J Po 0110 doc-1 de:
[ Pa 1:21 doy.
(A6)
Eq. A6 can be denoted in terms of operators: given values of (x,, Xt-1) M, x X Eq. A6 is
Lx -
Lx xx,e; Dx Xt 1----22-"
(A7)
Notably, Eq. A7 is the operator representation of the observed density function in 4 measurements. Furthermore, the structure of Markov process implies the following two equalities:
Jo Px 0x dC
11 [ Po -x- de,-1.
(A8)
For any fixed (x,, x1- E &e x X4-1, we notate Eq. A8 in terms of operators as follows:
Lx Lxx Lx 10
Txx,
Txx, Xx-1 0:-1 1
(A9)
Substituting the second line in Eq. A9 into R.H.S. of the first equation, we obtain
1x411888-1184 Lx /x., Lex,, 1.- Le, 1
1x ct-1 Le, 1.8 1
L xxt1/xe,0s 1x 1.9,--2"
(A10)
The second line above uses Assumption that I xx is injective. Next, we show how to eliminate Le.-, 11-1 from the above. Consider 3 measurements {x,, X1-1 have
M [
Jc. Pxolxx-- Po, 1ix 11x de
(A11)
which, in operator notation (for fixed x(-1), is denoted as
Lxx Lxx, x1 1.60- 1e- LEL _xx 1.0
Lxx x1-1-111112
(A12)
The R.H.S. applies Assumption 2. Hence, substituting the above into Eq. A10, we obtain:
Txx,,celee-1. T--1
1:6-1 Tuxx
I--1 x1+1 bes,, X 188-2
Lxx, Xe+1 be
Lxc 1x
x.x --1 1 bx,.bxxxxx-ee
(A13)
The second line applies Assumption 2 to post-multiply by L-1 L while in the third line, postmultiply both sides by Txx, 11--10
For all Xc, choose a x;-1 and neighborhood. ' around (xe, X:-1) to satisfy Assumption 2, and pick (&,, X,-1) i) within the neighborhood Nr. Because (Xt, Xt-1) N also know that (xt, 8 1), xx,xx111 E A", The joint distribution of of observations can be represented by Eq. A7:
Lx Xx Lx xx,er Dee.Dxilnn Dx.lxi Le.
(A14)
The first term on the R.H.S., ex does not depend on X;- and the last term Le. xx.xxx1xx.... does not depend x,. This feature suggests that, by evaluating Eq. A7 the four pairs of points
24

---

Published as a conference paper at ICLR 2026
x(-)). Xx-)) (xg, x1-1), (x,, Xx-1), each pair of equations will share the same operator representation in common. Specifically:
11 [x 0z Dxi(xi-1 10-1-111
(A15)
1e Lx 0 Dx (x, 1x,xx,1110
1, Lex.xexe. X1-1
(A16)
18 Lx xx. ae Dx x1188111.8 Le, 11
(A17)
1, [xxes1 0 D D2.181-116 Ice, 1:----
(A18)
Assumption implies that 11x441/8 1x. c, is injective. Moreover, Assumption implies Px, 1x (x²
Xt-1, cx) > for all Ct, that D Dx is invertible. We can then solve for x1-,--11,8-- from
Eq. A16 as
D & X Re B... 1 1.x--2 L0.
(A19)
Plugging this expression into Eq. A15 leads to
Ixx 80-1--1 - 1xx,
1. D 21 1.c (11,1 11
L-1
(A20)
At this point, we have decomposed the observable joint operator and expressed it in terms of latent- conditioned transitions, enabling spectral analysis for identifying latent structure.
Lemma 1 of (Hu & Schennach, 2008) shows that, given the injectivity of Lx--z N.,11 as in
Assumption 2, we can postmultiply by L Xx441 -1,1111-- to obtain:
M
L 18. X Lx Dx.1x-nnee x11xx nee Axi11x1no
Ne-2 Xi+ 1,
(A21)
Similarly, manipulations of Eq. A17 and A18 lead to
N = 1 11xx X Ne.
Exx_118s,c0 D 1 xe/ xx-i x 1xx,,. (A22)
D
D-1
Assumption 2 guarantees that, for any Xt, (5e, X1 (11 exist so that Eq. A21 and Eq. A22 are valid operations. Finally, we postmultiply Eq. A21 by Eq. A22 to obtain:
MN = lx,.o Dxi xx-
D=1 Diixi I.c. (Ixx44a (R (xxes1X(,02) x D2118 D-1 D-1
(x,,, (Dxilxu-1,oo) x11x1 x11x11100 noc
D-1
Dx xe 1,01
L-1
=
= Lx (x Dx..eexeei Be-1 Li
x1411x0,01
(A23)
where
(DxIxoxi (c:)
(Deilenn D-1 1e DR.I - D-1 D=/ hi) (ct)
Pxi x1-1,0; (xe
X1-1, Ct )x, (x,
(Xi
81-1, Ct)
(X: 8 Ct) h((:))
Px. x (8t. X Ct P X2.
= k(x,, X,, X4- 1, (c,,).
(A24)
This equation implies that the observed operator MN on the L.H.S. of Eq. A25 has an inher- decomposition, with the eigenvalues corresponding to the function kxx,, X,, Xt- C,) and the eigenfunctions corresponding to the density Px: e: Xx, Cz).
The decomposition in Eq. A25 is similar the decomposition in nonparametrie identification (Hu & Schennach, 2008; Carroll et al., 2010). First, Assumption 3 ensures this decomposition is unique. Second, the operator MN the I.H.S.. has the same spectrum the diagonal oper- Assumption guarantees that the spectrum of the diagonal operator is
ator Dx. x. bounded. Since an operator is bounded by the largest element of its spectrum, Assumption also implies that the operator MN is bounded, whence we can apply Theorem XV.4.3.5 from (Dunford & Schwartz, 1971) show the uniqueness of the spectral decomposition of bounded linear operators:
Texes1 ix,e, T1xem1lx,, P-1 Dx 1.44
PDx & 1 1 P-1
(A25)
25

---

Published as a conference paper at ICLR 2026
where C' is scalar accounting for scaling indeterminacy and P is a permutation the order of elements in ie, as discussed in (Dunford & Schwartz, 1971). These forms of indeterminacy are analogous to those in eigendecompooition, which can be viewed anite-dimensional special case.
We will show why the uniqueness of spectral decomposition is informative for identifications. First,
d
(A26)
Ret1
must hold for every č due to normalizing condition, one only solution is to set C'
Second, Assumption 3 implies that Eq. A25 imply that the eigenvalues k(x,, X1 Ct) are distinct for different values Ct. If several e₁ yield identical eigenvalues, the associated eigenfunctions cannot be uniquely identified, as any linear combination of them remains valid. Therefore, for each Xt, one can choose Xt, xt 1, 1 such that the cigenvalues differ for all Ct.
Ultimately, the unorder of eigenvaluex/eigenfuncii to the set [Px,)//x,0 (- X,, }t)} for all x, Ct admits unique solution (orderibng ambiguity of eigendecomposition only changes the entry position):
is left. The operator, Lx lx,, corresponding
{px*+ ex (- xx, cx)) (Px.+1 (&r. (X++1 X2, C,),, for all xx, Ct. &c, Ct
(A27)
Due to the set is unorder, the only way to match the R.H.S. with the L.H.S. in consistent order is to exchange the conditioning variables, that is,
(px, (1x (. x ,e()) Pxin x. x, , )...) (Pxupulte (.1 (2 al!)), Pxinis ( xx c,22), ...}
(2)
(A28)
Pe 11 ex (: x (a(1)) .et cft(1))), Px
x (2)) (
[xx4 (8, ( (: ((((())) a( Pre. x (2)) (n (2)))
where superscript (.) denotes the index of the conditioning variables [x,, cı], and π is reindexing the conditioning variables. We use a relabeling map H to represent its corresponding value mapping:
/x..x/ (xx,, c.)) 6x, (: X2, Et), for all X1, Ct, Xe,
(A29)
By Assumption 3, different C₁ corresponds to different 1x H(x,, c:)), which indicates that there is no repeated element in (px cx, (. H(,,, cx))) and (pxm es (: 8c, e.)]. Such uniqueness ensure that the relabelling map H is one-to-one.
Furthermore, Assumption 3 implies that Px (xx, c:)) corresponds unique H (Xe, cc).
& &c,, Eτ), implying that The same holds for the 1
Px44lxx.0 (: H Hxx,, cz)) I& (: &2, Cz) = x,, é. = Hxx, c.)
(A30)
Since the observation x, known and suppose &; = Xt, this relationship indeed represents an invertible transformation between è, and c,
è = h(c:).
(A31)
which ensures that p(c, x(-2:(4) can be identifiable up to an invertible transformation on the latent variables éN = =h(c.)
B.3
THEORY-ALGORTTE ALIGNMENT
Here, we provide a more detailed description of our theoretical foundations, model design guidance, and algorithmic implementation. This complements the high-level summary in the main paper and offers additional context about the technical depth behind our contributions.
B.3.1
IDENTIFIBBILTTY THEORY
We establish this non parametric identifiability result that gives sufficient conditions for recovering latent contexts from reinforcement learning (RL) trajectories using short temporal blocks. Each block
26

---

Published as a conference paper at ICLR 2026
includes small number of future steps, which allows the model to reason about both the immediate past and the near future. Formally, we prove that under mild and broadly applicable assumptions, the latent context C₁ driving the generative process of the observed states and actions (x,, a,) can be recovered up to an equivalence class. This identifiability guarantee is important because shows that latent-aware planning be theoretically justified even when the environment contains unobserved factors or task-dependent variations.
B.3.2 MODEL. DESIGN GUIDANCE
The theoretical result directly guides the design of our causal diffusion model. To leverage Theorem 1, the diffusion model must not only generate trajectories but also recover the true latent factors. Concretely, the model must:
1. capture temporal dependencies across short blocks of states and actions;
2. jointly model observable and latent variables:
3. enforce conditions for identifiability, ensuring that the latent c₂ can be isolated from the observed sequence.
These design requirements inform our noise schedule and the coupling of autoregressive denoising with latent refinement.
This provides guidance for the algorithm design:
Autoregressive Denoising. We model temporal dependencies over both observable and latent variables using an autoregressive diffusion process. At each step, x, is denoised while conditioning on partially denoised past states and inferred latent variables from a short temporal block (Section 4.1). This schedule results a structured temporal-lated modeling process that better preserves long-range dependencies.
Backward Refinement. To explicitly identify latent contexts whose posterior depends on future observations (e.g., 1, guided by Theorem 1), introduce a backward refinement step. At the second-to-last denoising stage, we refine a partial state x; using the initial estimate of é sampled from the prior and 2t+1 as additional evidence. The refined é is then used to produce the final denoised state z/. During training, this backward refinement is enforced satisfy the identifiability conditions. At inference time (zig-zag sampling), we substitute r++1 with predicted estimate to maintain efficiency.
Unification. The autoregressive denoising and backward refinement are integrated into a single noise schedule, enabling joint modeling of temporal dependencies and latent variables. Our imple- mentation follows four-step refinement scheme but can be extended to more steps if needed. Notably, despite the additional refinement, the method remains computationally efficient (see Appendix 1.1). Further acceleration is possible via Picard iteration, which parallelizes refinement steps and reduces inference runtime by about 25%.
B.3.3 DISCUSSION ON ASSUMPTIONS
B.3.4 RELAXING ASSUMPTION (BEYOND FIRST-ORDER MARKOV)
We relax the first-order Markov assumption to n-order Markov structure with delayed/cumu- lative influences without altering the core identifiability argument. Suppose the generative process satisfies
p( Xist, a1: C1a) = p(v(#++1 Xta- apt-n Ctstnnt)),
and that the conditioning sets across non overlapping lags exhibit block-wirs wise conditional inde- pendence (the same separation conditions used in Theorem 1). Then there exists finite window of observations whose statistics identify the contemporaneous block [c,, x,] up to an invertible reparameterization.
27

---

Published as a conference paper at ICLR 2026
Latent Type
Block size
Probing Acc
R²
Delayed
6
0.81
0.72
Delayed
8
0.85
0.78
Delayed
10
0.88
0.81
Delayed
20
0.91
0.86
Cumulative
6
0.84
0.75
Cumulative
8
0.87
0.79
Cumulative
10
0.89
0.83
Cumulative
20
0.93
0.88
Table A2: Identification under delayed and cumulative latent effects. Larger is better.
Concrete identification statement. Let W; = (x1-22:1+2m) denote 4n+1-lengget observation window.* Assume: (i) time direction is known (so [c,, x,] + [ce+1 V. 2++1l il is oriented); (ii) the variability (support) conditions from Theorem hold for the n-lag blocks; and (iii) block-wise independence across non-overlapping lags is satisfied. Then there exists an invertible map H such that
x,] H(W)),
so C₁ (and xx) are identifiable up to an invertible transformation from finite window of observations.
Illustration for n=2. When n=2, block- separations allow identification of the joint variables
Cr,, C4+1, X X from X (length 8+1). Knowing the temporal direction disambiguates (c, xe] from [C++1, Xe+1] Because x, observed, obtain c₁ 42443) for some invertible h, and thus the contemporaneous pair [ca, x1] identified.
Connection delayed/cumulative rewards. n-order view. For delay L
Delayed and cumulative effects naturally in
re+E p( (1Xt+e,8e+e,0 ct) (delayed effect),
while cumulative influence over a horizon T. can be written
re+h Ps (X++k: actk: Ci),
k O
both of which are encompassed by the n-order Markov factorization above. Our cheetah variants instantiate these with, e.g., T++e = -v/++ - cell² (delayed) and Ti+A = -||Ve+* -1:4 c; (cumulative), where denotes speed; the identification results remain valid.
Results. We evaluate identification under delayed and cumulative latent effects in the Cheetah environment using observation windows of length 6, 8, 10, and 20. In all cases, linear probes recover the latent with high accuracy, and performance improves monotonically with longer context. For delayed effects, probing accuracy rises from 0.81 to 0.91 and Rt from 0.72 0.86 block size increases from 6 to 20. For cumulative effects, probing accuracy increases from 0.84 to 0.93 and lt2 from 0.75 0.88
the same range. These results confirm that (i) the latent e is behaviorally consequential in non con-first-ore first-order settings and (ii) moderate temporal context suffices for accurate recovery, supporting our relaxed n-order Markov analysis.
B.3.5 CASES IN ASSUMPTION 2.
The assumption of the injectivity of linear operator is commonly employed the nonparametric identification (Hu& & Schennach, 2008; Carroll et al., 2010; Hu & Shum, 2012). Intuitively, it means that different input distributions a linear operator correspond to different output distributions that operator. For a better understanding, the mapping from Pa = Pb. where and are random variables:
provide several examples in Fu et al. (2025) that describe
Example (Invertible).
g(a), where g is an invertible function. "Any window of length least 4n++1 suffices; we
choice for clarity.
28

---

Published as a conference paper at ICLR 2026
Example 2 (Additive). a where plc) must vanish everywhere after the Fourier transform.
Example 3 (Nonlinear Additive). b = g(a) + where conditions from Examples 1-2 are required. 91 (92 + c), post-nonlinear model with invertible nonlinear functions 91, 92. combining the assumptions in Examples 1-3.
Example (Post-nonlinear).
Example (Nonlinear with Exponential Family), = g(a, E), where the joint distribution p( b) follows exponential family.
Example 6 (Nonparametrie). b= g(a, e), a general nonlinear formulation. Certain deviations from the nonlinear additive model (Example 3), e.g,, polynomial perturbations, can still be tractable.
B.4 ELBO
In this section, we provide analysis on the x" x -rerddiction Mean Squared Error (MSE) loss objectives used the Denoise-and Reffin Mechanism of Ada Dda-Diffuser. Our main argument establishes that minimizing the reconstruction losses Lpeier and Lpost corresponds to optimizing an ELBO on the conditional log-likell ligliiol of the clean observation x(! given noisy observation x² and an inferred latent context Ct.
Let x ~ (x{)) be clean data sample from the true data distribution at sequence time step Let Ca be the inferred latent context relevant to x².
The forward diffusion process gradually adds Gaussian noise to xf over K diffusion steps:
{(x] xxxxx11 =NN=)==(;;; v
k-1 (1-ayp (1--a))1) I)
for l E .I.... K}. where og E (0, 1) are predefined noise schedule parameters. This process allows sampling x directly from x²:
x = Vāux? + vI ā.c. where c ~ N(O, I), and 0g II
Oc..
1=1
xxf ce) that parameterized by aims to denoise x! to x k-1 conditioned The reverse process px{ on C₁.
The derivation of the ELBO for diffusion models is standard following DDPM related derivations (Ho et al., 2020; Chen et al., 2024). The conditional log-lik ligelikdd log pe (xf cz) can be bower-bounded using the ELBO:
pe (xt |xt cr)
pog (c.) 2 E [
bef) log (o (cx) +
Σ log
g(x;
| x9)
k=1
Assuming pa satisfies Markov Property (i.e., po (x+- x :........ X ct) ps (x ct)), which
a standard structural assumption for diffusion models, the ELBO can be rewritten as:
log Po (x.]cc)) 2 Ex(x [log (x |x c,)]
gx bx) DKr lxx x Xpollp 1x cx))]
Lu
DKulg(xF bx))llpo(x K lce)).
This inequality holds with equality if and only if the model's true posterior the latent diffusion path, pe (x{ ( x", ct), is identical to the approximate posterior used to derive the ELBO, which is the forward noising process qxx! [x?). This bound can also include additive constant C(x], c,) which does not depend on the model parameters and is thus typically omitted when focusing on terms relevant to parameter optimization.
To maximize log (x( (c), we aim to maximize this lower bound by optimizing Lo (i.c., maximizing this term) and each L term (i.e., minimizing these D terms, as they appear with negative
29

---

Published as a conference paper at ICLR 2026
sign). The term LK is often treated as constant (or absorbed into C(xp, cx) if (x (c,) is set to
standard Gaussian N (0, I) and ok == O.
We parameterize the reverse process (x xx Ct) as a Gaussian:
(x - (x ct) = N -x no k, ct),
The true posterior step x 8-1 ix xx.) is also Gaussian:
(x, # - bxf x = (x (x x,),,
where fiu (x/: xp)
(a 1 1-aa
x
y 1x x and o = T-âL - a4) is the variance.
1-
For an x" -prredictinn model, denoted co (x k, ca) in the main paper, that aims to predict x from the noisy input x and context Ct. the mean of the reverse model ps can be expressed as:
Va an)
co (xf k,);)))))))))))))))) van (1 an 1) X
px(xx;, k, ce)
Choosing oh oi the KL divergence term LA-1 simplifies to:
1
ji (x,+ xx)) - no(x ki ci)1' Ck
Lk-1 = 1x((
2
Va==1 - -aa)
||x; colVaxx + vI âre, k coll
-xx,,6
20
Oik
where C% are constants not depending on 0. The expectation Exq,e denotes averaging over clean data xP and the noise used to construct x Thus, maximizing the ELBO contribution from - is equivalent to minimizing the following weighted MSE term:
Exg.r.
(k()) ||x co(vaxx +
k, c
(A32)
where wkk)
1 (1-
is positive weighting factor.
20%
1-6e
The term To yx (x?) (log po (xq xx] ce) can also be made proportional to MSE if (x]]xx ct) a Gaussian centered at eo (x} 1, c;):
1
log (x11x), c.)
||x -0ff-coxx) co(x) c.) const
20
Maximizing Lo then equivalent to minimizing this MSE.
The diffusion model is typically trained by minimizing simplified objective (e.g., (Ho et al., 2020)), often an unweighted or equally weighted sum of these MSE terms over uniformly sampled diffusion steps & E [1, K] and data xf:
Csimpto (0) EE-U]I.K)x
[||]]
c Vāxx? + VI âge, k, c:) ct)
This simplification is justified by arguing that reweighting terms wi) in Equation A32 can be absorbed into the network or do not significantly alter the optimal solution for expressive models, allowing w<k) to be effectively set to I.
The Denoise-and- Refine losses are:
|x eo cx(x Kr, efrinr
Cprior = Exg epiii
xx eo (xfi Kis aporr 11'
Epost = E e
where x:
V + /1 āg, e, and k; is the specific input noise level for the observation X, determined by the causal denoising schedule kn + K. These losses, Cpicr and Lpaat specific
30

---

Published as a conference paper at ICLR 2026
(a)
(c)
1.
322
(c))
-3
-13
(e)
02
24
0.2
02
3.5
1.0
Figure Al: Verification of the assumptions. (a) Transition separability in Cheetah under the hyperparameter setting (m, n) = (5, 0.5). (b) Transition separability under a weak-context setting (m, n)) = (0.2, 0.2), where the context barely affects the dynamics. (c) Average reward drop when planning with vs. without conditioning on c, plotted against the transition separability. (d) k-distributions for (5, 0.5), (e) k d-disttibsttions for (0.2, 0.2), (f) reward drop versus k separability.
instances of the simplified MSE loss objective in equation A32 with wkk)) = 1, conditioned on the inferred contexts aprior and ageat respectively. Consequently, minimizing these MSE losses directly optimizes the corresponding terms in the ELBO for log po (x/(c..).
Therefore, we have proven that minimizing Cprior and Lpost as defined in the Denoise-ande and--edidd-dnn mech- anism serves to maximize a variational lower bound on the conditional logliikcldoood log pe pe(x{) (x]]c,). cc).
The underlying diffusion model ce(-, -) is trained to be proficient at denoising from range of noise levels k, as captured by objectives such as Cximple- The specific monotonically increasing noise schedule k; used in Sprir and Cpust represents a particular instance from this range of noise levels. Thus, these objectives are theoretically grounded in the principles of variational inference for diffusion models, adapted to conditioning on the inferred latent context c, and applied at specific noise levels relevant to the autoregressive denoising process of Ada- Di fuueer
B.5 ASSUMPTION VERIFICATION
Here, we test whether Assumption 2 and Assumption 3 hold in practice, explain why we view them as mild, and, importantly, analyze what happens when they fail. We use the Cheetah environment, where the latent context corresponds to time-varg vareing wind speed fw = 5 + m sin(ntt) that perturbs the agent's dynamics. We sweep over combinations of (m, n) to address two questions: (1) whether the setting used in the paper, (m.n) (5, 0.5), indeed satisfies these assumptions; and (2) how violations of the assumptions affect our method and the associated analyses.
B.5.1 ABOUT ASSUMPTION
To evaluate Assumption 2, which requires that the conditional dynamics P(x,-1 xx, c.) be injective in the context variable C,, we perform an empirical test to determine whether different contexts induce measurably different transition dynamics. Given a fitted probabilistic dynamics model p(x(+1 Xx, Ct), we estimate the distribution of next states under each context E (cI,..., C. by drawing samples from the replay buffer and computing p(x(+1 X1, c). For every pair of contexts (c,, c,), we quantify the difference between their induced transition distributions using the 1-Wasserstein distance: Inj(c,, c,) = Wi(p(xe+1 C;), p(x; Xx, 0z)).
(A33)
Large values of Inj(c,, cj) indicate that distinct contexts lead to distinct transition kernels, consistent with injectivity, while values near zero suggest that different contexts produce nearly indistinguishable
31

---

Published as a conference paper at ICLR 2026
dynamics. When (m, n) (5, (Fig. Al(a)). observe consistently non-zero Wasserstein distances across contexts, indicating that P(x++1 X3, Ct) is context-injecttvv. in the regime studied. In contrast, when we reduce the context variation (m, n) (0.2,0.2) (Fig. Al(b), the distances toward zero, showing the failure mode of the assumption. This shows that Assumption is mild in the latent-aware decision-making.
B.5.2 ABOUT ASSUMPTION
We provide an empirical test of the spectral ratio io to examine Assumption in the RL setting. Using the dynamics model on Cheetah p(x; X Cz), compute
p(xx 1.
Cc)
1(x,, X1, X;- X, C1)
(A34)
x X, cx) p(x; Cx)
We draw transitions (x.-1.X,) x,) from the replay buffer (600 samples) and form cross-paired transitions (x,-1, x2) by swapping endpoints across trajectories. For each context c E [c), C. ). this yields empirical distribution of k(-;c) We then quantify how well & separates contexts using the I- I-Wasserstein distance between pairs of k-distributions, i.c.,
Sep(c;, oÿ) = Wi(p(A cx), p(k cz)).
(A35)
When Assumption 3 holds, & remains bounded and its distribution varies across contexts. Empirically, we observe clear multi-modal separation across contexts in the paper's setting ((m, n) = (5, 0.5), Fig. Ad)), whereas in regimes where the dynamics become less context-dependent, the k- distributions overlap heavily.
When k is not distinguishable across contexts ((m, n) 00.2, Fig Al(e)) it implies that does not exert a noticeable effect on the transition dynamics. this regime, explicitly modeling the context is unnecessary, since the environment effectively behaves as single-context system. Hence, we believe Assumption 3 mild our main regime and also clarifies the failure mode when is violated.
B.5.3 POLICY LEARNING UNDER DIFFERENT SEPARABILITY
When the conditional transition P( xt,c) c) not injective c or & is nearly the same for different c, different contexts induce nearly identical transition kernels. This means the context not identifiable from the dynamics and does not meaningfully alter the environment; in such cases, explicitly modeling brings little benefit for policy learning. Figures Al(c) and (f) illustrate this effect. We vary (m, n.) to change the strength of the latent wind context, and compare policy performance when planning with the ground-trute context versus ignoring c. We then plot the resulting performance gap (reward-droy ratio) against transition separability and k-separability. The gap shrinks when separability is small, indicating that when both the transition and lo are weakly context-dependent, modeling e is unnecessary. Overall, these results verify that Assumption 2 and Assumption 3 are not only mild our setting, but also clarify why modeling the latent context important precisely in regimes where the dynamics are strongly context-dependent.
C SUMMARY ON DIFFERENT MDPs
Our work considers a contextual POMDP setting with an evolving latent process, which naturally relates to several established MDP formulations, including contextual MDPs (Hallak al., 2015), hidden-parameter MDPs (HiP-MDPs) (Doshi- Velez & Konidaris, 2016), and their variants. In this section, we provide formal definitions of these models and discuss their relationships and distinctions.
C.I CONTEXUTAL MDPs
contextual Markov decision process (CMDP) (Hallak et al., 2015) is defined by the tuple C.A.A,,M,, where C is the context space, S is the state space, and A is the action space. The mapping M assigns to each context e E C a set of MDP parameters M(c) = (R", T"), where R and 7 are the reward and transition functions associated with context
32

---

Published as a conference paper at ICLR 2026
Sodhani et al. (2021) and Liang et al. (2024a) extend the CMDP framework to settings in which the context variable с evolves according to its own Markovian dynamics p(c(+1 ct), closely aligning with our formulation of a latent process evolving over time.
C.2 HIDDEN-PARAMEN
MDPs
MDPs (HiP-MDPs) (Doshi- Dessi--Velzz & Konidaris, 2016) are defined by the tuple 7
. (S, A, + T, 7, Pa), where S is the state space, A is the action space, and θ is the space of task-specific latent parameters. For each 0 E O, the transition and reward functions are given by 8 A (SS) and Rg S x A IR, respectively. The parameter 0 is sampled from prior distribution Pe at the beginning of an episode and remains fixed during the episode. The discount factor is denoted by ? E [0. 1). This framework defines family of MDPs indexed by the latent parameter 0, with each 0 inducing a different set of dynamics and reward functions. It can be seen as a special case of a contextual MDP where the context is latent and fixed per episode. Xie al. (2021) further generalize this framework by allowing the task parameter & to evolve dynamically across episodes, rather than being fixed.
Bayes-Adaptive MDPs (BAMDPs) are closely related to both HiP-MDPs and contextual MDPs (CMDPs). In BAMDPs, the agent maintains a posterior distribution over MDPs based on its interaction history. Specifically, it maintains belief (R, T) PR,,7 T T.)), where (5o, ao, ro..... 5 denotes the trajectory observed up to time t. This belief captures the agent's uncertainty about the underlying transition and reward functions.
The transition and reward functions can then be defined in expectation over this posterior, effectively conditioning decision-make making on the belief b,. When the environment is driven by hidden contextual variables or latent task parameters, such as in CMDPs or HiP-MDPs-thh belief can be interpreted as distribution over these latent variables. In this view, BA provide non-parametric framework for reasoning over hidden structure, while approaches like ours explicitly model such latent variables and infer their posterior distributions using amortized inference. Both aim enable adaptive planning and learning under uncertainty, but differ in how latent structure is represented and inferred.
C.3 DISCUSSIONS AND COMPARISONS
The key distinction between contextual MDPs and hidden- paremneer MDPs lies in how the latent factors are represented: contextual MDPs explicitly treat them as latent variables, while HiP-MDPs model them implicitly as parameters governing the transition and reward functions. In our work, we adopt the contextual MDP perspective, where the latent process is modeled as random variable that evolves over time.
However, our identification theory, focused on recovering the posterior distribution over latent variables, also applies to the HiP-MDP setting. Once the posterior over the hidden parameters identified, the corresponding transition and reward functions can be recovered as well.
Additionally, our framework, which models a factorization over observed states and latent variables. conceptually related to factored MDPs (Guestrin et al., 2003). In factored MDP, the state space S is represented as a set of variables S = {s(1) s(2) s...,,(m)) are decomposed over these factors:
(n)) }, and the transition and reward functions
T(s' s,a) = II Ti (I( Pa ( (s.a)), R(s,a) = Σ Rg (Pay (Paf)((s,a) a)), 1=1
1=1
where f and Pag component, respectively. our framework, while not relying on an explicit graphical structure, shares conceptual similarities with factored MDPs (Guestrin et al., 2003) through its coarse-grained factor- ization over observed states and latent variables. Specifically, we distinguish between latent variables that affect the transition dynamics and those that affect the reward function. Formally, express the generative process as:
()) denote the parent variables (i.e., dependencies) for each transition and reward
T(s+1 St, a,, ci), R(r St, a,, c)),
where ci and distinct (or potentially overlapping) latent factors that influence transitions and rewards, respectively. This separation enables flexible modeling of partially observable environments where different unobserved processes govern the dynamics and task objectives.
33

---

Published as a conference paper at ICLR 2026
D DETAILS ON ADA-DIEFUSER
K
/+2
eprior 1-1
prior
aprior H+
arior I+2
/+1
&post 1-1
prior
eprior H+1
aprior I+2
X
X
1+2
apost
aprior 1+1
aprior 1+2
0
X
X
.2
ppost I+1
Eprior I+2
11
X
/-1
:+1
I+2
&pott
[+-
--1
0
X
:+11
X
/+2
Figure A2: An illustration of the zig-zag sampling process with a block of 4 time steps. J. and indicate denoising and identity mapping, respectively.
D.1 FULL ALGORITHM AND RESULTS
Our framework consists of two stages: latent factor identification and diffusion-based planning or policy learning. Below, we provide the algorithmic pseudocode for both stages. Specifically, Algorithm Al describes Stage 1: latent factor identification, while Algorithms A2 and A3 correspond to Sida-Diffusez Caciuser-Plannee Planner and Ada wwwff-ieff----eeeeee-ee--e Policy, respectively.
For clarity, we omit the detailed step-by-step procedures for denoise-and-refine and zig-zag sampling (Lines 7-8, 11, and 19-22 in Algorithm A2; Lines 6-7 and 13 in Algorithm A3), as these are fully described in Section 4.3. For
show Diffusion Policy (DP)-based algorithm, which provides general framework for multi-step action generation. the IDQL- variant, both the action execution horizon and observation horizon single-step policy inference conditioned only on the current observation.
1, corresponding to
Algorithm A1: Latent Factor Identification.
1: Input: offline dataset D
2: Randomly initialize decoder (0(5441,77 s,s,,, c),
encoder (c, st-T, T T1-1, and prior network po(c, cc-1),
3: while not done do
4: Sample batches of trajectories from D
5: Compute ELBO and update 0 1/.
6: end while
Additionally, we provide the full results for all experiments: Table A3 reports results for the action- free setting; Tables A4 and A5 present results for environments with latent factors affecting dynamics and rewards; and Tables A6, A7, A8, and A9 summarize results for environments without explicitly modeled latent factors.
34

---

Train Causal Diffusion Model (noise predictor 69) with X C-Ta and Aoost T-Tgtt and Published as a conference paper at ICLR 2026
Algorithm A2:
Input: Env, offline dataset D, pre-trained encoder 9s and prior network Pe
observation horizon T,, planning horizon Tp. action execution horizon Ta, condition // Training
2: Initialize noise predictor co, inverse dynamics model fø
3: while not done do
4: Sample XT-T.:t+,, from D
5: Sample cotior 1.t+,, and &Post 12:+T,2 from Ps and qu
6: if using inverse dynamics model then
7:
Train Causal Diffusion Model (noise predictor Eg) with X T-To:: and é t-Taa and other conditions y. target outputs are $1+1++T, Train encoder de with the contrastive improvement loss Ccontrast Train Inverse Dynamics Model fo to generate actions ar+1:
aorior
8:
9:
10:
else
11:
other conditions y, target outputs are ($t+1:147, ac+1:+T,,)
12:
Train encoder qe with the contrastive improvement loss Leontrau
13:
end if
14:
end while
// Execution
15:
Initialize environment: so Env, reset (), set
16:
while not done do
// Observe and infer latent factors
17:
Observe recent trajectory X;-T,it
18:
Sample latent variables eptir from Ps // Generate candidate trajectory
19:
if using inverse dynamics model then Generate future states (zig-zag sampling) S,4144T, conditioned on Toito C.TT, and via learned noise predictor Generate actions à++1:+T, +- fo Spr+T,-1)
20:
21:
22:
else
23:
Generate future trajectory {$1+1447, A41144r,) conditioned on X Xt-T,: CL+T,, and via learned noise predictor eg
prior
24:
end if
// Execute action(s) in environment for each step 1 T, do
26:
Execute â,ti in Env, observe st+++1, Ptti Append (st+i, âe+s, Pt4s) trajectory buffer end for
28:
29:
Update +- To
30:
end while
D.2 ARCHITECTURE CHOICES AND HYPER-PARAMENT
We detail the architectural design choices and hyperparameter settings used for model components, loss functions, and training procedures across all Ada Di fuuser variants under different environ- ments and benchmarks.
D.2.1 LATENT FACTOR IDENTIFICATION
Architectures We variational autoencoder (VAE) (Kingma & Welling, 2014) to optimize the evidence lower bound (ELBO). The same architectural design is used across all variants of Ada-Di coliffuser and all benchmark settings.
For the encoder, we first embed states, actions, and rewards using separate MLPs with ReLU activations. The resulting embeddings concatenated and passed through two-layer MI.P (each
35

---

Published as a conference paper at ICLR 2026
Algorithm A3: Ada-Di Ada-Diffuser-Policy (DP-based)
Input: Env, offline dataset D, pre-trained encoder do and prior network Pe
observation horizon T,, action generation horizon Tp, action execution horizon T.. condition // Training
2: Initialize noise predictor
3: while not done do
4: Sample Tt,T, from D
a 5: Sample latent variables tt++,, ~ Pas epost 1I+T,,2 ~ 9e
6: Train causal diffusion model (noise predictor Eg) to generate actions ap+1:r-Tp> conditioned
c CT.++, and y
on Tt Cct+T,
Train encoder qe with the contrastive improvement loss Ceontrast
8: end while
// Execution
9:
Initialize environment: 80 Env. reset (), set / <- 0
10:
while not done do
// Observe and infer latent factors
11:
Observe recent trajectory X;-T,:
12:
Sample latent variables ttt+, Ps // Generate actions using causal diffusion model Generate actions (zig-zag sampling) â conditioned on Xr-Taits CESTT, and y learned noise predictor // Execute action(s) in environment
aptior
13:
14:
for each step = 1 to To do
15:
Execute â,+z in Env, observe 1.4i+1.7t44
16:
Append (st+i: 0 (244) trajectory buffer end for
17:
Update + 8 To
19:
end while
Environment
LDP (AF)
Ours (AF)
LDP (AF, SubOpt)
Ours (AF, SubOpt)
Lift
0.67 1: 0.01
0.78 :: 0.05
1.00 0.00
0.98 1: 0.02
Can
0.78 = 0.04
0.85 = 0.07
0.98 0.00
0.98 : 0.02
Square
0.47 0.03
0.54 0.05
0.83 0.01
0.89 +: 0.03
Table A3: Results (success rate) action-free demonstrations. Here, AF and SubOpt indicate using Action-free and suboptimal demonstrations on Robomimic tasks, respectively (following the settings in L.P (Xie et al., 2025)).
layer of size 64) followed by a GRU. The GRU output used to parameterize a Gaussian distribution from which the latent variables are sampled.
The state and reward decoders are implemented as separate MLPs, each consisting of two fully connected layers of size 64 with ReLU activations. For the prior network, we use the output of the previous step's 's latent distribution embedding (shared GRU) and feed it into a two-layer MLP (each layer of size 32) to predict the parameters of the prior distribution. For the dimensionality of latents, we choose 20 for Cheetah, Walker, Ant, Maze; 64 for Robomimic, Kitchen, Libero.
Loss Function At each time step t, we optimize the following losses:
x,-1, ct)] + DRL (gs x,-T, Pp cr-1)). LELBO.t = E Eggles xx Tyaes) |- log (x
Reconstruction
KL regularization
36

---

Published as a conference paper at ICLR 2026
Here, x; may include different components depending the setting (e.g., {s,, ar) or st). and C2 denotes the latent context variable inferred from a temporal block of observations. The first term encourages accurate reconstruction of the current observation X, conditioned on its immediate past and the latent Ct, while the second term regularizes the posterior to remain close to the learned prior P& (ct Ct-1).
We implement the ELBO loss as a weighted combination of the reconstruction loss and the KL divergence:
T-2 LELO Σ
[| xx11x AKL. DKL. (1w(ct) XI-T,:+1) |po(c c--1))]
where & is the model's reconstruction the observation Xt, and ÅKL is weighting coefficient. The reconstruction is computed using mean squared error (MSE), and the KL divergence is computed in closed form for Gaussian posteriors and priors. The hyperparameter AKL is set to be 0.01 and the learning rate is set to be 3e 4.
D.2.2 PLANNER
For the planner, we consider two scenarios: (i) generating both states and actions, and (ii) generating states only. For the former, we build upon the Diffuser framework (Janner et al., 2022), which directly models full trajectories. For the latter, we adopt the Decision Diffuser (DD) framework (Ajay et al., 2022), where the model generates future states and uses an inverse dynamics model to recover the corresponding actions via inverse dynamics model.
For type (i) (full state-action trajectory generation), we apply our method to the Cheetah and Ant environments. For the noise predictor, we use 1D U-Net (Ronneberger al., 2015) with kernel size of 5, channel multipliers set to (1, 2, 2, 2), and a base channel width of 32, The model is trained using the Adam optimizer (Kingma, 2014) with a learning rate of 10- a batch size of 64, and for million training steps. We adopt classifier guidance (CG) (Ho et al., 2020) with gradient guidance on computed return, with a guidance scale w 1.5. The observation horizon set to 10 for both environments. The planning horizon Tp is set to 16 for Cheetah and 32 for Ant, with an action execution horizon of These hyperparameters are kept consistent across baselines, including Diffuser, DF, MetaDiffuser, and Diffuser combined with LILAC and DynaMITE for the Cheetah and Ant experiments (those in Table 1 and Appendix Table A4). For other components (c.g., VAE) in LDCQ, we employ all the hyperparameters in their original implementation (Venkatraman et al., 2024).
For type (ii) (state-only generation with inverse dynamics), we predictor with a hidden dimension of 256 and head dimension of 32. The architecture includes DiT blocks for Walker, Kitchen, and Maze2D, and 8 DiT blocks for LIBERO.Ths model trained using the Adam optimizer (Kingma, 2014) with a learning rate of 3 x 10-A, a batch size of 128, and for million training steps. The number of diffusion timesteps is 500. The observation horizon is set to for Kitchen, 2 for LIBERO, and 10 for the other environments. The planning horizon Tp is set to 16 for Kitchen, 10 for LIBERO, and 32 for the others. The action execution horizon is 8 for both Kitchen and LIBERO, and 10 for the remaining environments. For the inverse dynamics model, we use an MLPP-based diffusion model consisting 2-layer embedding module with 64 hidden units. This model is trained for million gradient steps.
Transformer-base noise
3-layer MLP with 128 hidden units, preceded by
For both cases, we set the coefficient of the contrastive improvement loss Leomtrst = max Lprior } to 0. The key hyper-parameters are summarized in Table A10.
D.2.3 POLICY
For the DP-based policy, we adopt the same architecture as the planner described earlier for Cheetah, Maze2D, Kitchen, Ant, and Walker. For LIBERO, we use a Transformer- based noise predictor with a decoder architecture comprising 12 layers, 12 attention heads, and hidden embedding dimension of 768. Following DP (Chi et al., 2023), we apply dropout with a rate of 0.1 to both the input embeddings and attention weights. The number of diffusion timesteps is 500. When conditioning used, we incorporate a Transformer encoder with 4 layers to encode the condition tokens, which include sinusoidal timestep embedding and projected observed trajectory tokens (all mapped to
37

---

Published as a conference paper at ICLR 2026
the same embedding dimension). In this encoder-decoder setup, causal masking is applied to ensure autoregressive generation. In the unconditioned case, we prepend the sinusoidal timestep embedding to the imput sequence and use BERT-style encoder-only Transformer. All environments (Cheetah, Ant, Kitchen, Maze2D, Walker, and LIBERO) are trained using the Adam W optimizer with a learning rate of 10-1, weight decay 10-3, B1 0.9, and 82 = 0.95. Layer normalization is applied before each Transformer block for stability. The observation, planning, and action horizons follow the same settings used for the planner in each environment.
For the IDQL-based policy, we align all hyperparameters for Cheetah and Ant with the original IDQL implementation. using an observation, planning, and action horizon of 1. Hence, in IDQL-based ones, we do not consider autoregressive modeling, Similarly, for both cases, we use consider the coefficient before the contrastive improvement loss as 0.1.
D.2.4 HYPERPARAMETERS OF CONTRASTIVE IMPROVEMENT LOSS
We set Apriors Arl to be fixed as 0.1 across all settings. is set to be the 0.05 x Sprior during the beginning of each epoch.
D.3 CONNECTION TO BAYESIAN FILTERING
In the absence of explicitly designed latent variables, model can be interpreted form of Bayesian filtering (Chen et al., 2003). Under a general formulation of the hidden Markov model (HMM) (Rabiner& & Juang, 1986) with an additional latent dependency on observation (c x), the latent process over с captures the underlying stochasticity present in the demonstration data, which arises from both the environment dynamics and the behavior policy. In this view, the latent variable acts as compact and expressive representation that summarizes the uncertainty in past observations, thereby improving the prediction of future observations. This, in turn, facilitates more robust policy learning and planning in the general settings.
E EXTENDED RELATED WORKS
E.1 DIFFUSION MODEL-BASED DECISION-MAKING
Recent advances use diffusion models the planner and policy for both reinforcement learning (RL) and imitation learning (IL). RL agent aims to learn a policy that maximizes cumulative rewards through interaction with environment (Sutton et al., 1998). The agent observes sequence of transitions (st, a,, Ft, (t, 1,+,), where St E S denotes the state, a, E A the action, rt E R the received reward, and $1+1 the next state. The goal is to learn policy iaaa s) that maximizes the expected return: π = arg max; Ex o 7're], where 7 E [0, 1) is the discount factor. In contrast, IL (Hussein al., 2017) focuses on learning policies from expert demonstrations, often without access to the reward signal. A common approach is behavior cloning (BC) (Pomerleau, 1991), which formulates IL as a supervised learning problem by maximizing the likelihood of expert actions given observed states, i.e., learning a policy x(a s) that closely imitates the expert policy т., (a) s).
planning methods are commonly used to approximate the sequence of future states and actions from given current state. By leveraging the conditional generation capabilities of diffusion models-suce as guidance techniques (Dhariwal Nichol, 2021; Ho & Salimans, 2022)-thest methods can generate plans (i.e., state trajectories) that satisfy desired properties, such as maximizing expected rewards. Taking Denoising Diffusion Probabilistic Models (DDPM (Ho et al., 2020))
Diffussio Planner Diffusion-
approaches as example, these methods learn generative model over expert trajectories T = {(Se, a....... (Sr, ar)) } by modeling a forward wwwiindd-iiiiggggggooo process: q(x* x = ~N(x;; / x (1 onk I), and parameterized denoising model po x() to reverse the process. Here, k denotes the diffusion step, x" is clean sub-sequence sampled from the expert trajectory т. and 0% controls the variance schedule at step kt
During inference. trajectories generated by starting from Gaussian noise and iteratively denoising through the learned reverse process. This generation can be optionally conditioned on the initial state or other guidance signals y, such as rewards, goals, or other constraints: f ~ po(T So: y).
38

---

Published as a conference paper at ICLR 2026
These methods generally fall into two main categories: (1) learning a joint distribution over state- action trajectories, as in Diffuser (Janner et al., 2022), or (2) learning only state trajectories via diffusion and using an inverse dynamics model recover actions, in Decision Diffuser (DD) (Ajay et al., 2022). Beyond these, several variants extend diffusion-based planning in different directions. For example, Latent Diffuser (Li, 2024) plans in high-level latent skill space improve generaliza- tion and LDP (Xie et al., 2025) plans with high-level latent actions directly from high-dimensional action-free demonstrations. Other approaches incorporate multi-task context enhance adaptation and performance in unseen tasks, including MetaDiffuser (Ni et al.. 2023), AdaptDiffuser (Liang et al., 2023), and MTDiff-p pp (He et al., 2023). In addition, recent efforts have explored various extensions of diffusion planning, such as ensuring safety during generation (Xiao et al., 2025), handling multi-agen scenarios (Jiang et al., 2023; Ajay et al., 2023b), learning skills (Liang et al., 2024b), and application in RI. from human feedback (RIHF) (Dong et al., 2024).
Diffusion Policy In contrast to liffusiionbbsed planners, Diffusion Policy methods directly parameter- ize the policy (a s) using diffusion models. For example, Diffusion Policy (Chi et al., 2023) uses a diffusion model to generate actions with expressive, multimodal distributions. DPPO (Ren et al., 2025) extends this idea by modeling two-layer MDP structure, where the inner MDP represents the denoising process and the outer MDP corresponds to the environment. This framework enables fine-tuning of diffusion-based policies in RL settings. Another line of work integrates diffusion models with model-free methods for offline RL by using diffusion models as to model the action distributions (Wang et al., 2022: Hansen- -etrsschh et al., 2023; Chen et al., 2023; Lu et al., 2023).
Recent explorations have also aimed to unify diffusion-based planning and policy learning within a single framework. For example, the Unified Video Action model (UVA) (Li et al., 2025) and Unified World Models (UWM) (Zhu et al., 2025) leverage diffusion models to jointly model planning and action generation, demonstrating scalability on large-scale robotic tasks with pre-training. spirit, Ada-Di fusser provides a general framework that can be integrated into both diffusion planners and diffusion-based policies. However, Ada-D f fuser differs in its explicit modeling of latent factors that influence the data generation process. By incorporating latent identification directly into the diffusion process, Ada- -Di E fuser enables more structured, context-aware decision- decision-making in partially observable and dynamically changing environments.
similar
E.2 LATENT BELIEF STATE LEARNING IN POMDP
In partially observable Markov decision processes (POMDPs), single-step observations are typically insufficient for making optimal decisions. A common strategy to overcome this limitation involves encoding an agent's 's history, encoding past observations and actions into belief state that captures a distribution over latent environmental states. Although such belief representations can, in theory, support optimal policy derivation (Kaelbling et al., 1998; Hauskrecht, 2000; Gangwani et al., 2020), their exact computation depends on full knowledge of the transition and observation models. This requirement quickly becomes intractable in high-dimensional settings.
To address this, recent work has focused on learning approximate belief representations directly from data. Notable approaches include those using recurrent neural networks (Guo et al., 2018) and variational inference methods (Igl et al., 2018; Gregor et al., 2018). which enable agents to encode temporal structure and uncertainty into compact latent embeddings. These representations are then used to inform downstream policy learning, optimizing for cumulative rewards.
This direction also aligns with developments in meta-reinforcement learning and non-stationarity, where belief states or Bayesian embeddings are used to capture hidden task contexts. Agents trained across distribution of tasks can use these latent variables to infer new environments and adapt quickly (Zintgraf et al., 2021; Nguyen et al., 2021; Huang et al., 2021; Rakelly et al., 2019; Xie et al., 2021; Feng et al., 2022; Feng & Magliacane, 2023; Liang et al., 2024a; Wang & Huang, 2025). For example. MetaDiffuser (Ni et al., 2023) incorporates task context conditioning input to diffusion-based decision models. Similarly, Pertsch et al. (2021) 1) and Zeng et al. (2023) use similar variational objectives (ELBO loss) to learn latent skill priors and predictive information for RL, where these latents greatly help policy learning.
Our approach diverges from these by offering theoretical guarantees on the identifiability of latent factors from minimal temporal observations. Rather than depending on diverse multi-environment data, we introduce framework that captures the full data generation process in RL using diffusion
39

---

Published as a conference paper at ICLR 2026
models. In contrast to MetaDiffuser, which assumes static task-level context, model treats the latent context as dynamic, time-evolving process that governs both environment transitions and agent behavior, capturing the underlying temporal structure of RL trajectories more faithfully.
E.3 AUTOREGRESSIVE DIFFUSION MODELS
To model temporal consistency and dynamics in sequential data such as videos and audios, recent work has incorporated autoregressive structures into diffusion models. These approaches differ in how they condition on prior time steps during generation and can be categorized into two main categories. (1) Conditioning on clean (denoised) inputs (ZZheng et al.. 2024; Gao al., 2024b; Blattmann et al., 2023)). At each time step i, the denoising model is conditioned on the previously denoised outputs {x2x)= (x(x x x²2), where x is the current noisy input, and x² denotes the clean (fully denoised) observations from earlier time steps. (2) Conditioning on noisy inputs ((Ho et al., 2022; Chen et al., 2024; Xie et al., 2024b; Sand-AI, 2025)). These methods instead condition on previous time steps at their corresponding noise levels. This setting can be further divided into two cases: (a) fully noisy conditioning (Ho et al., 2022): the model conditions on all prior time steps at the same noise level k: Pe (x 1
x xX-1 x x ). (b) partially noisy conditioning: each previous time step i < t conditioned at its own noise level ki which may vary over time: po(x. ko- x 11
- x x ...., Xy" Specifically, Diffusion Forcing (DF) (Chen et al., 2024) proposes general framework in which each time step x; assigns independent noise level. In contrast, other works adopt time- deeeddepndeent noise schedules that vary with the temporal index (Xie al., 2024b; Sand-AI, 2025; Wu et al., 2023).
To model the causal generative process of RL trajectories, our approach also employs time-dependent noise scheduling capture temporal dynamics. However, unlike prior work, further integrate the identification of latent factors directly into the denoising process. This is achieved through a structured reinforcement step during training and zig-zag inference procedure test time, enabling our model to more faithfully recover the underlying causal structure in sequential decision-making.
E.4 SUMMARY
To sum up. we compare our approach with representative diffusion- and meta-learning based base- lines (Table A11). Diffuser, DP. IDQL, and DD do not model or infer latent contexts; DF adopts autoregressive denoising but still lacks context inference. Meta-Diffuser, LIIAA,, and DynaMITE learn latents via meta-learning but omit our ciciimall-ufffciieee block design and backward refinement. LDCQ and LDP model only high-level latent actions/skills without explicit context identification. In contrast, our method jointly models latent factors, employs full autoregressive denoising with zig-zag sampling, and introduces a backward refinement mechanism that enables identifiable latent contexts.
Method
Latent Factors
AR Denoising
Min. & Suff. Obs.
Ours
Yes (dyn., rew., act.)
Yes
Yes (refine, zig-zag)
Diffuser/ DP/ DD/ / IDQL
No
No
No
DF
No
Yes
No
Meta-Diffuser / /LIACC/ DynaMITE
Yes (dyn., rew, only) Yes (hi-level act.)
No
No
LDCQ
No
No
LDP
Yes (hi-level act.)
No
No
Table A11: Comparison with representative baselines on whether they model latent contexts, use autoregressive (AR) denoising, and enforce minimal & sufficient observation blocks.
F BENCHMARK SETTINGS AND ILLUSTRATIONS
LATENT CHANGE FACTORS DESIGN
We consider the latent change factors dynamics and rewards. We consider the Half-Cheetah and Ant environments from the OpenAl Gym suite, which are widely used MuJoCo locomotion benchmarks (Brockman et al., 2016) for evaluating continuous control algorithms. In Half- Half-Cheetah,
40

---

Published as a conference paper at ICLR 2026
Figure A3: Illustrations of the Benchmarks. From left to right: Half-Cheetah, Ant, Walker, Franka- Kitchen, Maze2D, and LIBERO.
the agent is planar bipedal robot with 17-dimensional state space and a 6-dimensional continuous action space, where the goal move forward by applying torques to six actuated joints. In Ant, quadrupedal robot operates in 3D space with a 11 I-dimensional state space and an 8-dimensional action space, requiring more complex coordination across its four legs. In both environments, the reward encourages forward velocity while penalizing excessive control inputs and, in the case of Ant, also promotes stable contact with the ground. We consider variants of the Half-Cheetah environment to study changes in dynamics, specifically Cheetah-Windse and Cheetah-Wind-S, which introduce external wind forces applied to the agent. In Cheetah- Wind-E, an opposing wind force is applied at the beginming of each episode and remains constant throughout, defined as fw 10 + 5 sin where is the episode index. For this case, since с change over episode, we use data from several consecutive episodes to estimate it. In Cheetah- Sheeeta--WindSS, the wind force varies every time step according to the same formula f. = 5 + 5 (0..0.) with now representing the time step in each episode. We also consider variations in the reward function. In Cheetah-Dir-E, the reward depends on a time-varying goal direction, requiring the agent to alternate between moving forward and backward. Specifically, the reward at episode is defined
Tt dg U2 0.1 [a.]|
where i the agent's forward velocity, a, the action vector (torques applied), and de +1] indicates the target direction at time t. The direction signal d, changes, giving non-stationary reward function that challenges the policy to adapt to shifting goals. Specifically, we consider
d, (5 sin (2ml//200))
where (-) denotes the sigmoid function, a controls the sharpness of the transition, and determines the switching period. This formulation induces smooth periodic change in the preferred direction of movement, requiring the policy to adapt to gradually shifting objectives.
We also consider directional reward variant for the Ant environment, denoted as Ant-Dir-E, where the agent is required to alternate its movement direction over time. The reward function at time step t defined as
T2 = (2d, - of - 0.1||a,11", where of is the velocity of the agent's torso along the x-axis (forward direction), a, is the 8- dimensional action vector, and d, E [0, 1] is smooth directional signal. Similarly, we define de as:
dj (5 sin(2nt/200)),
where (.)) denotes the sigmoid function. This formulation causes the preferred movement direction to alternate approximately every 100 steps. Notably, for these settings with periodic changes (i.e., where latent factors do not evolve at every timestep), we estimate the latent variables periodically and perform refinement in the causal diffusion model only when changes are detected. This follows the same overall framework, but operates at a coarser temporal resolution aligned with the latent change frequency.
F.2 OVERVIEW ON OTHER BENCHMARKS
Fig. A3-A4 give the illustrations on the used benchmarks. Specifically, other than Cheetah and Ant introduced before, for others, we consider the basic settings in offline RL. Specifically,
Maze2D. Maze2D tasks focus on goal-din goal-directed navigation in a 2D plane, where the agent must traverse maze-like environment reach specified targets. These settings designed to evaluate an agent's ability to reason spatially and follow optimal trajectories based solely on positional and velocity observations.
41

---

Published as a conference paper at ICLR 2026
Lift
Can
Square
Figure A4: Illustrations of RoboMiiic Benchmark.
The Franka-Kitchen environment (Gupta et al., 2020) involves robotic arm interacting with series of articulated objects in a realistic kitchen setting. Tasks are composed of multiple stages, such as opening doors or toggling switches, and are intended to assess agent's capability in handling long-horizon, multi-step manipulation.
Kitchen
Walker. The Walker2D environment features a two-legged robot that must learn to walk and balance using continuous torque control. The agent's 's objective to maintain forward motion while remaining upright, which requires learning dynamic stability and coordination.
LIBERO (Liu et al., 2023). The Libero benchmark offers a diverse set of continual learning tasks focused object manipulation and generalization:
LIBERO-Object: The robot is required to manipulate a variety of novel objects through pick-and-piee operations. Each task introduces previously unseen objects, encouraging the agent to incrementally build knowledge about object-specific properties and behaviors.
LIBERO-Godd All tasks share common object set and spatial layout, but vary in goal specifications. This setup tests the agent's ' ability to continually adapt to new task intents and motion targets without changes in the visual scene.
ABBERO-Spatial: Tasks involve repositioning bowl onto different plate locations. Al- though the objects remain fixed, the spatial configurations vary across tasks, requiring the robot to incrementally acquire relational spatial understanding.
RoboMinic. RoboMimic (Mandlekar al., 2021) provides set of manipulation tasks based on human telcoperation demonstrations, varying in difficulty and required precision:
Lift: The robot arm is tasked with lifting small cube off the table. This task serves as foundational manipulation scenario focused on grasping and vertical motion.
Can: The robot must retrieve a cylindrical can from cluttered bin and place it into a designated smaller container. This task introduces greater complexity due to object shape and the need for accurate placement.
fine -nneggraied insertion task where the robot picks up a square nut and places onto vertical rod. This is the most challenging of the three, requiring precise alignment and control for successful completion.
Square:
G OTHER DETAILS ON ADA-DIFFUSER
G.1 LATENT ACTION PLANNER
For the latent action planner, we align our settings with those used in LDP (Xie et al., 2025), specifically focusing learning directly from image-based demonstrations. We first variational autoencoder (VAE) to extract latent representations from raw images via image encoders, An inverse dynamics model is then trained to recover actions a, from pairs of latent states (z,, Z(+1). A planner is subsequently trained to forecast future latents z. Hence, the objective function is
(IDM(E z) = Er,, |leg (āxi Ck2 2k. 2k+1. t) - <l|2] where where is the time step and is the diffusion step.
42

---

Published as a conference paper at ICLR 2026
In our framework, we treat the latent factors high-level latent actions that influence the evolution of z. These latent factors are jointly used with z to perform both inverse dynamics modeling and latent forecasting, enabling structured planning in the latent space.
We follow the experimental settings established LDP (Xie et al., 2025). Specifically, we use expert demonstrations alongside action-free and suboptimal demonstrations. All hyperparameters and architectural choices for the diffusion models are kept identical to those used in the original LDP implementation. We also directly utilize the pre-trained image encoder provided by LDP. The only modification in our framework is the introduction of an additional latent factor с trained by our latent factor identification stage, which incorporated into the model to enhance latent action planning.
G.2 NOISE SCHEDULING
In the autoregressive setting, we consider a monotonic increasing denoising schedule {k1..... ....... kr j. In practice, we use a linear schedule where k; = T step used in both training and sampling. We segment the sequence into temporal blocks of length T2 + 1 (T, = To in all settings), and slide the time window forward by one step at a time. This design ensures that the denoising steps progressively increase across the block, aligning the diffusion process with the underlying temporal structure. Such a schedule encourages early steps to rely more on strong priors and later steps refine based on more contextual information. Additionally, for better illustration, Fig. A2 provides a detailed illustration of the zig-Zag sampling process within temporal block of 4 timesteps.
: K. with denoting the maximum diffusion
H SPECIFIC DESIGN CHOICES FOR BASELINES
For all baselines, unless otherwise specified, we use the same set of diffusion parameters detailed in Appendix D.2.2-D.2.3. Below, we provide additional details how specific methods evaluated. While their diffusion backbones remain consistent as in Appendix 0.2.2-D.2.3, these methods include custom design choices and method-specific hyperparameters that are evaluated accordingly.
H.1 DETAILS ON LILAC AND DYNAMITE
In these settings, we extend both LILAC and DynaMITE by incorporating a context encoder to infer latent context variables c,, following their respective designs. Both methods learn belief state embeddings from historical observations. For a fair comparison, we use the same latent identification network architecture as in our framework, but modify the inputs according to each method's 's assumptions.
Specifically, LILAC and DynaMITE condition their inference networks solely on the historical trajectory X1a, without access to current and future information. Additionally, consistent with their original implementations, we do not include a separate prior head on top of the GRU; both methods share the encoder for posterior inference and prior prediction. And the primary difference (in implementation) between these two methods lies in the temporal context used: LILAC maintains the full belief over the entire history, i.e., it conditions X to infer Cr+11. while DynaMITE uses only the most recent context, i.e., it infers C++1 based solely on X4-
All other hyperparameters are aligned with those used in our Stage 1 training. The estimated context variables are then provided as additional conditioning inputs to the diffusion- based models.
H.2 DETAILS ON DIFFUSION FORCING
For Diffusion Forcing, we adopt the same autoregressive noise schedule as in our method, which accounts for causal uncertainty, similarly to the formulation in Eq. D.1 (Chen et al., 2024), to ensure fair comparison. Additionally, we the Monte Carlo Guidance (MCG) mechanism introduced in (Chen et al., 2024) for Maze2D, following the original setup. For all other environments, we use the same classifier guidance scheme as the other baselines maintain consistency in evaluation.
43

---

Published as a conference paper at ICLR 2026
I ABLATION ANALYSIS
I.1 TRAINING/INFFERENCE TIME ANALYSIS
We conduct all experiments on 4x NVIDIA A or 8x RTX 4090 GPUs, depending on the model scale and environment requirements. The main computational overhead in our framework arises from two components: (i) the latent factor identification network, and (ii) the in the diffusion model. During sampling, the additional cost comes from zig-zag latent exploration and latent variable sampling. However, these steps do not substantially increase either training or inference time.
steps
To quantify this, we report the training and inference speed of our method compared to the base models DD and DP across all environments (Table A12). Our framework introduces only moderate computational overhead - typically 1.2-1.3x the runtime of vanilla diffusion backbones, corre- sponding to roughly 20-30% extra training time and inference latency. This cost can be further reduced through parallel latent sampling, lightweight context encoders, or refinement only at infer- ence. Moreover, we additionally evaluate a Picard-accelerated variant (Table A13, Shih al. (2023)), where iterative refinement is parallelized by conditioning each denoising step on previously denoised nodes. With Picard iteration, inference time drops 0. 0.7-0.8x of our default iterative sampler while maintaining comparable performance, demonstrating the potential for further acceleration.
Training Time (sec/epoch) Ours vs DP
Inference Latency (ms)
Environment
Ours vs DD
Ours vs DD
Ours vs DP
Cheetah
72.1 60.1 (1.20)
6 (1.20)
182 114 (1.16)
160 125 (1.28)
Ant
79.5 64.3 (1.24)
76.0/62.00 (1.23)
148 / 125 (1.19)
172 139 (1.24)
Walker
85.3 67.1 (1.27)
81.5 (1.27)
182 / 144 (1.28)
170 130 (1.31)
Maze2D
90.2 72.0 (1.25)
88.3 (1.28)
184 / 149 (1.24)
196 152 (1.29)
Libero
104.0 81.0 (1.28)
(1.31)
209 169 (1.24)
219 162 (1.35)
Kitchen
117.8 88.1 (1.34)
115.3 / 85.0 (1.36)
228 / 180 (1.27)
211 168 (1.26)
Table A12: Training and inference time comparison for Ada-Di Saa-Difffuser-plante variants. We report absolute times (sec/epoch or sec/rollout) and relative
and
overheads.
Environment
Ours (sec)
Ours+Picard (sec)
Cheetah
1.51
1.15
Ant
1.67
1.25
Walker
1.83
1.40
Maze2D
1.94
1.47
Libero
2.18
1.62
Kitchen
2.45
1.84
Table A13: Picard-accelerated inference.
I.2 ABLATION RESULTS
1.2.1
FULL RESULTS SUPPLEMENT TO TABLE
Table A14 presents the full ablation results across all environments, as supplement to Table 2. Overall, the results highlight the importance of the key components in causal diffusion modeling: latent identification and autoregressive diffusion, both of which are critical for performance.
1.2.2 NOISE SCHEDULE: LINEAR VS. LOGISTIC vs. SIGMOID
We adopt a linear noise schedule by default since any monotonic, bounded schedule suffices to model the geteraeneration process and linear is simple and stable in practice. To validate this choice,
44

---

Clustering We assess whether the learned latent space organizes states by the underlying context on the Cheetah wind-change task, where the ground-truth latent evolves as (t) = + 5 sin We sample 1000 time steps, discretize fa (t) into five goul-ffee bins to define target clusters, embed the corresponding observations into the 20-dimensional learned latent representation, and run Published as a conference paper at ICLR 2026
compare linear, logistic, and sigmoid schedules three representative tasks. As shown in Table A15, performance remains stable across schedules with no significant differences, supporting our default choice.
Task
Schedule
Performance
Cheetah
linear
-68.9
logistic
63.6
sigmoid
-70.4
Maze2D
linear
161.4
logistic
157.6
sigmoid
168.5
Frache-Kitchen
linear
0.70
logistic
0.72
sigmoid 0.66
Table A15: Ablation noise schedules. "Performance" is the task score (higher is better for Maze2D/Kitchen; lower magnitude negative is better for Cheetah as per the benchmark).
1,2.3 EFFECT OF TEMPORAL BLOCK LENGTH ON LATENT IDENTIFICATION
We further analyze the impact of temporal block length on latent identification. As shown in Fig. A5, the results are consistent with findings reported the main paper. When the number of observations insufficient (e.g., 4), identification performance degrades. Performance improves when the block length is in moderate range (5-20), indicating that sufficient temporal context is beneficial. However, using overly long blocks (> 20) introduces redundancy and increases optimization difficulty, which in turn harms performance.
MSF J
0.8
0.8
MSE
0.1
0.4
0.2
25
30
100
Size
Size
Figure A5: Identification results (MSE linear probing and R²) versus the length of temporal
blocks. Left: Cheetah with viryiig wind; Right: Cheetah with time-varying rewards.
k-means with 5. We compare our method with LILAC and DynaMIE, together with ablation that without refinement. Results are given in Fig. A6.
1.2.4 LATENT PROBING: EFFECT OF BACKWARD REFINEMENT AND ZIG-ZAr
To test whether backward refinement and zig-zag primarily help by correcting posterior mismatch, perform latent probing analysis CHEETAH with changing wind. We linearly map the learned latent representation to the ground-truth wind variable using a simple least-squares probe (trained subset of blocks and evaluated held-out blocks). Table A A16 reports the mean squared error
45

---

Published as a conference paper at ICLR 2026
Ours
LILAC
Ours (w/o Refine)
DynaMITE
Figure A6: Clustering (I-SNE)results on Cheetah wind-change.
(MSE) of this probe for three variants: (i) the full model with backward refinement and zig-zag; (ii) without refinement; and (iii) without zig-zag.
Variant
MSE
Full (with refinement zig zag)
0.18
w/o refinement
0.28
w/o zig-7ag
0.23
Table A16:
Linear probing MSE for recovering the ground-truth wind latent CHEETAI (changing wind). Lower is better.
Analysis. The full model achieves the lowest MSE, indicating more accurate recovery of the latent wind. Removing backward refinement yields the largest degradation (0.18 -> 0.28), consistent with the role of refinement in letting future evidence within a block update the latent posterior and reduce temporal lag. Disabling zig-zag also harms accuracy (0.18 -> 0.23), suggesting that alternating conditioning helps align the denoising trajectory with the latent dynamics rather than purely following the forward temporal pass. Together, these results support our claim that both components reduce posterior mismatch and improve latent identifiability, which turn benefits planning and control in settings with evolving hidden factors.
I.2.5 ON THE EFFECT OF PLANNING AND EXECUTION HORIZONS: LONG-HORIZON PLANNING
We study the robustness of our approach under increased planning and execution horizons (Tp and T.). Specifically, we evaluate on two challenging tasks Franka- Kitchen- and Libero- Long, where the original settings are Kitchen (I, 16. T. 8) and Libero (T, 10, T, = 8). Results are in Fig. A7. When we increase these horizons, we observe that the baselines, DP and DF, suffer significant performance drops. In contrast, Ada-Di f fuse rmanntains relatively high
46

---

Published as a conference paper at ICLR 2026
Kitchen-partial
Kitchen-partes
70
70
Success Rate
Sucvess Rate
65
60
60
55
55
Ours
50
Qurs
DF
45
DF
50
Planning Horizon
Exccution Horizon
Libero-Long
Libero-Long
610
Success Rate
Sucoess Rate
50
60
40
40
Ours
Ours
DP
Planning Horizon
Execution Horizon
Figure A7: Results with different planning and execution horizons. We evaluate on Kitchen-partial and Libero-L experiments.
performance. This demonstrates that modeling the underlying causal generative process, through autoregressive structure and latent representations, enables better long-horizon planning. Although we do not explicitly impose latent variables, our model implicitly learns representations that can track stochasticity and support smooth control.
47

---

Published as a conference paper at ICLR 2026
-57.8 6.7
-58.1 9.0
962.1
+3
1: 23.6
Oracle
300.7
-58.
Meta
- -2..4 3.9
-65.3= 11.2
949.82241
24.1
-39.2+766
22.2
Ours
296.4
-68.9 7.6
8.7
25.6
+45.8
24.5
-73.5
Ours
943.3
285.3
Diffuser
+ LILAC
.95.3
-105.6 14.5
37.6 37.6
-62.6 11.1
1: 32.
00
229.4.
Diffuser DynaMITE
-79.2 = 11.0
-94.3=96
930.41 29.5
-76.3+ 11.7
27.5
27.5
262.8
MetaDiffuser
-89.7 6.5
-106.8 11.4
912.5 37.9
65..2
41.0
945.9
9.6
10.2
45.8
-85.6= 18.3
47.0
DF
-105.8
-102.0
195.4
902.1
-1204 12.
-148.5+ 9.8
850.8 = 54.2
-102.4 18.2
188.61 39.2
Diffuser
Cheetah- Wind-E (c*)
Cheetah- Wind-S (c*)
Cheetah- Dir-E (c")
Cheetah-Vel-EE66 Vel-E (c")
Ant-Dir-E (c")
Environment
Cheetah-
Table A4: Results (average returns) on Ada-Diffusez-PPannee with latent factors that affects dynamics and rewards. c* and c" indicate the changes on dynamics and reward, E and S represent the episodic and time-step changes. The results are with random seeds. The bold ones the best-performing ones, excluding meta-learning and oracle ones.
48

---

Published as a conference paper at ICLR 2026
+DQQ (Oracle)
-41.616.2
+44.7 6.1
969.8= 39.2
:33.8 6.5
25.4
Ours
309.6
+ IDQL
.48.5 7.9
-48.0 7.2
37.5
-38.6=7.77 7.7
32.7
Ours
965.0
295.8
+ynaMITE
-59.0& 11.2
-63.4 6.7
938.6=409 49.4
-59.4 6.5
29.5
IDQL
269.3
-97.5
-87.8 12.2
902.41:45 45.2
11.4
25.6
IDQL
-80.2
204.61:
+PP (Oracle)
-52.0
42.3
972.4 37.5
-39,8
312.5 37.2
Ours
-58.5 :.64 4.6
-52.9-95 988
7.2
49.4
960.7
-1.0=
Ours
+ DP
290.4
DynaMITE
-72.2
-76.5 15.6
949.6 36.1
-72.7+
27.0
2
DP
-104.8 + 10.9
-120.6 11.5
892.5 60.8
6.5
:82.5 41.2
DP
Cheetah- Wind-E (c*)
Cheetah- Wind-S (c*)
Cheetah- Dir-B (c))
Chectah- Velah--el-ECCe (c))
Ant-Dir-E (c")
Environment
Table A5: Results (average returns) on ada-Diffuser--Polic with latent factors. c* and c" indicate the changes dynamics and reward, E and S represent the episodic and time-step changes. The results are with 5 random seeds. The bold ones are the best-perforning ones, excluding meta-learning and oracle ones.
49

---

Published as a conference paper at ICLR 2026
Environment
Diffuser
DD
DF
LDCQ
Ours (DD)
Mixed
52.6 2.3
75.2
73.7
73.3 0.5
74.6 1.6
Partial
55.8+ 1.9
57.3# 1.2
68.6 2.4
67.8# 0.8
70.1
Table A6: Results (success rate (%)) on www.Difffuuee--laaaneeeee without explicit latent factors on Franka-kitchen environment. The results are with 5 random seeds.
Environment
Diffuser
DD
DF
LDCQ
Ours (DD)
umaze
113.5 2.8
114.8 3.2
116.7 = 2.0
134.2 + 4.1
148.6: 3.7
medium
121.5 5.6
129.6 2.
149.4 7.5
125.3
148.6
large
123.0 4.8
131.5 42
159.0
150.1 2.9
161.4 3.2
Table A7: Results Ada-Diffuser-Plamne without explicit latent factors Maze-2D environment. The results are averaged across 5 random seeds.
Environment
Diftuser
DD
DF
LDCQ
Ours (DD)
medium-cxpert
106.2 + 0.7
108.8 = 2.0
105.4 # 3.2
109.3 + 0.4
115.7 = 2.1
medium
79.6 9.8
82.5 1.6
66.2 19
69.4
83.6
medium-replay
70.6 0.6
75.0
72.2
68.5
74.3
Table A8: Results on Ada-Diffuser-Plante without explicit latent factors on Walker envi- ronment. The results are averaged across 5 random seeds.
Environment
Ours (DP)
Spatial
78.3 + 3.9
79.2 +: 4.2
Object
92.5 2.6
93.4= 2.8
Long
50.5
62.6
Table A9: Results on Ada-Di Policy without explicit latent factors on Libero environ- ment. The results are averaged across 5 random seeds.
Component
Type (i): Full Trajectory
Type (ii): State-Only
Model Backbone
ID U-Net (Ronneberger 2015)
Transformer (DIT)
Architecture
Kemel size: 5; chanmels: (1,2,2,2); base: 32
Hidden dim: 256; head dim: 32 (Walker, Kitchen, Maze2D), 8 (LIBERO)
DiT Blocks
Optimizer
Adam, Ir=3 14
Batch Size
64
128
Training Steps
Diffusion Timesteps
150
200
Observation Horizon
(Kitchen), (LIBERO), 10 (others)
Planning Horizon T
(Cheerah), (Amt)
16 (Kitchen), (TIRERO), 32 (others) 8 (Kitchen, LIBERO), 10 (others) CFU
Execution Horizon To
Guidance
Inverse Dynamics Model
2. layer embed (64), 3 layer MLP (128), IM steps
Lass
Table A10: Planner configurations for type (i): full trajectory generation and type (ii): state-only generation with inverse dynamics.
50

---

Published as a conference paper at ICLR 2026
Cases
Cheetah-1
Cheetah-2
Ant
Maze2D
Walker
Kitchen
RohoMimic
LIBERO
Original
-52.9
295.8
161.4
115.7
0.70
0.85
93.4
w/o refine
-82.0
-60.7
261.2
156.5
107.4
0.63
0.78
90.2
w/o zig-zag
-91.6
-56.1
258.3
147.6
107.9
0.59
0.75
91.6
same NS
-89.7
-62.4
259.7
140.1
105.8
0,56
0.72
85.2
random NS
-62.9
266.4
146.3
109.1
0.76
88.5
Table A14: Ablation on Design Choices. We conduct ablation studies across a diverse set tasks, including: Cheetah- Cheetah-Wind-S with planner-based approach (denoted Cheetah-1 1 in the table), Cheetah- Windd- with a diffusion policy (Cheetah-2), Ant-Dir-E (policy, IDQL-based), Maze2D- Large (planner), Valker2D-Medium-ESs (planner), Kitchen-Partial (planner), LIBERO-Object (diffusion policy), and RoboMimic-Can.
51


Reviews:
[{"review_id": "PJlrV5Iuy8", "reviewer": "Reviewer_SLzj", "summary": "The paper presents Ada-Diffuser, a diffusion-based sequence modeling framework for decision-making tasks. The proposed method has two stages: the first stage extracts latent variables that characterize environmental dynamics from minimal observations, through variational inference. In the second stage, the inferred latent factors will be used to condition the autoregressive denoising process of a diffusion planner or policy. During inference, the framework adopts a zig-zag sampling strategy that alternates between sequence denoising and latent refinement to further improve the performance. Experiments show that Ada-Diffuser achieves strong performance across various setups, including those without action labels or without explicit latent structures.", "questions": "- Line 422: Missing content “Table A4 and are“", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "- The framework integrates latent variables that capture environmental dynamics with a causal diffusion model to facilitate better decision making, with theoretical analysis on latent identification with minimal observations\n- The experiment extensively evaluates the proposed model across different task domains and settings, consistently highlighting its effectiveness and robustness.\n- The paper is well organized and clear to read.", "weaknesses": "Encoding state or state action sequences into latents using reconstruction or ELBO loss has been widely explored by prior work, in which the latent embeddings, a summarization of the input sequence dynamics, have versatile usages to guide decision making, such as skill prior [1] or predictive information [2]. It would be helpful to elaborate on why, intuitively, the extracted latents would work better when coupled with a causal diffusion planner/policy.\n\nIn this work, specifically, the latent variable is modeled using (s, a, r) triplets with an extra reward; however, when using the latent variable, the planner generates sequences of states or state-action pairs without explicitly accounting for accumulated rewards. Even though dynamics can be modeled by the latent, the planning performance may remain restricted by the quality of offline demonstrations.\n\n[1] Pertsch. et al. Accelerating Reinforcement Learning with Learned Skill Priors. CoRL. 2020\n\n[2] Zeng et al. Goal-Conditioned Predictive Coding for Offline Reinforcement Learning. NeurIPS \n2023.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "M2fFVeIAnu", "reviewer": "Reviewer_ZU5m", "summary": "This paper proposes Ada-Diffuser, a unified diffusion-based decision-making framework that explicitly learns latent dynamic factors from partial observations and adapts its planning/policy behavior accordingly.\nIt provides (1) identifiability theory showing latent states can be recovered from short temporal windows, and\n(2) a practical architecture combining latent inference + causal diffusion model for planning and control.\nExperiments on locomotion and robotic manipulation show improvements in adaptation to changing dynamics, rewards, and even action-free demos.", "questions": "see weaknesses", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "Novelty & Importance: Explicitly modeling latent dynamics inside diffusion-based decision-making is a meaningful step. Previous decision diffusers or trajectory diffusion methods ignored latent variables.\n\nTheoretical Grounding: The identifiability theorem supports the proposed latent inference mechanism, which adds credibility.\n\nUnified Framework: The modular latent inference + diffusion planner/policy is elegant and works for both planning and policy learning tasks.\n\nRobust Experiments: Demonstrates adaptation to dynamic changes, reward shifts, and even recovering hidden action variables.", "weaknesses": "1. Dataset diversity & realism feel limited.\nExperiments are described as locomotion and robotic manipulation benchmarks; however, I didn’t see evidence of large-scale, human demonstration datasets (e.g., MimicGen) or messy real-world logs. Without these, it’s hard to judge robustness of latent inference under noisy observations, embodiment shifts, or human suboptimality. (This is especially relevant since the paper claims adaptation to reward/dynamics changes and action-free demos.) \n\n\n2. Suggested text you can paste:\n“All evaluations are in simulation-style locomotion/manipulation. It would help to include a more realistic demonstration dataset (e.g., human-collected manipulation or MimicGen-like settings) to substantiate generalization of the latent inference mechanism.”\n\n\n3. Latent interpretability is under-analyzed.\nThe method identifies latent variables, but the paper doesn’t really show what they mean in practice (e.g., do they correlate with friction, mass, goal, terrain, tool state?). Some qualitative probes or counterfactual rollouts would strengthen the claim that the model captures meaningful hidden structure rather than just serving as an extra embedding. (The theory guarantees identifiability up to an invertible transform; interpretability needs separate evidence.) \n\n\n4. Assumptions may be optimistic for real deployments.\nAssumption 3 requires a distinguishability condition on a second-order kernel ratio over adjacent times; the paper argues it’s “mild,” but in high-noise real data or partial coverage, that separability can fail. It’d be helpful to discuss failure modes or provide stress tests when the assumptions are violated. \n\n\n5. Compute and latency reporting are thin.\nGiven the added latent-inference block plus causal diffusion, it would be good to see training/inference cost, planning latency (ms), and how the cost compares to strong baselines (e.g., Decision Diffuser variants, model-based world models). No clear numbers are reported alongside the main claims; a wall-clock profile would help practitioners decide whether the approach is deployable.\n\n\n6. Baselines could be broader and more diagnostic.\nThe text positions the method within diffusion-based decision-making, but I’d expect latent-variable or POMDP-aware baselines as well (e.g., strong world-model policies with explicit belief tracking) to isolate the contribution of latent inference inside diffusion. This would clarify whether the gains come from the diffusion prior or from better latent modeling per se.\n\n\n7. Ablations on the latent module are missing/limited.\nSince the contribution hinges on the latent factor identification block, I’d like to see ablations like: (a) remove the latent module; (b) freeze latents; (c) mis-specify latent dimension; (d) shorter/longer temporal windows for identifiability. This would show the sensitivity and where the modeling actually helps.\n\n\n8. Action-free demonstration results need stronger evidence.\nThe abstract claims recovery of hidden action variables from action-free demos. That’s a big promise; I couldn’t find a strong, quantitative study measuring action recovery accuracy or downstream policy benefits against alternatives (e.g., inverse dynamics baselines). Please add a clear metric and a head-to-head comparison. \n\n\n9. Generalization under distribution shift is unclear.\nIf the environment’s latent process changes abruptly (e.g., reward re-weighting, contact discontinuities), does the model adapt online in a few steps, or does it require re-training? Some online adaptation or few-shot tests would clarify how practical the approach is.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "WqVMfAcoL5", "reviewer": "Reviewer_DUvX", "summary": "The paper proposes Ada‑Diffuser, a diffusion‑based framework for decision‑making that explicitly models time‑varying latent factors influencing dynamics and/or rewards. Under a set of injectivity conditions, the posterior over the latent at time t is identified from a short temporal window. This motivates block‑wise latent identification via a VAE prior/posterior with an ELBO and a causal, autoregressive diffusion model.", "questions": "1. Can you provide an example of real-world scenarios indicating when Assumptions 2–3 plausibly hold (or fail)?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Incorporating explicit latent context into diffusion-based decision models addresses an important gap in decision making under uncertainty. The paper presents a unified approach for planning and policy learning. \n\n2. The identifiability result provides a principled justification for blockwise latent inference and motivates a denoise-then-refine design. The proof strategy follows the nonparametric identification literature and is clearly presented in the appendix. \n\n3. The empirical evaluation is broad, spanning (i) latents in dynamics and reward, (ii) latent actions and action-free settings, and (iii) settings without explicit latents.", "weaknesses": "1. The theorem uses $x_{t+1}$ (a future observation) for identifiability, but test-time planning cannot access the future. The paper proposes a zig-zag procedure that uses a predicted future observation during refinement. While ablations show empirical gains, the method would be challenging to quantify the error incurred by using predicted futures for posterior alignment.\n\n2. The discussion of the effectiveness of latents in environments without explicitly designed latents is not fully convincing. Doesn’t expert action noise still imply deterministic dynamics? How does the latent help a policy learned from datasets collected under stochastic behavior policies?", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "GsvG2lKpXQ", "reviewer": "Reviewer_5j5o", "summary": "This paper introduces a unified framework that explicitly incorporates latent dynamic inference into generative decision-making, particularly focusing on partially observable environments. The authors highlight that existing generative models (like Diffusion Models) often overlook evolving latent factors essential to environment transitions, reward structures, and high-level agent behavior. Algorithmic Innovations include causal autoregressive denoising schedule, denoise-and-refine training procedure and zig-zag sampling for inference. Extensive experiments across locomotion, navigation, and robotic manipulation benchmarks show improved performance in accurate latent inference, long-horizon planning, and adaptive policy learning compared to diffusion baselines and variants augmented with existing latent modeling techniques.", "questions": "- Questions on Identifiability and Online Inference\n\n1. Causal Timeline and Lookahead in Zig-Zag Sampling: During online inference (zig-zag sampling), the refined latent \\hat{c}_t^post is conditioned on x_{t+1}^{k_2}. In a standard sequential decision-making loop, s_t+1 (and thus x_t+1) is only observed after the action a_t has been executed. Please clarify the precise temporal availability of the conditioning inputs x_t^0 x_t^{k_1} x_{t+1}^{k_2} used to calculate \\hat{c}_t^post during online generation of a_t. Does x_{t+1}^{k_2} represent a predicted state from the generative model itself? If so, this suggests \\hat{c}_t^post is recursively dependent on the model's own (noisy) future prediction, which could introduce instability or drift, especially for long planning horizons (e.g., T_p=32).\n\n2. Practical Implication of Block Size (Beyond Stability): The ablation shows that performance degrades with blocks that are too small (≤4) or too large (>20). Given that the theory proves 4 steps are sufficient, what specific phenomena in the practical optimization environment (e.g., approximation error, high variance) necessitate using blocks in the range of 6 to 20 steps, rather than the theoretically suggested minimum? Does this suggest a fundamental limitation of using finite-capacity function approximators (like the GRU/MLP encoder) to capture the non-parametric identifiability conditions (Assumptions 2 and 3)?\n\n- Questions on Algorithmic Design\n\n3. Function and Tuning of the Contrastive Improvement Loss (L_{rel}): The refinement stage uses a contrastive improvement loss $L_{\\text {rel }}=$ softplus $\\left(\\log L_{\\text {post }}-\\log s g\\left(L_{\\text {prior }}\\right)+m\\right)$ to encourage the posterior latent (\\hat{c}_post) to produce better reconstructions than the prior latent (\\hat{c}_prior). This is crucial for avoiding posterior collapse, a common issue in VAE-based latent modeling. What is the impact of the margin hyperparameter m and the weighting coefficient λ_rel on the training process, particularly regarding the trade-off between maximizing reconstruction accuracy (low L_post) and maintaining latent expressiveness? Have the authors performed an ablation on m to demonstrate its necessity?\n\n4. Impact of Latent Dimensions: The details provided focus on the architecture (GRU encoder, MLP decoders) but not the specific dimensionality of the inferred latent context C for each environment. How does the choice of latent dimension d_c affect: (i) the satisfaction of the identifiability assumptions (e.g., Assumption 3 requires sufficient variability induced by c), and (ii) the required size of the minimal block T_x? Since the model captures high-level factors, rewards, or even latent actions, the complexity of c likely varies widely.\n\n- Clarification on the Optimization Roles of the Denoise-and-Refine Losses\n\n5. Parameter Update Paths: Please clarify which network parameters (θ for ϵ_θ, ψ for q_ψ, or ϕ for p_ϕ) are updated by which loss terms. Does minimizing L_post (reconstruction using the posterior latent c_post) update only θ, or does it also backpropagate through c_post to update the posterior encoder ψ? Similarly, does minimizing L_prior update θ and the prior network ϕ?\n\n6. Role of Stop-Gradient in L_{rel}: The contrastive loss L_rel uses a stop-gradient on L_prior(logsg(L_prior)). This typically means L_rel is designed to primarily improve the posterior network q_ψ by pushing c_post towards a representation that achieves better reconstruction compared to a frozen reconstruction accuracy baseline derived from the prior c_prior. Please confirm whether the optimization signal derived from L_rel is used solely to update the posterior encoder ψ, or if it also affects the denoising network θ.\n\n- Questions on Latent Correlation and Prior Quality\n\n7. Sensitivity to Latent Dynamics: The prior p_ϕ(c_t∣c_t−1) relies on the Markovian assumption of the latent context c. In cases where the latent dynamics are highly stochastic or non-Markovian (i.e., c_t is weakly dependent on c_t−1), the prior estimate c_prior may be very poor. Does the overall framework, particularly the contrastive loss L_rel, remain effective when the latent temporal correlation is weak, or is Ada-Diffuser implicitly relying on environments where the latent context exhibits strong temporal smoothness?\n\n8. Prior Collapse Mitigation: Since Stage 1 is an offline VAE optimization using L_ELBO, the term DKL(q_ψ||p_ϕ) serves to regularize the posterior towards the prior. Given the introduction of L_rel in Stage 2 (Causal Diffusion Model), which aims to make the posterior estimate significantly better than the prior estimate, how do the authors ensure that this competitive pressure does not induce posterior collapse (where the VAE minimizes the KL term to zero, ignoring the observations) in Stage 1, or that the training stages remain balanced?\n\n- Detailed Questions on Autoregressive Denoising and Complexity\n\n9. Sequential Steps vs. Total Diffusion Steps: Standard diffusion models perform K noise reduction steps. The autoregressive denoising process described appears to perform a sequential denoising across T time steps. If the trajectory length T is large (e.g., T=32), does the overall inference procedure involve T×K total denoising evaluations, or is the process structured such that the total number of ϵ_θ evaluations remains close to the standard K? Specifically, how does the sequential nature described in Equation (2) translate into the number of required calls to ϵ_θ(x^k, k, c)?\n\n10. Computational Overhead of Zig-Zag: The ablation shows that using a Picard-accelerated variant reduces inference time (e.g., 1.51s → 1.15s for Cheetah). Given that the zig-zag sampling scheme requires an additional \"update\" step where the posterior latent c_post is calculated using a predicted noisy next step x_{t+1}^{k_2}, what is the approximate percentage of the total inference time consumed specifically by the latent inference and refinement steps (non-denoising operations) in the default zig-zag sampler?", "limitations": "", "rating": 4, "confidence": 3, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. Originality. The paper demonstrates high originality, primarily through the unique integration of theoretical identifiability results into the architecture of a causal diffusion model. Specifically, establishing sufficient conditions under which latent factors can be identified from short temporal windows of RL trajectories is a novel theoretical contribution. This theoretical insight directly guides the algorithmic design, leading to the development of the minimal-sufficient block concept and the backward refinement step for online latent recovery, which contrasts with prior latent-augmented diffusion approaches. The combination of autoregressive denoising, latent factor identification, and the tailored zig-zag sampling scheme represents a creative and non-trivial synergy of existing ideas applied to sequential decision-making.\n\n2. Quality. The quality of the work is high, supported by rigorous technical details and thorough empirical validation. The identification theory provides a strong foundation for the method. The experimental evaluation is comprehensive, covering eight environments across 23 different settings, including latent factors affecting dynamics and rewards, latent actions from action-free data, and environments without explicit latents. Furthermore, the paper provides detailed ablation studies confirming the effectiveness of the key design choices, particularly the refinement and zig-zag sampling mechanisms, which significantly reduce probing MSE and improve performance.\n\n3. Clarity. The paper is well-structured. The framework's modular design, dividing the process into latent factor identification (Stage 1) and the causal diffusion model (Stage 2), enhances clarity. The authors clearly define the Latent Contextual POMDP with Time-dependent Context and formalize the data generation process using Structural Causal Models (SCMs). The purpose and justification for complex components like the latent identification loss and zig-zag sampling are explicitly tied back to the theoretical findings (Theorem 1), providing a clear narrative thread.\n\n4. Significance. The significance of Ada-Diffuser lies in its ability to handle time-varying, unobserved latent factors which are pervasive in real-world applications such as robotics and autonomous driving. By providing a unified generative framework for planning and policy learning that effectively adapts to these latent variations, the method offers substantial improvements over state-of-the-art diffusion baselines and meta-RL approaches in settings with explicit latent factors. Moreover, the finding that latent modeling can improve performance even in environments without explicit latent factors suggests broad utility beyond traditionally defined POMDPs.", "weaknesses": "1. Mismatch between Theoretical and Practical Block Sizes: The central theoretical result (Theorem 1) establishes that four consecutive time steps (x_t−2:t+1) are sufficient for latent identifiability. However, the authors state that in practice, they do not strictly limit the block size to four, finding that slightly larger blocks (typically between 6 and 20 steps) lead to more stable optimization and better performance. This practical necessity for larger blocks suggests that the assumptions required for the theoretical four-step result may be too stringent or idealized for the complexity of the empirical environments and the parametric function approximators used. The reliance on large blocks somewhat diminishes the immediate algorithmic guidance derived from the \"minimal-sufficient block\" result. This adds to the suspicion that the authors reverse-engineered the justification, deliberately finding an obscure mathematical theorem to back a model they already wanted, instead of allowing the model to emerge naturally from the theory.\n\n2. Dependency on Future Observations during Online Inference: The latent identifiability theory relies on incorporating future observations (x_t+1) to form an accurate posterior q_ψ(c_t∣x_{t−T_x:t+1}). In online inference, future observations are unavailable, creating a mismatch. The zig-zag sampling scheme attempts to resolve this by conditioning on a predicted next noisy step x_{t+1}^{k_2}​ to refine \\hat{c}_t^post. This means the quality of the online latent inference hinges critically on the accuracy of the diffusion model's prediction of x_{t+1}^{k_2}. If the future prediction is inaccurate, the latent refinement may degrade, potentially leading to cascading errors in long-horizon planning. This key dependency is not fully analyzed or quantified in the provided sources, particularly in terms of how errors propagate when using the noisy future estimate.\n\n3. Increased Model Complexity and Training Stage Separation: Ada-Diffuser requires training two distinct stages: (1) an offline VAE for latent factor identification (optimizing the ELBO loss), and (2) the causal diffusion model (optimizing along with the contrastive loss L_d−r). This two-stage training process is inherently more complex and potentially less robust than end-to-end training. While the authors claim the VAE introduces minimal computational overhead, the overall architectural complexity, including the GRU encoder, multiple MLPs for priors/posteriors, and the explicit contrastive loss L_rel demands careful tuning of multiple hyperparameters (λ_KL, λ_prior, λ_rel, m).\n\n4. Novelty of Component Integration: While the overall framework is original, the individual components draw heavily from existing literature. The concept of latent belief state learning in POMDPs is standard (e.g., LILAC, DynaMITE, VAE usage). The autoregressive noise schedule is derived from recent autoregressive diffusion models (e.g., DF). The novelty rests heavily on the specific way the blocks, refinement, and zig-zag work together. However, a deeper analysis justifying why this specific combination is necessary—rather than sufficient—to satisfy the identifiability constraints (especially compared to simpler VAE-in-loop approaches like DynaMITE/LILAC integrated with diffusion) would strengthen the contribution.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Poster)