Awesome-WAMs
GitHub
PAPER REPORTENAll readings ↗

World Models via Policy-Guided Trajectory Diffusion

English reading report: Method, equations, original figures, experiments and reproducibility.

Authors: Marc Rigter; Jun Yamada; Ingmar Posner

Source: TMLR 2024 · ref-da609a82e0cab3355273 ↗ · Catalog record

Reading: 441 / 558 · 6 original figures & tables · ~21 min ·

1. Paper overview

In one sentence: PolyGRAD refines whole imagined trajectories using a denoiser and policy-score action updates, gaining a non-autoregressive sampler while making action calibration a central training constraint. e02e04e05e09e11e15

At a glanceWhat to know
Research problem
Source description

On-policy world modeling normally alternates policy sampling and next-state prediction. Generating an entire trajectory simultaneously creates a circular dependency: future actions require future states, while those states depend on the actions. PolyGRAD addresses this dependency in fully observable continuous-control MDPs, aiming to supply synthetic data for policy optimization. e02e03e05

Core mechanism
Source description

The authors propose non-autoregressive on-policy trajectory generation using learned dynamics denoising and direct policy-score guidance of actions. e04e05

A key reported resultImagined RL on three MuJoCo control tasks: PolyGRAD exceeds PPO, TRPO, and A2C in aggregate final performance, but trails Dreamer-v3 in sample efficiency and final performance; exact scores and interval endpoints are not tabulated.

Smoothed episode return and aggregate normalized interquartile mean with 95% bootstrap confidence intervals. Residual-MLP PolyGRAD with h=10; five seeds; plots end at 1.5M steps for HalfCheetah/Walker2d and 1M for Hopper.

Final rewards are normalized by each environment's best average final algorithm reward. Confidence intervals use 50,000 stratified bootstrap repeats. Better state reconstruction does not imply better control. Comparisons with these on-policy model-free baselines do not establish superiority over off-policy model-free RL. e11e20

Reading caution
Source description

Slow policy updates are needed for stability, and low-entropy policies can produce incorrect action distributions. Appendix B also reports HalfCheetah NaNs before the 0.02 target and heavy tails without clipping, despite matching standard deviation. e13e18

Core contributions

  • Source description

    The authors propose non-autoregressive on-policy trajectory generation using learned dynamics denoising and direct policy-score guidance of actions. e04e05

  • Source description

    Langevin and classifier-guidance connections motivate the sampler, while an online scale adjustment approximately matches policy dispersion. These arguments do not establish convergence. e06e07

Figure 2. State predictions and policy-guided actions are refined together across diffusion steps. Original paper, p. 6 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Start at the upper-left random trajectory and follow the connecting arrows, which move right, down, right, up, right, and finally down. Colored dots represent states; red arrows represent actions. The lower-left panel separately illustrates the policy's action distribution throughout state space. Solid connecting arrows mark action-score updates, while hollow connecting arrows mark denoiser-driven state updates. The caption's phrase 'current state sequence' is shorthand: Algorithm 2 first estimates fully denoised states and conditions the policy on that estimate. Conversely, the denoiser conditions on the current actions. The dark-purple initial state comes from the dataset and is reset by inpainting at each diffusion step. e05e06e16

What it supports. The mechanism resolves a circular dependency by repeatedly revising both sides: states must fit actions, and actions must fit the policy at those states. The whole trajectory is refined without autoregressive environment-time rollout. Nevertheless, generating it still requires the full iterative diffusion process, rather than a single network evaluation.

Where the evidence stops. This is a conceptual illustration, not an evaluated navigation task. Its arrows agree with Algorithm 2, but they do not prove convergence or exact on-policy sampling. The policy-conditioned states change during diffusion, which limits the direct Langevin analogy.

2. Motivation

2.1 The problem and the proposed response

Source description

On-policy world modeling normally alternates policy sampling and next-state prediction. Generating an entire trajectory simultaneously creates a circular dependency: future actions require future states, while those states depend on the actions. PolyGRAD addresses this dependency in fully observable continuous-control MDPs, aiming to supply synthetic data for policy optimization. e02e03e05

2.2 What this reading follows

A world model must predict what happens under the policy being trained. Predicting all future states together makes that requirement awkward: the policy needs those states before it can choose their actions. PolyGRAD starts with noisy states and actions, then repeatedly brings them into agreement through an action-conditioned denoiser and the policy's action score. The resulting synthetic trajectories train a separate actor-critic policy. Read the evidence in three stages: how the coupled sampler works, whether its action distribution matches the policy, and whether its predictions support useful control. Strong short-horizon prediction is encouraging, but low-entropy failures and Dreamer-v3's stronger returns delimit the result. e02e04e05e09e11e15

