Amortizing Trajectory Diffusion with Keyed Drift Fields

This repository is the official (JAX/Flax) implementation of Keyed Drifting Policies (KDP) for D4RL locomotion datasets. It trains a one-step conditional trajectory generator in normalized trajectory-window space and uses a learned return scorer for receding-horizon action selection.

Website: https://keyed-drifting.github.io/
Paper: https://keyed-drifting.github.io/paper.pdf

• Code Layout
• Environment
• CUDA Troubleshooting
• Training Flow
• HalfCheetah Medium
• Hopper Medium
• Walker2d Medium
• Convenience Wrapper
• Checkpoints
• W&B
• Tests
• Acknowledgements
• Citing this work

rl/
  d4rl_dataset.py       D4RL loading, episode splitting, windows, RTG labels
  normalizer.py         Observation/action normalization helpers
  kdp_loss.py           Keyed drifting loss in full trajectory space
  train_scorer.py       Supervised return-to-go scorer training
  train_kdp.py          KDP generator training loop with EMA checkpoints
  planner.py            Receding-horizon KDP candidate planner
  eval_kdp.py           D4RL evaluation loop
  diagnose_kdp.py       Checkpoint/data diagnostics
  preflight.py          JAX/MuJoCo/D4RL setup check before training
  checkpoints.py        Pickle/JSON checkpoint helpers
  config_utils.py       YAML loading and schedules
  configs/              Canonical D4RL configs
  tests/                Unit tests for dataset, loss, clamp, planner
models/
  trajectory_generator.py
  return_scorer.py
scripts/
  run_kdp_pipeline.sh   Convenience scorer + generator training wrapper

Use Python 3.10. D4RL still depends on the old Gym/mujoco-py stack, so Python 3.11+ is not recommended. If your shell prompt shows Python 3.12, that is fine as long as uv sync creates and uses the project .venv with Python 3.10.

For NVIDIA GPU training, install the CUDA-enabled JAX extra. This uses the official jax[cuda12] package extra, which pulls in the CUDA/cuDNN/NCCL wheels JAX needs:

uv python install 3.10
uv sync --extra cuda12 --extra dev
./rl/kdp_env.sh python -c "import jax; print(jax.devices())"

The device check must print a CudaDevice. If it prints only CpuDevice, do not start training yet (I am highlighting this because it runs into a lot of downstream errors. PRs welcome to fix this). Then run the full preflight check:

./rl/kdp_env.sh python -m rl.preflight --env-name halfcheetah-medium-v2

Preflight checks Python 3.10, JAX GPU visibility, MuJoCo 2.1, mujoco-py, D4RL task registration, and env.get_dataset(). Run it once for each task before launching long training.

For a new machine, run this short scorer smoke test before the full training command:

./rl/kdp_env.sh python -m rl.train_scorer \
  --config rl/configs/halfcheetah_medium.yaml \
  --workdir /tmp/kdp_halfcheetah_scorer_smoke \
  --steps 1 \
  --batch-size 8

That command should load the D4RL dataset, print one JSON metric row, and save /tmp/kdp_halfcheetah_scorer_smoke/checkpoint.pkl.

For CPU-only development, omit --extra cuda12, but training will be very very very slow:

uv sync --extra dev

D4RL MuJoCo tasks also require MuJoCo 2.1 for mujoco-py. A typical Linux setup uses (you might need to change this based on your mujoco setup):

mkdir -p ~/.mujoco
# Put MuJoCo 2.1 at ~/.mujoco/mujoco210
export MUJOCO_PY_MUJOCO_PATH="$HOME/.mujoco/mujoco210"
export LD_LIBRARY_PATH="$MUJOCO_PY_MUJOCO_PATH/bin:/usr/lib/nvidia:/usr/lib/x86_64-linux-gnu:${LD_LIBRARY_PATH:-}"

The helper wrapper sets the D4RL/MuJoCo environment variables and activates .venv if it exists. It intentionally does not set D4RL_SUPPRESS_IMPORT_ERROR, because that variable can hide the real MuJoCo import failure and lead to Gym reporting only Environment ... does not exist.

If JAX logs Unable to load CUDA and falls back to CPU, the environment is missing a compatible NVIDIA driver/runtime path or the CUDA extra was not installed correctly.

Check the driver first:

nvidia-smi

Then recreate/sync the environment with the CUDA extra:

uv sync --extra cuda12 --extra dev --reinstall
./rl/kdp_env.sh python -c "import sys, jax; print(sys.executable); print(jax.devices())"

If that still prints CpuDevice, remove local CUDA paths that may override the pip-installed JAX CUDA wheels and retry:

unset CUDA_HOME CUDA_PATH
export LD_LIBRARY_PATH=""
./rl/kdp_env.sh python -c "import jax; print(jax.devices())"

JAX's official guidance is to install GPU support via jax[cuda12]/jax[cuda13] pip extras and to avoid incompatible LD_LIBRARY_PATH entries that override those wheel-provided CUDA libraries.

If preflight reports No module named 'mjrl', your environment was created from an older lockfile. Pull the latest pyproject.toml and uv.lock, then run:

uv sync --extra cuda12 --extra dev --reinstall

