DAWM: Diffusion Action World Models for Offline Reinforcement Learning via Action-Inferred Transitions

Technical Summary: DAWM – Diffusion Action World Models for Offline Reinforcement Learning via Action-Inferred Transitions

1. Problem Statement

Diffusion-based world models have demonstrated strong capabilities in synthesizing realistic long-horizon trajectories for offline reinforcement learning (RL). However, existing methods face a critical compatibility issue with standard value-based offline RL algorithms (e.g., TD3BC, IQL) that rely on one-step temporal difference (TD) learning.

• Limitation of Action-Free Models: Methods like the Diffusion World Model (DWM) generate future states and rewards but do not explicitly generate actions. Consequently, they cannot construct complete transitions $(s, a, r, s')$, making them incompatible with one-step TD learning, which requires consistency between states, actions, and rewards for effective critic regularization.
• Limitation of Joint Modeling: Approaches that attempt to jointly model states, rewards, and actions often suffer from increased training complexity and instability, leading to reduced practical performance.
• Limitation of Planning-Based Methods: Diffusion-based planners (e.g., Diffuser, Decision Diffuser) that recover actions via Inverse Dynamics Models (IDM) or joint modeling typically rely on iterative sampling and post-hoc reward selection. This results in high inference costs and bypasses the efficiency of one-step TD learning.

The core challenge is to generate synthetic data that supports stable, efficient one-step TD learning without incurring the training instability of joint modeling or the inference costs of iterative planning.

2. Methodology: DAWM

The authors propose DAWM (Diffusion Action World Model), a modular framework that decouples trajectory generation from action inference to produce complete, labeled transitions suitable for offline RL. The framework consists of two primary phases:

Phase 1: Conditional Diffusion Trajectory Generation

• A diffusion model $p_\theta$ is trained to generate future sequences of states and rewards conditioned on the current state $s_t$, action $a_t$, and a target return-to-go ($R_{tg}$).
• The model defines the conditional distribution $p_\theta(\hat{s}_{t+1:t+H+1}, \hat{r}_{t:t+H} \mid s_t, a_t, R_{tg})$.
• The training objective minimizes the expected error between predicted and true noise using a scoring denoising objective with classifier-free guidance.
• Crucially, this model generates action-free trajectories (states and rewards only), simplifying the generative task by avoiding the complexities of jointly modeling actions within the diffusion process.

Phase 2: Action Completion via Inverse Dynamics Model (IDM)

• To convert the action-free trajectories into complete transitions, a separately trained Inverse Dynamics Model (IDM), $f_\phi$, is employed.
• The IDM reconstructs the missing actions by modeling the distribution $f_\phi(a_t \mid s_{t-m:t-1}, s_t, s_{t+1})$ as a diagonal Gaussian distribution.
• The IDM is trained on real offline data to maximize the likelihood of observed actions given state transitions.
• At inference, the IDM is applied to the diffusion-generated trajectories to infer actions, yielding complete tuples $(s, a, r, s')$.

Offline Policy Training

• The synthesized complete transitions are used to train offline RL agents using standard value-based algorithms, specifically TD3BC (which adds a behavior cloning regularizer to TD3) and IQL (which performs advantage-weighted supervised regression).
• The Q-functions are updated using Bellman targets computed from the reconstructed transitions, enabling efficient one-step TD updates.

3. Key Contributions

1. Modular Framework: The authors propose an efficient framework combining diffusion-based state-reward generation with a separately trained IDM for action inference. This design avoids the training instability of joint modeling while producing complete transitions.
2. Compatibility with One-Step TD Learning: By generating full $(s, a, r, s')$ transitions, DAWM enables the use of conservative offline RL algorithms (TD3BC, IQL) that rely on one-step TD learning, offering a computationally efficient alternative to iterative diffusion planners.
3. Empirical Performance: The paper demonstrates that conservative offline RL algorithms benefit significantly from training on these augmented trajectories, consistently outperforming prior diffusion-based baselines (DWM, DD) across multiple D4RL locomotion tasks.
4. Synthetic Data Viability: Results indicate that agents trained on DAWM-generated synthetic data can achieve performance comparable to, or exceeding, agents trained on real offline datasets, highlighting the potential of synthetic trajectories as a practical alternative.

4. Experimental Results

Experiments were conducted on 9 standard D4RL locomotion benchmarks (Hopper, Walker2d, HalfCheetah) with varying data quality (Medium, Medium-Replay, Medium-Expert).

• Comparison with DWM: DAWM consistently outperformed the action-free DWM baseline. With TD3BC, DAWM achieved an average improvement of 9.3%, and with IQL, 9.5%. The gains were particularly pronounced in complex control tasks like Walker2d.
• Comparison with Planning-Based Methods (DD): DAWM outperformed the Decision Diffuser (DD) in terms of normalized return. Furthermore, DAWM demonstrated significantly higher efficiency, achieving approximately 4.0x faster inference compared to DD, as it avoids iterative sampling during inference.
• Comparison with Real Data: Agents trained on DAWM-generated data achieved normalized returns comparable to agents trained on real data using standard TD3BC/IQL, and in some cases (e.g., Hopper-MR with TD3BC), outperformed them. While SSORL (a baseline using IDM on real data) performed well, the authors note that SSORL benefits from filtering for the longest trajectories, a factor not present in the fixed-length generation of DAWM.
• Ablation Studies:
  • Data Volume (DAWM-T): Reducing the number of conditioning pairs by a factor of 8 (DAWM-T) resulted in only modest performance drops, suggesting that the action completion mechanism is a more critical factor than sheer data volume.
  • Generation Horizon: DAWM maintained robust performance across different generation horizons ($H=1, 3, 7$), indicating that the one-step TD learning framework is stable regardless of the rollout length.

5. Significance and Claims

The paper claims that DAWM bridges the gap between efficient generative modeling and TD-based policy learning. Its primary significance lies in:

• Enabling One-Step TD Learning: It provides a mechanism to use diffusion-generated data for standard, stable offline RL algorithms that were previously incompatible with action-free world models.
• Computational Efficiency: It offers a practical solution for large-scale generative offline RL by reducing inference costs compared to iterative planning methods.
• Practical Alternative: The results suggest that synthetic trajectories generated via this modular approach can serve as a viable, high-quality alternative to real offline datasets, potentially mitigating data scarcity issues in real-world applications where exploration is expensive or unsafe.

The authors conclude that while the method shows promise, future work will investigate its extension to visual control, high-dimensional robotics, and robustness under imperfect dynamics.