diff --git a/examples/bridge/demo_mujoco_bi_independent.py b/examples/bridge/demo_mujoco_bi_independent.py index 502506b..eee6935 100644 --- a/examples/bridge/demo_mujoco_bi_independent.py +++ b/examples/bridge/demo_mujoco_bi_independent.py @@ -15,12 +15,25 @@ def main(): """Demo 1: Independent dual-arm IK control.""" # Model path (Bessica is a dual-arm robot) - mjcf_path = get_model_path("Bessica_D", version="v1_1", variant="skeleton_interactive", model_format="mjcf") - # mjcf_path = get_model_path("Bessica_D", version="v1_0", variant="covered_interactive", model_format="mjcf") + mjcf_path = get_model_path("Alicia_M", version="v1_1", variant="bimanual_interactive", model_format="mjcf") + # mjcf_path = get_model_path("Bessica_D", version="v1_1", variant="covered_interactive", model_format="mjcf") + # End-effector links - left_end = "left_arm_link7" - right_end = "right_arm_link7" + # left_end = "left_arm_link7" + # right_end = "right_arm_link7" + + left_end = "tool0_site_l" + right_end = "tool0_site_r" + + + # # Model path (Bessica is a dual-arm robot) + # mjcf_path = get_model_path("Bessica_D", version="v1_1", variant="skeleton_interactive", model_format="mjcf") + # # mjcf_path = get_model_path("Bessica_D", version="v1_0", variant="covered_interactive", model_format="mjcf") + + # # End-effector links + # left_end = "left_arm_link7" + # right_end = "right_arm_link7" # Create interactive IK controller controller = InteractiveDualArmIK(mjcf_path, left_end, right_end) diff --git a/examples/bridge/demo_mujoco_bi_mirror.py b/examples/bridge/demo_mujoco_bi_mirror.py index f94ab49..e317e30 100644 --- a/examples/bridge/demo_mujoco_bi_mirror.py +++ b/examples/bridge/demo_mujoco_bi_mirror.py @@ -16,7 +16,8 @@ def main(): """Demo 3: Mirror symmetric dual-arm control.""" # Model path (Bessica is a dual-arm robot) - mjcf_path = get_model_path("Bessica_D", version="v1_0", variant="covered_interactive", model_format="mjcf") + mjcf_path = get_model_path("Bessica_D", version="v1_1", variant="skeleton_interactive", model_format="mjcf") + # mjcf_path = get_model_path("Bessica_D", version="v1_0", variant="covered_interactive", model_format="mjcf") # End-effector links left_end = "left_arm_link7" diff --git a/examples/bridge/demo_mujoco_bi_relative.py b/examples/bridge/demo_mujoco_bi_relative.py index 736569c..6177d52 100644 --- a/examples/bridge/demo_mujoco_bi_relative.py +++ b/examples/bridge/demo_mujoco_bi_relative.py @@ -16,12 +16,17 @@ def main(): """Demo 2: Cooperative dual-arm control with relative constraint.""" # Model path (Bessica is a dual-arm robot) - mjcf_path = get_model_path("Bessica_D", version="v1_0", variant="covered_interactive", model_format="mjcf") + # mjcf_path = get_model_path("Alicia_D", version="v5_6", variant="bimanual_interactive", model_format="mjcf") + mjcf_path = get_model_path("Bessica_D", version="v1_1", variant="skeleton_interactive", model_format="mjcf") + # mjcf_path = get_model_path("Bessica_D", version="v1_1", variant="covered_interactive", model_format="mjcf") + # End-effector links left_end = "left_arm_link7" right_end = "right_arm_link7" - + + # left_end = "tool0_site_l" + # right_end = "tool0_site_r" # Create interactive IK controller controller = InteractiveDualArmIK(mjcf_path, left_end, right_end) diff --git a/examples/bridge/demo_mujoco_bi_relativem.py b/examples/bridge/demo_mujoco_bi_relativem.py new file mode 100644 index 0000000..614e564 --- /dev/null +++ b/examples/bridge/demo_mujoco_bi_relativem.py @@ -0,0 +1,67 @@ +#!/usr/bin/env python3 +"""Demo 2: Cooperative Dual-Arm Control with Relative Constraint. + +1. Drag red/blue spheres to set grasp configuration (relative pose) +2. Drag green sphere to move both arms cooperatively maintaining grasp + +Copyright (c) 2025 Synria Robotics Co., Ltd. +""" + +import sys +from pathlib import Path + +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +from robocore.bridge.sim.mujoco.interactive_dual_arm import InteractiveDualArmIK +from synriard import get_model_path + +INITIAL_JOINT_STATE = { + "joint1_l": 1.57, + "joint2_l": -2.35, + "joint3_l": -1.57, + "joint4_l": 0,#-1.18, + "joint5_l": 0.0, + "joint6_l": 0.0, + "joint7_l": 0.0, + "joint8_l": 0.0, + "joint1_r": -1.57,#1.21, + "joint2_r": -2.35, + "joint3_r": -1.57, + "joint4_r": 0.0, + "joint5_r": 0.0, + "joint6_r": 0.0, + "joint7_r": 0.0, + "joint8_r": 0.0, +} + + +def main(): + """Demo 2: Cooperative dual-arm control with relative constraint.""" + + # Model path (Bessica is a dual-arm robot) + mjcf_path = get_model_path("Alicia_M", version="v1_1", variant="bi_interactive", model_format="mjcf") + + + # End-effector links + # left_end = "left_arm_link7" + # right_end = "right_arm_link7" + + left_end = "tool0_site_l" + right_end = "tool0_site_r" + # Create interactive IK controller + controller = InteractiveDualArmIK( + mjcf_path, + left_end, + right_end, + initial_joint_state=INITIAL_JOINT_STATE, + enable_gripper_ui=True, + ) + + # Run in cooperative mode + controller.run(mode='relative') + + +if __name__ == '__main__': + main() diff --git a/examples/kinematics/05a_demo_bimanual_fkm.py b/examples/kinematics/05a_demo_bimanual_fkm.py new file mode 100644 index 0000000..8df4def --- /dev/null +++ b/examples/kinematics/05a_demo_bimanual_fkm.py @@ -0,0 +1,272 @@ +"""Bimanual Forward Kinematics Demo + +This demo demonstrates bimanual forward kinematics computation for dual-arm systems. +It shows how to compute FK for both arms independently, with relative constraints, and mirror mode. + +Copyright (c) 2025 Synria Robotics Co., Ltd. + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU General Public License as published by +the Free Software Foundation, either version 3 of the License, or +(at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU General Public License for more details. + +You should have received a copy of the GNU General Public License +along with this program. If not, see . + +Author: Synria Robotics Team +Website: https://synriarobotics.ai +""" + +import numpy as np +import argparse +import time +import sys +from pathlib import Path + +# Prefer the in-repo RoboCore package when running this demo directly. +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +import robocore as rc +from robocore.modeling.robot_model import RobotModel +from robocore.kinematics.bimanual import bimanual_forward_kinematics +from robocore.utils.beauty_logger import beauty_print_array, beauty_print +from robocore.utils.backend import to_numpy +from robocore.transform.conversions import * + + +def compute_bimanual_fk(left_model, right_model, backend, q_left, q_right, mode='indep'): + """Compute bimanual forward kinematics for given backend. + + :param left_model: Left arm RobotModel + :param right_model: Right arm RobotModel + :param backend: Backend name ('numpy' or 'torch') + :param q_left: Left joint angles in radians + :param q_right: Right joint angles in radians + :param mode: FK mode ('indep', 'relative', 'mirror') + :return: Dictionary with results and computation time + """ + rc.set_backend(backend) + start_time = time.time() + + result = bimanual_forward_kinematics( + left_model, right_model, q_left, q_right, + return_end=True, mode=mode + ) + + elapsed_time = time.time() - start_time + + T_left = to_numpy(result['left']) + T_right = to_numpy(result['right']) + + pos_left = T_left[:3, 3] + pos_right = T_right[:3, 3] + rot_left = T_left[:3, :3] + rot_right = T_right[:3, :3] + euler_left = matrix_to_euler(rot_left, seq='XYZ') + euler_right = matrix_to_euler(rot_right, seq='XYZ') + quat_left = matrix_to_quaternion(rot_left) + quat_right = matrix_to_quaternion(rot_right) + + ret = { + 'left': { + 'position': pos_left, + 'rotation': rot_left, + 'euler': euler_left, + 'quat': quat_left, + 'transform': T_left, + }, + 'right': { + 'position': pos_right, + 'rotation': rot_right, + 'euler': euler_right, + 'quat': quat_right, + 'transform': T_right, + }, + 'time': elapsed_time + } + + if 'relative' in result: + T_rel = to_numpy(result['relative']) + rot_rel = T_rel[:3, :3] + euler_rel = matrix_to_euler(rot_rel, seq='XYZ') + quat_rel = matrix_to_quaternion(rot_rel) + ret['relative'] = { + 'transform': T_rel, + 'position': T_rel[:3, 3], + 'rotation': rot_rel, + 'euler': euler_rel, + 'quat': quat_rel, + } + + if 'mirror' in result: + T_mirror = to_numpy(result['mirror']) + rot_mirror = T_mirror[:3, :3] + euler_mirror = matrix_to_euler(rot_mirror, seq='XYZ') + quat_mirror = matrix_to_quaternion(rot_mirror) + ret['mirror'] = { + 'transform': T_mirror, + 'position': T_mirror[:3, 3], + 'rotation': rot_mirror, + 'euler': euler_mirror, + 'quat': quat_mirror, + } + + return ret + + +def main(args): + # Load robot model (Bessica is a dual-arm robot, use same model with different base/end links) + left_model = RobotModel(str(args.model_path), base_link=args.left_base_link, end_link=args.left_end_link) + right_model = RobotModel(str(args.model_path), base_link=args.right_base_link, end_link=args.right_end_link) + + if args.verbose: + beauty_print("Left Arm Model:", type="module") + left_model.summary(show_chain=True) + beauty_print("Right Arm Model:", type="module") + right_model.summary(show_chain=True) + + # Compute with both backends + results_np = compute_bimanual_fk(left_model, right_model, 'numpy', + args.q_left, args.q_right, mode=args.mode) + results_torch = compute_bimanual_fk(left_model, right_model, 'torch', + args.q_left, args.q_right, mode=args.mode) + + # Display results + beauty_print(f"Left Arm End-Effector Position (m):") + pos_left_np = to_numpy(results_np['left']['position']) + pos_left_torch = to_numpy(results_torch['left']['position']) + print(f" NumPy: {beauty_print_array(pos_left_np)}") + print(f" Torch: {beauty_print_array(pos_left_torch)}") + pos_diff_left = np.linalg.norm(pos_left_np - pos_left_torch) + print(f" Diff: {pos_diff_left:.6e}") + + beauty_print(f"Left Arm End-Effector Orientation (Euler XYZ, radians):") + euler_left_np = to_numpy(results_np['left']['euler']) + euler_left_torch = to_numpy(results_torch['left']['euler']) + print(f" NumPy: {beauty_print_array(euler_left_np)}") + print(f" Torch: {beauty_print_array(euler_left_torch)}") + euler_diff_left = np.linalg.norm(euler_left_np - euler_left_torch) + print(f" Diff: {euler_diff_left:.6e}") + + beauty_print(f"Left Arm End-Effector Orientation (Quaternion xyzw):") + quat_left_np = to_numpy(results_np['left']['quat']) + quat_left_torch = to_numpy(results_torch['left']['quat']) + print(f" NumPy: {beauty_print_array(quat_left_np, precision=6)}") + print(f" Torch: {beauty_print_array(quat_left_torch, precision=6)}") + quat_diff_left = np.linalg.norm(quat_left_np - quat_left_torch) + print(f" Diff: {quat_diff_left:.6e}") + + beauty_print(f"Right Arm End-Effector Position (m):") + pos_right_np = to_numpy(results_np['right']['position']) + pos_right_torch = to_numpy(results_torch['right']['position']) + print(f" NumPy: {beauty_print_array(pos_right_np)}") + print(f" Torch: {beauty_print_array(pos_right_torch)}") + pos_diff_right = np.linalg.norm(pos_right_np - pos_right_torch) + print(f" Diff: {pos_diff_right:.6e}") + + beauty_print(f"Right Arm End-Effector Orientation (Euler XYZ, radians):") + euler_right_np = to_numpy(results_np['right']['euler']) + euler_right_torch = to_numpy(results_torch['right']['euler']) + print(f" NumPy: {beauty_print_array(euler_right_np)}") + print(f" Torch: {beauty_print_array(euler_right_torch)}") + euler_diff_right = np.linalg.norm(euler_right_np - euler_right_torch) + print(f" Diff: {euler_diff_right:.6e}") + + beauty_print(f"Right Arm End-Effector Orientation (Quaternion xyzw):") + quat_right_np = to_numpy(results_np['right']['quat']) + quat_right_torch = to_numpy(results_torch['right']['quat']) + print(f" NumPy: {beauty_print_array(quat_right_np, precision=6)}") + print(f" Torch: {beauty_print_array(quat_right_torch, precision=6)}") + quat_diff_right = np.linalg.norm(quat_right_np - quat_right_torch) + print(f" Diff: {quat_diff_right:.6e}") + + if 'relative' in results_np: + beauty_print(f"Relative Transform Position (m):") + pos_rel_np = to_numpy(results_np['relative']['position']) + pos_rel_torch = to_numpy(results_torch['relative']['position']) + print(f" NumPy: {beauty_print_array(pos_rel_np)}") + print(f" Torch: {beauty_print_array(pos_rel_torch)}") + pos_diff_rel = np.linalg.norm(pos_rel_np - pos_rel_torch) + print(f" Diff: {pos_diff_rel:.6e}") + + if 'euler' in results_np['relative']: + beauty_print(f"Relative Transform Orientation (Euler XYZ, radians):") + euler_rel_np = to_numpy(results_np['relative']['euler']) + euler_rel_torch = to_numpy(results_torch['relative']['euler']) + print(f" NumPy: {beauty_print_array(euler_rel_np)}") + print(f" Torch: {beauty_print_array(euler_rel_torch)}") + euler_diff_rel = np.linalg.norm(euler_rel_np - euler_rel_torch) + print(f" Diff: {euler_diff_rel:.6e}") + + if 'quat' in results_np['relative']: + beauty_print(f"Relative Transform Orientation (Quaternion xyzw):") + quat_rel_np = to_numpy(results_np['relative']['quat']) + quat_rel_torch = to_numpy(results_torch['relative']['quat']) + print(f" NumPy: {beauty_print_array(quat_rel_np, precision=6)}") + print(f" Torch: {beauty_print_array(quat_rel_torch, precision=6)}") + quat_diff_rel = np.linalg.norm(quat_rel_np - quat_rel_torch) + print(f" Diff: {quat_diff_rel:.6e}") + + if 'mirror' in results_np: + beauty_print(f"Mirror Transform Position (m):") + pos_mirror_np = to_numpy(results_np['mirror']['position']) + pos_mirror_torch = to_numpy(results_torch['mirror']['position']) + print(f" NumPy: {beauty_print_array(pos_mirror_np)}") + print(f" Torch: {beauty_print_array(pos_mirror_torch)}") + pos_diff_mirror = np.linalg.norm(pos_mirror_np - pos_mirror_torch) + print(f" Diff: {pos_diff_mirror:.6e}") + + if 'euler' in results_np['mirror']: + beauty_print(f"Mirror Transform Orientation (Euler XYZ, radians):") + euler_mirror_np = to_numpy(results_np['mirror']['euler']) + euler_mirror_torch = to_numpy(results_torch['mirror']['euler']) + print(f" NumPy: {beauty_print_array(euler_mirror_np)}") + print(f" Torch: {beauty_print_array(euler_mirror_torch)}") + euler_diff_mirror = np.linalg.norm(euler_mirror_np - euler_mirror_torch) + print(f" Diff: {euler_diff_mirror:.6e}") + + if 'quat' in results_np['mirror']: + beauty_print(f"Mirror Transform Orientation (Quaternion xyzw):") + quat_mirror_np = to_numpy(results_np['mirror']['quat']) + quat_mirror_torch = to_numpy(results_torch['mirror']['quat']) + print(f" NumPy: {beauty_print_array(quat_mirror_np, precision=6)}") + print(f" Torch: {beauty_print_array(quat_mirror_torch, precision=6)}") + quat_diff_mirror = np.linalg.norm(quat_mirror_np - quat_mirror_torch) + print(f" Diff: {quat_diff_mirror:.6e}") + + beauty_print(f"Computation Time:") + print(f" NumPy: {results_np['time']*1000:.4f} ms") + print(f" Torch: {results_torch['time']*1000:.4f} ms") + print(f" Ratio: {results_torch['time'] / results_np['time']:.2f}x") + + +if __name__ == "__main__": + from synriard import get_model_path + + # Bessica is a dual-arm robot + model_path = get_model_path("Alicia_M", version="v1_1", variant="bimanual_interactive", model_format="mjcf") + # model_path = get_model_path("Bessica_D", version="v1_1", variant="covered_interactive", model_format="mjcf") + + parser = argparse.ArgumentParser(description="Bimanual Forward Kinematics Demo") + parser.add_argument('--model-path', type=str, default=model_path, help='Path to robot model file (default: Alicia-M)') + parser.add_argument('--left-base-link', type=str, default='world_base', help='Left arm base link name') + parser.add_argument('--left-end-link', type=str, default='tool0_l', help='Left arm end-effector link name') + parser.add_argument('--right-base-link', type=str, default='world_base', help='Right arm base link name') + parser.add_argument('--right-end-link', type=str, default='tool0_r', help='Right arm end-effector link name') + parser.add_argument('--q-left', type=float, nargs='+', default=[0.32980586939541284, 0.342077715698498, 0.3390097541227267, 0.34054373491061235, 0.3528155812136975, 0.3497476196379262], + # parser.add_argument('--q-left', type=float, nargs='+', default=[0.1, 0.2, -0.3, 0.0, 0.5, -0.2, 0.1], + help='Left joint angles in radians') + parser.add_argument('--q-right', type=float, nargs='+', default=[-0.1, -0.2, 0.3, 0.0, -0.5, 0.2], + help='Right joint angles in radians') + parser.add_argument('--mode', type=str, default='indep', choices=['indep', 'relative', 'mirror'], + help='FK mode: indep (independent), relative (relative transform), mirror (mirror mode)') + parser.add_argument('--verbose', action='store_true', help='Show detailed model information') + args = parser.parse_args() + main(args) diff --git a/examples/kinematics/07a_demo_bimanual_ikm.py b/examples/kinematics/07a_demo_bimanual_ikm.py new file mode 100644 index 0000000..c9ae98b --- /dev/null +++ b/examples/kinematics/07a_demo_bimanual_ikm.py @@ -0,0 +1,221 @@ +"""Bimanual Inverse Kinematics Demo + +Copyright (c) 2025 Synria Robotics Co., Ltd. + +This program is free software: you can redistribute it and/or modify +it under the terms of the GNU General Public License as published by +the Free Software Foundation, either version 3 of the License, or +(at your option) any later version. + +This program is distributed in the hope that it will be useful, +but WITHOUT ANY WARRANTY; without even the implied warranty of +MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the +GNU General Public License for more details. + +You should have received a copy of the GNU General Public License +along with this program. If not, see . + +Author: Synria Robotics Team +Website: https://synriarobotics.ai +""" + +import numpy as np +import argparse +import time +import sys +from pathlib import Path + +# Prefer the in-repo RoboCore package when running this demo directly. +REPO_ROOT = Path(__file__).resolve().parents[2] +if str(REPO_ROOT) not in sys.path: + sys.path.insert(0, str(REPO_ROOT)) + +import robocore as rc +from robocore.modeling import RobotModel +from robocore.kinematics.bimanual import bimanual_inverse_kinematics +from robocore.utils.beauty_logger import beauty_print_array, beauty_print +from robocore.utils.backend import to_numpy +from robocore.transform.conversions import * + + +def pose7_to_matrix(target_pose): + """Convert [px, py, pz, qx, qy, qz, qw] pose to 4x4 matrix.""" + target_pose = np.asarray(target_pose, dtype=np.float64) + T = np.eye(4, dtype=np.float64) + T[:3, 3] = target_pose[:3] + T[:3, :3] = np.asarray(quaternion_to_matrix(target_pose[3:]), dtype=np.float64) + return T + + +def compute_ik(left_model, right_model, backend, target_left, target_right, q0_left=None, q0_right=None, + coordination='indep', num_initial_guesses=1, initial_guess_strategy='random', + initial_guess_scale=1.0, random_seed=None): + """Compute bimanual inverse kinematics for given backend. + + :param left_model: Left arm RobotModel + :param right_model: Right arm RobotModel + :param backend: Backend name ('numpy' or 'torch') + :param target_left: Target left end-effector pose [px, py, pz, qx, qy, qz, qw] or 4x4 matrix + :param target_right: Target right end-effector pose [px, py, pz, qx, qy, qz, qw] or 4x4 matrix + :param coordination: Coordination mode ('indep', 'relative_pose', 'relative_pos', 'relative_ori', 'mirror') + :param num_initial_guesses: Number of initial guesses to try + :param initial_guess_strategy: Strategy for generating initial guesses + :param initial_guess_scale: Scale factor for joint limits + :param random_seed: Random seed for reproducibility + :return: Dictionary with results and computation time + """ + rc.set_backend(backend) + start_time = time.time() + + # Convert target poses to 4x4 matrices if needed. + if isinstance(target_left, (list, np.ndarray)) and len(target_left) == 7: + T_left = pose7_to_matrix(target_left) + else: + T_left = np.array(target_left, dtype=np.float64) + + if isinstance(target_right, (list, np.ndarray)) and len(target_right) == 7: + T_right = pose7_to_matrix(target_right) + else: + T_right = np.array(target_right, dtype=np.float64) + + ik_result = bimanual_inverse_kinematics( + left_model, + right_model, + target_left=T_left, + target_right=T_right, + q0_left=q0_left, + q0_right=q0_right, + method='dls', + coordination=coordination, + use_analytic_jacobian=True, + num_initial_guesses=num_initial_guesses, + initial_guess_strategy=initial_guess_strategy, + initial_guess_scale=initial_guess_scale, + random_seed=random_seed, + ) + + elapsed_time = time.time() - start_time + + # Extract error information from res_left and res_right + res_left = ik_result['res_left'] + res_right = ik_result['res_right'] + + # Handle case where res_left/res_right might be dict or list + if isinstance(res_left, list) and len(res_left) > 0: + res_left = res_left[0] + if isinstance(res_right, list) and len(res_right) > 0: + res_right = res_right[0] + + # Get iters (use max of both arms if available) + iters_left = res_left['iters'] + iters_right = res_right['iters'] + iters = max(iters_left, iters_right) + + return { + 'success_left': ik_result['success_left'], + 'success_right': ik_result['success_right'], + 'iters': iters, + 'pos_err_left': res_left['pos_err'], + 'pos_err_right': res_right['pos_err'], + 'ori_err_left': res_left['ori_err'], + 'ori_err_right': res_right['ori_err'], + 'q_left': ik_result['q_left'], + 'q_right': ik_result['q_right'], + 'time': elapsed_time + } + + +def main(args): + left_model = RobotModel(str(args.model_path), base_link=args.left_base_link, end_link=args.left_end_link) + right_model = RobotModel(str(args.model_path), base_link=args.right_base_link, end_link=args.right_end_link) + target_left = np.asarray(args.target_left, dtype=np.float64) + target_right = np.asarray(args.target_right, dtype=np.float64) + + # Compute with both backends + results_np = compute_ik( + left_model, right_model, 'numpy', target_left, target_right, + coordination=args.coordination, + num_initial_guesses=args.num_inits, + initial_guess_strategy=args.init_strategy, + initial_guess_scale=args.init_scale, + random_seed=args.seed + ) + results_torch = compute_ik( + left_model, right_model, 'torch', target_left, target_right, + coordination=args.coordination, + num_initial_guesses=args.num_inits, + initial_guess_strategy=args.init_strategy, + initial_guess_scale=args.init_scale, + random_seed=args.seed + ) + + # Convert to numpy for comparison + q_left_np = to_numpy(results_np['q_left']) + q_right_np = to_numpy(results_np['q_right']) + q_left_torch = to_numpy(results_torch['q_left']) + q_right_torch = to_numpy(results_torch['q_right']) + + beauty_print(f"IK Solution:") + print(f" Success Left: NumPy={results_np['success_left']}, Torch={results_torch['success_left']}") + print(f" Success Right: NumPy={results_np['success_right']}, Torch={results_torch['success_right']}") + print(f" Iterations: NumPy={results_np['iters']}, Torch={results_torch['iters']}") + print(f" Position Error Left: NumPy={results_np['pos_err_left']:.6e} m, Torch={results_torch['pos_err_left']:.6e} m") + print(f" Position Error Right: NumPy={results_np['pos_err_right']:.6e} m, Torch={results_torch['pos_err_right']:.6e} m") + print(f" Orientation Error Left: NumPy={results_np['ori_err_left']:.6e} rad, Torch={results_torch['ori_err_left']:.6e} rad") + print(f" Orientation Error Right: NumPy={results_np['ori_err_right']:.6e} rad, Torch={results_torch['ori_err_right']:.6e} rad") + + beauty_print(f"Solved Joint Angles - Left Arm (radians):") + print(f" NumPy: {beauty_print_array(q_left_np)}") + print(f" Torch: {beauty_print_array(q_left_torch)}") + + beauty_print(f"Solved Joint Angles - Right Arm (radians):") + print(f" NumPy: {beauty_print_array(q_right_np)}") + print(f" Torch: {beauty_print_array(q_right_torch)}") + + beauty_print(f"Computation Time:") + print(f" NumPy: {results_np['time']* 1000:.4f} ms") + print(f" Torch: {results_torch['time']* 1000:.4f} ms") + if results_np['time'] > 0: + ratio = results_torch['time'] / results_np['time'] + print(f" Ratio: {ratio:.2f}x") + + +if __name__ == "__main__": + from synriard import get_model_path + + # model_path = get_model_path("Bessica_D", version="v1_0", variant="covered", model_format="urdf") + model_path = get_model_path("Alicia_M", version="v1_1", variant="bimanual_interactive", model_format="mjcf") + + parser = argparse.ArgumentParser(description="Bimanual Inverse Kinematics Demo") + parser.add_argument('--model-path', type=str, + default=model_path, + help='Path to model file (default: Bessica-D)') + parser.add_argument('--left-base-link', type=str, default='world_base', help='Left arm base link name') + parser.add_argument('--left-end-link', type=str, default='tool0_l', help='Left arm end-effector link name') + parser.add_argument('--right-base-link', type=str, default='world_base', help='Right arm base link name') + parser.add_argument('--right-end-link', type=str, default='tool0_r', help='Right arm end-effector link name') + parser.add_argument('--target-left', type=float, nargs='+', + default=[-0.19449737591558663, -0.20039701467905618, 0.12041091987533925, + 0.9811121564712746, -0.022305916485078278, -0.1621691940912966, 0.10306568294939504], + help='Target left end-effector pose as 7 floats (px, py, pz, qx, qy, qz, qw)') + parser.add_argument('--target-right', type=float, nargs='+', + default=[0.09020439652403661, -0.07509374443188993, 0.08822214081723735, + -0.08784407246632346, 0.9749422854306018, -0.02894067453598843, -0.20232003452275402], + help='Target right end-effector pose as 7 floats (px, py, pz, qx, qy, qz, qw)') + parser.add_argument('--coordination', type=str, default='indep', + choices=['indep', 'relative_pose', 'relative_pos', 'relative_ori', 'mirror'], + help='Coordination mode: indep (independent), relative_pose, relative_pos, relative_ori, mirror') + parser.add_argument('--num-inits', type=int, default=1, + help='Number of initial guesses to try per target (default: 1)') + parser.add_argument('--init-strategy', type=str, default='random', + choices=['zero', 'random', 'sobol', 'latin', 'center', 'uniform'], + help='Strategy for generating initial guesses (default: random)') + parser.add_argument('--init-scale', type=float, default=1.0, + help='Scale factor for joint limits when generating guesses (0.0 to 1.0, default: 1.0)') + parser.add_argument('--seed', type=int, default=None, + help='Random seed for reproducibility (default: None)') + parser.add_argument('--backend', type=str, default='numpy', choices=['numpy', 'torch'], + help='Backend to use for computation (default: numpy, ignored - both are tested)') + args = parser.parse_args() + + main(args) diff --git a/robocore/bridge/sim/mujoco/interactive_dual_arm.py b/robocore/bridge/sim/mujoco/interactive_dual_arm.py index 3c5974b..1d98dae 100644 --- a/robocore/bridge/sim/mujoco/interactive_dual_arm.py +++ b/robocore/bridge/sim/mujoco/interactive_dual_arm.py @@ -11,6 +11,15 @@ from typing import Optional, Dict, Tuple import time +try: + import tkinter as tk + from tkinter import ttk + TK_AVAILABLE = True +except ImportError: + tk = None + ttk = None + TK_AVAILABLE = False + try: import mujoco import mujoco.viewer @@ -30,7 +39,15 @@ class InteractiveDualArmIK: :param right_end_link: Right arm end-effector link name """ - def __init__(self, mjcf_path: str, left_end_link: str, right_end_link: str): + def __init__( + self, + mjcf_path: str, + left_end_link: str, + right_end_link: str, + initial_joint_state: Optional[Dict[str, float]] = None, + initial_gripper_values: Optional[Dict[str, Optional[float]]] = None, + enable_gripper_ui: bool = False, + ): if not MUJOCO_AVAILABLE: raise ImportError("MuJoCo is required. Install with: pip install mujoco") @@ -43,21 +60,47 @@ def __init__(self, mjcf_path: str, left_end_link: str, right_end_link: str): self.mj_model = mujoco.MjModel.from_xml_path(str(self.mjcf_path)) self.mj_data = mujoco.MjData(self.mj_model) print(f"✓ MuJoCo model loaded: {self.mj_model.nq} DOF") - - # Gripper center offset (from link7 to gripper center, in link7 frame) - # Based on MJCF: gripper center at pos="0.14128 0 -0.00015" (both arms) - # Center is at 0.14128 m in X direction from link7 - self.gripper_offset = np.array([0.14128, 0.0, -0.00015, 1.0]) # Homogeneous coordinates - + + # Resolve requested end-effector references. The interactive demos may + # pass site names, while RoboCore kinematics works on body/link names. + self.left_end_name = left_end_link + self.right_end_name = right_end_link + self.left_end_body, self.left_end_offset = self._resolve_end_reference(left_end_link) + self.right_end_body, self.right_end_offset = self._resolve_end_reference(right_end_link) + self.left_end_offset_inv = np.linalg.inv(self.left_end_offset) + self.right_end_offset_inv = np.linalg.inv(self.right_end_offset) + + # Infer a shared base from the MuJoCo tree so Alicia-style world roots + # and Bessica-style base links both work. + self.base_link = self._infer_common_base(self.left_end_body, self.right_end_body) + # Load RoboCore bimanual model - self.robot = BimanualRobotModel(str(self.mjcf_path), left_end_link, right_end_link) + self.robot = BimanualRobotModel( + str(self.mjcf_path), + self.left_end_body, + self.right_end_body, + base_link=self.base_link, + ) self.left_model = self.robot.left_model self.right_model = self.robot.right_model + + # Map RoboCore chain joints to MuJoCo qpos indices instead of assuming + # a fixed left/right ordering in qpos. + self.left_qpos_indices = self._joint_qpos_indices(self.left_model) + self.right_qpos_indices = self._joint_qpos_indices(self.right_model) + # Optional startup joint configuration (joint name -> value in qpos units). + self.initial_joint_state = dict(initial_joint_state or {}) + + # Optional startup gripper configuration. + self._initialize_gripper_controls( + initial_gripper_values=initial_gripper_values, + enable_gripper_ui=enable_gripper_ui, + ) + # Current joint configuration - self.q_left = np.zeros(self.left_model.num_chain_dof) - self.q_right = np.zeros(self.right_model.num_chain_dof) + self.q_left, self.q_right = self._build_initial_joint_configuration() # Target poses self.T_left_target = None @@ -86,24 +129,303 @@ def __init__(self, mjcf_path: str, left_end_link: str, right_end_link: str): # T_rel = T_left^-1 @ T_right self.T_rel_grasp = None self._initialize_relative_transform() + + # Save startup joint values for reset behavior. + self.q_left_initial = self.q_left.copy() + self.q_right_initial = self.q_right.copy() self.is_running = False self.viewer = None self.reset_requested = False # Flag for reset request + + def _initialize_gripper_controls( + self, + initial_gripper_values: Optional[Dict[str, Optional[float]]], + enable_gripper_ui: bool, + ): + """Prepare gripper state for visualization and optional SDK streaming.""" + initial_gripper_values = dict(initial_gripper_values or {}) + + self.left_gripper_value = self._sanitize_gripper_value( + initial_gripper_values.get("left", 0.0) + ) + self.right_gripper_value = self._sanitize_gripper_value( + initial_gripper_values.get("right", 0.0) + ) + self.left_gripper_active = initial_gripper_values.get("left") is not None + self.right_gripper_active = initial_gripper_values.get("right") is not None + + self.left_gripper_initial = self.left_gripper_value + self.right_gripper_initial = self.right_gripper_value + self.left_gripper_initial_active = self.left_gripper_active + self.right_gripper_initial_active = self.right_gripper_active + + self.enable_gripper_ui = bool(enable_gripper_ui and TK_AVAILABLE) + if enable_gripper_ui and not TK_AVAILABLE: + print("⚠️ tkinter not available, gripper UI disabled") + + self.gripper_ui_root = None + self.left_gripper_var = None + self.right_gripper_var = None + self.left_gripper_value_label = None + self.right_gripper_value_label = None + self._updating_gripper_ui = False + + self.left_gripper_entries = self._resolve_gripper_joint_entries( + ("left_finger_l", "joint7_l"), + ("right_finger_l", "joint8_l"), + ) + self.right_gripper_entries = self._resolve_gripper_joint_entries( + ("left_finger_r", "joint7_r"), + ("right_finger_r", "joint8_r"), + ) + + if self.left_gripper_entries or self.right_gripper_entries: + print("✓ Gripper joints found for interactive finger control") + else: + print("⚠️ No gripper joints found in MuJoCo model; gripper UI will only affect SDK streaming") + + def _sanitize_gripper_value(self, value: Optional[float]) -> float: + """Clamp gripper command to the SDK's 0..1000 range.""" + if value is None: + return 0.0 + return float(np.clip(float(value), 0.0, 1000.0)) + + def _resolve_gripper_joint_entries(self, *joint_alias_groups) -> list: + """Resolve one arm's finger joints to qpos entries and open/closed limits.""" + entries = [] + for aliases in joint_alias_groups: + joint_id = -1 + for name in aliases: + joint_id = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_JOINT, name) + if joint_id != -1: + break + + if joint_id == -1: + continue + + lo, hi = sorted(self.mj_model.jnt_range[joint_id].tolist()) + if lo <= 0.0 <= hi: + closed_pos = 0.0 + else: + closed_pos = lo if abs(lo) <= abs(hi) else hi + open_pos = lo if abs(lo - closed_pos) >= abs(hi - closed_pos) else hi + + entries.append( + { + "qpos_index": int(self.mj_model.jnt_qposadr[joint_id]), + "closed_pos": float(closed_pos), + "open_pos": float(open_pos), + } + ) + return entries + + def _gripper_value_to_qpos(self, gripper_value: float, joint_entry: Dict[str, float]) -> float: + """Map SDK gripper command 0..1000 to one MuJoCo finger slide joint.""" + openness = self._sanitize_gripper_value(gripper_value) / 1000.0 + return joint_entry["closed_pos"] + openness * ( + joint_entry["open_pos"] - joint_entry["closed_pos"] + ) + + def _apply_gripper_targets_to_mujoco(self): + """Apply the current left/right gripper commands to MuJoCo finger joints.""" + for entry in self.left_gripper_entries: + self.mj_data.qpos[entry["qpos_index"]] = self._gripper_value_to_qpos( + self.left_gripper_value, entry + ) + for entry in self.right_gripper_entries: + self.mj_data.qpos[entry["qpos_index"]] = self._gripper_value_to_qpos( + self.right_gripper_value, entry + ) + + def _set_gripper_target(self, arm: str, value: float, activate: bool = True): + """Update one gripper target from the UI or reset flow.""" + sanitized = self._sanitize_gripper_value(value) + if arm == "left": + self.left_gripper_value = sanitized + self.left_gripper_active = activate + elif arm == "right": + self.right_gripper_value = sanitized + self.right_gripper_active = activate + else: + raise ValueError(f"Unknown arm '{arm}'") + + def get_gripper_targets_for_streaming(self) -> Tuple[Optional[int], Optional[int]]: + """Return SDK gripper targets, keeping inactive sides untouched.""" + left_value = int(round(self.left_gripper_value)) if self.left_gripper_active else None + right_value = int(round(self.right_gripper_value)) if self.right_gripper_active else None + return left_value, right_value + + def _create_gripper_slider_row(self, parent, title: str, initial_value: float, callback): + """Create one slider row for the sidecar gripper window.""" + row = ttk.Frame(parent) + row.pack(fill="x", pady=6) + + ttk.Label(row, text=title, width=12).pack(side="left") + value_var = tk.DoubleVar(value=initial_value) + slider = ttk.Scale( + row, + from_=0.0, + to=1000.0, + orient="horizontal", + variable=value_var, + command=callback, + ) + slider.pack(side="left", fill="x", expand=True, padx=8) + value_label = ttk.Label(row, width=12) + value_label.pack(side="left") + + return value_var, value_label + + def _refresh_gripper_ui_labels(self): + """Refresh the textual value labels in the sidecar gripper window.""" + if self.left_gripper_value_label is not None: + suffix = "" if self.left_gripper_active else " (idle)" + self.left_gripper_value_label.config( + text=f"{int(round(self.left_gripper_value))}{suffix}" + ) + if self.right_gripper_value_label is not None: + suffix = "" if self.right_gripper_active else " (idle)" + self.right_gripper_value_label.config( + text=f"{int(round(self.right_gripper_value))}{suffix}" + ) + + def _sync_gripper_ui_values(self): + """Push internal gripper state back into the UI without re-triggering control.""" + if self.gripper_ui_root is None: + return + + self._updating_gripper_ui = True + try: + if self.left_gripper_var is not None: + self.left_gripper_var.set(self.left_gripper_value) + if self.right_gripper_var is not None: + self.right_gripper_var.set(self.right_gripper_value) + self._refresh_gripper_ui_labels() + finally: + self._updating_gripper_ui = False + + def _on_gripper_slider_change(self, arm: str, value: str): + """Handle one gripper slider move from the sidecar window.""" + if self._updating_gripper_ui: + return + self._set_gripper_target(arm, float(value), activate=True) + self._refresh_gripper_ui_labels() + + def _close_gripper_ui(self): + """Destroy the sidecar gripper window if it exists.""" + root = self.gripper_ui_root + self.gripper_ui_root = None + self.left_gripper_var = None + self.right_gripper_var = None + self.left_gripper_value_label = None + self.right_gripper_value_label = None + if root is not None: + try: + root.destroy() + except Exception: + pass + + def _create_gripper_ui(self): + """Create a small sidecar window for left/right gripper control.""" + if not self.enable_gripper_ui or self.gripper_ui_root is not None: + return + + try: + root = tk.Tk() + except Exception as exc: + print(f"⚠️ Failed to create gripper UI: {exc}") + self.enable_gripper_ui = False + return + + root.title("MuJoCo Dual Gripper Control") + root.geometry("430x170") + root.protocol("WM_DELETE_WINDOW", self._close_gripper_ui) + + frame = ttk.Frame(root, padding=10) + frame.pack(fill="both", expand=True) + + ttk.Label( + frame, + text="Grippers: 0 = closed, 1000 = open", + ).pack(anchor="w", pady=(0, 8)) + + self.left_gripper_var, self.left_gripper_value_label = self._create_gripper_slider_row( + frame, + "Left gripper", + self.left_gripper_value, + lambda value: self._on_gripper_slider_change("left", value), + ) + self.right_gripper_var, self.right_gripper_value_label = self._create_gripper_slider_row( + frame, + "Right gripper", + self.right_gripper_value, + lambda value: self._on_gripper_slider_change("right", value), + ) + + ttk.Button( + frame, + text="Reset grippers", + command=self._reset_gripper_targets, + ).pack(anchor="e", pady=(10, 0)) + + self.gripper_ui_root = root + self._sync_gripper_ui_values() + + def _pump_gripper_ui(self): + """Keep the sidecar gripper window responsive during the MuJoCo loop.""" + if self.gripper_ui_root is None: + return + try: + self.gripper_ui_root.update_idletasks() + self.gripper_ui_root.update() + except Exception: + self._close_gripper_ui() + + def _reset_gripper_targets(self): + """Restore startup gripper values and activation flags.""" + self.left_gripper_value = self.left_gripper_initial + self.right_gripper_value = self.right_gripper_initial + self.left_gripper_active = self.left_gripper_initial_active + self.right_gripper_active = self.right_gripper_initial_active + self._sync_gripper_ui_values() + + def _build_initial_joint_configuration(self) -> Tuple[np.ndarray, np.ndarray]: + """Build initial left/right chain vectors from optional joint-name map.""" + q_left = np.zeros(self.left_model.num_chain_dof) + q_right = np.zeros(self.right_model.num_chain_dof) + + for i, js in enumerate(self.left_model._chain_actuated): + if js.name in self.initial_joint_state: + q_left[i] = float(self.initial_joint_state[js.name]) + for i, js in enumerate(self.right_model._chain_actuated): + if js.name in self.initial_joint_state: + q_right[i] = float(self.initial_joint_state[js.name]) + + # Keep MuJoCo state aligned with IK startup state. + self.mj_data.qpos[self.left_qpos_indices[: self.left_model.num_chain_dof]] = q_left + self.mj_data.qpos[self.right_qpos_indices[: self.right_model.num_chain_dof]] = q_right + self._apply_gripper_targets_to_mujoco() + mujoco.mj_forward(self.mj_model, self.mj_data) + return q_left, q_right def _initialize_mocap_ids(self): """Find mocap body IDs by name.""" + left_aliases = {'left_target', 'ik_target_left_gripper'} + right_aliases = {'right_target', 'ik_target_right_gripper'} + center_aliases = {'center_target', 'ik_target_center'} + # Iterate through all bodies to find mocap bodies for body_id in range(self.mj_model.nbody): # Check if this body is a mocap body mocap_id = self.mj_model.body_mocapid[body_id] if mocap_id >= 0: body_name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, body_id) - if body_name == 'left_target': + if body_name in left_aliases: self.left_marker_id = mocap_id - elif body_name == 'right_target': + elif body_name in right_aliases: self.right_marker_id = mocap_id - elif body_name == 'center_target': + elif body_name in center_aliases: self.center_marker_id = mocap_id print(f"✓ Mocap markers found: left={self.left_marker_id}, right={self.right_marker_id}, center={self.center_marker_id}") @@ -113,17 +435,15 @@ def _initialize_targets(self): # Get initial FK poses (at zero config, should be zero-config FK) result_left = self.left_model.fk(self.q_left) result_right = self.right_model.fk(self.q_right) - - # Get link7 poses (FK returns dict with 'end' key) - T_left_link7 = result_left['end'] - T_right_link7 = result_right['end'] - - # Convert to gripper center poses (T_gripper = T_link7 @ T_offset) - T_offset = np.eye(4) - T_offset[0:3, 3] = self.gripper_offset[0:3] - - self.T_left_target = T_left_link7 @ T_offset - self.T_right_target = T_right_link7 @ T_offset + + # Convert from the IK body pose to the requested end-effector reference + # (body or site). For Alicia this maps tool0_l/tool0_r bodies to the + # requested tool sites; for Bessica it maps link7 to gripper-center sites. + T_left_body = result_left['end'] + T_right_body = result_right['end'] + + self.T_left_target = T_left_body @ self.left_end_offset + self.T_right_target = T_right_body @ self.right_end_offset # Save initial poses for mirror mode reference self.T_left_initial = self.T_left_target.copy() @@ -141,9 +461,9 @@ def _initialize_targets(self): def _initialize_relative_transform(self): """Initialize relative grasp transform from current poses.""" - # Convert gripper targets to link7 targets - T_left_link7 = self._gripper_to_link7(self.T_left_target) - T_right_link7 = self._gripper_to_link7(self.T_right_target) + # Convert requested target frames back to the underlying IK body frames. + T_left_link7 = self._target_to_ik_link(self.T_left_target, arm='left') + T_right_link7 = self._target_to_ik_link(self.T_right_target, arm='right') # Compute relative transform self.T_rel_grasp = np.linalg.inv(T_left_link7) @ T_right_link7 @@ -171,6 +491,38 @@ def _update_markers(self): if self.center_marker_id is not None: self.mj_data.mocap_pos[self.center_marker_id] = self.T_center_target[0:3, 3] self.mj_data.mocap_quat[self.center_marker_id] = self._mat2quat(self.T_center_target[0:3, 0:3]) + + def _debug_marker_state(self, prefix: str = "Marker debug"): + """Print current target, mocap, and body poses for startup debugging.""" + marker_info = [ + ("left", self.left_marker_id, self.T_left_target), + ("right", self.right_marker_id, self.T_right_target), + ("center", self.center_marker_id, self.T_center_target), + ] + + print(f"\n[{prefix}]") + for label, mocap_id, target in marker_info: + if mocap_id is None: + print(f" {label}: mocap not found") + continue + + body_name = None + for body_id in range(self.mj_model.nbody): + if self.mj_model.body_mocapid[body_id] == mocap_id: + body_name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, body_id) + break + + body_pos = None + if body_name is not None: + body_id = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, body_name) + body_pos = self.mj_data.xpos[body_id].copy() + + target_pos = target[0:3, 3] if target is not None else None + mocap_pos = self.mj_data.mocap_pos[mocap_id].copy() + print( + f" {label}: body={body_name}, " + f"target={target_pos}, mocap={mocap_pos}, body_xpos={body_pos}" + ) def _mat2quat(self, R: np.ndarray) -> np.ndarray: """Convert rotation matrix to quaternion [w, x, y, z].""" @@ -243,14 +595,66 @@ def _get_marker_poses(self) -> Tuple[np.ndarray, np.ndarray, np.ndarray]: return T_left, T_right, T_center - def _gripper_to_link7(self, T_gripper: np.ndarray) -> np.ndarray: - """Convert gripper center pose to link7 pose. - - T_link7 = T_gripper @ inv(T_offset) - """ + def _resolve_end_reference(self, end_name: str) -> Tuple[str, np.ndarray]: + """Resolve a body/site end-effector reference to an IK body and offset.""" + body_id = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, end_name) + if body_id != -1: + return end_name, np.eye(4) + + site_id = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_SITE, end_name) + if site_id == -1: + raise ValueError(f"End-effector '{end_name}' is neither a MuJoCo body nor site") + + body_id = self.mj_model.site_bodyid[site_id] + body_name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, body_id) + if body_name is None: + raise ValueError(f"Could not resolve parent body for site '{end_name}'") + T_offset = np.eye(4) - T_offset[0:3, 3] = self.gripper_offset[0:3] - return T_gripper @ np.linalg.inv(T_offset) + T_offset[0:3, 3] = self.mj_model.site_pos[site_id] + if hasattr(self.mj_model, 'site_quat'): + T_offset[0:3, 0:3] = self._quat2mat(self.mj_model.site_quat[site_id]) + return body_name, T_offset + + def _infer_common_base(self, left_body: str, right_body: str) -> str: + """Infer the deepest common ancestor body of two end-effectors.""" + def body_ancestors(body_name: str): + body_id = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, body_name) + ancestors = [] + while body_id != -1: + name = mujoco.mj_id2name(self.mj_model, mujoco.mjtObj.mjOBJ_BODY, body_id) + if name is not None: + ancestors.append(name) + parent_id = self.mj_model.body_parentid[body_id] + if parent_id == body_id: + break + body_id = parent_id + return ancestors + + left_ancestors = body_ancestors(left_body) + right_ancestors = set(body_ancestors(right_body)) + for name in left_ancestors: + if name in right_ancestors: + return name + return left_ancestors[-1] + + def _joint_qpos_indices(self, model) -> np.ndarray: + """Map a RoboCore chain's actuated joints to MuJoCo qpos addresses.""" + indices = [] + for js in model._chain_actuated: + joint_id = mujoco.mj_name2id(self.mj_model, mujoco.mjtObj.mjOBJ_JOINT, js.name) + if joint_id == -1: + raise ValueError(f"MuJoCo joint '{js.name}' not found in model") + indices.append(self.mj_model.jnt_qposadr[joint_id]) + return np.asarray(indices, dtype=np.int32) + + def _target_to_ik_link(self, T_target: np.ndarray, arm: str) -> np.ndarray: + """Convert requested end-effector pose to the IK body pose.""" + if arm == 'left': + return T_target @ self.left_end_offset_inv + if arm == 'right': + return T_target @ self.right_end_offset_inv + raise ValueError(f"Unknown arm '{arm}'") def solve_ik_independent(self): """Solve IK for independent dual-arm control (Demo 1). @@ -259,8 +663,8 @@ def solve_ik_independent(self): """ # Use BimanualRobotModel.ik() with independent coordination res = self.robot.ik( - target_left=self._gripper_to_link7(self.T_left_target), - target_right=self._gripper_to_link7(self.T_right_target), + target_left=self._target_to_ik_link(self.T_left_target, arm='left'), + target_right=self._target_to_ik_link(self.T_right_target, arm='right'), q0_left=self.q_left, q0_right=self.q_right, method='dls', @@ -285,7 +689,7 @@ def solve_ik_relative(self): """ # Use BimanualRobotModel.ik() with relative_pose coordination res = self.robot.ik( - target_left=self._gripper_to_link7(self.T_left_target), + target_left=self._target_to_ik_link(self.T_left_target, arm='left'), target_right=None, # right is constrained by T_rel_grasp q0_left=self.q_left, q0_right=self.q_right, @@ -312,9 +716,9 @@ def solve_ik_mirror(self): - Mirror this rotation change across XZ plane - Apply mirrored rotation to right initial orientation """ - # Convert gripper center targets to link7 targets and solve independently - T_left_link7 = self._gripper_to_link7(self.T_left_target) - T_right_link7 = self._gripper_to_link7(self.T_right_target) + # Convert requested target frames to the underlying IK body frames. + T_left_link7 = self._target_to_ik_link(self.T_left_target, arm='left') + T_right_link7 = self._target_to_ik_link(self.T_right_target, arm='right') # Use BimanualRobotModel.ik() with independent coordination res = self.robot.ik( @@ -340,8 +744,9 @@ def _update_robot_pose(self): nq_right = self.right_model.num_chain_dof # Right arm first, left arm second - self.mj_data.qpos[0:nq_right] = self.q_right - self.mj_data.qpos[nq_right:nq_right+nq_left] = self.q_left + self.mj_data.qpos[self.left_qpos_indices[:nq_left]] = self.q_left + self.mj_data.qpos[self.right_qpos_indices[:nq_right]] = self.q_right + self._apply_gripper_targets_to_mujoco() # Forward kinematics in MuJoCo mujoco.mj_forward(self.mj_model, self.mj_data) @@ -417,8 +822,8 @@ def step(self): # CRITICAL: Update relative transform when red/blue balls are dragged # This redefines the relative grasp configuration - T_left_link7 = self._gripper_to_link7(self.T_left_target) - T_right_link7 = self._gripper_to_link7(self.T_right_target) + T_left_link7 = self._target_to_ik_link(self.T_left_target, arm='left') + T_right_link7 = self._target_to_ik_link(self.T_right_target, arm='right') self.T_rel_grasp = np.linalg.inv(T_left_link7) @ T_right_link7 print( @@ -520,9 +925,10 @@ def reset(self): """ print("\n🔄 Resetting to initial state...") - # Reset joint angles to zero - self.q_left = np.zeros(self.left_model.num_chain_dof) - self.q_right = np.zeros(self.right_model.num_chain_dof) + # Reset joint angles to startup configuration + self.q_left = self.q_left_initial.copy() + self.q_right = self.q_right_initial.copy() + self._reset_gripper_targets() # Reset target poses to initial values self.T_left_target = self.T_left_initial.copy() @@ -533,8 +939,8 @@ def reset(self): self.T_center_target[0:3, 0:3] = self.T_left_target[0:3, 0:3].copy() # Reset relative grasp transform - T_left_link7 = self._gripper_to_link7(self.T_left_target) - T_right_link7 = self._gripper_to_link7(self.T_right_target) + T_left_link7 = self._target_to_ik_link(self.T_left_target, arm='left') + T_right_link7 = self._target_to_ik_link(self.T_right_target, arm='right') self.T_rel_grasp = np.linalg.inv(T_left_link7) @ T_right_link7 # Update MuJoCo state @@ -573,11 +979,17 @@ def run(self, mode: str = 'independent'): print(f"\n - Press SPACE to reset") print(f" - Press ESC to exit\n") + if self.enable_gripper_ui: + print(" - Side window sliders: left/right gripper control") + print(" Move a slider once to start streaming that gripper") + print("") # Try passive viewer first, fallback on macOS mjpython error import platform try: + self._create_gripper_ui() + # Launch viewer with mujoco.viewer.launch_passive(self.mj_model, self.mj_data) as viewer: self.viewer = viewer @@ -591,6 +1003,7 @@ def run(self, mode: str = 'independent'): # This prevents the initial huge displacement that causes shaking self._update_markers() mujoco.mj_forward(self.mj_model, self.mj_data) + self._debug_marker_state(prefix="Startup marker state before viewer sync") # Sync viewer AFTER setting mocap positions viewer.sync() @@ -604,6 +1017,7 @@ def run(self, mode: str = 'independent'): while viewer.is_running() and self.is_running: step_start = time.time() + self._pump_gripper_ui() # Check for reset: either time went backwards OR qpos was reset to zero current_time = self.mj_data.time @@ -646,5 +1060,7 @@ def run(self, mode: str = 'independent'): raise else: raise + finally: + self._close_gripper_ui() print("\n✓ Visualization closed") diff --git a/robocore/kinematics/bimanual_m.py b/robocore/kinematics/bimanual_m.py new file mode 100644 index 0000000..34e877c --- /dev/null +++ b/robocore/kinematics/bimanual_m.py @@ -0,0 +1,301 @@ +"""Dual-arm cooperative kinematics. + +Author: Synria Robotics Team +License: GPL-3.0 +""" +from __future__ import annotations +from robocore.utils.backend import get_backend, set_backend +from typing import Dict, Any, Optional, Sequence, List +from dataclasses import dataclass +import numpy as np + +from .jacobian import jacobian as single_jacobian +from .fk import forward_kinematics as single_fk +from .utils import relative_pose_error, relative_jacobian + + +def bimanual_forward_kinematics( + left_model, + right_model, + q_left, + q_right, + *, + return_end: bool = True, + mode: str = 'indep', + device: Any | None = None, + dtype: Any | None = None, +): + """Compute bimanual forward kinematics. + + Supports both single and batch processing. Inputs are automatically detected: + - Single: [n] -> single 4x4 matrix + - Batch: [B, n] -> [B, 4, 4] array + + :param left_model: Left arm RobotModel + :param right_model: Right arm RobotModel + :param q_left: Left joint configuration(s) - [n] or [B, n] + :param q_right: Right joint configuration(s) - [n] or [B, n] + :param return_end: Return only end-effector poses + :param mode: 'indep'|'relative'|'mirror' + :param device: Torch device (only for torch backend) + :param dtype: Torch dtype (only for torch backend) + :return: {'left': T_left, 'right': T_right} or with additional fields + """ + b = get_backend() + if mode == 'indep': + if b == 'numpy': + from robocore.kinematics.fk_utils.bimanual_fk_solver_numpy import BiIndependentFKSolverNumpy + + solver = BiIndependentFKSolverNumpy(left_model, right_model) + return solver.fk(q_left, q_right, return_end=return_end) + elif b == 'torch': + from robocore.kinematics.fk_utils.bimanual_fk_solver_torch import BiIndependentFKSolverTorch + import torch # type: ignore + + solver = BiIndependentFKSolverTorch(left_model, right_model) + if dtype is None: + dtype = torch.float64 + return solver.fk(q_left, q_right, return_end=return_end, device=device, dtype=dtype) + else: + raise ValueError("Unsupported backend, expected 'auto'|'numpy'|'torch'") + elif mode == 'relative': + if b == 'numpy': + from robocore.kinematics.fk_utils.bimanual_fk_solver_numpy import BiRelativeFKSolverNumpy + + solver = BiRelativeFKSolverNumpy(left_model, right_model) + return solver.fk(q_left, q_right, return_end=return_end) + elif b == 'torch': + from robocore.kinematics.fk_utils.bimanual_fk_solver_torch import BiRelativeFKSolverTorch + import torch # type: ignore + + solver = BiRelativeFKSolverTorch(left_model, right_model) + if dtype is None: + dtype = torch.float64 + return solver.fk(q_left, q_right, return_end=return_end, device=device, dtype=dtype) + else: + raise ValueError("Unsupported backend, expected 'auto'|'numpy'|'torch'") + elif mode == 'mirror': + if b == 'numpy': + from robocore.kinematics.fk_utils.bimanual_fk_solver_numpy import BiMirrorFKSolverNumpy + + solver = BiMirrorFKSolverNumpy(left_model, right_model) + return solver.fk(q_left, q_right, return_end=return_end) + elif b == 'torch': + from robocore.kinematics.fk_utils.bimanual_fk_solver_torch import BiMirrorFKSolverTorch + import torch # type: ignore + + solver = BiMirrorFKSolverTorch(left_model, right_model) + if dtype is None: + dtype = torch.float64 + return solver.fk(q_left, q_right, return_end=return_end, device=device, dtype=dtype) + else: + raise ValueError("Unsupported backend, expected 'auto'|'numpy'|'torch'") + else: + raise ValueError("Unknown mode, expected 'indep'|'relative'|'mirror'") + + +def bimanual_inverse_kinematics( + left_model, + right_model, + *, + target_left=None, + target_right=None, + q0_left=None, + q0_right=None, + method: str = 'dls', + coordination: str = 'indep', + backend: str | None = None, + # Initial guess parameters (new system) + num_initial_guesses: int = 1, + initial_guess_strategy: str = 'random', + initial_guess_scale: float = 1.0, + random_seed: Optional[int] = None, + # Optional advanced params + T_rel_grasp=None, + T_left_initial=None, + T_right_initial=None, + return_all: bool = False, + # Torch specific + torch_device: Any | None = None, + torch_dtype: Any | None = None, + **solver_kwargs, +): + """Compute bimanual inverse kinematics. + + Supports both single and batch processing. Inputs are automatically detected: + - Single: 4x4 matrix -> single result + - Batch: [B, 4, 4] -> list of results + + :param left_model: Left arm RobotModel + :param right_model: Right arm RobotModel + :param target_left: Left target pose(s) - 4x4 or [B, 4, 4] + :param target_right: Right target pose(s) - 4x4 or [B, 4, 4] + :param q0_left: Initial left configuration (optional, used as base for strategies) + :param q0_right: Initial right configuration (optional, used as base for strategies) + :param method: 'dls'|'pinv'|'transpose' + :param coordination: 'indep'|'relative_pose'|'relative_pos'|'relative_ori'|'mirror' + :param backend: Optional backend override ('numpy' or 'torch'). If provided, + this is forwarded to the global backend manager so that IK solvers run + on the requested backend (matches usage in higher-level helpers such as + ``BimanualRobotModel.ik`` and interactive demos). + :param num_initial_guesses: Number of initial guesses to try (default: 1) + :param initial_guess_strategy: Strategy - 'zero'|'random'|'sobol'|'latin'|'center'|'uniform' + :param initial_guess_scale: Scale factor for joint limits (0.0 to 1.0) + :param random_seed: Seed for reproducibility + :param T_rel_grasp: Relative grasp transform (for relative modes) + :param T_left_initial: Left initial reference pose (for mirror mode) + :param T_right_initial: Right initial reference pose (for mirror mode) + :param return_all: Return all candidates if available + :param torch_device: Torch device (only for torch backend) + :param torch_dtype: Torch dtype (only for torch backend) + :return: Result dict + """ + # Allow callers (e.g. InteractiveDualArmIK, BimanualRobotModel.ik) to + # explicitly select the backend. If backend is None we keep the existing + # global backend setting. + if backend is not None: + set_backend(backend) + + b = get_backend() + # Prepare IK kwargs with new initial guess system + ik_kwargs = { + 'method': method, + 'num_initial_guesses': num_initial_guesses, + 'initial_guess_strategy': initial_guess_strategy, + 'initial_guess_scale': initial_guess_scale, + 'random_seed': random_seed, + **solver_kwargs, + } + if torch_device is not None: + ik_kwargs['torch_device'] = torch_device + if torch_dtype is not None: + ik_kwargs['torch_dtype'] = torch_dtype + + if coordination == 'indep': + if b == 'numpy': + from robocore.kinematics.ik_utils.bimanual_ik_solver_numpy import BiIndependentIKSolverNumpy + + solver = BiIndependentIKSolverNumpy(left_model, right_model) + return solver.solve(target_left, target_right, q0_left, q0_right, **ik_kwargs) + elif b == 'torch': + from robocore.kinematics.ik_utils.bimanual_ik_solver_torch import BiIndependentIKSolverTorch + solver = BiIndependentIKSolverTorch(left_model, right_model) + return solver.solve(target_left, target_right, q0_left, q0_right, **ik_kwargs) + else: + raise ValueError("Unsupported backend, expected 'auto'|'numpy'|'torch'") + elif coordination in ('relative_pose', 'relative_pos', 'relative_ori'): + constraint_type = 'pose' if coordination == 'relative_pose' else ('position' if coordination == 'relative_pos' else 'orientation') + if b == 'numpy': + from robocore.kinematics.ik_utils.bimanual_ik_solver_numpy import BiRelativeIKSolverNumpy + solver = BiRelativeIKSolverNumpy(left_model, right_model) + return solver.solve(target_left, target_right, q0_left, q0_right, + constraint_type=constraint_type, T_rel_grasp=T_rel_grasp, **ik_kwargs) + elif b == 'torch': + from robocore.kinematics.ik_utils.bimanual_ik_solver_torch import BiRelativeIKSolverTorch + solver = BiRelativeIKSolverTorch(left_model, right_model) + return solver.solve(target_left, target_right, q0_left, q0_right, + constraint_type=constraint_type, T_rel_grasp=T_rel_grasp, **ik_kwargs) + else: + raise ValueError("Unsupported backend, expected 'auto'|'numpy'|'torch'") + elif coordination == 'mirror': + if b == 'numpy': + from robocore.kinematics.ik_utils.bimanual_ik_solver_numpy import BiMirrorIKSolverNumpy + solver = BiMirrorIKSolverNumpy(left_model, right_model) + return solver.solve(target_left, target_right, q0_left, q0_right, + T_left_initial=T_left_initial, T_right_initial=T_right_initial, **ik_kwargs) + elif b == 'torch': + from robocore.kinematics.ik_utils.bimanual_ik_solver_torch import BiMirrorIKSolverTorch + solver = BiMirrorIKSolverTorch(left_model, right_model) + return solver.solve(target_left, target_right, q0_left, q0_right, + T_left_initial=T_left_initial, T_right_initial=T_right_initial, **ik_kwargs) + else: + raise ValueError("Unsupported backend, expected 'auto'|'numpy'|'torch'") + else: + raise ValueError("Unknown coordination mode") + + +def bimanual_jacobian( + left_model, + right_model, + q_left, + q_right, + *, + mode: str = 'indep', + row_mask=None, + device: Any | None = None, + dtype: Any | None = None, +): + """Compute bimanual Jacobian matrix. + + Supports both single and batch processing. Inputs are automatically detected: + - Single: [n] -> [12, nL+nR] or [6, nL+nR] matrix (depending on mode) + - Batch: [B, n] -> [B, 12, nL+nR] or [B, 6, nL+nR] array + + :param left_model: Left arm RobotModel + :param right_model: Right arm RobotModel + :param q_left: Left joint configuration(s) - [n] or [B, n] + :param q_right: Right joint configuration(s) - [n] or [B, n] + :param mode: 'indep'|'relative'|'mirror' + :param row_mask: Optional row selection mask + :param device: Torch device (only for torch backend) + :param dtype: Torch dtype (only for torch backend) + :return: Jacobian matrix - [12, nL+nR] or [6, nL+nR] (single) or [B, 12, nL+nR] / [B, 6, nL+nR] (batch) + """ + b = get_backend() + if mode == 'indep': + if b == 'numpy': + from robocore.kinematics.jacobian_utils.bimanual_jacobian_solver_numpy import BiIndependentJacobianSolverNumpy + import numpy as np + + solver = BiIndependentJacobianSolverNumpy(left_model, right_model) + J = solver.compute(q_left, q_right) + if row_mask is not None: + mask = np.array([bool(m) for m in row_mask]) + J = J[mask, :] + return J + elif b == 'torch': + from robocore.kinematics.jacobian_utils.bimanual_jacobian_solver_torch import BiIndependentJacobianSolverTorch + import torch # type: ignore + + solver = BiIndependentJacobianSolverTorch(left_model, right_model) + J = solver.compute(q_left, q_right) + if row_mask is not None: + mask = torch.tensor([bool(m) for m in row_mask], dtype=torch.bool) + J = J[mask, :] + return J + else: + raise ValueError("Unsupported backend, expected 'auto'|'numpy'|'torch'") + elif mode == 'relative': + if b == 'numpy': + from robocore.kinematics.jacobian_utils.bimanual_jacobian_solver_numpy import BiRelativeJacobianSolverNumpy + import numpy as np + + solver = BiRelativeJacobianSolverNumpy(left_model, right_model) + J = solver.compute(q_left, q_right) + if row_mask is not None: + mask = np.array([bool(m) for m in row_mask]) + J = J[mask, :] + return J + elif b == 'torch': + from robocore.kinematics.jacobian_utils.bimanual_jacobian_solver_torch import BiRelativeJacobianSolverTorch + import torch # type: ignore + + solver = BiRelativeJacobianSolverTorch(left_model, right_model) + J = solver.compute(q_left, q_right) + if row_mask is not None: + mask = torch.tensor([bool(m) for m in row_mask], dtype=torch.bool) + J = J[mask, :] + return J + else: + raise ValueError("Unsupported backend, expected 'auto'|'numpy'|'torch'") + elif mode == 'mirror': + raise NotImplementedError("Jacobian mode not implemented yet: mirror") + + +__all__ = [ + 'relative_pose_error', + 'relative_jacobian', + 'bimanual_forward_kinematics', + 'bimanual_jacobian', + 'bimanual_inverse_kinematics', +] diff --git a/robocore/kinematics/ik_utils/ik_solver_numpy.py b/robocore/kinematics/ik_utils/ik_solver_numpy.py index aa108c7..cd32143 100644 --- a/robocore/kinematics/ik_utils/ik_solver_numpy.py +++ b/robocore/kinematics/ik_utils/ik_solver_numpy.py @@ -140,6 +140,8 @@ def solve( if target_pose.shape[0] != q0.shape[0]: raise ValueError(f"Batch size mismatch: target_pose {target_pose.shape[0]} vs q0 {q0.shape[0]}") + q0 = self._apply_joint_limits_batch(q0) + result = self._solve_batch( target_pose, q0, pos_weight=pos_weight, @@ -467,6 +469,8 @@ def _solve_single( if q0.shape[0] != self.n: raise ValueError(f"Expected q0 with {self.n} elements, got {q0.shape[0]}") + + q0 = self._apply_joint_limits(q0) # Extract target position and rotation R_target = target_pose[:3, :3] @@ -918,3 +922,13 @@ def _apply_joint_limits(self, q: np.ndarray) -> np.ndarray: if js.limit_upper is not None: q_clamped[js.index] = min(js.limit_upper, q_clamped[js.index]) return q_clamped + + def _apply_joint_limits_batch(self, q_batch: np.ndarray) -> np.ndarray: + """Clamp a batch of joint configurations to limits.""" + q_clamped = q_batch.copy() + for js in self.model._chain_actuated: + if js.limit_lower is not None: + q_clamped[:, js.index] = np.maximum(js.limit_lower, q_clamped[:, js.index]) + if js.limit_upper is not None: + q_clamped[:, js.index] = np.minimum(js.limit_upper, q_clamped[:, js.index]) + return q_clamped diff --git a/robocore/kinematics/ik_utils/ik_solver_torch.py b/robocore/kinematics/ik_utils/ik_solver_torch.py index dbc8b1a..090ec9e 100644 --- a/robocore/kinematics/ik_utils/ik_solver_torch.py +++ b/robocore/kinematics/ik_utils/ik_solver_torch.py @@ -166,6 +166,8 @@ def solve( if target_pose.shape[0] != q0.shape[0]: raise ValueError(f"Batch size mismatch: target_pose {target_pose.shape[0]} vs q0 {q0.shape[0]}") + q0 = self._apply_joint_limits(q0) + # Common kwargs for _solve_single single_kwargs = dict( method=method, @@ -259,6 +261,7 @@ def _solve_single( torch.manual_seed(random_seed) base_q0 = q0.clone() if torch.is_tensor(q0) else torch.tensor(q0, dtype=self.dtype, device=self.device) + base_q0 = self._apply_joint_limits(base_q0) if base_q0.numel() != self.n: raise ValueError(f"q0 size {base_q0.numel()} != dof {self.n}") @@ -612,10 +615,16 @@ def _apply_joint_limits(self, q: Tensor) -> Tensor: out = q.clone() for js in self.model._chain_actuated: # type: ignore[attr-defined] if js.limit_lower is not None or js.limit_upper is not None: - if js.limit_lower is not None: - out[js.index] = torch.clamp(out[js.index], min=float(js.limit_lower)) - if js.limit_upper is not None: - out[js.index] = torch.clamp(out[js.index], max=float(js.limit_upper)) + if out.ndim == 1: + if js.limit_lower is not None: + out[js.index] = torch.clamp(out[js.index], min=float(js.limit_lower)) + if js.limit_upper is not None: + out[js.index] = torch.clamp(out[js.index], max=float(js.limit_upper)) + else: + if js.limit_lower is not None: + out[:, js.index] = torch.clamp(out[:, js.index], min=float(js.limit_lower)) + if js.limit_upper is not None: + out[:, js.index] = torch.clamp(out[:, js.index], max=float(js.limit_upper)) return out # ==================== Partial FK (single) ==================== @@ -695,7 +704,7 @@ def _solve_batch( :return: dict with q, success, iters, method, pos_err, ori_err """ B, n = q_init_batch.shape - q = q_init_batch.clone() + q = self._apply_joint_limits(q_init_batch.clone()) best_q = q.clone() best_err = torch.full((B,), float('inf'), dtype=self.dtype, device=self.device) success = torch.zeros(B, dtype=torch.bool, device=self.device) diff --git a/robocore/modeling/parser/mjcf_parser/wrapper.py b/robocore/modeling/parser/mjcf_parser/wrapper.py index 67b597d..ec73c22 100644 --- a/robocore/modeling/parser/mjcf_parser/wrapper.py +++ b/robocore/modeling/parser/mjcf_parser/wrapper.py @@ -55,8 +55,12 @@ def __init__(self, mjcf_path: str | Path): self.model = mujoco.MjModel.from_xml_path(str(self.mjcf_path)) self.data = mujoco.MjData(self.model) - # Parse joint information (requires self.model to be loaded) - self.joints = self._parse_joints() + # Parse joint information (requires self.model to be loaded). + # Keep two views: + # - actuated_joints: true MuJoCo joints that consume q entries + # - joints: full body-to-body kinematic tree including fixed segments + self.actuated_joints = self._parse_actuated_joints() + self.joints = self._build_kinematic_joints() self.link_mesh_map = {} def _get_joint_type(self, jnt_type: int) -> str: @@ -72,85 +76,115 @@ def _get_joint_type(self, jnt_type: int) -> str: } return type_map.get(jnt_type, "fixed") - def _parse_joints(self) -> List[JointSpec]: + def _body_name(self, body_id: int) -> str: + body_name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_BODY, body_id) + if body_name is None: + return f"body_{body_id}" + return body_name + + def _body_origin(self, body_id: int) -> tuple[list[float], list[float]]: + # MuJoCo stores a body's pose relative to its parent body. + origin_xyz = self.model.body_pos[body_id].copy().tolist() + quat = self.model.body_quat[body_id].copy() + origin_rpy = quaternion_to_rpy(quaternion_reorder(quat)) + return origin_xyz, origin_rpy + + def _parse_actuated_joints(self) -> List[JointSpec]: """ - :return: List of JointSpec objects + :return: List of MuJoCo joints that contribute configuration DOF """ - joints = [] - idx = 0 + joints: List[JointSpec] = [] for i in range(self.model.njnt): - # Get joint name joint_name = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_JOINT, i) if joint_name is None: joint_name = f"joint_{i}" - # Get joint type jnt_type = self.model.jnt_type[i] joint_type = self._get_joint_type(jnt_type) - # Get body IDs body_id = self.model.jnt_bodyid[i] parent_body_id = self.model.body_parentid[body_id] - # Get body names - child = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_BODY, body_id) - if child is None: - child = f"body_{body_id}" - - parent = mujoco.mj_id2name(self.model, mujoco.mjtObj.mjOBJ_BODY, parent_body_id) - if parent is None: - parent = f"body_{parent_body_id}" - - # Get joint axis (in body frame) - axis = self.model.jnt_axis[i].copy() - axis = axis.tolist() + parent = self._body_name(parent_body_id) + child = self._body_name(body_id) + axis = self.model.jnt_axis[i].copy().tolist() + origin_xyz, origin_rpy = self._body_origin(body_id) - # Get body's position and orientation relative to parent - # MuJoCo stores body_pos as position relative to parent body - # and body_quat as quaternion relative to parent body - pos = self.model.body_pos[body_id].copy() - origin_xyz = pos.tolist() - - # Get body's relative quaternion and convert to RPY - # MuJoCo uses [w, x, y, z] format - quat = self.model.body_quat[body_id].copy() - origin_rpy = quaternion_to_rpy(quaternion_reorder(quat)) - - # Get joint limits limit_lower = None limit_upper = None if self.model.jnt_limited[i]: limit_lower = float(self.model.jnt_range[i, 0]) limit_upper = float(self.model.jnt_range[i, 1]) - joint = JointSpec( - name=joint_name, - index=idx, - joint_type=joint_type, - parent=parent, - child=child, - axis=axis, - origin_xyz=origin_xyz, - origin_rpy=origin_rpy, - limit_lower=limit_lower, - limit_upper=limit_upper + joints.append( + JointSpec( + name=joint_name, + index=len(joints), + joint_type=joint_type, + parent=parent, + child=child, + axis=axis, + origin_xyz=origin_xyz, + origin_rpy=origin_rpy, + limit_lower=limit_lower, + limit_upper=limit_upper, + ) + ) + + return joints + + def _build_kinematic_joints(self) -> List[JointSpec]: + """ + :return: Full kinematic tree including fixed body-to-body segments + """ + joints: List[JointSpec] = [] + body_to_actuated: Dict[int, List[JointSpec]] = {} + + for joint in self.actuated_joints: + body_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_BODY, joint.child) + body_to_actuated.setdefault(body_id, []).append(joint) + + # Skip body 0 ("world"), add one edge per child body. + for body_id in range(1, self.model.nbody): + parent_body_id = self.model.body_parentid[body_id] + parent = self._body_name(parent_body_id) + child = self._body_name(body_id) + + explicit_joints = body_to_actuated.get(body_id) + if explicit_joints: + joints.extend(explicit_joints) + continue + + origin_xyz, origin_rpy = self._body_origin(body_id) + joints.append( + JointSpec( + name=f"{parent}_to_{child}_fixed", + index=-1, + joint_type="fixed", + parent=parent, + child=child, + axis=[0.0, 0.0, 0.0], + origin_xyz=origin_xyz, + origin_rpy=origin_rpy, + limit_lower=None, + limit_upper=None, + ) ) - joints.append(joint) return joints def get_joint_names(self) -> List[JointSpec]: """ :return: List of JointSpec objects """ - return self.joints + return self.actuated_joints def get_joint_limits(self) -> np.ndarray: """ :return: Array of joint limits with shape (n_joints, 2) where each row is [lower, upper] """ - return np.array([(j.limit_lower, j.limit_upper) for j in self.joints]) + return np.array([(j.limit_lower, j.limit_upper) for j in self.actuated_joints]) def get_link_names(self) -> List[str]: """