The lockfile pins mjrl from its upstream Git commit and pins setuptools<81 so D4RL's legacy pkg_resources imports work.

Each task uses two checkpoints:

1. Train a return scorer on normalized D4RL windows.
2. Train the KDP trajectory generator using the scorer checkpoint for validation metadata and later evaluation.
3. Evaluate with the scorer as a candidate ranker.

The canonical configs use horizon H=32, G=8 generated samples per condition, top-50% return window sampling, residual-noise generation, action upweighting, and EMA generator checkpoints.

./rl/kdp_env.sh python -m rl.preflight --env-name halfcheetah-medium-v2

./rl/kdp_env.sh python -m rl.train_scorer \
  --config rl/configs/halfcheetah_medium.yaml \
  --workdir runs/kdp_halfcheetah_medium/scorer

./rl/kdp_env.sh python -m rl.train_kdp \
  --config rl/configs/halfcheetah_medium.yaml \
  --scorer runs/kdp_halfcheetah_medium/scorer/checkpoint.pkl \
  --workdir runs/kdp_halfcheetah_medium/generator

./rl/kdp_env.sh python -m rl.eval_kdp \
  --config rl/configs/halfcheetah_medium.yaml \
  --generator runs/kdp_halfcheetah_medium/generator/best.pkl \
  --scorer runs/kdp_halfcheetah_medium/scorer/checkpoint.pkl \
  --episodes 10

Base KDP training:

./rl/kdp_env.sh python -m rl.preflight --env-name hopper-medium-v2

./rl/kdp_env.sh python -m rl.train_scorer \
  --config rl/configs/hopper_medium.yaml \
  --workdir runs/kdp_hopper_medium/scorer

./rl/kdp_env.sh python -m rl.train_kdp \
  --config rl/configs/hopper_medium.yaml \
  --scorer runs/kdp_hopper_medium/scorer/checkpoint.pkl \
  --workdir runs/kdp_hopper_medium/generator

You can finetune it like this:

./rl/kdp_env.sh python -m rl.train_kdp \
  --config rl/configs/hopper_medium_raw_finetune.yaml \
  --scorer runs/kdp_hopper_medium/scorer/checkpoint.pkl \
  --workdir runs/kdp_hopper_medium/generator_raw_ft \
  --init-generator runs/kdp_hopper_medium/generator/step_10000.pkl \
  --init-use-raw-params

./rl/kdp_env.sh python -m rl.eval_kdp \
  --config rl/configs/hopper_medium_raw_finetune.yaml \
  --generator runs/kdp_hopper_medium/generator_raw_ft/step_5000.pkl \
  --scorer runs/kdp_hopper_medium/scorer/checkpoint.pkl \
  --episodes 10 \
  --num-candidates 256 \
  --use-raw-params

./rl/kdp_env.sh python -m rl.preflight --env-name walker2d-medium-v2

./rl/kdp_env.sh python -m rl.train_scorer \
  --config rl/configs/walker2d_medium.yaml \
  --workdir runs/kdp_walker2d_medium/scorer

./rl/kdp_env.sh python -m rl.train_kdp \
  --config rl/configs/walker2d_medium.yaml \
  --scorer runs/kdp_walker2d_medium/scorer/checkpoint.pkl \
  --workdir runs/kdp_walker2d_medium/generator

./rl/kdp_env.sh python -m rl.eval_kdp \
  --config rl/configs/walker2d_medium.yaml \
  --generator runs/kdp_walker2d_medium/generator/best.pkl \
  --scorer runs/kdp_walker2d_medium/scorer/checkpoint.pkl \
  --episodes 10

All the config files for the above examples are provided (including the finetuning config).

The wrapper trains the scorer if missing and then trains the generator:

./scripts/run_kdp_pipeline.sh rl/configs/halfcheetah_medium.yaml runs/kdp_halfcheetah_medium
./scripts/run_kdp_pipeline.sh rl/configs/hopper_medium.yaml runs/kdp_hopper_medium
./scripts/run_kdp_pipeline.sh rl/configs/walker2d_medium.yaml runs/kdp_walker2d_medium

Generator checkpoints contain:

• raw generator params
• EMA generator params
• model config
• dataset normalizer statistics
• RTG normalization statistics
• original training config

Scorer checkpoints contain scorer params plus the same normalizer and RTG stats. Evaluation verifies shape and normalizer compatibility before planning.

The training scripts print JSON metrics to stdout. To capture them in W&B without changing training code, run training in one terminal and stream the log file in another:

./rl/kdp_env.sh python -m rl.train_kdp ... 2>&1 | tee runs/train.log
./rl/kdp_env.sh python -m rl.wandb_tail --logfile runs/train.log --project kdp-d4rl --name halfcheetah

The tests were generated by codex.

./rl/kdp_env.sh python -m pytest

These tests cover window boundaries, clamp behavior, KDP loss shapes/finite gradients, initial action diversity, and planner action unnormalization/clipping.

The code base is loosely based off of https://github.com/lambertae/drifting.

@article{puthumanaillam2026amortizing,
  title={Amortizing Trajectory Diffusion with Keyed Drift Fields},
  author={Puthumanaillam, Gokul and Ornik, Melkior},
  journal={arXiv preprint arXiv:2603.14056},
  year={2026}
}