Temporal Distance-aware Transition Augmentation for Offline Model-based Reinforcement Learning

• Offline RL must learn from fixed datasets while handling OOD actions that can corrupt state-action value estimation.Source: §1 · p. 1
• Offline MBRL synthesizes transitions with learned dynamics and behavior policies, but its prior successes are concentrated in dense-reward, short-horizon locomotion tasks.Source: §1 · p. 1
• Most offline MBRL methods achieve near-zero success on sparse-reward, long-horizon AntMaze tasks, including the simpler umaze variant.Source: Figure 1 · p. 1, §1 · p. 2
• Existing methods still struggle to synthesize transitions containing temporal information about distance from goals, limiting long-horizon behavior learning.Source: §1 · p. 2
• TempDATA combines an autoencoder, latent dynamics model, and offline policy to encode temporal distance at trajectory and transition levels and generate latent-space augmentations.Source: §1 · p. 2
• TempDATA outperforms previous baselines across AntMaze, Kitchen, CALVIN, and pixel-based Kitchen benchmarks, while matching or surpassing prior goal-conditioned MFRL methods.Source: §1 · p. 2

Offline Reinforcement Learning

• Offline MFRL regularizes policies or value functions toward the dataset distribution, whereas TempDATA augments the dataset using a learned dynamics model.Source: §2 · p. 2

Offline MBRL

• Prior offline MBRL learns environment dynamics and generates additional data, but reported methods do not solve goal-reaching tasks with pixel-based states.Source: §2 · p. 2

Goal-conditioned Reinforcement Learning

• Goal-conditioned RL develops policies for reaching specified or discovered goals through supervised, unsupervised, skill-discovery, hierarchical, contrastive, and distance-based methods.Source: §2 · p. 2
• TempDATA builds a state abstraction embedding temporal distance to any goal or next state, then optimizes a policy with latent representations and dynamics.Source: §2 · p. 2

Learning State Space Abstractions

• State abstraction compresses high-dimensional observations and can represent temporal differences, using objectives such as successor features, contrastive learning, metric learning, dynamics modeling, and distance learning.Source: §2 · pp. 2–3
• Temporal distance can better capture transition difficulty than Euclidean distance in environments with obstacles, varied objects, and image-based information.Source: §2 · p. 3
• TempDATA learns a geometric autoencoder that extracts temporal distance from raw states and reconstructs states from micro- and macroscopic levels.Source: §2 · p. 3

Markov Decision Process (MDP)

• An MDP is defined by states, actions, transition probabilities, initial-state distribution, rewards, and discounting; this work assumes a deterministic, discrete-time, infinite-horizon setting.Source: §3 · p. 3

Offline RL

• Offline RL learns policies and value functions from fixed trajectories generated by potentially optimal, suboptimal, or random behavioral policies.Source: §3 · p. 3
• Value-based offline RL estimates Q(s, a) or V(s) with Bellman optimality and trains neural-network parameters by minimizing temporal-difference error.Source: §3 · p. 3

Offline MBRL

• Offline MBRL fits a learned transition function by maximum likelihood and uses it to generate augmented transitions for subsequent policy and value learning.Source: §3 · p. 3

Problem Setting

• The problem setting uses fixed unlabeled reward-free trajectories and goal-conditioned rewards defined by whether the state equals a specified goal.Source: §3 · p. 3
• To avoid transitions in inaccessible regions, the method represents states with temporal-distance-aware latents and performs augmentation within that representation space.Source: §3 · p. 3

4.1. Temporal Distance-aware Representation

• TempDATA learns a temporal distance-aware representation that captures shortest-path information between states and goals while enforcing local temporal consistency between consecutive transitions.The autoencoder reconstructs states and regularizes representations at both trajectory and transition levels.Source: §4.1 · pp. 4–5

4.2. Latent Dynamic Transition Model

• The latent dynamic model predicts one-step representation transitions z → z′ by marginalizing the state space, supporting model learning in high-dimensional domains such as pixel-based environments.It is trained with a Gaussian likelihood objective, while reward modeling is omitted because goal-conditioned RL can use intrinsic goal-achieving rewards.Source: §4.2 · p. 5

4.3. Model-based Policy Training

• TempDATA iteratively rolls out the latent model, decodes synthetic transitions into state space, mixes them with offline data, and extracts a policy using an existing offline RL algorithm.The rollout and policy-extraction phases repeat until convergence or a predefined iteration limit.Source: §4.3 · pp. 5–6, §4.4 · p. 6

