From 51690633465d666d2dd130ce9874ac645ec7113b Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 29 Nov 2025 14:33:24 +0000 Subject: [PATCH 1/3] Initial plan From 27cd1f48dca2f0dd5cf988fe96857a806fb0e1f4 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 29 Nov 2025 14:52:21 +0000 Subject: [PATCH 2/3] Refactor duplicated code in solver modules - Add _get_tcp_as_numpy() helper method to SolverCfg base class - Add build_reduced_pinocchio_robot() utility function to solver_utils - Add validate_iteration_params() utility function to solver_utils - Add compute_pinocchio_fk() utility function to solver_utils - Update all solver init_solver methods to use _get_tcp_as_numpy() - Update PinocchioSolver and PinkSolver to use build_reduced_pinocchio_robot() - Update PinocchioSolver and PytorchSolver to use validate_iteration_params() - Update PinocchioSolver and PinkSolver to use compute_pinocchio_fk() Co-authored-by: yuecideng <42235877+yuecideng@users.noreply.github.com> --- embodichain/lab/sim/solvers/base_solver.py | 13 ++ .../lab/sim/solvers/differential_solver.py | 6 +- embodichain/lab/sim/solvers/opw_solver.py | 6 +- embodichain/lab/sim/solvers/pink_solver.py | 57 ++------ .../lab/sim/solvers/pinocchio_solver.py | 80 ++--------- embodichain/lab/sim/solvers/pytorch_solver.py | 27 +--- embodichain/lab/sim/solvers/srs_solver.py | 6 +- embodichain/lab/sim/utility/solver_utils.py | 131 +++++++++++++++++- 8 files changed, 173 insertions(+), 153 deletions(-) diff --git a/embodichain/lab/sim/solvers/base_solver.py b/embodichain/lab/sim/solvers/base_solver.py index af7433c..f04d40b 100644 --- a/embodichain/lab/sim/solvers/base_solver.py +++ b/embodichain/lab/sim/solvers/base_solver.py @@ -76,6 +76,19 @@ class SolverCfg: def init_solver(self, device: torch.device, **kwargs) -> "BaseSolver": pass + def _get_tcp_as_numpy(self) -> np.ndarray: + """Convert TCP to numpy array. + + This helper method handles the conversion of TCP from torch.Tensor to numpy + if needed. Used by subclass init_solver methods to set TCP on the solver. + + Returns: + np.ndarray: The TCP as a numpy array. + """ + if isinstance(self.tcp, torch.Tensor): + return self.tcp.cpu().numpy() + return self.tcp + @classmethod def from_dict(cls, init_dict: Dict[str, Any]) -> "SolverCfg": """Initialize the configuration from a dictionary.""" diff --git a/embodichain/lab/sim/solvers/differential_solver.py b/embodichain/lab/sim/solvers/differential_solver.py index 5bbb802..15f8518 100644 --- a/embodichain/lab/sim/solvers/differential_solver.py +++ b/embodichain/lab/sim/solvers/differential_solver.py @@ -90,11 +90,7 @@ def init_solver( ) # Set the Tool Center Point (TCP) for the solver - if isinstance(self.tcp, torch.Tensor): - tcp = self.tcp.cpu().numpy() - else: - tcp = self.tcp - solver.set_tcp(tcp) + solver.set_tcp(self._get_tcp_as_numpy()) return solver diff --git a/embodichain/lab/sim/solvers/opw_solver.py b/embodichain/lab/sim/solvers/opw_solver.py index 1dfd751..0ddecff 100644 --- a/embodichain/lab/sim/solvers/opw_solver.py +++ b/embodichain/lab/sim/solvers/opw_solver.py @@ -95,11 +95,7 @@ def init_solver( solver = OPWSolver(cfg=self, device=device, **kwargs) # Set the Tool Center Point (TCP) for the solver - if isinstance(self.tcp, torch.Tensor): - tcp = self.tcp.cpu().numpy() - else: - tcp = self.tcp - solver.set_tcp(tcp) + solver.set_tcp(self._get_tcp_as_numpy()) return solver diff --git a/embodichain/lab/sim/solvers/pink_solver.py b/embodichain/lab/sim/solvers/pink_solver.py index b9d74a5..b7ab7b0 100644 --- a/embodichain/lab/sim/solvers/pink_solver.py +++ b/embodichain/lab/sim/solvers/pink_solver.py @@ -24,6 +24,10 @@ lazy_import_pinocchio, lazy_import_pink, ) +from embodichain.lab.sim.utility.solver_utils import ( + build_reduced_pinocchio_robot, + compute_pinocchio_fk, +) from embodichain.utils import configclass, logger from embodichain.lab.sim.solvers import SolverCfg, BaseSolver @@ -111,12 +115,7 @@ def init_solver(self, **kwargs) -> "PinkSolver": solver = PinkSolver(cfg=self, **kwargs) # Set the Tool Center Point (TCP) for the solver - if isinstance(self.tcp, torch.Tensor): - tcp = self.tcp.cpu().numpy() - else: - tcp = self.tcp - - solver.set_tcp(tcp) + solver.set_tcp(self._get_tcp_as_numpy()) return solver @@ -162,7 +161,7 @@ def __init__(self, cfg: PinkSolverCfg, **kwargs): ) # Degrees of freedom of robot joints # Get reduced robot model - self.robot = self._get_reduce_robot() + self.robot = build_reduced_pinocchio_robot(self.entire_robot, self.joint_names) # Initialize Pink configuration self.pink_cfg = self.pink.configuration.Configuration( @@ -207,26 +206,6 @@ def __init__(self, cfg: PinkSolverCfg, **kwargs): self.dexsim_to_pink_ordering = None self.pink_to_dexsim_ordering = None - def _get_reduce_robot(self) -> "pin.RobotWrapper": - """Build a reduced robot model by locking all joints except those in self.joint_names. - - Returns: - pin.RobotWrapper: The reduced robot model with specified joints unlocked. - """ - pink_joint_names = self.entire_robot.model.names.tolist() - - # Lock all joints except those in self.joint_names and 'universe' - fixed_joint_names = [ - name - for name in pink_joint_names - if name not in self.joint_names and name != "universe" - ] - - reduced_robot = self.entire_robot.buildReducedRobot( - list_of_joints_to_lock=fixed_joint_names - ) - return reduced_robot - def reorder_array( self, input_array: List[float], reordering_array: List[int] ) -> List[float]: @@ -393,25 +372,7 @@ def _get_fk( Returns: torch.Tensor: The homogeneous transformation matrix (4x4) of the end-effector (after applying TCP). """ - if isinstance(qpos, torch.Tensor): - qpos_np = qpos.detach().cpu().numpy() - else: - qpos_np = np.array(qpos) - - qpos_np = np.squeeze(qpos_np) - if qpos_np.ndim != 1: - raise ValueError(f"qpos shape must be (nq,), but got {qpos_np.shape}") - - self.pin.forwardKinematics(self.robot.model, self.robot.data, qpos_np) - - # Retrieve the pose of the specified link - frame_index = self.robot.model.getFrameId(self.end_link_name) - joint_index = self.robot.model.frames[frame_index].parent - xpos_se3 = self.robot.data.oMi.tolist()[joint_index] - - xpos = np.eye(4) - xpos[:3, :3] = xpos_se3.rotation - xpos[:3, 3] = xpos_se3.translation.T - - result = np.dot(xpos, self.tcp_xpos) + result = compute_pinocchio_fk( + self.pin, self.robot, qpos, self.end_link_name, self.tcp_xpos + ) return torch.from_numpy(result) diff --git a/embodichain/lab/sim/solvers/pinocchio_solver.py b/embodichain/lab/sim/solvers/pinocchio_solver.py index 719fabb..7acdeba 100644 --- a/embodichain/lab/sim/solvers/pinocchio_solver.py +++ b/embodichain/lab/sim/solvers/pinocchio_solver.py @@ -29,6 +29,11 @@ lazy_import_casadi, # lazy_import_pinocchio_casadi, ) +from embodichain.lab.sim.utility.solver_utils import ( + build_reduced_pinocchio_robot, + validate_iteration_params, + compute_pinocchio_fk, +) if TYPE_CHECKING: @@ -72,12 +77,7 @@ def init_solver(self, **kwargs) -> "PinocchioSolver": solver = PinocchioSolver(cfg=self, **kwargs) # Set the Tool Center Point (TCP) for the solver - if isinstance(self.tcp, torch.Tensor): - tcp = self.tcp.cpu().numpy() - else: - tcp = self.tcp - - solver.set_tcp(tcp) + solver.set_tcp(self._get_tcp_as_numpy()) return solver @@ -121,7 +121,7 @@ def __init__(self, cfg: PinocchioSolverCfg, **kwargs): ) # Degrees of freedom of robot joints # Build reduced robot model (only relevant joints unlocked) - self.robot = self._get_reduce_robot() + self.robot = build_reduced_pinocchio_robot(self.entire_robot, self.joint_names) self.joint_names = self.robot.model.names.tolist()[ 1: ] # Exclude 'universe' joint @@ -221,26 +221,6 @@ def __init__(self, cfg: PinocchioSolverCfg, **kwargs): self.root_base_xpos[:3, :3] = root_base_pose.rotation self.root_base_xpos[:3, 3] = root_base_pose.translation.T - def _get_reduce_robot(self) -> "pin.RobotWrapper": - """Build a reduced robot model by locking all joints except those in self.joint_names. - - Returns: - pin.RobotWrapper: The reduced robot model with specified joints unlocked. - """ - all_joint_names = self.entire_robot.model.names.tolist() - - # Lock all joints except those in self.joint_names and 'universe' - fixed_joint_names = [ - name - for name in all_joint_names - if name not in self.joint_names and name != "universe" - ] - - reduced_robot = self.entire_robot.buildReducedRobot( - list_of_joints_to_lock=fixed_joint_names - ) - return reduced_robot - def set_tcp(self, tcp: np.ndarray): self.tcp = tcp @@ -290,25 +270,10 @@ def set_iteration_params( Returns: bool: True if all parameters are valid and set, False otherwise. """ - # TODO: Check which parameters are no longer needed. # Validate parameters - if pos_eps <= 0: - logger.log_warning("Pos epsilon must be positive.") - return False - if rot_eps <= 0: - logger.log_warning("Rot epsilon must be positive.") - return False - if max_iterations <= 0: - logger.log_warning("Max iterations must be positive.") - return False - if dt <= 0: - logger.log_warning("Time step must be positive.") - return False - if damp < 0: - logger.log_warning("Damping factor must be non-negative.") - return False - if num_samples <= 0: - logger.log_warning("Number of samples must be positive.") + if not validate_iteration_params( + pos_eps, rot_eps, max_iterations, dt, damp, num_samples + ): return False # Set parameters if all are valid @@ -620,25 +585,6 @@ def _get_fk( Returns: np.ndarray: The resulting end-effector pose as a (4, 4) homogeneous transformation matrix. """ - if isinstance(qpos, torch.Tensor): - qpos_np = qpos.detach().cpu().numpy() - else: - qpos_np = np.array(qpos) - - qpos_np = np.squeeze(qpos_np) - if qpos_np.ndim != 1: - raise ValueError(f"qpos shape must be (nq,), but got {qpos_np.shape}") - - self.pin.forwardKinematics(self.robot.model, self.robot.data, qpos_np) - - # Retrieve the pose of the specified link - frame_index = self.robot.model.getFrameId(self.end_link_name) - joint_index = self.robot.model.frames[frame_index].parent - xpos_se3 = self.robot.data.oMi.tolist()[joint_index] - - xpos = np.eye(4) - xpos[:3, :3] = xpos_se3.rotation - xpos[:3, 3] = xpos_se3.translation.T - - result = np.dot(xpos, self.tcp_xpos) - return result + return compute_pinocchio_fk( + self.pin, self.robot, qpos, self.end_link_name, self.tcp_xpos + ) diff --git a/embodichain/lab/sim/solvers/pytorch_solver.py b/embodichain/lab/sim/solvers/pytorch_solver.py index 238e891..d8da695 100644 --- a/embodichain/lab/sim/solvers/pytorch_solver.py +++ b/embodichain/lab/sim/solvers/pytorch_solver.py @@ -23,6 +23,7 @@ from embodichain.utils import configclass, logger from embodichain.lab.sim.solvers import SolverCfg, BaseSolver from embodichain.lab.sim.solvers.qpos_seed_sampler import QposSeedSampler +from embodichain.lab.sim.utility.solver_utils import validate_iteration_params if TYPE_CHECKING: from typing import Self @@ -90,11 +91,7 @@ def init_solver( solver = PytorchSolver(cfg=self, device=device, **kwargs) # Set the Tool Center Point (TCP) for the solver - if isinstance(self.tcp, torch.Tensor): - tcp = self.tcp.cpu().numpy() - else: - tcp = self.tcp - solver.set_tcp(tcp) + solver.set_tcp(self._get_tcp_as_numpy()) return solver @@ -227,23 +224,9 @@ def set_iteration_params( bool: True if all parameters are valid and set, False otherwise. """ # Validate parameters - if pos_eps <= 0: - logger.log_warning("Pos epsilon must be positive.") - return False - if rot_eps <= 0: - logger.log_warning("Rot epsilon must be positive.") - return False - if max_iterations <= 0: - logger.log_warning("Max iterations must be positive.") - return False - if dt <= 0: - logger.log_warning("Time step must be positive.") - return False - if damp < 0: - logger.log_warning("Damping factor must be non-negative.") - return False - if num_samples <= 0: - logger.log_warning("Number of samples must be positive.") + if not validate_iteration_params( + pos_eps, rot_eps, max_iterations, dt, damp, num_samples + ): return False # Set parameters if all are valid diff --git a/embodichain/lab/sim/solvers/srs_solver.py b/embodichain/lab/sim/solvers/srs_solver.py index 44fa894..6d57757 100644 --- a/embodichain/lab/sim/solvers/srs_solver.py +++ b/embodichain/lab/sim/solvers/srs_solver.py @@ -93,11 +93,7 @@ def init_solver( solver = SRSSolver(cfg=self, num_envs=num_envs, device=device, **kwargs) # Set the Tool Center Point (TCP) for the solver - if isinstance(self.tcp, torch.Tensor): - tcp = self.tcp.cpu().numpy() - else: - tcp = self.tcp - solver.set_tcp(tcp) + solver.set_tcp(self._get_tcp_as_numpy()) return solver diff --git a/embodichain/lab/sim/utility/solver_utils.py b/embodichain/lab/sim/utility/solver_utils.py index 976462f..70150e7 100644 --- a/embodichain/lab/sim/utility/solver_utils.py +++ b/embodichain/lab/sim/utility/solver_utils.py @@ -17,7 +17,7 @@ import torch from embodichain.lab.sim.utility.io_utils import suppress_stdout_stderr -from typing import Optional, Union, Tuple, Any, TYPE_CHECKING +from typing import Optional, Union, Tuple, Any, List, TYPE_CHECKING from copy import deepcopy from embodichain.utils import configclass, logger @@ -109,3 +109,132 @@ def create_pk_serial_chain( return pk.SerialChain( chain=chain, end_frame_name=end_link_name, root_frame_name=root_link_name ) + + +def build_reduced_pinocchio_robot( + entire_robot: "pin.RobotWrapper", + joint_names: List[str], +) -> "pin.RobotWrapper": + """Build a reduced robot model by locking all joints except those specified. + + This utility function creates a reduced Pinocchio robot model by locking all joints + except those in the provided joint_names list and the 'universe' joint. + + Args: + entire_robot: The full Pinocchio RobotWrapper model. + joint_names: List of joint names to keep unlocked in the reduced model. + + Returns: + pin.RobotWrapper: The reduced robot model with specified joints unlocked. + """ + all_joint_names = entire_robot.model.names.tolist() + + # Lock all joints except those in joint_names and 'universe' + fixed_joint_names = [ + name + for name in all_joint_names + if name not in joint_names and name != "universe" + ] + + reduced_robot = entire_robot.buildReducedRobot( + list_of_joints_to_lock=fixed_joint_names + ) + return reduced_robot + + +def validate_iteration_params( + pos_eps: float, + rot_eps: float, + max_iterations: int, + dt: float, + damp: float, + num_samples: int, +) -> bool: + """Validate iteration parameters for IK solvers. + + This helper validates common iteration parameters used by IK solvers. Returns + True if all parameters are valid, False otherwise, and logs warnings for invalid + parameters. + + Args: + pos_eps: Position convergence threshold, must be positive. + rot_eps: Rotation convergence threshold, must be positive. + max_iterations: Maximum number of iterations, must be positive. + dt: Time step size, must be positive. + damp: Damping factor, must be non-negative. + num_samples: Number of samples, must be positive. + + Returns: + bool: True if all parameters are valid, False otherwise. + """ + if pos_eps <= 0: + logger.log_warning("Pos epsilon must be positive.") + return False + if rot_eps <= 0: + logger.log_warning("Rot epsilon must be positive.") + return False + if max_iterations <= 0: + logger.log_warning("Max iterations must be positive.") + return False + if dt <= 0: + logger.log_warning("Time step must be positive.") + return False + if damp < 0: + logger.log_warning("Damping factor must be non-negative.") + return False + if num_samples <= 0: + logger.log_warning("Number of samples must be positive.") + return False + return True + + +def compute_pinocchio_fk( + pin_module: Any, + robot: "pin.RobotWrapper", + qpos: Union[torch.Tensor, "np.ndarray"], + end_link_name: str, + tcp_xpos: "np.ndarray", +) -> "np.ndarray": + """Compute forward kinematics using Pinocchio for the specified end-effector. + + This utility function computes FK using the Pinocchio library and applies + the TCP transformation. + + Args: + pin_module: The imported pinocchio module. + robot: The Pinocchio RobotWrapper model. + qpos: Joint positions, shape should be (nq,). + end_link_name: The name of the end-effector link. + tcp_xpos: The TCP transformation matrix (4x4). + + Returns: + np.ndarray: The resulting end-effector pose as a (4, 4) homogeneous + transformation matrix. + + Raises: + ValueError: If qpos shape is not (nq,). + """ + import numpy as np + + if isinstance(qpos, torch.Tensor): + qpos_np = qpos.detach().cpu().numpy() + else: + qpos_np = np.array(qpos) + + qpos_np = np.squeeze(qpos_np) + if qpos_np.ndim != 1: + raise ValueError(f"qpos shape must be (nq,), but got {qpos_np.shape}") + + pin_module.forwardKinematics(robot.model, robot.data, qpos_np) + + # Retrieve the pose of the specified link + frame_index = robot.model.getFrameId(end_link_name) + joint_index = robot.model.frames[frame_index].parent + xpos_se3 = robot.data.oMi.tolist()[joint_index] + + xpos = np.eye(4) + xpos[:3, :3] = xpos_se3.rotation + xpos[:3, 3] = xpos_se3.translation.T + + result = np.dot(xpos, tcp_xpos) + return result From c5dba4f07005bc42c61d8f16e3528ba3a96df544 Mon Sep 17 00:00:00 2001 From: "copilot-swe-agent[bot]" <198982749+Copilot@users.noreply.github.com> Date: Sat, 29 Nov 2025 14:55:13 +0000 Subject: [PATCH 3/3] Address code review comments - Move numpy import to top of solver_utils.py - Add validation for end_link_name in compute_pinocchio_fk() Co-authored-by: yuecideng <42235877+yuecideng@users.noreply.github.com> --- embodichain/lab/sim/utility/solver_utils.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/embodichain/lab/sim/utility/solver_utils.py b/embodichain/lab/sim/utility/solver_utils.py index 70150e7..d510cc2 100644 --- a/embodichain/lab/sim/utility/solver_utils.py +++ b/embodichain/lab/sim/utility/solver_utils.py @@ -15,6 +15,7 @@ # ---------------------------------------------------------------------------- import torch +import numpy as np from embodichain.lab.sim.utility.io_utils import suppress_stdout_stderr from typing import Optional, Union, Tuple, Any, List, TYPE_CHECKING @@ -212,10 +213,8 @@ def compute_pinocchio_fk( transformation matrix. Raises: - ValueError: If qpos shape is not (nq,). + ValueError: If qpos shape is not (nq,) or if the end_link_name is not found. """ - import numpy as np - if isinstance(qpos, torch.Tensor): qpos_np = qpos.detach().cpu().numpy() else: @@ -229,6 +228,8 @@ def compute_pinocchio_fk( # Retrieve the pose of the specified link frame_index = robot.model.getFrameId(end_link_name) + if frame_index >= robot.model.nframes: + raise ValueError(f"End link name '{end_link_name}' not found in robot model.") joint_index = robot.model.frames[frame_index].parent xpos_se3 = robot.data.oMi.tolist()[joint_index]