3. Research context

We place the paper in the collection through its world–action interface. The catalog labels and the reading’s assessment are shown separately.

Catalog dimensionRecorded classification
Major categoryFoundational work
ArchitectureNot applicable
Prediction paradigmNot applicable
QuadrantNot applicable

3.1 Evidence-based assessment

Supports the recorded classification

Reader analysis

Foundational world modeling/model-based RL and diffusion are supported. The source explicitly separates a denoising model from a policy; coupled trajectory generation does not establish a unified One Model architecture or an inverse-dynamics action decoder. Retaining the catalog's not-applicable architecture, prediction-paradigm, and quadrant fields is appropriate for its foundational placement. e02e04e05e07

This is the collection’s architectural analysis, not a new related-work survey. Benchmark comparisons and their protocols appear in Section 6.

4. Problem formulation

4.1 Inputs and outputs

InputsOutputs
  • Replay-buffer trajectories of fully observed states, actions, rewards, and terminal indicators
  • A buffer-sampled initial state and the current Gaussian policy
  • Synthetic trajectories for imagined on-policy actor-critic training
  • A separate feedforward policy for environment interaction

4.2 Equations and their role

∥ϵ−ϵθ(αiτsr+1−αiϵ∣i,τa)∥22\left\|\epsilon-\epsilon_\theta\left(\sqrt{\alpha_i}\tau^{sr}+\sqrt{1-\alpha_i}\epsilon\mid i,\tau^a\right)\right\|_2^2
Algorithm 1's sampled denoising loss: τ^{sr} contains states and rewards, τ^a contains actions, ε is sampled Gaussian noise, and ε_θ predicts it. The diffusion level is i; α_i=∏_{j=1}^{i}(1−β_j), with β_j the noise schedule. The objective updates denoiser parameters θ. e03e04
τ^i−1a←τ^ia+δ∇τ^ialog⁡πϕ(τ^ia∣τ^0s)+βiz,i>1\hat{\tau}_{i-1}^{a}\leftarrow\hat{\tau}_{i}^{a}+\delta\nabla_{\hat{\tau}_{i}^{a}}\log\pi_\phi(\hat{\tau}_{i}^{a}\mid\hat{\tau}_{0}^{s})+\sqrt{\beta_i}z,\qquad i>1
Algorithm 2's unclipped action update uses predicted clean states τ̂_0^s, current actions τ̂_i^a, action-update scale δ, policy parameters φ, and standard Gaussian z. The gradient changes actions, not policy parameters. Appendix A adds clipping before the noise term. e05e14
aˉi=ai−μϕ(si)σϕ(si),δ←δ+η(σaˉ−1)\bar a_i=\frac{a_i-\mu_\phi(s_i)}{\sigma_\phi(s_i)},\qquad\delta\leftarrow\delta+\eta(\sigma_{\bar a}-1)
Equations (9)–(10) standardize sampled actions using policy mean μ_φ and standard deviation σ_φ. Here i indexes batch state-action pairs, σ_{ā} is their standardized-action standard deviation, and η is a learning rate. Excess dispersion increases guidance; insufficient dispersion decreases it. e07

5. Method in detail

5.1 Learn dynamics first, then ask which policy generated the actions

Source description

The denoiser is trained on real buffer trajectories, with noise added to states and rewards and clean actions supplied as conditioning. Its job is to recover the noise, not to learn an action decoder. At generation time, PolyGRAD initializes an entire noisy trajectory and repeatedly predicts how the states should change given the current actions. A separate policy supplies the action-score direction. To query that policy sensibly, the sampler first estimates the fully denoised states: the policy was not trained to interpret diffusion-corrupted inputs. These two operations make a coupled sampler while keeping the learned components distinct. One reverse process spans the whole trajectory, but the implementation still performs 128 diffusion steps. Thus the architectural novelty concerns environment-time rollout, not elimination of iterative computation. e04e05e08e16

Table 1. The imagined-learning setup deliberately combines short rollouts with constrained policy changes. Original paper, p. 18 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Read the first two rows together: each actor-critic update receives 1,024 imagined trajectories, each of length 10. The final row sets one such training update per four environment steps. GAE parameter λ and discount γ specify the return-estimation setup. Keep the critic learning rate separate from the target policy-update row: Appendix A uses a line search on the policy learning rate to keep average absolute log-likelihood change within 20% of 0.01. The minimum policy standard deviation, 0.1, is another training constraint. This table describes the policy learner; denoiser architecture and diffusion settings appear in Tables 2–4. e07e09e15e16