4.4. Algorithm Summary

• The framework comprises autoencoder pretraining, latent dynamics training, and repeated rollout-based policy extraction, with goal labels assigned before network training.The algorithm samples offline and generated data using a preset ratio and supports actor-critic or weighted supervised-learning policy optimization.Source: §4.4 · p. 6, §4.3 · p. 6
• Table 1 evaluates TempDATA on D4RL AntMaze, listing eight task rows and reporting total scores without ultra of 350.5, 236.2, and 376.5, with best performance highlighted.Source: Table 1 · p. 7, §4.4 · p. 7

5. Experiments

• The evaluation covers AntMaze, FrankaKitchen, CALVIN, Visual FrankaKitchen, and dense-reward D4RL tasks using MBRL, GCRL, and diffusion-based trajectory-augmentation baselines.Scores are based on eight random seeds, with normalized average success rates over 50 test trials.Source: §5 · p. 7

5.1. Performance Comparison with MBRL and TARL

• TempDATA achieves the best score on 8 of 8 D4RL AntMaze tasks, with an approximately 200-point advantage over ROMI and reported gains of 90.79% over MBRL and 14.66% over TARL methods.It also reduces training time by avoiding the ensemble dynamics structure used by prior methods.Source: §5.1 · p. 8

5.2. Performance Comparison with GCRL

• TempDATA attains the highest overall score on FrankaKitchen and CALVIN multi-goal tasks, rivaling or surpassing the evaluated goal-conditioned RL baselines.Its advantage is consistent across kitchen-mixed, kitchen-partial, and CALVIN settings.Source: §5.2 · p. 8

5.3. Performance on Pixel-based Tasks

• TempDATA performs comparably to GC-IQL on pixel-based Visual FrankaKitchen, while the RepB-SDE baseline struggles in these settings.The evaluation uses two Visual FrankaKitchen datasets, as summarized in Table 3.Source: §5.3 · p. 8, Table 3 · p. 8

5.4. Generalizability on Dense Reward Tasks

• TempDATA matches or outperforms MOPO, GTA, and IQL across D4RL HalfCheetah and Walker2D medium-replay and medium-expert datasets.These dense-reward results test applicability beyond the method’s primary long-horizon, sparse-reward focus.Source: §5.4 · p. 8, Figure 7 · p. 8

5.5. Ablation Study

• Naive model-based rollouts consistently reduce average returns, whereas TempDATA maintains or improves performance across vanilla and skill-based policy variants.TempDATA preserves the gains of skill RL and can restore performance to match or exceed pure MFRL after naive rollouts degrade it.Source: §5.5 · p. 9, Figure 8 · p. 9

5.6. Scalability on Arbitrary Goal

• Heatmaps across four goal locations show learned latent distances following obstacle-aware shortest-path travel times rather than straight-line Euclidean paths.The consistent contours indicate a reusable temporal metric that supports arbitrary unseen goal locations.Source: §5.6 · p. 9, Figure 9 · p. 9

• TempDATA augments transitions in latent rather than raw state space, capturing temporal distance across microscopic and macroscopic levels for goal-reaching policies.Source: §6 · p. 9
• The framework is positioned as outperforming offline MBRL methods in challenging goal-reaching benchmarks while competing favorably with GCRL and TARL approaches.Source: §6 · p. 9
• Model-based rollouts may require further study regarding how far they extend beyond empirical support.Source: §6 · p. 9
• Partial, noisy, or stochastic observations may challenge global distance preservation.Source: §6 · p. 9
• Extending the approach to environments where other agents drive state transitions remains an open question.Source: §6 · p. 9
• The conclusion also questions symmetric metrics for temporally asymmetric tasks and strictly offline pre-training.Source: §6 · p. 9
• Future directions include uncertainty modeling, multi-agent reasoning, asymmetric geometry, and hybrid online–offline adaptation for more generalizable distance-aware control.Source: §6 · p. 9

• TempDATA generates temporal-distance-aware transitions in latent space without extra environment interactions, supporting long-horizon behavior learning.Source: Impact Statement · p. 10
• The framework is intended to improve sample efficiency and robustness in offline RL applications.Source: Impact Statement · p. 10

Convolutional encoder for the pixel-based task

