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


Sources