What it supports. PolyGRAD's control results use a carefully restricted actor-critic setup. The standard-deviation floor avoids the most problematic low-entropy region in the main diagnostic, while the small target update addresses policy-change instability. These are parts of the reported recipe, not merely incidental optimizer defaults that can be dropped in a reproduction.

Where the evidence stops. This is an original numerical configuration table, not a performance table or a controlled hyperparameter ablation. It does not isolate the causal benefit of each setting, and it does not supply numerical values for the initial guidance scale or its adaptation learning rate.

5.2 Treat on-policy generation as a calibration problem

Reader analysis

Moving actions toward a policy's high-density region is insufficient: excessive guidance collapses them toward its mean, whereas weak guidance leaves them too dispersed. PolyGRAD adjusts the action-update scale using the standard deviation of policy-standardized actions. The reader's inference is that this is a useful calibration signal but an incomplete distributional test. Appendix B makes the concern concrete: removing clipping can preserve standard deviation while producing a heavy-tailed, non-Gaussian distribution. Moreover, Appendix A clips the deterministic update before adding Gaussian noise, so the equation does not impose a strict final-action bound. Together with changing state estimates during diffusion, these details explain why the Langevin connection remains motivation rather than a convergence guarantee. The proposed conditional-calibration check below probes precisely what pooled moments can miss. e06e07e14e18

5.3 Separate a useful simulator from a strong control learner

Reader analysis

Prediction evaluation and policy evaluation answer different questions. The former replays generated action sequences in the simulator and compares predicted states with ground truth. The latter trains a policy using imagined rewards and transitions, then measures environment returns. PolyGRAD's control recipe chooses short rollouts and a residual MLP, whose short-horizon errors are similar to the transformer while its computation is cheaper. The policy receives A2C updates, constrained to avoid large changes that destabilize the sampler. Dreamer-v3 instead optimizes behavior through its dynamics model and achieves stronger control despite weaker state reconstruction in some comparisons. As reader analysis, this shows that state MSE alone cannot identify the best training simulator: policy optimization and the usefulness of its representation also matter. Their separate contributions remain unisolated here. e10e11e15e17e21

5.4 Training and inference

During training

Source description

The implemented Gaussian policy has state-dependent means and one state-independent standard-deviation parameter. A2C uses GAE, 1,024 imagined trajectories of length 10 per update, and one update per four environment steps. Policy standard deviation is floored at 0.1; a line search targets mean absolute log-likelihood change 0.01 within 20%. e15

Source description

Both denoisers have six layers: the residual MLP uses width 1,024 for h=10, while the transformer uses embedding dimension 312 and four causal attention heads. Diffusion uses 128 steps and a cosine schedule. Clean-action conditioning during training differs from initially noisy actions during generation. e08e16

During inference

Source description

A single diffusion pass means a complete reverse process over the trajectory, not one network evaluation. These imagined actions train the separate feedforward policy; environment actions are selected by querying that policy directly. e02e05e16

Reader analysis

Appendix A, Eq. (11), clips the deterministic guided action to policy mean ±3 standard deviations before adding Gaussian noise. This is more precise than the main text's shorthand about clipping after action updates; it does not guarantee that the final noisy action stays within those bounds. e08e14

5.5 Implementation flow

  1. Learn action-conditioned dynamics

    Sample a trajectory and diffusion level from the buffer. Corrupt the state-reward sequence and predict the added Gaussian noise conditioned on clean actions. Terminal-state prediction is implemented but omitted from the main notation. Initial-state inpainting is used during both training and generation. e04e15

  2. Initialize the whole rollout

    Initialize action and state-reward sequences with Gaussian noise. At each reverse step, replace the initial state with a uniformly sampled buffer state, then predict noise conditioned on the current action sequence. e05

  3. Couple denoising and policy guidance

    Estimate fully denoised states before evaluating the policy. For diffusion steps above one, move actions along the gradient of log policy probability with respect to actions, with added noise. Update states and rewards through the denoiser. The last diffusion step leaves actions unchanged. e05

  4. Train in imagination and recalibrate

    Alternate real data collection, denoiser training, synthetic rollout generation, and A2C policy/value updates. Standardize generated actions using the policy mean and standard deviation, then adjust guidance strength toward unit standardized dispersion. e07e15

6. Experiments & results

