Bring Your Own Policy

Learn how to integrate custom policies into the LeRobot framework

LeRobot is designed to be extensible. You can implement your own custom policy architectures and leverage LeRobot’s data collection, training infrastructure, and visualization tools.

Policy Interface

All LeRobot policies inherit from PreTrainedPolicy and implement a standard interface:

from lerobot.policies.pretrained import PreTrainedPolicy
import torch.nn as nn

class MyCustomPolicy(PreTrainedPolicy, nn.Module):
    """Custom policy implementation."""
    
    def __init__(self, config, dataset_stats=None):
        super().__init__(config)
        nn.Module.__init__(self)
        
        # Initialize your model architecture
        self.encoder = ...
        self.decoder = ...
    
    def forward(self, batch):
        """Forward pass for training.
        
        Args:
            batch: Dictionary with keys like 'observation.state', 'observation.image', 'action'
            
        Returns:
            loss: Training loss (scalar tensor)
            output_dict: Dictionary with additional outputs for logging
        """
        # Extract inputs
        state = batch['observation.state']
        images = batch['observation.image.side']
        actions = batch['action']
        
        # Forward pass
        encoding = self.encoder(state, images)
        predicted_actions = self.decoder(encoding)
        
        # Compute loss
        loss = nn.functional.mse_loss(predicted_actions, actions)
        
        # Additional outputs for logging
        output_dict = {
            'mse_loss': loss.item(),
        }
        
        return loss, output_dict
    
    def select_action(self, obs):
        """Select action for inference.
        
        Args:
            obs: Dictionary with observation keys
            
        Returns:
            action: Predicted action tensor
        """
        with torch.no_grad():
            state = obs['observation.state']
            images = obs['observation.image.side']
            
            encoding = self.encoder(state, images)
            action = self.decoder(encoding)
            
        return action

Configuration Class

Define a configuration class for your policy:

from lerobot.configs.policies import PreTrainedConfig
from dataclasses import dataclass, field

@dataclass
class MyCustomPolicyConfig(PreTrainedConfig):
    """Configuration for MyCustomPolicy."""
    
    # Required: policy type identifier
    type: str = "my_custom_policy"
    
    # Model architecture parameters
    hidden_dim: int = 256
    num_layers: int = 4
    dropout: float = 0.1
    
    # Training parameters
    optimizer_lr: float = 1e-4
    scheduler_decay_lr: float = 1e-5
    
    # Observation/action parameters (inherited from PreTrainedConfig)
    # input_features: dict
    # output_features: dict
    # device: str

Full Implementation Example

Here’s a complete example implementing a simple MLP policy:

import torch
import torch.nn as nn
from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.configs.policies import PreTrainedConfig
from dataclasses import dataclass

@dataclass
class MLPPolicyConfig(PreTrainedConfig):
    type: str = "mlp_policy"
    hidden_dims: list[int] = field(default_factory=lambda: [256, 256, 256])
    activation: str = "relu"
    dropout: float = 0.1

class MLPPolicy(PreTrainedPolicy, nn.Module):
    """Simple MLP policy for robot control."""
    
    def __init__(self, config: MLPPolicyConfig, dataset_stats=None):
        PreTrainedPolicy.__init__(self, config)
        nn.Module.__init__(self)
        
        self.config = config
        
        # Calculate input dimension from features
        input_dim = sum(ft.shape[0] for ft in config.input_features.values())
        output_dim = sum(ft.shape[0] for ft in config.output_features.values())
        
        # Build MLP layers
        layers = []
        prev_dim = input_dim
        
        for hidden_dim in config.hidden_dims:
            layers.append(nn.Linear(prev_dim, hidden_dim))
            layers.append(self._get_activation(config.activation))
            layers.append(nn.Dropout(config.dropout))
            prev_dim = hidden_dim
        
        # Output layer
        layers.append(nn.Linear(prev_dim, output_dim))
        
        self.network = nn.Sequential(*layers)
        
        # Normalization (optional but recommended)
        self.dataset_stats = dataset_stats
    
    def _get_activation(self, name):
        if name == "relu":
            return nn.ReLU()
        elif name == "tanh":
            return nn.Tanh()
        elif name == "gelu":
            return nn.GELU()
        else:
            raise ValueError(f"Unknown activation: {name}")
    
    def forward(self, batch):
        # Concatenate all input features
        inputs = []
        for key in sorted(self.config.input_features.keys()):
            inputs.append(batch[key])
        x = torch.cat(inputs, dim=-1)
        
        # Forward pass
        predicted_actions = self.network(x)
        
        # Get ground truth actions
        true_actions = batch['action']
        
        # Compute loss
        loss = nn.functional.mse_loss(predicted_actions, true_actions)
        
        output_dict = {
            'mse_loss': loss.item(),
        }
        
        return loss, output_dict
    
    def select_action(self, obs):
        """Select action during inference."""
        with torch.no_grad():
            # Concatenate inputs
            inputs = []
            for key in sorted(self.config.input_features.keys()):
                inputs.append(obs[key])
            x = torch.cat(inputs, dim=-1)
            
            # Predict action
            action = self.network(x)
            
        return action

