Summary: Value-guided diffusion uses a learned value function to steer the diffusion denoising process toward high-reward action sequences. This enables long-horizon planning (50+ steps) by combining diffusion's multi-modal generation with value-based optimization — the core of test-time compute scaling for robotics.
Core Idea
Diffusion Policy generates action sequences by denoising:
a_{0:H} ← Denoise(epsilon, obs, t)
Value-guided diffusion adds value gradient during denoising:
Score = ∇_{a} log p(a mid obs) + lambda ∇_{a} V(a, obs)
Effect: Denoising steps push samples toward high-value regions while preserving diversity.
Architecture
Observation → Encoder → Conditioning
↓
Action Chunk → DiT (Denoiser) → Clean Action Chunk
↑ ↑
Timestep t Value Gradient:
+ Noise λ ∇V(a)
Training
| Stage | Objective |
|---|---|
| 1. Diffusion | Standard denoising loss: ` |
| 2. Value | TD learning: V_psi(a, obs) ≈ r + gamma V_psi(a', obs') |
| 3. Joint (optional) | Combine denoising + value prediction in single model |
Inference (Planning)
# Value-guided diffusion sampling
def plan(obs, horizon=50, steps=20):
a = torch.randn(horizon, action_dim) # Start from noise
for t in reversed(range(steps)):
# Standard denoising step
a = denoise_step(a, t, obs)
# Value guidance
if t < guidance_steps:
with torch.enable_grad():
a.requires_grad_(True)
v = value_fn(a, obs)
grad = torch.autograd.grad(v.sum(), a)[0]
a = a + guidance_scale * grad
return a
Why It Works for Long Horizon
| Challenge | Standard Diffusion | + Value Guidance |
|---|---|---|
| Horizon | 16-32 (action chunking) | 50-100+ |
| Mode collapse | Averages modes | Preserves modes + biases to good |
| Contact precision | OK for short | Maintained via value |
| Credit assignment | None | Value = long-horizon signal |
With Mamba/SSM Dynamics
Stack:
Obs → Mamba Dynamics → Latent Trajectory → Value Function
↑
Diffusion Policy ← Action Diffusion ←──┘
Mamba/SSM provides:
- Linear recall (efficient long context)
- Parallel training (like Transformer)
- 50-100 step rollouts in latent space
Value-guided diffusion provides:
- Multi-modal action generation
- Test-time scaling (more denoising steps = better plans)
Related Work
| Method | Guidance | Horizon |
|---|---|---|
| Diffusion Policy (Chi et al.) | None | 16 |
| Value-Guided Diffusion (Janner et al.) | Value gradient | 50+ |
| Diffuser (Ajay et al.) | Reward classifier guidance | 100+ |
| VDP (Liu et al.) | Process reward (step-level) | 50+ |
| HRM + Diffusion | Hierarchical value | 50+ |
Related Concepts
- diffusion-policy — Base algorithm
- planning-in-latent-space — Latent planning
- inference-time-scaling — Test-time compute
- mamba-ssm — Efficient long-horizon dynamics
- value-function — Guidance signal
- diffusion-model — Generative base
Sources
- Janner et al., "Planning with Diffusion for Flexible Behavior Synthesis" (ICML 2022)
- Ajay et al., "Is Conditional Generative Modeling all you need for Decision Making?" (ICLR 2023)
- world-models-diffusion-ai-robotics — Decision guide
- planning-in-latent-space — Planning stack