feat(processors): Add ToBatchProcessor to multiple policy processors

- Integrated ToBatchProcessor into various policy processors to handle observation batching.
- Updated make functions for act, diffusion, pi0, pi0fast, sac, smolvla, tdmpc, and vqbet processors to include the new batching functionality.
- Ensured consistency across all processor implementations for improved data handling.
This commit is contained in:
Adil Zouitine
2025-07-24 10:39:46 +02:00
committed by Steven Palma
parent f5c6b03b61
commit 8b4a5368b3
8 changed files with 16 additions and 0 deletions
@@ -19,6 +19,7 @@ from lerobot.policies.act.configuration_act import ACTConfig
from lerobot.processor import (
NormalizerProcessor,
RobotProcessor,
ToBatchProcessor,
UnnormalizerProcessor,
)
@@ -33,6 +34,7 @@ def make_act_processor(
NormalizerProcessor(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
ToBatchProcessor(),
]
output_steps = [
UnnormalizerProcessor(
@@ -20,6 +20,7 @@ from lerobot.policies.diffusion.configuration_diffusion import DiffusionConfig
from lerobot.processor import (
NormalizerProcessor,
RobotProcessor,
ToBatchProcessor,
UnnormalizerProcessor,
)
@@ -34,6 +35,7 @@ def make_diffusion_processor(
NormalizerProcessor(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
ToBatchProcessor(),
]
output_steps = [
UnnormalizerProcessor(
@@ -20,6 +20,7 @@ from lerobot.policies.pi0.configuration_pi0 import PI0Config
from lerobot.processor import (
NormalizerProcessor,
RobotProcessor,
ToBatchProcessor,
UnnormalizerProcessor,
)
@@ -34,6 +35,7 @@ def make_pi0_processor(
NormalizerProcessor(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
ToBatchProcessor(),
]
output_steps = [
UnnormalizerProcessor(
@@ -20,6 +20,7 @@ from lerobot.policies.pi0.configuration_pi0 import PI0Config
from lerobot.processor import (
NormalizerProcessor,
RobotProcessor,
ToBatchProcessor,
UnnormalizerProcessor,
)
@@ -34,6 +35,7 @@ def make_pi0_processor(
NormalizerProcessor(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
ToBatchProcessor(),
]
output_steps = [
UnnormalizerProcessor(
@@ -21,6 +21,7 @@ from lerobot.policies.sac.configuration_sac import SACConfig
from lerobot.processor import (
NormalizerProcessor,
RobotProcessor,
ToBatchProcessor,
UnnormalizerProcessor,
)
@@ -35,6 +36,7 @@ def make_sac_processor(
NormalizerProcessor(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
ToBatchProcessor(),
]
output_steps = [
UnnormalizerProcessor(
@@ -19,6 +19,7 @@ from lerobot.policies.smolvla.configuration_smolvla import SmolVLAConfig
from lerobot.processor import (
NormalizerProcessor,
RobotProcessor,
ToBatchProcessor,
UnnormalizerProcessor,
)
@@ -33,6 +34,7 @@ def make_smolvla_processor(
NormalizerProcessor(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
ToBatchProcessor(),
]
output_steps = [
UnnormalizerProcessor(
@@ -20,6 +20,7 @@ from lerobot.policies.tdmpc.configuration_tdmpc import TDMPCConfig
from lerobot.processor import (
NormalizerProcessor,
RobotProcessor,
ToBatchProcessor,
UnnormalizerProcessor,
)
@@ -34,6 +35,7 @@ def make_tdmpc_processor(
NormalizerProcessor(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
ToBatchProcessor(),
]
output_steps = [
UnnormalizerProcessor(
@@ -21,6 +21,7 @@ from lerobot.policies.vqbet.configuration_vqbet import VQBeTConfig
from lerobot.processor import (
NormalizerProcessor,
RobotProcessor,
ToBatchProcessor,
UnnormalizerProcessor,
)
@@ -35,6 +36,7 @@ def make_vqbet_processor(
NormalizerProcessor(
features=config.output_features, norm_map=config.normalization_mapping, stats=dataset_stats
),
ToBatchProcessor(),
]
output_steps = [
UnnormalizerProcessor(