PaperlayerView exact revision on arXiv

Source-linked AI summary

Learning from Reward-Free Offline Data: A Case for Planning with Latent Dynamics Models

Vlad Sobal, Wancong Zhang, Kyunghyun Cho, Randall Balestriero, Tim G. J. Rudner, Yann LeCun

arXiv:2502.14819v4cs.LG

TL;DR

The paper asks how reinforcement learning and optimal control compare when learning diverse tasks from reward-free offline trajectories. It evaluates these paradigms on navigation datasets of varying quality, using PLDM as a latent-dynamics planning method, and finds that PLDM generalizes best while model-free RL benefits most from plentiful, high-quality data.

  • Problem

    The comparative strengths of reinforcement learning and optimal control for learning from reward-free offline trajectories remain underexplored, especially across datasets differing in quality and quantity.

  • Method

    The study evaluates reinforcement-learning and optimal-control methods on navigation tasks, including PLDM, which learns latent dynamics with JEPA and plans in latent space.

  • Results

    PLDM shows best-of-class generalization to new layouts and state avoidance, while model-free GCRL is strongest when data is plentiful and high quality.

  • Takeaways & Limitations

    Latent-dynamics planning is robust to suboptimal data, data-efficient, and effective for adapting to tasks beyond goal-reaching.

  • Takeaways & Limitations

    Experiments are limited to navigation environments, and PLDM is about 4 times slower during inference.

Abstract

from arXiv · show

A long-standing goal in AI is to develop agents capable of solving diverse tasks across a range of environments, including those never seen during training. Two dominant paradigms address this challenge: (i) reinforcement learning (RL), which learns policies via trial and error, and (ii) optimal control, which plans actions using a known or learned dynamics model. However, their comparative strengths in the offline setting - where agents must learn from reward-free trajectories - remain underexplored. In this work, we systematically evaluate RL and control-based methods on a suite of navigation tasks, using offline datasets of varying quality. On the RL side, we consider goal-conditioned and zero-shot methods. On the control side, we train a latent dynamics model using the Joint Embedding Predictive Architecture (JEPA) and employ it for planning. We investigate how factors such as data diversity, trajectory quality, and environment variability influence the performance of these approaches. Our results show that model-free RL benefits most from large amounts of high-quality data, whereas model-based planning generalizes better to unseen layouts and is more data-efficient, while achieving trajectory stitching performance comparable to leading model-free methods. Notably, planning with a latent dynamics model proves to be a strong approach for handling suboptimal offline data and adapting to diverse environments.

1 Introduction

The paper examines how reinforcement learning and optimal control can learn from reward-free, suboptimal offline trajectories for tasks and environments not seen during training. It systematically varies dataset quality and quantity, finding that latent-dynamics planning offers the strongest generalization.

  • Motivation: The study addresses learning systems that solve unseen combinations of tasks and environments from large collections of suboptimal, reward-free trajectories.It compares reinforcement learning and optimal control as the two main approaches.
  • Study design: The comparison spans reinforcement-learning and optimal-control paradigms, including methods trained on offline trajectories without reward annotations.The analysis is designed to isolate how dataset quality and quantity affect different approaches.
  • Study design: The authors introduce two navigation environments and 23 datasets varying in data quality, while evaluating six generalization properties across offline reward-free learning methods.The evaluation includes random-policy trajectories, trajectory stitching, limited data, unseen layouts, and tasks beyond goal-reaching.
  • Findings: Planning with a latent dynamics model achieves the highest level of generalization in the study.The paper presents PLDM as robust to suboptimal data quality and strong across environment variations.
  • Implications: The paper also provides practitioner guidelines for choosing methods according to available data and required generalization.Code, data, and environment visualizations are released for further research.

2 Related Work

Related work frames reward-free offline learning as task-agnostic behavior learning and distinguishes goal-conditioned, zero-shot, and model-based approaches. It also motivates systematic analysis of which offline-data properties affect performance.

  • Reward-free offline RL: Reward-free offline RL learns general behaviors from offline data without rewards so agents can solve varied downstream tasks.Goal-conditioned methods specify tasks through goals or observations, while zero-shot methods aim to solve arbitrary test-time tasks.
  • Reward-free offline RL: Zero-shot approaches include distance-preserving latent representations and Forward-Backward representations for handling tasks beyond ordinary goal-reaching.HILP learns a latent space whose distances correspond to the number of steps between states.
  • Optimal control: Optimal control selects actions by optimizing an objective with a known or learned dynamics model, whereas RL methods train policy functions using reward information.The paper uses this distinction to organize its comparison of learning paradigms.
  • Offline-data quality: Prior studies examined exploratory versus task-specific data or offline RL with rewards, but did not isolate which aspects of data quality affect reward-free learning.This motivates the paper’s systematic comparison across datasets varying in quality and quantity.

3 The Landscape of Available Methods