• For pixel-based tasks, an IMPALA CNN provides 512-dimensional features shared by the autoencoder, dynamics model, and policy, avoiding a raw image decoder.Source: A. Training Details · p. 14

Goal-conditioned Reinforcement Learning

• Goal-conditioned training samples a latent skill variable and learns π(a|s, w) from reward-free offline data using latent-space skill progress.Source: A. Training Details · p. 14
• Offline policy training combines IQL with AWR for stable in-sample improvement and does not use a hierarchical policy structure.Source: A. Training Details · p. 14

Policy execution

• At test time, skill adaptation selects a latent direction toward the goal without additional training and can operate zero-shot.Source: A. Training Details · p. 14

Goal labeling

• Hindsight goal relabeling assigns randomly selected, future, and completely random states as goals with probabilities 20%, 50%, and 30%, respectively.Source: A. Training Details · p. 14

Model rollout

• Model rollouts begin after the first 30% of training, then add synthetic samples equal to half the current offline buffer every 10% of the remaining period.Source: A. Training Details · p. 14

• Expectile regression replaces the usual squared Bellman error while learning dθ from dataset or replay-buffer transitions.Source: B. Proof of Theorem 4.2 · p. 15
• For bounded X, increasing τ toward 1 drives its τ-expectile toward the supremum x*, which transfers to the learned distance dτ.Source: B. Proof of Theorem 4.2 · p. 15
• Under deterministic dynamics, every path cost is at least the shortest-path cost, providing the lower-bound argument for the learned distance.Source: B. Proof of Theorem 4.2 · p. 15
• The Bellman-like constraints and expectile objective prevent dθ from collapsing below the maximum relevant cost as τ approaches 1.Source: B. Proof of Theorem 4.2 · p. 15
• In the limit, matching upper and lower bounds make dθ(s, sgoal) converge to the shortest-path or −V* cost.Source: B. Proof of Theorem 4.2 · p. 15

C. Implementation Details

• The implementation is based on JaxRL and runs on RTX 3090 GPUs, with experiments capped at 3 hours for state-based and 12 hours for pixel-based environments.These runtime limits apply per experiment.Source: C. Implementation Details · p. 16

C.1. Environments

• Experiments cover AntMaze, FrankaKitchen, CALVIN, and pixel-based Kitchen settings, with datasets drawn from D4RL or teleoperated demonstrations.AntMaze includes multiple environment scales, while Kitchen and CALVIN contain multi-subtask robotic manipulation trajectories.Source: §C.1 · p. 16
• AntMaze datasets contain 999 trajectories of 1000 steps, and evaluation assigns reward 1 when the agent reaches the designated goal.AntMaze-ultra is twice the size of AntMaze-large, and evaluation goals modify the first two state dimensions to the target position.Source: §C.1 · p. 16
• Kitchen episodes contain four subtasks, rewarding each completed subtask; pixel-based experiments render states as 64 × 64 × 3 camera images.The datasets are kitchen-{partial, mixed}-v0 from D4RL, using a camera configuration matching prior work.Source: §C.1 · p. 16
• CALVIN contains 499 trajectories of 1204 transitions, with four subtasks performed in arbitrary order and reward 1 for each completed subtask.The offline data originate from teleoperated demonstrations introduced in prior work.Source: §C.1 · p. 16

C.2. Hyperparameters

• Training uses 106 iterations for state-based and 5 × 10^5 for pixel-based experiments, with learning rate 3 × 10−4, Adam, and 50 evaluation episodes.Batch sizes are 512 for AntMaze, 256 for FrankaKitchen, and 128 for CALVIN.Source: §C.2 · p. 16
• The autoencoder uses latent dimension 32 for AntMaze and 10 for CALVIN and FrankaKitchen, with discount factor 0.99 and task-dependent expectile coefficients.Expectile coefficients are 0.95 for AntMaze, 0.97 for CALVIN and FrankaKitchen, and 0.7 for pixel-based settings.Source: §C.2 · p. 16
• Dynamic-model, critic, and actor networks each use three hidden layers of width 512, while model-based rollouts use 3 steps.These network dimensions are shared across the listed experiments.Source: §C.2 · p. 16
• Offline RL uses discount factor 0.99, inverse temperature 10 for AntMaze and 3 for CALVIN and FrankaKitchen, and expectile coefficient 0.9 for state-based versus 0.7 for pixel-based tasks.The expectile coefficient therefore differs by input modality, while inverse temperature differs by task group.Source: §C.2 · p. 16