mirror of
https://github.com/huggingface/lerobot.git
synced 2026-05-15 16:49:55 +00:00
919184d6f8
* docs(benchmarks): add benchmark integration guide and standardize benchmark docs Add a comprehensive guide for adding new benchmarks to LeRobot, and refactor the existing LIBERO and Meta-World docs to follow the new standardized template. Made-with: Cursor * refactor(envs): move dispatch logic from factory into EnvConfig subclasses Replace hardcoded if/elif chains in factory.py with create_envs() and get_env_processors() methods on EnvConfig. New benchmarks now only need to register a config subclass — no factory.py edits required. Net -23 lines: factory.py shrinks from ~200 to ~70 lines of logic. Made-with: Cursor * docs(benchmarks): clean up adding-benchmarks guide for clarity Rewrite for simpler language, better structure, and easier navigation. Move quick-reference table to the top, fold eval explanation into architecture section, condense the doc template to a bulleted outline. Made-with: Cursor * fix link * fix task count * fix: enable SmolVLA eval on LIBERO with custom camera mappings - Thread camera_name_mapping from LiberoEnv config through to gym envs - Sync features_map with camera_name_mapping in LiberoEnv.__post_init__ - Fix render() to use first available camera instead of hardcoded "image" - Handle non-dict final_info in rollout by falling back to info["is_success"] - Add use_peft legacy field to SmolVLAConfig for checkpoint compat - Add defaults to GR00TN15Config init=False fields for transformers 5.3 Made-with: Cursor * fix: use direct AutoresetMode import for gymnasium compat Made-with: Cursor * fix: handle gymnasium < 1.0 without AutoresetMode Made-with: Cursor * refactor: revert policy changes, keep env-only camera mapping fixes - Revert GR00T N1.5 default_factory/default changes (transformers compat) - Revert SmolVLA use_peft legacy field - Apply ruff formatting fixes - camera_name_mapping stays entirely in env/eval layer (no policy changes) Made-with: Cursor * Update docs/source/env_processor.mdx Co-authored-by: Khalil Meftah <khalil.meftah@huggingface.co> Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> * feat(envs): lazy env init + AsyncVectorEnv as default for n_envs > 1 LiberoEnv and MetaworldEnv previously allocated GPU resources (EGL context, OpenGL framebuffer) in __init__, before AsyncVectorEnv's fork(). Worker processes inherited stale GPU handles, causing EGL_BAD_CONTEXT crashes on first render. Fix: defer OffScreenRenderEnv / MT1 construction to _ensure_env(), called on first reset() or step() inside the worker subprocess. Each worker creates its own clean context after fork(). Also fixes lerobot_eval.py:170 (add_envs_task TODO): replace with env.call("task") which works with both SyncVectorEnv and AsyncVectorEnv. AsyncVectorEnv is now the default for n_envs > 1; auto-downgraded to SyncVectorEnv when n_envs=1 (no benefit, less overhead). Expected speedup: ~15-20x for LIBERO Spatial with batch_size=50. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix: close envs between tasks to prevent worker process accumulation eval_policy_all never closed environments after each task completed, causing AsyncVectorEnv worker processes to accumulate (N_tasks × n_envs). This led to OOM, BrokenPipeError and EOFError on multi-task benchmarks. Also fixes: - AsyncVectorEnv compat in envs/utils.py (use get_attr/call instead of .envs) - Tuple task handling in tokenizer_processor and lerobot_eval - _LazyAsyncVectorEnv for deferred worker spawning in LIBERO Made-with: Cursor * fix(eval): use task_description instead of task for language conditioning env.call("task") returns the LIBERO task name with underscores (e.g. "pick_up_the_black_bowl_...") instead of the natural language description ("pick up the black bowl ..."). The VLM tokenizes these completely differently, causing 0.0 reward across all episodes. Made-with: Cursor * docs: update adding_benchmarks for async env changes - Replace add_envs_task reference with env.call("task_description") - Update use_async_envs default to True - Add note about lazy GPU init for AsyncVectorEnv compatibility Made-with: Cursor * feat(eval): batch_size=auto + faster env loading - batch_size=0 (default) auto-tunes based on CPU cores, capped by n_episodes and 64. Removes the need for users to guess the right value. The old batch_size > n_episodes error is replaced by silently clamping to n_episodes. - _LazyAsyncVectorEnv accepts pre-computed spaces so only one temp env is created per suite (not per task). For libero_spatial (10 tasks) this avoids 9 redundant LiberoEnv instantiations during env setup. Made-with: Cursor * docs: add evaluation guide and update benchmarks doc - New docs/source/evaluation.mdx covering lerobot-eval usage, batch_size auto-tuning, AsyncVectorEnv performance, tuning tips, output format, multi-task evaluation, and programmatic usage. - Add evaluation page to _toctree.yml under Benchmarks section. - Update adding_benchmarks.mdx to reference batch_size auto default and link to the evaluation guide. Made-with: Cursor * docs(evaluation): remove benchmark table, rename section header Made-with: Cursor * perf(eval): shared memory, observation passthrough, task prefetch - AsyncVectorEnv now uses shared_memory=True for zero-copy observation transfer - LiberoEnvConfig.gym_kwargs passes observation_height/width to the env - eval_policy_all prefetches next task's workers while current task runs Made-with: Cursor * style: ruff format Made-with: Cursor * chore: revert env_processor.mdx changes (not part of this PR) Made-with: Cursor * ci(benchmarks): add isolated integration tests for libero and metaworld Each benchmark gets its own Docker image (lerobot[libero] / lerobot[metaworld] only) so incompatible dep trees cannot collide. A 1-episode smoke eval runs per benchmark on GPU runners. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * ci(benchmarks): pin action hashes and use uv sync --locked Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * ci(benchmarks): trigger only on envs/ or lerobot_eval.py changes Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(ci): set LIBERO_DATA_FOLDER to bypass interactive stdin prompt libero/__init__.py calls input() to ask about a custom dataset path, which raises EOFError when stdin is closed inside Docker. Setting LIBERO_DATA_FOLDER skips the prompt entirely. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * docs(benchmarks): add CI smoke test step to adding_benchmarks guide Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(ci): pre-create libero config in Dockerfile to bypass stdin prompt libero/__init__.py calls input() when ~/.libero/config.yaml is missing. We write the config at image build time (without importing libero) so the prompt never fires at runtime. Also trigger CI on pyproject.toml changes. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(ci): use shell to create libero config instead of multiline python -c The multiline RUN python -c "..." was being parsed as Dockerfile instructions. Use printf to write ~/.libero/config.yaml directly. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(ci): point libero config to bundled package init_files The config was pointing to /tmp/libero_init which doesn't exist. Use importlib.util.find_spec to locate the hf-libero package directory and write paths to the actual bundled bddl_files/init_files/assets. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(ci): add smolvla extra to benchmark Dockerfiles num2words (required by SmolVLM processor) is declared in lerobot[smolvla], not lerobot[libero/metaworld]. Install both extras together. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(eval): render_frame covers _LazyAsyncVectorEnv isinstance(env, AsyncVectorEnv) silently skipped _LazyAsyncVectorEnv, causing video rendering to produce no frames on the default async path. Switch to hasattr(env, "call") so any async-compatible env (including _LazyAsyncVectorEnv) hits the call("render") branch. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * refactor(envs): remove unused _get_sub_env_attr helper _get_sub_env_attr was defined but never called anywhere in the codebase. _sub_env_has_attr (its sibling) is kept — it is actively used in utils.py. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * chore: apply prettier formatting to docs Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * docs(env_processor): remove deprecated add_envs_task from pipeline example add_envs_task is replaced by env.call("task_description") in this PR. Remove it from the pipeline walkthrough and renumber the steps (8→7). Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * refactor(envs): remove __del__ from _LazyAsyncVectorEnv __del__ is unreliable as a cleanup mechanism. close() is already called explicitly in the eval loop's finally block, so the finalizer is redundant. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(eval): prefetch next task's workers after close to avoid GPU memory overlap Previously, next task's AsyncVectorEnv workers were spawned while the current task was still running, causing both tasks' GPU contexts to coexist. Moving the prefetch start into the finally block (after env.close()) ensures workers for task N+1 only spin up once task N has released GPU memory. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * refactor(envs): move _LazyAsyncVectorEnv to utils and apply to metaworld _LazyAsyncVectorEnv lived in libero.py but metaworld had the same OOM problem: all tasks' AsyncVectorEnv workers were spawned eagerly, wasting GPU memory for tasks not yet running. Move the class to envs/utils.py so both environments share it, then apply the same is_async + lazy wrapping pattern in create_metaworld_envs. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * chore: remove out-of-scope benchmark/CI/docs files from PR Benchmark CI workflow, Dockerfiles, benchmark docs, evaluation smoke-test doc, and dispatch tests belong in a separate PR. Scope this PR to the async env init changes only. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * chore: restore adding_benchmarks + test_dispatch, drop env_processor changes - Restore docs/source/adding_benchmarks.mdx (belongs in this PR) - Restore tests/envs/test_dispatch.py (belongs in this PR) - Revert docs/source/env_processor.mdx to main (out of scope for this PR) Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * docs(adding_benchmarks): remove CI smoke test step (coming in separate PR) Step 7 (Dockerfile + benchmark_tests.yml CI job) and its table rows are out of scope for this PR. The CI infrastructure will be added on top in a follow-up PR. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * refactor(envs): remove unused add_envs_task Replaced by env.call("task_description") in lerobot_eval.py. No callers remain in the codebase. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * style: fix prettier formatting in env_processor.mdx Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(eval): catch AttributeError and NotImplementedError explicitly for task description Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(envs): use forkserver context and close envs in test to prevent deadlock AsyncVectorEnv with default fork context leaks worker processes between test_policy parametrized cases; subsequent env creation deadlocks because new forked workers inherit stale pipe FDs from previous test's leaked workers. - configs.py: pass context="forkserver" to AsyncVectorEnv (matches _LazyAsyncVectorEnv) - test_policies.py: call close_envs(envs) at end of test_policy to clean up workers Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> * fix(envs): default use_async_envs=False in create_envs and make_env Tests that call make_env(n_envs=2) without passing use_async_envs were getting AsyncVectorEnv, whose forked workers can't resolve gym namespaces registered at runtime. Default to False (sync) so existing tests pass. lerobot_eval.py explicitly passes cfg.eval.use_async_envs, so the CLI async behaviour (controlled by EvalConfig.use_async_envs) is unchanged. Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com> --------- Signed-off-by: Pepijn <138571049+pkooij@users.noreply.github.com> Co-authored-by: Khalil Meftah <khalil.meftah@huggingface.co> Co-authored-by: Claude Sonnet 4.6 <noreply@anthropic.com>
438 lines
16 KiB
Plaintext
438 lines
16 KiB
Plaintext
# Environment Processors
|
|
|
|
Environment processors are a critical layer in LeRobot's data processing architecture that handle **environment-specific** transformations, separate from policy-specific processing. This separation of concerns enables cleaner code, better modularity, and easier experimentation with different environments and policies.
|
|
|
|
## Why Environment Processors?
|
|
|
|
When working with different robot environments (LIBERO, MetaWorld, Aloha, etc.), each environment often has unique data formats, coordinate systems, and conventions that need standardization **before** policy processing. Without environment processors, these transformations would be:
|
|
|
|
1. **Hardcoded in environment code** - Making it difficult to experiment with different state representations
|
|
2. **Duplicated across policies** - Each policy would need to handle environment-specific quirks
|
|
3. **Mixed with policy logic** - Violating separation of concerns and making debugging harder
|
|
|
|
Environment processors solve this by providing a **dedicated processing layer** between raw environment observations and policy inputs.
|
|
|
|
## The Processing Pipeline
|
|
|
|
Here's how data flows through the complete processing pipeline during evaluation:
|
|
|
|
```python
|
|
# In lerobot_eval.py rollout() function:
|
|
|
|
# 1. Raw environment observation (numpy arrays, various formats)
|
|
raw_observation = env.step(action)
|
|
|
|
# 2. Convert numpy to torch, normalize images [0,1]
|
|
observation = preprocess_observation(raw_observation)
|
|
|
|
# 3. Add task metadata (for multi-task environments)
|
|
observation = add_envs_task(env, observation)
|
|
|
|
# 4. ENVIRONMENT-SPECIFIC preprocessing (NEW!)
|
|
# - Flatten robot states
|
|
# - Rotate images to match dataset conventions
|
|
# - Handle environment-specific coordinate systems
|
|
observation = env_preprocessor(observation)
|
|
|
|
# 5. POLICY-SPECIFIC preprocessing
|
|
# - Normalize with dataset statistics
|
|
# - Add batch dimensions
|
|
# - Move to GPU
|
|
# - Tokenize language instructions
|
|
observation = preprocessor(observation)
|
|
|
|
# 6. Policy inference
|
|
action = policy.select_action(observation)
|
|
|
|
# 7. POLICY-SPECIFIC postprocessing
|
|
# - Unnormalize actions
|
|
# - Remove batch dimensions
|
|
action = postprocessor(action)
|
|
|
|
# 8. ENVIRONMENT-SPECIFIC postprocessing (NEW!)
|
|
# - Convert action formats if needed
|
|
# - Apply environment-specific constraints
|
|
action_transition = {"action": action}
|
|
action_transition = env_postprocessor(action_transition)
|
|
action = action_transition["action"]
|
|
|
|
# 9. Execute in environment
|
|
env.step(action)
|
|
```
|
|
|
|
## The Benefits
|
|
|
|
### 1. **Separation of Concerns**
|
|
|
|
Environment processors handle transformations specific to the **environment's data format**, while policy processors handle transformations specific to the **model's requirements**.
|
|
|
|
```python
|
|
# ❌ Before: Mixed concerns
|
|
class LiberoVLAPolicy:
|
|
def preprocess(self, obs):
|
|
# Environment-specific: Flatten robot state (shouldn't be in policy!)
|
|
state = self._flatten_robot_state(obs["robot_state"])
|
|
# Policy-specific: Normalize with dataset stats
|
|
state = self.normalizer(state)
|
|
return state
|
|
|
|
# ✅ After: Clear separation
|
|
# Environment processor: Handles LIBERO's nested robot state
|
|
env_preprocessor = LiberoProcessorStep() # Flattens robot_state
|
|
|
|
# Policy processor: Handles model requirements
|
|
policy_preprocessor = NormalizerProcessorStep(stats=dataset_stats)
|
|
```
|
|
|
|
### 2. **Flexibility and Reusability**
|
|
|
|
The same policy can work with different environment processors, and the same environment processor can work with different policies:
|
|
|
|
````python
|
|
# Use SmolVLA policy with LIBERO environment
|
|
# Use SmolVLA policy with LIBERO environment
|
|
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
|
env_cfg=libero_cfg,
|
|
policy_cfg=smolvla_cfg,
|
|
)
|
|
smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
|
|
# Or use ACT policy with the same LIBERO environment
|
|
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
|
env_cfg=libero_cfg,
|
|
policy_cfg=act_cfg,
|
|
)
|
|
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
|
|
```python
|
|
# Use SmolVLA policy with LIBERO environment
|
|
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
|
env_cfg=libero_cfg,
|
|
policy_cfg=smolvla_cfg,
|
|
)
|
|
smolvla_preprocessor, smolvla_postprocessor = make_pre_post_processors(smolvla_cfg)
|
|
|
|
# Or use ACT policy with the same LIBERO environment
|
|
libero_preprocessor, libero_postprocessor = make_env_pre_post_processors(
|
|
env_cfg=libero_cfg,
|
|
policy_cfg=act_cfg,
|
|
)
|
|
act_preprocessor, act_postprocessor = make_pre_post_processors(act_cfg)
|
|
|
|
### 3. **Easier Experimentation**
|
|
|
|
Want to try different state representations for LIBERO? Just create a new processor:
|
|
|
|
```python
|
|
# Original: 8D state (pos + quat→axisangle + gripper)
|
|
@ProcessorStepRegistry.register("libero_processor")
|
|
class LiberoProcessorStep(ObservationProcessorStep):
|
|
def _process_observation(self, obs):
|
|
eef_pos = robot_state["eef"]["pos"] # 3D
|
|
eef_axisangle = quat2axisangle(quat) # 3D
|
|
gripper = robot_state["gripper"]["qpos"] # 2D
|
|
state = torch.cat([eef_pos, eef_axisangle, gripper], dim=-1) # 8D
|
|
return state
|
|
|
|
# Experiment: Add velocity for better control
|
|
@ProcessorStepRegistry.register("libero_velocity_processor")
|
|
class LiberoVelocityProcessorStep(ObservationProcessorStep):
|
|
def _process_observation(self, obs):
|
|
# Include velocities for 14D state
|
|
eef_pos = robot_state["eef"]["pos"] # 3D
|
|
eef_axisangle = quat2axisangle(quat) # 3D
|
|
eef_vel = robot_state["eef"]["vel"] # 3D (NEW)
|
|
gripper_pos = robot_state["gripper"]["qpos"] # 2D
|
|
gripper_vel = robot_state["gripper"]["qvel"] # 3D (NEW)
|
|
state = torch.cat([eef_pos, eef_axisangle, eef_vel,
|
|
gripper_pos, gripper_vel], dim=-1) # 14D
|
|
return state
|
|
````
|
|
|
|
### 4. **Cleaner Environment Code**
|
|
|
|
Environments expose **all available data** without needing to know what downstream models will use:
|
|
|
|
```python
|
|
# LIBERO environment exposes full robot state
|
|
observation = {
|
|
"pixels": {"image": img, "image2": img2},
|
|
"robot_state": {
|
|
"eef": {"pos": ..., "quat": ..., "vel": ..., "mat": ..., "axisangle": ...},
|
|
"gripper": {"qpos": ..., "qvel": ...},
|
|
"joints": {"pos": ..., "vel": ...}
|
|
}
|
|
}
|
|
|
|
# Environment processor decides what to use
|
|
# Policy processor handles model-specific transformations
|
|
```
|
|
|
|
## Using Environment Processors
|
|
|
|
### Factory Function
|
|
|
|
The `make_env_pre_post_processors` function follows the same pattern as `make_pre_post_processors` for policies:
|
|
|
|
```python
|
|
from lerobot.envs.factory import make_env_pre_post_processors
|
|
from lerobot.envs.configs import LiberoEnv, PushtEnv
|
|
|
|
# For LIBERO: Returns LiberoProcessorStep in preprocessor
|
|
libero_cfg = LiberoEnv(task="libero_spatial", camera_name=["agentview"])
|
|
env_preprocessor, env_postprocessor = make_env_pre_post_processors(libero_cfg)
|
|
|
|
# For other environments: Returns identity processors (no-op)
|
|
pusht_cfg = PushtEnv()
|
|
env_preprocessor, env_postprocessor = make_env_pre_post_processors(pusht_cfg)
|
|
```
|
|
|
|
### Implementation in `envs/factory.py`
|
|
|
|
```python
|
|
def make_env_pre_post_processors(
|
|
env_cfg: EnvConfig,
|
|
) -> tuple[
|
|
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
|
PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
|
|
]:
|
|
"""
|
|
Create preprocessor and postprocessor pipelines for environment observations.
|
|
|
|
Args:
|
|
env_cfg: The configuration of the environment.
|
|
|
|
Returns:
|
|
A tuple containing:
|
|
- preprocessor: Pipeline that processes environment observations
|
|
- postprocessor: Pipeline that processes environment outputs
|
|
"""
|
|
# For LIBERO environments, add the LiberoProcessorStep to preprocessor
|
|
if isinstance(env_cfg, LiberoEnv) or "libero" in env_cfg.type:
|
|
preprocessor = PolicyProcessorPipeline(steps=[LiberoProcessorStep()])
|
|
else:
|
|
# For all other environments, return an identity preprocessor
|
|
preprocessor = PolicyProcessorPipeline(steps=[])
|
|
|
|
# Postprocessor is currently identity for all environments
|
|
# Future: Could add environment-specific action transformations
|
|
postprocessor = PolicyProcessorPipeline(steps=[])
|
|
|
|
return preprocessor, postprocessor
|
|
```
|
|
|
|
### Integration in Evaluation
|
|
|
|
In `lerobot_eval.py`, the environment processors are created once and used throughout:
|
|
|
|
```python
|
|
def eval_main(cfg: EvalPipelineConfig):
|
|
# Create environment
|
|
envs = make_env(cfg.env, n_envs=cfg.eval.batch_size)
|
|
|
|
# Create policy
|
|
policy = make_policy(cfg=cfg.policy, env_cfg=cfg.env)
|
|
|
|
# Create policy processors
|
|
preprocessor, postprocessor = make_pre_post_processors(
|
|
policy_cfg=cfg.policy,
|
|
pretrained_path=cfg.policy.pretrained_path,
|
|
)
|
|
|
|
# Create environment processors (NEW!)
|
|
env_preprocessor, env_postprocessor = make_env_pre_post_processors(env_cfg=cfg.env)
|
|
|
|
# Run evaluation with both processor types
|
|
eval_policy_all(
|
|
envs=envs,
|
|
policy=policy,
|
|
env_preprocessor=env_preprocessor, # Environment-specific
|
|
env_postprocessor=env_postprocessor, # Environment-specific
|
|
preprocessor=preprocessor, # Policy-specific
|
|
postprocessor=postprocessor, # Policy-specific
|
|
n_episodes=cfg.eval.n_episodes,
|
|
)
|
|
```
|
|
|
|
## Example: LIBERO Environment Processor
|
|
|
|
The `LiberoProcessorStep` demonstrates a real-world environment processor:
|
|
|
|
```python
|
|
from lerobot.processor.pipeline import ObservationProcessorStep
|
|
|
|
@dataclass
|
|
@ProcessorStepRegistry.register(name="libero_processor")
|
|
class LiberoProcessorStep(ObservationProcessorStep):
|
|
"""
|
|
Processes LIBERO observations into the LeRobot format.
|
|
|
|
**State Processing:**
|
|
- Extracts end-effector position (3D)
|
|
- Converts quaternion to axis-angle representation (3D)
|
|
- Extracts gripper joint positions (2D)
|
|
- Concatenates into 8D state vector
|
|
|
|
**Image Processing:**
|
|
- Rotates images 180° to match HuggingFaceVLA/libero convention
|
|
"""
|
|
|
|
def _process_observation(self, observation):
|
|
processed_obs = observation.copy()
|
|
|
|
# Process images: Flip 180° for camera convention
|
|
for key in list(processed_obs.keys()):
|
|
if key.startswith("observation.images."):
|
|
img = processed_obs[key]
|
|
img = torch.flip(img, dims=[2, 3]) # Flip H and W
|
|
processed_obs[key] = img
|
|
|
|
# Process robot_state: Flatten to 8D vector
|
|
if "observation.robot_state" in processed_obs:
|
|
robot_state = processed_obs.pop("observation.robot_state")
|
|
|
|
eef_pos = robot_state["eef"]["pos"] # (B, 3)
|
|
eef_quat = robot_state["eef"]["quat"] # (B, 4)
|
|
gripper_qpos = robot_state["gripper"]["qpos"] # (B, 2)
|
|
|
|
# Convert quaternion to axis-angle
|
|
eef_axisangle = self._quat2axisangle(eef_quat) # (B, 3)
|
|
|
|
# Concatenate into single state vector
|
|
state = torch.cat((eef_pos, eef_axisangle, gripper_qpos), dim=-1)
|
|
state = state.float()
|
|
|
|
processed_obs["observation.state"] = state
|
|
|
|
return processed_obs
|
|
```
|
|
|
|
### Why These Transformations?
|
|
|
|
1. **Image Rotation**: The HuggingFaceVLA/libero dataset has images rotated 180° from the raw LIBERO simulator. The processor handles this convention mismatch so policies trained on the dataset work seamlessly.
|
|
|
|
2. **State Flattening**: The raw LIBERO environment exposes nested dictionaries with all available state information (position, quaternion, velocity, matrix representation, etc.). The processor:
|
|
- Selects the relevant components (pos, quat, gripper)
|
|
- Converts quaternion to axis-angle (more suitable for learning)
|
|
- Flattens to a single 8D vector that policies expect
|
|
|
|
3. **Flexibility**: The environment still exposes **all** raw data. If you want to try different state representations (e.g., including velocities, using matrix representation instead of axis-angle), you can create a new processor without modifying the environment code.
|
|
|
|
## Adding Environment Processors for New Environments
|
|
|
|
To add environment processors for a new environment:
|
|
|
|
### 1. Create the Processor Step
|
|
|
|
```python
|
|
# In src/lerobot/processor/env_processor.py
|
|
|
|
@dataclass
|
|
@ProcessorStepRegistry.register(name="myenv_processor")
|
|
class MyEnvProcessorStep(ObservationProcessorStep):
|
|
"""Process observations from MyEnv."""
|
|
|
|
def _process_observation(self, observation):
|
|
processed = observation.copy()
|
|
|
|
# Your environment-specific transformations
|
|
if "myenv.specific.state" in processed:
|
|
state = processed.pop("myenv.specific.state")
|
|
# Transform to standard format
|
|
processed["observation.state"] = self._transform_state(state)
|
|
|
|
return processed
|
|
```
|
|
|
|
### 2. Update Your `EnvConfig` Subclass
|
|
|
|
```python
|
|
# In src/lerobot/envs/factory.py
|
|
|
|
def make_env_pre_post_processors(env_cfg: EnvConfig):
|
|
if isinstance(env_cfg, LiberoEnv) or "libero" in env_cfg.type:
|
|
preprocessor = PolicyProcessorPipeline(steps=[LiberoProcessorStep()])
|
|
elif isinstance(env_cfg, MyEnvConfig) or "myenv" in env_cfg.type:
|
|
preprocessor = PolicyProcessorPipeline(steps=[MyEnvProcessorStep()])
|
|
else:
|
|
preprocessor = PolicyProcessorPipeline(steps=[])
|
|
|
|
postprocessor = PolicyProcessorPipeline(steps=[])
|
|
return preprocessor, postprocessor
|
|
```
|
|
|
|
### 3. Use in Evaluation
|
|
|
|
No changes needed! The evaluation script automatically uses the appropriate processor:
|
|
|
|
```bash
|
|
lerobot-eval \
|
|
--policy.path=lerobot/my_policy \
|
|
--env.type=myenv \ # Automatically uses MyEnvProcessorStep
|
|
--eval.n_episodes=10
|
|
```
|
|
|
|
## Future: Environment Postprocessors
|
|
|
|
Currently, postprocessors are identity (no-op) for all environments. Future use cases include:
|
|
|
|
### Action Space Transformations
|
|
|
|
```python
|
|
@dataclass
|
|
class MyEnvActionPostprocessor(ProcessorStep):
|
|
"""Convert policy actions to environment-specific format."""
|
|
|
|
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
|
action = transition["action"]
|
|
|
|
# Example: Convert from Cartesian to joint space
|
|
if self.action_space == "joint":
|
|
action = self.ik_solver(action)
|
|
|
|
# Example: Apply environment-specific safety limits
|
|
action = torch.clamp(action, self.min_action, self.max_action)
|
|
|
|
transition["action"] = action
|
|
return transition
|
|
```
|
|
|
|
### Coordinate System Conversions
|
|
|
|
```python
|
|
@dataclass
|
|
class CoordinateTransformPostprocessor(ProcessorStep):
|
|
"""Transform actions between coordinate systems."""
|
|
|
|
def __call__(self, transition: EnvTransition) -> EnvTransition:
|
|
action = transition["action"]
|
|
|
|
# Example: Policy outputs in world frame, env expects base frame
|
|
action = self.world_to_base_transform(action)
|
|
|
|
transition["action"] = action
|
|
return transition
|
|
```
|
|
|
|
## Best Practices
|
|
|
|
1. **Keep environment processors simple**: They should only handle environment-specific data format issues, not complex learning-related transformations.
|
|
|
|
2. **Use policy processors for model requirements**: Normalization, batching, device placement, and tokenization belong in policy processors.
|
|
|
|
3. **Expose all data from environments**: Let processors decide what to use rather than hardcoding choices in the environment.
|
|
|
|
4. **Document conventions**: Clearly document any coordinate system conventions, camera orientations, or data formats that your processor handles.
|
|
|
|
5. **Test independently**: Environment processors should be testable without loading full policies or environments.
|
|
|
|
## Summary
|
|
|
|
Environment processors provide a **clean separation** between environment-specific data transformations and policy-specific model requirements. This architecture:
|
|
|
|
- ✅ Enables easy experimentation with different state representations
|
|
- ✅ Allows policies to work seamlessly across different environments
|
|
- ✅ Keeps environment code focused on simulation/hardware interface
|
|
- ✅ Makes processor pipelines more maintainable and debuggable
|
|
- ✅ Follows the single responsibility principle
|
|
|
|
The key insight: **Environments define data formats, processors standardize them, policies consume standardized data.** Each layer has a clear, focused responsibility.
|