TorchRL 0.15: Dataset Exports, TD-MPC2, Robot Deployment, and Flow Q-Learning

TorchRL v0.15.0

TorchRL 0.15 introduces powerful new capabilities for robot learning, offline RL, data interoperability, and training at scale. New primitives include ReplayBuffer.export with Hugging Face and LeRobot adapters, export_onnx_policy, MenagerieEnv, LowLevelController, TransformerModule, RewardModelLoss, StopOnSignal, and the TD-MPC2, Flow Q-Learning and Transitive RL objectives.

Highlights

• Hugging Face and LeRobot dataset exports: ReplayBuffer.export writes stored data in storage order or by trajectory and step through exporter adapters, with HuggingFaceDatasetAdapter for Hugging Face datasets and LeRobotDatasetAdapter for the LeRobot format. ReplayBuffer.as_dataset() lets a torch.utils.data.DataLoader read any replay buffer, and a tutorial walks through exporting MicroDuck rollouts to LeRobot. #4592, #4611 and #4628 by @bsprenger, #4387 by @theap06.
• TD-MPC2 ships as a modular trainer stack: world model, Q-ensemble, loss, planner, optimization stepper and TdMpc2Trainer, each with a Hydra config and usable on its own, plus a runnable recipe under sota-implementations/tdmpc2_trainer. #4457 and the stack it builds on by @bsprenger.
• ONNX export and robot deployment: export_onnx_policy writes feed-forward and recurrent TensorDict policies to ONNX with inputs and outputs named after their keys, MicroDuckEnv(policy_observation=True) matches the observation the robot runtime feeds its policy, MenagerieEnv loads any MuJoCo Menagerie robot by name, and LowLevelController lets a high-level policy drive frozen skills. #4626, #4625 and #4386 by @theap06, #4329 by @vmoens.
• Flow Q-Learning: FQLLoss, FQLTrainer, flow-matching and one-step policies, and a recipe under sota-implementations/fql. OfflineToOnlineTrainer is generalized so FQL can pretrain offline before training online. 0bd9892b and 5e402a4e by @neuriv.

Breaking changes

• TorchRL now requires PyTorch 2.11 or later and tensordict>=0.15.2,<0.16.0. #4595 and #4600 by @theap06.
• PyTorch 2.15 requires Python 3.11 or later, so the Python 3.10 wheels of TorchRL 0.15 are built against PyTorch 2.14.1. pip install torchrl on Python 3.10 still installs 0.15, with PyTorch 2.14. #4706 by @theap06.
• EnvBase.reset and EnvBase.rollout now default to set_state=False, completing the change announced in v0.14. Pass set_state=True if you relied on state entries in the input being applied. #4501 by @coder-jayp.

Deprecations and future defaults

• The default changes v0.14 announced for v0.15 move to v0.16: auto_register_policy_transforms=True, the AsyncBatchedCollector and AsyncEnvPool backend defaults, and the unified Trainer checkpoint format. Their FutureWarnings stay until then. #4698 by @theap06.
• SharedMemWeightSyncScheme(atomic=True) gives workers private weight copies that switch only on update, so they never act on half-written weights. It becomes the default in v0.16; until then the default warns. #4695 by @theap06.
• PyTorch 2.11 and 2.12 emit a DeprecationWarning; v0.17 will require PyTorch 2.13. #4595 by @theap06.
• PairwiseDataset and GPT2RewardModel are deprecated in favor of RewardModelLoss and will be removed in v0.17. #4434 and #4442 by @theap06.
• A stateful transform placed after a transform whose inverse drops its keys now warns, and will raise in v0.17. #4571 by @itwasabhi.

Features

Algorithms and objectives