PolyGRAD generates approximately on-policy synthetic state, reward, and action trajectories by combining an action-conditioned denoiser with a separate policy's action score. It refines the whole trajectory through diffusion instead of rolling out one transition at a time. Short-horizon prediction is competitive, but imagined-policy learning remains weaker than Dreamer-v3, and low-entropy action calibration is fragile.

Source and visual limitations
Source description

The supplied paper reports predictive accuracy, returns, confidence intervals, and action-distribution diagnostics as plots. Its four tables contain hyperparameters rather than measured performance. This edition therefore pairs the original result plots with the original A2C configuration table; exact MSE, return, and confidence-interval endpoints are not invented from unlabelled curve positions. e10e11e12e15e16e17

6.1 Read the original evidence

Figure 5. Short-horizon accuracy is competitive; whole-trajectory diffusion does not secure a long-horizon advantage. Original paper, p. 10 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Read each environment separately, with lower prediction MSE indicating better state reconstruction. The red curves use different markers for separately trained PolyGRAD horizon settings; they are not successive checkpoints of one model. The other colors identify the four autoregressive baselines, and the axes use logarithmic spacing. All models receive the same one-million-transition dataset for each environment. Evaluation starts from 500 dataset states, generates actions under the same policy, and replays those actions in the simulator to obtain ground truth. Shaded bands represent standard deviation across five seeds. Compare short-horizon red curves first, then examine how the longer-trained variants fare as the evaluation horizon grows. e10e19e21

What it supports. The paper reports h=10 PolyGRAD as best on Walker2d and second to autoregressive diffusion on HalfCheetah and Hopper. Longer-horizon training weakens its relative position. This supports short-rollout utility, while the broader motivation of avoiding compounding error does not become an empirical claim that PolyGRAD wins at long horizons.

Where the evidence stops. Exact MSE values are not tabulated. Starting states come from the training dataset, and model selection uses five-step error within a capped training budget; a separate held-out selection split is unspecified. These curves therefore do not establish novel-state generalization or equal-compute superiority.

Figures 6 and 7. Imagined on-policy learning beats the tested model-free baselines in aggregate, while Dreamer-v3 remains stronger. Original paper, p. 11 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. The left three panels show average episode return against actual environment interactions. Hopper is plotted to one million steps; HalfCheetah and Walker2d extend to one and a half million. The curves use five seeds and standard-deviation shading. Each evaluation averages ten episodes, and the displayed curves smooth ten consecutive evaluations. The right panel summarizes final performance using normalized interquartile mean, with 95% bootstrap intervals rather than standard-deviation bands. Normalization divides by the best algorithm's average final return within each environment; the vertical scale is not a success percentage. PolyGRAD here uses its residual-MLP denoiser and short h=10 imagined trajectories to train A2C. e10e11e15e20

What it supports. The red aggregate interval lies above the tested on-policy model-free methods and below Dreamer-v3. PolyGRAD therefore demonstrates useful policy learning from imagined trajectories, but the stronger prediction accuracy seen in some state-space comparisons does not translate into the strongest control agent. The distinction between reconstructing states and optimizing behavior is essential to interpreting the paper.

Where the evidence stops. The aggregate result does not imply a clear win over every baseline in every task. Smoothing obscures short-term instability, and the paper does not tabulate exact interval endpoints. This comparison also omits stronger off-policy model-free algorithms discussed by the authors.

6.2 Results and evaluation conditions

Task & protocolReported resultComparison & interpretation
On-policy action-distribution calibration

Walker2d, h=50; policy standard deviation decayed from 1 to 0.02 during Algorithm 3.

Close agreement reported for σ_φ≥0.1; deviations below 0.1, with additional modes near ±3 standard deviations at 0.02.

Aggregate density of action minus policy mean versus the target Gaussian

Generated batch density versus the policy density; Figure 3 supplies no numerical distribution-distance score.

This supports approximate aggregate calibration, not exact conditional or joint trajectory-distribution matching. e09

MuJoCo trajectory prediction

HalfCheetah-v3, Hopper-v3, Walker2d-v3; shared 1M-transition datasets and final collection policy; 500 buffer-sampled starting states; generated actions replayed in the simulator; five seeds.

At h=10, PolyGRAD is reported best on Walker2d and second to autoregressive diffusion on HalfCheetah and Hopper. Longer-trained h=50/200 models lose relative accuracy; exact MSE values are plotted, not tabulated.

State prediction MSE versus horizon; lower is better

Probabilistic MLP ensemble, autoregressive transformer, Dreamer-v3, and autoregressive diffusion.

