Muninn: Your Trajectory Diffusion Model But Faster

🏆 News: Muninn won the Outstanding Student Paper Award at RSS 2026!

Diffusion-based trajectory planners can synthesize rich, multimodal robot motions, but their iterative denoising makes online planning and control prohibitively slow. Existing accelerations either modify the sampler or compress the network — sacrificing plan quality or requiring retraining without accounting for downstream control risk. Our key insight is that diffusion trajectory planners expose two signals we can exploit: a cheap probe of how their internal trajectory representation changes across steps, and analytic coefficients that describe how denoiser errors affect the sampler's state update. By calibrating the first signal against the second on offline runs, we obtain a per-step score that upper-bounds how far the final trajectory can deviate when we reuse a cached denoiser output, and we treat this bound as an uncertainty budget spent over the denoising process. Muninn is a training-free caching wrapper that tracks this budget during sampling and, at each diffusion step, chooses between reusing a cached denoiser output when the predicted deviation is small and recomputing the denoiser when it is not — delivering large wall-clock speedups while preserving task performance, and certifying, at a user-chosen deviation tolerance and risk level, that cached rollouts remain within a specified distance of their full-compute counterparts.

This repository is the official implementation of Muninn: the model-agnostic muninn package together with its Diffuser integration, with worked examples on D4RL locomotion (HalfCheetah and Walker2d) — a ~90 line adapter, a drop-in guided policy, and the full calibrate → certify → evaluate pipeline.

Integrations: Diffuser (Janner et al., 2022) — value-guided planning on halfcheetah-medium-v2 and walker2d-medium-v2.

teaser.mp4

Table of contents

• Results
• Installation
• Quickstart: the Diffuser benchmark
• Integrating your own diffusion model
• Repository layout
• Acknowledgments
• Citation

Results

See RESULTS.md for the result figures.

Installation

The core package is pure Python and depends only on numpy, torch, and scikit-learn:

uv add muninn                 # or, from a checkout of this repo:
uv pip install -e .

That is all you need to wrap your own diffusion planner (see Integrating your own diffusion model).

To reproduce the Diffuser benchmark you additionally need the MuJoCo / D4RL stack. This is a one-shot script:

# prerequisites (cannot be automated):
#   - uv                (https://docs.astral.sh/uv/)
#   - MuJoCo 2.1 binary at ~/.mujoco/mujoco210
#   - an NVIDIA GPU (the script installs CUDA 12.8 torch wheels)
cd integrations/diffuser
bash setup_env.sh             # .venv at repo root: torch, mujoco-py, d4rl, diffuser, muninn
source env.sh                 # activate + MuJoCo env vars

Pretrained Diffuser checkpoints go into integrations/diffuser/logs/pretrained (see integrations/diffuser/README_DIFFUSER.md, "Downloading weights"). D4RL datasets download automatically on first use.

Quickstart: the Diffuser benchmark

Everything below runs from integrations/diffuser/ (details in its README). The same five commands work for both environments — swap --dataset halfcheetah-medium-v2 / walker2d-medium-v2 and the η value (4.2 / 10.3).

# 0. correctness checks: full-compute mode must match the original Diffuser
#    guided sampler bit-for-bit, and the budget mechanics must be exact
python scripts/muninn_test.py      --dataset halfcheetah-medium-v2 --logbase logs/pretrained --vis_freq -1

# 1. calibrate: full-compute rollouts with ghost-reuse bookkeeping
#    -> per-timestep split-conformal error bounds U_t (saved as .npz)
python scripts/muninn_calibrate.py --dataset halfcheetah-medium-v2 --logbase logs/pretrained --vis_freq -1

# 2. inspect the calibrated budget scale and pick eta (the speed-fidelity knob)
python scripts/muninn_eta_probe.py --dataset halfcheetah-medium-v2 --logbase logs/pretrained --vis_freq -1

# 3. evaluate Full vs Muninn (paired deviation, latency, evals/step, D4RL score)
python scripts/muninn_eval.py      --dataset halfcheetah-medium-v2 --logbase logs/pretrained --vis_freq -1 --mode full
python scripts/muninn_eval.py      --dataset halfcheetah-medium-v2 --logbase logs/pretrained --vis_freq -1 --mode muninn --eta 4.2