• TD-MPC2 components: TdMpc2QEnsemble, TdMpc2Loss, TdMpc2Planner, TdMpc2OptimizationStepper and TdMpc2Trainer, following arXiv:2310.16828. Trainer.optim_steps() accepts per-call overrides for schedules such as pretraining. #4449, #4451, #4452, #4453, #4454, #4455, #4456, #4457 and #4424 by @bsprenger.
• Flow Q-Learning: FlowMatchingPolicy, OneStepPolicy, FQLLoss, FQLTrainer and a recipe under sota-implementations/fql. 0bd9892b and 5e402a4e by @neuriv.
• Transitive RL: TransitiveQLoss (arXiv:2510.22512), matched numerically against the authors' JAX code, the SubgoalTriplet sample unit and an OGBench recipe. #4596, #4597 and #4598 by @theap06.
• RewardModelLoss, a Bradley-Terry pairwise loss for any sequence-scoring module, with a reward-model training recipe. #3922 and #4434 by @theap06.
• PPOTrainer.from_env builds a standard PPO setup from an env and actor/critic networks; GAE accepts a validity mask and GAE.for_recurrent(...) builds a recurrent-safe estimator. #4391, #4404 and #4413 by @vmoens.

Policies, robots and environments

• TransformerModule runs a causal transformer over [B, T] windows for training or one step at a time against a key/value cache for collection, with a reference CausalTransformer backbone. #4193 and #4432 by @theap06.
• export_onnx_policy exports feed-forward, LSTMModule and GRUModule policies to ONNX with inputs and outputs named after their keys (pip install torchrl[onnx]), and MicroDuckEnv(policy_observation=True) matches the robot runtime's observation. #4626 and #4625 by @theap06.
• Hierarchical control: LowLevelController, ClosedLoopMultiAction, MicroDuckSkillEnv and a skill zoo with MicroDuckSkills.from_pretrained, plus new MicroDuck head-level reward, turning and jumping tasks. #4329, #4361, #4392, #4393, #4394, #4395 and #4396 by @vmoens.
• MenagerieEnv loads any MuJoCo Menagerie robot by name, with a task config and a zero default reward that a transform can replace. #4386 by @theap06.
• New transforms: LastAction and a public DoneTransform for group-level done flags. #4360 and #4356 by @YeonwooSung.

Replay buffers and data

• Storage.as_dataset() and ReplayBuffer.as_dataset() for torch.utils.data.DataLoader, with tensordict_collate. #4387 by @theap06.
• ReplayBuffer.export with HuggingFaceDatasetAdapter and LeRobotDatasetAdapter. #4592 and #4611 by @bsprenger.
• RateLimitedReplayBuffer(samples_per_insert=...), ReplayBuffer.wait_until_sampleable() and sample(wait=True), and SliceSampler(output_layout="batch_time"). #4418, #4408 and #4414 by @vmoens.

Collectors, evaluation and training runs

• Collectors take replay_write_mode="rollout" or "trajectory", report in-flight statistics, and expose worker_policy_versions(). #4412, #4416 and #4410 by @vmoens.
• Evaluator adds per-env episode sampling, seeded paired episodes and bootstrap confidence intervals; EvaluatorHook runs it on a schedule from a Trainer; LogValidationReward can record complete episodes. #4500 by @theap06, #4415 by @vmoens and #4482 by @quinnarnold.
• StopOnSignal turns SIGINT/SIGTERM into a stop and a final checkpoint; Trainer checkpoints include RNG state and load the newest file from a rotation directory; recipes resume with resume=<dir>. #4380, #4381, #4384 and #4385 by @vmoens.
• On-policy trainers log standard telemetry by default (frame and episode counts, reward summaries, learning rate, gradient norm and throughput), Logger.with_prefix("name") places metrics under a namespace, and the SOTA recipes log under training/, evaluation/ and timing/. #4417, #4440, #4403 and #4420 by @vmoens.

Performance

• ParallelEnv workers that report their own metadata now use shared buffers instead of pipes. #4341 by @vmoens.
• Bounded epsilon-greedy exploration and truncated-normal log-probabilities compile with fullgraph=True. #4645 and #4648 by @bsprenger.

Bug fixes