Registering Your Policy

To use your policy with LeRobot’s training and evaluation scripts, register it in the policy factory:

# In lerobot/policies/factory.py

from lerobot.policies.my_custom_policy.modeling_my_custom_policy import MyCustomPolicy
from lerobot.policies.my_custom_policy.configuration_my_custom_policy import MyCustomPolicyConfig

def get_policy_class(name: str):
    # ... existing policies ...
    
    if name == "my_custom_policy":
        return MyCustomPolicy
    
    # ... rest of function ...

def make_policy_config(policy_type: str, **kwargs):
    # ... existing configs ...
    
    if policy_type == "my_custom_policy":
        return MyCustomPolicyConfig(**kwargs)
    
    # ... rest of function ...

Using Your Custom Policy

Training

Once registered, use your policy with the training CLI:

lerobot-train \
  --policy.type=my_custom_policy \
  --policy.hidden_dims=[512,512,512] \
  --policy.dropout=0.1 \
  --dataset.repo_id=your_username/your_dataset \
  --steps=50000 \
  --batch_size=64

Evaluation

Evaluate your trained policy:

lerobot-eval \
  --policy.path=your_username/my_custom_policy \
  --env.type=gym \
  --env.task=PandaPickPlace-v3 \
  --eval.n_episodes=10

Programmatic Usage

from lerobot.policies.my_custom_policy.modeling_my_custom_policy import MyCustomPolicy
from lerobot.policies.my_custom_policy.configuration_my_custom_policy import MyCustomPolicyConfig

# Create config
config = MyCustomPolicyConfig(
    input_features=input_features,
    output_features=output_features,
    hidden_dims=[512, 512, 512]
)

# Initialize policy
policy = MyCustomPolicy(config)

# Train
for batch in dataloader:
    loss, outputs = policy.forward(batch)
    loss.backward()
    optimizer.step()

# Inference
policy.eval()
action = policy.select_action(observation)

Advanced Features

Vision Encoders

Integrate vision encoders for image observations:

from lerobot.policies.utils import get_image_encoder

class MyVisionPolicy(PreTrainedPolicy, nn.Module):
    def __init__(self, config):
        super().__init__(config)
        
        # Create vision encoder
        self.vision_encoder = get_image_encoder(
            encoder_type="resnet18",
            input_channels=3,
            output_dim=512
        )
        
        # Create policy head
        self.policy_head = nn.Linear(512, action_dim)
    
    def forward(self, batch):
        # Encode images
        images = batch['observation.image.side']
        features = self.vision_encoder(images)
        
        # Predict actions
        actions = self.policy_head(features)
        
        loss = nn.functional.mse_loss(actions, batch['action'])
        return loss, {}

Action Chunking

Implement action chunking for temporal consistency:

class ChunkedPolicy(PreTrainedPolicy, nn.Module):
    def __init__(self, config):
        super().__init__(config)
        self.chunk_size = 16  # Predict 16 future actions
        self.action_dim = config.output_features['action'].shape[0]
        
        self.network = nn.Sequential(
            nn.Linear(obs_dim, 512),
            nn.ReLU(),
            nn.Linear(512, self.action_dim * self.chunk_size)
        )
    
    def forward(self, batch):
        # Predict action chunk
        obs = batch['observation.state']
        action_chunk = self.network(obs)
        action_chunk = action_chunk.reshape(-1, self.chunk_size, self.action_dim)
        
        # Supervise with ground truth chunk
        true_actions = batch['action']  # Shape: [batch, chunk_size, action_dim]
        loss = nn.functional.mse_loss(action_chunk, true_actions)
        
        return loss, {}
    
    def select_action(self, obs):
        """Returns next action from chunk."""
        with torch.no_grad():
            action_chunk = self.network(obs['observation.state'])
            action_chunk = action_chunk.reshape(-1, self.chunk_size, self.action_dim)
            # Return first action in chunk
            return action_chunk[:, 0, :]

Normalization

Use dataset statistics for input/output normalization:

class NormalizedPolicy(PreTrainedPolicy, nn.Module):
    def __init__(self, config, dataset_stats=None):
        super().__init__(config)
        
        # Store normalization statistics
        if dataset_stats:
            self.register_buffer(
                'state_mean',
                dataset_stats['observation.state']['mean']
            )
            self.register_buffer(
                'state_std',
                dataset_stats['observation.state']['std']
            )
            self.register_buffer(
                'action_mean',
                dataset_stats['action']['mean']
            )
            self.register_buffer(
                'action_std',
                dataset_stats['action']['std']
            )
        
        self.network = nn.Sequential(...)
    
    def normalize_state(self, state):
        return (state - self.state_mean) / (self.state_std + 1e-8)
    
    def unnormalize_action(self, action):
        return action * self.action_std + self.action_mean
    
    def forward(self, batch):
        # Normalize inputs
        state = self.normalize_state(batch['observation.state'])
        
        # Predict normalized action
        pred_action_norm = self.network(state)
        
        # Loss in normalized space
        true_action_norm = (batch['action'] - self.action_mean) / (self.action_std + 1e-8)
        loss = nn.functional.mse_loss(pred_action_norm, true_action_norm)
        
        return loss, {}
    
    def select_action(self, obs):
        with torch.no_grad():
            state = self.normalize_state(obs['observation.state'])
            action_norm = self.network(state)
            action = self.unnormalize_action(action_norm)
        return action

Best Practices

1. Follow the standard interface
2. Implement forward() for training and select_action() for inference:
3. def forward(self, batch) -> tuple[torch.Tensor, dict]:
       """Returns (loss, output_dict)."""
       pass
   
   def select_action(self, obs) -> torch.Tensor:
       """Returns action tensor."""
       pass
   

4. Use configuration dataclasses
5. Define all hyperparameters in a config class:
6. @dataclass
   class MyPolicyConfig(PreTrainedConfig):
       type: str = "my_policy"
       hidden_dim: int = 256
       # ... other params
   

7. Support saving/loading
8. Inherit from PreTrainedPolicy for automatic checkpoint handling:
9. # Save
   policy.save_pretrained("outputs/my_policy")
   
   # Load  
   policy = MyPolicy.from_pretrained("outputs/my_policy")
   

10. Add logging outputs
11. Return useful metrics in output_dict for monitoring:
12. output_dict = {
        'mse_loss': mse_loss.item(),
        'action_mean': actions.mean().item(),
        'action_std': actions.std().item(),
    }
    return loss, output_dict
    

Testing Your Policy

Test your policy implementation:

import torch
from lerobot.datasets.lerobot_dataset import LeRobotDataset

# Create dummy config
config = MyCustomPolicyConfig(
    input_features=input_features,
    output_features=output_features
)

# Initialize policy
policy = MyCustomPolicy(config)

# Test forward pass
dataset = LeRobotDataset("lerobot/pusht")
batch = next(iter(torch.utils.data.DataLoader(dataset, batch_size=4)))

loss, outputs = policy.forward(batch)
print(f"Loss: {loss.item()}")
print(f"Outputs: {outputs}")

# Test inference
policy.eval()
obs = batch
action = policy.select_action(obs)
print(f"Action shape: {action.shape}")

Next Steps

• Train Your First Policy - Train your custom policy
• Evaluate Policies - Test your policy
• Multi-GPU Training - Scale up training
• Policy Concepts - Learn about existing policies

Examples

See existing policy implementations:

• Diffusion Policy
• ACT Policy
• SAC Policy