The paper surveys reward-free offline learning approaches spanning model-free RL and latent-dynamics planning. It contrasts goal-conditioned and zero-shot policies with PLDM, which learns latent dynamics and plans in latent space.

  • Problem setting: Reward-free offline learning seeks policies that solve varied tasks and layouts from state-action trajectories without reward annotations.The setting treats the task as a downstream objective specified at evaluation.
  • Method comparison: Across generalization stress tests, HILP and PLDM are identified as the best-performing methods, with PLDM competitive across all tested settings.The roadmap compares four goal-conditioned methods, one zero-shot RL method, and PLDM using three performance levels.
  • Reinforcement learning: The evaluated model-free methods include goal-conditioned RL methods and the zero-shot representation-based method HILP.HILP learns a distance-preserving representation and a direction-conditioned policy, while goal-conditioned methods learn policies for specified goals.
  • Latent-dynamics planning: PLDM learns latent dynamics from offline trajectories using a reconstruction-free JEPA-based objective, then plans actions in the learned latent space.Its training minimizes distances between predicted and encoded latents; VICReg-inspired and inverse-dynamics objectives prevent representation collapse.
  • Latent-dynamics planning: PLDM plans toward goal observations using a goal cost and an uncertainty cost that penalizes transitions deviating from the training distribution.The method uses model-predictive-control replanning and MPPI; by default, it replans every interaction and is approximately 4x slower than model-free baselines.

4 Every Method Can Excel but Few Generalize

Across navigation tasks, methods perform well with abundant, high-quality data, but their robustness diverges under scarce, short, noisy, or distribution-shifted data. PLDM and selected model-free methods handle suboptimal data effectively, while PLDM generalizes best to new tasks and unseen layouts.

  • In-distribution performance: With 3 million transitions and high-quality trajectories, all methods perform well, with HIQL, GCIQL, HILP, and PLDM nearing 100% success.The topline setting provides abundant data, good state coverage, and trajectories long enough to traverse both rooms.
  • Data efficiency: PLDM and GCIQL are the most sample-efficient methods, outperforming other model-free methods when only a few thousand transitions are available.HILP requires more data but reaches perfect performance when enough data is provided.
  • Trajectory stitching: When trajectories are short, most goal-conditioned model-free methods struggle to stitch them into long-horizon behavior, whereas PLDM and HILP remain effective.The experiments vary episode lengths while keeping the total number of transitions fixed; PLDM learns a model and plans at test time.
  • Trajectory stitching: Under imperfect coverage, HILP and GCIQL achieve perfect performance, while PLDM declines but remains stronger than most offline model-free GCRL methods.The constrained trajectories remain within one room, although evaluation requires crossing the connecting door.
  • Data quality: With low-quality random-policy data, HILP, GCIQL, and PLDM outperform the other goal-conditioned RL methods.Random trajectories stay near their starting points, making distant goals out-of-distribution for many goal-conditioned methods; PLDM uses them to learn dynamics.
  • Generalization: PLDM generalizes beyond goal-reaching by changing the planning objective and generalizes best to unseen obstacle layouts, even when trained on only 5 maps.In the chasing task, PLDM outperforms HILP; in held-out layouts, other methods degrade as layouts become more out-of-distribution.

5 Conclusion

The study identifies PLDM and HILP as the strongest approaches for learning from reward-free offline data, with PLDM especially robust, data-efficient, and effective on unseen layouts and tasks beyond goal-reaching.

  • The evaluation uses navigation environments to study spatial reasoning, generalization to new layouts, and trajectory stitching.
  • PLDM combines robustness to data quality, high data efficiency, strong generalization to new layouts, and adaptation to tasks beyond goal-reaching.The conclusion presents these as PLDM’s main strengths.
  • HILP’s well-structured latent space supports trajectory stitching and robustness to data quality, but requires more data than other methods.
  • Model-free GCRL methods are preferable when offline data is plentiful and high quality.

C.2.1 Dataset Generation

The experiments use controlled navigation datasets and evaluate generalization across unseen layouts, distribution shifts, and environments with more complex control dynamics.

  • The dataset-generation setup includes four training datasets for the Diverse PointMaze experiments.
  • Test layouts are disjoint from training layouts, with randomly sampled starts and goals at least three cells apart.
  • Agents are evaluated both across varying numbers of training layouts and across test maps with different distribution shifts.
  • Distribution shift is quantified by the minimum edit distance between a test layout and the closest training layout.
  • The Ant U-Maze study uses an 8-dimensional action space and 29-dimensional state space to test trajectory stitching under more complex dynamics.
  • Ant-U-Maze data contains 5M exploratory transitions generated with a pretrained directional policy, direction resampling every 10 steps, and Gaussian action noise.

D.1.2 Ablations of Objective Components

