Refactor observation preprocessing to use a modular pipeline system

- Introduced `RobotPipeline` and `ObservationProcessor` for handling observation transformations.
- Updated `preprocess_observation` to maintain backward compatibility while leveraging the new pipeline.
- Added tests for the new processing components and ensured they match the original functionality.
- Removed hardcoded logic in favor of a more flexible, composable architecture.
This commit is contained in:
Adil Zouitine
2025-07-02 17:29:58 +02:00
parent 945e1ff266
commit f6c7287ae7
9 changed files with 1472 additions and 50 deletions
+9 -2
View File
@@ -22,7 +22,9 @@ from gymnasium.utils.env_checker import check_env
import lerobot
from lerobot.envs.factory import make_env, make_env_config
from lerobot.envs.utils import preprocess_observation
from lerobot.processor.pipeline import RobotPipeline, TransitionIndex
from lerobot.processor.observation_processor import ObservationProcessor
from tests.utils import require_env
OBS_TYPES = ["state", "pixels", "pixels_agent_pos"]
@@ -48,7 +50,12 @@ def test_factory(env_name):
cfg = make_env_config(env_name)
env = make_env(cfg, n_envs=1)
obs, _ = env.reset()
obs = preprocess_observation(obs)
# Process observation using pipeline
obs_pipeline = RobotPipeline([ObservationProcessor()])
transition = (obs, None, None, None, None, None, None)
processed_transition = obs_pipeline(transition)
obs = processed_transition[TransitionIndex.OBSERVATION]
# test image keys are float32 in range [0,1]
for key in obs: