admin/lerobot
1
0
Fork 0
mirror of https://github.com/huggingface/lerobot.git synced 2026-07-23 17:56:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

refactor(pipeline): Transition from tuple to dictionary format for EnvTransition

Browse Source 
- Updated the EnvTransition structure to use a dictionary format instead of a tuple, enhancing readability and maintainability.
- Replaced instances of TransitionIndex with TransitionKey for accessing transition components.
- Adjusted related processing functions and tests to accommodate the new dictionary format, ensuring consistent handling of transitions across the codebase.
This commit is contained in:
Adil Zouitine
2025-07-21 14:54:31 +02:00
parent 14c2ece004
commit f2b79656eb
16 changed files with 828 additions and 650 deletions
Download Patch File Download Diff File Expand all files Collapse all files
docs/source/processor_tutorial.mdx
+257 -175
View File
@@ -112,24 +112,23 @@ RobotProcessor solves these issues by providing a declarative pipeline approach
RobotProcessor works with two data formats: RobotProcessor works with two data formats:
### 1. EnvTransition Tuple Format ### 1. EnvTransition Dictionary Format
An `EnvTransition` is a 7-tuple that represents a complete transition in the environment: An `EnvTransition` is a dictionary that represents a complete transition in the environment:
```python ```python
from lerobot.processor.pipeline import TransitionIndex from lerobot.processor.pipeline import TransitionKey
# EnvTransition structure: # EnvTransition structure:
# (observation, action, reward, done, truncated, info, complementary_data) transition = {
transition = ( TransitionKey.OBSERVATION: {"observation.image": ..., "observation.state": ...}, # observation at time t
{"observation.image": ..., "observation.state": ...}, # observation at time t TransitionKey.ACTION: [0.1, -0.2, 0.3], # action taken at time t
[0.1, -0.2, 0.3], # action taken at time t TransitionKey.REWARD: 1.0, # reward received
1.0, # reward received TransitionKey.DONE: False, # episode done flag
False, # episode done flag TransitionKey.TRUNCATED: False, # episode truncated flag
False, # episode truncated flag TransitionKey.INFO: {"success": True}, # additional info from environment
{"success": True}, # additional info from environment TransitionKey.COMPLEMENTARY_DATA: {"step_idx": 42} # complementary_data for inter-step communication
{"step_idx": 42} # complementary_data for inter-step communication }
)
``` ```
### 2. Batch Dictionary Format ### 2. Batch Dictionary Format
@@ -160,9 +159,17 @@ from lerobot.processor.observation_processor import ImageProcessor
processor = RobotProcessor([ImageProcessor()]) processor = RobotProcessor([ImageProcessor()])
# Works with EnvTransition tuples # Works with EnvTransition dictionaries
transition = ({"pixels": image_array}, None, 0.0, False, False, {}, {}) transition = {
processed_transition = processor(transition) # Returns EnvTransition tuple TransitionKey.OBSERVATION: {"pixels": image_array},
TransitionKey.ACTION: None,
TransitionKey.REWARD: 0.0,
TransitionKey.DONE: False,
TransitionKey.TRUNCATED: False,
TransitionKey.INFO: {},
TransitionKey.COMPLEMENTARY_DATA: {}
}
processed_transition = processor(transition) # Returns EnvTransition dictionary
# Also works with batch dictionaries # Also works with batch dictionaries
batch = { batch = {
@@ -176,25 +183,25 @@ batch = {
processed_batch = processor(batch) # Returns batch dictionary processed_batch = processor(batch) # Returns batch dictionary
``` ```
### Using TransitionIndex ### Using TransitionKey
Instead of using magic numbers to access tuple elements, use the `TransitionIndex` enum: Use the `TransitionKey` enum to access dictionary elements:
```python ```python
from lerobot.processor.pipeline import TransitionIndex from lerobot.processor.pipeline import TransitionKey
# Bad - using magic numbers # Good - using TransitionKey
obs = transition[0] obs = transition[TransitionKey.OBSERVATION]
action = transition[1] action = transition[TransitionKey.ACTION]
reward = transition[TransitionKey.REWARD]
done = transition[TransitionKey.DONE]
truncated = transition[TransitionKey.TRUNCATED]
info = transition[TransitionKey.INFO]
comp_data = transition[TransitionKey.COMPLEMENTARY_DATA]
# Good - using TransitionIndex # Alternative - using .get() for optional access
obs = transition[TransitionIndex.OBSERVATION] obs = transition.get(TransitionKey.OBSERVATION)
action = transition[TransitionIndex.ACTION] action = transition.get(TransitionKey.ACTION)
reward = transition[TransitionIndex.REWARD]
done = transition[TransitionIndex.DONE]
truncated = transition[TransitionIndex.TRUNCATED]
info = transition[TransitionIndex.INFO]
comp_data = transition[TransitionIndex.COMPLEMENTARY_DATA]
``` ```
### Default Conversion Functions ### Default Conversion Functions
@@ -203,43 +210,49 @@ RobotProcessor uses these default conversion functions:
```python ```python
def _default_batch_to_transition(batch): def _default_batch_to_transition(batch):
"""Default conversion from batch dict to EnvTransition tuple.""" """Default conversion from batch dict to EnvTransition dictionary."""
# Extract observation keys (anything starting with "observation.") # Extract observation keys (anything starting with "observation.")
observation_keys = {k: v for k, v in batch.items() if k.startswith("observation.")} observation_keys = {k: v for k, v in batch.items() if k.startswith("observation.")}
observation = observation_keys if observation_keys else None
observation = None # Extract padding and task keys for complementary data
if observation_keys: pad_keys = {k: v for k, v in batch.items() if "_is_pad" in k}
observation = {} task_key = {"task": batch["task"]} if "task" in batch else {}
# Keep observation.* keys as-is (don't remove "observation." prefix) complementary_data = {**pad_keys, **task_key} if pad_keys or task_key else {}
for key, value in observation_keys.items():
observation[key] = value
return ( transition = {
observation, TransitionKey.OBSERVATION: observation,
batch.get("action"), TransitionKey.ACTION: batch.get("action"),
batch.get("next.reward", 0.0), # Note: "next.reward" not "reward" TransitionKey.REWARD: batch.get("next.reward", 0.0),
batch.get("next.done", False), # Note: "next.done" not "done" TransitionKey.DONE: batch.get("next.done", False),
batch.get("next.truncated", False), # Note: "next.truncated" not "truncated" TransitionKey.TRUNCATED: batch.get("next.truncated", False),
batch.get("info", {}), TransitionKey.INFO: batch.get("info", {}),
{}, # Empty complementary_data TransitionKey.COMPLEMENTARY_DATA: complementary_data,
) }
return transition
def _default_transition_to_batch(transition): def _default_transition_to_batch(transition):
"""Default conversion from EnvTransition tuple to batch dict.""" """Default conversion from EnvTransition dictionary to batch dict."""
obs, action, reward, done, truncated, info, _ = transition
batch = { batch = {
"action": action, "action": transition.get(TransitionKey.ACTION),
"next.reward": reward, # Note: "next.reward" not "reward" "next.reward": transition.get(TransitionKey.REWARD, 0.0),
"next.done": done, # Note: "next.done" not "done" "next.done": transition.get(TransitionKey.DONE, False),
"next.truncated": truncated, # Note: "next.truncated" not "truncated" "next.truncated": transition.get(TransitionKey.TRUNCATED, False),
"info": info, "info": transition.get(TransitionKey.INFO, {}),
} }
# Flatten observation dict (keep observation.* keys as-is) # Add padding and task data from complementary_data
if isinstance(obs, dict): complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
for key, value in obs.items(): if complementary_data:
batch[key] = value pad_data = {k: v for k, v in complementary_data.items() if "_is_pad" in k}
batch.update(pad_data)
if "task" in complementary_data:
batch["task"] = complementary_data["task"]
# Handle observation - flatten dict to observation.* keys if it's a dict
observation = transition.get(TransitionKey.OBSERVATION)
if isinstance(observation, dict):
batch.update(observation)
return batch return batch
``` ```
@@ -250,33 +263,32 @@ You can customize how RobotProcessor converts between formats:
```python ```python
def custom_batch_to_transition(batch): def custom_batch_to_transition(batch):
"""Custom conversion from batch dict to EnvTransition tuple.""" """Custom conversion from batch dict to EnvTransition dictionary."""
# Extract observation keys (anything starting with "observation.") # Extract observation keys (anything starting with "observation.")
observation = {k: v for k, v in batch.items() if k.startswith("observation.")} observation = {k: v for k, v in batch.items() if k.startswith("observation.")}
return ( return {
observation, TransitionKey.OBSERVATION: observation,
batch.get("action"), TransitionKey.ACTION: batch.get("action"),
batch.get("reward", 0.0), # Use "reward" instead of "next.reward" TransitionKey.REWARD: batch.get("reward", 0.0), # Use "reward" instead of "next.reward"
batch.get("done", False), # Use "done" instead of "next.done" TransitionKey.DONE: batch.get("done", False), # Use "done" instead of "next.done"
batch.get("truncated", False), TransitionKey.TRUNCATED: batch.get("truncated", False),
batch.get("info", {}), TransitionKey.INFO: batch.get("info", {}),
batch.get("complementary_data", {}) TransitionKey.COMPLEMENTARY_DATA: batch.get("complementary_data", {})
) }
def custom_transition_to_batch(transition): def custom_transition_to_batch(transition):
"""Custom conversion from EnvTransition tuple to batch dict.""" """Custom conversion from EnvTransition dictionary to batch dict."""
obs, action, reward, done, truncated, info, comp_data = transition
batch = { batch = {
"action": action, "action": transition.get(TransitionKey.ACTION),
"reward": reward, # Use "reward" instead of "next.reward" "reward": transition.get(TransitionKey.REWARD), # Use "reward" instead of "next.reward"
"done": done, # Use "done" instead of "next.done" "done": transition.get(TransitionKey.DONE), # Use "done" instead of "next.done"
"truncated": truncated, "truncated": transition.get(TransitionKey.TRUNCATED),
"info": info, "info": transition.get(TransitionKey.INFO),
} }
# Flatten observation dict # Flatten observation dict
obs = transition.get(TransitionKey.OBSERVATION)
if obs: if obs:
batch.update(obs) batch.update(obs)
@@ -292,21 +304,21 @@ processor = RobotProcessor(
### Advanced: Controlling Output Format with `to_output` ### Advanced: Controlling Output Format with `to_output`
The `to_output` function determines what format is returned when you call the processor with a batch dictionary. Sometimes you want to output `EnvTransition` tuples even when you input batch dictionaries: The `to_output` function determines what format is returned when you call the processor with a batch dictionary. Sometimes you want to output `EnvTransition` dictionaries even when you input batch dictionaries:
```python ```python
# Identity function to always return EnvTransition tuples # Identity function to always return EnvTransition dictionaries
def keep_as_transition(transition): def keep_as_transition(transition):
"""Always return EnvTransition tuple regardless of input format.""" """Always return EnvTransition dictionary regardless of input format."""
return transition return transition
# Processor that always outputs EnvTransition tuples # Processor that always outputs EnvTransition dictionaries
processor = RobotProcessor( processor = RobotProcessor(
steps=[ImageProcessor(), StateProcessor()], steps=[ImageProcessor(), StateProcessor()],
to_output=keep_as_transition # Always return tuple format to_output=keep_as_transition # Always return dictionary format
) )
# Even when called with batch dict, returns EnvTransition tuple # Even when called with batch dict, returns EnvTransition dictionary
batch = { batch = {
"observation.image": image_tensor, "observation.image": image_tensor,
"action": action_tensor, "action": action_tensor,
@@ -316,13 +328,13 @@ batch = {
"info": info_dict "info": info_dict
} }
result = processor(batch) # Returns EnvTransition tuple, not batch dict! result = processor(batch) # Returns EnvTransition dictionary, not batch dict!
print(type(result)) # <class 'tuple'> print(type(result)) # <class 'dict'>
``` ```
### Real-World Example: Environment Interaction ### Real-World Example: Environment Interaction
This is particularly useful for environment interaction where you want consistent tuple output: This is particularly useful for environment interaction where you want consistent dictionary output:
```python ```python
from lerobot.processor.observation_processor import VanillaObservationProcessor from lerobot.processor.observation_processor import VanillaObservationProcessor
@@ -332,7 +344,7 @@ from lerobot.processor.observation_processor import VanillaObservationProcessor
env_processor = RobotProcessor( env_processor = RobotProcessor(
[VanillaObservationProcessor()], [VanillaObservationProcessor()],
to_transition=lambda x: x, # Pass through - no conversion needed to_transition=lambda x: x, # Pass through - no conversion needed
to_output=lambda x: x, # Always return EnvTransition tuple to_output=lambda x: x, # Always return EnvTransition dictionary
) )
# Environment interaction loop # Environment interaction loop
@@ -340,12 +352,20 @@ env = make_env()
obs, info = env.reset() obs, info = env.reset()
for step in range(1000): for step in range(1000):
# Create transition - input is already in tuple format # Create transition - input is already in dictionary format
transition = (obs, None, 0.0, False, False, info, {"step": step}) transition = {
TransitionKey.OBSERVATION: obs,
TransitionKey.ACTION: None,
TransitionKey.REWARD: 0.0,
TransitionKey.DONE: False,
TransitionKey.TRUNCATED: False,
TransitionKey.INFO: info,
TransitionKey.COMPLEMENTARY_DATA: {"step": step}
}
# Process - output is guaranteed to be EnvTransition tuple # Process - output is guaranteed to be EnvTransition dictionary
processed_transition = env_processor(transition) processed_transition = env_processor(transition)
processed_obs = processed_transition[TransitionIndex.OBSERVATION] processed_obs = processed_transition[TransitionKey.OBSERVATION]
# Use with policy # Use with policy
action = policy.select_action(processed_obs) action = policy.select_action(processed_obs)
@@ -357,12 +377,12 @@ for step in range(1000):
### When to Use Different Output Formats ### When to Use Different Output Formats
**Use EnvTransition tuple output when:** **Use EnvTransition dictionary output when:**
- Environment interaction and real-time control - Environment interaction and real-time control
- You need to access individual transition components frequently - You need to access individual transition components frequently
- Performance is critical (avoids dictionary creation overhead) - Working with gym environments that expect structured data
- Working with gym environments that expect tuple format - You need the flexibility of dictionary operations
**Use batch dictionary output when:** **Use batch dictionary output when:**
@@ -372,10 +392,10 @@ for step in range(1000):
- You need the standardized "next.\*" key format - You need the standardized "next.\*" key format
```python ```python
# For environment interaction - use tuple output # For environment interaction - use dictionary output
env_processor = RobotProcessor( env_processor = RobotProcessor(
steps=[ImageProcessor(), StateProcessor()], steps=[ImageProcessor(), StateProcessor()],
to_output=lambda x: x # Return EnvTransition tuple to_output=lambda x: x # Return EnvTransition dictionary
) )
# For training - use batch output (default) # For training - use batch output (default)
@@ -391,9 +411,17 @@ for batch in dataloader:
# Environment loop # Environment loop
for step in range(1000): for step in range(1000):
transition = (obs, None, 0.0, False, False, info, {}) transition = {
processed_transition = env_processor(transition) # Returns EnvTransition tuple TransitionKey.OBSERVATION: obs,
obs = processed_transition[TransitionIndex.OBSERVATION] TransitionKey.ACTION: None,
TransitionKey.REWARD: 0.0,
TransitionKey.DONE: False,
TransitionKey.TRUNCATED: False,
TransitionKey.INFO: info,
TransitionKey.COMPLEMENTARY_DATA: {}
}
processed_transition = env_processor(transition) # Returns EnvTransition dictionary
obs = processed_transition[TransitionKey.OBSERVATION]
action = policy.select_action(obs) action = policy.select_action(obs)
``` ```
@@ -426,7 +454,7 @@ batch = {
Let's create a processor that properly handles image and state preprocessing: Let's create a processor that properly handles image and state preprocessing:
```python ```python
from lerobot.processor.pipeline import RobotProcessor, TransitionIndex from lerobot.processor.pipeline import RobotProcessor, TransitionKey
from lerobot.processor.observation_processor import ImageProcessor, StateProcessor from lerobot.processor.observation_processor import ImageProcessor, StateProcessor
import numpy as np import numpy as np
@@ -440,7 +468,15 @@ observation = {
} }
# Create a full transition # Create a full transition
transition = (observation, None, 0.0, False, False, {}, {}) transition = {
TransitionKey.OBSERVATION: observation,
TransitionKey.ACTION: None,
TransitionKey.REWARD: 0.0,
TransitionKey.DONE: False,
TransitionKey.TRUNCATED: False,
TransitionKey.INFO: {},
TransitionKey.COMPLEMENTARY_DATA: {}
}
# Create and use the processor # Create and use the processor
processor = RobotProcessor([ processor = RobotProcessor([
@@ -449,7 +485,7 @@ processor = RobotProcessor([
]) ])
processed_transition = processor(transition) processed_transition = processor(transition)
processed_obs = processed_transition[TransitionIndex.OBSERVATION] processed_obs = processed_transition[TransitionKey.OBSERVATION]
# Check the results # Check the results
print("Original keys:", observation.keys()) print("Original keys:", observation.keys())
@@ -541,11 +577,19 @@ obs, info = env.reset()
for step in range(1000): for step in range(1000):
# Raw environment observation # Raw environment observation
transition = (obs, None, 0.0, False, False, info, {}) transition = {
TransitionKey.OBSERVATION: obs,
TransitionKey.ACTION: None,
TransitionKey.REWARD: 0.0,
TransitionKey.DONE: False,
TransitionKey.TRUNCATED: False,
TransitionKey.INFO: info,
TransitionKey.COMPLEMENTARY_DATA: {}
}
# Process for policy input # Process for policy input
processed_transition = online_processor(transition) processed_transition = online_processor(transition)
processed_obs = processed_transition[TransitionIndex.OBSERVATION] processed_obs = processed_transition[TransitionKey.OBSERVATION]
# Get action from policy # Get action from policy
action = policy.select_action(processed_obs) action = policy.select_action(processed_obs)
@@ -585,15 +629,16 @@ class ImagePadder:
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
"""Main processing method - required for all steps.""" """Main processing method - required for all steps."""
obs = transition[TransitionIndex.OBSERVATION] obs = transition.get(TransitionKey.OBSERVATION)
if obs is None: if obs is None:
return transition return transition
# Process all image observations # Process all image observations
for key in list(obs.keys()): processed_obs = dict(obs) # Create a copy
for key in list(processed_obs.keys()):
if key.startswith("observation.images."): if key.startswith("observation.images."):
img = obs[key] img = processed_obs[key]
# Calculate padding # Calculate padding
_, _, h, w = img.shape _, _, h, w = img.shape
pad_h = max(0, self.target_height - h) pad_h = max(0, self.target_height - h)
@@ -609,10 +654,12 @@ class ImagePadder:
img = F.pad(img, (pad_left, pad_right, pad_top, pad_bottom), img = F.pad(img, (pad_left, pad_right, pad_top, pad_bottom),
mode='constant', value=self.pad_value) mode='constant', value=self.pad_value)
obs[key] = img processed_obs[key] = img
# Return modified transition # Return modified transition
return (obs, *transition[1:]) new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = processed_obs
return new_transition
def get_config(self) -> Dict[str, Any]: def get_config(self) -> Dict[str, Any]:
"""Return JSON-serializable configuration - required for save/load.""" """Return JSON-serializable configuration - required for save/load."""
@@ -694,8 +741,8 @@ class ImageStatisticsCalculator:
"""Calculate image statistics and pass to next steps.""" """Calculate image statistics and pass to next steps."""
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
obs = transition[TransitionIndex.OBSERVATION] obs = transition.get(TransitionKey.OBSERVATION)
comp_data = transition[TransitionIndex.COMPLEMENTARY_DATA] or {} comp_data = transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}
if obs is None: if obs is None:
return transition return transition
@@ -714,18 +761,13 @@ class ImageStatisticsCalculator:
image_stats[key] = stats image_stats[key] = stats
# Store in complementary_data for next steps # Store in complementary_data for next steps
comp_data = dict(comp_data) # Make a copy
comp_data["image_statistics"] = image_stats comp_data["image_statistics"] = image_stats
# Return transition with updated complementary_data # Return transition with updated complementary_data
return ( new_transition = transition.copy()
obs, new_transition[TransitionKey.COMPLEMENTARY_DATA] = comp_data
transition[TransitionIndex.ACTION], return new_transition
transition[TransitionIndex.REWARD],
transition[TransitionIndex.DONE],
transition[TransitionIndex.TRUNCATED],
transition[TransitionIndex.INFO],
comp_data # Updated complementary_data
)
@dataclass @dataclass
class AdaptiveBrightnessAdjuster: class AdaptiveBrightnessAdjuster:
@@ -734,8 +776,8 @@ class AdaptiveBrightnessAdjuster:
target_brightness: float = 0.5 target_brightness: float = 0.5
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
obs = transition[TransitionIndex.OBSERVATION] obs = transition.get(TransitionKey.OBSERVATION)
comp_data = transition[TransitionIndex.COMPLEMENTARY_DATA] or {} comp_data = transition.get(TransitionKey.COMPLEMENTARY_DATA) or {}
if obs is None or "image_statistics" not in comp_data: if obs is None or "image_statistics" not in comp_data:
return transition return transition
@@ -743,15 +785,18 @@ class AdaptiveBrightnessAdjuster:
# Use statistics from previous step # Use statistics from previous step
image_stats = comp_data["image_statistics"] image_stats = comp_data["image_statistics"]
for key in obs: processed_obs = dict(obs) # Create a copy
for key in processed_obs:
if key.startswith("observation.images.") and key in image_stats: if key.startswith("observation.images.") and key in image_stats:
current_mean = image_stats[key]["mean"] current_mean = image_stats[key]["mean"]
brightness_adjust = self.target_brightness - current_mean brightness_adjust = self.target_brightness - current_mean
# Adjust brightness # Adjust brightness
obs[key] = torch.clamp(obs[key] + brightness_adjust, 0, 1) processed_obs[key] = torch.clamp(processed_obs[key] + brightness_adjust, 0, 1)
return (obs, *transition[1:]) new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = processed_obs
return new_transition
# Use them together # Use them together
processor = RobotProcessor([ processor = RobotProcessor([
@@ -782,7 +827,7 @@ class ActionRepeatStep:
env: gym.Env = None # This can't be serialized to JSON! env: gym.Env = None # This can't be serialized to JSON!
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
obs, action, reward, done, truncated, info, comp_data = transition action = transition.get(TransitionKey.ACTION)
if self.env is not None and action is not None: if self.env is not None and action is not None:
# Repeat action multiple times in environment # Repeat action multiple times in environment
@@ -792,9 +837,13 @@ class ActionRepeatStep:
total_reward += r total_reward += r
if d or t: if d or t:
break break
reward = total_reward
return (obs, action, reward, done, truncated, info, comp_data) # Update reward in transition
new_transition = transition.copy()
new_transition[TransitionKey.REWARD] = total_reward
return new_transition
return transition
def get_config(self) -> dict[str, Any]: def get_config(self) -> dict[str, Any]:
# Note: env is NOT included because it's not serializable # Note: env is NOT included because it's not serializable
@@ -1211,7 +1260,7 @@ This enables sharing of preprocessing logic while allowing each user to provide
Here's a complete example showing proper device management and all features: Here's a complete example showing proper device management and all features:
```python ```python
from lerobot.processor.pipeline import RobotProcessor, ProcessorStepRegistry, TransitionIndex from lerobot.processor.pipeline import RobotProcessor, ProcessorStepRegistry, TransitionKey
import torch import torch
import torch.nn.functional as F import torch.nn.functional as F
import numpy as np import numpy as np
@@ -1224,23 +1273,29 @@ class DeviceMover:
device: str = "cuda" device: str = "cuda"
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
obs = transition[TransitionIndex.OBSERVATION] obs = transition.get(TransitionKey.OBSERVATION)
if obs is None: if obs is None:
return transition return transition
# Move all tensor observations to device # Move all tensor observations to device
for key, value in obs.items(): processed_obs = dict(obs) # Create a copy
for key, value in processed_obs.items():
if isinstance(value, torch.Tensor): if isinstance(value, torch.Tensor):
obs[key] = value.to(self.device) processed_obs[key] = value.to(self.device)
# Also handle action if present # Also handle action if present
action = transition[TransitionIndex.ACTION] action = transition.get(TransitionKey.ACTION)
if action is not None and isinstance(action, torch.Tensor): if action is not None and isinstance(action, torch.Tensor):
action = action.to(self.device) action = action.to(self.device)
return (obs, action, *transition[2:]) new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = processed_obs
new_transition[TransitionKey.ACTION] = action
return new_transition
return (obs, *transition[1:]) new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = processed_obs
return new_transition
def get_config(self) -> Dict[str, Any]: def get_config(self) -> Dict[str, Any]:
return {"device": str(self.device)} return {"device": str(self.device)}
@@ -1260,7 +1315,7 @@ class RunningNormalizer:
self.initialized = False self.initialized = False
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
obs = transition[TransitionIndex.OBSERVATION] obs = transition.get(TransitionKey.OBSERVATION)
if obs is None or "observation.state" not in obs: if obs is None or "observation.state" not in obs:
return transition return transition
@@ -1284,9 +1339,12 @@ class RunningNormalizer:
# Normalize # Normalize
state_normalized = (state - self.running_mean) / (self.running_var + 1e-8).sqrt() state_normalized = (state - self.running_mean) / (self.running_var + 1e-8).sqrt()
obs["observation.state"] = state_normalized processed_obs = dict(obs) # Create a copy
processed_obs["observation.state"] = state_normalized
return (obs, *transition[1:]) new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = processed_obs
return new_transition
def get_config(self) -> Dict[str, Any]: def get_config(self) -> Dict[str, Any]:
return { return {
@@ -1317,7 +1375,7 @@ class RunningNormalizer:
processor = RobotProcessor([ processor = RobotProcessor([
ImageProcessor(), # Convert images to float32 [0,1] ImageProcessor(), # Convert images to float32 [0,1]
StateProcessor(), # Convert states to torch tensors StateProcessor(), # Convert states to torch tensors
ImagePadder(224, 224), # Pad images to standard size ImagePadder(target_height=224, target_width=224), # Pad images to standard size
DeviceMover("cuda"), # Move everything to GPU DeviceMover("cuda"), # Move everything to GPU
RunningNormalizer(7), # Normalize states RunningNormalizer(7), # Normalize states
], name="CompletePreprocessor") ], name="CompletePreprocessor")
@@ -1330,11 +1388,19 @@ obs = {
"pixels": {"cam": np.random.randint(0, 255, (200, 300, 3), dtype=np.uint8)}, "pixels": {"cam": np.random.randint(0, 255, (200, 300, 3), dtype=np.uint8)},
"agent_pos": np.random.randn(7).astype(np.float32) "agent_pos": np.random.randn(7).astype(np.float32)
} }
transition = (obs, None, 0.0, False, False, {}, {}) transition = {
TransitionKey.OBSERVATION: obs,
TransitionKey.ACTION: None,
TransitionKey.REWARD: 0.0,
TransitionKey.DONE: False,
TransitionKey.TRUNCATED: False,
TransitionKey.INFO: {},
TransitionKey.COMPLEMENTARY_DATA: {}
}
# Everything is processed and on GPU # Everything is processed and on GPU
processed = processor(transition) processed = processor(transition)
print(processed[TransitionIndex.OBSERVATION]["observation.images.cam"].device) # cuda:0 print(processed[TransitionKey.OBSERVATION]["observation.images.cam"].device) # cuda:0
``` ```
## Solving Real-World Problems with RobotProcessor ## Solving Real-World Problems with RobotProcessor
@@ -1358,22 +1424,25 @@ class KeyRemapper:
}) })
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
obs = transition[TransitionIndex.OBSERVATION] obs = transition.get(TransitionKey.OBSERVATION)
if obs is None: if obs is None:
return transition return transition
# Create new observation with renamed keys # Create new observation with renamed keys
processed_obs = dict(obs) # Create a copy
renamed_obs = {} renamed_obs = {}
for old_key, new_key in self.key_mapping.items(): for old_key, new_key in self.key_mapping.items():
if old_key in obs: if old_key in processed_obs:
renamed_obs[new_key] = obs[old_key] renamed_obs[new_key] = processed_obs[old_key]
# Keep any unmapped keys as-is # Keep any unmapped keys as-is
for key, value in obs.items(): for key, value in processed_obs.items():
if key not in self.key_mapping: if key not in self.key_mapping:
renamed_obs[key] = value renamed_obs[key] = value
return (renamed_obs, *transition[1:]) new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = renamed_obs
return new_transition
``` ```
### Workspace-Focused Image Processing ### Workspace-Focused Image Processing
@@ -1390,13 +1459,14 @@ class WorkspaceCropper:
output_size: Tuple[int, int] = (224, 224) output_size: Tuple[int, int] = (224, 224)
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
obs = transition[TransitionIndex.OBSERVATION] obs = transition.get(TransitionKey.OBSERVATION)
if obs is None: if obs is None:
return transition return transition
for key in list(obs.keys()): processed_obs = dict(obs) # Create a copy
for key in list(processed_obs.keys()):
if key.startswith("observation.images."): if key.startswith("observation.images."):
img = obs[key] img = processed_obs[key]
# Crop to workspace # Crop to workspace
x1, y1, x2, y2 = self.crop_bbox x1, y1, x2, y2 = self.crop_bbox
img_cropped = img[:, :, y1:y2, x1:x2] img_cropped = img[:, :, y1:y2, x1:x2]
@@ -1407,9 +1477,11 @@ class WorkspaceCropper:
mode='bilinear', mode='bilinear',
align_corners=False align_corners=False
) )
obs[key] = img_resized processed_obs[key] = img_resized
return (obs, *transition[1:]) new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = processed_obs
return new_transition
``` ```
### Building Complete Pipelines for Different Robots ### Building Complete Pipelines for Different Robots
@@ -1471,7 +1543,7 @@ The beauty of this approach is that:
```python ```python
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
obs = transition[TransitionIndex.OBSERVATION] obs = transition.get(TransitionKey.OBSERVATION)
# Always check if observation exists # Always check if observation exists
if obs is None: if obs is None:
@@ -1496,7 +1568,7 @@ return (modified_obs, None, 0.0, False, False, {}, {})
```python ```python
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
obs = transition[TransitionIndex.OBSERVATION] obs = transition.get(TransitionKey.OBSERVATION)
if self.store_previous: if self.store_previous:
# Good - clone to avoid reference issues # Good - clone to avoid reference issues
@@ -1522,7 +1594,7 @@ def state_dict(self) -> Dict[str, torch.Tensor]:
Here's how to use RobotProcessor in a real robot control loop, showing both tuple and batch formats: Here's how to use RobotProcessor in a real robot control loop, showing both tuple and batch formats:
```python ```python
from lerobot.processor.pipeline import RobotProcessor, ProcessorStepRegistry, TransitionIndex from lerobot.processor.pipeline import RobotProcessor, ProcessorStepRegistry, TransitionKey
from lerobot.policies.act.modeling_act import ACTPolicy from lerobot.policies.act.modeling_act import ACTPolicy
from pathlib import Path from pathlib import Path
import time import time
@@ -1545,11 +1617,13 @@ class ActionClipper:
max_value: float = 1.0 max_value: float = 1.0
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
action = transition[TransitionIndex.ACTION] action = transition.get(TransitionKey.ACTION)
if action is not None: if action is not None:
action = torch.clamp(action, self.min_value, self.max_value) action = torch.clamp(action, self.min_value, self.max_value)
return (transition[TransitionIndex.OBSERVATION], action, *transition[2:]) new_transition = transition.copy()
new_transition[TransitionKey.ACTION] = action
return new_transition
return transition return transition
@@ -1578,28 +1652,36 @@ for episode in range(10):
for step in range(1000): for step in range(1000):
# Create transition with raw observation # Create transition with raw observation
transition = (obs, None, 0.0, False, False, info, {"step": step}) transition = {
TransitionKey.OBSERVATION: obs,
TransitionKey.ACTION: None,
TransitionKey.REWARD: 0.0,
TransitionKey.DONE: False,
TransitionKey.TRUNCATED: False,
TransitionKey.INFO: info,
TransitionKey.COMPLEMENTARY_DATA: {"step": step}
}
# Preprocess - works with tuple format # Preprocess - works with dictionary format
processed_transition = preprocessor(transition) processed_transition = preprocessor(transition)
processed_obs = processed_transition[TransitionIndex.OBSERVATION] processed_obs = processed_transition.get(TransitionKey.OBSERVATION)
# Get action from policy # Get action from policy
with torch.no_grad(): with torch.no_grad():
action = policy.select_action(processed_obs) action = policy.select_action(processed_obs)
# Postprocess action # Postprocess action
action_transition = ( action_transition = {
processed_obs, TransitionKey.OBSERVATION: processed_obs,
action, TransitionKey.ACTION: action,
0.0, TransitionKey.REWARD: 0.0,
False, TransitionKey.DONE: False,
False, TransitionKey.TRUNCATED: False,
info, TransitionKey.INFO: info,
{"raw_action": action.clone()} # Store raw action in complementary_data TransitionKey.COMPLEMENTARY_DATA: {"raw_action": action.clone()} # Store raw action in complementary_data
) }
processed_action_transition = postprocessor(action_transition) processed_action_transition = postprocessor(action_transition)
final_action = processed_action_transition[TransitionIndex.ACTION] final_action = processed_action_transition.get(TransitionKey.ACTION)
# Execute action # Execute action
obs, reward, terminated, truncated, info = env.step(final_action.cpu().numpy()) obs, reward, terminated, truncated, info = env.step(final_action.cpu().numpy())
@@ -1667,7 +1749,7 @@ Use the full power of `RobotProcessor` for debugging:
```python ```python
# Enable detailed logging # Enable detailed logging
def log_observation_shapes(step_idx: int, transition: EnvTransition): def log_observation_shapes(step_idx: int, transition: EnvTransition):
obs = transition[TransitionIndex.OBSERVATION] obs = transition.get(TransitionKey.OBSERVATION)
if obs: if obs:
print(f"Step {step_idx} observations:") print(f"Step {step_idx} observations:")
for key, value in obs.items(): for key, value in obs.items():
@@ -1679,7 +1761,7 @@ processor.register_after_step_hook(log_observation_shapes)
# Monitor complementary data flow # Monitor complementary data flow
def monitor_complementary_data(step_idx: int, transition: EnvTransition): def monitor_complementary_data(step_idx: int, transition: EnvTransition):
comp_data = transition[TransitionIndex.COMPLEMENTARY_DATA] comp_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
if comp_data: if comp_data:
print(f"Step {step_idx} complementary_data: {list(comp_data.keys())}") print(f"Step {step_idx} complementary_data: {list(comp_data.keys())}")
return None return None
@@ -1688,7 +1770,7 @@ processor.register_before_step_hook(monitor_complementary_data)
# Validate data integrity # Validate data integrity
def validate_tensors(step_idx: int, transition: EnvTransition): def validate_tensors(step_idx: int, transition: EnvTransition):
obs = transition[TransitionIndex.OBSERVATION] obs = transition.get(TransitionKey.OBSERVATION)
if obs: if obs:
for key, value in obs.items(): for key, value in obs.items():
if isinstance(value, torch.Tensor): if isinstance(value, torch.Tensor):
@@ -1705,21 +1787,21 @@ processor.register_after_step_hook(validate_tensors)
RobotProcessor provides a powerful, modular approach to data preprocessing in robotics: RobotProcessor provides a powerful, modular approach to data preprocessing in robotics:
- **Dual format support**: Works seamlessly with both EnvTransition tuples and batch dictionaries - **Dual format support**: Works seamlessly with both EnvTransition dictionaries and batch dictionaries
- **Automatic format conversion**: Converts between tuple and batch formats as needed - **Automatic format conversion**: Converts between dictionary and batch formats as needed
- **LeRobot integration**: Native support for LeRobotDataset and ReplayBuffer formats - **LeRobot integration**: Native support for LeRobotDataset and ReplayBuffer formats
- **Clear separation of concerns**: Each transformation is a separate, testable unit - **Clear separation of concerns**: Each transformation is a separate, testable unit
- **Proper state management**: Clear distinction between config (JSON) and state (tensors) - **Proper state management**: Clear distinction between config (JSON) and state (tensors)
- **Device-aware**: Seamless GPU/CPU transfers with `.to(device)` - **Device-aware**: Seamless GPU/CPU transfers with `.to(device)`
- **Inter-step communication**: Use `complementary_data` for passing information - **Inter-step communication**: Use `complementary_data` for passing information
- **Easy sharing**: Push to Hugging Face Hub for reproducibility - **Easy sharing**: Push to Hugging Face Hub for reproducibility
- **Type safety**: Use `TransitionIndex` instead of magic numbers - **Type safety**: Use `TransitionKey` instead of magic numbers
- **Debugging tools**: Step through transformations and add monitoring hooks - **Debugging tools**: Step through transformations and add monitoring hooks
- **Flexible conversion**: Customize `to_transition` and `to_output` functions for specific needs - **Flexible conversion**: Customize `to_transition` and `to_output` functions for specific needs
Key advantages of the dual format approach: Key advantages of the dual format approach:
- **Environment interaction**: Use tuple format for real-time robot control - **Environment interaction**: Use dictionary format for real-time robot control
- **Training/evaluation**: Use batch format for dataset processing and model training - **Training/evaluation**: Use batch format for dataset processing and model training
- **Seamless integration**: Same processor works with both formats automatically - **Seamless integration**: Same processor works with both formats automatically
- **Backward compatibility**: Existing code using either format continues to work - **Backward compatibility**: Existing code using either format continues to work
src/lerobot/envs/utils.py
+14 -6
View File
@@ -36,17 +36,25 @@ def preprocess_observation(observations: dict[str, np.ndarray]) -> dict[str, Ten
Returns: Returns:
Dictionary of observation batches with keys renamed to LeRobot format and values as tensors. Dictionary of observation batches with keys renamed to LeRobot format and values as tensors.
""" """
from lerobot.processor import RobotProcessor, TransitionIndex, VanillaObservationProcessor from lerobot.processor import RobotProcessor, TransitionKey, VanillaObservationProcessor
# Create processor with observation processor # Create processor with observation processor
processor = RobotProcessor([VanillaObservationProcessor()]) processor = RobotProcessor([VanillaObservationProcessor()])
# Create transition tuple and process # Create transition dictionary and process
transition = (observations, None, None, None, None, None, None) transition = {
processed_transition = processor(transition) TransitionKey.OBSERVATION: observations,
TransitionKey.ACTION: None,
TransitionKey.REWARD: None,
TransitionKey.DONE: None,
TransitionKey.TRUNCATED: None,
TransitionKey.INFO: None,
TransitionKey.COMPLEMENTARY_DATA: None,
}
result = processor(transition)
# Return processed observations # Extract and return the processed observation
return processed_transition[TransitionIndex.OBSERVATION] return result[TransitionKey.OBSERVATION]
def env_to_policy_features(env_cfg: EnvConfig) -> dict[str, PolicyFeature]: def env_to_policy_features(env_cfg: EnvConfig) -> dict[str, PolicyFeature]:
src/lerobot/processor/__init__.py
+2 -2
View File
@@ -32,7 +32,7 @@ from .pipeline import (
ProcessorStepRegistry, ProcessorStepRegistry,
RewardProcessor, RewardProcessor,
RobotProcessor, RobotProcessor,
TransitionIndex, TransitionKey,
TruncatedProcessor, TruncatedProcessor,
) )
from .rename_processor import RenameProcessor from .rename_processor import RenameProcessor
@@ -54,7 +54,7 @@ __all__ = [
"RewardProcessor", "RewardProcessor",
"RobotProcessor", "RobotProcessor",
"StateProcessor", "StateProcessor",
"TransitionIndex", "TransitionKey",
"TruncatedProcessor", "TruncatedProcessor",
"VanillaObservationProcessor", "VanillaObservationProcessor",
] ]
src/lerobot/processor/device_processor.py
+32 -21
View File
@@ -18,7 +18,7 @@ from typing import Any
import torch import torch
from lerobot.processor.pipeline import EnvTransition, TransitionIndex from lerobot.processor.pipeline import EnvTransition, TransitionKey
@dataclass @dataclass
@@ -35,30 +35,41 @@ class DeviceProcessor:
self.non_blocking = "cuda" in self.device self.non_blocking = "cuda" in self.device
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
observation: dict[str, torch.Tensor] = transition[TransitionIndex.OBSERVATION] # Create a copy of the transition
action = transition[TransitionIndex.ACTION] new_transition = transition.copy()
reward = transition[TransitionIndex.REWARD]
done = transition[TransitionIndex.DONE]
truncated = transition[TransitionIndex.TRUNCATED]
info = transition[TransitionIndex.INFO]
complementary_data = transition[TransitionIndex.COMPLEMENTARY_DATA]
# Process observation tensors
observation = transition.get(TransitionKey.OBSERVATION)
if observation is not None: if observation is not None:
observation = { new_observation = {
k: v.to(self.device, non_blocking=self.non_blocking) for k, v in observation.items() k: v.to(self.device, non_blocking=self.non_blocking) if isinstance(v, torch.Tensor) else v
for k, v in observation.items()
} }
if action is not None: new_transition[TransitionKey.OBSERVATION] = new_observation
action = action.to(self.device)
return ( # Process action tensor
observation, action = transition.get(TransitionKey.ACTION)
action, if action is not None and isinstance(action, torch.Tensor):
reward, new_transition[TransitionKey.ACTION] = action.to(self.device, non_blocking=self.non_blocking)
done,
truncated, # Process reward tensor
info, reward = transition.get(TransitionKey.REWARD)
complementary_data, if reward is not None and isinstance(reward, torch.Tensor):
) new_transition[TransitionKey.REWARD] = reward.to(self.device, non_blocking=self.non_blocking)
# Process done tensor
done = transition.get(TransitionKey.DONE)
if done is not None and isinstance(done, torch.Tensor):
new_transition[TransitionKey.DONE] = done.to(self.device, non_blocking=self.non_blocking)
# Process truncated tensor
truncated = transition.get(TransitionKey.TRUNCATED)
if truncated is not None and isinstance(truncated, torch.Tensor):
new_transition[TransitionKey.TRUNCATED] = truncated.to(
self.device, non_blocking=self.non_blocking
)
return new_transition
def get_config(self) -> dict[str, Any]: def get_config(self) -> dict[str, Any]:
"""Return configuration for serialization.""" """Return configuration for serialization."""
src/lerobot/processor/normalize_processor.py
+18 -24
View File
@@ -1,8 +1,8 @@
from __future__ import annotations from __future__ import annotations
from collections.abc import Mapping
from dataclasses import dataclass, field from dataclasses import dataclass, field
from typing import Any from typing import Any
from collections.abc import Mapping
import numpy as np import numpy as np
import torch import torch
@@ -10,7 +10,7 @@ from torch import Tensor
from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature from lerobot.configs.types import FeatureType, NormalizationMode, PolicyFeature
from lerobot.datasets.lerobot_dataset import LeRobotDataset from lerobot.datasets.lerobot_dataset import LeRobotDataset
from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionIndex from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
def _convert_stats_to_tensors(stats: dict[str, dict[str, Any]]) -> dict[str, dict[str, Tensor]]: def _convert_stats_to_tensors(stats: dict[str, dict[str, Any]]) -> dict[str, dict[str, Tensor]]:
@@ -166,17 +166,14 @@ class NormalizerProcessor:
raise ValueError("Action stats must contain either ('mean','std') or ('min','max')") raise ValueError("Action stats must contain either ('mean','std') or ('min','max')")
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
observation = self._normalize_obs(transition[TransitionIndex.OBSERVATION]) observation = self._normalize_obs(transition.get(TransitionKey.OBSERVATION))
action = self._normalize_action(transition[TransitionIndex.ACTION]) action = self._normalize_action(transition.get(TransitionKey.ACTION))
return (
observation, # Create a new transition with normalized values
action, new_transition = transition.copy()
transition[TransitionIndex.REWARD], new_transition[TransitionKey.OBSERVATION] = observation
transition[TransitionIndex.DONE], new_transition[TransitionKey.ACTION] = action
transition[TransitionIndex.TRUNCATED], return new_transition
transition[TransitionIndex.INFO],
transition[TransitionIndex.COMPLEMENTARY_DATA],
)
def get_config(self) -> dict[str, Any]: def get_config(self) -> dict[str, Any]:
config = { config = {
@@ -297,17 +294,14 @@ class UnnormalizerProcessor:
raise ValueError("Action stats must contain either ('mean','std') or ('min','max')") raise ValueError("Action stats must contain either ('mean','std') or ('min','max')")
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
observation = self._unnormalize_obs(transition[TransitionIndex.OBSERVATION]) observation = self._unnormalize_obs(transition.get(TransitionKey.OBSERVATION))
action = self._unnormalize_action(transition[TransitionIndex.ACTION]) action = self._unnormalize_action(transition.get(TransitionKey.ACTION))
return (
observation, # Create a new transition with unnormalized values
action, new_transition = transition.copy()
transition[TransitionIndex.REWARD], new_transition[TransitionKey.OBSERVATION] = observation
transition[TransitionIndex.DONE], new_transition[TransitionKey.ACTION] = action
transition[TransitionIndex.TRUNCATED], return new_transition
transition[TransitionIndex.INFO],
transition[TransitionIndex.COMPLEMENTARY_DATA],
)
def get_config(self) -> dict[str, Any]: def get_config(self) -> dict[str, Any]:
return { return {
src/lerobot/processor/observation_processor.py
+9 -21
View File
@@ -21,7 +21,7 @@ import numpy as np
import torch import torch
from torch import Tensor from torch import Tensor
from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionIndex from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
@dataclass @dataclass
@@ -36,7 +36,7 @@ class ImageProcessor:
""" """
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
observation = transition[TransitionIndex.OBSERVATION] observation = transition.get(TransitionKey.OBSERVATION)
if observation is None: if observation is None:
return transition return transition
@@ -60,15 +60,9 @@ class ImageProcessor:
processed_obs[key] = value processed_obs[key] = value
# Return new transition with processed observation # Return new transition with processed observation
return ( new_transition = transition.copy()
processed_obs, new_transition[TransitionKey.OBSERVATION] = processed_obs
transition[TransitionIndex.ACTION], return new_transition
transition[TransitionIndex.REWARD],
transition[TransitionIndex.DONE],
transition[TransitionIndex.TRUNCATED],
transition[TransitionIndex.INFO],
transition[TransitionIndex.COMPLEMENTARY_DATA],
)
def _process_single_image(self, img: np.ndarray) -> Tensor: def _process_single_image(self, img: np.ndarray) -> Tensor:
"""Process a single image array.""" """Process a single image array."""
@@ -124,7 +118,7 @@ class StateProcessor:
""" """
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
observation = transition[TransitionIndex.OBSERVATION] observation = transition.get(TransitionKey.OBSERVATION)
if observation is None: if observation is None:
return transition return transition
@@ -150,15 +144,9 @@ class StateProcessor:
del processed_obs["agent_pos"] del processed_obs["agent_pos"]
# Return new transition with processed observation # Return new transition with processed observation
return ( new_transition = transition.copy()
processed_obs, new_transition[TransitionKey.OBSERVATION] = processed_obs
transition[TransitionIndex.ACTION], return new_transition
transition[TransitionIndex.REWARD],
transition[TransitionIndex.DONE],
transition[TransitionIndex.TRUNCATED],
transition[TransitionIndex.INFO],
transition[TransitionIndex.COMPLEMENTARY_DATA],
)
def get_config(self) -> dict[str, Any]: def get_config(self) -> dict[str, Any]:
"""Return configuration for serialization.""" """Return configuration for serialization."""
src/lerobot/processor/pipeline.py
+132 -152
View File
@@ -18,39 +18,42 @@ from __future__ import annotations
import importlib import importlib
import json import json
import os import os
from dataclasses import dataclass, field
from enum import IntEnum
from pathlib import Path
from typing import Any, Protocol
from collections.abc import Callable, Iterable, Sequence from collections.abc import Callable, Iterable, Sequence
from dataclasses import dataclass, field
from enum import Enum
from pathlib import Path
from typing import Any, Protocol, TypedDict
import torch import torch
from huggingface_hub import ModelHubMixin, hf_hub_download from huggingface_hub import ModelHubMixin, hf_hub_download
from safetensors.torch import load_file, save_file from safetensors.torch import load_file, save_file
class TransitionIndex(IntEnum): class TransitionKey(str, Enum):
"""Explicit indices for EnvTransition tuple components.""" """Keys for accessing EnvTransition dictionary components."""
OBSERVATION = 0 OBSERVATION = "observation"
ACTION = 1 ACTION = "action"
REWARD = 2 REWARD = "reward"
DONE = 3 DONE = "done"
TRUNCATED = 4 TRUNCATED = "truncated"
INFO = 5 INFO = "info"
COMPLEMENTARY_DATA = 6 COMPLEMENTARY_DATA = "complementary_data"
# (observation, action, reward, done, truncated, info, complementary_data) class EnvTransition(TypedDict, total=False):
EnvTransition = tuple[ """Environment transition data structure.
dict[str, Any] | None, # observation
Any | torch.Tensor | None, # action All fields are optional (total=False) to allow flexible usage.
float | torch.Tensor | None, # reward """
bool | torch.Tensor | None, # done
bool | torch.Tensor | None, # truncated observation: dict[str, Any] | None
dict[str, Any] | None, # info action: Any | torch.Tensor | None
dict[str, Any] | None, # complementary_data reward: float | torch.Tensor | None
] done: bool | torch.Tensor | None
truncated: bool | torch.Tensor | None
info: dict[str, Any] | None
complementary_data: dict[str, Any] | None
class ProcessorStepRegistry: class ProcessorStepRegistry:
@@ -165,10 +168,9 @@ class ProcessorStep(Protocol):
def _default_batch_to_transition(batch: dict[str, Any]) -> EnvTransition: # noqa: D401 def _default_batch_to_transition(batch: dict[str, Any]) -> EnvTransition: # noqa: D401
"""Convert a *batch* dict coming from Learobot replay/dataset code into an """Convert a *batch* dict coming from Learobot replay/dataset code into an
``EnvTransition`` tuple. ``EnvTransition`` dictionary.
The function is intentionally **strictly positional** – it maps well known The function maps well known keys to the EnvTransition structure. Missing keys are
keys to the fixed slot order used inside the pipeline. Missing keys are
filled with sane defaults (``None`` or ``0.0``/``False``). filled with sane defaults (``None`` or ``0.0``/``False``).
Keys recognised (case-sensitive): Keys recognised (case-sensitive):
@@ -193,15 +195,16 @@ def _default_batch_to_transition(batch: dict[str, Any]) -> EnvTransition: # noq
task_key = {"task": batch["task"]} if "task" in batch else {} task_key = {"task": batch["task"]} if "task" in batch else {}
complementary_data = {**pad_keys, **task_key} if pad_keys or task_key else {} complementary_data = {**pad_keys, **task_key} if pad_keys or task_key else {}
return ( transition: EnvTransition = {
observation, TransitionKey.OBSERVATION: observation,
batch.get("action"), TransitionKey.ACTION: batch.get("action"),
batch.get("next.reward", 0.0), TransitionKey.REWARD: batch.get("next.reward", 0.0),
batch.get("next.done", False), TransitionKey.DONE: batch.get("next.done", False),
batch.get("next.truncated", False), TransitionKey.TRUNCATED: batch.get("next.truncated", False),
batch.get("info", {}), TransitionKey.INFO: batch.get("info", {}),
complementary_data, TransitionKey.COMPLEMENTARY_DATA: complementary_data,
) }
return transition
def _default_transition_to_batch(transition: EnvTransition) -> dict[str, Any]: # noqa: D401 def _default_transition_to_batch(transition: EnvTransition) -> dict[str, Any]: # noqa: D401
@@ -209,25 +212,16 @@ def _default_transition_to_batch(transition: EnvTransition) -> dict[str, Any]:
the canonical field names used throughout *LeRobot*. the canonical field names used throughout *LeRobot*.
""" """
(
observation,
action,
reward,
done,
truncated,
info,
complementary_data,
) = transition
batch = { batch = {
"action": action, "action": transition.get(TransitionKey.ACTION),
"next.reward": reward, "next.reward": transition.get(TransitionKey.REWARD, 0.0),
"next.done": done, "next.done": transition.get(TransitionKey.DONE, False),
"next.truncated": truncated, "next.truncated": transition.get(TransitionKey.TRUNCATED, False),
"info": info, "info": transition.get(TransitionKey.INFO, {}),
} }
# Add padding and task data from complementary_data # Add padding and task data from complementary_data
complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
if complementary_data: if complementary_data:
pad_data = {k: v for k, v in complementary_data.items() if "_is_pad" in k} pad_data = {k: v for k, v in complementary_data.items() if "_is_pad" in k}
batch.update(pad_data) batch.update(pad_data)
@@ -236,6 +230,7 @@ def _default_transition_to_batch(transition: EnvTransition) -> dict[str, Any]:
batch["task"] = complementary_data["task"] batch["task"] = complementary_data["task"]
# Handle observation - flatten dict to observation.* keys if it's a dict # Handle observation - flatten dict to observation.* keys if it's a dict
observation = transition.get(TransitionKey.OBSERVATION)
if isinstance(observation, dict): if isinstance(observation, dict):
batch.update(observation) batch.update(observation)
@@ -293,33 +288,35 @@ class RobotProcessor(ModelHubMixin):
def __call__(self, data: EnvTransition | dict[str, Any]): def __call__(self, data: EnvTransition | dict[str, Any]):
"""Process data through all steps. """Process data through all steps.
The method accepts either the classic EnvTransition tuple or a batch dictionary The method accepts either the classic EnvTransition dict or a batch dictionary
(like the ones returned by ReplayBuffer or LeRobotDataset). If a dict is supplied (like the ones returned by ReplayBuffer or LeRobotDataset). If a dict is supplied
it is first converted to the internal tuple format using to_transition; after all it is first converted to the internal dict format using to_transition; after all
steps are executed the tuple is transformed back into a dict with to_batch and the steps are executed the dict is transformed back into a batch dict with to_batch and the
result is returned – thereby preserving the caller's original data type. result is returned – thereby preserving the caller's original data type.
Args: Args:
data: Either an EnvTransition tuple or a batch dictionary to process. data: Either an EnvTransition dict or a batch dictionary to process.
Returns: Returns:
The processed data in the same format as the input (tuple or dict). The processed data in the same format as the input (EnvTransition or batch dict).
Raises: Raises:
ValueError: If the transition is not a valid 7-tuple format. ValueError: If the transition is not a valid EnvTransition format.
""" """
called_with_batch = isinstance(data, dict) # Check if data is already an EnvTransition or needs conversion
if isinstance(data, dict) and not all(isinstance(k, TransitionKey) for k in data.keys()):
# It's a batch dict, convert it
called_with_batch = True
transition = self.to_transition(data)
else:
# It's already an EnvTransition
called_with_batch = False
transition = data
transition = self.to_transition(data) if called_with_batch else data # Basic validation
if not isinstance(transition, dict):
# Basic validation with helpful error message for tuple input raise ValueError(f"EnvTransition must be a dictionary. Got {type(transition).__name__}")
if not isinstance(transition, tuple) or len(transition) != 7:
raise ValueError(
"EnvTransition must be a 7-tuple of (observation, action, reward, done, "
"truncated, info, complementary_data). "
f"Got {type(transition).__name__} with length {len(transition) if hasattr(transition, '__len__') else 'unknown'}."
)
for idx, processor_step in enumerate(self.steps): for idx, processor_step in enumerate(self.steps):
for hook in self.before_step_hooks: for hook in self.before_step_hooks:
@@ -339,25 +336,28 @@ class RobotProcessor(ModelHubMixin):
def step_through(self, data: EnvTransition | dict[str, Any]) -> Iterable[EnvTransition | dict[str, Any]]: def step_through(self, data: EnvTransition | dict[str, Any]) -> Iterable[EnvTransition | dict[str, Any]]:
"""Yield the intermediate results after each processor step. """Yield the intermediate results after each processor step.
Like __call__, this method accepts either EnvTransition tuples or batch dictionaries Like __call__, this method accepts either EnvTransition dicts or batch dictionaries
and preserves the input format in the yielded results. and preserves the input format in the yielded results.
Args: Args:
data: Either an EnvTransition tuple or a batch dictionary to process. data: Either an EnvTransition dict or a batch dictionary to process.
Yields: Yields:
The intermediate results after each step, in the same format as the input. The intermediate results after each step, in the same format as the input.
""" """
called_with_batch = isinstance(data, dict) # Check if data is already an EnvTransition or needs conversion
transition = self.to_transition(data) if called_with_batch else data if isinstance(data, dict) and not all(isinstance(k, TransitionKey) for k in data.keys()):
# It's a batch dict, convert it
called_with_batch = True
transition = self.to_transition(data)
else:
# It's already an EnvTransition
called_with_batch = False
transition = data
# Basic validation with helpful error message for tuple input # Basic validation
if not isinstance(transition, tuple) or len(transition) != 7: if not isinstance(transition, dict):
raise ValueError( raise ValueError(f"EnvTransition must be a dictionary. Got {type(transition).__name__}")
"EnvTransition must be a 7-tuple of (observation, action, reward, done, "
"truncated, info, complementary_data). "
f"Got {type(transition).__name__} with length {len(transition) if hasattr(transition, '__len__') else 'unknown'}."
)
# Yield initial state # Yield initial state
yield self.to_output(transition) if called_with_batch else transition yield self.to_output(transition) if called_with_batch else transition
@@ -684,7 +684,7 @@ class ObservationProcessor:
Subclasses should override the `observation` method to implement custom observation processing. Subclasses should override the `observation` method to implement custom observation processing.
This class handles the boilerplate of extracting and reinserting the processed observation This class handles the boilerplate of extracting and reinserting the processed observation
into the transition tuple, eliminating the need to implement the `__call__` method in subclasses. into the transition dict, eliminating the need to implement the `__call__` method in subclasses.
Example: Example:
```python ```python
@@ -696,7 +696,7 @@ class ObservationProcessor:
return observation * self.scale_factor return observation * self.scale_factor
``` ```
By inheriting from this class, you avoid writing repetitive code to handle transition tuple By inheriting from this class, you avoid writing repetitive code to handle transition dict
manipulation, focusing only on the specific observation processing logic. manipulation, focusing only on the specific observation processing logic.
""" """
@@ -712,10 +712,12 @@ class ObservationProcessor:
return observation return observation
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
observation = transition[TransitionIndex.OBSERVATION] observation = transition.get(TransitionKey.OBSERVATION)
observation = self.observation(observation) processed_observation = self.observation(observation)
transition = (observation, *transition[TransitionIndex.ACTION :]) # Create a new transition dict with the processed observation
return transition new_transition = transition.copy()
new_transition[TransitionKey.OBSERVATION] = processed_observation
return new_transition
class ActionProcessor: class ActionProcessor:
@@ -723,7 +725,7 @@ class ActionProcessor:
Subclasses should override the `action` method to implement custom action processing. Subclasses should override the `action` method to implement custom action processing.
This class handles the boilerplate of extracting and reinserting the processed action This class handles the boilerplate of extracting and reinserting the processed action
into the transition tuple, eliminating the need to implement the `__call__` method in subclasses. into the transition dict, eliminating the need to implement the `__call__` method in subclasses.
Example: Example:
```python ```python
@@ -736,7 +738,7 @@ class ActionProcessor:
return np.clip(action, self.min_val, self.max_val) return np.clip(action, self.min_val, self.max_val)
``` ```
By inheriting from this class, you avoid writing repetitive code to handle transition tuple By inheriting from this class, you avoid writing repetitive code to handle transition dict
manipulation, focusing only on the specific action processing logic. manipulation, focusing only on the specific action processing logic.
""" """
@@ -752,10 +754,12 @@ class ActionProcessor:
return action return action
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
action = transition[TransitionIndex.ACTION] action = transition.get(TransitionKey.ACTION)
action = self.action(action) processed_action = self.action(action)
transition = (transition[TransitionIndex.OBSERVATION], action, *transition[TransitionIndex.REWARD :]) # Create a new transition dict with the processed action
return transition new_transition = transition.copy()
new_transition[TransitionKey.ACTION] = processed_action
return new_transition
class RewardProcessor: class RewardProcessor:
@@ -763,7 +767,7 @@ class RewardProcessor:
Subclasses should override the `reward` method to implement custom reward processing. Subclasses should override the `reward` method to implement custom reward processing.
This class handles the boilerplate of extracting and reinserting the processed reward This class handles the boilerplate of extracting and reinserting the processed reward
into the transition tuple, eliminating the need to implement the `__call__` method in subclasses. into the transition dict, eliminating the need to implement the `__call__` method in subclasses.
Example: Example:
```python ```python
@@ -775,7 +779,7 @@ class RewardProcessor:
return reward * self.scale_factor return reward * self.scale_factor
``` ```
By inheriting from this class, you avoid writing repetitive code to handle transition tuple By inheriting from this class, you avoid writing repetitive code to handle transition dict
manipulation, focusing only on the specific reward processing logic. manipulation, focusing only on the specific reward processing logic.
""" """
@@ -791,15 +795,12 @@ class RewardProcessor:
return reward return reward
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
reward = transition[TransitionIndex.REWARD] reward = transition.get(TransitionKey.REWARD)
reward = self.reward(reward) processed_reward = self.reward(reward)
transition = ( # Create a new transition dict with the processed reward
transition[TransitionIndex.OBSERVATION], new_transition = transition.copy()
transition[TransitionIndex.ACTION], new_transition[TransitionKey.REWARD] = processed_reward
reward, return new_transition
*transition[TransitionIndex.DONE :],
)
return transition
class DoneProcessor: class DoneProcessor:
@@ -807,7 +808,7 @@ class DoneProcessor:
Subclasses should override the `done` method to implement custom done flag processing. Subclasses should override the `done` method to implement custom done flag processing.
This class handles the boilerplate of extracting and reinserting the processed done flag This class handles the boilerplate of extracting and reinserting the processed done flag
into the transition tuple, eliminating the need to implement the `__call__` method in subclasses. into the transition dict, eliminating the need to implement the `__call__` method in subclasses.
Example: Example:
```python ```python
@@ -824,7 +825,7 @@ class DoneProcessor:
self.steps = 0 self.steps = 0
``` ```
By inheriting from this class, you avoid writing repetitive code to handle transition tuple By inheriting from this class, you avoid writing repetitive code to handle transition dict
manipulation, focusing only on the specific done flag processing logic. manipulation, focusing only on the specific done flag processing logic.
""" """
@@ -840,16 +841,12 @@ class DoneProcessor:
return done return done
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
done = transition[TransitionIndex.DONE] done = transition.get(TransitionKey.DONE)
done = self.done(done) processed_done = self.done(done)
transition = ( # Create a new transition dict with the processed done flag
transition[TransitionIndex.OBSERVATION], new_transition = transition.copy()
transition[TransitionIndex.ACTION], new_transition[TransitionKey.DONE] = processed_done
transition[TransitionIndex.REWARD], return new_transition
done,
*transition[TransitionIndex.TRUNCATED :],
)
return transition
class TruncatedProcessor: class TruncatedProcessor:
@@ -857,7 +854,7 @@ class TruncatedProcessor:
Subclasses should override the `truncated` method to implement custom truncated flag processing. Subclasses should override the `truncated` method to implement custom truncated flag processing.
This class handles the boilerplate of extracting and reinserting the processed truncated flag This class handles the boilerplate of extracting and reinserting the processed truncated flag
into the transition tuple, eliminating the need to implement the `__call__` method in subclasses. into the transition dict, eliminating the need to implement the `__call__` method in subclasses.
Example: Example:
```python ```python
@@ -870,7 +867,7 @@ class TruncatedProcessor:
return truncated or some_condition > self.threshold return truncated or some_condition > self.threshold
``` ```
By inheriting from this class, you avoid writing repetitive code to handle transition tuple By inheriting from this class, you avoid writing repetitive code to handle transition dict
manipulation, focusing only on the specific truncated flag processing logic. manipulation, focusing only on the specific truncated flag processing logic.
""" """
@@ -886,17 +883,12 @@ class TruncatedProcessor:
return truncated return truncated
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
truncated = transition[TransitionIndex.TRUNCATED] truncated = transition.get(TransitionKey.TRUNCATED)
truncated = self.truncated(truncated) processed_truncated = self.truncated(truncated)
transition = ( # Create a new transition dict with the processed truncated flag
transition[TransitionIndex.OBSERVATION], new_transition = transition.copy()
transition[TransitionIndex.ACTION], new_transition[TransitionKey.TRUNCATED] = processed_truncated
transition[TransitionIndex.REWARD], return new_transition
transition[TransitionIndex.DONE],
truncated,
*transition[TransitionIndex.INFO :],
)
return transition
class InfoProcessor: class InfoProcessor:
@@ -904,7 +896,7 @@ class InfoProcessor:
Subclasses should override the `info` method to implement custom info processing. Subclasses should override the `info` method to implement custom info processing.
This class handles the boilerplate of extracting and reinserting the processed info This class handles the boilerplate of extracting and reinserting the processed info
into the transition tuple, eliminating the need to implement the `__call__` method in subclasses. into the transition dict, eliminating the need to implement the `__call__` method in subclasses.
Example: Example:
```python ```python
@@ -922,7 +914,7 @@ class InfoProcessor:
self.step_count = 0 self.step_count = 0
``` ```
By inheriting from this class, you avoid writing repetitive code to handle transition tuple By inheriting from this class, you avoid writing repetitive code to handle transition dict
manipulation, focusing only on the specific info dictionary processing logic. manipulation, focusing only on the specific info dictionary processing logic.
""" """
@@ -938,18 +930,12 @@ class InfoProcessor:
return info return info
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
info = transition[TransitionIndex.INFO] info = transition.get(TransitionKey.INFO)
info = self.info(info) processed_info = self.info(info)
transition = ( # Create a new transition dict with the processed info
transition[TransitionIndex.OBSERVATION], new_transition = transition.copy()
transition[TransitionIndex.ACTION], new_transition[TransitionKey.INFO] = processed_info
transition[TransitionIndex.REWARD], return new_transition
transition[TransitionIndex.DONE],
transition[TransitionIndex.TRUNCATED],
info,
*transition[TransitionIndex.COMPLEMENTARY_DATA :],
)
return transition
class ComplementaryDataProcessor: class ComplementaryDataProcessor:
@@ -957,7 +943,7 @@ class ComplementaryDataProcessor:
Subclasses should override the `complementary_data` method to implement custom complementary data processing. Subclasses should override the `complementary_data` method to implement custom complementary data processing.
This class handles the boilerplate of extracting and reinserting the processed complementary data This class handles the boilerplate of extracting and reinserting the processed complementary data
into the transition tuple, eliminating the need to implement the `__call__` method in subclasses. into the transition dict, eliminating the need to implement the `__call__` method in subclasses.
""" """
def complementary_data(self, complementary_data): def complementary_data(self, complementary_data):
@@ -972,18 +958,12 @@ class ComplementaryDataProcessor:
return complementary_data return complementary_data
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
complementary_data = transition[TransitionIndex.COMPLEMENTARY_DATA] complementary_data = transition.get(TransitionKey.COMPLEMENTARY_DATA)
complementary_data = self.complementary_data(complementary_data) processed_complementary_data = self.complementary_data(complementary_data)
transition = ( # Create a new transition dict with the processed complementary data
transition[TransitionIndex.OBSERVATION], new_transition = transition.copy()
transition[TransitionIndex.ACTION], new_transition[TransitionKey.COMPLEMENTARY_DATA] = processed_complementary_data
transition[TransitionIndex.REWARD], return new_transition
transition[TransitionIndex.DONE],
transition[TransitionIndex.TRUNCATED],
transition[TransitionIndex.INFO],
complementary_data,
)
return transition
class IdentityProcessor: class IdentityProcessor:
src/lerobot/processor/rename_processor.py
+7 -11
View File
@@ -18,7 +18,7 @@ from typing import Any
import torch import torch
from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionIndex from lerobot.processor.pipeline import EnvTransition, ProcessorStepRegistry, TransitionKey
@dataclass @dataclass
@@ -29,7 +29,7 @@ class RenameProcessor:
rename_map: dict[str, str] = field(default_factory=dict) rename_map: dict[str, str] = field(default_factory=dict)
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
observation = transition[TransitionIndex.OBSERVATION] observation = transition.get(TransitionKey.OBSERVATION)
if observation is None: if observation is None:
return transition return transition
@@ -39,15 +39,11 @@ class RenameProcessor:
processed_obs[self.rename_map[key]] = value processed_obs[self.rename_map[key]] = value
else: else:
processed_obs[key] = value processed_obs[key] = value
return (
processed_obs, # Create a new transition with the renamed observation
transition[TransitionIndex.ACTION], new_transition = transition.copy()
transition[TransitionIndex.REWARD], new_transition[TransitionKey.OBSERVATION] = processed_obs
transition[TransitionIndex.DONE], return new_transition
transition[TransitionIndex.TRUNCATED],
transition[TransitionIndex.INFO],
transition[TransitionIndex.COMPLEMENTARY_DATA],
)
def get_config(self) -> dict[str, Any]: def get_config(self) -> dict[str, Any]:
return {"rename_map": self.rename_map} return {"rename_map": self.rename_map}
src/lerobot/scripts/eval.py
+3 -3
View File
@@ -72,7 +72,7 @@ from lerobot.envs.utils import add_envs_task, check_env_attributes_and_types
from lerobot.policies.factory import make_policy from lerobot.policies.factory import make_policy
from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.policies.utils import get_device_from_parameters from lerobot.policies.utils import get_device_from_parameters
from lerobot.processor import RobotProcessor, TransitionIndex, VanillaObservationProcessor from lerobot.processor import RobotProcessor, TransitionKey, VanillaObservationProcessor
from lerobot.utils.io_utils import write_video from lerobot.utils.io_utils import write_video
from lerobot.utils.random_utils import set_seed from lerobot.utils.random_utils import set_seed
from lerobot.utils.utils import ( from lerobot.utils.utils import (
@@ -160,7 +160,7 @@ def rollout(
# Numpy array to tensor and changing dictionary keys to LeRobot policy format. # Numpy array to tensor and changing dictionary keys to LeRobot policy format.
transition = (observation, None, None, None, None, None, None) transition = (observation, None, None, None, None, None, None)
processed_transition = obs_processor(transition) processed_transition = obs_processor(transition)
observation = processed_transition[TransitionIndex.OBSERVATION] observation = processed_transition[TransitionKey.OBSERVATION]
if return_observations: if return_observations:
all_observations.append(deepcopy(observation)) all_observations.append(deepcopy(observation))
@@ -211,7 +211,7 @@ def rollout(
if return_observations: if return_observations:
transition = (observation, None, None, None, None, None, None) transition = (observation, None, None, None, None, None, None)
processed_transition = obs_processor(transition) processed_transition = obs_processor(transition)
observation = processed_transition[TransitionIndex.OBSERVATION] observation = processed_transition[TransitionKey.OBSERVATION]
all_observations.append(deepcopy(observation)) all_observations.append(deepcopy(observation))
# Stack the sequence along the first dimension so that we have (batch, sequence, *) tensors. # Stack the sequence along the first dimension so that we have (batch, sequence, *) tensors.
tests/envs/test_envs.py
+2 -2
View File
@@ -22,7 +22,7 @@ from gymnasium.utils.env_checker import check_env
import lerobot import lerobot
from lerobot.envs.factory import make_env, make_env_config from lerobot.envs.factory import make_env, make_env_config
from lerobot.processor import RobotProcessor, TransitionIndex, VanillaObservationProcessor from lerobot.processor import RobotProcessor, TransitionKey, VanillaObservationProcessor
from tests.utils import require_env from tests.utils import require_env
OBS_TYPES = ["state", "pixels", "pixels_agent_pos"] OBS_TYPES = ["state", "pixels", "pixels_agent_pos"]
@@ -53,7 +53,7 @@ def test_factory(env_name):
obs_processor = RobotProcessor([VanillaObservationProcessor()]) obs_processor = RobotProcessor([VanillaObservationProcessor()])
transition = (obs, None, None, None, None, None, None) transition = (obs, None, None, None, None, None, None)
processed_transition = obs_processor(transition) processed_transition = obs_processor(transition)
obs = processed_transition[TransitionIndex.OBSERVATION] obs = processed_transition[TransitionKey.OBSERVATION]
# test image keys are float32 in range [0,1] # test image keys are float32 in range [0,1]
for key in obs: for key in obs:
tests/policies/test_policies.py
+2 -2
View File
@@ -39,7 +39,7 @@ from lerobot.policies.factory import (
) )
from lerobot.policies.normalize import Normalize, Unnormalize from lerobot.policies.normalize import Normalize, Unnormalize
from lerobot.policies.pretrained import PreTrainedPolicy from lerobot.policies.pretrained import PreTrainedPolicy
from lerobot.processor import RobotProcessor, TransitionIndex, VanillaObservationProcessor from lerobot.processor import RobotProcessor, TransitionKey, VanillaObservationProcessor
from lerobot.utils.random_utils import seeded_context from lerobot.utils.random_utils import seeded_context
from tests.artifacts.policies.save_policy_to_safetensors import get_policy_stats from tests.artifacts.policies.save_policy_to_safetensors import get_policy_stats
from tests.utils import DEVICE, require_cpu, require_env, require_x86_64_kernel from tests.utils import DEVICE, require_cpu, require_env, require_x86_64_kernel
@@ -188,7 +188,7 @@ def test_policy(ds_repo_id, env_name, env_kwargs, policy_name, policy_kwargs):
obs_processor = RobotProcessor([VanillaObservationProcessor()]) obs_processor = RobotProcessor([VanillaObservationProcessor()])
transition = (observation, None, None, None, None, None, None) transition = (observation, None, None, None, None, None, None)
processed_transition = obs_processor(transition) processed_transition = obs_processor(transition)
observation = processed_transition[TransitionIndex.OBSERVATION] observation = processed_transition[TransitionKey.OBSERVATION]
# send observation to device/gpu # send observation to device/gpu
observation = {key: observation[key].to(DEVICE, non_blocking=True) for key in observation} observation = {key: observation[key].to(DEVICE, non_blocking=True) for key in observation}
tests/processor/test_batch_conversion.py
+59 -65
View File
@@ -2,7 +2,7 @@ import torch
from lerobot.processor.pipeline import ( from lerobot.processor.pipeline import (
RobotProcessor, RobotProcessor,
TransitionIndex, TransitionKey,
_default_batch_to_transition, _default_batch_to_transition,
_default_transition_to_batch, _default_transition_to_batch,
) )
@@ -63,27 +63,27 @@ def test_batch_to_transition_observation_grouping():
transition = _default_batch_to_transition(batch) transition = _default_batch_to_transition(batch)
# Check observation is a dict with all observation.* keys # Check observation is a dict with all observation.* keys
assert isinstance(transition[TransitionIndex.OBSERVATION], dict) assert isinstance(transition[TransitionKey.OBSERVATION], dict)
assert "observation.image.top" in transition[TransitionIndex.OBSERVATION] assert "observation.image.top" in transition[TransitionKey.OBSERVATION]
assert "observation.image.left" in transition[TransitionIndex.OBSERVATION] assert "observation.image.left" in transition[TransitionKey.OBSERVATION]
assert "observation.state" in transition[TransitionIndex.OBSERVATION] assert "observation.state" in transition[TransitionKey.OBSERVATION]
# Check values are preserved # Check values are preserved
assert torch.allclose( assert torch.allclose(
transition[TransitionIndex.OBSERVATION]["observation.image.top"], batch["observation.image.top"] transition[TransitionKey.OBSERVATION]["observation.image.top"], batch["observation.image.top"]
) )
assert torch.allclose( assert torch.allclose(
transition[TransitionIndex.OBSERVATION]["observation.image.left"], batch["observation.image.left"] transition[TransitionKey.OBSERVATION]["observation.image.left"], batch["observation.image.left"]
) )
assert transition[TransitionIndex.OBSERVATION]["observation.state"] == [1, 2, 3, 4] assert transition[TransitionKey.OBSERVATION]["observation.state"] == [1, 2, 3, 4]
# Check other fields # Check other fields
assert transition[TransitionIndex.ACTION] == "action_data" assert transition[TransitionKey.ACTION] == "action_data"
assert transition[TransitionIndex.REWARD] == 1.5 assert transition[TransitionKey.REWARD] == 1.5
assert transition[TransitionIndex.DONE] assert transition[TransitionKey.DONE]
assert not transition[TransitionIndex.TRUNCATED] assert not transition[TransitionKey.TRUNCATED]
assert transition[TransitionIndex.INFO] == {"episode": 42} assert transition[TransitionKey.INFO] == {"episode": 42}
assert transition[TransitionIndex.COMPLEMENTARY_DATA] == {} assert transition[TransitionKey.COMPLEMENTARY_DATA] == {}
def test_transition_to_batch_observation_flattening(): def test_transition_to_batch_observation_flattening():
@@ -94,15 +94,15 @@ def test_transition_to_batch_observation_flattening():
"observation.state": [1, 2, 3, 4], "observation.state": [1, 2, 3, 4],
} }
transition = ( transition = {
observation_dict, # observation TransitionKey.OBSERVATION: observation_dict,
"action_data", # action TransitionKey.ACTION: "action_data",
1.5, # reward TransitionKey.REWARD: 1.5,
True, # done TransitionKey.DONE: True,
False, # truncated TransitionKey.TRUNCATED: False,
{"episode": 42}, # info TransitionKey.INFO: {"episode": 42},
{}, # complementary_data TransitionKey.COMPLEMENTARY_DATA: {},
) }
batch = _default_transition_to_batch(transition) batch = _default_transition_to_batch(transition)
@@ -137,14 +137,14 @@ def test_no_observation_keys():
transition = _default_batch_to_transition(batch) transition = _default_batch_to_transition(batch)
# Observation should be None when no observation.* keys # Observation should be None when no observation.* keys
assert transition[TransitionIndex.OBSERVATION] is None assert transition[TransitionKey.OBSERVATION] is None
# Check other fields # Check other fields
assert transition[TransitionIndex.ACTION] == "action_data" assert transition[TransitionKey.ACTION] == "action_data"
assert transition[TransitionIndex.REWARD] == 2.0 assert transition[TransitionKey.REWARD] == 2.0
assert not transition[TransitionIndex.DONE] assert not transition[TransitionKey.DONE]
assert transition[TransitionIndex.TRUNCATED] assert transition[TransitionKey.TRUNCATED]
assert transition[TransitionIndex.INFO] == {"test": "no_obs"} assert transition[TransitionKey.INFO] == {"test": "no_obs"}
# Round trip should work # Round trip should work
reconstructed_batch = _default_transition_to_batch(transition) reconstructed_batch = _default_transition_to_batch(transition)
@@ -162,15 +162,15 @@ def test_minimal_batch():
transition = _default_batch_to_transition(batch) transition = _default_batch_to_transition(batch)
# Check observation # Check observation
assert transition[TransitionIndex.OBSERVATION] == {"observation.state": "minimal_state"} assert transition[TransitionKey.OBSERVATION] == {"observation.state": "minimal_state"}
assert transition[TransitionIndex.ACTION] == "minimal_action" assert transition[TransitionKey.ACTION] == "minimal_action"
# Check defaults # Check defaults
assert transition[TransitionIndex.REWARD] == 0.0 assert transition[TransitionKey.REWARD] == 0.0
assert not transition[TransitionIndex.DONE] assert not transition[TransitionKey.DONE]
assert not transition[TransitionIndex.TRUNCATED] assert not transition[TransitionKey.TRUNCATED]
assert transition[TransitionIndex.INFO] == {} assert transition[TransitionKey.INFO] == {}
assert transition[TransitionIndex.COMPLEMENTARY_DATA] == {} assert transition[TransitionKey.COMPLEMENTARY_DATA] == {}
# Round trip # Round trip
reconstructed_batch = _default_transition_to_batch(transition) reconstructed_batch = _default_transition_to_batch(transition)
@@ -189,13 +189,13 @@ def test_empty_batch():
transition = _default_batch_to_transition(batch) transition = _default_batch_to_transition(batch)
# All fields should have defaults # All fields should have defaults
assert transition[TransitionIndex.OBSERVATION] is None assert transition[TransitionKey.OBSERVATION] is None
assert transition[TransitionIndex.ACTION] is None assert transition[TransitionKey.ACTION] is None
assert transition[TransitionIndex.REWARD] == 0.0 assert transition[TransitionKey.REWARD] == 0.0
assert not transition[TransitionIndex.DONE] assert not transition[TransitionKey.DONE]
assert not transition[TransitionIndex.TRUNCATED] assert not transition[TransitionKey.TRUNCATED]
assert transition[TransitionIndex.INFO] == {} assert transition[TransitionKey.INFO] == {}
assert transition[TransitionIndex.COMPLEMENTARY_DATA] == {} assert transition[TransitionKey.COMPLEMENTARY_DATA] == {}
# Round trip # Round trip
reconstructed_batch = _default_transition_to_batch(transition) reconstructed_batch = _default_transition_to_batch(transition)
@@ -256,33 +256,27 @@ def test_custom_converter():
# Custom converter that modifies the reward # Custom converter that modifies the reward
tr = _default_batch_to_transition(batch) tr = _default_batch_to_transition(batch)
# Double the reward # Double the reward
reward = tr[TransitionIndex.REWARD] * 2 if tr[TransitionIndex.REWARD] is not None else 0.0 reward = tr.get(TransitionKey.REWARD, 0.0)
return ( new_tr = tr.copy()
tr[TransitionIndex.OBSERVATION], new_tr[TransitionKey.REWARD] = reward * 2 if reward is not None else 0.0
tr[TransitionIndex.ACTION], return new_tr
reward,
tr[TransitionIndex.DONE],
tr[TransitionIndex.TRUNCATED],
tr[TransitionIndex.INFO],
tr[TransitionIndex.COMPLEMENTARY_DATA],
)
def to_batch(tr): def to_batch(tr):
# Custom converter that adds a custom field
batch = _default_transition_to_batch(tr) batch = _default_transition_to_batch(tr)
batch["custom_field"] = "custom_value"
return batch return batch
proc = RobotProcessor([], to_transition=to_tr, to_output=to_batch) processor = RobotProcessor(steps=[], to_transition=to_tr, to_output=to_batch)
batch = _dummy_batch()
out = proc(batch)
# Check that custom modifications were applied batch = {
assert out["next.reward"] == batch["next.reward"] * 2 "observation.state": torch.randn(1, 4),
assert out["custom_field"] == "custom_value" "action": torch.randn(1, 2),
"next.reward": 1.0,
"next.done": False,
}
# Check that observation.* keys are still preserved result = processor(batch)
original_obs_keys = {k: v for k, v in batch.items() if k.startswith("observation.")}
output_obs_keys = {k: v for k, v in out.items() if k.startswith("observation.")}
assert set(original_obs_keys.keys()) == set(output_obs_keys.keys()) # Check the reward was doubled by our custom converter
assert result["next.reward"] == 2.0
assert torch.allclose(result["observation.state"], batch["observation.state"])
assert torch.allclose(result["action"], batch["action"])
tests/processor/test_normalize_processor.py
+94 -35
View File
@@ -1,3 +1,18 @@
#!/usr/bin/env python
# Copyright 2025 The HuggingFace Inc. team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
from unittest.mock import Mock from unittest.mock import Mock
import numpy as np import numpy as np
@@ -10,7 +25,22 @@ from lerobot.processor.normalize_processor import (
UnnormalizerProcessor, UnnormalizerProcessor,
_convert_stats_to_tensors, _convert_stats_to_tensors,
) )
from lerobot.processor.pipeline import RobotProcessor, TransitionIndex from lerobot.processor.pipeline import RobotProcessor, TransitionKey
def create_transition(
observation=None, action=None, reward=None, done=None, truncated=None, info=None, complementary_data=None
):
"""Helper to create an EnvTransition dictionary."""
return {
TransitionKey.OBSERVATION: observation,
TransitionKey.ACTION: action,
TransitionKey.REWARD: reward,
TransitionKey.DONE: done,
TransitionKey.TRUNCATED: truncated,
TransitionKey.INFO: info,
TransitionKey.COMPLEMENTARY_DATA: complementary_data,
}
def test_numpy_conversion(): def test_numpy_conversion():
@@ -120,10 +150,10 @@ def test_mean_std_normalization(observation_normalizer):
"observation.image": torch.tensor([0.7, 0.5, 0.3]), "observation.image": torch.tensor([0.7, 0.5, 0.3]),
"observation.state": torch.tensor([0.5, 0.0]), "observation.state": torch.tensor([0.5, 0.0]),
} }
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
normalized_transition = observation_normalizer(transition) normalized_transition = observation_normalizer(transition)
normalized_obs = normalized_transition[TransitionIndex.OBSERVATION] normalized_obs = normalized_transition[TransitionKey.OBSERVATION]
# Check mean/std normalization # Check mean/std normalization
expected_image = (torch.tensor([0.7, 0.5, 0.3]) - 0.5) / 0.2 expected_image = (torch.tensor([0.7, 0.5, 0.3]) - 0.5) / 0.2
@@ -134,10 +164,10 @@ def test_min_max_normalization(observation_normalizer):
observation = { observation = {
"observation.state": torch.tensor([0.5, 0.0]), "observation.state": torch.tensor([0.5, 0.0]),
} }
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
normalized_transition = observation_normalizer(transition) normalized_transition = observation_normalizer(transition)
normalized_obs = normalized_transition[TransitionIndex.OBSERVATION] normalized_obs = normalized_transition[TransitionKey.OBSERVATION]
# Check min/max normalization to [-1, 1] # Check min/max normalization to [-1, 1]
# For state[0]: 2 * (0.5 - 0.0) / (1.0 - 0.0) - 1 = 0.0 # For state[0]: 2 * (0.5 - 0.0) / (1.0 - 0.0) - 1 = 0.0
@@ -157,10 +187,10 @@ def test_selective_normalization(observation_stats):
"observation.image": torch.tensor([0.7, 0.5, 0.3]), "observation.image": torch.tensor([0.7, 0.5, 0.3]),
"observation.state": torch.tensor([0.5, 0.0]), "observation.state": torch.tensor([0.5, 0.0]),
} }
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
normalized_transition = normalizer(transition) normalized_transition = normalizer(transition)
normalized_obs = normalized_transition[TransitionIndex.OBSERVATION] normalized_obs = normalized_transition[TransitionKey.OBSERVATION]
# Only image should be normalized # Only image should be normalized
assert torch.allclose(normalized_obs["observation.image"], (torch.tensor([0.7, 0.5, 0.3]) - 0.5) / 0.2) assert torch.allclose(normalized_obs["observation.image"], (torch.tensor([0.7, 0.5, 0.3]) - 0.5) / 0.2)
@@ -176,10 +206,10 @@ def test_device_compatibility(observation_stats):
observation = { observation = {
"observation.image": torch.tensor([0.7, 0.5, 0.3]).cuda(), "observation.image": torch.tensor([0.7, 0.5, 0.3]).cuda(),
} }
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
normalized_transition = normalizer(transition) normalized_transition = normalizer(transition)
normalized_obs = normalized_transition[TransitionIndex.OBSERVATION] normalized_obs = normalized_transition[TransitionKey.OBSERVATION]
assert normalized_obs["observation.image"].device.type == "cuda" assert normalized_obs["observation.image"].device.type == "cuda"
@@ -220,10 +250,10 @@ def test_state_dict_save_load(observation_normalizer):
# Test that it works the same # Test that it works the same
observation = {"observation.image": torch.tensor([0.7, 0.5, 0.3])} observation = {"observation.image": torch.tensor([0.7, 0.5, 0.3])}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result1 = observation_normalizer(transition)[0] result1 = observation_normalizer(transition)[TransitionKey.OBSERVATION]
result2 = new_normalizer(transition)[0] result2 = new_normalizer(transition)[TransitionKey.OBSERVATION]
assert torch.allclose(result1["observation.image"], result2["observation.image"]) assert torch.allclose(result1["observation.image"], result2["observation.image"])
@@ -271,10 +301,10 @@ def test_mean_std_unnormalization(action_stats_mean_std):
) )
normalized_action = torch.tensor([1.0, -0.5, 2.0]) normalized_action = torch.tensor([1.0, -0.5, 2.0])
transition = (None, normalized_action, None, None, None, None, None) transition = create_transition(action=normalized_action)
unnormalized_transition = unnormalizer(transition) unnormalized_transition = unnormalizer(transition)
unnormalized_action = unnormalized_transition[TransitionIndex.ACTION] unnormalized_action = unnormalized_transition[TransitionKey.ACTION]
# action * std + mean # action * std + mean
expected = torch.tensor([1.0 * 1.0 + 0.0, -0.5 * 2.0 + 0.0, 2.0 * 0.5 + 0.0]) expected = torch.tensor([1.0 * 1.0 + 0.0, -0.5 * 2.0 + 0.0, 2.0 * 0.5 + 0.0])
@@ -290,10 +320,10 @@ def test_min_max_unnormalization(action_stats_min_max):
# Actions in [-1, 1] # Actions in [-1, 1]
normalized_action = torch.tensor([0.0, -1.0, 1.0]) normalized_action = torch.tensor([0.0, -1.0, 1.0])
transition = (None, normalized_action, None, None, None, None, None) transition = create_transition(action=normalized_action)
unnormalized_transition = unnormalizer(transition) unnormalized_transition = unnormalizer(transition)
unnormalized_action = unnormalized_transition[TransitionIndex.ACTION] unnormalized_action = unnormalized_transition[TransitionKey.ACTION]
# Map from [-1, 1] to [min, max] # Map from [-1, 1] to [min, max]
# (action + 1) / 2 * (max - min) + min # (action + 1) / 2 * (max - min) + min
@@ -315,10 +345,10 @@ def test_numpy_action_input(action_stats_mean_std):
) )
normalized_action = np.array([1.0, -0.5, 2.0], dtype=np.float32) normalized_action = np.array([1.0, -0.5, 2.0], dtype=np.float32)
transition = (None, normalized_action, None, None, None, None, None) transition = create_transition(action=normalized_action)
unnormalized_transition = unnormalizer(transition) unnormalized_transition = unnormalizer(transition)
unnormalized_action = unnormalized_transition[TransitionIndex.ACTION] unnormalized_action = unnormalized_transition[TransitionKey.ACTION]
assert isinstance(unnormalized_action, torch.Tensor) assert isinstance(unnormalized_action, torch.Tensor)
expected = torch.tensor([1.0, -1.0, 1.0]) expected = torch.tensor([1.0, -1.0, 1.0])
@@ -332,7 +362,7 @@ def test_none_action(action_stats_mean_std):
features=features, norm_map=norm_map, stats={"action": action_stats_mean_std} features=features, norm_map=norm_map, stats={"action": action_stats_mean_std}
) )
transition = (None, None, None, None, None, None, None) transition = create_transition()
result = unnormalizer(transition) result = unnormalizer(transition)
# Should return transition unchanged # Should return transition unchanged
@@ -396,23 +426,31 @@ def test_combined_normalization(normalizer_processor):
"observation.state": torch.tensor([0.5, 0.0]), "observation.state": torch.tensor([0.5, 0.0]),
} }
action = torch.tensor([1.0, -0.5]) action = torch.tensor([1.0, -0.5])
transition = (observation, action, 1.0, False, False, {}, {}) transition = create_transition(
observation=observation,
action=action,
reward=1.0,
done=False,
truncated=False,
info={},
complementary_data={},
)
processed_transition = normalizer_processor(transition) processed_transition = normalizer_processor(transition)
# Check normalized observations # Check normalized observations
processed_obs = processed_transition[TransitionIndex.OBSERVATION] processed_obs = processed_transition[TransitionKey.OBSERVATION]
expected_image = (torch.tensor([0.7, 0.5, 0.3]) - 0.5) / 0.2 expected_image = (torch.tensor([0.7, 0.5, 0.3]) - 0.5) / 0.2
assert torch.allclose(processed_obs["observation.image"], expected_image) assert torch.allclose(processed_obs["observation.image"], expected_image)
# Check normalized action # Check normalized action
processed_action = processed_transition[TransitionIndex.ACTION] processed_action = processed_transition[TransitionKey.ACTION]
expected_action = torch.tensor([(1.0 - 0.0) / 1.0, (-0.5 - 0.0) / 2.0]) expected_action = torch.tensor([(1.0 - 0.0) / 1.0, (-0.5 - 0.0) / 2.0])
assert torch.allclose(processed_action, expected_action) assert torch.allclose(processed_action, expected_action)
# Check other fields remain unchanged # Check other fields remain unchanged
assert processed_transition[TransitionIndex.REWARD] == 1.0 assert processed_transition[TransitionKey.REWARD] == 1.0
assert not processed_transition[TransitionIndex.DONE] assert not processed_transition[TransitionKey.DONE]
def test_processor_from_lerobot_dataset(full_stats): def test_processor_from_lerobot_dataset(full_stats):
@@ -466,13 +504,21 @@ def test_integration_with_robot_processor(normalizer_processor):
"observation.state": torch.tensor([0.5, 0.0]), "observation.state": torch.tensor([0.5, 0.0]),
} }
action = torch.tensor([1.0, -0.5]) action = torch.tensor([1.0, -0.5])
transition = (observation, action, 1.0, False, False, {}, {}) transition = create_transition(
observation=observation,
action=action,
reward=1.0,
done=False,
truncated=False,
info={},
complementary_data={},
)
processed_transition = robot_processor(transition) processed_transition = robot_processor(transition)
# Verify the processing worked # Verify the processing worked
assert isinstance(processed_transition[TransitionIndex.OBSERVATION], dict) assert isinstance(processed_transition[TransitionKey.OBSERVATION], dict)
assert isinstance(processed_transition[TransitionIndex.ACTION], torch.Tensor) assert isinstance(processed_transition[TransitionKey.ACTION], torch.Tensor)
# Edge case tests # Edge case tests
@@ -482,7 +528,7 @@ def test_empty_observation():
norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
transition = (None, None, None, None, None, None, None) transition = create_transition()
result = normalizer(transition) result = normalizer(transition)
assert result == transition assert result == transition
@@ -493,11 +539,13 @@ def test_empty_stats():
norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats={}) normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats={})
observation = {"observation.image": torch.tensor([0.5])} observation = {"observation.image": torch.tensor([0.5])}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = normalizer(transition) result = normalizer(transition)
# Should return observation unchanged since no stats are available # Should return observation unchanged since no stats are available
assert torch.allclose(result[0]["observation.image"], observation["observation.image"]) assert torch.allclose(
result[TransitionKey.OBSERVATION]["observation.image"], observation["observation.image"]
)
def test_partial_stats(): def test_partial_stats():
@@ -507,9 +555,9 @@ def test_partial_stats():
norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD} norm_map = {FeatureType.VISUAL: NormalizationMode.MEAN_STD}
normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats) normalizer = NormalizerProcessor(features=features, norm_map=norm_map, stats=stats)
observation = {"observation.image": torch.tensor([0.7])} observation = {"observation.image": torch.tensor([0.7])}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
processed = normalizer(transition)[TransitionIndex.OBSERVATION] processed = normalizer(transition)[TransitionKey.OBSERVATION]
assert torch.allclose(processed["observation.image"], observation["observation.image"]) assert torch.allclose(processed["observation.image"], observation["observation.image"])
@@ -551,14 +599,25 @@ def test_serialization_roundtrip(full_stats):
"observation.state": torch.tensor([0.5, 0.0]), "observation.state": torch.tensor([0.5, 0.0]),
} }
action = torch.tensor([1.0, -0.5]) action = torch.tensor([1.0, -0.5])
transition = (observation, action, 1.0, False, False, {}, {}) transition = create_transition(
observation=observation,
action=action,
reward=1.0,
done=False,
truncated=False,
info={},
complementary_data={},
)
result1 = original_processor(transition) result1 = original_processor(transition)
result2 = new_processor(transition) result2 = new_processor(transition)
# Compare results # Compare results
assert torch.allclose(result1[0]["observation.image"], result2[0]["observation.image"]) assert torch.allclose(
assert torch.allclose(result1[1], result2[1]) result1[TransitionKey.OBSERVATION]["observation.image"],
result2[TransitionKey.OBSERVATION]["observation.image"],
)
assert torch.allclose(result1[TransitionKey.ACTION], result2[TransitionKey.ACTION])
# Verify features and norm_map are correctly reconstructed # Verify features and norm_map are correctly reconstructed
assert new_processor.features.keys() == original_processor.features.keys() assert new_processor.features.keys() == original_processor.features.keys()
tests/processor/test_observation_processor.py
+49 -33
View File
@@ -23,6 +23,22 @@ from lerobot.processor import (
StateProcessor, StateProcessor,
VanillaObservationProcessor, VanillaObservationProcessor,
) )
from lerobot.processor.pipeline import TransitionKey
def create_transition(
observation=None, action=None, reward=None, done=None, truncated=None, info=None, complementary_data=None
):
"""Helper to create an EnvTransition dictionary."""
return {
TransitionKey.OBSERVATION: observation,
TransitionKey.ACTION: action,
TransitionKey.REWARD: reward,
TransitionKey.DONE: done,
TransitionKey.TRUNCATED: truncated,
TransitionKey.INFO: info,
TransitionKey.COMPLEMENTARY_DATA: complementary_data,
}
def test_process_single_image(): def test_process_single_image():
@@ -33,10 +49,10 @@ def test_process_single_image():
image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8) image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8)
observation = {"pixels": image} observation = {"pixels": image}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[0] processed_obs = result[TransitionKey.OBSERVATION]
# Check that the image was processed correctly # Check that the image was processed correctly
assert "observation.image" in processed_obs assert "observation.image" in processed_obs
@@ -60,10 +76,10 @@ def test_process_image_dict():
image2 = np.random.randint(0, 256, size=(48, 48, 3), dtype=np.uint8) image2 = np.random.randint(0, 256, size=(48, 48, 3), dtype=np.uint8)
observation = {"pixels": {"camera1": image1, "camera2": image2}} observation = {"pixels": {"camera1": image1, "camera2": image2}}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[0] processed_obs = result[TransitionKey.OBSERVATION]
# Check that both images were processed # Check that both images were processed
assert "observation.images.camera1" in processed_obs assert "observation.images.camera1" in processed_obs
@@ -82,10 +98,10 @@ def test_process_batched_image():
image = np.random.randint(0, 256, size=(2, 64, 64, 3), dtype=np.uint8) image = np.random.randint(0, 256, size=(2, 64, 64, 3), dtype=np.uint8)
observation = {"pixels": image} observation = {"pixels": image}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[0] processed_obs = result[TransitionKey.OBSERVATION]
# Check that batch dimension is preserved # Check that batch dimension is preserved
assert processed_obs["observation.image"].shape == (2, 3, 64, 64) assert processed_obs["observation.image"].shape == (2, 3, 64, 64)
@@ -98,7 +114,7 @@ def test_invalid_image_format():
# Test wrong channel order (channels first) # Test wrong channel order (channels first)
image = np.random.randint(0, 256, size=(3, 64, 64), dtype=np.uint8) image = np.random.randint(0, 256, size=(3, 64, 64), dtype=np.uint8)
observation = {"pixels": image} observation = {"pixels": image}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
with pytest.raises(ValueError, match="Expected channel-last images"): with pytest.raises(ValueError, match="Expected channel-last images"):
processor(transition) processor(transition)
@@ -111,7 +127,7 @@ def test_invalid_image_dtype():
# Test wrong dtype # Test wrong dtype
image = np.random.rand(64, 64, 3).astype(np.float32) image = np.random.rand(64, 64, 3).astype(np.float32)
observation = {"pixels": image} observation = {"pixels": image}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
with pytest.raises(ValueError, match="Expected torch.uint8 images"): with pytest.raises(ValueError, match="Expected torch.uint8 images"):
processor(transition) processor(transition)
@@ -122,10 +138,10 @@ def test_no_pixels_in_observation():
processor = ImageProcessor() processor = ImageProcessor()
observation = {"other_data": np.array([1, 2, 3])} observation = {"other_data": np.array([1, 2, 3])}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[0] processed_obs = result[TransitionKey.OBSERVATION]
# Should preserve other data unchanged # Should preserve other data unchanged
assert "other_data" in processed_obs assert "other_data" in processed_obs
@@ -136,7 +152,7 @@ def test_none_observation():
"""Test processor with None observation.""" """Test processor with None observation."""
processor = ImageProcessor() processor = ImageProcessor()
transition = (None, None, None, None, None, None, None) transition = create_transition()
result = processor(transition) result = processor(transition)
assert result == transition assert result == transition
@@ -167,10 +183,10 @@ def test_process_environment_state():
env_state = np.array([1.0, 2.0, 3.0], dtype=np.float32) env_state = np.array([1.0, 2.0, 3.0], dtype=np.float32)
observation = {"environment_state": env_state} observation = {"environment_state": env_state}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[0] processed_obs = result[TransitionKey.OBSERVATION]
# Check that environment_state was renamed and processed # Check that environment_state was renamed and processed
assert "observation.environment_state" in processed_obs assert "observation.environment_state" in processed_obs
@@ -188,10 +204,10 @@ def test_process_agent_pos():
agent_pos = np.array([0.5, -0.5, 1.0], dtype=np.float32) agent_pos = np.array([0.5, -0.5, 1.0], dtype=np.float32)
observation = {"agent_pos": agent_pos} observation = {"agent_pos": agent_pos}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[0] processed_obs = result[TransitionKey.OBSERVATION]
# Check that agent_pos was renamed and processed # Check that agent_pos was renamed and processed
assert "observation.state" in processed_obs assert "observation.state" in processed_obs
@@ -211,10 +227,10 @@ def test_process_batched_states():
agent_pos = np.array([[0.5, -0.5], [1.0, -1.0]], dtype=np.float32) agent_pos = np.array([[0.5, -0.5], [1.0, -1.0]], dtype=np.float32)
observation = {"environment_state": env_state, "agent_pos": agent_pos} observation = {"environment_state": env_state, "agent_pos": agent_pos}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[0] processed_obs = result[TransitionKey.OBSERVATION]
# Check that batch dimensions are preserved # Check that batch dimensions are preserved
assert processed_obs["observation.environment_state"].shape == (2, 2) assert processed_obs["observation.environment_state"].shape == (2, 2)
@@ -229,10 +245,10 @@ def test_process_both_states():
agent_pos = np.array([0.5, -0.5], dtype=np.float32) agent_pos = np.array([0.5, -0.5], dtype=np.float32)
observation = {"environment_state": env_state, "agent_pos": agent_pos, "other_data": "keep_me"} observation = {"environment_state": env_state, "agent_pos": agent_pos, "other_data": "keep_me"}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[0] processed_obs = result[TransitionKey.OBSERVATION]
# Check that both states were processed # Check that both states were processed
assert "observation.environment_state" in processed_obs assert "observation.environment_state" in processed_obs
@@ -251,10 +267,10 @@ def test_no_states_in_observation():
processor = StateProcessor() processor = StateProcessor()
observation = {"other_data": np.array([1, 2, 3])} observation = {"other_data": np.array([1, 2, 3])}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[0] processed_obs = result[TransitionKey.OBSERVATION]
# Should preserve data unchanged # Should preserve data unchanged
np.testing.assert_array_equal(processed_obs, observation) np.testing.assert_array_equal(processed_obs, observation)
@@ -275,10 +291,10 @@ def test_complete_observation_processing():
"agent_pos": agent_pos, "agent_pos": agent_pos,
"other_data": "preserve_me", "other_data": "preserve_me",
} }
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[0] processed_obs = result[TransitionKey.OBSERVATION]
# Check that image was processed # Check that image was processed
assert "observation.image" in processed_obs assert "observation.image" in processed_obs
@@ -303,10 +319,10 @@ def test_image_only_processing():
image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8) image = np.random.randint(0, 256, size=(64, 64, 3), dtype=np.uint8)
observation = {"pixels": image} observation = {"pixels": image}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[0] processed_obs = result[TransitionKey.OBSERVATION]
assert "observation.image" in processed_obs assert "observation.image" in processed_obs
assert len(processed_obs) == 1 assert len(processed_obs) == 1
@@ -318,10 +334,10 @@ def test_state_only_processing():
agent_pos = np.array([1.0, 2.0], dtype=np.float32) agent_pos = np.array([1.0, 2.0], dtype=np.float32)
observation = {"agent_pos": agent_pos} observation = {"agent_pos": agent_pos}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[0] processed_obs = result[TransitionKey.OBSERVATION]
assert "observation.state" in processed_obs assert "observation.state" in processed_obs
assert "agent_pos" not in processed_obs assert "agent_pos" not in processed_obs
@@ -332,10 +348,10 @@ def test_empty_observation():
processor = VanillaObservationProcessor() processor = VanillaObservationProcessor()
observation = {} observation = {}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[0] processed_obs = result[TransitionKey.OBSERVATION]
assert processed_obs == {} assert processed_obs == {}
@@ -369,8 +385,8 @@ def test_equivalent_to_original_function():
original_result = preprocess_observation(observation) original_result = preprocess_observation(observation)
# Process with new processor # Process with new processor
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
processor_result = processor(transition)[0] processor_result = processor(transition)[TransitionKey.OBSERVATION]
# Compare results # Compare results
assert set(original_result.keys()) == set(processor_result.keys()) assert set(original_result.keys()) == set(processor_result.keys())
@@ -396,8 +412,8 @@ def test_equivalent_with_image_dict():
original_result = preprocess_observation(observation) original_result = preprocess_observation(observation)
# Process with new processor # Process with new processor
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
processor_result = processor(transition)[0] processor_result = processor(transition)[TransitionKey.OBSERVATION]
# Compare results # Compare results
assert set(original_result.keys()) == set(processor_result.keys()) assert set(original_result.keys()) == set(processor_result.keys())
tests/processor/test_pipeline.py
+107 -73
View File
@@ -18,7 +18,7 @@ import json
import tempfile import tempfile
from dataclasses import dataclass from dataclasses import dataclass
from pathlib import Path from pathlib import Path
from typing import Any, Dict from typing import Any
import numpy as np import numpy as np
import pytest import pytest
@@ -26,6 +26,22 @@ import torch
import torch.nn as nn import torch.nn as nn
from lerobot.processor import EnvTransition, ProcessorStepRegistry, RobotProcessor from lerobot.processor import EnvTransition, ProcessorStepRegistry, RobotProcessor
from lerobot.processor.pipeline import TransitionKey
def create_transition(
observation=None, action=None, reward=0.0, done=False, truncated=False, info=None, complementary_data=None
):
"""Helper to create an EnvTransition dictionary."""
return {
TransitionKey.OBSERVATION: observation,
TransitionKey.ACTION: action,
TransitionKey.REWARD: reward,
TransitionKey.DONE: done,
TransitionKey.TRUNCATED: truncated,
TransitionKey.INFO: info if info is not None else {},
TransitionKey.COMPLEMENTARY_DATA: complementary_data if complementary_data is not None else {},
}
@dataclass @dataclass
@@ -45,14 +61,16 @@ class MockStep:
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
"""Add a counter to the complementary_data.""" """Add a counter to the complementary_data."""
obs, action, reward, done, truncated, info, comp_data = transition comp_data = transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
comp_data = {} if comp_data is None else dict(comp_data) # Make a copy comp_data = {} if comp_data is None else dict(comp_data) # Make a copy
comp_data[f"{self.name}_counter"] = self.counter comp_data[f"{self.name}_counter"] = self.counter
self.counter += 1 self.counter += 1
return (obs, action, reward, done, truncated, info, comp_data) # Create a new transition with updated complementary_data
new_transition = transition.copy()
new_transition[TransitionKey.COMPLEMENTARY_DATA] = comp_data
return new_transition
def get_config(self) -> dict[str, Any]: def get_config(self) -> dict[str, Any]:
# Return all JSON-serializable attributes that should be persisted # Return all JSON-serializable attributes that should be persisted
@@ -79,12 +97,14 @@ class MockStepWithoutOptionalMethods:
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
"""Multiply reward by multiplier.""" """Multiply reward by multiplier."""
obs, action, reward, done, truncated, info, comp_data = transition reward = transition.get(TransitionKey.REWARD)
if reward is not None: if reward is not None:
reward = reward * self.multiplier new_transition = transition.copy()
new_transition[TransitionKey.REWARD] = reward * self.multiplier
return new_transition
return (obs, action, reward, done, truncated, info, comp_data) return transition
@dataclass @dataclass
@@ -105,7 +125,7 @@ class MockStepWithTensorState:
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
"""Update running statistics.""" """Update running statistics."""
obs, action, reward, done, truncated, info, comp_data = transition reward = transition.get(TransitionKey.REWARD)
if reward is not None: if reward is not None:
# Update running mean # Update running mean
@@ -143,7 +163,7 @@ def test_empty_pipeline():
"""Test pipeline with no steps.""" """Test pipeline with no steps."""
pipeline = RobotProcessor() pipeline = RobotProcessor()
transition = (None, None, 0.0, False, False, {}, {}) transition = create_transition()
result = pipeline(transition) result = pipeline(transition)
assert result == transition assert result == transition
@@ -155,15 +175,15 @@ def test_single_step_pipeline():
step = MockStep("test_step") step = MockStep("test_step")
pipeline = RobotProcessor([step]) pipeline = RobotProcessor([step])
transition = (None, None, 0.0, False, False, {}, {}) transition = create_transition()
result = pipeline(transition) result = pipeline(transition)
assert len(pipeline) == 1 assert len(pipeline) == 1
assert result[6]["test_step_counter"] == 0 # complementary_data assert result[TransitionKey.COMPLEMENTARY_DATA]["test_step_counter"] == 0
# Call again to test counter increment # Call again to test counter increment
result = pipeline(transition) result = pipeline(transition)
assert result[6]["test_step_counter"] == 1 assert result[TransitionKey.COMPLEMENTARY_DATA]["test_step_counter"] == 1
def test_multiple_steps_pipeline(): def test_multiple_steps_pipeline():
@@ -172,46 +192,46 @@ def test_multiple_steps_pipeline():
step2 = MockStep("step2") step2 = MockStep("step2")
pipeline = RobotProcessor([step1, step2]) pipeline = RobotProcessor([step1, step2])
transition = (None, None, 0.0, False, False, {}, {}) transition = create_transition()
result = pipeline(transition) result = pipeline(transition)
assert len(pipeline) == 2 assert len(pipeline) == 2
assert result[6]["step1_counter"] == 0 assert result[TransitionKey.COMPLEMENTARY_DATA]["step1_counter"] == 0
assert result[6]["step2_counter"] == 0 assert result[TransitionKey.COMPLEMENTARY_DATA]["step2_counter"] == 0
def test_invalid_transition_format(): def test_invalid_transition_format():
"""Test pipeline with invalid transition format.""" """Test pipeline with invalid transition format."""
pipeline = RobotProcessor([MockStep()]) pipeline = RobotProcessor([MockStep()])
# Test with wrong number of elements # Test with wrong type (tuple instead of dict)
with pytest.raises(ValueError, match="EnvTransition must be a 7-tuple"): with pytest.raises(ValueError, match="EnvTransition must be a dictionary"):
pipeline((None, None, 0.0)) # Only 3 elements pipeline((None, None, 0.0, False, False, {}, {})) # Tuple instead of dict
# Test with wrong type # Test with wrong type (string)
with pytest.raises(ValueError, match="EnvTransition must be a 7-tuple"): with pytest.raises(ValueError, match="EnvTransition must be a dictionary"):
pipeline("not a tuple") pipeline("not a dict")
def test_step_through(): def test_step_through():
"""Test step_through method with tuple input.""" """Test step_through method with dict input."""
step1 = MockStep("step1") step1 = MockStep("step1")
step2 = MockStep("step2") step2 = MockStep("step2")
pipeline = RobotProcessor([step1, step2]) pipeline = RobotProcessor([step1, step2])
transition = (None, None, 0.0, False, False, {}, {}) transition = create_transition()
results = list(pipeline.step_through(transition)) results = list(pipeline.step_through(transition))
assert len(results) == 3 # Original + 2 steps assert len(results) == 3 # Original + 2 steps
assert results[0] == transition # Original assert results[0] == transition # Original
assert "step1_counter" in results[1][6] # After step1 assert "step1_counter" in results[1][TransitionKey.COMPLEMENTARY_DATA] # After step1
assert "step2_counter" in results[2][6] # After step2 assert "step2_counter" in results[2][TransitionKey.COMPLEMENTARY_DATA] # After step2
# Ensure all results are tuples (same format as input) # Ensure all results are dicts (same format as input)
for result in results: for result in results:
assert isinstance(result, tuple) assert isinstance(result, dict)
assert len(result) == 7 assert all(isinstance(k, TransitionKey) for k in result.keys())
def test_step_through_with_dict(): def test_step_through_with_dict():
@@ -279,7 +299,7 @@ def test_hooks():
pipeline.register_before_step_hook(before_hook) pipeline.register_before_step_hook(before_hook)
pipeline.register_after_step_hook(after_hook) pipeline.register_after_step_hook(after_hook)
transition = (None, None, 0.0, False, False, {}, {}) transition = create_transition()
pipeline(transition) pipeline(transition)
assert before_calls == [0] assert before_calls == [0]
@@ -292,15 +312,16 @@ def test_hook_modification():
pipeline = RobotProcessor([step]) pipeline = RobotProcessor([step])
def modify_reward_hook(idx: int, transition: EnvTransition): def modify_reward_hook(idx: int, transition: EnvTransition):
obs, action, reward, done, truncated, info, comp_data = transition new_transition = transition.copy()
return (obs, action, 42.0, done, truncated, info, comp_data) new_transition[TransitionKey.REWARD] = 42.0
return new_transition
pipeline.register_before_step_hook(modify_reward_hook) pipeline.register_before_step_hook(modify_reward_hook)
transition = (None, None, 0.0, False, False, {}, {}) transition = create_transition()
result = pipeline(transition) result = pipeline(transition)
assert result[2] == 42.0 # reward modified by hook assert result[TransitionKey.REWARD] == 42.0 # reward modified by hook
def test_reset(): def test_reset():
@@ -316,7 +337,7 @@ def test_reset():
pipeline.register_reset_hook(reset_hook) pipeline.register_reset_hook(reset_hook)
# Make some calls to increment counter # Make some calls to increment counter
transition = (None, None, 0.0, False, False, {}, {}) transition = create_transition()
pipeline(transition) pipeline(transition)
pipeline(transition) pipeline(transition)
@@ -335,7 +356,7 @@ def test_profile_steps():
step2 = MockStep("step2") step2 = MockStep("step2")
pipeline = RobotProcessor([step1, step2]) pipeline = RobotProcessor([step1, step2])
transition = (None, None, 0.0, False, False, {}, {}) transition = create_transition()
profile_results = pipeline.profile_steps(transition, num_runs=10) profile_results = pipeline.profile_steps(transition, num_runs=10)
@@ -397,10 +418,10 @@ def test_step_without_optional_methods():
step = MockStepWithoutOptionalMethods(multiplier=3.0) step = MockStepWithoutOptionalMethods(multiplier=3.0)
pipeline = RobotProcessor([step]) pipeline = RobotProcessor([step])
transition = (None, None, 2.0, False, False, {}, {}) transition = create_transition(reward=2.0)
result = pipeline(transition) result = pipeline(transition)
assert result[2] == 6.0 # 2.0 * 3.0 assert result[TransitionKey.REWARD] == 6.0 # 2.0 * 3.0
# Reset should work even if step doesn't implement reset # Reset should work even if step doesn't implement reset
pipeline.reset() pipeline.reset()
@@ -419,7 +440,7 @@ def test_mixed_json_and_tensor_state():
# Process some transitions with rewards # Process some transitions with rewards
for i in range(10): for i in range(10):
transition = (None, None, float(i), False, False, {}, {}) transition = create_transition(reward=float(i))
pipeline(transition) pipeline(transition)
# Check state # Check state
@@ -466,7 +487,7 @@ class MockModuleStep(nn.Module):
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
"""Process transition and update running mean.""" """Process transition and update running mean."""
obs, action, reward, done, truncated, info, comp_data = transition obs = transition.get(TransitionKey.OBSERVATION)
if obs is not None and isinstance(obs, torch.Tensor): if obs is not None and isinstance(obs, torch.Tensor):
# Process observation through linear layer # Process observation through linear layer
@@ -509,7 +530,7 @@ def test_to_device_with_state_dict():
# Process some transitions to populate state # Process some transitions to populate state
for i in range(10): for i in range(10):
transition = (None, None, float(i), False, False, {}, {}) transition = create_transition(reward=float(i))
pipeline(transition) pipeline(transition)
# Check initial device (should be CPU) # Check initial device (should be CPU)
@@ -551,7 +572,7 @@ def test_to_device_with_module():
# Process some data # Process some data
obs = torch.randn(2, 5) obs = torch.randn(2, 5)
transition = (obs, None, 1.0, False, False, {}, {}) transition = create_transition(observation=obs, reward=1.0)
pipeline(transition) pipeline(transition)
# Check initial device # Check initial device
@@ -575,7 +596,7 @@ def test_to_device_with_module():
# Verify the module still works after transfer # Verify the module still works after transfer
obs_cuda = torch.randn(2, 5, device="cuda:0") obs_cuda = torch.randn(2, 5, device="cuda:0")
transition = (obs_cuda, None, 1.0, False, False, {}, {}) transition = create_transition(observation=obs_cuda, reward=1.0)
pipeline(transition) # Should not raise an error pipeline(transition) # Should not raise an error
@@ -589,7 +610,7 @@ def test_to_device_mixed_steps():
# Process some data # Process some data
for i in range(5): for i in range(5):
transition = (torch.randn(2, 10), None, float(i), False, False, {}, {}) transition = create_transition(observation=torch.randn(2, 10), reward=float(i))
pipeline(transition) pipeline(transition)
# Check initial state # Check initial state
@@ -630,7 +651,7 @@ def test_to_device_preserves_functionality():
# Process initial data # Process initial data
rewards = [1.0, 2.0, 3.0] rewards = [1.0, 2.0, 3.0]
for r in rewards: for r in rewards:
transition = (None, None, r, False, False, {}, {}) transition = create_transition(reward=r)
pipeline(transition) pipeline(transition)
# Check state before transfer # Check state before transfer
@@ -645,7 +666,7 @@ def test_to_device_preserves_functionality():
assert step.running_count == initial_count assert step.running_count == initial_count
# Process more data to ensure functionality # Process more data to ensure functionality
transition = (None, None, 4.0, False, False, {}, {}) transition = create_transition(reward=4.0)
_ = pipeline(transition) _ = pipeline(transition)
assert step.running_count == 4 assert step.running_count == 4
@@ -700,7 +721,8 @@ class MockNonModuleStepWithState:
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
"""Process transition using tensor operations.""" """Process transition using tensor operations."""
obs, action, reward, done, truncated, info, comp_data = transition obs = transition.get(TransitionKey.OBSERVATION)
comp_data = transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
if obs is not None and isinstance(obs, torch.Tensor) and obs.numel() >= self.feature_dim: if obs is not None and isinstance(obs, torch.Tensor) and obs.numel() >= self.feature_dim:
# Perform some tensor operations # Perform some tensor operations
@@ -718,7 +740,12 @@ class MockNonModuleStepWithState:
comp_data[f"{self.name}_mean_output"] = output.mean().item() comp_data[f"{self.name}_mean_output"] = output.mean().item()
comp_data[f"{self.name}_steps"] = self.step_count.item() comp_data[f"{self.name}_steps"] = self.step_count.item()
return (obs, action, reward, done, truncated, info, comp_data) # Return updated transition
new_transition = transition.copy()
new_transition[TransitionKey.COMPLEMENTARY_DATA] = comp_data
return new_transition
return transition
def get_config(self) -> dict[str, Any]: def get_config(self) -> dict[str, Any]:
return { return {
@@ -763,9 +790,9 @@ def test_to_device_non_module_class():
# Process some data to populate state # Process some data to populate state
for i in range(3): for i in range(3):
obs = torch.randn(2, 5) obs = torch.randn(2, 5)
transition = (obs, None, float(i), False, False, {}, {}) transition = create_transition(observation=obs, reward=float(i))
result = pipeline(transition) result = pipeline(transition)
comp_data = result[6] comp_data = result[TransitionKey.COMPLEMENTARY_DATA]
assert f"{non_module_step.name}_steps" in comp_data assert f"{non_module_step.name}_steps" in comp_data
# Verify all tensors are on CPU initially # Verify all tensors are on CPU initially
@@ -811,9 +838,9 @@ def test_to_device_non_module_class():
# Test that step still works on GPU # Test that step still works on GPU
obs_gpu = torch.randn(2, 5, device="cuda") obs_gpu = torch.randn(2, 5, device="cuda")
transition = (obs_gpu, None, 1.0, False, False, {}, {}) transition = create_transition(observation=obs_gpu, reward=1.0)
result = pipeline(transition) result = pipeline(transition)
comp_data = result[6] comp_data = result[TransitionKey.COMPLEMENTARY_DATA]
# Verify processing worked # Verify processing worked
assert comp_data[f"{non_module_step.name}_steps"] == 4 assert comp_data[f"{non_module_step.name}_steps"] == 4
@@ -835,7 +862,7 @@ def test_to_device_module_vs_non_module():
# Process some data # Process some data
obs = torch.randn(2, 5) obs = torch.randn(2, 5)
transition = (obs, None, 1.0, False, False, {}, {}) transition = create_transition(observation=obs, reward=1.0)
_ = pipeline(transition) _ = pipeline(transition)
# Check initial devices # Check initial devices
@@ -860,7 +887,7 @@ def test_to_device_module_vs_non_module():
# Process data on GPU # Process data on GPU
obs_gpu = torch.randn(2, 5, device="cuda") obs_gpu = torch.randn(2, 5, device="cuda")
transition = (obs_gpu, None, 2.0, False, False, {}, {}) transition = create_transition(observation=obs_gpu, reward=2.0)
_ = pipeline(transition) _ = pipeline(transition)
# Verify both steps processed the data # Verify both steps processed the data
@@ -889,7 +916,8 @@ class MockStepWithNonSerializableParam:
self.env = env # Non-serializable parameter (like gym.Env) self.env = env # Non-serializable parameter (like gym.Env)
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
obs, action, reward, done, truncated, info, comp_data = transition reward = transition.get(TransitionKey.REWARD)
comp_data = transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
# Use the env parameter if provided # Use the env parameter if provided
if self.env is not None: if self.env is not None:
@@ -897,10 +925,14 @@ class MockStepWithNonSerializableParam:
comp_data[f"{self.name}_env_info"] = str(self.env) comp_data[f"{self.name}_env_info"] = str(self.env)
# Apply multiplier to reward # Apply multiplier to reward
new_transition = transition.copy()
if reward is not None: if reward is not None:
reward = reward * self.multiplier new_transition[TransitionKey.REWARD] = reward * self.multiplier
return (obs, action, reward, done, truncated, info, comp_data) if comp_data:
new_transition[TransitionKey.COMPLEMENTARY_DATA] = comp_data
return new_transition
def get_config(self) -> dict[str, Any]: def get_config(self) -> dict[str, Any]:
# Note: env is intentionally NOT included here as it's not serializable # Note: env is intentionally NOT included here as it's not serializable
@@ -928,13 +960,15 @@ class RegisteredMockStep:
device: str = "cpu" device: str = "cpu"
def __call__(self, transition: EnvTransition) -> EnvTransition: def __call__(self, transition: EnvTransition) -> EnvTransition:
obs, action, reward, done, truncated, info, comp_data = transition comp_data = transition.get(TransitionKey.COMPLEMENTARY_DATA, {})
comp_data = {} if comp_data is None else dict(comp_data) comp_data = {} if comp_data is None else dict(comp_data)
comp_data["registered_step_value"] = self.value comp_data["registered_step_value"] = self.value
comp_data["registered_step_device"] = self.device comp_data["registered_step_device"] = self.device
return (obs, action, reward, done, truncated, info, comp_data) new_transition = transition.copy()
new_transition[TransitionKey.COMPLEMENTARY_DATA] = comp_data
return new_transition
def get_config(self) -> dict[str, Any]: def get_config(self) -> dict[str, Any]:
return { return {
@@ -993,18 +1027,18 @@ def test_from_pretrained_with_overrides():
assert loaded_pipeline.name == "TestOverrides" assert loaded_pipeline.name == "TestOverrides"
# Test the loaded steps # Test the loaded steps
transition = (None, None, 1.0, False, False, {}, {}) transition = create_transition(reward=1.0)
result = loaded_pipeline(transition) result = loaded_pipeline(transition)
# Check that overrides were applied # Check that overrides were applied
comp_data = result[6] comp_data = result[TransitionKey.COMPLEMENTARY_DATA]
assert "env_step_env_info" in comp_data assert "env_step_env_info" in comp_data
assert comp_data["env_step_env_info"] == "MockEnvironment(test_env)" assert comp_data["env_step_env_info"] == "MockEnvironment(test_env)"
assert comp_data["registered_step_value"] == 200 assert comp_data["registered_step_value"] == 200
assert comp_data["registered_step_device"] == "cuda" assert comp_data["registered_step_device"] == "cuda"
# Check that multiplier override was applied # Check that multiplier override was applied
assert result[2] == 3.0 # 1.0 * 3.0 (overridden multiplier) assert result[TransitionKey.REWARD] == 3.0 # 1.0 * 3.0 (overridden multiplier)
def test_from_pretrained_with_partial_overrides(): def test_from_pretrained_with_partial_overrides():
@@ -1024,13 +1058,13 @@ def test_from_pretrained_with_partial_overrides():
# Both steps will get the override # Both steps will get the override
loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides)
transition = (None, None, 1.0, False, False, {}, {}) transition = create_transition(reward=1.0)
result = loaded_pipeline(transition) result = loaded_pipeline(transition)
# The reward should be affected by both steps, both getting the override # The reward should be affected by both steps, both getting the override
# First step: 1.0 * 5.0 = 5.0 (overridden) # First step: 1.0 * 5.0 = 5.0 (overridden)
# Second step: 5.0 * 5.0 = 25.0 (also overridden) # Second step: 5.0 * 5.0 = 25.0 (also overridden)
assert result[2] == 25.0 assert result[TransitionKey.REWARD] == 25.0
def test_from_pretrained_invalid_override_key(): def test_from_pretrained_invalid_override_key():
@@ -1082,10 +1116,10 @@ def test_from_pretrained_registered_step_override():
loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides)
# Test that overrides were applied # Test that overrides were applied
transition = (None, None, 0.0, False, False, {}, {}) transition = create_transition()
result = loaded_pipeline(transition) result = loaded_pipeline(transition)
comp_data = result[6] comp_data = result[TransitionKey.COMPLEMENTARY_DATA]
assert comp_data["registered_step_value"] == 999 assert comp_data["registered_step_value"] == 999
assert comp_data["registered_step_device"] == "cuda" assert comp_data["registered_step_device"] == "cuda"
@@ -1110,13 +1144,13 @@ def test_from_pretrained_mixed_registered_and_unregistered():
loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides) loaded_pipeline = RobotProcessor.from_pretrained(tmp_dir, overrides=overrides)
# Test both steps # Test both steps
transition = (None, None, 2.0, False, False, {}, {}) transition = create_transition(reward=2.0)
result = loaded_pipeline(transition) result = loaded_pipeline(transition)
comp_data = result[6] comp_data = result[TransitionKey.COMPLEMENTARY_DATA]
assert comp_data["unregistered_env_info"] == "MockEnvironment(mixed_test)" assert comp_data["unregistered_env_info"] == "MockEnvironment(mixed_test)"
assert comp_data["registered_step_value"] == 777 assert comp_data["registered_step_value"] == 777
assert result[2] == 8.0 # 2.0 * 4.0 assert result[TransitionKey.REWARD] == 8.0 # 2.0 * 4.0
def test_from_pretrained_no_overrides(): def test_from_pretrained_no_overrides():
@@ -1133,10 +1167,10 @@ def test_from_pretrained_no_overrides():
assert len(loaded_pipeline) == 1 assert len(loaded_pipeline) == 1
# Test that the step works (env will be None) # Test that the step works (env will be None)
transition = (None, None, 1.0, False, False, {}, {}) transition = create_transition(reward=1.0)
result = loaded_pipeline(transition) result = loaded_pipeline(transition)
assert result[2] == 3.0 # 1.0 * 3.0 assert result[TransitionKey.REWARD] == 3.0 # 1.0 * 3.0
def test_from_pretrained_empty_overrides(): def test_from_pretrained_empty_overrides():
@@ -1153,10 +1187,10 @@ def test_from_pretrained_empty_overrides():
assert len(loaded_pipeline) == 1 assert len(loaded_pipeline) == 1
# Test that the step works normally # Test that the step works normally
transition = (None, None, 1.0, False, False, {}, {}) transition = create_transition(reward=1.0)
result = loaded_pipeline(transition) result = loaded_pipeline(transition)
assert result[2] == 2.0 assert result[TransitionKey.REWARD] == 2.0
def test_from_pretrained_override_instantiation_error(): def test_from_pretrained_override_instantiation_error():
@@ -1185,7 +1219,7 @@ def test_from_pretrained_with_state_and_overrides():
# Process some data to create state # Process some data to create state
for i in range(10): for i in range(10):
transition = (None, None, float(i), False, False, {}, {}) transition = create_transition(reward=float(i))
pipeline(transition) pipeline(transition)
with tempfile.TemporaryDirectory() as tmp_dir: with tempfile.TemporaryDirectory() as tmp_dir:
tests/processor/test_rename_processor.py
+41 -25
View File
@@ -13,14 +13,28 @@
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
import tempfile import tempfile
from pathlib import Path from pathlib import Path
import numpy as np import numpy as np
import torch import torch
from lerobot.processor import ProcessorStepRegistry, RenameProcessor, RobotProcessor, TransitionIndex from lerobot.processor import ProcessorStepRegistry, RenameProcessor, RobotProcessor, TransitionKey
def create_transition(
observation=None, action=None, reward=None, done=None, truncated=None, info=None, complementary_data=None
):
"""Helper to create an EnvTransition dictionary."""
return {
TransitionKey.OBSERVATION: observation,
TransitionKey.ACTION: action,
TransitionKey.REWARD: reward,
TransitionKey.DONE: done,
TransitionKey.TRUNCATED: truncated,
TransitionKey.INFO: info,
TransitionKey.COMPLEMENTARY_DATA: complementary_data,
}
def test_basic_renaming(): def test_basic_renaming():
@@ -36,10 +50,10 @@ def test_basic_renaming():
"old_key2": np.array([3.0, 4.0]), "old_key2": np.array([3.0, 4.0]),
"unchanged_key": "keep_me", "unchanged_key": "keep_me",
} }
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[TransitionIndex.OBSERVATION] processed_obs = result[TransitionKey.OBSERVATION]
# Check renamed keys # Check renamed keys
assert "new_key1" in processed_obs assert "new_key1" in processed_obs
@@ -63,10 +77,10 @@ def test_empty_rename_map():
"key1": torch.tensor([1.0]), "key1": torch.tensor([1.0]),
"key2": "value2", "key2": "value2",
} }
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[TransitionIndex.OBSERVATION] processed_obs = result[TransitionKey.OBSERVATION]
# All keys should be unchanged # All keys should be unchanged
assert processed_obs.keys() == observation.keys() assert processed_obs.keys() == observation.keys()
@@ -78,7 +92,7 @@ def test_none_observation():
"""Test processor with None observation.""" """Test processor with None observation."""
processor = RenameProcessor(rename_map={"old": "new"}) processor = RenameProcessor(rename_map={"old": "new"})
transition = (None, None, None, None, None, None, None) transition = create_transition()
result = processor(transition) result = processor(transition)
# Should return transition unchanged # Should return transition unchanged
@@ -98,10 +112,10 @@ def test_overlapping_rename():
"b": 2, "b": 2,
"x": 3, "x": 3,
} }
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[TransitionIndex.OBSERVATION] processed_obs = result[TransitionKey.OBSERVATION]
# Check that renaming happens correctly # Check that renaming happens correctly
assert "a" not in processed_obs assert "a" not in processed_obs
@@ -124,10 +138,10 @@ def test_partial_rename():
"reward": 1.0, "reward": 1.0,
"info": {"episode": 1}, "info": {"episode": 1},
} }
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[TransitionIndex.OBSERVATION] processed_obs = result[TransitionKey.OBSERVATION]
# Check renamed keys # Check renamed keys
assert "observation.proprio_state" in processed_obs assert "observation.proprio_state" in processed_obs
@@ -178,10 +192,12 @@ def test_integration_with_robot_processor():
"pixels": np.zeros((32, 32, 3), dtype=np.uint8), "pixels": np.zeros((32, 32, 3), dtype=np.uint8),
"other_data": "preserve_me", "other_data": "preserve_me",
} }
transition = (observation, None, 0.5, False, False, {}, {}) transition = create_transition(
observation=observation, reward=0.5, done=False, truncated=False, info={}, complementary_data={}
)
result = pipeline(transition) result = pipeline(transition)
processed_obs = result[TransitionIndex.OBSERVATION] processed_obs = result[TransitionKey.OBSERVATION]
# Check renaming worked through pipeline # Check renaming worked through pipeline
assert "observation.state" in processed_obs assert "observation.state" in processed_obs
@@ -191,8 +207,8 @@ def test_integration_with_robot_processor():
assert processed_obs["other_data"] == "preserve_me" assert processed_obs["other_data"] == "preserve_me"
# Check other transition elements unchanged # Check other transition elements unchanged
assert result[TransitionIndex.REWARD] == 0.5 assert result[TransitionKey.REWARD] == 0.5
assert result[TransitionIndex.DONE] is False assert result[TransitionKey.DONE] is False
def test_save_and_load_pretrained(): def test_save_and_load_pretrained():
@@ -229,10 +245,10 @@ def test_save_and_load_pretrained():
# Test functionality after loading # Test functionality after loading
observation = {"old_state": [1, 2, 3], "old_image": "image_data"} observation = {"old_state": [1, 2, 3], "old_image": "image_data"}
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = loaded_pipeline(transition) result = loaded_pipeline(transition)
processed_obs = result[TransitionIndex.OBSERVATION] processed_obs = result[TransitionKey.OBSERVATION]
assert "observation.state" in processed_obs assert "observation.state" in processed_obs
assert "observation.image" in processed_obs assert "observation.image" in processed_obs
@@ -306,17 +322,17 @@ def test_chained_rename_processors():
"img": "image_data", "img": "image_data",
"extra": "keep_me", "extra": "keep_me",
} }
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
# Step through to see intermediate results # Step through to see intermediate results
results = list(pipeline.step_through(transition)) results = list(pipeline.step_through(transition))
# After first processor # After first processor
assert "agent_position" in results[1][TransitionIndex.OBSERVATION] assert "agent_position" in results[1][TransitionKey.OBSERVATION]
assert "camera_image" in results[1][TransitionIndex.OBSERVATION] assert "camera_image" in results[1][TransitionKey.OBSERVATION]
# After second processor # After second processor
final_obs = results[2][TransitionIndex.OBSERVATION] final_obs = results[2][TransitionKey.OBSERVATION]
assert "observation.state" in final_obs assert "observation.state" in final_obs
assert "observation.image" in final_obs assert "observation.image" in final_obs
assert final_obs["extra"] == "keep_me" assert final_obs["extra"] == "keep_me"
@@ -343,10 +359,10 @@ def test_nested_observation_rename():
"observation.proprio": torch.randn(7), "observation.proprio": torch.randn(7),
"observation.gripper": torch.tensor([0.0]), # Not renamed "observation.gripper": torch.tensor([0.0]), # Not renamed
} }
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[TransitionIndex.OBSERVATION] processed_obs = result[TransitionKey.OBSERVATION]
# Check renames # Check renames
assert "observation.camera.left_view" in processed_obs assert "observation.camera.left_view" in processed_obs
@@ -378,10 +394,10 @@ def test_value_types_preserved():
"old_dict": {"nested": "value"}, "old_dict": {"nested": "value"},
"old_list": [1, 2, 3], "old_list": [1, 2, 3],
} }
transition = (observation, None, None, None, None, None, None) transition = create_transition(observation=observation)
result = processor(transition) result = processor(transition)
processed_obs = result[TransitionIndex.OBSERVATION] processed_obs = result[TransitionKey.OBSERVATION]
# Check that values and types are preserved # Check that values and types are preserved
assert torch.equal(processed_obs["new_tensor"], tensor_value) assert torch.equal(processed_obs["new_tensor"], tensor_value)