Masked Skill Token Training for Hierarchical Off-Dynamics Transfer
Authors: Zeyu Feng
Haiyan Yin
Yew-Soon Ong
Harold Soh
Venue: ICLR 2026 Poster (2026)
Keywords: Tranfser Learning
Skills
Hierarchical RL
Embodied AI
Primary area: reinforcement learning

Abstract: Generalizing policies across environments with altered dynamics remains a key challenge in reinforcement learning, particularly in offline settings where direct interaction or fine-tuning is impractical. We introduce Masked Skill Token Training (MSTT), a fully offline hierarchical RL framework that enables policy transfer using observation-only demonstrations. MSTT constructs a discrete skill space via unsupervised trajectory tokenization and trains a skill-conditioned value function using masked Bellman updates, which simulate dynamics shifts by selectively disabling skills. A diffusion-based trajectory generator, paired with feasibility-based filtering, enables the agent to execute valid, temporally extended actions without requiring action labels or access to the target environment. Our results in both discrete and continuous domains demonstrate the potential of mask-guided planning for robust generalization under dynamics shifts. To our knowledge, MSTT is the first work to explore masking as a mechanism for simulating and generalizing across off-dynamics environments. It marks a promising step toward scalable, structure-aware transfer and opens avenues to explore multi-goal conditioning, and extensions to more complex, real-world scenarios.

Published as a conference paper at ICLR 2026
MASKED SKILL TOKEN TRAINING FOR HIERARCHICAL OFF-DYNAMIC TRANSFER
Zeyu Feng', Haiyan Yin', Yew-Soon Ong! Harold Soh
CFAR and IHPC, Agency for Science, Technology and Research (A*STAR), Singapore
Department of Computer Science, School of Computing, National University Singapore, Singapore Nanyang Technological University (NTU), Singapore "Smart Systems Institute, National University of Singapore, Singapore {feng zeyu, yin hainhaiyan haahaaaea-staaaaesaaaaaeeaaaaaaa ea-star. edu. sg asysongentu. ha haroldecomp comop nus. edu.sg
ABSTRACT
Generalizing policies across environments with altered dynamics remains key challenge in reinforcement learning, particularly in offline settings where direct interaction or fine-tuning is impractical. We introduce Masked Skill Token Train- ing (MSTT), fully offline hierarchical RL framework that enables policy transfer using observation-only demonstrations. MSTT constructs a discrete skill space via unsupervised trajectory tokenization and trains a skill-conditioned value function using masked Bellman updates, which simulate dynamics shifts by selectively dis- abling skills. diffusion-based trajectory generator, paired with fcasibility- based filtering, enables the agent to execute valid, temporally extended actions without requiring action labels or access to the target environment. Our results both discrete and continuous domains demonstrate the potential of mask-guided plan- ning for robust generalization under dynamics shifts. To our knowledge, MSTT is the first work to explore masking mechanism for simulating and generalizing across off-dynamics environments. It marks a promising step toward scalable, moneture-awant transfer and opens avenues explore multi-goad conditioning, and extensions to more complex, real-world scenarios.
INTRODUCTION
Reinforcement learning (RL) agents often struggle generalize when deployed in environments that differ from their training conditions (Pinto et al. 2017 Clavera et al. 2019 Abdolshah et al. 2021 [ansen al. 202I; Feng et al. (023). In many practical scenarios involving navigation and manipulation, deployment environments exhibit structural changes (e.g.. blocked passages or new obstacles) that invalidate parts the agent's learned behavior repertoire. The changes not typically require new tasks or capabilities, but render some previously executable behaviors infeasible. A natural question is: can agents adapt to such dynamics shifts with minimal additional supervision, i.e, without extensive interaction full retraining?
In this work, address this question in practical setting where the agent is trained offline in source environment and deployed zero-shot in a target environment with altered dynamics. At deploy- ment, the agent receives single observation-only demonstrationn--itbuu action annotations -that illustrates feasible behavior in the new environment. These demonstrations may come from human demonstrators be inferred from video, and are far cheaper obtain than full action-labeles trajec- tories or interactive feedback. This setting enables lightweight, low-overhead adaptation and is well suited to real-world applications where interaction is expensive or infeasible.
We introduce Masked Skill Token Training (MSTT), a hierarchical RL framework for adaptation to feasibility constraints using only offline data. The key insight is that dynamics shifts can be abstracted as constraints over a learned space of temporally extended skills. MSTT learns these discrete skills from souree trajectories using vector-quantized variational autoencoder (VQ- VA-) (van den Oord et al., 2017] Mentzer et al.] 2024), resulting in a compact and reusable skill vocabulary. During training, MSTT simulates environment changes by randomly masking subsets of skills, and learns

---