These are in-distribution branched rollouts; models are selected by five-step error, with a 1M-update or 72-hour cap. Long-horizon superiority is unestablished. e10e19e21

Imagined RL on three MuJoCo control tasks

Residual-MLP PolyGRAD with h=10; five seeds; plots end at 1.5M steps for HalfCheetah/Walker2d and 1M for Hopper.

PolyGRAD exceeds PPO, TRPO, and A2C in aggregate final performance, but trails Dreamer-v3 in sample efficiency and final performance; exact scores and interval endpoints are not tabulated.

Smoothed episode return and aggregate normalized interquartile mean with 95% bootstrap confidence intervals

Final rewards are normalized by each environment's best average final algorithm reward. Confidence intervals use 50,000 stratified bootstrap repeats.

Better state reconstruction does not imply better control. Comparisons with these on-policy model-free baselines do not establish superiority over off-policy model-free RL. e11e20

Policy-guidance ablations and training cost

Same three tasks, five seeds; variants alter action updates, clipping, state conditioning, or state guidance.

Random Actions and Policy Sampling perform poorly; No Clipping and Noisy State Conditioning degrade performance. Add State Update is comparable but costs about 80 hours versus 54 hours per 1M environment steps.

Return, normalized final-performance IQM, and reported wall-clock training time

The base run's 54-hour figure is specified for an RTX 3090 GPU; the 80-hour variant figure has no separate hardware statement in Section 5.4.

Guidance details matter. Comparable returns do not establish equivalence, and the reported costs have no uncertainty intervals. e12e22e23

6.3 Ablations and diagnostic examples

Read component removals and qualitative examples within their stated evaluation conditions.

Figure 3. Approximate action matching deteriorates as the policy becomes nearly deterministic. Original paper, p. 9 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Move left to right through decreasing policy standard deviations. Blue curves are generated-action densities; dashed black curves are the target policy distributions. Although the horizontal label says 'Action,' the caption defines the plotted quantity as action minus the policy mean at the corresponding state. Both horizontal and vertical scales change across panels, so compare shape agreement within each panel rather than raw peak heights across panels. These batches come from Walker2d with imagined horizon 50 while the policy standard deviation is decayed during training. The first four panels track the target closely; inspect the last panel for the sharper center and small additional outer modes. e09e15e18

What it supports. The main experiment supports approximate aggregate action matching down to policy standard deviation 0.1. Below that level, the mismatch becomes visible; at 0.02, the extra modes are attributed to clipping near three standard deviations. The reported control configuration consequently keeps its policy standard deviation at or above 0.1.

Where the evidence stops. These are pooled batch distributions, not state-conditional or temporal-independence tests. The threshold is an empirical observation for this setup. Appendix B's h=10 diagnostics vary by environment and include HalfCheetah numerical failure before reaching 0.02.

Figures 8 and 9. Gradual action guidance and clean-state conditioning matter more than adding a state-score update. Original paper, p. 12 ↗

Excerpt from the authors’ paper; cropped without altering the figure or table.

How to read it. Use the shared legend to compare each modification with red base PolyGRAD. Random Actions leaves the initial Gaussian action sequence unchanged. Policy Sampling replaces it with fresh policy samples at every diffusion step. No Clipping removes the action constraint, and Noisy State Conditioning evaluates the policy on noisy states instead of the denoised estimate. Add State Update also differentiates policy likelihood with respect to states and modifies the denoising prediction accordingly. The left panels display smoothed task returns and five-seed standard deviations; the right panel aggregates normalized final performance with 95% IQM intervals. Focus first on the low black and blue curves, then compare purple against red. e05e12e20e22e23

What it supports. Direct resampling is not an adequate substitute for gradual guidance in these experiments. Removing clipping or clean-state conditioning also harms aggregate returns. Additional state guidance gives comparable performance, but the text reports about 80 training hours per million environment steps versus 54 for the base method. More policy gradients therefore bring extra cost without an established control gain here.

Where the evidence stops. The noisy-state descriptions on pp. 12 and 25 mistakenly use an action superscript while naming states; Algorithm 2 resolves the intended conditioning. Overlapping intervals do not establish equivalence. Normalized values here should not be treated as raw returns or directly compared with Figure 7's scale.

7. Analysis & limitations

7.1 What the evidence leaves open

Source description

Slow policy updates are needed for stability, and low-entropy policies can produce incorrect action distributions. Appendix B also reports HalfCheetah NaNs before the 0.02 target and heavy tails without clipping, despite matching standard deviation. e13e18