• Training correctness (these change loss values or sampled data): CQLLoss computes the CQL(H) regularizer and Lagrange update of the paper and reference code; CQLLoss, DiscreteCQLLoss, IQLLoss, DiscreteIQLLoss, REDQLoss, CrossQLoss and QMixerLoss apply prioritized-replay weights; HardUpdate copies every value_network_update_interval steps instead of every interval + 1; and PrioritizedSliceSampler respects episode boundaries after a multi-env storage wraps and sets is_init at slice starts. #4693, #4694 and #4700 by @theap06.
• Collectors and weight sync: MultiProcessWeightSyncScheme workers no longer read the trainer's live CPU weights and see new weights only after update_policy_weights_(), Collector(AsyncEnvPool(...)) no longer crashes on the first batch, AsyncEnvPool(exchange="shm") works with outer transforms and partial resets, syncing a non-policy model no longer hangs, and multi-process collectors with a replay buffer and a batched env no longer crash. #4701 by @theap06. Ray collector.stats() responds while collecting. #4400 by @coder-jayp.
• Environments and transforms: offline RewardSum sums along time (#4539 by @aalquwayfili), StepCounter with several done groups (#4587 by @mottopanikeiku), RandomCropTensorDict last window (#4578 by @jackiectl2), spec fixes (#4548 by @coder-jayp, #4556 by @YeonwooSung, #4591 by @mottopanikeiku).
• Replay buffers: SliceSampler span edge cases (#4700 by @theap06), compact checkpoints of partially filled buffers (#4644 by @theap06), memory-mapped checkpoints under TensorDict 0.16+ (#4551 by @peterdsharpe).
• Losses and recipes: MultiAgentMLP export with shared parameters (#4700 by @theap06), MCAdvantage on constant groups (#4610 by @catherinelee274), compiled DreamerV3 reward decoding on CUDA (#4657 by @arun-elr), TD3 with CUDA graphs (#4636 by @peterdsharpe), detached optimizer metrics (#4561 by @bsprenger), IncrementalTokenizer prompts after tool calls (#4674 by @arun-elr).
• TensorDict 0.17 compatibility. #4576, #4637 and #4638 by @peterdsharpe.

Documentation

• A MicroDuck LeRobot export tutorial, a checkpointing and resume guide, and a corrected TRL interoperability tutorial. #4628 by @bsprenger, #4230 by @YeonwooSung and #4161 by @coder-jayp.
• Refreshed README highlights. #4678 and #4703 by @bsprenger.

Upgrade notes

• Install matching release lines together: pip install "torchrl==0.15.0" "tensordict>=0.15.2,<0.16.0".
• Linux CUDA wheels are attached to this release rather than published to PyPI. Install one by URL, for example pip install "https://github.com/pytorch/rl/releases/download/v0.15.0/torchrl-0.15.0+cu132-cp312-cp312-manylinux_2_28_x86_64.whl".
• CQL users should expect different loss_cql values. Users of the CQL, IQL, REDQ, CrossQ and QMixer losses with prioritized replay can pass use_prioritized_weights=False to reproduce v0.14 numbers.
• Dashboards keyed on r_training, done_percentage, train/, eval/ or time/ need the new metric names, or pass telemetry="minimal" to the on-policy trainers.
• Multi-process collector users can pass SharedMemWeightSyncScheme(atomic=True) to opt into the v0.16 behavior now, or atomic=False to keep the current one without a warning.

Full changelog

For the complete commit-by-commit history, see v0.14.0...v0.15.0.

Contributors

Thanks to @vmoens, @theap06, @bsprenger, @peterdsharpe, @coder-jayp, @huydhn, @YeonwooSung, @neuriv, @arun-elr, @mottopanikeiku, @sramireddy2, @aalquwayfili, @catherinelee274, @itwasabhi, @jackiectl2 and @quinnarnold for the changes included in this release.

We especially welcome first-time contributors @peterdsharpe, @neuriv, @arun-elr, @mottopanikeiku, @sramireddy2, @aalquwayfili, @catherinelee274 and @jackiectl2.