Published as a conference paper at ICLR 2026
a feasiblityy-condttonedd critic via masked Bellman operator that propagates value only through feasible skill sequences.
At test time, MSTT infers a binary skill mask from a single observation-only demonstration in the target environment, and plans using the learned critic under the inferred constraints. Unlike prior off-dynamics RL approaches such as DARC (Eysenbach et al. 2021) and VGDF (Xu et al. 2023a), which require action-labeled data target interaction, MSTT operates fully offline and uses only observational signals to guide transfer. In contrast to standard hierarchical RL methods (Bagaria & Konidaris 2020 Qiao al. 2025), which assume that learned skills remain valid during deployment, MSTT explicitly models and adapts to shifts in skill feasibility. MSTT uses a diffusion based trajectory generator trained on source data. At deployment, it samples candidate trajectories and discards those whose encoded skill tokens violate the inferred mask. This enables behavior synthesis aligned with test-time constraints, without needing to explicitly condition the generative model on the mask or retraining it for each deployment scenario.
We evaluate MSTT on continuous domains, specifically the Maze2D in D4RL, FetchReach and Habi- tat environments with test-time modifications and requirements. MSTT significantly outperforms offline RL and transfer RL baselines, demonstrating robust adaptation using only single observation- only demonstration per environment. These results indicate the promise MSTT as a step toward enabling structure-aware transfer in dynamic environments with minimal supervision. To our knowl- edge, this work is the first to tackle the challenging setting of skill-level adaptation under structural dynamics shifts using only observation-onl only demonstrations and fully offline training.
To summarize, our main contributions are:
We formulate new off-dynamics transfer setting where agents adapt using only observation-only demonstrations, without access to target actions or environment interaction.
We propose Masked Skill Token Training (MSTT), a hierarchical RI framework that models synamics-inducee skill constraints using binary masks, and learns feasiblityyccondiioneed critic via masked Bellman updates.
We integrate a diffusion-basse trajectory generator to enable robust skill execution under test-time constraints, and demonstrate strong generalization across structurally altered environments.
2 RELATED WORKS
Policy Transfer under Dynamics Mismatch. Generalizing RL policies under dynamics mismatch is key challenge for real-world deployment. DARC (Fysenbach et al.) (02T) addresses Off- Off-Dynamics RL. (ODRL) (Lyu et al.] 2024) by learning reward corrections using discriminators trained on transition dynamics. VGDF (Xu u et al. (203a) filters transferable transitions using model ensembles, while H2O (Niu et al., 2022 reweights offline datasets to better match target distributions. Domain
randomization (Peng et al. 2018 trains robust policies via diverse simulated dynamics, but often requires extensive environment access. In the fully offline setting, DARA (Liu et al. 2022) modifies source rewards based on dlsssifier-sstmmatee transition alignment between source and target domains. BOSA (Liu et al.) 2024) proposes supported value optimization to better mitigate out-of-distribution transitions. IGDF (Wen et al., (0244 filters source domain data using contrastive cross-domain representation learning on states, action and rewards. xDIO (Raychaudhuri et al. (021)) imitate skills across different morphologies. Offline Meta-RL methods (Nam al.] 2022; Ni al.) 2023 Xu et al. 2022; 2023b) tackle this problem by meta-learning on multiple source tasks with different dynamics and rewards. For example, PDT (Xu et al. 2022) achieves task transfer by training decision transformer on inputs that contain both snippet of task demonstration and offline trajectory. HDP (Xu) jet al., 2023b) transfers adapter layers' parameters decision transformer by learning on multiple source tasks with task demonstration for each task as conditioning factor, generalizing to new task demonstration. Unlike these approaches, this work aims to achieve generalization across structural dynamics shifts without target interaction or explicit modeling of transition mismatches.
Hierarchical Skills and Discrete Latent Policies. Hierarchical RL (HRL) facilitates long-horizon decision makinn by structuring policies around temporally extended skills, improving exploration, credit assignment, and policy reuse. While classical HRL relies predefined subgoals options, recent works learn skill abstractions from data, often through discrete latent policies that represent reusable primitives compact codes. OPAL (Ajay al. 2021) learns discrete skill dictionary using
2

---

Published as a conference paper at ICLR 2026
02
02
$1
S
8
$
01
01
02
01
01
blocked
d1
actioos
S4
S3
84
S3
01
0/
(a) Four-room
(b) MDP Abstraction
(c) MDP with sink
Figure 1: The Four-room environment and its skill-level MDP abstraction. (a) The original Four-rou Four-room grid world, where the agent must reach specified goal rooms from varying start locations. (b) Abstract MDP representation, where rooms are states and inter-room transitions are treated as high-level skills (c) Modified MDP with blocked transition (e.g., the door from room 2 to room 3 is impassable), where exccuting the corresponding skill leads to sink state. MSTT Iearns a crasibility-condiiiee critic through simulated skill masking, enabling transfer to such target domains without interaction.
VQ- VAE vvan den Oord et al. 2017: Mentzer et al. 2024 Ozair et al., 2021) and Que T (Mete et al.
2024) encodes variable-lengee trajectory tokens for multitask reuse. DADS (Sharma et al. 2020) and
VALOR (Achiam et al.) 2018) focus on skill predictability and diversity to enhance generalization.
SiMPI. (Nam et al. 2022) learns skill embedding and high-level policy over the skill conditioned on task embedding. Li et al. 20021 2023) propose theoretically grounded and sample-efficient ways to learn options and embeddings.
particular, HiT-MDP (Li et al.) 2023) theoretically proves that opdiue- SMDP is homomorphic equivalents to MDP with option embeddings, greatly improving sample efficiency and convergence stability of policy optimization in the corresponding MDP. Furthermore, VMOC (Li al. 2025) extends this idea the of latent thoughts as options in an abstract reasoning space for efficient and effective reasoning. These approaches offer modularity and efficiency but implicitly assume that skills remain executable during deployment. Our method removes this assumption by introducing a skill-masking mechanism that enables policies to reason over partially-availabl skills.
Diffusion-Basse Generative Policy Learning. Diffusion models have emerged an expressive policy class for offline RL (Kang et al. 2023). Diffuser (Janner et al. 2022) formulates planning as trajectory denoising through conditional diffusion. Decision Diffuser (Ajay al., 2023) extends this idea to decision-man makinn, modeling return- conuiicondti trajectories. Diffusion Policies (Chi et al. 2023) apply diffusion to visuomotor control, producing pixel-to-action policies with high fidelity. Latent Diffusion RL /Venkatraman et al. 2024 leverages www-dimensional latent spaces to improve efficiency in trajectory generation. L.T.DoG (Feng al.) 2024) and DOPPLER (Feng et al., 2025) generates novel trajectories and actions to satisfy symbolic temporal logic specifications with diffusion models. These methods achieve strong results static environments, capable of composing hierarchical offline knowledge (Zhou et al. 2023; Du et al. 2023 Liang et al. 2024 Ma et al. 2024 but assume fixed dynamics and do not exphiitly handle feasibility constraints introduced by changes in environment structure. MetaDiffuser (Ni et al. (023) trains context encoder from trajectories with transition and reward information from different tasks. Our method augments this line of work by conditioning diffusion -aaed skill generation on masked skill availability, enabling policy rollout under dynamics-induced constraints.
3 SKILL-LEVEL MDP ABSTRACTION WITH MASKED FEASIBILITY
Recall that our goal is to enable diciiiionmaaingg under dynamics shifts, where certain skills may become infeasible due changes in transition feasibility (c.g.. due physical obstacles or altered connectivity D. In this section, we formalize the connection between skill feasibility and transition dynamics using a skill-level abstraction of the MDP.
In this work we primarily consider structural dynamics changes such locked passages or new obstacles, not continuous changes parameters of transition, e.g. mass, friction damping. Transfer across continuous changes requires training on trajectories from randomly sampled parameters, and pursue this future work.
3

---

Published as a conference paper at ICLR 2026
We introduce the notion of skill mask, binary indicator over the discrete skill space, which captures environment-speciinn feasibility constraints. This mask provides a compact mechanism to encode target environment dynamics at the skill level, and allows to analyze how such feasibility constraints affect value estimation and policy behavior. Our formulation supports off-dynamics policy transfer by explicitly modeling which skills remain executable in the target environment, without requiring direct interaction to capture its underlying dynamics or action-labeled trajectories.
To illustrate this abstraction, we use the Four-room navigation domain an example of discrete hierarchical reinforcement learning (Sutton et al. 1999) task involving structured decision-making with reusable behavioral primitives. As shown in Fig. I (a) the environment comprises four inter- connected rooms with doors at fixed positions. The agent is tasked with navigating from different start locations to designated goal rooms using previously learned skills that enable room-to-room transitions. These skills can be modeled as temporally extended actions.
Skill Abstraction of the Source MDP. We assume that a set of temporally extended skills are available, which enables the agent transition between high-level regions of the environment (e.g.. adjacent rooms). Based on these behaviors, we construct skill-level abstraction of the environment in the form of discrete augmented action-space Markov Decision Process (MDP) (Puterman 2014)
M ;= (S, A. p. 7), where S denotes the abstract state space (c.g., room indices), A denotes the set of skill-level actions, is the reward function defined over abstract states, p captures the skill kkill-induced transition dynamics, and y is the discount factor This high-level abstraction is illustrated in Fig. [(b) and serves as the foundation for reasoning over skill composition in our framework.
Skill-Level Feasibility Masking. In deployment settings, environment dynamics may differ from those observed during training, rendering certain skill-level transitions infeasible (e.g., a blocked passage between room and room 3). To capture such structural constraints, we augment the abstract state space with a dedicated sink state s1, which absorbs any transition resulting from the exccution of an invalid skill (Fig. [o. This sink state formulation provides a conservative but practical model of failure, where executing an unavailable skill results in termination-tiin reflects scenarios such as safety violations or irreversible errors. While our current formulation adopts this absorbing failure model, alternative semantics (e.g., remaining in the same state upon executing an invalid skill) are also possible and are left for future work.
The resulting target MDP is denoted MB == (S A,r,), where B {0, |]ls|((s] is blocking matrix indicating the feasibility of transitions. Big = denotes that the transition from s; to sj blocked under the current dynamics:
pn (m(s)|s a) 0 if Brj O; Pn (8,, a)
Σ (1 8
where transitions to blocked states are redirected to a sink state 81, which models failure due to infeasible skill execution. We also introduce a binary availability mask m(s, a) E {0, 1) 1)IiiAA1, where m(s, a) = denotes that action a is blocked at state
Masked Bellman Operator. To support learning in the presence of unknown blocked dynamics, we introduce masked Bellman operator that uses the binary availability mask m(s, a) E (0, 1] 1|(81||1 where m(s a) = 0 denotes that action is blocked at state s. The masked Bellman operator modifies the standard backup by restricting the value propagation to feasible transitions:
(TT (s, a) == r()) + + (m(s,a)p(s,)) a), Vm)
where 1I (8) a masked value operator that computes the maximum over available actions:
Q(s, a), if such exists
max Vi ae Amm.m)) r(s).
otherwise.
In other words, this masked operator enables the critic to account for feasibility constraints during learning by preventing value propagation through blocked or invalid transitions.
Theoretical Justification. Under blocked dynamics, iterative updates using the masked Bellman operator can converge to a near-optimal policy. Specifically, for any target MDP with infeasible transitions, there exists an availability mask that enables convergence to a policy whose value closely approximates that of the optimal policy under the modified dynamics:
4

---

Published as a conference paper at ICLR 2026
Skill Tokenization
Feasibilit-AAwar Option Sampling
Trajectory Encoder
Skill Codebook
X
S1
Feasbilit Check
S₂
D
Og
:
(' m))
TI-H
Sn
Resemple
Figure 2: Overview of our masked skill token learning framework, (Left) Skill tokenization encodes trajectory segments into discrete latent tokens via a VQ- VA-VAE encoder dg. (Right) At execution time, the diffusion model De samples candidate sub-trajectories, which are encoded into skill tokens (z). Feasibility is enforced via the binary mask m((z;; only valid skills are propagated for Q-learning.
Theorem 1. Let Qo: S A R, be an initial action-value function. There exists a mask such that the iterative update Qx+1 T Qs converges policy xm whose value under the blocked MDP MB satisfies:
th MB
VMm I1. say + 3(11 Pmin),
(3)
where Pmii == minl,a max,/ p(s'|s (s, a) and a, 3 are constants dependent on ".
For more detailed derivation and proof, please see Appendix
3.1 VALIDATING MASKED BELLMAN UPDATES
We empirically validate Theorem 1 using the example discrete environment shown in Fig- ure [(a] We simulate masked updates in the unblocked source MDP (Figure ((b) using feasibility mask where action 2 is unavailable state $2. The resulting value function VK is computed through K iterations of the masked Bellman operator. We evaluate performance across different levels of Pmin (which reflects the entropy of the transition dynamics).
(a) Est. & optimal values
(b) Gap (f -noom)
Figure 3: Performance of the masked Bellman operator using Eq. ( and Eq. (2). (a) Converged value estimates (Est.) at state 81 compared with the analytically computed optimal values under blocked dynamics. (b) loc -nomm of the difference between the estimated and optimal value functions. Both plots show that the masked Bellman operator produces accurate estimates, particularly as Pmii approaches 1.
Figure coomprres the value function estimated via the masked Bellman operator to the analyt- ically derived optimal values under blocked dy- namies. Notably, in low-entropy regimes where Pmin & 1, the largest gap (Figure between estimated and optimal values becomes negligi- ble. These findings empirically validate our the- oretical findings and highlight the potential of our approach in structured settings where skill transitions are predictable.
In summary, Theorem shows that masked value learning via the Bellman operator T can approxi- mate optimal behavior in target MDP MB, provided suitable feasibility mask m is applied. In hierarchical settings, where skills are temporally extended and induce low-entropy transitions, the approximation error remains small (e.g., when Pmii RS 1). This property allows dynamics shifts to be simulated in the source domain by masking skill availability during training. As a result, value functions can be learned entirely offline and still generalize to test environments without requiring access to their true dynamics. We leverage this insight to design the training framework described in the next section.
5

---

Published as a conference paper at ICLR 2026
Algorithm 1 Offline Masked Hierarchical Reinforcement Learning
Require: Trajectory dataset D, total training steps E, batch size B
Initialize critics Qu Q2 and targets Q, Q44
2: for e == =0 to E do
3: Sample [TiH J bE[M) uniformly randomly from D
4: Sample [m / randomly from {0,
5: Sample N target skills at 8'H and m using Algorithm
(b)
Select best skill o/(h) by Qui
6:
7:
Get clipped Gaussian noised version 5*(b)
Get Bellman loss ((0) from Eq. (5) with double Q-learning
8:
Update - 5jj - V P3 B Σ, (b) S.jj (1,2)
9:
10:
if E mod Ep = 0 then
11:
4 X4p (1 2) p/. [1,2}
12:
end
13:
end for
14:
return Up
MASKED SKILL TOKEN LEARNING FOR OFF-DYNAMICS RL
In this section, transition to learning continuous state and action spaces, where the discrete skill abstractions assumed in Section 3 are not directly available.
We introduce Masked Skill Token Training (MSTT), a hierarchical RL framework that learns compact skill vocabulary from offline data and supports policy transfer under altered dynamics. MSTT comprises three main components: (1) it encodes sub-trajectories into discrete skill tokens using VQ-VAE trained on offline demonstrations; (2) learns fcasibility-aware critic using masked Bellman updates that simulate skill infeasibility through random masking; and (3) it uses diffusion based policy generate candidate skills, filtering them according to an inferred skill feasibility mask obtained from single observation-only demonstration in the target environment. now describe each component in detail.
4.1 SKILL TOKENIZATION VIA UNSUPERVISED TRAJECTORY ENCODING
In continuous environments where the state-action space cannot be directly abstracted into a finite MDP, we obtain reusable skill representations by learning discrete encodings of raw trajectory segments from offline data. Specifically, we adapt the VQ-VAE discover latent behavioral primitives tub-ttrajectorie) that form a compact symbolic skill vocabulary.
Let TI:H ( 8 denote a sub-trajectory with horizon H, where we let 8h denote an observation with slight abuse of notation. We encode this trajectory using the VQ-VAE encoder deg to produce discrete latent token:
(o(TT.))
L),
where is the size of the learned codebook Z. Each token E Z indexes a cluster of similar uub-trajectories in the dataset and thus, the codebook Z provides a discrete abstraction over the continuous behavior space, with each token capturing consistent transition patterns. We find that even when trained with observation-only sub-trajectories, the learned codebook captures meaningful spatial and behavioral structure in the environment.
The VQ- VAE model is trained by minimizing the following loss function over sub-trajectories:
Lakill-enc
||T (=))l + (eg(@o(T1.H)) 211 +33 10
sg[=]]|
reconstruction
wodebook loss
commitment
where &a the decoder, sgl.1 denotes the stop-gradiee operator, and / is a hyperparameter that balances codebook commitment against the other loss terms.

---

Published as a conference paper at ICLR 2026
4.2 DIFFUSION- BARDED OFF-DYNAMICS SKILL ADAPTATION
To enable hierarchical decision-making under dynamics shifts, we consider each discrete skill token E 2 temporally extended action in an augmented action-space MDP. Each skill is instantiated using a trajectory segment ($1. A1..... (HH)
and the internal policy of the skill corresponds to the action sequence. Our formulation of the skill as fixed-horizon temporally extended actions can be viewed as simplified version of the options framework (Sutton et al. 1999. /Sutton & Barto 2018), where option o = (I, o B), I C S is the initiation set, To is the intra-option policy, and B:S + [0, 1] is the termination condition. In our case, the skill terminates deterministically after H steps.
................................
To support planning, we extend the masked Bellman operator (Eq, ① to operate over skills: max
0, = R(o) + (
(s',o',
(5)
where O(s') is the of valid skills at the successor state s'. The return R(o) is computed as the discounted sum of per-step rewards. If no feasible skill is available, the value defaults to the terminal
reward as in Eq. 2
We apply the VQ-VAE trajectory encoder de (Section 4.1 to associate each skill with a skill token by encoding the sub-trajectory of obser- vations 2 = do (Tr.)) Following the abstrac- tion introduced in Section [3 we define a skill feasibility mask m(z), which models which skills remain executable under altered dynamics: ie., m22) = indicates a feasible skill, while m(2) denotes infeasible one.
Algorithm 2 Off-Dynamics Skill Sampling
Require: Start Diffusson model Dv, encoder Øy. mask m
1:
< null
2: while 2 = null do
3: Sample T1.H D (s)
4:
f- 20 (T:H)
5:
if m(()) 0 then
We sample skills using a trajectory- level diffu- sion policy model Da., trained generate plau- sible observation-action sub-trajectories. Unlike
6:
2 null
Reject and resample
7:
end if
8:
end while
the VQ-VAE decoder, the diffusion model sup- ports conditioning on continuous start states
9: return Feasible skill T and
enabling flexible skill generation in diverse environments. Rather than conditioning the model directly on the feasibility mask, we apply a simple filtering procedure: sampled trajectories whose encoded skill tokens are masked out are discarded (Algorithm{)) 2). The accepted trajectory then executed as a temporally extended skill. This mechanism allows for zero-shot composition of valid high-level behaviors under novel dynamics, guided solely by the feasibility mask.
4.3 MODEL TRAINING AND INFERENCE
The overall training procedure is summarized in Algorithm At each iteration, we randomly sample trajectory segments and skill masks from the offline dataset, and update the Q-function by minimizing the squared Bellman eтrоr under the masked update rule (Eq. (5). Although the number of possible skill-mask combinations grows combinatorially with the size of the skill vocabulary, we find that the learned skill space is sufficiently structured and random masking provides effective coverage in practice. To stabilize training, we incorporate several standard techniques from deep Q-learning, including target networks (Mnih et al.) (016). clipped double Q-learning (Hasselt 2010) Fujimoto et al.) 2018), and target smoothing via noise injection (Simmons-Edie et al. 2019) similar
to DOPPLER (Feng et al. 2025).
At test time, MSTT infers the feasibility mask from a single observation-onf, demonstration. The demonstration can be non-expert as long as it covers path towards the goal. MSTT is capable of compositing skills present in the demonstration to achieve the goal. We begin by initializing the skill mask m=)) = 0 for all z E Z. Then, we extract all sub-trajectories T1- from the demonstration and encode them using the trajectory encoder to obtain skill tokens = = the (T(-H) ). For each observed
skill token, we set m(z) = 1, marking the corresponding skill as feasible. This inferred mask is passed to the learned critic and used to guide sampling via Algorithm 2
5 EXPERIMENTS
We evaluate MSTT's ability to achieve observation-oni -nnly demonstration following in previously unseen environments with altered dynamics, without additional training fine-tuning. Our primary
7

---

Published as a conference paper at ICLR 2026
Table Performance on Maze2D test suites.
Env Metric
BC
Diffuser
Diffuser
BC
DARA
MSTT (ours)
Return
0.0==0.0
83.5+:50.78
107 107.36 277.44147
17.76+18.24
39.13+31.95
163.97
Steps (L)
350.0+0.0
280.65++42.17
162.84427.02
157.49+15.61
310.87+31.93
186.92-66.85
Goal (t)
0%
73%
99%
99%
65%
89%
Return (1)
1.19=1889
15.86:+39.31
36.05=42.49
19.32+7.11
63.67+22.20
145.51=50.93
T2
Steps (U)
349.8=:1.98
336.7+32.96
298. 198.5466.766
187
187.37=57.37
286.33++22.20
211.01146.74
Goal (1)
1%
14%
44%
93%
93%
95%
Return (1)
8.44:12.88
0.0+0.0
27.94+22.04
22.94+21.73
61 31.83337.60
111,59-34.74
T3
Steps (L)
310.85=157.88
350.010.0
259, 43.3
280
290
288. 17+37.77660
249
349.33335.02
Goal (1)
32%
0%
73%
54%
78%
99%
Return (1)
2.87
33.12
57.11
33.34
54.87
140.35
Avg Steps (4)
336.88
322.45
240.27
208.38
295.12
215.75
Goal (1)
11%
29%
72%
82%
78.66%
94.33%
t: fine--tuuenn sene-tundd on target state-only demonstrations. ta: fine-tuned on target action demonstrations.
Table 2: FetchReach test results.
Metric
BC
Diffuser
Diffusert,
BCto
DARA
MSTT (ours)
Return (1)
-5.0010.0
46.95
-30.85 11.44
50,0+0.00
-49 49.64410.75
37.26617.43
Cost (1)
0.0+00.0
200.0+206.51
112.0--99.02
0.0+00.0
28.6-85665
0.0=0.0
Goal (t)
0%
99%
50%
0%
23%
88%
Failure (4)
61
0%
93%
0%
t: fine-tuned on target state-only demonstrations. ta: fine-tuned on target action demonstrations.
objective is to benchmark MSTT against state-of-the-arr offline and transfer RL methods in popular simulation environments (Figures 4, and 7 The empirical study includes visualizations of skill masks (see also appendix for more details) and generalization results to provide insights into MSTT's mask sampling behavior. We also include a case study showcasing the use of a Vision-Language Model (VLM) to infer demonstrations from visual observations.
EXPERIMENTAL SETUP
Environments. We evaluate MSTT in three simulated environments from Gymnasium-Ron Robottccs
suite de Lazcano et al. 2024) and Habitat-Lab (Savva et al. 2019 Szot et al.) 2021 Puig et al.
2024): Maze2D and FetchReach, commonly used in offline Janner et al.) 2022) and hierarchical
RL & Kim 2023) research, and a photorealistic 3D indoor navigation task in ReplicaCAD (Zzot
al. [202]) commonly used in embodied AI. Maze2D (Figure 4(a) requires long-horizon navigation with continuous control and sparse rewards (1 at the goal). Learning is difficult from non-expert,
fully offline data without hierarchical modeling. At test time, block certain paths (Eysenbach)
et al. 202T creating target environments where the original policy fails. We evaluate on three such variants (Figures ((d)), each allowing only a single feasible route. The agent receives a single locations-only demonstration per environment and must infer the required skills to follow it to the
(a) Maze2D
(b) Target
(c) Target 2
(d) Target
(c) Target (f) Target 2 (g) Target
Figure 4: Environments and test-time performance of MSTT Maze2D. (a) Source environment with continuous state and action spaces. Three distinct long-horizon paths connect the start location (green ball) to the goal (red ball), each requiring hundreds of low-level control steps. (b)-(d) Test-time variants with altered dynamics, where only one path remains feasible. The agent must navigate to the goal given observation-only demonstration. (c)-(g) Trajectorics generated by MSTT in target domains. The agents successfully finished navigation by selecting the correct path.
8

---

Published as a conference paper at ICLR 2026
Figure 5: Demonstration in Habitat ReplicaCAD. The agent must navigate the goal in this studio with pixel observations, following the demonstration images in order from left to right, then top to bottom.
Table 3: Habitat ReplicaCAD test results. Cost measures the average deviation of the trajectory coordinates from the demonstrated trajectory.
Table 4: MSTT test results with VLM generated state demonstration.
Metric
Maze TI
Maze T2
Maze T3
FetchReach
Metric
Diffuser Diffusert,
ID
MSTT
Return (f) 188. 188.6:447.9 152. 42. 111.
3555.6 -3,00=6644 239.4++35.5
Cost (1) 2.4 2.2:1.99
1811.
0.7=0.66
Steps ())
162.34447.6
198.7
35.0=16.4
Goal (1) 10%
16%
23%
75%
Goal (f)
94%
97%
98%
93%
Failure (L)
6%
3%
2%
0%
fine-tuned on target state-only demonstrations.
goal. In FetchReach (Figure (aa). the robot controls a gripper to reach target positions, receiving reward of 11 per step. In the target variant (Figure ((b)), ), a red cuboid defines a constrained region that incurs heavy penalty if entered. The agent must imitate safe coordinates-only demonstration to reach the goal while avoiding the penalty zone. In ReplicaCAD (Figure the agent must navigate to goal in photorealistic indoor environment using only RGB observations, following sequence of images as a demonstration.
Compared Methods and Evaluation Metrics. We compare MSTT against several popular of- fline RL. and imitation learning baselines: Diffuser (Janner et al., 2022) and Behavior Cloning (BC) (Pomerleau) 1988 Levine al. 22020). both of which learn from offline data but lack built-in mechanisms for adapting to new environments. To enable adaptation, we fine-tune Diffuser on target observation-only demonstrations, denoted as Diffuser,. We also include DARA (Liu et al. 2022). an offline dynamics-awe aware transfer RL method that extends DARC (Eysenbach et al. 2002 to the offline setting using action-annotated target data. Additionally, BCta denotes behavior cloning fine-tuned with action-labeled target demonstrations. Importantly, MSTT uses only state-only demonstrations in the test environment, while DARA and BCt. rely action annotations. Although these methods have more information about the target environment, we include their results for completeness and comparison. Another relevant method that can transfer with observation-only demonstrations is inverse dynamics (ID) modeling (Radosavovic et al. 2021). However, we find that it performs poorly the continuous control tasks Maze2D and FetchReach, due the difficulty of accurately learning the actions from locations alone without robot velocities. Therefore, experiment, Evaluation metrics include total episode return, number of steps to reach the goal, and success/failure rates over 100 trials. We report the mean and standard deviation across trials for each target task.
only include it in the Habitat
5.2 RESULTS
Can MSTT achieve transfer in the off--dynamics scenario with only observation-only demon- strations? MSTT achieves near-expert performance in all Maze2D variants, FetchReach, and ReplicaCAD, outperforming all baselines on average, as shown in Tables and 3, Offline RI. methods, such as Diffuser and BC, fail to adapt to the changed dynamics and perform poorly in transfer settings. Fine-tuning Diffuser on observation-only demonstrations (Diffuser,) improves performance, particularly in Target Env 1, where the model benefits from pretraining bias toward the correct path. However, it still lags behind MSTT overall. DARA performs inconsistently. showing moderate success in some settings (c.g.. Target Env 2), but struggles to generalize; this is likely due to limited hierarchical abstraction and a poorly trained domain classifier affected by dataset imbalance. Planning with inverse dynamics models performs poorly in ReplicaCAD with overfitting to high-

---

Published as a conference paper at ICLR 2026
8.30-
300
1003 1903
2000
(a) Fetch
(b) Target
(c) Demo
(d) Result
(a) Validation
(b) Skill transition
Figure 7:
Environments and test-time performance of
Figure 6: Analysis of MSTT's train- ing with random masks on skills in Maze2D, (a) Training loss on training masks and validation mask set. (b) The empirical skill indices transition matrix estimated from the dataset.
MSTT in Fetch. (a) The source environment, where the robot gripper must reach the goal location (red ball). (b) The target environment with modified dynamics: the red box indicates restricted region that the gripper must avoid. (c) Demonstration, where the robot gripper is tasked with reaching the goal with avoidance. (d) MSTT successfully executes a trajectory that reaches the goal while avoiding the red box.
dimensional inputs. These results highlight MSTT's strength in leveraging skill-level abstractions and feasibility- -aware planning for robust adaptation and demonstration following.
How does MSTT learn generalizable critic across diffferntt skills with random mask sampling?
The critic learned in MSTT can generalize across different skill combinations, as evidenced by the training and validation loss curves in Figure [(a] The training loss decreases steadily, while the
validation loss on a fixed set of masks that are not used for sampling during training also decreases
and remains close to the training loss, indicating good generalization.
While the skill library can grow large with a bigger codebook size, we find that the skills learned by VQ-VAE are often well-structured and the transitions between skills sparse. Figure 6(b) (skill transition matrix) illustrates the transition matrix of the skill tokens extracted from the source dataset, revealing structured, banded patterns in skill encoding and usage: transitions concentrate along patterned indices, indicating temporally consistent and semantically related motion primitives. Therefore, many masking patterns become irrelevant due to the sparse transitions in training, and the random masking strategy provides sufficient coverage of the skill space during training, allowing the critic to generalize well at test time. More visualizations in Section C.2 in the appendix further demonstrates the semantic meanings of the learned skill embeddings. MSTT effectively identifies and samples skills relevant to the demonstrated trajectory, while masking out irrelevant ones.
Case study: Using VLMs with MSTT. We investigate the use of modern vasion-language models (VLMs) to infer observation-only demonstrations directly from visual observations. This approach reduces annotation overhead and illustrates how foundation models can be used with MSTT for off-dynamics transfer. Specifically, we prompt VLM to generate coordinates sequences. Full prompt details, model outputs, and analysis are provided in the appendix. As shown in Table 4 MSTT achieves strong performance using these VIMM-inferred demonstrations, demonstrating potential of combining pretrained VLMs with our framework for low-effort generalization to new environments.
6 CONCLUSION AND FUTURE WORK
In this work, we introduced MSTT, an offline hierarchical RL framework that uses masked skill tokens and diffusion-based skills to enable skill transfer across off-dynamics environments. By conditioning value learning on binary skill masks, MSTT simulates skill infeasibility and avoids the need for action annotations or explicit dynamics modeling. Experiments in both discrete and continuous settings demonstrate robust transfer performance under dynamics shifts. Future work includes improving mask sampling strategies for efficiency, incorporating diverse dynamics parameters with domain randomization, extending to multi-goal transfer via goal-conditioned critics, and further leveraging I.LMs or VL.Ms to infer demonstrations from visual input to reduce reliance on human supervision. We believe MSTT forms an important step towards scalable skill transfer in complex, real-world settings.
10

---

Published as a conference paper at ICLR 2026
ACKNOWLEDGEMENTS
This research is supported by the National Research Foundation, Singapore and Infocomm Media Development Authority under its Trust Tech Funding Initiative, and the National Research Foundation, Singapore, under its Thematic Competitive Research Programme 2025 This research is partly supported by the MTI under its Al Centre of Excellence for Manufacturing (AIMfg) (Award W25MCMF014). Any opinions, findings and conclusions recommendations expressed in this material are those of the author(s) and do not reflect the views of National Research Foundation, Singapore and Infocomm Media Development Authority.
REPRODUCIBILITY STATEMENT
We have taken the following steps to enhance the reproducibility of our results. For our theoretical results we state required assumptions and problem setups Section [A.1] and include full proofs
in Section A.2 Our simulated environmental setup, dataset processing and experimental pipelines are documented in Section B We provide implementation details of the proposed algorithm and
computing environments in Section B of the appendix.
REFERENCES
Majid Abdolshah, Hung Le, Thommen Karimpanal George, Sunil Gupta, Santu Rana, and Svetha Venkatesh. A new representation of successor features for transfer across dissimilar environments. 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-9. PMLR, 18-24 Jul 2021.
Joshua Achiam. Harrison Edwards. Dario Amodei, and Pieter Abbeel. Variational option discovery algorithms. CoRR, abs/1807.10299 2018. URL
www.oo/gg/bbbb///ee /107 10299
Anurag Ajay, Aviral Kumar, Pulkit Agrawal, Sergey Levine, and Ofrr Nachum. OPAI:: offline primitive discovery for accelerating offline reinforcement learning. In 9th International Conference on Learning Representaations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021. OpenReview.net, 2021. pttps::::///peneeee
net/ htt.a//eeee.......ooooooooooooon
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 Inter- national Conference on Learning Representations, ICLR 2023, Kigali, Rwanda, May 1-5, 2023. OpenReview.net, 2023. URL
net www.ee//ooommddsssoooooooo.
Akhil Bagaria and George Konidaris. Option discovery using deep skill chaining. In International Conference on Learning Representations, 2020.
Cheng Chi, Siyuan Feng, Yilun Du, Zhenjia Xu, Eric Cousineau, Benjamin Burchfiel, and Shuran Song. Diffusion policy: Visuomotor policy learning via action diffusion. In Robotics: Science and Systems XIX, Daegu, Republic of Korea, July 10-14, 2023, 2023. URL https://doii org/
10. 15607 /SSS. 1123XXXX0222
Ignasi Clavera, Anusha Nagabandi, Simin Liu, Ronald S. Fearing, Picter Abbeel, Sergey Levine, and Chelsea Finn. Learning to adapt dynamic, real-world environments through meta-reinforcemmnnt learning. In International Conference Learning Representations, 2019.
Rodrigo de Lazcano, Kallinteris Andreas, Jun Jet Tai, Scungjae Ryan Lee, and Jordan Terry. Gymnasium roboties, 2024. URL htttp://githbb.... wtt.aaaa..oo//aaaaaaaaaaa Foundation/ Gymnas i umm-Roboiics
Yilun Du, Conor Durkan, Robin Strudel, Joshua B. Tenenbaum, Sander Dieleman, Rob Fergus, Jascha Sohl-Dickstein, Arnaud Doucet, and Will Sussman Grathwohl. Reduce, reuse, recycle: Compositional generation with energy-based diffusion models and MCMC. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato. and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning. volume 202 of Proceedings of Machine Learning Research, pp. 8489-8510. PMLR, 23-29 Jul 2023.

---

Published as a conference paper at ICLR 2026
Benjamin Eysenbach, Shreyas Chaudhari, Swapnil Asawa, Sergey Levine, and Ruslan Salakhutdinov. Off-dynamics reinforcement learning: Training for transfer with domain classifiers. In 9th Interna- tional Conference Learning Representations, ICLR 2021, Virtual Event, Austria, May 3-7, 2021, 2021.
Zeyu Feng, Bowen Zhang, Jianxin Bi, and Harold Soh. successor features. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 7219-7225, 2023. doi: N..100//CR.4488991002...000 22023.101612255
policy transfer with
Zeyu Feng, Hao Luan, Pranav Goyal, and Harold Soh. LTLDoG: Satisfying temporally-extended symbolic constraints for safe diffusion based planning. IEEE Robotics and Automation Letters, (10):8571-8577 2024. doi: N0.1109/LRA.2224.4444
Zeyu Feng, Hao Luan, Kevin Yuchen Ma, and Harold Soh. Diffusion meets options: Hierarchical generative skill composition for lemporally-extended tasks. In 2025 IEEE International Conference on Robotics and Automation (ICRA), pp. 11127641.
2025, doi: 00.109/1CRA55443.02555
Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine. D4RL: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:200..07219, 2020.
Scott Fujimoto, Herke van Hoof, and David Meger. Addressing function approximation error in actor- critic methods. In Jennifer Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, volume 80 of Proceedings of Machine Learning Research, pp. 1587-1596. PMLR, 10-15 Jul 2018.
Nicklas Hansen, Rishabh Jangir, Yu Sun, Guillem Alenyà, Pieter Abbeel, Alexei Efros, Lerrel Pinto, and Xiaolong Wang. Self-supervised policy adaptation during deployment. In International Conference on Learning Representations, 2021. URL
net/ forum?
id=o_ V-MjyyVVVVVV
Hado Hasselt. Double q-learning. In J. Lafferty, C. Williams, J.
Shawe-Taylor, R. Zemel, and A. Culotta (eds.), Advances in Neural Information Processing Systems, volume 23. Curran Asso- ciates, Inc., 2010. URL httttss:/pprcceediiiii www.ips.rrccomm ww/ppper.pppp...eee files/paper/
Michael Janner, Yilun Du, Joshua B. Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In International Conference on Machine Learning, ICML 2022, 17-23 July 2022, Baltimore, Maryland, USA, volume 162, pp. 9902-9915, PMLR, 2022.
Bingyi Kang, Xiao Ma. Chao Du Tianyu Pang, and Shuicheng Yan. Efficient diffusion policies for offline reinforcement learning. In A. Oh, T. Naumann, A. Globerson, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 67195-67211 Curran Associates, Inc., 2023.
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. arXiv preprint arXiv:2005.0164 2020.
Chang Li, Dongjin Song, and Dacheng Tao. The skill-action architecture: Learning abstract action embeddings for reinforcement learning. 2021.
Chang Li, Dongjin Song, and Dacheng Tao. Hit-MDP: Learning the SMDP option framework on
MDPs with hidden temporal embeddings. In The Eleventh International Conference on Learning Representations, 2023. httptttttss:sppeeeeeeeeeee
wtwp://orummiimmm Fi jAc
Chang Li, Yaren Zhang, Haoran Lv. Qiong Cao, Chao Xue, and Xiaodong He. Learning temporal abstractions via variational homomorphisms in option-induced induced abstract mdps. arXiv preprint arXiv:2507. 16473, 2025.
Zhixuan Liang, Yao Mu, Hengbo Ma, Masayoshi Tomizuka, Mingyu Ding, and Ping Luo. Skilldif- fuser: Interpretable hierarchical planning via skill abstractions in diffusion-basels task execution. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR),
pp. 16467-16476. June 2024.
12

---

Published as a conference paper at ICLR 2026
Jinxin Liu, Zhang Hongyin, and Donglin Wang. DARA: Dynamics aware reward augmentation in offline reinforcement learning. In International Conference on Learning Representations, 2022.
Jinxin Liu, Ziqi Zhang. Zhenyu Wei, Zifeng Zhuang, Yachen Kang, Sibo Gai, and Donglin Wang. Beyond ood state actions: Supported cross-domain offline reinforcement learning. Proceedings of the AAAI Conference on Artificial Intelligence, 38(12): 13944,133,,, Mar. 2024. doi: 10. 10.1609/ naai.v38i12.29302.
Jiafei Lyu, Kang Xu, Jiacheng Xu, Menghei Yan, Jingwen Yang, Zongzhang Zhang. Chenjia Bai, Zongqing Lu, and Xiu Li. Odrl: A benchmark for off off-dynamics reinforcement learning. In The Thirty-eight Conference on Neural Information Processing Systems Datasets and Benchmarks wtt.:/eewwwee/fouurr.m. aid=ap4xlkArry kArGy
Truck, 2024, URL
net forum www.oooumdddppppo aiddaaax1llaa
Xiao Ma, Sumit Patidar, lain Haughton, and Stephen James. Hierarchical diffusion policy for kinematics-awarc multi-task robotic manipulation. In Proceedings of the IEEE/CVE Conference on Computer Vision and Pattern Recognition (CVPR), pp. 18081-18090, June 2024.
Fabian Mentzer, David Minnen, Eirikur Agustsson, and Michael Tschamnen. Finite scalar quantization: VQ-VAE made simple. In The Twelfth International Conference on Learning Representations. 2024.
Atharva Mete. Haotian Xue, Albert Wilcox, Yongxin Chen, and Animesh Garg. Quest: Self- supervised skill abstractions for learning continuous control. Advances in Neural Information Processing Systems 38: Annual Conference on Neural Information Processing Systems 2024, NeurlPS 2024, Vancouver, BC, Canada, December 10 15, 2024, 2024.
Volodymyr Mnih, Adria Puigdomenech Badia, Mehdi Mirza, Alex Graves, Timothy Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. Maria Florina Balcan and Kilian Q. Weinberger (eds.), Proceedings of The 33rd International Conference Machine Learning, volume 48 of Proceedings of Machine Learning Research, pp. 1928-1931 New York, New York, USA, 20-22 Jun 2016. PMLR.
Taewook Nam, Shao-Hua Sun, Karl Pertsch, Sung Ju Hwang, and Joseph J. Lim. Skill-based meta-reinforrement learning. In International Conference on Learning Representations (ICLR), 2022.
Fei Ni, Jianye Hao, Yao Mu, Yifu Yuan, Yan Zheng, Bin Wang. and Zhixuan Liang. MetaDiffuser: Diffusion model as conditional planner for offline me. In Andreas Krause, Emma Brunskill, Kyunghyun Cho, Barbara Engelhardt, Sivan Sabato, and Jonathan Scarlett (eds.), Proceedings of the 40th International Conference on Machine Learning, volume 202 of Proceedings of Machine Learning Research, pp. 26087-2605... PMLR, 23-29 Jul 2023, hRLpst mlr.
html
Haoyi Niu, Shubham Sharma, Yiwen Qiu, Ming Li, Guyue Zhou, Jianming Hu, and Xianyuan Zhan. When to trust your simulator: Dynamics- aware hybrid offline- and-online reinforcement learning. In Advances Neural Information Processing Systems 35: Anmual Conference on Neural Information Processing Systems 2022, NewrlPS 2022, New Orleans, LA, USA, November 28- December 9, 2022, 2022. URL http:: ed3cd25201480 4802220488677777 (bbec3000.6577110cceee2626440o.3
nips wwwp:niwpppppp.pneeeeeeeeee http:pperr..llleeepprr/p////a/ /
html]
Sherjil Ozair, Yazhe Li, Ali Razavi, Ioannis Antonoglou, Aaron Van Den Oord, and Oriol Vinyals. Vector quantized models for planning. Marina Meila and Tong Zhang (eds.), Proceedings of the 38th International Conference on Machine Learning, volume 139 of Proceedings of Machine Learning Research, pp. $302-8313, PMLR, 18-24 Jul 2021.
Xue Bin Peng, Marcin Andrychowicz, Wojciech Zaremba, and Pieter Abbeel. Sim-to-real transfer of robotic control with dynamics randomization. In 2018 IEEE International Conference on Robotics and Automation (ICRA), pp. 3803-3810, 2018. doi: 10. 20.1100/CRA.2008..44460 8.660528.
Lerrel Pinto, James Davidson, Rahul Sukthankar, and Abhinav Gupta. Robust adversarial reinforce- ment learning. In Doina Precup and Yee Whye Teh (eds.), Proceedings of the 34th International Conference on Machine Learning, volume 70 of Proceedings of Machine Learning Research, pp. 2817-2826. PMLR, 06-11 Aug 2017.
13

---

Published as a conference paper at ICLR 2026
Dean A. Pomerleau. Alvinn: An autonomous land vehicle in neural network. In D. Touretzky (ed.), Advances in Neural Information Processing Systems, volume 1. Morgano Kaufmann, 1988.
Xavier Puig, Eric Undersander, Andrew Szot, Mikael Dallaire Cote, Tsung- Yen Yang, Ruslan Partsey, Ruta Desai, Alexander Clegg, Michal Hlavac, Yeon Min, Vladimír Vondruš, Theophile Gervet, Vincent- Pierre Berges, John M Turner, Oleksandr Maksymets, Zsolt Kira, Mrinal Kalakrishnan, Jitendra Malik, Devendra Singh Chaplot, Unnat Jain, Dhruv Batra, Akshara Rai, and Roozbeh Mottagii. Habitat 3.0: A co-habitat for humans, avatars, and robots. In The Twelfth International Conference on Learning Representations, 2024. URL 1d=42nw2G92CH
net. four
Martin I. Puterman. Markav decision processes: discrete stochastic dynamic programming. John Wiley & Sons, 2014.
RuiXi Qiao, Jie Cheng, Xingyuan Dai, Yonglin Tian, and Yisheng Lv. Offline reinforcement learning with discrete diffusion skills. arXiv preprint S.kiv:25533.0/222 76, 2025.
Ilija Radosavovic, Xiaolong Wang. Lerrel Pinto, and Jitendra Malik. State-only imitation learning for dexterous manipulation. In 202/ TEFF/RSJ International Conference on Intelligent Robots and Systems (IROS), pp. 7865-7871, 2021. doi: 10.1 10.1109/RRO5| 16..20009611111222299996693553
Dripta S. Raychaudhuri, Sujoy Paul, Jeroen Vanbaar, and Amit K. Koy-Chowdhury. Cross- domain imitation from observations. 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. 8902-8912. PMLR, 18-24 Jul 2021. URL atttss:/pprocediin
mmlr
press/ Ftees/v139/raate htt.::///yycccccuurui2222
html
Manolis Savva, Abhishek Kadian, Oleksandr Maksymets, Yili Zhao, Erik Wijmans, Bhavana Jain, Julian Straub, Jia Liu, Vladlen Koltun, Jitendra Malik, Devi Parikh, and Dhruv Batra. Habitat: A Platform for Embodied AI Research. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), 2019.
Archit Sharma, Shixiang Gu, Sergey Levine, Vikash Kumar, and Karol Hausman. Dynamics-aware unsupervised discovery of skills. In 8th International Conference on Learning Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. | /apennevieww httpe////fo.......o....o........
Wonchul Shin and Yusung Kim. Guide to control: Offline hierarchical reinforcement learning using subgoal generation for long-horizon and sparse-reward tasks. In Edith Elkind (ed.), Proceedings of the Thirty-Second International Joint Conference on Artificial Intelligence. IJCAI-23, pp.
4217 4225. International Joint Conferences on Artificial Intelligence Organization, 2023. doi: 10.249963/jcci.22223// 2023/466. URL Track.
o0t.o//010...9333eec.... 2023/469 Main
Riley Simmons-Edlerr, Ben Eisner, Eric Mitchell, Sebastian Seung, and Daniel Lee. Q-learning for continuous actions with cross-entropy guided policies, 2019.
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, second edition, 2018
Richard S Sutton, Doina Precup, and Satinder Singh. Between mdps and semi-mdps: A framework for temporal abstraction in reinforcement learning. Artificial intelligence, 112(1-2):181-21, 1999.
Andrew Szot, Alex Clegg, Eric Undersander, Erik Wijmans, Yili Zhao, John Turner, Noah Maestre, Mustafa Mukadam, Devendra Chaplot, Oleksandr Maksymets, Aaron Gokaslan, Vladimir Vondrus, Sameer Dharur, Franziska Meier, Wojciech Galuba, Angel Chang, Zsolt Kira, Vladlen Koltun, Jitendra Malik, Manolis Savva, and Dhruv Batra. Habitat 2.0: Training home assistants rearrange their habitat. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Mark Towers, Ariel Kwiatkowski, Jordan Terry, John U Balis, Gianluca De Cola, Tristan Deleu, Manuel Goulao, Andreas Kallinteris, Markus Krimmel, Arjun KG, et al. Gymnasium: A standard interface for reinforcement learning environments. arXiv preprint 2407 17032, 2024.
14

---

Published as a conference paper at ICLR 2026
Aaron van den Oord, Oriol Vinyals, and koray kavukcuoglu. Neural discrete representation Iearn- ing. In Guyon, U. Von Luxburg, S. Bengio. H. Wallach, R. Fergus, S. Vishwanathan, and R. Garnett (eds.), Advances in Neural Information Processing Systems, volume 30. Curran Asso- ciates, Inc., 2017. 2017/file/ 2a www.eeaaeeeeee000000000000 220777/f887.888111100000002e2230
httptttppppppppp.peeeeeeeeeeeeee neuri wwtp://ggg.eiiiis.cc.o..ocoooo cc/ppper fil les/
Siddarth Venkatraman, Shivesh Khaitan, Ravi Tej Akella, John M. Dolan, Jeff Schneider, and Glen Berseth. Reasoning with latent diffusion in offline reinforcement learning. In The Twelfth Inter- national Conference Learning Representations, ICLR 2024, Vienna, Austria, May 7-11, 2024. OpenReview.net, 2024. URL /ttp::/s////peeeeiee.....e ottp://wwweeeeeeeetooooraaaaaaaa
Xiaoyu Wen, Chenjia Bai, Kang Xu, Xudong Yu, Yang Zhang, Xuelong Li, and Zhen Wang. Contrastive representation for data filtering in cross-domain offline reinforcement learning. In Forty-first Internatiional Conference on Machine Learning, 2024. URL URLLAttps://ppenree net net/Eorunn Eorum/irummideoo rReWho166R
Kang Xu, Chenjia Bai, Xiaoteng Ma, Dong Wang, Bin Zhao, Zhen Wang, Xuelong Li, and Wei Li. Cross-domain policy adaptation via value-guided data filtering. In Advances in Nearal Information Processing Systems 36: Annual Conference on Neural Informa- tion Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA, December 10 16. ni is wwtppapee..iiieeeeeeeeeeeeeeeepe les pttp:/ppae//222//aaases Abst 2023, 2023a. URL -tt.r//aeeeeetet-e-----aa------- BBad8 7 www.111110011111 /ww.:/00055599944555
wwwftteoeoooeeee.eeeooeeee oat.aaao-eebttttttaoooreeeeeeeee html
Mengdi Xu, Yikang Shen, Shun Zhang, Yuchen Lu, Ding Zhao, B. Joshua Tenenbaum, and Chuang Gan. Prompting decision transformer for few-shot policy generalization. In Thirty-ninth Interna- tional Conference on Machine Learning, 2022.
Mengdi Xu, Yuchen Lu, Yikang Shen, Shun Zhang, Ding Zhao, and Chuang Gan. Hyper- decision transformer for efficient online policy adaptation. In The Eleventh International Confer- on Learning Representations, 2023b. Utthttps: http://ppeervviieeeee net forum?id
Aat W
Siyuan Zhou, Yilun Du, Shun Zhang, Mengdi Xu, Yikang Shen, Wei Xiao, Dit-Yan Yeung, and Chuang Gan. Adaptive online replanning with diffusion models. In A. Oh, T. Naumann, A. Glober- son, K. Saenko, M. Hardt, and S. Levine (eds.), Advances in Neural Information Processing Systems, volume 36, pp. 44000-44016. Curran Associates, Inc., 2023.
APPENDIX
A THEORETICAL ANALYSIS
In this section, we present the theoretical results that support the claims made in the main text. We begin by outlining the detailed problem setup, followed by the statement of the main theorem. The proof of the theorem is provided thereafter.
A.I NOTATIONS AND PROBLEM FORMULATION
Let S-1 denote the space of probability distributions over a discrete set S and S = |S]. For a conditional probability mass function p (Y (YlX), we denote the probability vector given X x as Sa (x) E A Let R, = {i E IR x 2 0}. The L of a vector v is denoted as || := max, [[v]]], where [v], refers to the i-th entry of v. The inner product between two vectors and is (u, v).
In this analysis we consider finite Markov decision process (MDP) (Puterman 2014) denoted as M = (S,A, p, 7). which consists of finite state space S, finite action space A. transition probability function p: S A A5-1, reward function r: S - [0, 1] and discount factor [0, 1). The fixed discounted factor here is motivated by our practical choice that skills represented as temporally extended actions in augmented action-space MDP with diffusion model have fixed horizon as shown in Section 4.2 making the original semi-MDP equivalent an abstraction MDP. A policy :: S A-1 maps states to distributions over actions. An agent equipped with a policy can interact with the environment by taking action a, at time step according 8,, and observing the resulting state transitions and rewards. The environment transitions to a new state 8441 according p (St, at) conditioned the current state and action. The agent receives
based the current state
15

---

Published as a conference paper at ICLR 2026
reward (st) based on the current state. Our analysis can be generalized to the reward function r: x A - [0, 1] by modifying the masked Bellman operator below. But for sake of simplicity in analytical representation, we assume the reward function is state dependent only.
The agent's goal is to maximize the expected sum of discounted rewards over time, or expected return. The objective at s is exactly the value function given by V (s) = Ex, [R|s0 s] Ex o 7'r (st) ) |so = s). where R is the return for a trajectory and y is the discount fac- tor that determines the importance of future rewards. The action-value function Q (s,a) Ex.p [R so s. au = al represents this expected return conditioned on a specific initial and a. We denote the vectors of a value function and an action-value function as V E R and Q" E RS respectively. We define the operator y for a function f: S x A + R. such that for each s, (Vf) == max f(,. The Bellman operator for is denoted as (Tf) a) a= r + (p (8,a), a), Vf). Here we also overload Vf and f to denote the correspond- ing vector obtained after applying the operators. Let n² be the optimal policy and Q" be the optimal Q function. The optimal Bellman equation for Q' can be written as Q' TQQ which lays the foundation for many reinforcement learning algorithms where p and r are unknown (Sutton & Barto 2018) to the learner.
In this work, the target test environments may exhibit altered dynamics due to unavailability of certain transitions, e.g., physical traps obstructions that block previously accessible paths. For each transition, we can define a binary function b: S x S + {0, 1}. where (s,, sy) = 1 indicates that the transition from 8; to 8j is available, and unavailable otherwise. Its matrix form is denoted as B E {0. 1xsss where bi.j == [Bl,, = b (8,, 8y). dij = 0 indicates that the transition path from s; to 8y is blocked, and any actions applied at 8; will not lead to 8,. For example, in Figure I of the main text, this corrssoonss to the case where there is one-way door from room 3 to room 2. When such transition attempts are made following the original MDP transition, the agent will collide with the obstruction and fail the episode, which is equivalent to having rewards in the future.
Therefore, we can model a target environment as an MDP MB == (SU(s)), PB. 7) with same action space and reward, but different state space and transition probability function pn: SU ( (s x
A A A For all bijj = O, we have pn (sy 8i, a) = 0 for all a E A, indicating that this transition is not possible. Since the agent fails the episode when it was about to have an unavailable transition, we can introduce sink state $1 where (s1) = 0 and pB (si|s1. a) = 1, Va E A. The probability of transitioning other states at T is O. The probability of transitioning to s_ from s; for any is PB (81|s), a) Exjes (1 bú.)) p (s51s; a). PB remains same as for all other transitions. The related values in MB are denoted as VM QMg and VMaa QiM for optimal policy.
During training on the source environment M, we cannot efficiently modify the real dynamics to either learn from samples distributed as from the target environment, or learn for many possible target environments at the same time. We aim to rely on the trajectories collected from the source environment to approximately learn policy for target environment. One way to achieve this is directly modifying the transition tuple according to B during the Bellman update by changing the target state to a sink state if the intent transition is unavailable. However, it is difficult to scale to continuous state environments as it would require comparing state similarity with a suitable state space partition. Instead, here we tackle this problem based on and a that lead the unavailable transition. Our algorithm in practice deals with partial observations of the state. The formulation here can align with this practice by considering the aggregations of states with same observation as an abstract state.
Given B, we define mask m (s;, a:) = 0 for any transition (si, 4; 85) if bis and 8, E
arg es (s|s;, and m (st, = 1 otherwise. With this mask vector m E 10, 13 SA we introduce the following masked Bellman operator
(TTQ) (s, (5) + 7 (m (8, (s, ((s))) Vm)Q),
(6)
where
max dEA:m(s,a)=1 (() 3a E E A::mms m(s =1
(yrh ())
(r(s)
otherwise
is a masked operator for a function f: S X + R.. Eq. (6) is general in that it can be developed into
a loss function based on this temporal difference for learning in continuous environment. Compared
16

---

Published as a conference paper at ICLR 2026
with modifying the transition tuple according B, T is more efficient it does not require checking the next state in the transition tuple. Moreover, during deployment stage, V q can filter out unavailable skills actions based the mask, whereas the B approach would require learning a dynamics model to predict the next state for checking. Our objective here is to examine the difference of the proposed masked Bellman operator T compared with learning under the exact B by theoretically analyzing its behavior and performance.
A.2 THEOREM AND PROOF
We restate Theorem I below which now includes the detailed formulation.
Theorem 1. Let Q
S A - R+ be the initial action-value function randomly initialized as The masked Bellman operator TTh iteratively updates Q² for k = 0, 0,
10 E 0, 1-
by Q = =T Q¹ The performance of the masked greedy policy T (s), which chooses action == arg Ammisa QK (s, a) if 3a E A :m{s,a) = 1 and randomly otherwise, close to the optimal value function VMg in the target environment after K iterations:
according to Tm (+)
V - VMnll
27
2x (1 + 7) (1 (1)=))=- + 7 -) (1 - Puiin),
7)²
+
(3)
(1 (1-
71
where Punin == min,es, P (8a(8, a)ls, a) and sa(s, a) E arg max P (s'(8,0). a).
To proof this theorem start with the following lemmas. Since the value function in MM has an additional sink state L1, we extend the definition of both 7 and ym by letting (T f) (s1, and (vmh f) (s1) == max.cA (s 1,,a) for any function f: SUS {8 x A-+R. R+.
Lemma 1. [ VQMaills 1 (1 Pmin).
Proof. Based on the definitions, for any E S,
(v (VT'Qua)) (s) (VQ (s)l
(LHS)
Qia (s. () - max QMD (s if A: m(s,a))==
max (EAAm((.a)=1
acA
(8)
(()) maxx max and
QMa a))
otherwise.
Consider the first case when 3a E A a) We analyze two subcases based on the value of m (s, (s)). Let n (s) E arg ma Q'm (s, a), and the lemma is clearly true when m (8, n" (s)) = 1 as LHS = O. In the other subcase where m (s, x- (s)) O, by letting n (s) e arg max
QMg as, we have
LHS max Q (s,a) +EA:((,a)=]1 1MM2(,,0
(9)
aEA
r(s) + 7 (pB (s, ( (s)), VQMa)) r(s) - (PB ( (s)), VQMa)
(10)
Σ PB (6'1,, (s)) max Uma (d',d)
(11)
('SU({_)
S7 Σ PB (5 16 x))))*(())) 1-
(12)
eES
?
Σ
PB (5'(s rs)))
(13)
1 1- 7
Σ
( p n (s))
(14)
WeS (eSt
1 7 (1 (1-p(a)) (sa (s))
(15)
(1 Pmii) )
(16)
In Eq. (T]) the optimal action-value function at the sink state is 0. Eq. (33 is due to the fact that we
in the case m (s, x² (s)) = O. where there exists a 8a (8, x²) such that b (8, Sa (8, n = 0.
17

---

Published as a conference paper at ICLR 2026
In the other case when Va A: m(s, a) = 0, we have
LHS [r(s) r(s) 7 (pD (s, x (6)), vQMa))
(17)
=
Σ
Pn (s'|s, (s)) max a's (5',
(18)
a'EA
('ESU(+L)
Pmin)
(19)
by following similar steps from Eq.
The inequality holds s since the optimal Q function is 0 at $1. The lemma is then proved by summarizing the above analysis under the L
The masked Bellman operator Th is contraction mapping with respect to the T, -norm, i.e., for any two vectors f and g. we have 1 f - -7 gl1x I IIf -glloo
Lemma 2.
Proof. To handle vectors in the space of I (S+11) similar to the extension of TTh also denote Pe as an extended version of by letting Pe (s |s, a) = 0 for all s E SU((1) and a E A. Based on Eq. (6) and Eq. 77, we have
TTh T g)aa
(LHS)
-I((s) + (m(s, a) (8,
- r(s) )m
f)
a)po (s, (m(s,
vm Vmmgi
(20)
=y Kmn(s, a)pa a), vm >mmg)) g)l
(21)
Σ ESS
m(s. a)p (s (s|ss,0) a) (>TH) (5) (V g) (s'))
(22)
For any 8, we can bound
1( f) (8) (s)-~ (s))
(23)
{(5,a) if 3a E A: m(s,a)
max nE.A:m(s.)== =1
max
aE.am(s.a)
(24)
(r(s) rs))
otherwise
max aEAmm(s.)=]
[f(s,a) -g(s,a)
if E A: m(,
(25)
0
otherwise
< max [f(s a) - ((s,a) aEA
(26)
Therefore,
LHS < Σ m(s, (s'ls, a) max (f (6',a') - (s',d'))
(27)
d'E
S'ES
S
max vesufaiha'ee
If (s' - (s',a'))
(28)
=> gllooc
(29)
Lemma 3. For any function f: S A R+, VM Vinll < p11s-Qiu 1l, +7(1-pman) 1-y
Proof. As stated in Theorem [ J.m. is defined the greedy policy with respect to and Jim arg Amx,, (s. a) if 3a e A m(s, a) = and randomly otherwise. We can analyze the difference in two cases based on the value of m (s, (s))) (s)). For any E S, consider the case where m (s, x² (s)) 0,
VMa (s) V. (s)
(30)
18

---

Published as a conference paper at ICLR 2026
+ (s))) (s), - (pB (s, (s. (s)), Vilam)
(31)
Kn Σ PB (s'ls, (s))
max
UMa (s',a)
(32)
a'EA
e'E(U(L))
(1 (1-Pmin)
(33)
similar the steps in Lemma
In the other case where m (8, x (s)) =
VMa (8) - VMa (s)
(34)
QlMB (s, QMa (s. + QMa (s, (s)())))) QMx (s))
(35)
SQUB (s, f(, x' (s)) + (s, Qia (s))).m(s))
+7 (PB (. (s)): Vi V.rm)
(36)
2 $2 Qinnlllll + Vin
MB
(37)
Noting that the value function at the sink state is 0, arranging the above inequality gives MB ViMBl llce under the norm.
2
The lemma is proved by summarizing the above analysis
1-7
Now we are ready to prove Theorem
Proof of Theorem well-defined on si. Thus, by Lemma 2
4 Similar to the extension of the operators T and Vmm, also extend Q* be
| k - Q0velaa
1 Q K- QMell
(38)
@" -1 - -T QimB + +T TTMMu Qmallx
(39)
W r Q kk -Qim 1loo + | -Qme 1lx
(40)
1 2² 12 xQ=22 Quallo + +(1 ?) T 1mQmm Quello
(41)
(42)
K
S 5n Qut Qia 1lo0 +
Σ Th
QiM llx
(43)
K
The first term in Eq. term,
due to the choice of initialization. To bound the second 43) is bounded by 1-
TTM Qin - Pianl s.a)
(LHS)
= |r(s) + 7 (m(s, (s a (01e)-r(() - (p) VQMa))
(44)
= (m(s, a) (s, a), 1m QMn) - (PB(s, VQ'un)l
(45)
If (mmssaa) a)
0,
LHS = | (pe(s, a), VQQMa))
(46)
=7
Σ
PB (s ss, a) max Ou,, (s', a')
(47)
a'EAA
s'Es
Σ
PB (5'1s,))
(48)
?
< 1 (1 Pmii).
(49)
by following similar steps from Eq. 13 In the other case, let Ss(s) = {s' (b,,s'))))) = 0}. and
m(s,a) a) implies that as(s,)) & (s)s). Thus, by Lemma
LHS
Σ
Pe (5'1s, ()Tm Qima)
Σ
pB ( (s)a) (VQina)
(50)
(ES
19

---

Published as a conference paper at ICLR 2026
Σ ('E))
p (VTQum)) (6)
Σ ( (VQnn) (8)
d'E8),
Σ ((
E PB (s'ls, a) (VQ's) (5')
+7
d'SL(()
d'ES((a)
v Ipo (s (um) ($1) -BL) Pp (VQiee) (s1)
(51)
SY Σ (s
1-7 (1- Pmmia) (1-Pmiu) +y
Σ p())ss,)) 1
(52)
1-7
(
#'ES((a)
x2 - (1-Pmin) + 7 (1 Pmin)
(53)
1-7
1 1-7
7 (1 + +22)
- -PPmin).
(54)
1-7
Therefore, by Lemma 3
21|Q" Qume 10 +7 -Pmin)
Iu -VMB
(55)
1-7
2
nk
K-1
k (1++)) Σ
(1 (1-Pmin) + 2 (1 Pmiin
+
(56)
7
1 1-7
1 1-7
k=0
27 k
27 (1 + (1 ))) +7 71)
(1 Pmmin)
(57)
(1 7)²
(1 - 7)
which concludes the proof.
B IMPLEMENTATION DETAILS
In this section we provide details on the environments, datasets and computational resources used in our experiments.
We the Maze2D, Fetch and ReplicaCAD domains from the Gymnasium Robotics suite (Towers) let al. 2024) and Habitat-Labb (Savva et af.] 2019 Szot et al.) 2021 Puig et al., 2024). In Maze2D, a point-mass actuated in the x-y plane to navigate a goal location. Observations consist of the agent's 's (x, y) position while the full state includes velocity as well. Reward is sparse and r(s,a) 1.0 when the agent is within 0.5-unit radius of the goal, and 0 otherwise. The maximal episode length is 350 steps (trajectory is truncated at 350 if goal is not reached) in our modified map. The map is modified to have three paths and in each test environment only one path feasible for the agent to reach the goal. We follow the D4RL setup (Fu et al. (020) to collect trajectories of agents reaching randomly sampled goals. Fetch, we use FetchReach- v4 from Gymnasium, a 7-DoF Fetch manipulator tasked with moving its gripper to goal position in 3D space. The reward is sparse, r = -1 at each timestep until the goal is achieved (within tolerance), upon which r = 0. The maximal episode length is 50 steps. In the test environment, we add constraint region to the workspace to limit the agent's 's reachability. The agent will incur a loss of -50 if it moves into the constraint region. The constraint region is box (prism) defined by its center [1.4,0.85,0.5] and scene
half-size [0.1,0.15,00 5,0099 for each edge. In Habitat-Lab, we use the w33.sc_sttaiinn. from ReplicaCAD (Szot al. 2021) create visual navigation task. The agent is tasked with navigating to goal location with basic navigation actions including move forward, turn left, and turn right. The observation is 44x64 RGB image from the agent's first-person view. The maximal episode length is 500 steps. During testing, the agent is required to follow trajectory not directly demonstrated in the training dataset with visual observations only.
For offline training in Maze2D and Fetch, we generate static dataset of 200,000 interactions by repeatedly sampling a uniformly random goal and recording the full trajectory until the next goal is reached (with no automatic resets upon success). In Maze2D, the process follows the D4RL data collection protocol (Fu et al. 2020) by letting the agent navigate under a simple controller (open-loop waypoints) to the goal. In FetchReach, we pretrain goal-conditioned actor-critic model and use this
20

---

Published as a conference paper at ICLR 2026
policy to collect data. In ReplicaCAD, we collect 5,000 interactions similar to Maze2D using the shortest-path planner provided by Habitat-Lab, All experiments were conducted on a workstation featuring an NVIDIA GeForce RTX 3090 graphics card with 24 GB of GDDR6X VRAM, alongside a 32-core CPU and 128 GB of system RAM. Training the skill autoencoder and the diffusion model each took approximately 6 hours to converge, while the MSTT algorithm required about 8 hours to reach full convergence. In our implementation of MSTT, we use a codebook of size 240 to represent the latent skill space. The length of all skills (horizon of the skills from dataset) for the Maze2D environment is 32, and 8 for the FetchReach and ReplicaCAD environments. The diffusion model is trained conditioned on a start state using 16 denoising steps for the first two environments and 256 steps for the latter. During testing, the agent conducts closed-loop control by executing the skill with their corresponding horizon steps before decision on the next skill.
C ADDITIONAL ANALYSIS AND RESULTS
C.I VALIDATION OF MASKED BELLMAN UPDATE
In this section, we empirically validate the masked Bellman update and the results brought by Theo- rem by comparing the -nrm gaps between dif- ferent estimated and true evaluated value function. We simulate the masked Bellman update in the un- blocked source MDP using all possible feasibility masks. There are 8 abstracted skills, so the total 255.
Prompt
Observation:
You art gives map (Fig. 1) for a comtinou 2D maxe. The sroen ball the shet locatio and the red ball is the goel. Thei coondaans ar [1.0. 1.0] and (400 4.0).
e
Instruction:
You sre roquired to rosse pub from the start as the asell and cuppi all state coondmares ncoded to complette daz uak in I.
Example:
number of possible target MDPs is 2 We evaluate the performance of the masked Bellman update in all these target MDPs. We keep the Pmin at 0.95 for all MDPs. The resulting value function VK is computed through K iterations of the masked Bellman operator.
tininng ervinommnt shows 2, suberjectery poist) foo.
the blue und cnding
following coostinates:
[5.0803, 199299
[5,0253, 19965]
[5,0016, 1.9434]
[5.01511
The unvinomente hove the ume size hut diffirone wall configuntioss Note specing
eagectory should consiseen the spacing evanpte.
Figure 8(a) shows the foc -nom gap between Vx and the optimal values. We can see that the gap is small for most target MDPs. We also find that when Pmin is very close to 1 (e.g. 0.999), all the gaps are close to 0 (< 0.01). Figure [(b) shows the -nom gap between the true evaluation of the policy derived from VK and the optimal values. We can see that the gap is also small for most target MDPs.
Goneration
Vislon Lang language Model
is colitiion flee, piecewise linesr tnjoetorry from
occoss top coridor the [1.0. 10 up lett path length units, and brokze goal [4.0, 4.0)
201 point.
[1,0000, 1.0000)
move staight 11.6, 1.0 11.0, 4.0j Poinss 101-2001 move strighe right L1.., 4.0.)) [4.1, 4.0 Sepp monchas youressmple sparng
Figure 9: Diagram of generating observations with VLM.
To validate the proposed training method MSTT on datasets, we collect a dataset of 1000 interactions in the unblocked source MDP using random policy. We then train MSTT on this dataset using randomly sampled masks and discrete state/action features instead of skill representation learning and diffusion model as they readily available. Figure 8(c) shows the lo -orm gap between the true evaluation of the MSTT policy and the optimal values. We see that the feasibility conditioned value function learned by MSTT
induce policy close to
the optimal values on most target tasks. In Figure 8(b) [ there are some target MDPs where the gap is non-zero. This is because the induced policy is derived from the learned value function by breaking ties randomly (due to numerical errors) when the actions' values are same. This action deviates from the optimal action and can lead to a different value at some states.
C.2 VISUALIZATION OF SKILL SPACES
10 visualizes the entire set of skills decoded from the pre-trained VQ-VAE in Maze2D, revealing a diverse collection of discrete motion primitives that cover different directions and local maneuvers. The plots illustrate that the VQ-VAE discovers structured, reusable skill space. Fig- ures [I and environment masks for the three target environments. Target environment I has 21 valid skills under the inferred mask:
Figure
and [3 show the subsets of those skills that remain available after applying the inferred
12, 19, 20, 65, 66, 67, 97, 105, 113, 144, 145, 192, 193, 201, 202, 203, 204, 205, 214, 230.
21

---

Published as a conference paper at ICLR 2026
2.0
1.5
1.0
0.5
0.0
20
40
60
80
100
140 Test tasks
120
160
180
200
220
240
255
(a) Estimated values v.s. optimal values
1.0
0.8
0.6
00.4
0.2
0.0
1
20
40
60
80
100
120
140
160
180
200
220
240
255
Test tasks
(b) True value evaluations v.s. optimal values
1.0
0.8
VV V
0.6
0.4
0.2
0.0
20
40
60
80
100
120
140
160
180
200
220
240
255
Test tasks
(c) True values of MSTT v.s. optimal values
Figure 8: Comparison of the -orm gaps between different estimated and evaluated value functions. (a) The -nom gap between estimated values after convergence of the masked Bellman update compared with the optimal values. (b) The Loo -nom gap between the evaluations of the converged policy and the optimal values. (c) The loc noom gap between the true evaluation values of MSTT policy and the optimal values.
22

---

Published as a conference paper at ICLR 2026
Figure 10: All skills decoded from the pre-trained VQ- VAE in Maze2D.
Figure 11: Skills that are available under the inferred mask environment target 1,
HH H
Figure 12: Skills that are available under the inferred mask in environment target 2.
23

---

Published as a conference paper at ICLR 2026
Figure 13:
Skills that are available under the inferred mask environment target
Table 5: Performance on Maze2D with different codebook sizes.
Env T1
Env T2
Env T3
Codebook Size
Return (f) Steps ()) Goal (t) Return (f)
Steps (L) Goal (1)
Return (f)
Steps (4)
Goal (t)
15
204.57
146.43
100%
0.0
350.0
0%
84.32
266.28
60%
64
120.78
230.10
88%
23.06
327.24
30%
0.82
349.19
1%
240
163.97
186.92
89%
145.51
211.01
95%
111.59
249.33
99%
512
191.04
159.90
94%
135.39
215.51
90%
124.19
226.80
99%
1000
206.36
144.64
100%
168.31
182.66
97%
92.04
258.82
86%
Target enviromment 2 has 34 valid skills under the inferred mask:
15, 23, 31, 39, 53, 54, 74, 94, 95, 102, 103, 117, 120, 121, 122, 123, 124, 125, 126, 134, 151, 160,
161, 162, 163, 164, 165, 168, 172, 174, 199, 207, 222, 230.
Target environment 3 has 31 valid skills under the inferred mask: 15, 23, 28, 31, 36, 37, 39, 40, 41, 42, 43, 44, 45, 46, 47, 54, 85, 88, 92, 93, 102, 103, 141, 151, 184, 199, 207, 232, 233, 234, 235.
The mask inference selects compact, task-relevant subsets that preserve sufficient coverage for navigation and goal-directed behavior.
C.3 ANALYSIS ON THE CODEBOOK SIZE
We conduct an ablation study on the codebook size of the VQ-VAE Maze2D. We train the VQ- VQ-VAE with different codebook sizes and train MSTT with the same hyperparameters the resulting skill spaces. The results are shown Table 5 The codebook size determines the granularity of the skills learned by the VQ-VAE. We see that smaller codebook size (e.g. 15) leads to worse performance, which is expected as may not be able to capture sufficient skills for solving the task. Performance generally improves with larger codebook sizes, and the benefits plateau after certain points, which may be due to the increased difficulty in learning value function over a larger skill space.
C.4 STUDY ON TARGET DEMONSTRATION
Figure 14 shows suboptimal demonstration that detours around the maze reach the goal. The right panel shows the MSTT rollout ing the mask inferred from that demonstration: MSTT reaches the goal with short path while adhering to the inferred skill subset, indicat- ing the effectiveness of the learned value func- tion. Figure 15 shows the MSTT rollout using mask inferred from incomplete demonstra- tion. The inferred mask supports safe local be- havior within the demonstrated region, produc- ing only partial progress following the incom- plete demonstration.
Figure 14: Left: subopti- mal demonstration that detours around the maze to reach the goal. Right: rollout of MSTT using the inferred mask from the suboptimal demonstration.
Figure Roll- out of MSTT us- ing the inferred mask from an in- complete demon- stration.
C.5 EXAMPLE
PROMPT AND OUTPUT FROM VL.M
In this section, we provide the example prompt and output from the VLM reasoning for
24

---

Published as a conference paper at ICLR 2026
Maze2D task and FetchReach task. We prompt GPT-o4-minii-oe with images of the enviromment and an example trajectory to generate state coordinate sequences. The example diagram is shown in Figure When prompting the VLM, we give many views of the task environment that the VLM can understand the spatial relationship between the robot and different objects in the environments. We also provide example trajectory that is used to generate the trajectory shown in the example image. This can help the VLM to understand the number of steps required for similar tasks in the environment. From Table[and) the given tasks. In future work, we will explore directly generating image frames from a generated model for visual demonstration following.
can that the VLM can generate satisfactory trajectories for
25

---

Published as a conference paper at ICLR 2026
Table 6: Prompts and Responses from spt-o4-mini-high for Demonstration Generation (Maze2D).
Prompt
The first picture you is map for continuous 2-dimensional maze viewed from 3d renderer. The green ball the start location where you are right now. The red ball the goal location that you need reach, The coordinates green ball is [1.0, 1.0] and red ball is [4.0, 4.0J. You are required to reason path from start to goal and output the all state coordinates.
Example:
source training environment as shown picture in the top down angle, the state coordinates of subtrajectory with horizon 32 (starting from blue to red)
5,0303545, 1.9299148
5.025338, 1.9365628
5.02067, 1.9434459
5.0151925, 1.9526943
5,009231, 1,9643023
5,003044, 1.9782642
4.9968333, 1.9945747
4.990752, 2.01
4.9849145, 2.0342183
4,9794, 2,0575402
4.974263, 2,083188
4. 4,9695354, 2.111157
4.9652305, 2.1414409
4,9613495, 2. 1.4703339
4.9578824, 2,2089312
4.954813, 2.246127
4.9521 184, 2.2856157
4.949773, 2.327392
4.9477477, 2.3714507
4.9460144, 2.417786
4.944544, 2.4650137
4.9433084, 2.5117943
4.9422803, 2.5571213
4.941435, 2.6002645
4,9407473, 2.6409237
4.940197, 2.6791043
4.9397635, 2.7148123
4,939429, 2.7480536
78, 2.778834
4.938996, 2,8071597
4.938871, 2,833036
4.938792, 2.8564692
The two environments have same map size but diffcrent wall configurations. Please output full trajectory of continuous state coordinates to complete the task in the first picture. Note that the spacing between consecutive states in your trajectory should be consistent with the spacing in the given example.
Response
Below is collision-free piecewise-linee trajectory from the start [1.0,1.0 0] up the left corridor and then across the top corridor the goal [4.0,4.0]. The total path length 6.0 units, we's broken it into 201 points (so each step 6s0.03 units, similar to your example)
1. 1.0000, 1.0000
1. 1.0000, 1.0300
1.0000, 1.0600
1.0000, 1.0900
1. 1.0000, 1.1200
1.0000, 1.1500
26

---

Published as a conference paper at ICLR 2026
Table continued from previous page
1,0000, 1.1800
1.0000, 1.2100
1.0000, 1.2400
1.0000, 1.2700
1.0000, 1.3000
1.0000, 1.3300
1.0000, 1.3600
1. 1.0000, 1.3900
1 1,0000, 1,4200
1. 1.0000, 1.4500
1. 1.0000, 1.4800
.0000,, 1.5100
1. 1,0000, 1.5400
1 1,0000, 1.5700
1. 1.0000, 1.6000
1,0000, 1.6300
1. .0000,, 1.6600
1. 1,000,, 1.6900
1. 1.000,, 1.7200
1.0000, 1.7500
1,0000, 1.7800
1. 1,000,, 1.8100
1. 1,0000, 1.8400
1. 1,0000, 1.8700
1,0000, 1,9000
1,000,, 1.9300
1. L,000,, 1.9600
1. 1.0000, 1.9900
1. 1,0000, 2,0200
1,0000, 2.0500
1. 1.0000, 2.0800
1. 1.0000, 2.1 100
1. 1.0000, 2.1400
1,0000, 2.1700
1 1,0000, 2.2000
1. 1.0000, 2.2300
1. 1,0000, 2.2600
1. 1.0000, 2.2900
1,0000, 2.3200
1. 1,0000, 2.3500
1. 1.0000, 2.3800
1. 1.0000, 2,4100
1. 1.000,, 2.4400
1. 1,0000, 2.4700
1. 1,0000, 2.5000
1. 1.0000 2.5300
1. 1.0000, 2.5600
1.0000, 2.5900
1. 1,0000, 2.6200
1. 1,0000, 2.6500
1. 1.000,, 2.6800
1. 1.0000, 2.7100
1. 1,0000, 2.7400
1. 1.0000, 2.7700
1. 1,0000, 2,8000
1.0000, 2.8300
1. 1.0000, 2.8600
1.0000, 2,8900
1.0000, 2.9200
27

---

Published as a conference paper at ICLR 2026
Table continued from previous page
1.0000, 2.9500
1.0000, 2.9800
1.0000, 3.0100
1. 1.0000, 3.0400
1.0000, 3.0700
1.0000, 3.1000
1.0000, 3.1300
1. 1,0000, 3.1600
1 1,000,, 3,1900
1. 1,0000, 3.2200
1. 1,0000, 3.2500
10000, 3.2800
1. 1,000,, 3.3100
1 1,000, 3.3400
1. 1,0000, 3.3700
1,0000, 3.4000
1. .0000, 3.4300
1. 1,0000, 3.4600
1. 1.0000, 3.4900
1. 1.0000, 3.5200
1,0000, 3.5500
1. 1,000, 3.5800
1 1,0000, 3.6100
1. 1,0000, 3,6400
1,0000, 3,6700
1. 1,0000, 3.7000
1. 1,000,, 3.7300
1. 1.0000, 3.7600
1. 1,0000, 3,7900
1,0000, 3,8200
1. 1,000,, 3.8500
1. 1,0000, 3.8800
1. 10000, 3.9100
1.0000, 3,9400
1. 1,000,, 3.9700
1. 1.0000, 4.0000
1 1,0300, 4.0000
1. 1.0000, 4.0000
1.0900, 4,0000
1,1200, 4.0000
1.1500, 4.0000
1. 1.1800, 4.0000
1,2100, 4.0000
1,2400, 4.0000
1,2700, 4.0000
1.3000, 4,0000
1.3300, 4,0000
1,3600, 4,0000
1,3900, 4,0000
1, 1,4200, 4,0000
1,4500, 4,0000
1,4800, 4.0000
1.5100, 4.0000
1,5400, 4,0000
1.5700, 4,0000
1.6000, 4.0000
1.6300, 4.0000
28

---

Published as a conference paper at ICLR 2026
Table 6 continued from previous page
1.6600, 4.0000
1.6900, 4.0000
1.7200, 4.0000
1.7500, 4.0000
1.7800, 4.0000
1.8100, 4.0000
1.8400, 4.0000
1,8700, 4.0000
1. 1,9000, 4.0000
1,9300, 4.0000
1. 1,9600, 4.0000
1,9900, 4.0000
2,0200, 4_0000
2,0500, 4,0000
2,0800, 4.0000
2.1100, 4.0000
2.1400, 4.0000
2.1700, 4.0000
2,2000, 4.0000
2,2300, 4.0000
2,2600, 4.0000
2,2900, 4.0000
2,3200, 4,0000
2,3500, 4.0000
2,3800, 4.0000
2.4100, 4.0000
2.4400, 4.0000
2.4700, 4.0000
2,5000, 4.0000
2,5300, 4.0000
2.5600, 4.0000
2,5900, 4.0000
2,6200, 4.0000
2.6500, 4.0000
2,6800, 4.0000
2.7100, 4.0000
2.7400, 4.0000
2,7700, 4.0000
2,8000, 4.0000
2,8300, 4.0000
2,8600, 4,0000
2,8900, 4.0000
2,9200, 4.0000
2,9500, 4.0000
2,9800, 4.0000
3,0100, 4.0000
3,0400, 4.0000
3,0700, 4.0000
3,1000, 4.0000
3,1300, 4,0000
3,1600. 4.0000
3,1900, 4.0000
3,2200, 4.0000
3,2500, 4.0000
3,2800, 4.0000
3,3100, 4.0000
3,3400, 4.0000
3,3700, 4.0000
29

---

Published as a conference paper at ICLR 2026
Table 6 continued from previous page
3,4000, 4.0000
3,4300, 4.0000
3,4600, 4.0000
3,4900, 4.0000
3,5200, 4.0000
3,5500, 4.0000
3,5800, 4.0000
3,6100, 4.0000
3 3,6400, 4,0000
3,6700, 4.0000
3,7000, 4.0000
3.7300, 4.0000
3.7600, 4.0000
3,7900, 4,0000
3,8200, 4.0000
3,8500, 4.0000
3,8800, 4.0000
3,9100, 4.0000
3 3,9400, 4.0000
3,9700, 4.0000
4,0000, 4,0000
Points 1-101 straight up from [1.1] [1,4].
Points 102-201 move straight right from [1.4] to [4.4].
Step size za0.03 units matches your example spacing.
30

---

Published as a conference paper at ICLR 2026
Table 7: Prompts and Responses from spt-o4-mini-hige for Demonstration Generation (Fetch).
Prompt
You can see details of a test environment for a continuous FetchReach robot arm from Front, Left and Top views as shown in 1st, 2nd, 3rd pictures as Test-F.png Test-L.png and Test-T.png. The end effector of the robot is the gripper located at the start location. The red ball the goal location that the gripper needs reach. The coordinates of the gripper [1.5, 0.45, 0.45] and red ball are meters. There is a red rectangular prism represents the region that the gripper should always avoid when reaching to the goal. Example:
[1.1, 1.1, 0.6]. the units here and below
source training env shown picture 4, 5, Front, Left and Top views, the state coordinates of the subtrajectory that available the dataset with horizon 24 (starting red ball to gripper) is
[1 11.0914499 0.66841 225, 0.44147953]
[1.1227403. 0.66162956, 0.44578195)
[1.1558249 0.66254006 0.44903293].
[1.1889781, 0..71188605, 0.4525502 J.
[1.2226527 0.74391514, 0,4652187 ),
[1.2558049 00.7763036 0.45177814),
[1.290576. 0.80899537, 0.47786778),
[1.3192018 0.8403678 0.44858858),
[1.2866951
0.8561721, 0.45177984),
[1.2538819
0,88425666, 0.44807455).
[1.2229187 0.91265494, 0.45137948).
[1.2152855 0.8794606. 0.45907718).
[1.2098249 0,8462405 0.45971522].
[1.1989532, 0.81283104, 0.4467045 ),
11.2004328, 0.779887366, 0.45744836),
[1.2142384 0.7472065 0.47155502),
[1.2419714 0.71475255, 0.48558167),
[1.2744763 0.68229145, 0.49823788].
[1.3076947 0.64975435, 0.51746166).
[1.3411665 0,,61711025, 0.5364092
[1.3745239 0.5844047 0.5114905
[1.408034, 0,5516487, 0,48911682).
[1.4419813, 0.5182403, 0.44622924],
[1.472218 0,4902828, 0.42677632).
The two environments have same size but differ the obstacle constraints. The example trajectory may violate test time constraint.
You are required to reason a path from the gripper start location to the goal location and output all the state coordinates. You should choose the safest path away from the constraint far possible since the actions may have enrors in local range. The available range of the gripper is [1.0, 1.5] for axis, [0.4, 1.1] for y axis and [0.4, 0.6] for the axis. Based on the top view, the x axis left to right, the axis is from bottom to top, and the axis is perpendicular to the view. Please output a full trajectory of continuous state coordinates to complete the task in the test environment shown in different angles in the first four pictures. Note that the planned trajectory does need to have same length as the example trajectory but ensure that the spacing of the states in the trajectory should be similar as the given example. Keep in mind that choose the safest path away from the constraint far as possible.
31

---

Published as a conference paper at ICLR 2026
Table continued from previous page
Response
Below one possible safe, collision-free trajectory that drives the gripper from its start at [1.50, 0,45, 0.45] out around the left side of the red obstacle, then up to the goal at [1.10, 1.10, 0.60]. We keep all I 1.2 whenever 0.6 I S 0.85 to stay well clear of the step-lengths (==0.05 m) par with
prism, and we roughly 5 cm increments example's spacing.
[1,500, 0,450, 0,450]
[1.450, 0,450, 0.450]
[1.400, 0.450, 0.450]
[1.350, 0.450, 0.450]
[1.300, 0.450, 0.450]
[1.250, 0.450, 0.450]
[1.200. 0,450, 0.450]
[1.150. 0.450, 0.450]
[1.100. 0.450, 0.450]
[1.050, 0.450, 0.450)
[1.000, 0.450, 0.450)
[1.000, 0.500, 0.450)
[1.000, 0.550, 0.450]
[1.000, 0.600, 0.450]
[1.000, 0.650, 0.450]
[1.000, 0.700, 0.450]
[1.000, 0.750, 0.450]
[1.000, 0,800, 0.450]
[1.000, 0,850, 0.450]
[1.000, 0.900, 0.450)
[1.000, 0.950, 0.450]
[1.000, 1.000, 0.450)
[1.050. 1.000. 0.450]
[1.100, 1.000, 0.450]
[1.100. 1.050. 0.450]
[1.100, 1.100, 0.450]
[1.100, 1.100, 0.500]
[1,100, 1.100, 0.550)
[1.100, 1.100, 0.600]
Explanation the segments:
1. Pull back along x from 1.50 - 1.00 at y=0.45, z=0.45, staying below the obstacle's 's y-range.
Climb up in from 0.45 1.00 x=1.00, 2=0.45, keeping the left of the obstacle (all x.0.)
Move in and to the target column/row: two steps in x (1.00~1.100) and two steps in y (1.00-+1.10) at z=0.455.
4. Raise in z from 0.45++0.6 over three steps meet the ball height.
This keeps comfortable margin from the red prism (which lies roughly between x 1.45, y=.0.6-0.85 2=.0.400555 and uses roughly em per step, comparable to the example.
D USE OF LARGE LANGUAGE MODELS
We did not use LLMs for research ideation, methodology development, analysis or writing. The only involvement of language model in this paper the usage of -pt-o--min generate coordinates sequences from images for case study, All prompts and responses are included in Section C.5 of the appendix.
32


Reviews:
[{"review_id": "2AiAaIOABu", "reviewer": "Reviewer_Z4mu", "summary": "This paper proposes MSTT, an offline hierarchical learning method for off-dynamics transfer. It treats the dynamics discrepancy as blocking some learnt skills and thus masking those out  at deployment. The target skill mask is inferred based on single observation-only demonstration which makes it suitable for domains where action annotations or environment interactions are costly. Empirical results show MSTT shows higher return and less complete steps than DARC and BC in Maze2d transfer.", "questions": "1. Related to W1, is it possible that dynamics discrepancy leads to new skills that have not been learnt with source domain dataset, in which case the application of MSTT may lead to suboptimal performance?\n\n2. How does the codebook size affect the downstream performance?", "limitations": "", "rating": 6, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "1. The method proposed is novel to me. Modeling dynamics discrepancy as masked out skill tokens alleviate the problem presented in specific transfer learning problems.\n\n2. Proposed MSTT shows better performance in the demonstration-only offline transfer learning setting over some imitation and offline learning baselines.", "weaknesses": "1. The assumption that source and target domains only differ in feasibility distinction of specific skills is strong to me. This limits the generalizability of the method proposed, and raises my concern to the setting proposed that target dataset should have similar levels of accessibility compared to source dataset.\n\n2. Weak baseline comparison. For cross-domain imitation learning, [1] suggests a method which also learns from cross domain state-only demonstrations, which explicitly deals with domain discrepancy upon ID methods. Moreover, as a minor concern, DARA is not a recent offline transfer method though it doesn't share the same setting with MSTT. I would suggest replace it with some stronger baselines such as [2].\n\n[1] Cross-domain Imitation from Observations Dripta S. Raychaudhuri, Sujoy Paul, Jeroen Vanbaar, Amit K. Roy-Chowdhury Proceedings of the 38th International Conference on Machine Learning, PMLR 139:8902-8912, 2021.\n\n[2] Contrastive Representation for Data Filtering in Cross-Domain Offline Reinforcement Learning Xiaoyu Wen, Chenjia Bai, Kang Xu, Xudong Yu, Yang Zhang, Xuelong Li, Zhen Wang Proceedings of the 41st International Conference on Machine Learning, PMLR 235:52720-52743, 2024.", "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": "aZxVMI1RU5", "reviewer": "Reviewer_J5hB", "summary": "The paper proposes Masked Skill Token Training (MSTT), an offline hierarchical reinforcement learning framework for off-dynamics transfer using only observation-only demonstrations. MSTT learns a discrete skill space via VQ-VAE and trains a feasibility-conditioned critic through masked Bellman updates. At deployment, it infers a binary skill mask from a single demonstration and uses a diffusion-based trajectory generator to execute only feasible skills. Experiments on Maze2D, FetchReach, and Habitat environments show that MSTT achieves robust off-dynamics transfer.", "questions": "1. In the VLM + MSTT case study, does the VLM function as a world model that generates visual observation trajectories to serve as demonstration inputs for the policy transfer?", "limitations": "", "rating": 4, "confidence": 2, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "This paper addresses the important problem of transferring policies using only observation-based demonstrations. By introducing a simple masking mechanism over discretized skills, the method enables efficient offline adaptation. The proposed framework demonstrates strong generalization and robustness for off-dynamics transfer, and further presents experiments combining MSTT with VLMs, showing its potential for low-effort policy adaptation.", "weaknesses": "1. The masking mechanism relies on random sampling during training, which seems inefficient. Could the masking process be learned? Moreover, MSTT's scalability is limited as the skill vocabulary grows, might cause the critic to overfit to partial masking patterns and weaken generalization.\n\n2. The proposed method mainly addresses structural dynamics shifts (e.g., blocked transitions) and does not handle continuous or parametric changes in dynamics (e.g., friction, mass, damping). This narrows its applicability.\n\n3. Visualizing which discretized skills become masked in each target domain could provide strong qualitative and interpretable evidence of proper adaptation. By examining how feasibility patterns change across domains, the authors could further demonstrate the effectiveness of the proposed method.\n\n4. The baselines, such as BC and Diffuser, are not specifically designed for off-dynamics transfer, limiting the fairness of comparison. It would strengthen the evaluation to include methods such as SiMPL[1], PDT[2], MetaDiffuser[3], HDP[4], which explicitly target policy generalization or meta-adaptation across dynamics shifts. Furthermore, the hierarchical skill-based approaches are missing.\n    \n    [1] Skill-based Meta-Reinforcement Learning, ICLR 2022\n    \n    [2] Prompting Decision Transformer for Few-Shot Policy Generalization, ICML 2022 \n\n    [3] MetaDiffuser: Diffusion Model as Conditional Planner for Offline Meta-RL, ICML 2023\n\n    [4] Hyper-Decision Transformer for Efficient Online Policy Adaptation, ICLR 2023", "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": "42Lr60i8iH", "reviewer": "Reviewer_DAdX", "summary": "This paper introduces Masked Skill Token Training (MSTT), a hierarchical reinforcement learning framework for off-dynamics policy transfer in a fully offline setting. The core idea is to adapt an agent trained in a source environment to a target environment with new constraints (like a blocked passage) using only a single, observation-only demonstration for adaptation.\n\nMSTT works in three phases:\n1.  **Skill Tokenization:** It uses a VQ-VAE to learn a discrete codebook of \"skills\" from fixed-length, unsupervised trajectory segments.\n2.  **Masked Critic Training:** It trains a \"feasibility-conditioned\" $Q$-function. This is the clever part: it uses a novel masked Bellman operator and randomly samples binary skill masks during training to simulate a wide variety of potential dynamics shifts.\n3.  **Adaptation & Execution:** At test time, it infers a single, fixed feasibility mask by tokenizing the provided demonstration. A diffusion-based generative policy then proposes candidate trajectories, which are filtered by this mask before being executed.\n\nThe authors provide a theoretical analysis for their masked Bellman operator and demonstrate strong empirical performance on several control and navigation tasks.", "questions": "1.  Given that the theory and implementation rely on fixed-horizon skills $H$, would the authors be willing to clarify this in the main text (e.g., by framing it as an \"augmented action-space MDP\") to avoid confusion with the more general SMDP \"options\" framework? \n2.  How sensitive is the mask inference to the quality of the single demonstration? What happens if the demo is suboptimal (e.g., takes a long but feasible path) or noisy (e.g., contains a skill that is, in fact, infeasible)? \n3.  Can you report the test-time sampling efficiency? Specifically, what is the average rejection rate (or number of samples) in Algorithm 2 for the different target environments? One might expect the highly-constrained Maze2D tasks to have a high rejection rate.", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "1.  **Novel and Practical Problem:** The paper tackles a well-motivated and important problem: adapting offline-trained agents to new environmental constraints without costly interaction or action labels.\n2.  **Elegant Core Idea:** The \"masked Bellman operator\" is the paper's strongest asset. The insight to train a single critic conditioned on a binary feasibility mask $m$, and to generalize this critic by randomly sampling $m$ during training, is a new idea.\n3.  **Sound Theoretical Core:** The method is justified by a novel theoretical analysis. The core contribution is Lemma 1 (we have checked the correctness), which bounds the approximation error of the masking approach and links it directly to skill determinism ($p_{min} \\to 1$). While the associated proofs for Lemma 2 and Lemma 3 are based on standard Q-learning convergence analyses, they are essential for the completeness of the main theorem, which is both novel and well-motivated.\n4.  **Dominant Empirical Success:** The framework is not just a theory. It achieves near-expert performance and massively outperforms strong baselines across all tasks. The FetchReach result, where MSTT is the *only* method to both succeed and avoid the new constraint (0% failure rate), is particularly striking.", "weaknesses": "1.  **Misleading Presentation of \"Options\":** The paper's framing around the \"options\" framework is a bit oversell. The method's theory and implementation do not address true temporal abstraction (i.e., SMDPs) but are instead restricted to a fixed-horizon, augmented action-space MDP. While the proof for the masked MDP's bounded convergence is novel, this limitation of the theoretical contribution to a standard MDP, rather than a more general SMDP, constrains the work's theoretical novelty. This critical simplifying assumption—the fixed skill horizon—might be better stated clearly upfront in the main text, not relegated to a footnote or appendix. We would encourage the authors to adopt the more rigorous and accurate term \"augmented action-space MDP\" throughout the paper; doing so would significantly improve the clarity and precision of the contribution.\n\n2.  **Simplistic Mask Inference:** The test-time mask inference is a simple \"set $m(z)=1$ if $z$ is in the demo.\" This assumes the single demonstration is both correct and comprehensive. The paper seems doesn't explore how robust this is to suboptimal or noisy demonstrations. And a soft version seems to be more plausible than hard assignment for real world application.\n\n3.  **Potential Inefficiency:** The pipeline of a VQ-VAE encoder, a diffusion model policy and the complexity of sampling mechanism seems limited to scale to large scale applications. The test-time adaptation (Algorithm 2) is a rejection sampler. In highly constrained target environments (where most $m(z)=0$), this could be extremely inefficient, requiring many samples from the diffusion model to find one valid option. So the scalability feels a bit limited.\n\n4. A closely related work with potential theoretical improvement is if the masked value theorem applied to HiT-MDP then the claim of options will hold. Also skill discovery algo will be simplified and generalized. Authors may refer to \n\nLi, Chang, Dongjin Song, and Dacheng Tao. \"The skill-action architecture: Learning abstract action embeddings for reinforcement learning.\" (2021).\n\nLi, Chang, Dongjin Song, and Dacheng Tao. \"Hit-MDP: learning the SMDP option framework on MDPs with hidden temporal embeddings.\" The Eleventh International Conference on Learning Representations. 2023.\n\nfor an efficient and scalable option embedding discovery framework. An offline algo vmoc is derived here:\n\nLi, Chang, et al. \"Learning Temporal Abstractions via Variational Homomorphisms in Option-Induced Abstract MDPs.\" arXiv preprint arXiv:2507.16473 (2025).", "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": "ZX94ufwqDk", "reviewer": "Reviewer_VHqp", "summary": "This paper introduces MSTT, a hierarchical RL framework for off-dynamics policy transfer. The agent is trained entirely offline on data from a source environment and must adapt to a target environment with structural changes (e.g., blocked paths) using only a single, observation-only demonstration. MSTT first learns a discrete skill vocabulary by tokenizing trajectory segments with a VQ-VAE. It then trains a skill-conditioned Q-function using a novel masked Bellman operator, which simulates dynamics shifts during training by randomly making subsets of skills unavailable. At test time, a feasibility mask is inferred from the skills observed in the demonstration, and the agent plans using the learned Q-function conditioned on this mask. A diffusion model generates trajectories corresponding to feasible skills. The method is shown to significantly outperform baselines on tasks with altered dynamics in Maze2D, FetchReach, and Habitat.", "questions": "- Robustness of Mask Inference: How sensitive is the method to the quality of the demonstration? For example, in Maze2D, what if the demonstration covers only the first half of the only feasible path? Can the agent still reach the goal, or does the incomplete mask render the problem unsolvable?\n- On the Sufficiency of Random Masking: Could you provide more intuition as to why random mask sampling during training is effective, given the combinatorial explosion of the mask space? Does the value function learn to generalize across masks with similar properties, or does the structure of the skill space make many masks irrelevant?\n- Beyond Binary Feasibility: The current model assumes skills are either feasible or not. Have you considered extending this to handle dynamics shifts that make skills more costly or stochastic, rather than impossible? For example, by learning a mask that represents a cost multiplier instead of a binary indicator.\n- Influence of VQ-VAE: The performance seems highly dependent on the quality of the learned skill codebook. Could you comment on how performance changes with the codebook size and provide any qualitative analysis of the learned skills to show they are semantically meaningful?", "limitations": "", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- Originality: The primary strength is the novel formulation of the masked Bellman operator for learning a feasibility-conditioned value function. This provides an elegant way to simulate a wide range of dynamics shifts during offline training without requiring multi-environment data.\n\n- Significance: The paper tackles the critical problem of generalization in RL under a practical, low-supervision setting. The ability to adapt from a single, action-free demonstration makes the approach applicable to real-world scenarios where target environment interaction is costly or unsafe.\n\n- Clarity: The paper's exposition is good. Complex ideas are broken down into understandable components, and the writing is concise.\n\n- Empirical Quality: The experiments are thorough and convincingly demonstrate the superiority of MSTT over strong baselines in diverse and challenging environments, including continuous control and visual navigation.", "weaknesses": "- Mask Inference: The test-time mask is inferred by assuming any skill observed in the single demonstration is feasible, and all others are not. This is a strong heuristic that may be brittle. If the demonstration is suboptimal or does not exhaustively cover all feasible and necessary skills for a given task, the agent's capability will be unnecessarily constrained.\n\n- Scalability of Mask Sampling: The training relies on random sampling of skill masks from a space of size 2^L, where L is the codebook size. For the reported L=240, this space is astronomically large. While the paper shows this works, there is no analysis on why random sampling is sufficient. It is plausible the method's effectiveness could degrade for problems requiring larger or less structured skill vocabularies where random sampling provides poor coverage.\n\n- Scope of Dynamics Shifts: The method is tailored for structural dynamics changes that render skills binary (feasible/infeasible). It is less clear how it would adapt to continuous dynamics shifts (e.g., changes in mass or friction), which is a limitation on the \"off-dynamics\" claim. The authors acknowledge this in a footnote.", "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)