| name | shortcut-trajectory-planning-offline-rl |
| description | Single-stage shortcut-model trajectory planner for efficient offline model-based RL. Replaces two-stage consistency-distillation with one-stage step-size-conditioned shortcut models + feasibility-aware critic for fast one/few-step generative planning. Use when building fast diffusion-based planners for offline RL (D4RL) without the training cost/instability of teacher-student distillation. |
Shortcut Trajectory Planning (STP) — from arXiv:2607.09336
Offline model-based RL framework that uses shortcut models as efficient trajectory generators,
avoiding the expensive two-stage teacher–student consistency-distillation pipeline common in
diffusion planners.
When to use
- You want diffusion-quality trajectory planning but the iterative denoising (50+ steps) is too slow
at inference.
- You want to avoid consistency-model distillation (two-stage, unstable, needs a teacher).
- Target domains: D4RL locomotion / navigation / manipulation / dexterous control.
Core idea
- Train a conditional shortcut trajectory model in a single stage. The model is conditioned on a
step-size parameter
s so it can do one-step (s large) or few-step (s small) inference from
the same weights — no separate distilled student.
- Generate K candidate plans (one/few-step), then select the best with a critic augmented by a
feasibility-aware correction term that penalizes physically infeasible trajectories.
Implementation steps
- Shortcut model: parameterize the trajectory generator
G(z, s; θ) where s ∈ (0,1] is the
step-size. Train with a single-stage objective that maps noise → trajectory such that large s
yields a coarse one-step guess and small s refines it. (Single-stage = no teacher needed.)
- Feasibility-aware critic: standard Q/critic
Q(s, a) plus a correction c(τ) that scores
trajectory feasibility from the dynamics model (e.g. constraint violation, terminal-state
reachability). Selection score = Q + λ·c(τ).
- Planning: sample
K latent z, generate τ_k = G(z_k, s) at chosen s (1 or few steps),
pick argmax_k selection score. Roll out in the environment / model.
- Sweep
s to trade latency vs plan quality.
Key advantages over consistency distillation
- One training stage instead of two → simpler, cheaper, more stable.
s-conditioning gives adjustable compute at inference with NO extra models.- Feasibility correction prevents the shortcut from emitting infeasible plans that a raw diffusion
sampler might.
Pitfalls
- Shortcut models can under-perform full diffusion at very few steps on hard dynamics; calibrate
s
per task (few-step often best tradeoff).
- Feasibility correction must be cheap (use the learned dynamics model, not ground-truth sim rolls).
- Offline distribution shift: shortcut model may extrapolate; keep
K modest and rely on the critic.
- Single-stage training needs careful noise scheduling; reuse the diffusion scheduler's sigma curve.
Verification
- Report planning latency (steps) vs return on D4RL benchmarks vs a consistency-distilled baseline
at matched compute.
- Ablate the feasibility correction: show it recovers infeasible-plan rejection.
