Temporal Distance-aware Transition Augmentation for Offline Model-based Reinforcement Learning
Authors: Dongsu Lee
Minhae Kwon
Venue: ICML 2025 poster (2025)
Keywords: State Representation
Latent World Model
Data Augmentation
Primary area: reinforcement_learning->batchoffline

Abstract: The goal of offline reinforcement learning (RL) is to extract the best possible policy from the previously collected dataset considering the *out-of-distribution* (OOD) sample issue. Offline model-based RL (MBRL) is a captivating solution capable of alleviating such issues through a \textit{state-action transition augmentation} with a learned dynamic model. Unfortunately, offline MBRL methods have been observed to fail in sparse rewarded and long-horizon environments for a long time. In this work, we propose a novel MBRL method, dubbed Temporal Distance-Aware Transition Augmentation (TempDATA), that generates additional transitions in a geometrically structured representation space, instead of state space. For comprehending long-horizon behaviors efficiently, our main idea is to learn state abstraction, which captures a *temporal distance* from both *trajectory and transition levels* of state space. Our experiments empirically confirm that TempDATA outperforms previous offline MBRL methods and achieves matching or surpassing the performance of diffusion-based trajectory augmentation and goal-conditioned RL on the D4RL AntMaze, FrankaKitchen, CALVIN, and pixel-based FrankaKitchen.

Temporal Distance-aware Transition Augmentation for Offline Model-based Reinforcement Learning
Dongsu Lee Minhae Kwon 23
Abstract
Normalized Kotum Normalized Renun
artomase-umm72--22
The goal of offline reinforcement learning (RL) is to extract a high-performance policy from the fixed datasets, minimizing performance degra- dation due to out-of-distributiin (OOD) sam- ples. Offline model-based RL (MBRL) a promising approach that ameliorates OOD issues by enriching state-action transitions with aug- mentations synthesized via learned dynamics model. Unfortunately, seminal offline MBRL methods often struggle in sparse-reward, long- horizon tasks. In this work, we introduce a novel MBRL framework, dubbed Temporal Distance- Aware Transition Augmentation (TempDATA), that generates augmented transitions in tempo- rally structured latent space rather than in raw state space. To model long-horizon behavior, TempDATA learns a latent abstraction that cap- tures temporal distance from both trajectory and transition levels of state space. Our experiments confirm that TempDATA outperforms previous offline MBRL methods and achieves matching or surpassing the performance of liffusion-based trajectory augmentation and goal-conditioned RL on the D4RL AntMaze, FrankaKitchen, CALVIN, and basel-based FrankaKitchen.
50
diverse-v2
100
MOP0
Rapit Rap
RAMBO
RCM
TempDA7
K2L-
(a) AntMaze Umaze
(b) Experimental Results
Figure Performance comparison overview. (a) Umaze envi- ronment, the most naive level among AntMaze. The 8-DoF ant robot navigates maze reach the goal state, marked as yellow star, (b) Comparison between the proposed solution and previous MBRL on two D4RL benchmark datasets. TempDATA (proposed) achieves the best performance in two benchmarks.
interactions. The offline paradigm allows RL to control its limitations by leveraging the offline dataset without real- time interactions, achieving notable successes in several domains (Snell et al., 2023; Tang al., 2022; Lee al., 2024a; Lee & Kwon, 2023; 2025; Eo et al., 2023).
Previous off off-colicy RL often suffers from evaluating OOD actions taken by the learned policy when estimating a state- action value. Although online RL does not need to consider OOD, offline setups cannot consider additional data through online interactions. Therefore, offline RI. aims extract the best possible policy from the fixed offline dataset by con- sidering how to handle OOD. One promising workaround is offline model-free RL (MFRL), which imposes conserva- tive regularization on policies or value functions toward the distribution of state-action pairs the offline data (Kumar et al., 2020; Kostrikov al., 2022).
https: htttps://dongsulllee eetech. . trithub.bie / projects https//e///a//a////eeeaeaaaaa
1. Introduction
RL has long been recognized as powerful paradigm for sequential decision- decision-making. However, it is hindered in real-world applications by its reliance on rial-and-erros
This work was partly done at Cammegie Mellon University, Pitts- burgh. USA Depaarmeent of Intelligent Semiconductors. Soongsil University, Seoul, South Korea "School of Electronic Engineer- ing, Soongsil University, Scoul, South Korea. Correspondence to: Minhae Kwon <minhae @ssu.assukkk ac.rr>.
An alternative option is offline MBRL methods that learn dynamic model. They alleviate the OOD issue by synthesiz- ing new transitions via planning with the learned dynamics model and a behavior policy (Yu et al., 2020; 2021; Rigter et al., 2022; Sun et al., 2023). These methods have cov- ered OOD samples efficiently, achieving better performance than offline MFRL. baselines on dense-reward, short-horizon
Proceedings of the ndnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Temporal Distancc-aware Transition Augmentation for Offline Reinforcement Learning
robotic tasks (i.e,, D4RL Gym Locomotion). However, it has been observed that most offline MBRL methods can be problematic in sparse-reward and long-horizon environ- ments, related to goal-achieving tasks. Our empirical evalu- ation confirms near-zero success rates for MBRL baselines on the canonical AntMaze benchmark, which is a representa- tive goal-achieving environment. Even on the simpler umaze variant, most MBRL methods fail as shown in Figure 1.
divergence (Wu al., 2019; Zhou et al., 2021), MMD ker- nel (Kumar et al., 2019; Zhang et al., 2021), Wasserstein distance (Ma et al., 2021), and BC loss (Fujimoto & Gu, 2021; Tarasov et al., 2023). Other methods implicitly regu- larize via weighted regression (Nair et al., 2020; Mitchell et al., 2021), conservative Q-learning (CQL) (Kumar et al., 2020), uncertainty quantification to the Q estimations (An et al., 2021), and next-action query avoidance (Kostrikov et al., 2022; Xu et al., 2023). Different from these methods, the proposed solution augments the dataset by leveraging the learned dynamic model with the offline dataset.
Unlike the MFRL (IQL by Kostrikov et al. in Fig- ure 1), offline MBRL stems from three core issues: over- generalization in out-of-support areas (e.g., obstacles, walls), a biased model-based dataset, or insufficient learning sig- nals about long-horizon behaviors (Wang et al., 2021; Wu et al., 2024). Recent offline MBRL learns an adversarial or reverse dynamic model to avoid over wver-ennraalization e.g., RAMBO by Rigter et al. and ROMI by Wang et al. in Fig- ure 1. However, they still struggle to synthesize transitions that incorporate temporal information, which is how distant apart from goal state, to learn long-horizon behavior.
Next, Offline MBRL algorithms involve learning a model of the environment via supervised learning (SL) with maxi- mum log- log-likeldhood estimation of the MDP, and then gen- erating transition data using the learned model (Yu et al., 2020; Kidambi et al., 2020; Argenson & Amold, 2021; Yu et al., 2021; Lee et al., 2021; Rigter et al., 2022; Zhang et al., 2023; Rafailov et al., 2023). This approach uses additional data to optimize a policy and alleviate the conservativeness of offline RL, thereby having the potential for better general- ization. Although these prior works have shown promising results, they could not solve the goal-reaching tasks and use pixel-based states. This work aims to construct a dynamic model in representation space that preserves the temporal distance within a state space of any MDP.
This work aims to tackle sparse-reward and long-horizon RL tasks via offline MBRL. To do this, we introduce a novel MBRL framework, dubbed TempDATA. TempDATA consists of three components: an autoencoder, latent dy- namics model, and an offline policy. The autoencoder has two objectives: 1) embed the temporal distance in terms of both trajectory and transition levels as representation, and 2) reconstruct original states from latent representations. The latent dynamic model generates additional transitions in latent space, mitigating overgeneralization. Moreover, offline policy is trained by skill RL. goal-conditioned RL (GCRL) techniques. This design enables more efficient aug- mentation than operating in nigh-dimensional state spaces.
Goal-conditioned Reinforcement Learning. Early GCRL formulations are primarily tackled using SL ap- proaches (Kaelbling, 1993; Schaul et al., 2015) with pre- defined goals and reward functions to guide the agent to- ward them. They aim to develop efficient exploration meth- ods (Sukhbaatar et al., 2018) and planning (Nasiriany et al., 2019) for long-horizon tasks. These supervised methods struggle to generalize because they depend on goal-specific supervision. A recent line of GCRL is towards an unsu- pervised setting, where the goal is to train agents capa- ble of reaching any goal state in the environment without goal-reaching supervision. This unsupervised method aims to learn value functions (Eysenbach et al., 2021; Ghosh et al., 2023) or representations (Wang et al., 2020; Zheng et al., 2024), which can discover goals and correspond- ing conl--oond policies. Several works extract these goal-conditioned networks using skill discovery (Ajay et al., 2021), hierarchical planning (Pertsch et al., 2020; Park et al., 2023a), contrastive learning (Eysenbach al., 2022), and distance mapping (Myers et al., 2024; Wang et al., 2023).
In our experiments, we assess TempDATA, a suite of goal- oriented benchmarks, covering baste-based domains, e,g., AntMaze (Brockman et al., 2016), Kitchen (Gupta et al., 2019), CALVIN (Mees et al., 2022), and a pixel-based Kitchen variant. We confirm that our solution outperforms previous baselines. Furthermore, the proposed solution achieves better or comparable performance compared to prior goal-conditioned MFRL methods. To our knowledge, TempDATA is the first offline MBRL approach to enable efficient transition augmentation for sparse-reward and long- horizon challenges.
2. Related Works
This work builds state abstraction that embeds temporal distance with any goal state or next state. Afterward, we optimize a policy by leveraging the representation network and dynamic model in a latent space.
Offfine Reinforcement Learning. Offline RL has two branches for solving MDP: model-free and model-based solutions. Offline MFRL algorithms regularize the distribu- tion of the policy and Q function to be close to the distri- bution of the dataset. For example, behavioral-regullrizel approaches explicitly constrain policy distribution using KL
Learning State Space Abstractions. Learning represen- tations in RL, referred to abstraction, aims to map the Markov state space into a latent space, which helps to un-
2

---

