diff --git a/src/lerobot/configs/policies.py b/src/lerobot/configs/policies.py index 0ecfa169b..7271d7c9a 100644 --- a/src/lerobot/configs/policies.py +++ b/src/lerobot/configs/policies.py @@ -52,7 +52,6 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno output_normalization_modes: Similar dictionary as `input_normalization_modes`, but to unnormalize to the original scale. """ - n_obs_steps: int = 1 input_features: dict[str, PolicyFeature] = field(default_factory=dict) @@ -203,7 +202,7 @@ class PreTrainedConfig(draccus.ChoiceRegistry, HubMixin, abc.ABC): # type: igno with open(config_file) as f: config = json.load(f) - + config.pop("type") with tempfile.NamedTemporaryFile("w+", delete=False, suffix=".json") as f: json.dump(config, f) diff --git a/src/lerobot/policies/xvla/__init__.py b/src/lerobot/policies/xvla/__init__.py index 42ead21df..afc7e0142 100644 --- a/src/lerobot/policies/xvla/__init__.py +++ b/src/lerobot/policies/xvla/__init__.py @@ -1,2 +1,2 @@ # add domainid -from lerobot.policies.xvla.processor_xvla import XVLAAddDomainIdProcessorStep +from lerobot.policies.xvla.processor_xvla import XVLAAddDomainIdProcessorStep, XVLAImageNetNormalizeProcessorStep, XVLAImageToFloatProcessorStep diff --git a/src/lerobot/policies/xvla/action_hub.py b/src/lerobot/policies/xvla/action_hub.py index e7e6485cd..0de151ec7 100644 --- a/src/lerobot/policies/xvla/action_hub.py +++ b/src/lerobot/policies/xvla/action_hub.py @@ -1,5 +1,5 @@ # ------------------------------------------------------------------------------ -# Copyright 2025 2toINF (https://github.com/2toINF) +# Copyright 2025 2toINF and HuggingFace Inc. (https://github.com/2toINF) # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. @@ -293,6 +293,217 @@ class FrankaJoint7ActionSpace(BaseActionSpace): return action +@register_action("so101_bimanual_old") +class BimanualSO101OldActionSpace(BaseActionSpace): + """ + Bimanual SO101 robot: 2 arms with 5 joints each + gripper. + + Layout: [left_arm (5 joints + gripper), right_arm (5 joints + gripper)] + - Left arm: shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper + - Right arm: shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper + Total: 12 dimensions + """ + + dim_action = 12 + gripper_idx = (5, 11) # left_gripper at idx 5, right_gripper at idx 11 + GRIPPER_SCALE = 1.0 + JOINTS_SCALE = 1.0 + + # Indices for left and right arm joints (excluding grippers) + LEFT_ARM_JOINTS = (0, 1, 2, 3, 4) + RIGHT_ARM_JOINTS = (6, 7, 8, 9, 10) + + def __init__(self): + super().__init__() + self.mse = nn.MSELoss() + self.bce = nn.BCEWithLogitsLoss() + + def compute_loss(self, pred, target): + assert pred.shape == target.shape, "pred/target shapes must match" + batch_size, seq_len, action_dim = pred.shape + _ensure_indices_valid(action_dim, self.gripper_idx, "gripper_idx") + + # Gripper BCE loss (binary classification for open/close) + g_losses = [self.bce(pred[:, :, gi], target[:, :, gi]) for gi in self.gripper_idx] + gripper_loss = sum(g_losses) / len(self.gripper_idx) * self.GRIPPER_SCALE + + # Joint positions MSE (all non-gripper dimensions) + joints_idx = tuple(i for i in range(action_dim) if i not in set(self.gripper_idx)) + joints_loss = self.mse(pred[:, :, joints_idx], target[:, :, joints_idx]) * self.JOINTS_SCALE + + # Separate losses for left and right arms for better monitoring + left_arm_loss = self.mse(pred[:, :, self.LEFT_ARM_JOINTS], target[:, :, self.LEFT_ARM_JOINTS]) + right_arm_loss = self.mse(pred[:, :, self.RIGHT_ARM_JOINTS], target[:, :, self.RIGHT_ARM_JOINTS]) + + return { + "joints_loss": joints_loss, + "gripper_loss": gripper_loss, + "left_arm_loss": left_arm_loss, + "right_arm_loss": right_arm_loss, + } + + def preprocess(self, proprio, action, mode="train"): + """Zero-out gripper channels in proprio/action to focus learning on continuous joint control.""" + proprio_m = proprio.clone() + action_m = action.clone() + proprio_m[..., self.gripper_idx] = 0.0 + action_m[..., self.gripper_idx] = 0.0 + return proprio_m, action_m + + def postprocess(self, action: torch.Tensor) -> torch.Tensor: + """Apply sigmoid to gripper logits to convert to [0, 1] range.""" + if action.size(-1) > max(self.gripper_idx): + action[..., self.gripper_idx] = torch.sigmoid(action[..., self.gripper_idx]) + return action + +@register_action("so101_bimanual") +class BimanualSO101ActionSpace(BaseActionSpace): + """ + Bimanual SO101 robot: 2 arms with 5 joints each + gripper. + + Layout (real robot): + [left_arm (5 joints + gripper), right_arm (5 joints + gripper)] + - Left arm: shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper + - Right arm: shoulder_pan, shoulder_lift, elbow_flex, wrist_flex, wrist_roll, gripper + + Real action dim: 12 + Model-facing dim: 20 (extra 8 dummy dims at the end) + """ + + # Model output / training dimension (to match pretrained policy) + dim_action = 20 + + # Real robot action dimension + REAL_DIM = 12 + + # Indices of real vs dummy channels + REAL_IDXS = tuple(range(REAL_DIM)) # 0..11 + DUMMY_IDXS = tuple(range(REAL_DIM, dim_action)) # 12..19 + + # Grippers live in the real part + gripper_idx = (5, 11) # left_gripper at idx 5, right_gripper at idx 11 + GRIPPER_SCALE = 1.0 + JOINTS_SCALE = 1.0 + + # Indices for left and right arm joints (excluding grippers) + LEFT_ARM_JOINTS = (0, 1, 2, 3, 4) + RIGHT_ARM_JOINTS = (6, 7, 8, 9, 10) + + def __init__(self): + super().__init__() + self.mse = nn.MSELoss() + self.bce = nn.BCEWithLogitsLoss() + + # ---------- helpers ---------- + + def _pad_to_model_dim(self, x: torch.Tensor) -> torch.Tensor: + """If last dim is REAL_DIM (12), pad zeros to reach dim_action (20).""" + if x is None: + return None + if x.size(-1) == self.dim_action: + return x + if x.size(-1) != self.REAL_DIM: + raise ValueError( + f"Expected last dim to be {self.REAL_DIM} or {self.dim_action}, got {x.size(-1)}" + ) + pad_shape = list(x.shape[:-1]) + [self.dim_action - self.REAL_DIM] + pad = x.new_zeros(pad_shape) + return torch.cat([x, pad], dim=-1) + + def _trim_to_real_dim(self, x: torch.Tensor) -> torch.Tensor: + """Keep only the first REAL_DIM (12) dims for the real robot.""" + return x[..., :self.REAL_DIM] + + # ---------- loss ---------- + + def compute_loss(self, pred, target): + """ + pred: [B, T, 20] from the model + target: [B, T, 12] or [B, T, 20] + We pad target → 20 and compute loss only on the real dims. + """ + # Ensure both are [B, T, 20] + pred = self._pad_to_model_dim(pred) + target = self._pad_to_model_dim(target) + assert pred.shape == target.shape, "pred/target shapes must match" + batch_size, seq_len, action_dim = pred.shape + _ensure_indices_valid(action_dim, self.gripper_idx, "gripper_idx") + + # --- Gripper BCE loss (only real gripper indices) --- + g_losses = [self.bce(pred[:, :, gi], target[:, :, gi]) for gi in self.gripper_idx] + gripper_loss = sum(g_losses) / len(self.gripper_idx) * self.GRIPPER_SCALE + + # --- Joint positions MSE (all non-gripper *real* dims) --- + real_set = set(self.REAL_IDXS) + joints_idx = tuple(i for i in real_set if i not in set(self.gripper_idx)) + + joints_loss = self.mse( + pred[:, :, joints_idx], + target[:, :, joints_idx], + ) * self.JOINTS_SCALE + + # Separate losses for left and right arms for better monitoring + left_arm_loss = self.mse( + pred[:, :, self.LEFT_ARM_JOINTS], + target[:, :, self.LEFT_ARM_JOINTS], + ) + right_arm_loss = self.mse( + pred[:, :, self.RIGHT_ARM_JOINTS], + target[:, :, self.RIGHT_ARM_JOINTS], + ) + return { + "joints_loss": joints_loss, + "gripper_loss": gripper_loss, + "left_arm_loss": left_arm_loss, + "right_arm_loss": right_arm_loss, + } + + # ---------- preprocess / postprocess ---------- + + def preprocess(self, proprio, action, mode="train"): + """ + - If proprio/action are 12-dim, pad them to 20 for the model. + - Zero-out gripper channels in proprio/action to focus learning on joints. + """ + proprio_m = self._pad_to_model_dim(proprio.clone()) + action_m = self._pad_to_model_dim(action.clone()) if action is not None else None + + proprio_m[..., self.gripper_idx] = 0.0 + if action_m is not None: + action_m[..., self.gripper_idx] = 0.0 + + return proprio_m, action_m + + def postprocess(self, action: torch.Tensor) -> torch.Tensor: + """ + - Model outputs [*, 20] + - Apply sigmoid to gripper logits + - Return only the first 12 dims for the real robot: + ["left_shoulder_pan.pos", + "left_shoulder_lift.pos", + "left_elbow_flex.pos", + "left_wrist_flex.pos", + "left_wrist_roll.pos", + "left_gripper.pos", + "right_shoulder_pan.pos", + "right_shoulder_lift.pos", + "right_elbow_flex.pos", + "right_wrist_flex.pos", + "right_wrist_roll.pos", + "right_gripper.pos"] + """ + # Ensure we at least have the real dims + grippers + if action.size(-1) < self.REAL_DIM: + raise ValueError(f"Expected at least {self.REAL_DIM} dims in action, got {action.size(-1)}") + + # Apply sigmoid on gripper channels in model space (indices 5 and 11) + if action.size(-1) > max(self.gripper_idx): + action[..., self.gripper_idx] = torch.sigmoid(action[..., self.gripper_idx]) + + # Return only the real 12-dim control vector for the env + return self._trim_to_real_dim(action) + + # ============================================================================= # Exports # ============================================================================= @@ -304,5 +515,6 @@ __all__ = [ "JointActionSpace", "AGIBOTEE6DActionSpace", "FrankaJoint7ActionSpace", + "BimanualSO101ActionSpace", "ACTION_REGISTRY", ] diff --git a/src/lerobot/policies/xvla/modeling_xvla.py b/src/lerobot/policies/xvla/modeling_xvla.py index 26454ab24..8ae0cb1ca 100644 --- a/src/lerobot/policies/xvla/modeling_xvla.py +++ b/src/lerobot/policies/xvla/modeling_xvla.py @@ -359,8 +359,8 @@ class XVLAPolicy(PreTrainedPolicy): - skip list for layers that should remain randomly initialized """ import safetensors.torch - # step 1: load config + #TODO: jadechoghari, fix this if config is None: config = PreTrainedConfig.from_pretrained( pretrained_name_or_path=pretrained_name_or_path, @@ -373,6 +373,7 @@ class XVLAPolicy(PreTrainedPolicy): revision=revision, **kwargs, ) + model_id = str(pretrained_name_or_path) instance = cls(config, **kwargs) diff --git a/src/lerobot/policies/xvla/processor_xvla.py b/src/lerobot/policies/xvla/processor_xvla.py index 388f04e95..374c9f4ac 100644 --- a/src/lerobot/policies/xvla/processor_xvla.py +++ b/src/lerobot/policies/xvla/processor_xvla.py @@ -68,6 +68,8 @@ def make_xvla_pre_post_processors( padding=config.pad_language_to, padding_side=config.tokenizer_padding_side, ), + XVLAImageToFloatProcessorStep(), + XVLAImageNetNormalizeProcessorStep(), DeviceProcessorStep(device=config.device), XVLAAddDomainIdProcessorStep(), NormalizerProcessorStep( @@ -266,6 +268,77 @@ class XVLAImageScaleProcessorStep(ProcessorStep): } +@dataclass +@ProcessorStepRegistry.register(name="xvla_image_to_float") +class XVLAImageToFloatProcessorStep(ProcessorStep): + """Convert image observations from [0, 255] to [0, 1] range. + + This processor step divides image observations by 255 to convert from uint8-like + range [0, 255] to float range [0, 1]. This is typically used when loading images + that are stored as uint8 values. + + Args: + image_keys: List of observation keys that contain images to convert. + If None, will automatically detect keys starting with "observation.images." + validate_range: If True, validates that input values are in [0, 255] range (default: True) + + Raises: + ValueError: If validate_range is True and image values are not in [0, 255] range. + """ + + image_keys: list[str] | None = None + validate_range: bool = True + + def __call__(self, transition: EnvTransition) -> EnvTransition: + """Convert image observations from [0, 255] to [0, 1].""" + new_transition = transition.copy() + obs = new_transition.get(TransitionKey.OBSERVATION, {}) + if obs is None: + return new_transition + + # Make a copy of observations to avoid modifying the original + obs = obs.copy() + + # Determine which keys to convert + keys_to_convert = self.image_keys + if keys_to_convert is None: + # Auto-detect image keys + keys_to_convert = [k for k in obs if k.startswith("observation.images.")] + + # Convert each image + for key in keys_to_convert: + if key in obs and isinstance(obs[key], torch.Tensor): + tensor = obs[key] + + # Validate that values are in [0, 255] range if requested + if self.validate_range: + min_val = tensor.min().item() + max_val = tensor.max().item() + if min_val < 0.0 or max_val > 255.0: + raise ValueError( + f"Image '{key}' has values outside [0, 255] range: " + f"min={min_val:.4f}, max={max_val:.4f}. " + f"Cannot convert to [0, 1] range." + ) + + # Convert to float and divide by 255 + obs[key] = tensor.float() / 255.0 + + new_transition[TransitionKey.OBSERVATION] = obs + return new_transition + + def transform_features(self, features): + """Image conversion doesn't change feature structure.""" + return features + + def get_config(self) -> dict[str, Any]: + """Return serializable configuration.""" + return { + "image_keys": self.image_keys, + "validate_range": self.validate_range, + } + + @dataclass @ProcessorStepRegistry.register(name="xvla_imagenet_normalize") class XVLAImageNetNormalizeProcessorStep(ProcessorStep): diff --git a/src/lerobot/scripts/lerobot_eval.py b/src/lerobot/scripts/lerobot_eval.py index 65b6cb144..db93dbe0d 100644 --- a/src/lerobot/scripts/lerobot_eval.py +++ b/src/lerobot/scripts/lerobot_eval.py @@ -510,7 +510,6 @@ def eval_main(cfg: EvalPipelineConfig): envs = make_env(cfg.env, n_envs=cfg.eval.batch_size, use_async_envs=cfg.eval.use_async_envs) logging.info("Making policy.") - policy = make_policy( cfg=cfg.policy, env_cfg=cfg.env, diff --git a/src/lerobot/scripts/lerobot_train.py b/src/lerobot/scripts/lerobot_train.py index 7e02f879f..fd200a254 100644 --- a/src/lerobot/scripts/lerobot_train.py +++ b/src/lerobot/scripts/lerobot_train.py @@ -200,12 +200,12 @@ def train(cfg: TrainPipelineConfig, accelerator: Accelerator | None = None): if is_main_process: logging.info("Creating policy") + policy = make_policy( cfg=cfg.policy, ds_meta=dataset.meta, rename_map=cfg.rename_map, ) - # Wait for all processes to finish policy creation before continuing accelerator.wait_for_everyone()