"""Pinocchio-based robot model with FK, joint limits, gravity compensation, and reduced-DOF views. RobotModel wraps a URDF + optional supplemental info to provide forward kinematics, Jacobians, gravity torques, and joint-group queries. ReducedRobotModel maps between a full configuration and an actuated-joint subset. """ from typing import List, Optional, Set, Union import numpy as np import pinocchio as pin from gear_sonic.data.robot_model.supplemental_info import RobotSupplementalInfo class RobotModel: def __init__( self, urdf_path, asset_path, set_floating_base=False, supplemental_info: Optional[RobotSupplementalInfo] = None, ): self.pinocchio_wrapper = pin.RobotWrapper.BuildFromURDF( filename=urdf_path, package_dirs=[asset_path], root_joint=pin.JointModelFreeFlyer() if set_floating_base else None, ) self.is_floating_base_model = set_floating_base self.joint_to_dof_index = {} # Assume we only have single-dof joints # First two names correspond to universe and floating base joints names = ( self.pinocchio_wrapper.model.names[2:] if set_floating_base else self.pinocchio_wrapper.model.names[1:] ) for name in names: j_id = self.pinocchio_wrapper.model.getJointId(name) jmodel = self.pinocchio_wrapper.model.joints[j_id] self.joint_to_dof_index[name] = jmodel.idx_q # Store joint limits only for actual joints (excluding floating base) # if set floating base is true and the robot can move in the world # then we don't want to impose joint limits for the 7 dofs corresponding # to the floating base dofs. root_nq = 7 if set_floating_base else 0 self.upper_joint_limits = self.pinocchio_wrapper.model.upperPositionLimit[root_nq:].copy() self.lower_joint_limits = self.pinocchio_wrapper.model.lowerPositionLimit[root_nq:].copy() # Set up supplemental info if provided self.supplemental_info = supplemental_info if self.supplemental_info is not None: # Cache indices for body and hand actuated joints separately self._body_actuated_joint_indices = [ self.dof_index(name) for name in self.supplemental_info.body_actuated_joints ] self._left_hand_actuated_joint_indices = [ self.dof_index(name) for name in self.supplemental_info.left_hand_actuated_joints ] self._right_hand_actuated_joint_indices = [ self.dof_index(name) for name in self.supplemental_info.right_hand_actuated_joints ] self._hand_actuated_joint_indices = ( self._left_hand_actuated_joint_indices + self._right_hand_actuated_joint_indices ) # Cache indices for joint groups, handling nested groups self._joint_group_indices = {} for group_name, group_info in self.supplemental_info.joint_groups.items(): indices = [] # Add indices for direct joints indices.extend([self.dof_index(name) for name in group_info["joints"]]) # Add indices from subgroups for subgroup_name in group_info["groups"]: indices.extend(self.get_joint_group_indices(subgroup_name)) self._joint_group_indices[group_name] = sorted(set(indices)) # Update joint limits from supplemental info if available if ( hasattr(self.supplemental_info, "joint_limits") and self.supplemental_info.joint_limits ): for joint_name, limits in self.supplemental_info.joint_limits.items(): if joint_name in self.joint_to_dof_index: # joint_to_dof_index is in full-space (includes floating base DOFs), # but limits arrays are indexed from 0 starting at the first real joint idx = self.joint_to_dof_index[joint_name] - root_nq self.lower_joint_limits[idx] = limits[0] self.upper_joint_limits[idx] = limits[1] # Initialize default body pose self.default_body_pose = self.q_zero # Update with supplemental info if available if self.supplemental_info is not None: default_joint_q = self.supplemental_info.default_joint_q for joint, joint_values in default_joint_q.items(): # Get the joint name mapping for this type joint_mapping = self.supplemental_info.joint_name_mapping[joint] # Handle both single joint names and left/right mappings if isinstance(joint_mapping, str): # Single joint (e.g., waist joints) if joint_mapping in self.joint_to_dof_index: joint_idx = self.dof_index(joint_mapping) self.default_body_pose[joint_idx] = ( joint_values # joint_values is the value for single joints ) else: # Left/right mapping (e.g., arm joints) for side, value in joint_values.items(): if side in joint_mapping and joint_mapping[side] in self.joint_to_dof_index: joint_idx = self.dof_index(joint_mapping[side]) self.default_body_pose[joint_idx] = value # Initialize initial body pose self.initial_body_pose = self.default_body_pose.copy() @property def num_dofs(self) -> int: """Get the number of degrees of freedom of the robot (floating base pose + joints).""" return self.pinocchio_wrapper.model.nq @property def q_zero(self) -> np.ndarray: """Get the zero pose of the robot.""" return self.pinocchio_wrapper.q0.copy() @property def joint_names(self) -> List[str]: """Get the names of the joints of the robot.""" return list(self.joint_to_dof_index.keys()) @property def num_joints(self) -> int: """Get the number of joints of the robot.""" return len(self.joint_to_dof_index) def dof_index(self, joint_name: str) -> int: """ Get the index in the degrees of freedom vector corresponding to the single-DoF joint with name `joint_name`. """ if joint_name not in self.joint_to_dof_index: raise ValueError( f"Unknown joint name: '{joint_name}'. " f"Available joints: {list(self.joint_to_dof_index.keys())}" ) return self.joint_to_dof_index[joint_name] def get_body_actuated_joint_indices(self) -> List[int]: """ Get the indices of body actuated joints in the full configuration. Ordering is that of the actuated joints as defined in the supplemental info. Requires supplemental_info to be provided. """ if self.supplemental_info is None: raise ValueError("supplemental_info must be provided to use this method") return self._body_actuated_joint_indices def get_hand_actuated_joint_indices(self, side: str = "both") -> List[int]: """ Get the indices of hand actuated joints in the full configuration. Ordering is that of the actuated joints as defined in the supplemental info. Requires supplemental_info to be provided. Args: side: String specifying which hand to get indices for ('left', 'right', or 'both') """ if self.supplemental_info is None: raise ValueError("supplemental_info must be provided to use this method") if side.lower() == "both": return self._hand_actuated_joint_indices elif side.lower() == "left": return self._left_hand_actuated_joint_indices elif side.lower() == "right": return self._right_hand_actuated_joint_indices else: raise ValueError("side must be 'left', 'right', or 'both'") def get_joint_group_indices(self, group_names: Union[str, Set[str]]) -> List[int]: """ Get the indices of joints in one or more groups in the full configuration. Requires supplemental_info to be provided. The returned indices are sorted in ascending order, so that the joint ordering of the full model is preserved. Args: group_names: Either a single group name (str) or a set of group names (Set[str]) Returns: List of joint indices in sorted order with no duplicates """ if self.supplemental_info is None: raise ValueError("supplemental_info must be provided to use this method") # Convert single string to set for uniform handling if isinstance(group_names, str): group_names = {group_names} # Collect indices from all groups all_indices = set() for group_name in group_names: if group_name not in self._joint_group_indices: raise ValueError(f"Unknown joint group: {group_name}") all_indices.update(self._joint_group_indices[group_name]) return sorted(all_indices) def cache_forward_kinematics(self, q: np.ndarray, auto_clip=True) -> None: """ Perform forward kinematics to update the pose of every joint and frame in the Pinocchio data structures for the given configuration `q`. :param q: A numpy array of shape (num_dofs,) representing the robot configuration. """ if q.shape[0] != self.num_dofs: raise ValueError(f"Expected q of length {self.num_dofs}, got {q.shape[0]} instead.") # Apply auto-clip if enabled if auto_clip: q = self.clip_configuration(q) pin.framesForwardKinematics(self.pinocchio_wrapper.model, self.pinocchio_wrapper.data, q) def compute_gravity_compensation_torques( self, q: np.ndarray, joint_groups: Union[str, List[str], Set[str]] = None, auto_clip=True ) -> np.ndarray: """ Compute gravity compensation torques for specified joint groups using pinocchio. :param q: Robot configuration (joint positions) :param joint_groups: Joint groups to compensate (e.g., "arms", ["left_arm", "waist"], {"left_arm", "waist"}). If None, compensates all joints :param auto_clip: Whether to automatically clip joint values to limits :return: Array of gravity compensation torques for all DOFs (zero for non-compensated joints) """ if q.shape[0] != self.num_dofs: raise ValueError(f"Expected q of length {self.num_dofs}, got {q.shape[0]} instead.") # Apply auto-clip if enabled if auto_clip: q = self.clip_configuration(q) try: # Cache forward kinematics for the current configuration self.cache_forward_kinematics(q, auto_clip=False) # Already clipped if needed # RNEA with zero velocity and acceleration isolates the gravity term: # tau = M(q)*0 + C(q,0)*0 + g(q) = g(q), i.e. pure gravity compensation v = np.zeros(self.num_dofs) a = np.zeros(self.num_dofs) gravity_torques_full = pin.rnea( self.pinocchio_wrapper.model, self.pinocchio_wrapper.data, q, v, a ) # If no joint groups specified, return full gravity torques if joint_groups is None: return gravity_torques_full # Convert list to set for get_joint_group_indices compatibility if isinstance(joint_groups, list): joint_groups = set(joint_groups) # Get joint indices for specified groups - get_joint_group_indices handles str and Set[str] try: compensated_joint_indices = self.get_joint_group_indices(joint_groups) except ValueError as e: raise ValueError(f"Error resolving joint groups {joint_groups}: {e}") # Create mask for joints that should receive gravity compensation compensation_mask = np.zeros(self.num_dofs, dtype=bool) for joint_idx in compensated_joint_indices: if 0 <= joint_idx < len(compensation_mask): compensation_mask[joint_idx] = True # Apply mask to only compensate specified joints compensated_torques = np.zeros_like(gravity_torques_full) compensated_torques[compensation_mask] = gravity_torques_full[compensation_mask] return compensated_torques except Exception as e: raise RuntimeError(f"Error computing gravity compensation: {e}") def clip_configuration(self, q: np.ndarray, margin: float = 1e-6) -> np.ndarray: """ Clip the configuration to stay within joint limits with a small tolerance. :param q: Configuration to clip :param margin: Tolerance to keep away from joint limits :return: Clipped configuration """ q_clipped = q.copy() # Only clip joint positions, not floating base root_nq = 7 if self.is_floating_base_model else 0 q_clipped[root_nq:] = np.clip( q[root_nq:], self.lower_joint_limits + margin, self.upper_joint_limits - margin ) return q_clipped def frame_placement(self, frame_name: str) -> pin.SE3: """ Returns the SE3 transform of the specified frame in the world coordinate system. Note: make sure cache_forward_kinematics() has been previously called. :param frame_name: Name of the frame, e.g. "link_elbow_frame", "hand_imu_frame", etc. :return: A pin.SE3 object representing the pose of the frame. """ model = self.pinocchio_wrapper.model data = self.pinocchio_wrapper.data frame_id = model.getFrameId(frame_name) if frame_id < 0 or frame_id >= len(model.frames): valid_frames = [f.name for f in model.frames] raise ValueError(f"Unknown frame '{frame_name}'. Valid frames: {valid_frames}") # Pinocchio's data.oMf[frame_id] is a pin.SE3. return data.oMf[frame_id].copy() def frame_jacobian( self, frame_name: str, q: np.ndarray, reference_frame: pin.ReferenceFrame = pin.LOCAL_WORLD_ALIGNED, ) -> np.ndarray: """ Compute the Jacobian of the specified frame. :param frame_name: Name of the frame, e.g. "fingertip_frame", "hand_frame", etc. :param q: Configuration vector (joint positions) :param reference_frame: Reference frame for the Jacobian. Options: - pin.LOCAL: Jacobian expressed in the local frame - pin.WORLD: Jacobian expressed in the world frame - pin.LOCAL_WORLD_ALIGNED: Local frame with world orientation (default, best for IK) :return: A 6xN Jacobian matrix where N is the number of DOFs. First 3 rows are linear velocity, last 3 rows are angular velocity. """ model = self.pinocchio_wrapper.model data = self.pinocchio_wrapper.data # Get frame ID frame_id = model.getFrameId(frame_name) if frame_id < 0 or frame_id >= len(model.frames): valid_frames = [f.name for f in model.frames] raise ValueError(f"Unknown frame '{frame_name}'. Valid frames: {valid_frames}") # Compute the Jacobian J = pin.computeFrameJacobian(model, data, q, frame_id, reference_frame) return J.copy() def get_body_actuated_joints(self, q: np.ndarray) -> np.ndarray: """ Get the configuration of body actuated joints from a full configuration. :param q: Configuration in full space :return: Configuration of body actuated joints """ indices = self.get_body_actuated_joint_indices() return q[indices] def get_hand_actuated_joints(self, q: np.ndarray, side: str = "both") -> np.ndarray: """ Get the configuration of hand actuated joints from a full configuration. Args: q: Configuration in full space side: String specifying which hand to get joints for ('left', 'right', or 'both') """ indices = self.get_hand_actuated_joint_indices(side) return q[indices] def get_configuration_from_actuated_joints( self, body_actuated_joint_values: np.ndarray, hand_actuated_joint_values: Optional[np.ndarray] = None, left_hand_actuated_joint_values: Optional[np.ndarray] = None, right_hand_actuated_joint_values: Optional[np.ndarray] = None, ) -> np.ndarray: """ Get the full configuration from the body and hand actuated joint configurations. Can specify either both hands together or left and right hands separately. Args: body_actuated_joint_values: Configuration of body actuated joints hand_actuated_joint_values: Configuration of both hands' actuated joints (optional) left_hand_actuated_joint_values: Configuration of left hand actuated joints (optional) right_hand_actuated_joint_values: Configuration of right hand actuated joints (optional) Returns: Full configuration including body and hand joints """ q = self.pinocchio_wrapper.q0.copy() q[self.get_body_actuated_joint_indices()] = body_actuated_joint_values # Handle hand configurations if hand_actuated_joint_values is not None: # Use combined hand configuration q[self.get_hand_actuated_joint_indices("both")] = hand_actuated_joint_values else: # Use separate hand configurations if left_hand_actuated_joint_values is not None: q[self.get_hand_actuated_joint_indices("left")] = left_hand_actuated_joint_values if right_hand_actuated_joint_values is not None: q[self.get_hand_actuated_joint_indices("right")] = right_hand_actuated_joint_values return q def reset_forward_kinematics(self) -> None: """ Reset the forward kinematics to the initial configuration. """ self.cache_forward_kinematics(self.q_zero) def get_initial_upper_body_pose(self) -> np.ndarray: """ Get the initial upper body pose of the robot. """ return self.initial_body_pose[self.get_joint_group_indices("upper_body")] def get_default_body_pose(self) -> np.ndarray: """ Get the default body pose of the robot. """ return self.default_body_pose.copy() def set_initial_body_pose(self, q: np.ndarray, q_idx=None) -> None: """ Set the initial body pose of the robot. """ if q_idx is None: self.initial_body_pose = q else: self.initial_body_pose[q_idx] = q class ReducedRobotModel(RobotModel): """ A class that creates a reduced order robot model by fixing certain joints. This class maintains a mapping between the reduced state space and the full state space. """ def __init__( self, full_robot_model: RobotModel, fixed_joints: List[str], fixed_values: Optional[List[float]] = None, ): """ Create a reduced order robot model by fixing specified joints. :param full_robot_model: The original robot model :param fixed_joints: List of joint names to fix :param fixed_values: Optional list of values to fix the joints to. If None, uses the initial joint positions (q0) from the full robot model. """ self.full_robot = full_robot_model self.supplemental_info = full_robot_model.supplemental_info # If fixed_values is None, use q0 from the full robot model if fixed_values is None: fixed_values = [] for joint_name in fixed_joints: full_idx = full_robot_model.dof_index(joint_name) fixed_values.append(full_robot_model.pinocchio_wrapper.q0[full_idx]) elif len(fixed_joints) != len(fixed_values): raise ValueError("fixed_joints and fixed_values must have the same length") # Store fixed joints and their values self.fixed_joints = fixed_joints self.fixed_values = fixed_values # reduced_to_full[i] = full-space index of the i-th reduced-space DOF # full_to_reduced[j] = reduced-space index of the j-th full-space DOF (active joints only) self.reduced_to_full = [] self.full_to_reduced = {} # Initialize with floating base indices if present if full_robot_model.is_floating_base_model: self.reduced_to_full.extend(range(7)) # Floating base indices for i in range(7): self.full_to_reduced[i] = i # Add active joint indices for joint_name in full_robot_model.joint_names: if joint_name not in fixed_joints: full_idx = full_robot_model.dof_index(joint_name) reduced_idx = len(self.reduced_to_full) self.reduced_to_full.append(full_idx) self.full_to_reduced[full_idx] = reduced_idx # Create a reduced Pinocchio model using buildReducedModel # First, get the list of joint IDs to lock locked_joint_ids = [] for joint_name in fixed_joints: joint_id = full_robot_model.pinocchio_wrapper.model.getJointId(joint_name) # Pinocchio reserves id=0 for "universe" and id=1 for floating base (if present). # Only lock actual robot joints, not these special entries. if (full_robot_model.is_floating_base_model and joint_id > 1) or ( not full_robot_model.is_floating_base_model and joint_id > 0 ): locked_joint_ids.append(joint_id) # First build the reduced kinematic model reduced_model = pin.buildReducedModel( full_robot_model.pinocchio_wrapper.model, locked_joint_ids, full_robot_model.pinocchio_wrapper.q0, ) # Then build the reduced geometry models using the reduced kinematic model self.pinocchio_wrapper = pin.RobotWrapper( model=reduced_model, ) # Create joint to dof index mapping self.joint_to_dof_index = {} # Assume we only have single-dof joints # First two names correspond to universe and floating base joints names = ( self.pinocchio_wrapper.model.names[2:] if self.full_robot.is_floating_base_model else self.pinocchio_wrapper.model.names[1:] ) for name in names: j_id = self.pinocchio_wrapper.model.getJointId(name) jmodel = self.pinocchio_wrapper.model.joints[j_id] self.joint_to_dof_index[name] = jmodel.idx_q # Initialize joint limits root_nq = 7 if self.full_robot.is_floating_base_model else 0 self.lower_joint_limits = self.pinocchio_wrapper.model.lowerPositionLimit[root_nq:].copy() self.upper_joint_limits = self.pinocchio_wrapper.model.upperPositionLimit[root_nq:].copy() # Update joint limits from supplemental info if available if self.supplemental_info is not None: if ( hasattr(self.supplemental_info, "joint_limits") and self.supplemental_info.joint_limits ): for joint_name, limits in self.supplemental_info.joint_limits.items(): if joint_name in self.joint_to_dof_index: idx = self.joint_to_dof_index[joint_name] - root_nq self.lower_joint_limits[idx] = limits[0] self.upper_joint_limits[idx] = limits[1] # Get full indices for body and hand actuated joints full_body_indices = full_robot_model.get_body_actuated_joint_indices() full_hand_indices = full_robot_model.get_hand_actuated_joint_indices("both") full_left_hand_indices = full_robot_model.get_hand_actuated_joint_indices("left") full_right_hand_indices = full_robot_model.get_hand_actuated_joint_indices("right") # Map to reduced indices self._body_actuated_joint_indices = [] for idx in full_body_indices: if idx in self.full_to_reduced: self._body_actuated_joint_indices.append(self.full_to_reduced[idx]) self._hand_actuated_joint_indices = [] for idx in full_hand_indices: if idx in self.full_to_reduced: self._hand_actuated_joint_indices.append(self.full_to_reduced[idx]) self._left_hand_actuated_joint_indices = [] for idx in full_left_hand_indices: if idx in self.full_to_reduced: self._left_hand_actuated_joint_indices.append(self.full_to_reduced[idx]) self._right_hand_actuated_joint_indices = [] for idx in full_right_hand_indices: if idx in self.full_to_reduced: self._right_hand_actuated_joint_indices.append(self.full_to_reduced[idx]) # Cache indices for joint groups in reduced space self._joint_group_indices = {} for group_name in self.supplemental_info.joint_groups: full_indices = full_robot_model.get_joint_group_indices(group_name) reduced_indices = [] for idx in full_indices: if idx in self.full_to_reduced: reduced_indices.append(self.full_to_reduced[idx]) self._joint_group_indices[group_name] = sorted(set(reduced_indices)) # Initialize default body pose in reduced space self.default_body_pose = self.full_to_reduced_configuration( full_robot_model.default_body_pose ) # Initialize initial body pose in reduced space self.initial_body_pose = self.full_to_reduced_configuration( full_robot_model.initial_body_pose ) @property def num_joints(self) -> int: """Get the number of active joints in the reduced model.""" return len(self.joint_names) @property def joint_names(self) -> List[str]: """Get the names of the active joints in the reduced model.""" return [name for name in self.full_robot.joint_names if name not in self.fixed_joints] @classmethod def from_fixed_groups( cls, full_robot_model: RobotModel, fixed_group_names: List[str], fixed_values: Optional[List[float]] = None, ) -> "ReducedRobotModel": """ Create a reduced order robot model by fixing all joints in specified groups. :param full_robot_model: The original robot model :param fixed_group_names: List of joint group names to fix :param fixed_values: Optional list of values to fix the joints to. If None, uses the initial joint positions (q0) from the full robot model. :return: A ReducedRobotModel instance """ if full_robot_model.supplemental_info is None: raise ValueError("supplemental_info must be provided to use this method") # Get all joints in the groups, including those from subgroups fixed_joints = set() # Use a set to avoid duplicates for group_name in fixed_group_names: if group_name not in full_robot_model.supplemental_info.joint_groups: raise ValueError(f"Unknown joint group: {group_name}") group_info = full_robot_model.supplemental_info.joint_groups[group_name] # Add direct joints fixed_joints.update(group_info["joints"]) # Add joints from subgroups for subgroup_name in group_info["groups"]: subgroup_joints = full_robot_model.get_joint_group_indices(subgroup_name) fixed_joints.update([full_robot_model.joint_names[idx] for idx in subgroup_joints]) # Convert set back to list for compatibility with the original constructor return cls(full_robot_model, list(fixed_joints), fixed_values) @classmethod def from_fixed_group( cls, full_robot_model: RobotModel, fixed_group_name: str, fixed_values: Optional[List[float]] = None, ) -> "ReducedRobotModel": """ Create a reduced order robot model by fixing all joints in a specified group. This is a convenience method that calls from_fixed_groups with a single group. :param full_robot_model: The original robot model :param fixed_group_name: Name of the joint group to fix :param fixed_values: Optional list of values to fix the joints to. If None, uses the initial joint positions (q0) from the full robot model. :return: A ReducedRobotModel instance """ return cls.from_fixed_groups(full_robot_model, [fixed_group_name], fixed_values) @classmethod def from_active_group( cls, full_robot_model: RobotModel, active_group_name: str, fixed_values: Optional[List[float]] = None, ) -> "ReducedRobotModel": """ Create a reduced order robot model by fixing all joints EXCEPT those in the specified group. This is a convenience method that calls from_active_groups with a single group. :param full_robot_model: The original robot model :param active_group_name: Name of the joint group to keep active (all other joints will be fixed) :param fixed_values: Optional list of values to fix the joints to. If None, uses the initial joint positions (q0) from the full robot model. :return: A ReducedRobotModel instance """ return cls.from_active_groups(full_robot_model, [active_group_name], fixed_values) @classmethod def from_active_groups( cls, full_robot_model: RobotModel, active_group_names: List[str], fixed_values: Optional[List[float]] = None, ) -> "ReducedRobotModel": """ Create a reduced order robot model by fixing all joints EXCEPT those in the specified groups. This is useful when you want to keep multiple groups active and fix everything else. :param full_robot_model: The original robot model :param active_group_names: List of joint group names to keep active (all other joints will be fixed) :param fixed_values: Optional list of values to fix the joints to. If None, uses the initial joint positions (q0) from the full robot model. :return: A ReducedRobotModel instance """ if full_robot_model.supplemental_info is None: raise ValueError("supplemental_info must be provided to use this method") # Get all joints in the active groups, including those from subgroups active_joints = set() def add_group_joints(group_name: str): if group_name not in full_robot_model.supplemental_info.joint_groups: raise ValueError(f"Unknown joint group: {group_name}") group_info = full_robot_model.supplemental_info.joint_groups[group_name] # Add direct joints if "joints" in group_info: active_joints.update(group_info["joints"]) # Add joints from subgroups if "groups" in group_info: for subgroup_name in group_info["groups"]: add_group_joints(subgroup_name) for group_name in active_group_names: add_group_joints(group_name) # Get all joints from the model all_joints = set(full_robot_model.joint_names) # The fixed joints are all joints minus the active joints fixed_joints = list(all_joints - active_joints) return cls(full_robot_model, fixed_joints, fixed_values) def reduced_to_full_configuration(self, q_reduced: np.ndarray) -> np.ndarray: """ Convert a reduced configuration to the full configuration space. :param q_reduced: Configuration in reduced space :return: Configuration in full space with fixed joints set to their fixed values """ if q_reduced.shape[0] != self.num_dofs: raise ValueError( f"Expected q_reduced of length {self.num_dofs}, got {q_reduced.shape[0]} instead" ) q_full = np.zeros(self.full_robot.num_dofs) # Set active joints for reduced_idx, full_idx in enumerate(self.reduced_to_full): q_full[full_idx] = q_reduced[reduced_idx] # Set fixed joints for joint_name, value in zip(self.fixed_joints, self.fixed_values): full_idx = self.full_robot.dof_index(joint_name) q_full[full_idx] = value return q_full def full_to_reduced_configuration(self, q_full: np.ndarray) -> np.ndarray: """ Convert a full configuration to the reduced configuration space. :param q_full: Configuration in full space :return: Configuration in reduced space """ if q_full.shape[0] != self.full_robot.num_dofs: raise ValueError( f"Expected q_full of length {self.full_robot.num_dofs}, got {q_full.shape[0]} instead" ) q_reduced = np.zeros(self.num_dofs) # Copy active joints for reduced_idx, full_idx in enumerate(self.reduced_to_full): q_reduced[reduced_idx] = q_full[full_idx] return q_reduced def cache_forward_kinematics(self, q_reduced: np.ndarray, auto_clip=True) -> None: """ Perform forward kinematics using the reduced configuration. :param q_reduced: Configuration in reduced space """ # First update the full robot's forward kinematics q_full = self.reduced_to_full_configuration(q_reduced) self.full_robot.cache_forward_kinematics(q_full, auto_clip) # Then update the reduced model's forward kinematics pin.framesForwardKinematics( self.pinocchio_wrapper.model, self.pinocchio_wrapper.data, q_reduced ) def clip_configuration(self, q_reduced: np.ndarray, margin: float = 1e-6) -> np.ndarray: """ Clip the reduced configuration to stay within joint limits with a small tolerance. :param q_reduced: Configuration to clip :param margin: Tolerance to keep away from joint limits :return: Clipped configuration """ q_full = self.reduced_to_full_configuration(q_reduced) q_full_clipped = self.full_robot.clip_configuration(q_full, margin) return self.full_to_reduced_configuration(q_full_clipped) def reset_forward_kinematics(self): """ Reset the forward kinematics to the initial configuration. """ # Reset full robot's forward kinematics self.full_robot.reset_forward_kinematics() # Reset reduced model's forward kinematics self.cache_forward_kinematics(self.q_zero)