Source description

The classifier-guidance derivation assumes clean trajectories and a uniform background action distribution; a correct noisy-trajectory gradient and formal sampler convergence remain open. Image-based, non-Markovian, and improved long-horizon settings are future work. e06e13

Reader analysis

Appendix C.1 does not specify an independent held-out rollout split for model selection. Aggregate action histograms also leave state-conditional errors unresolved. Neither protocol establishes out-of-distribution control reliability. e09e19

7.2 Questions for discussion

  1. Can conditional calibration fail while aggregate standardized actions appear Gaussian? (e09, e18)
  2. How much of Dreamer's control advantage comes from its policy optimizer rather than prediction quality? (e11, e21)

8. Reproducibility audit

8.1 Requirements and known gaps

Source description

Match the trajectory-generating policy, dataset, rollout initialization, and five-step model-selection rule. Dreamer receives five real state-action pairs to initialize its latent state; the autoregressive transformer has maximum context 15. These baseline-specific contexts must be retained. e19e21

Reader analysis

Match Appendix A's optimizer settings, clipping order, and policy line search. Hopper with the residual MLP uses cosine schedule parameter 0.1 instead of 1. The tables do not supply initial δ or numerical η, and the cited implementations lack pinned versions in the PDF. e07e14e15e16e21

Reader analysis

Proposed checks: test conditional action calibration with and without clipping while holding the denoiser fixed, and compare gradual score guidance against direct policy resampling with matched initial noise. Measure both action mismatch and simulator replay MSE before attributing return changes to either mechanism. e09e12e18e19e22

8.2 Proposed reproduction checks

The following checks are proposals motivated by the paper. They have not been run as part of this reading.

Check 1: Can correct aggregate dispersion hide conditional action errors?

Reader-proposed check, not performed: freeze a trained denoiser and policy mean network, then compare standard deviations 0.2, 0.1, 0.06, and 0.02 using the same initial states and random seeds. Run the stated guidance adaptation with and without clipping, preserving Eq. (11)'s noise-after-clipping order. Report pooled standardized-action mean, variance, tail mass, and quantiles, then repeat the diagnostics within state regions defined before sampling. Compare each with direct policy samples at the corresponding predicted states. The incomplete-calibration concern is supported if pooled variance is near one while conditional quantiles or tails remain wrong; agreement across both controls and state regions would weaken it. e07e09e14e18

Check 2: Does gradual guidance improve both action consistency and dynamics accuracy?

Reader-proposed check, not performed: hold the h=10 denoiser, policy, 128-step diffusion schedule, clipping rule, and 500 buffer-sampled starting states fixed. Compare base score guidance, direct Policy Sampling, and unchanged Random Actions using matched initial noise. Record final conditional action-distribution error and state MSE after replaying each generated action sequence in MuJoCo, plus sampling time. Keep denoiser training data and checkpoint selection identical. The paper's explanation predicts that direct resampling disrupts consistency and worsens prediction; similar calibration and MSE under resampling would challenge that explanation. Only a subsequent matched A2C training comparison could establish whether these sampler differences cause the reported return gap. e05e12e16e19e22

8.3 Reading coverage

Visual audit: Visually inspected the title/version page; all main-body method, equation, result, ablation, and discussion pages; and all appendix pages containing implementation, additional diagnostics, baseline protocols, result processing, or compute claims. This includes Figures 1–10, 12, 14–18 and Tables 1–4. Cross-checked Figure 2's solid/hollow arrow meanings with Algorithm 2, clipping order with Eq. (11), and noisy-state notation with Algorithm 2 and Appendix C.4. All six final crops were opened and inspected; the control-performance crop was expanded and reinspected to preserve panel labels. Reference-only pages 14–17 were read in the complete text but were not visually inspected. Separate supplements and linked code remain unverified.

PDF pages inspected for this edition: 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 18, 19, 20, 21, 22, 23, 24, 25. Appendix coverage: reviewed.

Original figures and tables remain the work of the source’s authors. Extractions preserve their scientific content; any HTML wrapper layout is disclosed with each figure. The surrounding reading notes are our own.

