From 33969a03373860b325631d8041bddb4b4b8c16ba Mon Sep 17 00:00:00 2001
From: Adil Zouitine <adilzouitinegm@gmail.com>
Date: Tue, 8 Jul 2025 15:13:41 +0200
Subject: [PATCH] refactor(pipeline): Simplify observation and padding data
 handling in batch transitions

---
 src/lerobot/processor/pipeline.py | 26 ++++++++++++++------------
 1 file changed, 14 insertions(+), 12 deletions(-)

diff --git a/src/lerobot/processor/pipeline.py b/src/lerobot/processor/pipeline.py
index d2beed3c3..038711339 100644
--- a/src/lerobot/processor/pipeline.py
+++ b/src/lerobot/processor/pipeline.py
@@ -183,15 +183,13 @@ def _default_batch_to_transition(batch: dict[str, Any]) -> EnvTransition:  # noq
     metadata without breaking the processor.
     """
 
-    # Handle observation and observation.* keys
+    # Extract observation keys
     observation_keys = {k: v for k, v in batch.items() if k.startswith("observation.")}
+    observation = observation_keys if observation_keys else None
 
-    observation = None
-    if observation_keys:
-        observation = {}
-        # Add observation.* keys to the observation dict
-        for key, value in observation_keys.items():
-            observation[key] = value
+    # Extract padding keys for complementary data
+    pad_keys = {k: v for k, v in batch.items() if "_is_pad" in k}
+    complementary_data = pad_keys if pad_keys else {}
 
     return (
         observation,
@@ -200,7 +198,7 @@ def _default_batch_to_transition(batch: dict[str, Any]) -> EnvTransition:  # noq
         batch.get("next.done", False),
         batch.get("next.truncated", False),
         batch.get("info", {}),
-        {},
+        complementary_data,
     )
 
 
@@ -216,7 +214,7 @@ def _default_transition_to_batch(transition: EnvTransition) -> dict[str, Any]:
         done,
         truncated,
         info,
-        _,
+        complementary_data,
     ) = transition
 
     batch = {
@@ -227,11 +225,15 @@ def _default_transition_to_batch(transition: EnvTransition) -> dict[str, Any]:
         "info": info,
     }
 
+    # Add padding data from complementary_data
+    if complementary_data:
+        pad_data = {k: v for k, v in complementary_data.items() if "_is_pad" in k}
+        batch.update(pad_data)
+
     # Handle observation - flatten dict to observation.* keys if it's a dict
     if isinstance(observation, dict):
-        # Check if this looks like a dict that was created from observation.* keys
-        for key, value in observation.items():
-            batch[key] = value
+        batch.update(observation)
+
     return batch