Address GROOT relative action review feedback

src/lerobot/policies/factory.py

@@ -288,6 +288,7 @@ def make_pre_post_processors(
                 config=policy_cfg,
                 pretrained_path=pretrained_path,
                 dataset_stats=kwargs.get("dataset_stats"),
+                dataset_meta=kwargs.get("dataset_meta"),
                 preprocessor_overrides=kwargs.get("preprocessor_overrides"),
                 postprocessor_overrides=kwargs.get("postprocessor_overrides"),
                 preprocessor_config_filename=kwargs.get(
@@ -402,6 +403,7 @@ def make_pre_post_processors(
         processors = make_groot_pre_post_processors(
             config=policy_cfg,
             dataset_stats=kwargs.get("dataset_stats"),
+            dataset_meta=kwargs.get("dataset_meta"),
         )
 
     elif isinstance(policy_cfg, XVLAConfig):

src/lerobot/policies/groot/configuration_groot.py

@@ -268,7 +268,6 @@ class GrootConfig(PreTrainedConfig):
     )
 
     # Groot-specific model parameters
-    model_version: str = GROOT_N1_7
 
     # Path or HuggingFace model ID for the base GR00T N1.7 model whose backbone weights and
     # checkpoint sidecars (statistics.json, processor_config.json, ...) are loaded. This is the
@@ -325,11 +324,10 @@ class GrootConfig(PreTrainedConfig):
     # Set to True only after installing a flash-attn build matching your torch/CUDA env.
     use_flash_attention: bool = False
 
-    # Train on state-relative action chunks. The listed joints stay absolute, which is normally used
-    # for gripper channels whose command frame is not the arm joint state.
+    # Enable GR00T-style state-relative action chunks. Prefer deriving action representation from
+    # embodiment metadata; relative_exclude_joints is a flat-vector override for datasets without it.
     use_relative_actions: bool = False
     relative_exclude_joints: list[str] = field(default_factory=list)
-    action_feature_names: list[str] | None = None
 
     # Training parameters
     optimizer_lr: float = 1e-4
@@ -365,8 +363,6 @@ class GrootConfig(PreTrainedConfig):
     resume: bool = False
 
     def __post_init__(self):
-        self.model_version = normalize_groot_model_version(self.model_version)
-
         if self.tokenizer_assets_repo is not None:
             raise ValueError(
                 "Config sets 'tokenizer_assets_repo', which only existed for GR00T N1.5; this looks "
@@ -417,9 +413,9 @@ class GrootConfig(PreTrainedConfig):
                 setattr(self, field_name, n1_7_value)
 
         inferred_version = infer_groot_model_version(self.base_model_path)
-        if inferred_version is not None and inferred_version != self.model_version:
+        if inferred_version is not None and inferred_version != GROOT_N1_7:
             message = (
-                f"GR00T model_version '{self.model_version}' does not match base_model_path "
+                f"GR00T model_version '{GROOT_N1_7}' does not match base_model_path "
                 f"'{self.base_model_path}', which looks like '{inferred_version}'."
             )
             if inferred_version == GROOT_N1_5:

src/lerobot/policies/groot/processor_groot.py

@@ -434,6 +434,7 @@ def make_groot_pre_post_processors_from_pretrained(
     pretrained_path: str,
     *,
     dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+    dataset_meta: Any | None = None,
     preprocessor_overrides: dict[str, Any] | None = None,
     postprocessor_overrides: dict[str, Any] | None = None,
     preprocessor_config_filename: str = f"{POLICY_PREPROCESSOR_DEFAULT_NAME}.json",
@@ -456,6 +457,7 @@ def make_groot_pre_post_processors_from_pretrained(
         preprocessor, postprocessor = make_groot_pre_post_processors(
             config=processor_cfg,
             dataset_stats=dataset_stats,
+            dataset_meta=dataset_meta,
         )
         # Raw checkpoints have no serialized pipelines to load overrides into,
         # so apply the caller overrides (e.g. device and rename_map from
@@ -545,8 +547,20 @@ def _reconnect_groot_n1_7_pack_decode_steps(
             step.pack_step = pack_step
 
 
+def _resolve_action_feature_names_from_dataset_meta(dataset_meta: Any | None) -> list[str] | None:
+    features = getattr(dataset_meta, "features", {}) or {}
+    action_feature = features.get(ACTION) if isinstance(features, dict) else None
+    if isinstance(action_feature, dict):
+        names = action_feature.get("names")
+    else:
+        names = getattr(action_feature, "names", None)
+    return list(names) if names is not None else None
+
+
 def make_groot_pre_post_processors(
-    config: GrootConfig, dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None
+    config: GrootConfig,
+    dataset_stats: dict[str, dict[str, torch.Tensor]] | None = None,
+    dataset_meta: Any | None = None,
 ) -> tuple[
     PolicyProcessorPipeline[dict[str, Any], dict[str, Any]],
     PolicyProcessorPipeline[PolicyAction, PolicyAction],
@@ -659,7 +673,7 @@ def make_groot_pre_post_processors(
         relative_step = RelativeActionsProcessorStep(
             enabled=True,
             exclude_joints=list(config.relative_exclude_joints or []),
-            action_names=config.action_feature_names,
+            action_names=_resolve_action_feature_names_from_dataset_meta(dataset_meta),
         )
         input_steps.insert(2, relative_step)
 

src/lerobot/scripts/lerobot_train.py

@@ -256,11 +256,7 @@ def _iter_action_state_training_samples(dataset: Any):
         yield item.get(ACTION), item.get(OBS_STATE), item.get(f"{ACTION}_is_pad")
 
 
-def _resolve_action_feature_names(active_cfg: Any, dataset: Any) -> list[str] | None:
-    config_names = getattr(active_cfg, "action_feature_names", None)
-    if config_names is not None:
-        return list(config_names)
-
+def _resolve_action_feature_names(dataset: Any) -> list[str] | None:
     features = getattr(getattr(dataset, "meta", None), "features", {}) or {}
     action_feature = features.get(ACTION) if isinstance(features, dict) else None
     if isinstance(action_feature, dict):
@@ -475,14 +471,14 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
         processor_stats = _make_relative_action_training_stats(
             dataset,
             exclude_joints=getattr(active_cfg, "relative_exclude_joints", []),
-            action_names=_resolve_action_feature_names(active_cfg, dataset),
+            action_names=_resolve_action_feature_names(dataset),
         )
 
     processor_kwargs = {}
     if (processor_pretrained_path and not cfg.resume) or not processor_pretrained_path:
         processor_kwargs["dataset_stats"] = processor_stats
 
-    if cfg.is_reward_model_training:
+    if cfg.is_reward_model_training or getattr(active_cfg, "use_relative_actions", False):
         processor_kwargs["dataset_meta"] = dataset.meta
 
     if not cfg.is_reward_model_training and processor_pretrained_path is not None:
@@ -506,7 +502,7 @@ def train(cfg: TrainPipelineConfig, accelerator: "Accelerator | None" = None):
             preprocessor_overrides["relative_actions_processor"] = {
                 "enabled": True,
                 "exclude_joints": getattr(active_cfg, "relative_exclude_joints", []),
-                "action_names": _resolve_action_feature_names(active_cfg, dataset),
+                "action_names": _resolve_action_feature_names(dataset),
             }
             postprocessor_overrides["absolute_actions_processor"] = {"enabled": True}
         processor_kwargs["preprocessor_overrides"] = preprocessor_overrides