Text reading scope & known omissions
  • Abstract
  • 1 Introduction
  • 2 Related Work
  • 3 Preliminaries
  • 4 World Models via Policy-Guided Trajectory Diffusion
  • 4.1 Denoising Model Training (Algorithm 1)
  • 4.2 Policy-Guided Trajectory Diffusion (Algorithm 2)
  • 4.3 Imagined RL in PolyGRAD World Models (Algorithm 3)
  • 5.1 Action distributions
  • 5.2 Trajectory accuracy
  • 5.3 Imagined RL performance
  • 5.4 Ablations and modifications
  • 6 Discussion
  • References
  • A Implementation Details
  • B.1 Prediction Error Plots with Transformer and MLP Denoising Networks
  • B.2 Confidence Intervals for Mean Performance
  • B.3 Extended Plots of PolyGRAD Action Distributions for All Environments
  • B.4 PolyGRAD Action Distributions with Action Clipping Ablated
  • C.1 Experimental Setup
  • C.2 Results Processing
  • C.3 Baselines
  • C.4 Ablations and Modifications
  • C.5 Computational Requirements

Outside the original text pass

  • Text extraction does not reconstruct figure images; inspect the retained PDF for figures and equation/table layout.
  • Separate supplemental material availability has not been fully verified.
  • Identity and edition scope: the inspected title page identifies arXiv:2312.08533v4, 27 March 2024, with a TMLR 02/2024 publication header. Title and all three authors match the catalog. Earlier revisions and a separate OpenReview PDF were not supplied, so differences between their contents and v4 cannot be established.
  • Text extraction does not reconstruct figure images; this limitation was addressed by inspecting the retained PDF's scientific pages and all final crops. References on PDF pp. 14–17 were read as text only.
  • Separate supplemental material availability has not been fully verified. All appendices included in the supplied 25-page PDF were read.
  • No code, linked repositories, or external sources were inspected, and no experiments were reproduced. The title page provides no affiliations; none are inferred.

The visual audit above records the subsequent illustrated pass.

8.4 Traceable evidence

e01PDF p. 1, title block, publication header, and arXiv margin stampInspect

Title and authors Marc Rigter, Jun Yamada, and Ingmar Posner match the catalog. The artifact identifies v4, 27 March 2024, and TMLR 02/2024; no affiliations appear in the title block.

Go to primary source ↓
e02PDF pp. 1–3, Introduction, Figure 1 and caption, Related Work: Diffusion Models for Decision-makingInspect

Whole-trajectory generation must reconcile future states and policy actions. A separate feedforward policy serves environment rollouts; diffusion produces its synthetic training data.

Go to primary source ↓
e03PDF p. 4, Section 3, Equations (1)–(4) and MDP definitionsInspect

Defines diffusion indices, noise schedule, noise-prediction objective, fully observable MDPs, and state-reward/action trajectory subsequences.

Go to primary source ↓
e04PDF p. 5, Section 4.1 and Algorithm 1Inspect

The denoiser predicts corruption of states and rewards conditioned on clean actions and diffusion level. The implementation also predicts terminal states, omitted from the notation.

Go to primary source ↓
e05PDF pp. 5–6, Section 4.2; p. 6, Algorithm 2 lines 2–14, Figure 2 and captionInspect

Initial-state inpainting anchors random trajectories. Denoised states condition policy-score action updates for i>1; the denoiser updates states/rewards. Solid arrows depict action changes and hollow arrows depict state changes.

Go to primary source ↓
e06PDF pp. 7–8, Section 4.2, Equations (5)–(8)Inspect

Langevin sampling motivates policy-score updates without a convergence proof. Classifier guidance assumes a uniform background action distribution and clean trajectories; it would guide both states and actions, whereas base PolyGRAD guides actions only.

Go to primary source ↓
e07PDF p. 8, Section 4.3, Algorithm 3 and Equations (9)–(10)Inspect

The loop collects real data, trains the denoiser, generates imagined data, and trains the policy. Guidance scale is adjusted using the standard deviation of policy-standardized actions.

Go to primary source ↓
e08PDF p. 9, Implementation DetailsInspect

Describes action clipping, unstable large policy updates, A2C/GAE, and an empirical preference for training on clean action conditioning despite noisy actions during generation.

Go to primary source ↓
e09PDF p. 9, Section 5.1, Figure 3 panels (a)–(f) and captionInspect

Walker2d h=50 action-minus-mean densities approximately match Gaussian targets for standard deviations at least 0.1; lower entropy produces mismatches and clipping-related modes at 0.02.

Go to primary source ↓
e10PDF pp. 10–11, Section 5.2; p. 10, Figures 4–5 and captionsInspect

Shared-data state-MSE comparisons use five seeds and h=10/50/200 PolyGRAD. Short-horizon PolyGRAD ranks best on Walker2d and second on the other tasks; long-horizon relative performance deteriorates. Timing uses 1,000 Walker2d trajectories on a V100; autoregressive diffusion repeats diffusion per transition.