# 4. aggregate the results table
python scripts/muninn_report.py logs/pretrained/*/muninn/eval_*.json

Useful knobs (all CLI flags of muninn_eval.py / muninn_calibrate.py):

flag	meaning	default
--eta	trajectory deviation budget η (higher = more reuse)	env-specific
--alpha	risk level α of the certificate	0.05
--prefix_forbid / --suffix_forbid	first/last steps where reuse is forbidden	2 / 2
--n_calib_episodes	calibration rollouts	3
--n_eval_episodes	evaluation rollouts (unseen seeds)	10
--paired_every	measure paired deviation every k-th planning call	10

Integrating your own diffusion model

Muninn is model-agnostic: the core loop only talks to your sampler through muninn.DiffusionAdapter. Implement five methods (plus two optional ones) and you get calibration, the certificate, and cached sampling for free:

import muninn

class MyAdapter(muninn.DiffusionAdapter):
    n_timesteps = ...      # T
    device = ...
    gamma = ...            # metric constant: d(a, b) = gamma * ||a - b||_F

    def sample_shape(self, batch_size): ...   # e.g. (B, horizon, transition_dim)
    def alphas_cumprod(self): ...             # DDPM schedule, numpy [T]
    def probe(self, x, t): ...                # cheap stem features [B, d_F]
    def compute(self, x, t, ctx): ...         # denoiser (+ guidance) -> (x, model_out, aux)
    def posterior(self, x, model_out, t): ... # affine update -> (mean, log_var)
    # optional:
    # condition(x, ctx)  – projection applied after every update (inpainting)
    # subset(ctx, idx)   – restrict the context to a sub-batch

Calibrate once, offline — run your planner at full compute on rollouts from the deployment distribution; the loop records (score, would-be reuse error) pairs for you:

adapter = MyAdapter(...)

records = []
for ctx in calibration_contexts:
    muninn.cached_sample_loop(adapter, batch_size, ctx, record=records)

calib = muninn.Calibrator.fit(records, alpha=0.05,
                              t_cache=range(2, adapter.n_timesteps - 2))
calib.save('calibrator.npz')

Deploy — same loop, now with a budget:

result, info = muninn.cached_sample_loop(
    adapter, batch_size, ctx,
    calibrator=muninn.Calibrator.load('calibrator.npz'),
    eta=..., return_info=True)

trajectories = result.trajectories        # [B, ...] in original batch order
print(info.n_evals.mean(), 'denoiser evals instead of', adapter.n_timesteps)

Design notes that make an adapter easy to get right:

• probe should be the cheapest prefix of your denoiser that still sees the trajectory and timestep — the stem / first block, pooled to a vector. A few percent of a full forward is the target cost.
• compute is everything expensive in one diffusion step (denoiser plus any guidance): a Muninn reuse step skips it entirely, and calibration measures exactly that reuse rule, so guidance is covered by the bound.
• gamma fixes the deviation metric. We use d = ||·||_F / √H (gamma = 1/√H), the per-step RMS deviation of the trajectory.
• For paired experiments (deviation of cached vs full-compute plans), pass the same noise_pack (from muninn.make_noise_pack) to both calls.

The complete worked example — including value-function guidance and the state-inpainting projection — is integrations/diffuser/muninn_diffuser/adapter.py.

Repository layout

muninn/                        the core library (model-agnostic)
├── adapter.py                 DiffusionAdapter interface
├── sampler.py                 budgeted caching loop (Algorithm 1)
├── calibration.py             split-conformal per-step error bounds
└── coefficients.py            analytic DDPM sensitivity coefficients
integrations/
└── diffuser/                  Diffuser benchmark
    ├── muninn_diffuser/       DiffuserAdapter + drop-in MuninnGuidedPolicy
    ├── scripts/muninn_*.py    test / calibrate / eta_probe / eval / report
    ├── setup_env.sh, env.sh   one-shot uv environment
    └── diffuser/, config/     original Diffuser code (compat patches only)
RESULTS.md                     result figures
assets/                        figures

Acknowledgments

The benchmark backend is Diffuser by Janner et al. (MIT license), included under integrations/diffuser with only minor compatibility patches.

Citation

@article{puthumanaillam2026muninn,
  title  = {Muninn: Your Trajectory Diffusion Model But Faster},
  author = {Puthumanaillam, Gokul and Jiang, Hao and Hernandez, Ruben and
            Fuentes, Jose and Padrao, Paulo and Bobadilla, Leonardo and
            Ornik, Melkior},
  year   = {2026},
}