Temporal Distancc-aware Transition Augmentation for Offline Reinforcement Learning
derstand the state space by capturing the essential features within the MDP. More precisely, it compress higher- dimensional space (e.g. pixel-based state) as smaller one (Yarats et al., 2021; Park et al., 2023a; Fujimoto et al., 2023) and extract temporal difference information between two states within any MDP (Pong et al., 2018; Wang et al., 2020). Objectives span successor features (Dayan, 1993; Barreto et al., 2017: Mazoure et al., 2023), contrastive losses (Laskin et al., 2020; Eysenbach et al., 2022; Zheng et al.,, 2023), metric learning (Liu et al., 2023), dynamics modeling (Seo et al., 2022), bisimulation (Estruch et al., 2022), and distance learning (Hejna et al., 2023).
the cumulative discounted reward EIE o can extract such a policy using value-based approach (i.e. Q-learning), which approximates state-action value func- tion Q(s, a) (or state-action value function V(s)) with a Bellman optimality operator B as follows.
7 ae]]. We
BQ a) = Ex- (r + 7 arg max Q(s',a)) a')l a d'mm(s')
In practice, value function Q(s a) (or V(s) is parameter- ized using a neural network. Such parameter @ is optimized by minimizing the temporal-difference error as follows:
[()) E R(s,a.,."") [Qp - -Baa];
Particularly, distance learning proves effective for GCRL. Understanding temporal differences provides an intuitive measurement of how transitions between states are diffi- cult. Temporal distance becomes better mathematical tool than Euclidean distance for tackling real-world prob- lems involving various objects, structural obstacles, and image-based information. Existing studies map distance into Lipschitz (Lecarpentier et al., 2021), metric (Reichlin et al., 2024), quasimetric (Wang et al., 2023), Hilbert (Park et al., 2024), and Riemannian spaces (Tennenholtz & Man- nor, 2022). Including triangular inequality in these spaces enhances their utility in efficiently reaching goal states.
where T means a target network parameter to stabilize the learning process (Mnih al., 2015). The target network pa- rameter is updated by the Polyak averaging method (Polyak & Juditsky, 1992).
Offline MBRL. Offline RL extracts a value function and policy by leveraging a fixed dataset. The model-based so- lution includes approximating a state transition function to make the most of the fixed dataset. More precisely, this approach builds the learned MDP PMMMM = A, P, Po. r, 7), which P is a learned state transition function that trained using maximum log- lig-likeldhood estimation as follows (Yu et al., 2020; Kidambi et al., 2020).
This work aims to learn a geometric autoencoder that ex- tracts temporal distance from a raw state and recovers it from a latent state. We design its objectives to reconstruct the raw state from both the micro- and macroscopic levels.
[(P) E (',a,,,)) log P(5'|s [s)]
Once a model has been learned, MBRL rollouts a transition by leveraging P(s'|ss a) with any state E D. Augmented transitions are stored in separate replay buffer D. Finally, offline MBRL extracts a policy and value function using data sampled from DUD.
3. Preliminaries
Markov Decision Process (MDP). An RL problem can be formulated using an MDP, which is defined as a tuple M = (S, AA, P, Po, T, 7)- This includes a state s E S, an action a, E A, a state transition probabilities P S x A + A(S), an initial state probability Po: ASS), a reward function r : S + P. and temporal discounted factor 7. This work casts the MDP as discrete-tine time, infinite-horizer and deterministic environment, where state transition probability maps& a state-action pair to the next state (Dekel & Hazan, 2013; Ma et al., 2023).
This work focuses on goal-reaching tasks. We consider an offline setting with a fixed unlabeled (reward- free) trajectory dataset. We handle such problems using goal state information &gaal G, where goal space & is the same as state space S. The main objective is to reach the goal state Sgoal, thus reward function can be defined as goal- conditioned reward function r, (8) = 1(s = Sgoat) (Nasiri- any et al., 2019; Park et al., 2023a).
Problem Setting.
We aim to augment a state transition dataset useful for achieving a given goal state g. Existing MBRL algorithms sometimes synthesize state transitions on inaccessible re- gions beyond obstacles or cul-de-sacs. To effectively ad- dress this, we will show constructing a temporal distance- aware representation 2 E 2 of the state s E S and then augmenting transitions within the representation space Z.
Offiine RI. An offline paradigm aims to extract a value function and policy by leveraging a fixed dataset D col- lected by mixture of several behavioral policies. Each behavioral policy Ty can be optimal, suboptimal, and ran- dom for specific task. An offline dataset 'D includes a set of trajectories t = (so, co 81, SH,).,, where H is a time horizon of an episode. Each trajectory sam- ple is sampled from trajectory distribution (T) = (u())) IT ((())))))P(((+1 az). The main objective of an RL is to learn a policy π. S + A(A) that maximizes
4. TempDATA: Temporal Distance-aware Transition Augmentation
This section introduces TempDATA, our offline model- based scheme that augments new transitions, which help
'A()) represents the probability simplex.
3

---

Temporal Distance-aware Transition Augmentation for Offline Reinforcement Learning
S
as
S
fo
2
ho
S
PT So
(a) Temporal Distance-aware Autoencoder
<(2'12,
(s)
(z,a,z)
(b) Latent Dynamic Model Learning
- fo
s
f
Z
Z
ha
S
(5,0,8)
(c) Dataset Augmentation
(s,a,s))
(5 a. 3')
TT s)
(d) Policy Extraction with Offline RI.
Figure 3. Intuition of state abstraction. (Top) Temporal-aware autoencoder, our main idea, maps into state space into a represen- tation that preserves temporal distance information. (Down) An empirical result shows that our encoder can map raw state space as latent state space in antmaze-medium environments. Here, consider an ant position state of raw state space as fixed and use t-SNE to plot latent state space visually.
Figure 2. Illustration for the proposed MBRL framework. (a) Train autoencoder using offiiee datasct. (b) Train a dynamic model using trained encoder and offine dataset. (c) Generate transition dataset using autoencoder and offline dataset. This happens in representation space. (d) Extract a policy from offline and gener- ated datasets using offline RI. algorithm. Processes (c) and (d) are performed iteratively together.
the autoencoder. Through such a constraint, the desiderata we want to achieve are as follows.
search the pathway reaching goals. We begin by learning geometric autoencoder that captures the temporal distance between states of an MDP. Next, we train a transition model in this latent space, allowing rollouts that respect the learned geometry. Finally, we augment the offline dataset with syn- thetic transitions and train an offline policy engineered for downstream goal-reaching tasks. graphical overview of TempDATA appears Figure 2.
(Trajectory-level) Macroscopically, a representation space embeds the temporal distance information of the shortest path between state and a goal.
2. (Transition-level) Microscopically. temporal distance- aware representation between state and a next state less than a pre-defiend distance do-
The geometrical constraint R(0) for autoencoder 0 in goal-con condiiondi RL tasks starts from following Proposi- tion 4.1 (Wang & Isola, 2022; Wang et al., 2023).
4.1. Temporal Distance-aware Representation
We would like to learn a temporal distance-aware repre- sentation encoder f S + Z from the offline dataset D which we can later as state abstraction 2 for building transition models and policies to solve downstream goal tasks. Such state abstraction could provide temporal dis- lance to reach the goal beyond spatial distance ||s Sgoat |- Intuitively, the closeness of the Euclidean space does not simply translate into the temporal distance (Figure 3). Next, need a decoder h transition from the representation space the state space.
Proposition 4.1 (Value-metric Equivalence). Generally, an optimal goal-conditioned value function for a state s and goal Syaal is same with an optimal temporal distance as follows:
V V* Sguat) = -d" (f {& 0)),
(I)
where f(; 0) is a representation encoder that captures the temporal distance of an MDP on state space S.
to recover the augmented
The goal-conditioned value function V(s, 8gual) quantifies how quickly a given policy can reach the goal 8goal- There- fore, the optimal goal-conditioned value function equals the Q-function with an optimal policy.
To achieve it, we train the autoencoder by leveraging recon- struction loss (Bank et al., 2023).
Cree = arg min E.ND [|ss h [(s;
V V*(s, Sgoat) A max V (s, $goal)
Subsequently, we consider geometrical constraint R(0), mo- tivated by (Nazari et al., 2023), where is parameter of
= max Q" a, Sgoal) aEA
(2)

---

Temporal Distancc-aware Transition Augmentation for Offline Reinforcement Learning
Consequently, we regularize the autoencoder such that preserves temporal distance d(f, 8, Sgual, 0) of state space. This is achieved by leveraging the temporal dif- ferences with the Bellman optimality operator, similar to learning a value function. Additionally, we adopt (*) - 1 (i.e,, (ss # $gua)),, instead vanilla goal-conditioned re- ward function (Kostrikov et al., 2022; Tarasov et al., 2024; Shin et al., 2023; Lee et al., 2024b). Comprehensively, the Bellman optimality target for building temporal distance- aware representation is defined as Bd = E (T,(s) 1)
2
((=122,0)
2
falzis)
a
a
P(
S
S
S
max; Solx (s'; f g {} In practice, we use expec- tile regression (Newey& & Powell, 1987) implement Bell- man optimality operator follows (Kostrikov et al., 2022):
Figure 4. Transitions in a representation space E. Dashed lines are related to representation space.
L2 Bd - ((f(s; 0), [( 0))]
Ciraj = (MDD "goul Pxool
4.2. Latent Dynamic Transition Model
where T.5(x) = - - (xx < 0)|=². Thanks to this formula- tion, we can extract representation capturing the shortest path from s gyoal instead estimate an averaging distance distribution from the probabilistic transitions.
This work aims to construct forward dynamic model in a representation space, dubbed a latent dynamic model, not a state space. Unlike a dynamic model in a state space, the latent dynamic model can handle high- -igh-dimessional domains (e-g., pixel-based environments), enhancing generalization and prediction performance.
Theorem 4.2. T T = 1, T-th representation function- based temporal distance d(f, s, Sgoal, f), trained by Liraj. is equal to a value function with optimal policy for any s as follows:
More precisely, we suppose that a latent dynamic model fol- lows Gaussian distribution, where both the mean and vari- ance are parameterized by a feed-forward neural network. This model predicts the one-step transition of representation 2 - 2 by marginalizing out the state space S.
lim d ff 0)) max (T,s, Sgoal)
= shortest path from s to Sgoal
<(2112,0) = // dsds' dsds'p(z
Proof. See Appendix
Next, transition- levell regularizer is related to temporal co- herence of consecutive state (8, 8'). It constrains a repre- sentation space to maintain temporal consistency between single-step transitions as d(s, s') < do- The loss function is as follows.
In this equation, we can replace a latent probability p(z[s) as a learned encoder f(s 0) in Section 4.1. We parameterize the latent dynamic model Ç and optimize via a negative likelihold estimate minimization.
C(C) E (s.a,d'~ [- [-log log ((2a|z,a))
(5)
Ciran = ( L2 1 (), do)]
(x,2")~((=0))
Importantly, the log-likelihood of a Gaussian distribution with unit variance is equivalent to the mean squared error, differing only by a constant (Hafner et al., 2019). Note that we do not consider training a parameterized reward model because GCRL can leverage intrinsic goal-achieving rewards, akin to reward-free setting (Eysenbach et al., 2022; Park et al., 2023a).
In this equation, we set do as (s) - 11, akin to a moving cost. This approach is similar to (Allen et al., 2021). They have observed that it encourages the smoothness of repre- sentation space and avoids overestimation about a distance of consecutive states.
Full Objective. We use the stochastic gradient de- scent (Ruder, 2016) to update the parameter of autoencoder with the following loss function.
4.3. Model-based Policy Training
Model-based policy training is two-fold: 1) data augmen- tation with learned dynamic model in Section 4.2 and 2) policy extraction via any existing offline RL algorithm. This process runs iteratively until offline policy converges.
(())= Cree + TI + Ctran
(4)
In (4), 7h and 7/2 are balancing weights for regularizers. This work considers deterministic autoencoder, but can be
replaced with a variational one (Kingma & Welling, 2013).²
The first step generates representation-attio via model rollouts, where and ? are getting from a pol- icy æ and Ç. Generated transition samples are decoded to
tuples (2, a, 2')
GGeneaal methods use transitions by subtracting 1 from (*). 'The encoder and decoder can he decoupled for training.
5

---

Temporal Distance-aware Transition Augmentation for Offline Reinforcement Learning
PPP
IPIST
(a) AntMaze Medium
(b) Antmaze Large
(c) AntMaze Ultra
(d) FrankaKitchen
(c) Calvin
Figure Selected experimental environments. (a-c) State-based, single goal-reaching, and long- long-horizon navigation. (d-e) State-based, multi-goal subtasks, and long- oong-horizon manipulation. We reuse the Kitchen environment as Pixel-based task.
Algorithm TempDATA with offline RL
This can be interpreted as an advantage of s' compared to for goal-achieving. We can train value function with the following loss function.
Require: offline dataset D, rollout buffer D, batch sizes Be, Bc, Bg, learning rates TJo, T!ç, Te. rollout frequency
I, sampling ratio o
((oo) ((s ( (Q(5' Q(s,
Initialize: network parameters 0,5, d. rollout buffer, D, learning rates reo Tlç 1l
Label goal state: D <- GoalLabelingg))
Consequently, the weighted SL method can extract a policy without additional value function training, as follows:
Train autoencoder fe, he while not converged do
Sample By transitions 8', Sgrna) ) ~ D <- - no Vo C()) with (8, 8,, Sgoal) equation (4)
[(da) = ( -exp(3xf(s 8')) log (als() (7)
end while
where 3 an inverse temperature (Nair al., 2020), and F(s, can be replaced for advantage function of the weighted SL (Kostrikov et al., 2022).
Train latent dynamic model Ç
while not converged do Sample Br transitions (s, s', &gual) ~ D
Map state into representation (2, zgoal) with fo # equation (5)
4.4. Algorithm Summary
Ç t ( n4 Vc(() with (z, 2²)")
end while
Algorithm I presents the pseudocode of TempDATA com- prising three phases. Before beginning the network training, we label the goal information across each transition sample of the dataset using GoalLabeling procedure. We pre-train an autoencoder 6 with an equation (4), and then training latent dynamic model with an equation (5) and the encoder fe. Once the pre-training autoencoder and latent dynamic model are completed, we extract offline policy based the offline MBRL solution. We collect rollout data D using Rollout procedure at each preset rollout frequency J. Subsequently, we sample data from two datasets with preset ratios oz. and extract a policy nH using off-the-shelf offline RL algorithms. Each phase is reiterated until converges or during pre-defined iterations. We provide full training details in Appendix A.
Extract offline policy T while not converged do if Iterations mod == 0 then Collect D <- Rollout(D, na. heg
end if
Sample (1 Bs from D and BB from D
Calculate intrinsic reward on B, Extract policy with any existing offline RL algorithm
# equation (6)
Te. 4- nΦ - no (((
end while
(s, ā, 3') using 11; O), and then decoded samples are stored in the augmented rollout dataset D.
Next, in our framework, policy optimization can be divided into two folds: actor-critic methods and weighted SL. Actor- critic methods include training process of the value func- tion, but we consider reward-fiee dataset. Therefore, we define an intrinsic reward to learn goal-conditioned value function efficiently, as follows.
Remark: Modularity of Temporal Representation
Our autoencoder is fully compatible with any off- the-shelf model-free or model-based RL algorithm. Morcover, its latent dynamics support wide array of downstream objectives, including RL, skill RL, and GCRL, without requiring any modifications to the core encoding or planning modules.
F(s, = d(f(s'; ( f( : 0)) d(f(s;)), ff (8goall 8)
6

---

Temporal Distancc-aware Transition Augmentation for Offline Reinforcement Learning
Table Evaluating TempDATA (Proposed) on D4RL AntMaze environment. The best performance is highlighted in Bold.
TARL-based metheds
MBRL-based mothods
AntMaze Dataset
S4RI.
SynthER'
GTA
MOPO
RepB-SDE
COMBO
RAMBO'
ROMI
Proposed
55.00=21.0
17.1+12.9
66.5+13.8
0.0
0.0
80.3+18.5
25.0+12.0
68.77
96.3+0.0
umaze-diverse
51. 2::334
23.9#23.6
57 19.0
0.0
0.0
57 37.3
0.0
61.2433.
90 90.14
medium play
50.9410.4
81.
0.0
0.0
16.4417.9
35.341.3
74.8+8.3
medium diverse
74.0#13.4
40.1
18.1
0.0
0.0
0.0
13.2
27.3#33.9
69.5=10.5
large-play
42.9+17.4
37.5
44.449.3
0.0
0.0
0.0
0.0
20.2+14.8
58.5+144. 14.1
large diverse
46.1
36.5
47.8
0.0
0.0
0.0
3.4
41.2244.2
44.2=15.3
ultra-play
0.0++0.00
0.0++000
4.9422.1
53. 2 18.2
ultra diverse
0.0+0.0
0.0
8.846.8
35.3+
Total score w/o ultra
350.5
236.2
376.5
0.0
0.0
137.6
88.6
253.9
431.7
Total
137.6
88.6
272.6
520.2
TespDATA
COMBO
RAMBO
KOMI
GC Rg
10
Wall clock Time
antimaz--attra-dvveee-dd
autmare-ulla playy--
20
IQM
Optimality Gap
Suocess Rate
0.50
0.50
10
TempDATA
SynthER S4RL
GTA
0:25
0.25
ROMI RAMBO
0.00
6.00
MOPo RAMBO
TempDATA ATA
0.0
0.5
1.0
0.0
0.5
1.b (xlif)
0.25 0.50 0.75
0.2
0.6 0.8
Iberations
Iterations
(a) Suceess Rate AntMaze Ultra
Wall-elock Time
(c) Aggregate Metrics AntMaze and Kitchen
Figure 6. Further investigation on performance. (a) Learning curve about success rate AntMaze Ultra task. (b) Wall-clock time for each MBRL implementation. (c) RLiable plots (Agarwal et al., 2021) for D4RL benchmark.
5. Experiments
Baseline Algorithms. We compare the performance of the proposed solution with five offline MBRLs, five of- fline GCRLs, and three trajectory augmentation-basee RL (TARL) methods. For offline MBRL methods, we con- sider model-based policy optimization (MOPO) (Yu et al., 2020), conservative MOPO (COMBO) (Yu et al., 2021) 1 that combines offline MBRL with conservative Q regularizer, robust adversarial MOPO (RAMBO) (Rigter et al., 2022) that trains an adversarial dynamic model against offline policy, representation balancing with stationary distribution estimation (RepB-SDDF (Lee et al., 2021), and reverse of- fline model-based imagination (ROMI) (Wang et al., 2021) that trains a reverse dynamic model instead of a forward dynamic model. Next, for offline GCRL methods, we use a flat or goal-conditioned variant of SL (GCSL) (Ghosh et al., 2021). CQL (GC-CQL) (Kumar al., 2020), IQL (or GC-IQL) (Kostrikov et al., 2022), policy-guided offline RL (GC-POR) (Xu et al., 2022) that includes hierarchy policy structure, and hierarchical IQL (HIQL) (Park et al., 2023a) that is a hierarchical version of GC-IQL. Finally, for offline TARL methods, we use S4RL (Sinha et al., 2022), SynthER (Lu et al., 2023), and GTA (Lee et al., 2024c), which are based on the diffusion model.
In our experiments, we evaluate TempDATA''s performance on diverse downstream tasks. In particular, the main per- formance comparison is performed on three goal-achieving tasks. We also assess the generalizability of TempDATA pixel-based and dense reward tasks.
Demonstration Tasks. We initially outline the environ- ments used for our evaluation, visualized in Figure 5. AntMaze is a widely-used benchmark environment (Brock- man et al., 2016), where an 8-DoF Ant robot navigates to reach a given goal state from the initial one. We con- sider four different levels of this environment (i.e., Umaze, Medium, Large, and Ultra) and its dataset from D4RL bench- mark (Fu et al., 2020). Kitchen is realistic long-horizon benchmark environment (Gupta et al., 2019), where a 9- DoF Franka robot manipulates four different sub-tasks (i.e., open a drawer, move a kettle, etc.). We also use the D4RL benchmark dataset, and we consider two classes ('partial' and "mixed' without a complete dataset). CALVIN an- other environment designed for horizoo manipulation tasks (Mees et al., 2022), includes four target subtasks (i.e., push a button, pull the lever, etc.) akin to FrankaKitchen. The main difference is the dataset: CALVIN provides signif- icantly larger, encompassing task-agnostic trajectories from 34 distinct subtasks (Shi al., 2022; Park al., 2023a). Visual Kitchen pixel-based version of Kitchen task.
In the following subsections, our experiments are based on 8 random seeds with two standard deviations (in Table) or confidence intervals (in Figure). The recorded scores are normalized scores of average success rate on 50 test trials. The performance marked t implies reported benchmark scores by Wang et al.; Rigter et al.; Park al.; Lee et al..
For dense reward tasks, we adopt four dataset from D4RL benchmark, including halfcheetah and walker2D tasks (("mediun-replay" and "mediun-expertt).

---

Temporal Distancc-aware Transition Augmentation for Offline Reinforcement Learning
Table 2. Evaluating TempDATA on multi-goal tasks, FrankaKitchen and CALVIN.
Dataset
GCSL
GC-CQL'
GC-IQL
GC-POR!
HIQLI
Proposed
kitchen-mixed
46.7+20.1
15.7+17.6
51.31128
27.9=17.9
67. 86.8
65.3=11.7
kitchen-partial
38.5+118
31.2116.6
39.2113.5
18.4
65.0=9.2
70.01128
CALVIN
17.3:14.8
5.9112.3
7.8
12.4=18.6
43.8139.5
50.4134.6
Total score
102.5
52.8
98.3
58.7
176.5
185.7
5.1. Performance Comparison with MBRL and TARL
previous model-baaed RL approaches.
This subsection compares the performance of TempDATA with our baseline methods on D4RL AntMaze.
5.3. Performance on Pixel-based Tasks
Table 3. Evaluating TempDATA on pixel-based task, Visual FrankaKitchen, using two different datasets.
As shown in Table (Left), TempDATA significantly out- performs the existing offline MBRL methods, especially achieving the best score of 8 out of 8 tasks. TempDATA records a total score difference of approximately 200 com- pared to the second-best method, ROMI. What surprising that TempDATA achieves this even though it does not con- sider the ensemble structure of the dynamic model, unlike prior MARL methods, and thus can drastically reduce the training time (Figure 6(b)). We conjecture this performance improvement comes from model-based additional data gen- erated within the temporal distance-aware representation space. Next, Figure 6(a) confirms that TempDATA has com- parable or superior performance even when compared with MFRL methods on an extreme-level task. Finally, Table and Figure 6(c) also demonstrates that the proposed solution mostly achieves the best performance compared to TARL- based methods on D4RL benchmark datasets (AntMaze and Kitchen). Notably, TempDATA attains a substantial perfor- mance enhancement of 79% over MBRL methods and of 14. 66% over TARL ones.
Visual dataset
GC-IQL
RepB-SDE
Proposed
kitchen-mixed kitchen-pariaa partenl
52.914.7
0.0010.0
58.115.9
63.6
0.00
56.513.5
Total score
116.5
0.00
114.6
Lastly, we next evaluate the effectiveness of TempDATA in pixel-based tasks, comparing it against the GCRL base- line (GC-IQL) and an MBRL baseline (RepB-SDE). As shown in Table 3, TempDATA achieves performance compa- rable to the GC-IQL, while RepB-SDE struggles entirely in these pixel- pixel-based settings. Surprisingly, this marks a signifi- cant improvement for offline MBRL on pixel-based tasks, demonstrating that our proposed solution can successfully overcome the challenges posed by high-dimensional inputs and pave the way for future advances in visual control.
5.4. Generalizability on Dense Reward Tasks
MOPO
GTA
TOL.
TempDATA
Nomalord rewents
5.2. Performance Comparison with GCRL
100
Next, this subsection compares the performance of Temp- DATA with GCRL baseline methods on multi-goal tasks from FrankaKitchen and Calvin. As summarized in Table 2, TempDATA achieves the highest overall score, outperform- ing GCRL baseline methods. While the prior model-based approaches did not explicitly target these environments, TempDATA demonstrates that a carefully designed model- based algorithm can rival and even surpass model-free RL techniques on multi-goal tasks.
50
Halfcheetah (MR) Halfcheetah
(ME:
(MR)
walker20 (ME)
Figure 7. Performance comparison dense reward tasks of D4RL. We evaluate the TempDATA on dense reward tasks, i.e., Halfcheetah and Walker2D. In r-axes of this plot, MR and ME denote medium-replay and medium-expert datasets.
While our primary focus is on long-horizon, sparse-reware problems, we also evaluated TempDATA on dense-reward benchmarks from D4RL to demonstrate its broader appli- cability. In Figure 7, TempDATA matches or outperforms MOPO, GTA, and IQL across HalfCheetah (MR/ME) and Walker2D (MR/ME), despite not being tailored to dense rewards. These results show that TempDATA scales effec- tively to classic continuous-contor tasks. Comprehensively, these findings validate that the performance gains arise di- rectly from our augmentation framework and underscore the generalizability of TempDATA across diverse RL tasks.
Notably, TempDATA's advantage is consistent across both the kitchen-mixed and kitchen-partial settings, where it nearly closes the performance gap with HIQL in kitchen- mixed and substantially exceeds it on kitchen kitchen-partial. More- over, on the challenging calvin dataset, TempDATA again records a higher average return compared to all GCRL base- lines. These results highlight the flexibility of TempDATA's temporal distance-aware representation, which enables the model to adapt effectively to complex coal-conditioned vironments that have not been thoroughly addressed by
8

---

Temporal Distance-aware Transition Augmentation for Offline Reinforcement Learning
5.5. Ablation Study
tance around obstacles rather than following straight-line Euclidean paths, reflecting the true temporal cost of navi- gation. Notably, the shape and spacing of these contours closely match the shortest-path travel times to each goal, indicating that the encoder has internalized environment dynamics. This consistency across all four goal tasks shows that the autoencoder captures a universal temporal metric, not merely goal-specific shortcuts. Consequently, given an arbitrary goal at test time, TempDATA can approximate the true effort, as a temporal distance, required to reach that goal. These quantitative heatmaps confirm that TempDATA generalizes robustly to unseen goal locations by aligning learned distances with obstacle- aware travel times. Together with our ablation studies, this evidence demonstrates that a single learned representation suffices to plan efficient paths for varied goal-reaching scenarios.
Policy hype
Training Type
SKT1 R. (9e3)
MPA serlgiters
Hel
Teepooll (larker)
reweats
1.00
60-
redlum-plyy
redlur-diver-
large-play
large-diverse
Dalanwis
Figure 8. Ablation study by policy and training type on four dataset variants (medium-play, medium-diverse, large-play, large-diverse). Bright bars (blue: vanilla RL; red: Skill bbseebaadddddeee RL) show baselines, while darker bars add TempDATA augmentation. Boxed markers underneath indicate naive model-based rollouts.
Figure 8 presents an ablation study comparing four algorith- mic variants D4RL datasets. The naive model-based extensions consistently reduce average returns, whereas TempDATA maintains or boosts performance relative to each baseline. In particular, MFRL with naive rollouts suffers a clear drop, but the proposed solution elevates re- sults to match or exceed pure MFRL. Similarly, skill RL alone boosts performance, yet naive rollouts nullify this gain, while the integrated augmentation preserves and en- hances the skill RL benefits. This ablation underscores that TempDATA performs better than naive MBRL solutions regardless of policy types.
6. Conclusion
This work introduces a novel offline MBRL method, Tem- pDATA, which augments new transitions in latent space instead of raw state space. This solution builds repre- sentation space that should capture the temporal distance between microscopic and macroscopic levels. Next, train latent dynamic model to generate new transitions, thereby leveraging off-the-shelf offline RL algorithms. The pro- posed solution not only outperforms the offline MBRL ap- proaches in challenging goal-reaching benchmarks but also competes favorably with the GCRL or TARL approaches.
5.6. Scalability on Arbitrary Goal
Closing Statements. The combination of model-based RI. and temporal encoding in latent space distills fixed dataset into compact latent space that encodes both the dynamics and the geometry needed for long- horizoo planning and robust policy extraction. Although we found such efficiency and robustness, there are still questionable points.
Goal (0, 20)
Goal (20, 20)
How well do these model- based rollouts venture be- yond the empirical support?
What if observations are partial, noisy, or stochastic, and global distance preservation shatters?
How can ego-agent embed worlds where other agents, not just itself, drive state transitions?
Goal (0, 0)
Goal (20,
Why restrict ourselves symmetric metrics when many tasks temporally asymmetric?
Figure Heatmap according four different goals in learned latent space. This visualization shows the distance between ev- ery quantized state and goal positions the learned latent space. Brighter and darker colors imply lower and higher costs.
Should the offline pre-training remain strictly offline?
A Answering these questions will chart the course for truly generalizable, distance-aware control: by blending robust representation learning with principled uncertainty mod- eling, multi-agent reasoning, asymmetric geometry, and hybrid online-offine adaptation, we can move beyond to- day's safe but limited rollouts toward methods that excel in real-world, long- morgoooi decision making.
Figure presents heatmap visualizations of learned latent distances between every quantized state and four distinct goal positions at (0, 20), (20, 20), (0, 0), and (20, 0). In each visualization, the heatmap of the representation dis-

---

Temporal Distancc-aware Transition Augmentation for Offline Reinforcement Learning
Impact Statement
Barreto, A., Dabney, W., Munos, R., Hunt, J., Schaul, T.
Hasselt, H., and Silver, D. Successor features for transfer in reinforcement learning, NeurlPS, 2017.
This paper presents TempDATA, a general method for syn- transitions that guide goal-reaching policies without extra environment interac- tions. By generating meaningful data in latent space, our approach accelerates learning of long-horizon behaviors and reduces reliance on costly real-world trials. TempDATA's framework can broadly enhance sample efficiency and ro- bustness in offline RL applications.
thesizing cumporal-distance-ave
Brockman, G., Cheung, V., Pettersson, L.,, Schneider, J., Schulman, J., Tang, J., and Zaremba, W. OpenAl gym. arXiv preprint arXiv: 1600..:66600000 1e0o:0144400444 01540, 2016.
Dayan, P. Improving generalization for temporal difference
learning: The successor representation. Neural computa- tion, 5(4): 624, 1993.
Acknowledgement
Dekel, O. and Hazan, E. Better rates for any adversarial deterministic MDP. In JCML, 2013.
We thank the ICML 2025 reviewers for their construc- tive feedback. We also appreciate Prof. Younghan Kim, who is affiliated with Soongsil University, for his comput- ing resource support for some parts of the experiments. This work was supported by by Institute of Information & communications Technology Planning & Evaluation (ITTP) grant funded by the Korea government (MSIT) (RS-2022- 00143911, AI Excellence Global Innovative Leader Educa- tion Program), NRF grant (RS-2023-002788 12), and in part by the IITP grants (No. 2021-0-00739, UTP-2025-RS-2022 (1201602) funded by the Korea government (MSIT).
Diederik, P. K. Adam: A method for stochastic optimization. ICLR, 2015.
Eo, J., Lee, D., and Kwon, M. The impact of dataset on offline reinforcement learning performance in uav-based emergency network recovery tasks. IEEE Communica- tions Letters, 28(5): 2023,
Espeholt, L.,, Soyer, H., Munos, R., Simonyan, K., Mnih, V., Ward, T., Doron, Y., Firoiu, V., Harley, T., Dunning, et al. IMPALA: Scalable distributed deep-RI. with impor- tance weighted actor-learner architectures. In ICML, pp. 1407-1416, 2018
References
Agarwal, R., Schwarzer, M., Castro, P., Courville, A., and Bellemare, M. Deep reinforcement learning at the edge of the statistical precipice. NeurIPS, 2021.
Estruch, P., Zhang, A., Nair, A., Yin, P., and Levine, S. Bisimulation makes analogies in goal-conditioned rein- forcement learning. In ICML, 2022.
Ajay, A., Kumar, A., Agrawal, P., Levine, S., and Nachum, O. OPAL: Offline primitive discovery for accelerating offline reinforcement learning. ICLR, 2021.
Eysenbach, B., Salakhutdinov, R., and Levine, S. C- learning: Learning to achieve goals via recursive clas- sification. ICLR, 2021.
Eysenbach, B., Zhang, T., Levine, S., and Salakhutdinov, R. Contrastive learning as goal-conditionel reinforcement learning. NeurlPS, 2022.
Allen, C., Parikh, N., Gottesman, O., and Konidaris, G. Learning Markov state abstractions for deep reinforce- ment learning. NeurlPS, 2021.
Fu, J.. Kumar, A., Nachum, O., Tucker, G., and Levine, S. D4RL: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.072/9 2020.
An, G., Moon, S., Kim, J., and Song, H. Uncertainty- based offline reinforcement learning with diversified Q- ensemble. NeurlPS, 2021.
Fujimoto, S. and Gu, S. A minimalist approach to offline reinforcement learning. NeurlPS, 2021.
Andrychowiez, M., Wolski, F., Ray, A., Schneider, J., Fong, R., Welinder, P., McGrew, B., Tobin, J., Pieter Abbeel, O., and Zaremba, W. Hindsight experience replay. NeurlPS, 30, 2017.
Fujimoto, S., Chang, W., Smith, E., Gu, S., Precup, D., and Meger, D. For SALE: State-action representation learning for deep reinforcement learning. NeurIPS, 2023.
Argenson, A. and Arnold, G. Model-based offline planning. ICLR, 2021.
Ghosh, D., Gupta, A., Reddy, A., Fu, J., Devin, C., Ey- senbach, B., and Levine, Learning to reach goals via iterated supervised learning. ICLR, 2021.
Bank, D., Koenigstein, N., and Giryes, R. Autoencoders. Machine learning for data science handbook: data min- ing and knowledge discovery handbook, pp. 353-374, 2023.
Ghosh, D., Bhateja, C., and Levine, S. Reinforcement learning from passive data via latent intentions. ICML, 2023.
10

---

Temporal Distance-aware Transition Augmentation for Offline Reinforcement Learning
Gupta, A., Kumar, V., Lynch, C., Levine, S., and Hausman, K. Relay policy learning: Solving long-horizon tasks via imitation and reinforcement learning. CoRL, 2019.
Lee, D., Yoo, M., Kim, W. K., Choi, W., and Woo, H. Incre- mental learning of retrievable skills for efficient continual task adaptation. NeurIPS, 2024b.
Hafner, D., Lillicrap, T., Fischer, I,, Villegas, R., Ha, D., Lee, H., and Davidson, J. Learning latent dynamics for planning from pixels. In ICMMI,, 2019.
Lee, J., Yun, S., Yun, T., and Park, J. GTA: Generative trajectory augmentation with guidance for offline rein- forcement learning. NeurIPS, 2024c.
Hejna, J., Gao, J., and Sadigh, D. Distance weighted su- pervised learning for offline interaction data. ICML., 2023.
Liu, B., Feng, Y., Liu, Q., and Stone, P. Metric residual net- work for sample efficient goal-conditioned reinforcement learning. In AAAI, 2023.
Jiang, Z., Zhang, T., Janner, M., Li, Y., Rocktäschel, T., Grefenstette, E., and Tian, Y. Efficient planning in compact latent action space. ICLR, 2022.
Lu, C., Ball, P., Teh, Y., and Holder, J. Synthetic experience replay. NearlPS, 2023.
Ma, Y., Jayaraman, D., and Bastani, O. Conservative offline distributional reinforcement learning. NeurlPS, 2021.
Kaelbling, L. Learning achieve goals. In IJCAL, 1993.
Kidambi, R., Rajeswaran, A., Netrapalli, P., and Joachims, T. MOReL: Model-based offline reinforcement learning. NeurIPS, 2020.
Ma, Y., Sodhani, S., Jayaraman, D., Bastani, O., Kumar, V., and Zhang, A. VIP: Towards universal visual reward and representation via value-implicit pre-training. ICLR, 2023.
Kingma, D. and Welling, M. Auto-encoding variational Bayes. arXiv preprint arXiv: 2013.
Mazoure, B.., Eysenbach, B., Nachum, O., and Tompson, J. Contrastive value learning: Implicit models for simple offline RL. In CoRL, 2023.
Kostrikov, I., Nair, A., and Levine, S. Offline reinforcement learning with implicit gelearning. ICLR, 2022.
Kumar, A., Fu, J., Soh, M., Tucker, G., and Levine, S. Stabilizing off-policy Q-learning via hootstrapping error reduction. NeurlPS, 2019.
Mees, O., Hermann, L., Rosete, E., and Burgard, W. CALVIN: benchmark for language-conditione pol- icy learning for long-horizon robot manipulation tasks. IEEE Robotics and Automation Letters, 733)732777734, 2022.
Kumar, A., Zhou, A., Tucker, G., and Levine, S. Con- servative Q-learning for offline reinforcement learning. NeurlPS, 2020.
Mendonca, R., Rybkin, O., Daniilidis, K., Hafner, D.,, and Pathak, D. Discovering and achieving goals via world models. NeurlPS, 34:24379-24391, 2021.
Laskin, M., Srinivas, A., and Abbeel, CURL: Contrastive unsupervised representations for reinforcement learning. In ICML, 2020.
Mitchell, E., Rafailov, R., Peng, X., Levine, S., and Finn,
Offline meta-reinfforeement learning with advantage weighting. In ICML, 2021.
Lecarpentier, E., Abel, D., Asadi, K., Jinnai, Y., Rachelson, E., and Littman, M. Lipschitz lifelong reinforcement learning. AAAI, 2021.
Mnih, V., Kavukcuoglu, K., Silver, D., Rusu, A., Veness, J., Bellemare, M., Graves, A., Riedmiller, M., Fidjeland,
Lee, B., Lee, J., and Kim, K. Representation balancing offline model-based reinforcement learning. In ICLR, 2021.
A., Ostrovski, G., Petersen, S., Beattie, C., Sadik, A., Antonoglou, I., King, H., Kumaran, D., Wierstra, D., Legg, S., and Hassabis, D. Human-level control through deep reinforcement learning. Nature, 518(7540):529-53: 2015.
Lee, D. and Kwon, M. Foresighted decisions for inter- vehicle interactions: An offline reinforcement learning approach. In /TSC, 2023.
Myers, V., Zheng, C., Dragan, A., Levine, S., and Eysen- bach, B. Learning temporal distances: Contrastive succes- sor features can provide metric structure for decision- making, ICML, 2024.
Lee, D. and Kwon, M. Episodic future thinking with offline reinforcement learning for autonomous driving. IEEE Internet of Things Journal. 2025.
Lee, D., Eom, C., and Kwon, M. AD4RL: Autonomous driving benchmarks for offline reinforcement learning with value-based dataset. In ICRA. IEEE, 2024a.
Nair, A., Gupta, A., Dalal, M., and Levine, S. AWAC: Accelcrating online reinforcement learning with offline datasets. arXiv preprint arXiv:2006.00 09359, 2020.

---

Temporal Distance-aware Transition Augmentation for Offline Reinforcement Learning
Nasiriany, S., Pong, V., Lin, S., and Levine, S. Planning with goal-conditioned policies. NeurIPS, 32, 2019.
Sinha, S., Mandlekar, A., and Garg, A. S4RL: Surprisingly simple self-supervision for offline reinforcement learning in robotics. In CoRL, 2022.
Nazari, P., Damrich, S., and Hamprecht, F. Geometric autoencoders-wiie you see what you decode. ICML, 2023.
Snell, C., Kostrikov, I., Su, Y., Yang, M., and Levine, S. Offline RL for natural language generation with implicit language Q learning, ICLR, 2023.
Newey, W. and Powell, J. Asymmetric least squares estima- tion and testing. Econometrica: Journal of the Economet- ric Society, pp. 819-847, 1987.
Sukhbaatar, S., Lin, Z., Kostrikov, I., Synnaeve, G., Szlam, A., and Fergus, R. Intrinsic motivation and automatic curricula via asymmetric self-play. ICLR, 2018.
Park, S., Ghosh, D., Eysenbach, B.,, and Levine, S. HIQI:: Offline goal-conditioned RL with latent states actions. NeurlPS, 36, 2023a.
Sun, Y., Zhang, J., Jia, C., Lin, H., Ye, J., and Yu, Y. Model- Bellman inconsistency for model-basee offline reinforce- ment learning. In ICML, 2023.
Park, S., Kreiman, T., and Levine, S. Foundation policies with Hilbert representations. ICML, 2024.
Tang, S., Makar, M., Sjoding, M., Doshi, F., and Wiens, J. Leveraging factored action spaces for efficient offline reinforcement learning in healthcare. NeurIPS, 2022.
Pertsch, K., Rybkin, O., Ebert, F., Zhou, S., Jayaraman, D., Finn, C., and Levine, S. Long-horizon visual planning with goal-conditioned hierarchical predictors. NeurIPS, 2020.
Tarasov, D., Kurenkov, V., Nikulin, A., and Kolesnikov, S. Revisiting the minimalist approach to offline reinforce- ment learning. NeurlPS, 2023.
Polyak, B. and Juditsky, A. Acceleration of stochastic ap- proximation by averaging. SIAM Journal on Control and Optimization, 3084))838-855 1992.
Tarasov, D., Nikulin, A., Akimov, D., Kurenkov, V., and Kolesnikov, S. CORL: Research-oriented deep offline reinforcement learning library. NeurlPS, 2024.
Pong, V., Gu, S., Dalal, M., and Levine, S. Temporal dif- ference models: Model-free deep RL for model-based control. ICLR, 2018.
Temnenholtz, G. and Mannor, S. Uncertainty estimation using Riemannian model dynamics for offline reinforce- ment learning. NeurIPS, 2022.
Rafailov, R., Hatch, K., Kolev, V., Martin, J., Phielipp, M., and Finn, C. MOTO: Offline pre-training online fine- tuning for model-basen robot learning. In CoRL, 2023.
Wang. J., Li, W., Jiang, H., Zhu, G., Li, S., and Zhang, Offline reinforcement learning with reverse model-based imagination. NeurlPS, 2021.
Reichlin, A., Vasco, M., Yin, H., and Kragic, D. Goal- conditioned offline reinforcement learning via metric learning. arXiv preprint arXiv:2402. 10820, 2024.
Wang. and Isola, On the learning and learnability of quasimetrics. ICLR, 2022.
Wang. T., Torralba, A., Isola, P., and Zhang, A. Opti- mal goal-reaching reinforcement learning via quasimetric learning. In ICML, 2023.
Rigter, M., Lacerda, B., and Hawes, N. RAMBO-RL: Ro- bust adversarial model- based offline reinforcement learn- ing. NeurlPS, 2022.
Wang, Z., Novikov, A., Zolna, K., Merel, J., Springenberg. J., Reed, S., Shahriari, B., Siegel, N., Gulcehre, C., Heess, N., and Freitas, N. Critic regularized regression. NeurIPS, 2020.
Ruder, An overview of gradient descent optimization algorithms. arXiv preprint arXhv:1609.04747, 2016.
Schaul, T., Horgan, D., Gregor, K., and Silver, D. Universal value function approximators. In ICML, 2015.
Wu, F., Zhang, R., Yi, Q., Gao, Y., Guo, J., Peng, S., Lan, S.. Han, H., Pan, Y., Yuan, K., et al. OCEAN-MBRL: Offline conservative exploration for model-based offline reinforcement learning. In AAAI, 2024.
Seo, Y. Lee, James, S., and Abbeel, P. Reinforcement learning with action-free pre-training from videos. In ICML, 2022.
Wu, Y., Tucker, G., and Nachum, O. Behavior regu- larized offline reinforcement learning. arXiv preprint arXiv: 1911. 1161,, 2019.
Shi, L., Lim, J., and Lee, Y. Skill-based model-based rein- forcement learning. CoRL, 2022.
Shin, S., Lee, D., Yoo, M., Kim, W. K., and Woo, H. One- shot imitation a non non-stationary environment via multi- modal skill. In /CML, 2023.
Xu, H., Jiang, L., Jianxiong, L., and Zhan, X. A policy- guided imitation approach for offline reinforcement leamn ing. NewrIPS, 2022.
12

---

Temporal Distance-aware Transition Augmentation for Offline Reinforcement Learning
Xu, H., Jiang, L., Li, J., Yang, Z.,, Wang, Z., Chan, V. W. K., and Zhan, X. Offline RI. with no OOD actions: In-sample learning via implicit value regularization. ICLR, 2023.
Yarats, D., Kostrikov, I., and Fergus, R. Image augmentation is all you need: Regularizing deep reinforcement learning from pixels. In ICLRR 2021.
Yu, T., Thomas, G., Yu, L., Ermon, S., Zou, J., Levine, S., Finn, C., and Ma, T. MOPO: Model-based offline policy optimization. NeurIPS, 2020.
Yu, T., Kumar, A., Rafailov, R., Rajeswaran, A., Levine, S., and Finn, C. COMBO: Conservative offline mosel--bsed policy optimization. NeurIPS, 2021.
Zhang, C., Kuppannagari, S., and Viktor, P. BRAC+: Im- proved behavior regularized actor critic for offline rein- forcement learning. In ACML, 2021.
Zhang, W., Xu, H., Niu, H., Cheng, P., Li, M., Zhang, H., Zhou, G., and Zhan, X. Discriminator-uuia guudde model- based offline imitation learning. In CoRL, 2023.
Zheng, C., Salakhutdinov, R., and Eysenbach, B. Con- trastive difference predictive coding. ICLR, 2024.
Zheng, R., Wang, X., Sun, Y., Ma, S., Zhao, J., Xu, H., Daumé, H., and Huang, TACO: Temporal latent action- driven contrastive loss for visual reinforcement learning. NeurlPS, 2023.
Zhou, W., Bajracharya, S., and Held, D. PLAS: Latent action space for offline reinforcement learning. In CoRL, 2021.
13

---

Temporal Distancc-aware Transition Augmentation for Offline Reinforcement Learning
A. Training Details
Convolutional encoder for the pixel-based task. For pixel-based environments, we handle image input using the IMPALA CNN architecture (Espeholt et al., 2018). The TempDATA unto-encoder uses the 112-dimensional output features of Impala CNN. In addition to the auto-encoder, our dynamic model and RL policy also use such compressed image features of IMPALA CNN. Therefore, we do not need to build raw image decoder model.
Goal-conditioned Reinforcement Learning. We train a skill-conditioned policy in an unsupervised manner to enable efficient learning in goal-achieving environmental settings. Moreover, this solution can be applied to multiple goal scenarios instead of fixed goal scenarios. We introduce a latent skill variable w, which is randomly sampled from p(w) inspired by DIYAN, and learn policy T 8, w) to explore and capture diverse behaviors from a reward-free offline dataset. This unsupervised phase leverages pre-trained encoder compute rewards and measure skill progress in the latent space. Specifically, we discussed in Section 4.3, we define an intrinsic reward, F(s, s') - d(f(s';)), [[ggoaa; 0)) d({(;; 0): (*golll 0)). Next, we implement our offline RI. training using the IQI. combined with AWR objective function, ensuring stable policy improvement in an in-sample learning manner. Our overall pipeline integrated within the TempDATA codebase, which extends HIQL implementations (Park et al., 2023a). However, we do not consider hierarchical policy structure in this work.
Policy execution. At the test phase, the trained policy requires a specific skill variable according to the given goal state. To do this, we adopt the test-time skill adaptation, which is proposed in HILP (Park et al., 2024), as follows:
{Sgooll 0) (s:
w²
[ 0) - [s; 0)
which guides the policy to decrease the distance to (S 0) by moving the latent direction f( 0) - f(; This approach requires no additional training and can be performed in zero-shot manner.
Goal labeling. To train autoencoder and offline RL networks, relabel goal distributions using hindsight experience relabeling (Andrychowicz al., 2017). Specifically, for randomly selected state, assign it the goal with 20% probability, future state from the trajectory with a 50% probability, and completely random state with a 30% probability. The reward is set to -1 for every timestep until the goal is achieved, at which point the done mask is set to True. During the encoding process, we ensure that at least one of the samples contains the goal state.
Procedures:
Model rollout. We adopt a standard model rollout approach (Yu et al., 2020) which a fully trained dynamics model is used to simulate additional trajectories during policy learning. Specifically, the agent performs rollouts of length steps in the learned model to generate synthetic state-action-rewart transitions that complement the original offline dataset. To ensure stable and well groundeoddd initialization, we train the policy for the first 30% of the total training duration without any rollouts. This phase provides the agent with direct understanding of the environment's dynamics before synthetic samples are introduced. Subsequently, at every 10% iteration of the remaining training period, we generate synthetic samples equivalent to half the current size of the offline buffer and add them to the training pool. By gradually incorporating godeelggeerrtte data, we allow the agent to refine its decision- decision-making while mitigating potential inaccuracies from extensive model rollouts.
Procedures:
14

---

Temporal Distance-aware Transition Augmentation for Offline Reinforcement Learning
B. Proof of Theorem 4.2
Proof. We use a Bellman operator with expectile regression to learn ds(s, 8gont). Concretely, each update enforces
(1];;),, ($gaal 8)) = + 7 min d(f(s'; f(goal: 0)),
where is the discount factor, and (s, a, s are transitions sampled from the dataset or replay buffer. The loss itself is formulated via expectile regression, which replaces the usual squared Bellman error with 7-expectile operator.
property expectile regression states that if X is a bounded random variable with supremum x", then
lim m, (X) = x=;
where mx is the r-expectile of X. Translating this to our Bellman setting, driving toward forces the learned distance d, to match the supremum of returns along feasible trajectories.
In a deterministic environment, the total cost from to Sgoal along any path is at least as large as the shortest-path cost. From standard arguments, we have:
do(s, Sgoal) S (shortest-pan path cost from to Sgoal)
V (s, 8goal).
Meanwhile, the Bellman- like constraints in our training also ensure that da cannot collapse below the maximum relevant cost, since the expectile objective pushes toward the worst-case outcome ->>
Because do(s, 8goal) is bounded both above and below by the same shortest-path - -V* cost in the limit, it converges exactly to
do(s, Sgoal) -V*(ss Sgoal).
Equivalently,
lim (), r-1
max V* Sgoal);
completing the proof.
15

---

Temporal Distance-aware Transition Augmentation for Offline Reinforcement Learning
C. Implementation Details
Our implementation for TempDATA based on JaxRL and is available at the following repository. We run our experiments on RTX 3090 GPUs. Each experiment runs for no more than hours in state-based environments and no more than 12 hours in pixel-based environments.
C.1. Environments
www.u--mizzeeeeeeeeeeeee play-v2', gaatmaze-lare-diveerr play-v2', and autmaaz-uutta-ddieess play-v0'. The datasets for the 'umaze', 'medium', and 'large' scale environments originate from the D4RL benchmark (Fu et al., 2020), while the datasets for AntMaze-ultra are sourced from separate work. Notably, the AntMaze-ultra environment (Jiang et al., 2022) is twice as large as AntMaze-large Each dataset comprises 999 trajectories, each with length of 1000 steps, where the Ant agent moves from a randomly chosen starting position to a goal location, which is not necessarily the evaluation target. During testing, to define goal for the policy, we modify the first two state dimensions correspondigg to x-y coordinates-to the designated target position the environment, while the remaining proprioceptive state variables are to those from the first observation in the dataset. The agent receives a reward of upon successfully reaching the goal during evaluation.
AntMaze. We examine eight distinct scenarios within the AntMaze task: uantaaze-umaze-vee cetaazze-mdiuum-dfe
The Kitchen datasets from D4RL (Fu et al., 2020) consists of two datasets 'kitchen- Kitthen-(partial, mixed)-vO'. These datasets capture trajectories of a robotic arm interacting with various objects in different sequences within its environment. For this task, the agent gets a reward of upon completing each subtask, with each episode comprising four subtasks. In addition to state-based work, for pixel-based Kitchen experiments, we transform each state into a 64 x 64 x camera image through rendering, employing the same camera configuration used by Mendonca al. and Park al..
Kitchen.
CALVIN. The CALVIN offline datasets are rooted in the teleoperated demonstrations by Mees et al. and, thereby being introduced by Shi al. and Park et al.. This dataset comprises 499 trajectories; each trajectory has 1204 transitions, capturing various subtasks performed in an arbitrary order. Similar to the FrankaKitchen task, the CALVIN task consists of four subtasks and the agent obtains a reward of upon completing each subtask.
C.2. Hyperparameters
Hyperparameter
Value
Iterations
106 (state-based), x 10 (pixel-based)
Learning rate
3 x 10 (all networks)
Optimizer
Adam (Diederik, 2015) 512 (AntMaze), 256 (FrankaKitchen), 128 (CALVIN) 50 (all tasks)
Batch size
The number of evaluation episodes
[512, 512, 512, {32, 10), 512, 512] 32 (AntMaze), 10 (CALVIN, FrankaKitchen) 0.99 (all tasks)
Dimensions for autoencoder network
Discount factor for autoencoder
Expectile coefficient for Autoencoder
0.95 (AntMaze), 0.97 (CALVIN, FrankaKitchen), 0.7 (Pixcl-based)
Dimensions for dynamic model network The number of rollout steps
[122, 512,
3
Dimensions for critic network Dimensions for actor network
[512, 512, 512) [512, 512, 512)
Target smoothing coefficient Discount factor for offline RI. Inverse temperature for offline RL Expectile coefficient for offline RL
10
0.99
10 (AntMaze), 3 (CALVIN, FrankaKitchen)
0.9 basedd, 0.7 (pixel-basee based)
16


Reviews:
[{"review_id": "Klmo41nWWQ", "reviewer": "Reviewer_p9eS", "summary": "This paper presents an offline model-based reinforcement learning algorithm called TempDATA. The algorithm aims to tackle goal-conditioned tasks with long horizon and sparse task-completion reward.\n\n----\nThe main idea is to learn an embedding space which that enables the computation of a temporal distance measure between pairs of states. The temporal distance allows the authors to perform reward shaping to decrease the distance to the goal.\n\nAside from reconstruction, the latent space is trained to also (1) satisfy temporal difference in the same trajectory, and (2) maintain proximity of consecutive states (Section 4.1). Next, the dynamics model learns to perform forward prediction in the latent space, and the synthetic rollouts are combined with the offline dataset for policy training.\n\n----\n\nThe experiments consist of comparisons with various offline RL methods on a few domains, including state-based AntMaze, FrankaKitchen, Calvin and pixel-based Kitchen. The quantitative results demonstrate that (1) TempDATA outperform existing MBRL and TARL methods in the single-goal AntMaze environment. (2) TempDATA outperforms existing goal-conditioned offline RL methods in the multi-task kitchen and calvin environments. (3) TempDATA achieves similar performance as the state-of-the-art model-free offline RL method IQL.\n\n----\n\n## Update after Rebuttal\n\nI appreciate the authors for sharing additional experiment results. The ablation study and new results in respond to other reviewers are sound.\n\nI raised my score from 1 to 2 in recognition of the empirical results and the authors' explanations on the math formulations under my and reviewer DPTt's reviews. I did not raise my score further because I'm concerned that the paper might need substantial edits to clarify the math formulation.", "questions": "", "limitations": "", "rating": 2, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's valuable feedback.\n\n---\n# Main contribution and reward shaping\nWe agree with the reviewer’s observation that TempDATA involves multiple components and that it is important to delineate our main contribution clearly. Our main contribution lies in temporal-distance-aware autoencoder. This encoder serves as a unified representation module within our framework, enabling effective policy learning and data augmentation through temporal-distance-aware representations.\n\nRegarding the reward-shaping component, we clarify that our strategy follows widely used practices in skill-conditioned RL settings [1, 2, 3]. Distance-based reward shaping from encoded representations of state transitions is general. Therefore, our main contribution is not the reward shaping method itself but the accurate encoding of temporal distances between states, which empirically demonstrates a substantial improvement.\n||Laplacian|Contrastive learning|Random|TempDATA|\n|-|-|-|-|-|\n|antmaze-large-play|$0\\pm0$|$0\\pm0$|$0\\pm0$| $56.5\\pm14.1$|\n|antmaze-large-medium|$0\\pm0$| $0\\pm0$|$0\\pm0$|$44.2\\pm15.3$|\n|kitchen-partial|$42.9\\pm10.3$|$55.5\\pm10.6$|$44.0\\pm6.4$|$70.0\\pm12.8$|\n|kitchen-mixed|$51.1\\pm12.0$|$56.2\\pm9.0$|$47.5\\pm8.5$|$65.3\\pm11.7$|\n\nTo further substantiate the effectiveness of our solution, we ran additional experiments comparing TempDATA’s encoder with other representation-learning approaches. Specifically, we learn successor features with three alternative feature learners [4, 5] (Laplacian, Contrastive Learning, and Random Feature) as baselines, and compared their performance against TempDATA. The results demonstrate that our solution significantly outperforms alternative feature learners.\n\n---\n# Generalizability\nAlthough our primary focus has been long-horizon tasks, we acknowledge the reviewer’s valid concern regarding generalization to broader task settings. To address the reviewer's concern, we ran additional experiments in dense rewarded task of D4RL. Our additional experiments demonstrate that the **proposed solution achieves performance comparable to widely used D4RL benchmark** (full details provided in response to Reviewer Akpp). Furthermore, we ran additional image-based experiments (see our detailed response to Reviewer DPTt), confirming that TempDATA performs well in visual domain, further supporting its generalizability.\n\n---\n# Ablation study\nTo alleviate the reviewer's concern, we conducted additional ablation experiments on AntMaze datasets with the following variants:\n- Skill-conditioned RL: without representation learning and dynamics learning\n- Skill-conditioned RL with MOPO: without representation learning\n- TempDATA with IQL: without reward engineering and skill learning\n| |Skill RL|Skill RL w MOPO|TempDATA w IQL|Full TempDATA|\n|-|-|-|-|-|\n|medium-paly|$66.5\\pm14.6$|$21.9\\pm13.2$|$76.0\\pm10.4$|$74.8\\pm8.3$|\n|medium-diverse|$65.9\\pm11.2$|$26.0\\pm10.8$|$65.5\\pm8.9$|$69.5\\pm10.8$|\n|large-play|$55.5\\pm14.3$|$6.1\\pm3.9$|$44.4\\pm12.1$|$56.5\\pm14.1$|\n|large-diverse|$49.0\\pm18.8$|$9.3\\pm6.0$|$47.5\\pm9.5$|$44.2\\pm15.3$|\n\nThese results highlight the substantial contribution of the representation learning component, as performance noticeably decreases when it is removed (Skill RL with MOPO). Furthermore, even without reward engineering (TempDATA with IQL), the method maintains performance comparable to the full TempDATA, suggesting that our representation learning module plays a central role in TempDATA’s effectiveness.\n\n---\n# Other questions\n1. **Unit for wall clock time**.  We set the unit for wall clock time as hours in Figure 6 (b)\n2. **How to define $d$ function**. We used Euclidean norm to calculate the distance between $z$.\n3. **Pixel-based outperforms state-based**. In our experiments, GC-IQL with pixel-based tasks occasionally outperforms state-based ones. We believe this may be attributed to the strong performance of the IMPALA encoder [6], which can extract richer visual features that are not captured by low-dimensional state representations. We have also observed a similar experience in [7]; GCIQL, HIQL, and FQL are sometimes better in pixel-based than state-based tasks.\n\nWe will update all references. We would be happy to continue the discussion if you have any other questions or comments that could raise your score.\n\n---\n# References\n[1] S. Park, et al. Lipschitz-constrained unsupervised skill discovery. ICLR 2022.\n\n[2] K. Frans, et al. Unsupervised zero-shot reinforcement learning via functional reward encodings. ICML 2024.\n\n[3] R. Yang, et al. Behavior contrastive learning for unsupervised skill discovery. ICML 2023.\n\n[4] C. Zheng, et al. Contrastive difference predictive coding. ICLR 2024.\n\n[5] A. Touati et al. Does zero-shot reinforcement learning exist?. ICML 2023.\n\n[6] L. Espeholt, et al. IMPALA: Scalable distributed deep-RL with importance weighted actor-learner architectures. ICML 2018.\n\n[7] S. Park, et al. OGbench: Benchmarking offline goal-conditioned RL. ICLR 2025.", "claims_and_evidence": "\"The latent dynamics model alleviates the overgeneralization issue with MBRL, and is more efficient than high-dimensional state space\": this is only partially supported by the experiments. Yes, the offline RL results with the AntMaze is better than relevant baselines, but it is unclear whether the performance gain comes from synthetic transitions or the reward engineering trick. Moreover, the authors use a skill-conditioned policy. It is unclear how much this helped with training.", "methods_and_evaluation_criteria": "The evaluation domains are diverse. However, only one domain is used to test the main offline MBRL contribution.", "theoretical_claims": "I checked the arguments in section 4.1 as well as in the appendix. The mathematical derivations are sloppy:\n- In equation (1), $V = -d$ which makes sense because the value is higher the closer to the goal. However, in Theorem 4.2, the minus sign is dropped.\n- In equation (2), the second row should also have a max over $\\pi$.\n- In $\\mathcal{L}_{traj}$, the loss has $\\mathcal{B}d + d$ but should be $\\mathcal{B}d - d$.", "experimental_designs_or_analyses": "The choice of baseline methods for experimenting with single-goal, multi-goal and image-based settings are sound. However, the strongest offline MBRL results are only shown in the AntMaze environment. The other two settings don't show a clear advantage over the existing SOTA.\n\nA key issue is that the algorithm has a lot of moving parts, including representation learning, dynamics learning, reward engineering and skill learning. Ablation studies are definitely needed to justify these design choices and understand their individual contributions.", "relation_to_broader_scientific_literature": "Offline RL and MBRL are hot fields. A new SOTA offline MBRL approach could be impactful. Additionally, a good pixel-based offline MBRL method could be a good baseline.", "essential_references_not_discussed": "[Embed to Control: A Locally Linear Latent Dynamics Model for Control from Raw Images (NeurIPS 2015)] and [Accelerating Visual Sparse-Reward Learning with Latent Nearest-Demonstration-Guided Explorations (CoRL 2024)]: prior works have explored learning an embedding space and dynamics model to obtain a distance measure for reward shaping, although not in the offline MBRL setting.", "other_strengths_and_weaknesses": "- The visualization in Figure 3 is helpful to justify the representation learning.", "other_comments_or_suggestions": "The authors should consider making their core contributions more focused. TempDATA involves a lot of moving parts. For example, if the reward shaping  aspect is the most important/beneficial, they should compare with other value-based reward shaping methods.", "questions_for_authors": "- What is the unit for wall clock time in Figure 6(b)?\n- Is $d$ just the L2 distance between $z$ vectors?\n- How is GC-IQL performance better in pixel-based compared to state-based?", "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": "Cz0kqG1OTN", "reviewer": "Reviewer_Akpp", "summary": "This paper addresses the challenges of offline model-based reinforcement learning, particularly in sparse reward and long-horizon environments. The authors propose Temporal Distance-Aware Transition Augmentation (TempDATA), a novel method that generates additional transitions in a geometrically structured representation space rather than the state space. By learning state abstraction that captures temporal distance at both trajectory and transition levels, TempDATA enhances the ability to comprehend long-horizon behaviors efficiently. Experimental results demonstrate that TempDATA outperforms previous offline MBRL methods and achieves comparable or superior performance to diffusion-based trajectory augmentation and goal-conditioned RL across multiple benchmark environments, including D4RL AntMaze, FrankaKitchen, CALVIN, and pixel-based FrankaKitchen.\n\n## update after rebuttal\nThe additional experiments on locomotion tasks improve the quality of this work. Currently, I have no additional questions. I would like to maintain my original rating, considering the overall algorithm novelty and contribution.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are grateful for the reviewer's thorough review and valuable suggestions about this work. Below, we outline how we have revised the paper to address the reviewer's concerns and questions.\n\n-----\n# Generalizability\nWe acknowledge the importance of demonstrating applicability beyond long-horizon tasks. Following the reviewer's helpful suggestion, we ran additional experiments on dense-reward tasks from the widely recognized D4RL benchmark (e.g., halfcheetah and walker2d), achieving competitive performance as detailed below:\n\n| |MOPO (MBRL) |GTA (TARL) |IQL (MFRL) |TempDATA|\n|-|-|-|-|-|\n|halfcheetah-medium-replay|$68.2\\pm20.8$|$50.0\\pm0.8$|$43.4\\pm0.5$|$62.8\\pm8.9$|\n|halfcheetah-medium-expert|$63.3\\pm38.0$|$93.1\\pm3.1$ |$94.6\\pm1.9$|$95.8\\pm2.5$|\n|walker2d-medium-replay|$69.4\\pm18.8$|$93.8\\pm1.7$|$69.6\\pm10.8$|$91.8\\pm1.9$|\n|walker2d-medium-expert|$44.6\\pm12.9$|$110.9\\pm0.3$|$105.2\\pm4.9$|$111.8\\pm1.1$|\n\nThese experiments demonstrate that the **proposed solution has competitive performance in widely used D4RL tasks** [1]. It further confirms our solution's effectiveness across various task types and reward structures.\n\nAdditionally, given that our primary focus is on sparse-rewarded long-horizon tasks, we have presented strong empirical results across a diverse set of environments, e.g., AntMaze, FrankaKitchen, and CALVIN, highlighting our method's capabilities in subgoal navigation, manipulation, and decomposition tasks.\n\n-----\nWe hope this response addresses your suggestion. Please let us know if any additional clarifications are required.\n\n# References\n[1] J. Fu, et al. D4RL: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219 (2020).", "claims_and_evidence": "The algorithm design is validated by both theoretical and empirical analysis.", "methods_and_evaluation_criteria": "The methods and evaluation metric are motivating.", "theoretical_claims": "I checked the proofs presented in the appendix.", "experimental_designs_or_analyses": "The experiments in four different goal-achieving tasks", "relation_to_broader_scientific_literature": "The paper is related to model-based trajectory augmentation in offline rl.", "essential_references_not_discussed": "No missing essential references.", "other_strengths_and_weaknesses": "The theoretical proofs enhance the soundness of the proposed method, and the method design is motivating itself. However, the experiments are limited to only 4 goal-achieving tasks.", "other_comments_or_suggestions": "See the Questions section.", "questions_for_authors": "1. The current experiments are conducted in four goal-achieving tasks. Is the proposed method applicable to non-goal-achieving tasks, such as locomotion tasks?", "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": "vFvwf3Jmr6", "reviewer": "Reviewer_DPTt", "summary": "This paper proposes TempDATA, a new offline model-based reinforcement learning (MBRL) method that learns a temporal-distance-aware autoencoder model, a latent dynamic model, and an offline policy. Specifically, TempDATA first trains the autoencoder and latent dynamic model to possess the temporal-distance attribute. Then, it generates augmented rollouts using the learned autoencoder and latent dynamic model. Lastly, it learns the offline policy by leveraging both original and augmented transitions. The proposed TempDATA is evaluated on tasks from multiple benchmarks, including AntMaze, FrankaKitchen, and CALVIN. From the results, the paper concludes that TempDATA achieves competitive or improved performance compared to various baselines.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the reviewer's time and effort. Here are our answers to the reviewer's comments.\n\n---\n# Performance according to dataset\nWe agree with the concern raised by the reviewer. To address the concern, we ran additional experiments by gradually reducing the offline dataset coverage for the antmaze-medium tasks.\n\n| | 100%| 80%| 60%| 40%| 20%|\n|-|-|-|-|-|-|\n| antmaze-medium-play| $74.8\\pm8.3$  | $78.6\\pm10.0$ | $65.5\\pm15.5$ | $41.4\\pm6.6$ | $11.9\\pm5.0$ |\n| antmaze-medium-diverse | $69.5\\pm10.8$ | $69.1\\pm9.3$  | $49.8\\pm18.2$ | $34.8\\pm8.2$ | $18.0\\pm7.6$ |\n\nWhile performance drops significantly when the dataset is extremely limited at $20$% coverage, experimental results demonstrate that performance remains relatively robust with up to $60$% coverage. These results indicate that the representation can still be learned effectively under moderate data scarcity. To sum up, TempDATA is robust under moderate dataset scarcity, but there is room for improvement under severe limitations or bias in the dataset.\n\n---\n# TempDATA with pixel-based tasks\nTo verify the competence of TempDATA in the pixel-based RL, we ran additional experiments on a pixel-based benchmark: Visual AntMaze. Visual AntMaze is a variant of AntMaze with camera images and proprioceptive states, requiring the agent to navigate based on visual cues. This dataset is rendered by using D4RL antmaze-medium and -large datasets. \n\n| visual-antmaze dataset | RepB-SDE | GC-IQL| TempDATA|\n|-|-|-|-|\n| medium-play| $0\\pm0$  | $58.8\\pm12.0$ | $60.8\\pm11.6$ |\n| medium-diverse| $0\\pm0$  | $65.6\\pm19.8$ | $62.0\\pm12.2$ |\n| large-play| $0\\pm0$  | $35.7\\pm10.2$ | $42.8\\pm9.7$  |\n| large-diverse| $0\\pm0$  | $29.0\\pm8.8$  | $40.8\\pm7.0$  |\n \nThese results confirm that TempDATA effectively handles pixel-based tasks, achieving notable performance compared to GC-IQL and RepB-SDE. **Despite challenges inherent in augmenting visual data augmentation, our solution underscores its competence unlike other MBRL solution**. \n\n---\n# Reward for skill-conditioned RL\nWe appreciate the reviewer’s careful observation regarding the reward design $r_g(s) - 1$, and we acknowledge that our wording causes confusion between the terms ``vanilla`` and ``general``. Specifically, the vanilla reward refers to the binary scenario ($r_g(s) = 1$ for goal states, $r_g(s) = 0$  otherwise). our footnote inadvertently associated this definition with the general method (subtracting $1$ from $r_g(s)$).\n\nTo clarify, the vanilla reward yields sparse signals, hindering effective learning of long-horizon tasks. Therefore, we adopted the general alternative $r_g(s) - 1$, as it provides clearer temporal information and better learning signals [1, 2].\n\nIn our framework, this transformation plays a crucial role during the pre-training stage, where the agent learns a temporally meaningful value function. By consistently assigning a penalty $(-1)$ for non-goal states and 0 for goal states, **we enable the Q-function to represent a form of temporal distance to the goal, as $Q(s, a) \\approx -d(s, g)$, where $d(s, g)$ is the minimum number of steps to reach the goal from state $s$**. This mathematical property allows our Q-function to encode structured temporal knowledge, which is especially valuable in downstream tasks requiring generalization over time and space.\n\n---\n# Challenge of MBRL\nWe fully agree that the coverage of the offline dataset significantly impacts the performance of offline MBRL methods, particularly regarding their ability to generalize to OOD states and actions. However, the original statement aimed at comparing offline MBRL with MFRL methods under offline datasets in dense-rewarded or short-horizon robotic tasks. Even under an identical offline dataset, MFRL methods rely exclusively on observed trajectories and thus inherently lack the inductive biases that enable generalization beyond the seen data. In contrast, offline MBRL can leverage learned dynamics to partially alleviate mild OOD challenges through structured forward prediction. Nevertheless, we acknowledge that this statement was overly optimistic and could be misleading. We will make efforts to tone down the over-claim about MBRL.\n \n---\nThank you for your helpful comments on the clarity and correctness of our writing. We will revise our manuscript and include additional information, e.g., IQM and optimality Gap.\n\nPlease let us know if you have any additional concerns or questions.\n\n# Reference\n[1] A. Kumar, et al. Conservative Q-learning for offline reinforcement learning. NeurIPS 2020.\n\n[2] K. Frans, et al. Unsupervised zero-shot reinforcement learning via functional reward encodings. ICML 2024.", "claims_and_evidence": "The main claim made in this paper is that the policy learns with augmented transitions, generated by the temporal-distance-aware system (autoencoder + latent dynamic models), and can achieve better results.\n\n[+] From the theoretical aspect, the proposition is sound, though it has some confusing parts. From the empirical aspect, the proposed method significantly outperforms baselines in state-based environments, and the visualization result (Figure 3) further supports this claim qualitatively.\n\n[-] However, I believe there is a prerequisite for the paper to obtain this result: the offline dataset itself must cover a sufficient range of transitions to effectively build/approximate the dynamic model. More details will be summarized below.", "methods_and_evaluation_criteria": "[+] While the components in TempDATA, namely temporal-distance-aware latents and the latent dynamic model, may have been explored individually in previous work, their combination and usage in MBRL make sense to me. The motivation to introduce the temporal-distance-aware autoencoder is also supported by theoretical backing.\n\n[+] The proposed method is evaluated on benchmarks across various domains, such as robot arm and maze navigation. These benchmarks include both state-based and image-based environments.\n\n[-] Some evaluation metrics lack descriptions; for instance, the IQM and Optimality gap in Figure 6 (c).", "theoretical_claims": "[+] As mentioned above, Proposition 4.1 and its proof are sound in general.\n\n[-] It is unclear why the reward function is designed as $r_g(s)-1$. Especially in line 225, it is specifically noted that this design differs from the vanilla goal-conditioned reward function, but the associated footnote indicates, \"General methods use transitions by subtracting 1 from $r_g(s)$.\" Isn't this exactly $r_g(s)-1$? Why is the reward function designed this way? Is there any physical meaning to support it?\n\n[-] While the proposition is sound, I believe that in practice, how well the offline dataset covers the space of possible states and actions will significantly affect the method's performance—i.e., how well the dynamic/world model can predict the observation/state/latent for the next time step. \n\nI believe this is a common challenge within the MBRL approach, which is why I am skeptical of the argument in lines 42-46 (right) of the paper, as quoted: \"*The offline MBRL methods have covered OOD samples efficiently, achieving a better performance than offline MFRL in dense rewarded or short-horizon robotic manipulation tasks.*\" \n\nIf there is no such state-action pair in the offline dataset, it is inherently impossible for the dynamic model to learn and make correct predictions.", "experimental_designs_or_analyses": "[+] The results for each environment are tested over multiple rounds, and the standard deviation of the method’s performance is also reported.\n\n[+] TempDATA is compared with several baseline methods, showing a significant performance improvement in state-based environments.\n\n[-] In image-based environments, TempDATA achieves sub-optimal performance, raising concerns about its ability to learn the dynamic model when working with data that contains only partial information—something commonly encountered in practice. The explanation for why TempDATA performs sub-optimally is insufficient. Additional experiments in image-based environments, along with a more detailed analysis, could help address this concern.\n\n[-] As mentioned earlier, I believe the proposed method outperforms MFRL methods because the offline dataset provides enough state-action coverage to learn an effective dynamic model. However, when the offline dataset is small or lacks data in critical areas, I feel that MBRL methods will be significantly impacted. Conducting experiments on this aspect is essential to alleviate this concern.", "relation_to_broader_scientific_literature": "[+] In my view, offline RL is a challenging yet valuable problem, especially for real-world applications where data collection is costly, risky, or impractical. Advancing offline RL methods has the potential to significantly impact related fields, such as robotic learning.", "essential_references_not_discussed": "[+] The paper discusses and cites sufficient works. The literature review is well summarized in the related work section. Additionally, the proposed method is compared against baselines of various types and attributes.", "other_strengths_and_weaknesses": "[+] The training details provided in the supplementary materials improve the reproducibility of this work.\n\n[-] While the paper is generally easy to follow, I found the following points that need correction or revision:\n- [line 107 (right)]: \"image-based state\" → From a rigorous perspective, an image is an observation of a state that only contains partial information, rather than a type of state.\n- [line 113 (right)]: \"or state-action value function $V(s)$\" → Should this be \"state value function\" ?\n- [line 116 (right)]: $\\arg  \\underset{a' \\sim \\pi(s')}{\\max} \\, Q(s', a')$ -> $\\underset{a' \\sim \\pi(s')}{\\arg \\max} \\, Q(s', a')$\n- [line 221]: The parameters required by the function $d$ are inconsistent between Eq. 1 and line 221.\n- [line 252 (right)]: In the formula, $ds$ and $ds^{'}$ are usually placed at the end, not at the beginning.", "other_comments_or_suggestions": "All my comments and suggestions are listed in the appropriate fields above. Regarding my recommendation, I actually feel the paper is around the threshold; I would choose \"borderline\" if there were such an option. Based on the observed pros and cons, I will set my initial score as \"weak accept.\" I would be happy to further adjust it if my concerns are well-addressed or clarified.\n\n---\n\n**Comments after author-reviewer discussion**\n\nThank you for providing additional experiments on data scarcity. As promised, I will increase my score to \"accept.\" That said, if the paper is accepted, please make sure to incorporate these updates and revisions into the camera-ready version. I also agree with Reviewer p9eS that the writing in the method section—especially the mathematical formulations—can be further improved. For future work, it would be valuable to evaluate the proposed TempDATA on long-horizon tasks and assess whether it still demonstrates effectiveness in estimating the temporal distance to task completion.", "questions_for_authors": "Please consider addressing the concerns listed in the previous sections. I will highlight some important ones here:\n\n1. How much does the performance of the dynamic model influence the proposed method? Specifically, if the offline dataset is limited or biased, what impact does this have on the method’s effectiveness?\n2. What explains the sub-optimal performance of the proposed method in image-based environments?\n3. Could you provide further clarification regarding the design of the reward function $r_g(s)$?", "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": "4ApdVa6cd7", "reviewer": "Reviewer_rU4d", "summary": "This paper focuses on the failure of model-based reinforcement learning (MBRL) in sparse-reward and long-horizon environments, emphasizing that the key to addressing this issue lies in generating data that incorporates temporal information. To tackle this challenge, this paper introduces a novel MBRL framework, TempDATA, which consists of three components: an autoencoder, a latent dynamic model, and an offline policy. The autoencoder is trained to learn state abstractions that capture temporal distances at both the trajectory and transition levels in a representative space and then reconstruct these abstractions in the original state space. The latent dynamic model is trained to augment the dataset with temporal distance-aware transition, while the offline policy is extracted from the augmented dataset.\n\n## update after rebuttal\n\nThank the authors for the detailed response, and I have reviewed the additional comments. The results of these two experiments are truly impressive. The first experiment provides excellent insights into the alignment between distances in the encoded representation space and obstacle-induced temporal distances, while the second experiment further offers quantitative evidence to support this point. I have no additional questions. Given the innovation and thoroughness of the paper, I believe a score of 3 is appropriate, and I will maintain this rating.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We thank the reviewer for the positive reviews and insightful feedback about this work. Below, we describe how we have revised the paper to address the reviewer's concerns and questions.\n\n---\n# Experimental coverage of baseline algorithms\nWe thank the reviewer for pointing this out. In our work, we designed the experimental evaluation around three types of tasks: \n1) **AntMaze**, which is a single-goal long-horizon task\n2) **CALVIN and FrankaKitchen**, which require sequential multi-goal manipulation (sub-goal decomposition tasks)\n3) **Visual FrankaKitchen**, which is based on image-rendered dataset version of \bFrankaKitchen (pixel-based goal-conditioned tasks)\n\nWe mainly compare with offline MBRL and TARL baselines on the AntMaze task, where such model-based approaches have been previously studied extensively. However, for the sub-goal decomposition tasks (e.g., CALVIN and FrankaKitchen) and pixel-based tasks, **prior offline MBRL and TARL methods have not been explored** in the existing literature. \n\nTo investigate their applicability further, we ran above-mentioned experiments by applying representative offline MBRL and TARL baselines to these tasks. However, **we observed that the resulting performance was consistently near-zero**, with success rates failing to exceed random behavior. We attribute this to the severe sparsity of rewards and the difficulty of long-horizon goal decomposition in these environments.\n\nGiven these empirical observations, **we concluded that such methods do not serve as meaningful baselines in these domains, and instead chose to compare against goal-conditioned model-free RL (GCRL) methods**, which are more capable of handling these challenges. We will explicitly highlight this reasoning in the revised manuscript to avoid confusion.\n\n---\n# Evidence of temporal representation\n\nWe appreciate the reviewer’s suggestion. We would like to clarify that Figure 3 in our manuscript directly addresses this concern; Reviewer p9eS and DPTt agree that Figure 3 justifies representation encoding qualitatively. The figure provides an empirical visualization of the learned latent space using t-SNE embeddings, which demonstrates that **the trained encoder maps states according to their temporal distance rather than their spatial similarity**.\n\nSpecifically, if the encoder fails to capture temporal distance, the t-SNE projection results in an entangled and structureless cloud. In contrast, our result shows a meaningful geometric progression, where temporally distant states are mapped further apart, confirming that the encoder preserves temporal information. We strongly believe that the provided visualization qualitatively validates the temporal property of the embedding. \bWe will explicitly emphasize and clarify this qualitative validation in the revised manuscript to strengthen our claims further. \n\n---\nPlease let us know if there are any additional concerns or questions.", "claims_and_evidence": "The claims made in the paper are supported by clear evidence.", "methods_and_evaluation_criteria": "It makes sense for problem.", "theoretical_claims": "I’ve checked the correctness of the theorem proof mentioned in the main paper.", "experimental_designs_or_analyses": "There’s a doubt about the experiment. The TempDATA is tested in four datasets: AntMaze, Kitchen, Calvin, and Visual Kitchen with thirteen baselines. However, some baselines were only tested for part of the four datasets. I suppose that the reason for this experimental design isn’t mentioned in the paper. Additionally, while I acknowledge the superior performance of TempDATA across various benchmarks, the paper would be even stronger if additional experiments were included to demonstrate the autoencoder’s ability to capture temporal information.", "relation_to_broader_scientific_literature": "The paper is well-grounded in the broader literature.", "essential_references_not_discussed": "No", "other_strengths_and_weaknesses": "I appreciate the effort put into providing an intuitive understanding of TempDATA. These illustrative figures were helpful in clarifying the key points of the paper. However, the writing of this paper could be improved. Spelling errors and grammatical mistakes hinder the reading experience and make some of the claims unclear at first glance. Therefore, the author should revise the paper to enhance its clarity.", "other_comments_or_suggestions": "No", "questions_for_authors": "See the experimental designs or analyses part.", "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:

Meta-review: This paper was discussed controversially and even after rebuttal and discussion the assessments ranged from Weak Reject to Accept. The rebuttal was rated as very strong and the results of the two additional experiments are really impressive.

The criticism remains that the initial submission did not reach this level of quality and that extensive changes were necessary and were made, but can no longer be assessed by the reviewers, and that substantial edits are needed to clarify the math formulation.
Decision: Accept (poster)