Go to primary source ↓
e11PDF p. 11, Section 5.3, Figures 6–7, axes, legends, and captionsInspect

Residual-MLP h=10 imagined RL beats on-policy model-free baselines in aggregated final return but trails Dreamer-v3. Curves show five-seed standard deviations; IQM intervals are 95%. Hopper ends at 1M steps and the other tasks at 1.5M.

Go to primary source ↓
e12PDF pp. 11–12, Section 5.4; p. 12, Figures 8–9 and captionsInspect

Random actions and direct resampling perform poorly; removing clipping or using noisy state conditioning harms returns. State guidance is comparable to base PolyGRAD but raises reported cost from 54 to about 80 hours per 1M steps.

Go to primary source ↓
e13PDF pp. 12–13, Section 6, Limitations and Future WorkInspect

Slow policy updates and low-entropy action matching limit the method. Image-based and non-Markovian environments, convergence analysis, and improved long-horizon behavior are future directions.

Go to primary source ↓
e14PDF p. 18, Appendix A, Action Clipping During Diffusion, Equation (11)Inspect

The implemented action update clips the deterministic policy-score step to mean ±3 policy standard deviations and adds Gaussian noise outside the clip operation.

Go to primary source ↓
e15PDF p. 18, Appendix A, Policy Parameterisation, Initial State Conditioning, Imagined RL Training, Table 1Inspect

Specifies state-independent policy standard deviation, inpainting in training and inference, 1,024 h=10 imagined rollouts, GAE λ=0.9, critic LR 3e−4, Adam, γ=0.99, target log-probability change 0.01, entropy weight 1e−5, minimum std 0.1, and update ratio 0.25. Line search accepts changes within 20% of target.

Go to primary source ↓
e16PDF pp. 18–19, Appendix A, Denoising Model; p. 19, Tables 2–4 and Hyperparameter TuningInspect

Residual MLP: six layers, widths 1,024/2,048 for short/long horizons. Transformer: six layers, embedding 312, four causal attention heads. Both use batch 256; diffusion has 128 steps and cosine schedule, with Hopper/MLP schedule parameter 0.1 instead of 1.

Go to primary source ↓
e17PDF pp. 20–21, Appendices B.1–B.2, Figures 10, 12, and 14Inspect

MLP and transformer denoisers have similar h=10 prediction errors; transformer fares better for longer horizons. Mean-performance intervals support the same broad comparisons as IQM.

Go to primary source ↓
e18PDF pp. 21–23, Appendices B.3–B.4, Figures 15–18 and captionsInspect

h=10 action diagnostics span the three environments. Figure 17 reports HalfCheetah NaNs before std 0.02. Removing clipping yields non-Gaussian heavy tails even when overall standard deviation matches the policy.

Go to primary source ↓
e19PDF p. 23, Appendix C.1, Error Evaluation ExperimentsInspect

Models share 1M Algorithm-3 transitions and its final policy; evaluation starts from 500 uniformly sampled dataset states. Training selects best five-step error within 1M gradient steps or 72 RTX 3090 hours. A separate held-out rollout-selection split is not specified.

Go to primary source ↓
e20PDF p. 23, Appendix C.2, Smoothing of Reward Curves and Confidence IntervalsInspect

Ten-episode evaluations occur every 10,000 environment steps and are smoothed over ten evaluations. Final scores are divided by each environment's best algorithm mean and aggregated using 50,000 stratified bootstrap repeats; IQM uses the middle 50%.

Go to primary source ↓
e21PDF pp. 23–24, Appendix C.3, BaselinesInspect

Details the MLP ensemble, maximum-15-context autoregressive transformer, action-conditioned one-step diffusion, Dreamer's five-real-step latent initialization, and model-free implementations. Dreamer policy training backpropagates through dynamics; cited code versions are not pinned.

Go to primary source ↓
e22PDF pp. 24–25, Appendix C.4; p. 12, Noisy State Conditioning bullet; p. 6, Algorithm 2 line 10Inspect

Specifies each ablation and the additional state-score gradient. The noisy-state descriptions on pp. 12 and 25 print an action superscript while referring to states; Algorithm 2 explicitly conditions the policy on denoised states.

Go to primary source ↓
e23PDF p. 25, Appendix C.5, Computational RequirementsInspect

One base PolyGRAD RL run to 1M environment steps requires 54 hours on an RTX 3090 GPU.

Go to primary source ↓

8.5 Primary sources

Scroll across the image to inspect details. Press Esc to close.