from wall_x.infer.infer_config import InferConfig from typing import Optional, List from dataclasses import dataclass, field import numpy as np import torch import wall_x.infer.data_utils as data_utils from wall_x.infer.logger import InferLogger dof_dims = { "left_ee_cartesian_pos": 3, "left_ee_cartesian_pos_relative": 3, "left_ee_rotation": 3, "left_ee_rotation_relative": 3, "left_ee_rotation_6D": 6, "left_ee_rotation_6D_relative": 6, "left_arm_joint_pos": 7, "left_gripper": 1, "left_gripper_cur": 1, "left_arm_joint_cur": 1, "right_ee_cartesian_pos": 3, "right_ee_cartesian_pos_relative": 3, "right_ee_rotation": 3, "right_ee_rotation_relative": 3, "right_ee_rotation_6D": 6, "right_ee_rotation_6D_relative": 6, "right_arm_joint_pos": 7, "right_gripper": 1, "right_gripper_cur": 1, "right_arm_joint_cur": 1, "head_actions": 2, "height": 1, "car_pose": 3, "velocity_decomposed": 3, } class ComputedDict(dict): """Smart dictionary that supports registering computation rules and auto-computes None values on get""" def __init__(self, *args, **kwargs): super().__init__(*args, **kwargs) self._compute_rules = {} # key -> compute_function def register_compute_rule(self, key, compute_func): """ Register a computation rule. Args: key: The key that needs computation compute_func: Computation function that takes self as argument and returns the computed result """ self._compute_rules[key] = compute_func def get(self, key, default=None): """Override get method to support auto-computation""" value = super().get(key, default) # If value is None and there's a compute rule, try to compute if value is None and key in self._compute_rules: try: computed_value = self._compute_rules[key](self) if computed_value is not None: # Cache the computed result self[key] = computed_value return computed_value except Exception: pass # If computation fails, return None or default return value if value is not None else default def __getitem__(self, key): """Override [] operator to support auto-computation""" value = super().get(key, None) # If value is None and there's a compute rule, try to compute if value is None and key in self._compute_rules: try: computed_value = self._compute_rules[key](self) if computed_value is not None: # Cache the computed result self[key] = computed_value return computed_value except Exception: pass # If computation fails, raise original KeyError or return None if key in self: return super().__getitem__(key) raise KeyError(key) @dataclass class RobotStateActionData: config: InferConfig = None data: ComputedDict = field( default_factory=lambda: ComputedDict( { # State (formerly pose) - using state_ prefix "state_left_ee_cartesian_pos": None, # (1, 3) "state_left_ee_rotation": None, # (1, 3) "state_left_ee_rotation_6D": None, "state_left_arm_joint_pos": None, "state_left_gripper": None, # (1, 1) "state_left_arm_joint_cur": None, "state_left_gripper_cur": None, "state_right_ee_cartesian_pos": None, # (1, 3) "state_right_ee_rotation": None, "state_right_ee_rotation_6D": None, # (1, 6) "state_right_arm_joint_pos": None, "state_right_gripper": None, "state_right_gripper_cur": None, "state_right_arm_joint_cur": None, # (1, 1) "state_head_actions": None, "state_height": None, "state_car_pose": None, "state_velocity_decomposed": None, # Action - using action_ prefix "action_left_ee_cartesian_pos": None, "action_left_ee_cartesian_pos_relative": None, "action_left_ee_rotation": None, "action_left_ee_rotation_relative": None, "action_left_ee_rotation_6D": None, "action_left_ee_rotation_6D_relative": None, "action_left_gripper": None, "action_left_arm_joint_pos": None, "action_right_ee_cartesian_pos": None, "action_right_ee_cartesian_pos_relative": None, "action_right_ee_rotation": None, "action_right_ee_rotation_relative": None, "action_right_ee_rotation_6D": None, "action_right_ee_rotation_6D_relative": None, "action_right_gripper": None, "action_right_arm_joint_pos": None, "action_head_actions": None, "action_height": None, "action_car_pose": None, "action_velocity_decomposed": None, } ) ) dof_mask: np.ndarray = None logger = InferLogger.get_robot_logger("RobotStateActionData") def __post_init__(self): """Register computation rules""" # State computation rules - euler angles -> 6D rotation self.data.register_compute_rule( "state_left_ee_rotation_6D", lambda d: ( data_utils.euler_to_matrix_zyx_6d_nb(d["state_left_ee_rotation"]) if d.get("state_left_ee_rotation") is not None else None ), ) self.data.register_compute_rule( "state_right_ee_rotation_6D", lambda d: ( data_utils.euler_to_matrix_zyx_6d_nb(d["state_right_ee_rotation"]) if d.get("state_right_ee_rotation") is not None else None ), ) # Action computation rules - absolute position computed from relative + state self.data.register_compute_rule( "action_left_ee_cartesian_pos", lambda d: ( d.get("state_left_ee_cartesian_pos") + d.get("action_left_ee_cartesian_pos_relative") if d.get("state_left_ee_cartesian_pos") is not None and d.get("action_left_ee_cartesian_pos_relative") is not None else None ), ) self.data.register_compute_rule( "action_right_ee_cartesian_pos", lambda d: ( d.get("state_right_ee_cartesian_pos") + d.get("action_right_ee_cartesian_pos_relative") if d.get("state_right_ee_cartesian_pos") is not None and d.get("action_right_ee_cartesian_pos_relative") is not None else None ), ) # Action computation rules - get absolute rpy self.data.register_compute_rule( # delta rpy -> abs rpy "action_left_ee_rotation", lambda d: ( data_utils.compose_state_and_delta_to_abs_rpy( d["action_left_ee_rotation_relative"], d["state_left_ee_rotation"][0], ) if d.get("action_left_ee_rotation_relative") is not None and d.get("state_left_ee_rotation") is not None else None ), ) self.data.register_compute_rule( # abs 6D -> abs rpy "action_left_ee_rotation", lambda d: ( data_utils.so3_to_euler_zyx_batch_nb(d["action_left_ee_rotation_6D"]) if d.get("action_left_ee_rotation_6D") is not None else None ), ) self.data.register_compute_rule( # delta 6D -> abs 6D -> abs rpy "action_left_ee_rotation_6D", lambda d: ( data_utils.compose_state_and_delta_to_abs_rpy( d["action_left_ee_rotation_6D_relative"], d["state_left_ee_rotation_6D"][0], ) if d.get("action_left_ee_rotation_6D_relative") is not None and d.get("state_left_ee_rotation_6D") is not None else None ), ) self.data.register_compute_rule( # delta rpy -> abs rpy "action_right_ee_rotation", lambda d: ( data_utils.compose_state_and_delta_to_abs_rpy( d["action_right_ee_rotation_relative"], d["state_right_ee_rotation"][0], ) if d.get("action_right_ee_rotation_relative") is not None and d.get("state_right_ee_rotation") is not None else None ), ) self.data.register_compute_rule( # abs 6D -> abs rpy "action_right_ee_rotation", lambda d: ( data_utils.so3_to_euler_zyx_batch_nb(d["action_right_ee_rotation_6D"]) if d.get("action_right_ee_rotation_6D") is not None else None ), ) self.data.register_compute_rule( # delta 6D -> abs 6D -> abs rpy "action_right_ee_rotation_6D", lambda d: ( data_utils.compose_state_and_delta_to_abs_rpy( d["action_right_ee_rotation_6D_relative"], d["state_right_ee_rotation_6D"][0], ) if d.get("action_right_ee_rotation_6D_relative") is not None and d.get("state_right_ee_rotation_6D") is not None else None ), ) def get_agent_pos(self, obs_action_keys=None): if obs_action_keys is None: obs_action_keys = self.config.train_config["data"]["obs_action_keys"] agent_pose_data = [] for key in obs_action_keys: # Remove follow_ or master_ prefix if key.startswith("follow_"): key = key.replace("follow_", "") elif key.startswith("master_"): key = key.replace("master_", "") # Add state_ prefix to access state data state_key = f"state_{key}" if state_key in self.data: # Use get method, which will auto-handle None value computation value = self.data.get(state_key) if value is None: # If still None after computation, use zero vector agent_pose_data.append(np.zeros((1, dof_dims[key]))) else: agent_pose_data.append(value) else: raise ValueError(f"Key {state_key} not found in data") agent_pose_data = np.concatenate(agent_pose_data, axis=1)[None] # (1, 1, D) return agent_pose_data def get_agent_pos_mask(self, obs_action_keys=None): if obs_action_keys is None: obs_action_keys = self.config.train_config["data"]["obs_action_keys"] agent_pos_mask_data = [] for key in obs_action_keys: # Remove follow_ or master_ prefix if key.startswith("follow_"): key = key.replace("follow_", "") elif key.startswith("master_"): key = key.replace("master_", "") # Add state_ prefix to access state data state_key = f"state_{key}" if state_key in self.data: # Use get method, which will auto-handle None value computation value = self.data.get(state_key) if value is None: agent_pos_mask_data.append(np.zeros((1, dof_dims[key]))) else: agent_pos_mask_data.append(np.ones((1, dof_dims[key]))) else: raise ValueError(f"Key {state_key} not found in data") return np.concatenate(agent_pos_mask_data, axis=1)[None] # (1, 1, D) def save_state_data_with_key(self, value, key): # Remove follow_ or master_ prefix key = key.replace("follow_", "") key = key.replace("master_", "") # if torch, convert to numpy if isinstance(value, torch.Tensor): value = value.detach().cpu().numpy() if f"state_{key}" not in self.data: # TODO: joint angle control self.logger.warning(f"{key} is not a valid state key, not recorded") return # Shape validation for value, expected shape is (1, D) if value.shape == (1, dof_dims[key]): self.data[f"state_{key}"] = value elif value.shape == (1, 1, dof_dims[key]): self.data[f"state_{key}"] = value[0] elif value.shape == (dof_dims[key],): self.data[f"state_{key}"] = value[None] else: raise ValueError(f"Value shape {value.shape} is not legal") def save_action_data_with_key(self, value, key): key = key.replace("follow_", "") key = key.replace("master_", "") if isinstance(value, torch.Tensor): value = value.detach().cpu().numpy() if value.shape == (dof_dims[key],): self.data[f"action_{key}"] = value[None] else: self.data[f"action_{key}"] = value def save_action_data( self, predict_action, predict_action_keys: Optional[List[str]] = None ): if predict_action_keys is None: predict_action_keys = self.config.data_config["predict_action_keys"] if isinstance(predict_action, torch.Tensor): predict_action = predict_action.detach().cpu().numpy() if predict_action.ndim == 3: predict_action = predict_action[0] dof_start = 0 for action_key in predict_action_keys: action_key = action_key.replace("follow_", "") action_key = action_key.replace("master_", "") dof_dim = dof_dims[action_key] action_key = f"action_{action_key}" self.data[action_key] = predict_action[:, dof_start : dof_start + dof_dim] dof_start += dof_dim # For compatibility, provide convenient property access @property def agent_pos(self): return self.get_agent_pos() @property def agent_pos_mask(self): return self.get_agent_pos_mask() @property def action(self): pass # TODO: support action access