docs/source/bring_your_own_policies.mdx
This guide walks you through implementing a custom policy and getting it to work with LeRobot's training, evaluation, and deployment tools. There are two paths:
lerobot_policy_* package. Faster, no PR required, easy to iterate. Right for experimentation, internal use, or when you want to publish independently.src/lerobot/policies/. Requires a PR, but makes your policy a first-class citizen of the library.The plugin route is usually the right starting point — promote to in-tree once the policy has stabilized and there's clear value in shipping it with the library.
Either way, the building blocks are the same: a configuration class, a policy class, and a processor factory. The first half of this guide covers those shared pieces; the second half covers the path-specific scaffolding (Path A, Path B).
A note on tone: robot-learning is an actively evolving field, and "what a policy looks like" can shift with each new architecture. The conventions described here exist because they let lerobot-train and lerobot-eval work uniformly across very different models. When a new policy genuinely doesn't fit them, raise it (in your PR, or an issue) — the conventions are not sacred.
Three building blocks make up every policy. The names below use my_policy as a placeholder — replace with your policy's name. That name is load-bearing: it must match the string you pass to @PreTrainedConfig.register_subclass, the MyPolicy.name class attribute, and the make_<name>_pre_post_processors factory function (more on each below).
Inherit from PreTrainedConfig and register your policy type. Here is a template — customize the parameters and methods as needed for your policy's architecture and training requirements.
# configuration_my_policy.py
from dataclasses import dataclass, field
from lerobot.configs import PreTrainedConfig
from lerobot.optim import AdamWConfig
from lerobot.optim import CosineDecayWithWarmupSchedulerConfig
@PreTrainedConfig.register_subclass("my_policy")
@dataclass
class MyPolicyConfig(PreTrainedConfig):
"""Configuration class for MyPolicy.
Args:
n_obs_steps: Number of observation steps to use as input
horizon: Action prediction horizon
n_action_steps: Number of action steps to execute
hidden_dim: Hidden dimension for the policy network
# Add your policy-specific parameters here
"""
horizon: int = 50
n_action_steps: int = 50
hidden_dim: int = 256
optimizer_lr: float = 1e-4
optimizer_weight_decay: float = 1e-4
def __post_init__(self):
super().__post_init__()
if self.n_action_steps > self.horizon:
raise ValueError("n_action_steps cannot exceed horizon")
def validate_features(self) -> None:
"""Validate input/output feature compatibility.
Call this explicitly from your policy's __init__ — the base class does not.
"""
if not self.image_features:
raise ValueError("MyPolicy requires at least one image feature.")
if self.action_feature is None:
raise ValueError("MyPolicy requires 'action' in output_features.")
def get_optimizer_preset(self) -> AdamWConfig:
return AdamWConfig(lr=self.optimizer_lr, weight_decay=self.optimizer_weight_decay)
def get_scheduler_preset(self):
"""Return a LRSchedulerConfig from lerobot.optim, or None."""
return None
@property
def observation_delta_indices(self) -> list[int] | None:
"""Relative timestep offsets the dataset loader provides per observation.
Return `None` for single-frame policies. For temporal policies that consume
multiple past or future frames, return a list of offsets, e.g. `[-20, -10, 0, 10]` for
3 past frames at stride 10 and 1 future frame at stride 10.
"""
return None
@property
def action_delta_indices(self) -> list[int]:
"""Relative timestep offsets for the action chunk the dataset loader returns."""
return list(range(self.horizon))
@property
def reward_delta_indices(self) -> None:
return None
The string you pass to @register_subclass must match MyPolicy.name (next section) and is what users supply as --policy.type on the CLI. Default to AdamW from lerobot.optim for get_optimizer_preset unless you genuinely need otherwise.
Inherit from PreTrainedPolicy and set two class attributes — both are checked by __init_subclass__:
# modeling_my_policy.py
import torch
import torch.nn as nn
from typing import Any
from lerobot.policies import PreTrainedPolicy
from lerobot.utils.constants import ACTION
from .configuration_my_policy import MyPolicyConfig
class MyPolicy(PreTrainedPolicy):
config_class = MyPolicyConfig # must match the string in @register_subclass
name = "my_policy"
def __init__(self, config: MyPolicyConfig, dataset_stats: dict[str, Any] = None):
super().__init__(config, dataset_stats)
config.validate_features() # not called automatically by the base class
self.config = config
self.model = ... # your nn.Module here
def reset(self):
"""Reset per-episode state. Called by lerobot-eval at the start of each episode."""
...
def get_optim_params(self) -> dict:
"""Return parameters to pass to the optimizer (e.g. with per-group lr/wd)."""
return {"params": self.parameters()}
def predict_action_chunk(self, batch: dict[str, torch.Tensor], **kwargs) -> torch.Tensor:
"""Return the full action chunk (B, chunk_size, action_dim) for the current observation."""
...
def select_action(self, batch: dict[str, torch.Tensor], **kwargs) -> torch.Tensor:
"""Return a single action for the current timestep (called every step at inference)."""
...
def forward(self, batch: dict[str, torch.Tensor]) -> tuple[torch.Tensor, dict | None]:
"""Compute the training loss.
Returns `(loss, output_dict)`. `output_dict` may be `None`; everything in it must be
logging-friendly Python natives (no tensors with gradients).
`batch["action_is_pad"]` is a bool mask of shape (B, horizon) that marks
timesteps padded because the episode ended before `horizon` steps; you
can exclude those from your loss.
"""
actions = batch[ACTION]
action_is_pad = batch.get("action_is_pad")
...
return loss, {"some_loss_component": some_loss_component.item()}
The methods called by the train/eval loops:
| Method | Used by | What it does |
|---|---|---|
reset() -> None | lerobot-eval | Clear per-episode state at the start of each episode. |
select_action(batch, **kwargs) -> Tensor | lerobot-eval | Return the next action (B, action_dim). Called every step. |
predict_action_chunk(batch, **kwargs) -> Tensor | the policy itself | Return an action chunk (B, chunk_size, action_dim). Currently abstract on the base class — raise NotImplementedError if your policy doesn't chunk. |
forward(batch, reduction="mean") -> tuple[Tensor, dict | None] | lerobot-train | Return (loss, output_dict). Accept reduction="none" if you want to support per-sample weighting. |
get_optim_params() -> dict | the optimizer | Return self.parameters() for simple policies; return a named parameter dict for multi-optimizer policies. |
update() -> None (optional) | lerobot-train | Called after each optimizer step if defined. Use for EMA, target nets, replay buffers (TDMPC uses this). |
Batches are flat dictionaries keyed by the constants in lerobot.utils.constants: OBS_STATE (observation.state.<motor>), OBS_IMAGES (observation.images.<camera>), OBS_LANGUAGE, ACTION, etc. Reuse the constants — don't invent new prefixes.
LeRobot uses PolicyProcessorPipelines to normalize inputs and de-normalize outputs around your policy. For a concrete reference, see processor_act.py or processor_diffusion.py.
# processor_my_policy.py
from typing import Any
import torch
from lerobot.processor import PolicyAction, PolicyProcessorPipeline
def make_my_policy_pre_post_processors(
config,
dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
) -> tuple[
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
PolicyProcessorPipeline[PolicyAction, PolicyAction],
]:
preprocessor = ... # build your PolicyProcessorPipeline for inputs
postprocessor = ... # build your PolicyProcessorPipeline for outputs
return preprocessor, postprocessor
Important — function naming: LeRobot discovers your processor by name. The function must be called make_{policy_name}_pre_post_processors (matching the string you passed to @PreTrainedConfig.register_subclass).
The fastest way to ship a policy: package it as a standalone Python distribution and install it alongside LeRobot. No PR required, you own the release cycle, and you can publish to PyPI under your own namespace.
Create a package with the prefix lerobot_policy_ (IMPORTANT!) followed by your policy name:
lerobot_policy_my_policy/
├── pyproject.toml
└── src/
└── lerobot_policy_my_policy/
├── __init__.py
├── configuration_my_policy.py
├── modeling_my_policy.py
└── processor_my_policy.py
pyproject.toml[project]
name = "lerobot_policy_my_policy"
version = "0.1.0"
dependencies = [
# your policy-specific dependencies
]
requires-python = ">= 3.12"
[build-system]
build-backend = # your-build-backend
requires = # your-build-system
__init__.pyExpose your classes in the package's __init__.py and guard against missing lerobot:
# __init__.py
"""Custom policy package for LeRobot."""
try:
import lerobot # noqa: F401
except ImportError:
raise ImportError(
"lerobot is not installed. Please install lerobot to use this policy package."
)
from .configuration_my_policy import MyPolicyConfig
from .modeling_my_policy import MyPolicy
from .processor_my_policy import make_my_policy_pre_post_processors
__all__ = [
"MyPolicyConfig",
"MyPolicy",
"make_my_policy_pre_post_processors",
]
cd lerobot_policy_my_policy
pip install -e .
# Or install from PyPI if published
pip install lerobot_policy_my_policy
Once installed, your policy automatically integrates with LeRobot's training and evaluation tools:
lerobot-train \
--policy.type my_policy \
--env.type pusht \
--steps 200000
When your policy has stabilized and there's clear value in shipping it with the library, you can land it directly in LeRobot. Read the general contribution guide and the PR template first — that's where you'll find the testing/quality expectations every PR has to meet (pre-commit run -a, pytest, the community-review rule, etc.). What's below is the policy-specific layer on top of that.
src/lerobot/policies/my_policy/
├── __init__.py # re-exports config + modeling + processor factory
├── configuration_my_policy.py # MyPolicyConfig + @register_subclass
├── modeling_my_policy.py # MyPolicy(PreTrainedPolicy)
├── processor_my_policy.py # make_my_policy_pre_post_processors
└── README.md # symlink → ../../../../docs/source/policy_my_policy_README.md
Two notes:
README.md next to the source is a symlink into docs/source/policy_<name>_README.md — the actual file lives under docs/. Existing policies (act, smolvla, diffusion, …) all do this; copy one of those symlinks. The policy README is conventionally minimal: paper link + BibTeX citation.docs/source/<my_policy>.mdx and is registered in _toctree.yml under "Policies".The file names are load-bearing: the factory does lazy imports by name, and the processor is discovered by the make_<policy_name>_pre_post_processors convention.
Four places need to know about your policy. All by name.
policies/__init__.py — re-export MyPolicyConfig and add it to __all__. Don't re-export the modeling class; it loads lazily through the factory (so import lerobot stays fast).factory.py:get_policy_class — add a branch returning MyPolicy from a lazy import.factory.py:make_policy_config and factory.py:make_pre_post_processors — same idea, two more branches.templates/lerobot_modelcard_template.md and the root README.md — the template is what push_model_to_hub renders into the model card of every checkpoint trained with your policy: add a one-line description of your policy in the model_name branches, map it in policy_docs so cards link to your MDX guide, and optionally add an architecture image to diagrams. Then add your policy to the models table in the root README.md, under the right category, linking to your doc page.Mirror an existing policy that's structurally similar to yours; the diff is small.
Most policies need a heavy backbone (transformers, diffusers, a specific VLM SDK). The convention is two-step gating: a TYPE_CHECKING-guarded import at module top, and a require_package runtime check in the constructor. modeling_diffusion.py is the canonical reference:
from typing import TYPE_CHECKING
from lerobot.utils.import_utils import _diffusers_available, require_package
if TYPE_CHECKING or _diffusers_available:
from diffusers.schedulers.scheduling_ddim import DDIMScheduler
else:
DDIMScheduler = None # keeps the symbol bindable at import time
class DiffusionPolicy(PreTrainedPolicy):
def __init__(self, config):
require_package("diffusers", extra="diffusion")
super().__init__(config)
...
This way:
import lerobot.policies keeps working without the extra installed (the symbol is just bound to None).ImportError pointing at pip install 'lerobot[diffusion]'.Add a matching extra to pyproject.toml [project.optional-dependencies] and include it in the all extra so pip install 'lerobot[all]' keeps installing everything.
A new policy is much easier to review — and far more useful — when it ships with a working checkpoint and at least one number you can reproduce.
Pick at least one in-tree benchmark. LeRobot ships sim benchmarks with per-benchmark Docker images (LIBERO, LIBERO-plus, Meta-World, RoboTwin 2.0, RoboCasa365, RoboCerebra, RoboMME, VLABench and more). Pick the one that matches your policy's modality — VLAs usually go to LIBERO or VLABench; image-only BC to LIBERO or Meta-World. The full list lives under Benchmarks in the docs sidebar.
Push the checkpoint & processors to the Hub under lerobot/<policy>_<benchmark> (or your namespace if you don't have write access; a maintainer can mirror it). Use PreTrainedPolicy.push_model_to_hub so the repo gets config.json, model.safetensors, and a model card.
Report results in your policy's MDX, with the exact lerobot-eval command and hardware so anyone can re-run:
## Results
Evaluated on LIBERO with `lerobot/<policy>_libero`:
| Suite | Success rate | n_episodes |
| -------------- | -----------: | ---------: |
| libero_spatial | 87.5% | 50 |
| libero_object | 93.0% | 50 |
| libero_goal | 81.5% | 50 |
| libero_10 | 62.0% | 50 |
| **average** | **81.0%** | 200 |
Reproduce: `lerobot-eval --policy.path=lerobot/<policy>_libero --env.type=libero --env.task=libero_spatial --eval.n_episodes=50` (1× A100 40 GB).
Use n_episodes ≥ 50 per suite for stable success-rate estimates.
If your policy is real-robot-only and no sim benchmark applies, swap the sim eval for: a public training dataset on the Hub, the lerobot-train command, the checkpoint, and a real-robot success rate over ≥10 episodes via lerobot-rollout --policy.path=....
The general expectations are in CONTRIBUTING.md and the PR template. On top of those, reviewers will look for:
MyPolicy and MyPolicyConfig cover the surface above; __init_subclass__ accepts the class.factory.py and policies/__init__.py are wired (lazy imports for modeling).make_my_policy_pre_post_processors follows the naming convention.[project.optional-dependencies] extra and the TYPE_CHECKING + require_package guard.tests/policies/ updated; backward-compat artifact committed & policy-specific tests.src/lerobot/policies/<name>/README.md symlinked into docs/source/policy_<name>_README.md; user-facing docs/source/<name>.mdx written and added to _toctree.yml.templates/lerobot_modelcard_template.md has a description entry and a policy_docs link for your policy.README.md lists your policy in the right category, linking to your doc page.The fastest way to get a clean PR is to copy the directory of the existing policy closest to yours, rename, and replace contents method by method. Don't wait until everything is polished — open a draft PR early and iterate with us; reviewers would much rather give feedback on a half-finished branch than a fully-merged one.
Check out these example policy implementations:
Thanks for taking the time to bring a new policy into LeRobot. Every architecture that lands in main — and every plugin published by the community — makes the library a little more useful for the next person, and a little more representative of where robot learning is going. We're looking forward to seeing what you ship. 🤗