The ablation and efficiency analyses examine PLDM’s objective components, architecture, and planning frequency, showing that less frequent replanning can preserve performance while reducing inference cost.

  • The ablations remove each loss component under fixed optimal Two-Rooms and Diverse Maze settings to assess its contribution.Two-Rooms uses sequence length 90, dataset size 3M, and expert data; Diverse Maze uses five training maps.
  • The model uses an Impala Small Encoder and a 2-layer GRU predictor with 512 hidden dimensions for the standard setting.
  • For Diverse PointMaze, the model encodes maze images and velocity maps convolutionally, then predicts representations from state and expanded action inputs.
  • Replanning every 4 steps makes PLDM close to other methods’ inference speed, with only slight performance degradation in Two-Rooms and diverse mazes.

F.1 Evaluating PLDM Performance With Adjusted Inference Compute Budget

Under a fixed inference-time compute budget, PLDM remains close to its per-step-replanning performance when replanning less frequently.

  • Replanning every 4 steps causes only slight degradation compared with replanning every step in Two-Rooms and diverse mazes.

G Extended Baselines: Input Reconstruction Learning and TD-MPC2

The paper compares reconstruction-based latent-dynamics learning and TD-MPC2 with PLDM and other baselines on good-quality Two-Rooms data. Reconstruction objectives perform poorly, while TD-MPC2 collapses without additional inverse dynamics modeling.

  • Reconstruction-Based Methods: Reconstruction-based methods perform significantly worse than PLDM and other baselines on good-quality Two-Rooms data.The comparison includes a Dreamer-based approach and a reconstruction variant using PLDM’s encoder and dynamics architecture.
  • Comparison Caveat: The DreamerV3 comparison is acknowledged as flawed because the architecture was not designed for this reward-free setting.The adaptation removes rewards and policy learning, then plans through latent representation distance using KL divergence.
  • Reconstruction-Based Methods: Pixel observation reconstruction is a poor objective for learning representations and produces poor planning performance.The evaluation compares pixel reconstruction against PLDM’s reconstruction-free objective.
  • TD-MPC2: TD-MPC2 collapses altogether, achieving 0% success rate in the Two-Rooms evaluation.TD-MPC2 was adapted by removing reward prediction and retaining consistency between predictor and encoder outputs.
  • TD-MPC2: Adding inverse dynamics modeling somewhat improves TD-MPC2 performance, although the adapted method remains poor.The inverse-dynamics addition was introduced to prevent collapse.

H Analyzing Statistical Significance of Results

The paper assesses statistical significance by pooling results and by running selected evaluations with ten seeds. It also analyzes HILP’s out-of-distribution failure through its learned distance representation.

  • Statistical Analysis: Welch’s t-test shows that PLDM’s results are significant overall, except for certain settings.The pooled analysis combines results across settings because five seeds were insufficient for statistical testing.
  • Statistical Analysis: The selected-dataset analysis uses statistical tests across 10 seeds per method.Additional seeds were run for selected settings and reported separately from the pooled analysis.
  • HILP Generalization: HILP represents distances so that the latent distance to a goal corresponds to the lowest number of transitions needed to reach it.The analysis uses this representation to investigate HILP’s poor generalization to out-of-distribution maze layouts.
  • HILP Generalization: The paper hypothesizes that HILP fails on unseen maze layouts because its learned representation predicts incorrect distances there.The resulting distance errors are linked in the analysis to failure of the goal-conditioned policy.

J.2.1 Two-Rooms

The Two-Rooms and related experiments vary dataset and training settings while documenting implementation and baseline context. The paper also situates the methods within offline RL, foundation-model, and representation-learning research.

  • Experimental Settings: The reported best-case setting uses sequence length 91, dataset size 3M, 100% non-random data, and approximately 35% wall crossing.The experiments vary each of these parameters individually.
  • Experimental Settings: The experiments include dataset-agnostic and dataset-specific hyperparameters for Two-Rooms and other navigation environments.The supplied tables document these settings for Two-Rooms, Diverse PointMaze, and Ant-U-Maze.
  • Related Work: Offline RL learns behaviors purely from offline data without online interactions, while seeking policies that can outperform the data-collecting policy.A central challenge is avoiding actions that were not observed in the dataset.
  • Related Work: Foundation-model work in RL has been enabled by large datasets such as Open-X-Embodiment and DROID.The passage connects these datasets to models including RT-2 and Octo.
  • Related Work: Self-supervised representation learning is another route for using large amounts of data to improve RL agents and planning.Prior work includes SSL objectives combined with RL, separate pre-training, and visual representations for planning.
  • Computational Cost: The experiments require a single GPU and can take up to one day on an Nvidia V100, with estimated total research cost of 500–2000 GPU days.The estimate includes experiments and the broader research process.

L Limitations

The evaluation is limited to navigation environments and excludes robot manipulation and partially observable settings. PLDM also incurs substantially higher inference cost than model-free methods.

  • Scope: All experiments use navigation environments, excluding robot manipulation and partially observable settings.The authors state that the conceptual findings may still apply beyond this relatively simple setting.
  • Inference Cost: PLDM is about 4 times slower during inference than model-free methods.Under an equal inference-compute budget, PLDM retains most of its performance according to the appendix analysis.
Loading 2502.14819v4…