From 14b65db9525cb7f22d014166aa862d72c1d4cc7b Mon Sep 17 00:00:00 2001 From: haixuantao Date: Mon, 23 Jun 2025 16:05:41 +0200 Subject: [PATCH] Add robot descriptions --- examples/urdf/broken_fanuc.yml | 30 +++ examples/urdf/broken_poppy.yml | 30 +++ examples/urdf/franka.yml | 32 +++ examples/urdf/gen3.yml | 30 +++ examples/urdf/kuka.yml | 30 +++ examples/urdf/piper.yml | 32 +++ examples/urdf/so_arm101.yml | 32 +++ examples/urdf/ur5.yml | 30 +++ examples/urdf/z1.yml | 30 +++ .../dora_pytorch_kinematics/main.py | 168 ++++++------- node-hub/dora-rerun/Cargo.toml | 2 +- node-hub/dora-rerun/pyproject.toml | 5 +- node-hub/dora-rerun/src/urdf.rs | 38 ++- node-hub/gamepad/gamepad/main.py | 234 +++++++++--------- 14 files changed, 510 insertions(+), 213 deletions(-) create mode 100644 examples/urdf/broken_fanuc.yml create mode 100644 examples/urdf/broken_poppy.yml create mode 100644 examples/urdf/franka.yml create mode 100644 examples/urdf/gen3.yml create mode 100644 examples/urdf/kuka.yml create mode 100644 examples/urdf/piper.yml create mode 100644 examples/urdf/so_arm101.yml create mode 100644 examples/urdf/ur5.yml create mode 100644 examples/urdf/z1.yml diff --git a/examples/urdf/broken_fanuc.yml b/examples/urdf/broken_fanuc.yml new file mode 100644 index 00000000..469cd0d3 --- /dev/null +++ b/examples/urdf/broken_fanuc.yml @@ -0,0 +1,30 @@ +nodes: + - id: plot + build: pip install -e ../../node-hub/dora-rerun + path: dora-rerun + inputs: + jointstate_m710ic70: pytorch_kinematics/cmd_vel + env: + m710ic70_urdf: "fanuc_m710ic_description" + m710ic70_transform: "0. 0. 0. 1. 0. 0. 0." + + - id: gamepad + build: pip install -e ../../node-hub/gamepad + path: gamepad + outputs: + - cmd_vel + - raw_control + inputs: + tick: dora/timer/millis/10 + + - id: pytorch_kinematics + build: pip install -e ../../node-hub/dora-pytorch-kinematics + path: dora-pytorch-kinematics + inputs: + cmd_vel: gamepad/cmd_vel + outputs: + - cmd_vel + env: + MODEL_NAME: "fanuc_m710ic_description" + END_EFFECTOR_LINK: "tool0" + TRANSFORM: "0. 0. 0. 1. 0. 0. 0." diff --git a/examples/urdf/broken_poppy.yml b/examples/urdf/broken_poppy.yml new file mode 100644 index 00000000..f8c39a81 --- /dev/null +++ b/examples/urdf/broken_poppy.yml @@ -0,0 +1,30 @@ +nodes: + - id: plot + build: pip install -e ../../node-hub/dora-rerun + path: dora-rerun + inputs: + jointstate_poppy_ergo_jr: pytorch_kinematics/cmd_vel + env: + poppy_ergo_jr_urdf: "poppy_ergo_jr_description" + poppy_ergo_jr_transform: "0. 0. 0. 1. 0. 0. 0." + + - id: gamepad + build: pip install -e ../../node-hub/gamepad + path: gamepad + outputs: + - cmd_vel + - raw_control + inputs: + tick: dora/timer/millis/10 + + - id: pytorch_kinematics + build: pip install -e ../../node-hub/dora-pytorch-kinematics + path: dora-pytorch-kinematics + inputs: + cmd_vel: gamepad/cmd_vel + outputs: + - cmd_vel + env: + MODEL_NAME: "poppy_ergo_jr_description" + END_EFFECTOR_LINK: "section_5" + TRANSFORM: "0. 0. 0. 1. 0. 0. 0." diff --git a/examples/urdf/franka.yml b/examples/urdf/franka.yml new file mode 100644 index 00000000..4aef1345 --- /dev/null +++ b/examples/urdf/franka.yml @@ -0,0 +1,32 @@ +nodes: + - id: plot + build: pip install -e ../../node-hub/dora-rerun + path: dora-rerun + inputs: + jointstate_panda: pytorch_kinematics/cmd_vel + env: + panda_urdf: "panda_description" + panda_transform: "0. 0. 0. 1. 0. 0. 0." + + - id: gamepad + build: pip install -e ../../node-hub/gamepad + path: gamepad + outputs: + - cmd_vel + - raw_control + inputs: + tick: dora/timer/millis/10 + + - id: pytorch_kinematics + build: pip install -e ../../node-hub/dora-pytorch-kinematics + path: dora-pytorch-kinematics + inputs: + cmd_vel: gamepad/cmd_vel + outputs: + - cmd_vel + env: + MODEL_NAME: "panda_description" + END_EFFECTOR_LINK: "panda_link8" + TRANSFORM: "0. 0. 0. 1. 0. 0. 0." + POSITION_TOLERANCE: 0.01 + ROTATION_TOLERANCE: 0.01 diff --git a/examples/urdf/gen3.yml b/examples/urdf/gen3.yml new file mode 100644 index 00000000..71a8a1e4 --- /dev/null +++ b/examples/urdf/gen3.yml @@ -0,0 +1,30 @@ +nodes: + - id: plot + build: pip install -e ../../node-hub/dora-rerun + path: dora-rerun + inputs: + jointstate_jaco: pytorch_kinematics/cmd_vel + env: + jaco_urdf: "gen3_description" + jaco_transform: "0. 0. 0. 1. 0. 0. 0." + + - id: gamepad + build: pip install -e ../../node-hub/gamepad + path: gamepad + outputs: + - cmd_vel + - raw_control + inputs: + tick: dora/timer/millis/10 + + - id: pytorch_kinematics + build: pip install -e ../../node-hub/dora-pytorch-kinematics + path: dora-pytorch-kinematics + inputs: + cmd_vel: gamepad/cmd_vel + outputs: + - cmd_vel + env: + MODEL_NAME: "gen3_description" + END_EFFECTOR_LINK: "j2n6s300_end_effector" + TRANSFORM: "0. 0. 0. 1. 0. 0. 0." diff --git a/examples/urdf/kuka.yml b/examples/urdf/kuka.yml new file mode 100644 index 00000000..a7668be7 --- /dev/null +++ b/examples/urdf/kuka.yml @@ -0,0 +1,30 @@ +nodes: + - id: plot + build: pip install -e ../../node-hub/dora-rerun + path: dora-rerun + inputs: + jointstate_iiwa14_primitive_collision: pytorch_kinematics/cmd_vel + env: + iiwa14_primitive_collision_urdf: "iiwa14_description" + iiwa14_primitive_collision_transform: "0. 0. 0. 1. 0. 0. 0." + + - id: gamepad + build: pip install -e ../../node-hub/gamepad + path: gamepad + outputs: + - cmd_vel + - raw_control + inputs: + tick: dora/timer/millis/10 + + - id: pytorch_kinematics + build: pip install -e ../../node-hub/dora-pytorch-kinematics + path: dora-pytorch-kinematics + inputs: + cmd_vel: gamepad/cmd_vel + outputs: + - cmd_vel + env: + MODEL_NAME: "iiwa14_description" + END_EFFECTOR_LINK: "iiwa_link_7" + TRANSFORM: "0. 0. 0. 1. 0. 0. 0." diff --git a/examples/urdf/piper.yml b/examples/urdf/piper.yml new file mode 100644 index 00000000..e2ca4385 --- /dev/null +++ b/examples/urdf/piper.yml @@ -0,0 +1,32 @@ +nodes: + - id: plot + build: pip install -e ../../node-hub/dora-rerun + path: dora-rerun + inputs: + jointstate_piper_description: pytorch_kinematics/cmd_vel + env: + piper_description_urdf: "piper_description" + piper_description_transform: "0. 0. 0. 1. 0. 0. 0." + + - id: gamepad + build: pip install -e ../../node-hub/gamepad + path: gamepad + outputs: + - cmd_vel + - raw_control + inputs: + tick: dora/timer/millis/10 + + - id: pytorch_kinematics + build: pip install -e ../../node-hub/dora-pytorch-kinematics + path: dora-pytorch-kinematics + inputs: + cmd_vel: gamepad/cmd_vel + outputs: + - cmd_vel + env: + MODEL_NAME: "piper_description" + END_EFFECTOR_LINK: "link6" + TRANSFORM: "0. 0. 0. 1. 0. 0. 0." + POSITION_TOLERANCE: 0.001 + ROTATION_TOLERANCE: 0.001 diff --git a/examples/urdf/so_arm101.yml b/examples/urdf/so_arm101.yml new file mode 100644 index 00000000..d61e5ef4 --- /dev/null +++ b/examples/urdf/so_arm101.yml @@ -0,0 +1,32 @@ +nodes: + - id: plot + build: pip install -e ../../node-hub/dora-rerun + path: dora-rerun + inputs: + jointstate_so101_new_calib: pytorch_kinematics/cmd_vel + env: + so101_new_calib_urdf: "so_arm101_description" + so101_new_calib_transform: "0. 0. 0. 1. 0. 0. 0." + + - id: gamepad + build: pip install -e ../../node-hub/gamepad + path: gamepad + outputs: + - cmd_vel + - raw_control + inputs: + tick: dora/timer/millis/10 + + - id: pytorch_kinematics + build: pip install -e ../../node-hub/dora-pytorch-kinematics + path: dora-pytorch-kinematics + inputs: + cmd_vel: gamepad/cmd_vel + outputs: + - cmd_vel + env: + MODEL_NAME: "so_arm101_description" + END_EFFECTOR_LINK: "gripper" + TRANSFORM: "0. 0. 0. 1. 0. 0. 0." + POSITION_TOLERANCE: 0.01 + ROTATION_TOLERANCE: 0.01 diff --git a/examples/urdf/ur5.yml b/examples/urdf/ur5.yml new file mode 100644 index 00000000..cf00a4a4 --- /dev/null +++ b/examples/urdf/ur5.yml @@ -0,0 +1,30 @@ +nodes: + - id: plot + build: pip install -e ../../node-hub/dora-rerun + path: dora-rerun + inputs: + jointstate_ur5_robot: pytorch_kinematics/cmd_vel + env: + ur5_robot_urdf: "ur5_description" + ur5_robot_transform: "0. 0. 0. 1. 0. 0. 0." + + - id: gamepad + build: pip install -e ../../node-hub/gamepad + path: gamepad + outputs: + - cmd_vel + - raw_control + inputs: + tick: dora/timer/millis/10 + + - id: pytorch_kinematics + build: pip install -e ../../node-hub/dora-pytorch-kinematics + path: dora-pytorch-kinematics + inputs: + cmd_vel: gamepad/cmd_vel + outputs: + - cmd_vel + env: + MODEL_NAME: "ur5_description" + END_EFFECTOR_LINK: "tool0" + TRANSFORM: "0. 0. 0. 1. 0. 0. 0." diff --git a/examples/urdf/z1.yml b/examples/urdf/z1.yml new file mode 100644 index 00000000..bdfd8a0d --- /dev/null +++ b/examples/urdf/z1.yml @@ -0,0 +1,30 @@ +nodes: + - id: plot + build: pip install -e ../../node-hub/dora-rerun + path: dora-rerun + inputs: + jointstate_z1: pytorch_kinematics/cmd_vel + env: + z1_urdf: "z1_description" + z1_transform: "0. 0. 0. 1. 0. 0. 0." + + - id: gamepad + build: pip install -e ../../node-hub/gamepad + path: gamepad + outputs: + - cmd_vel + - raw_control + inputs: + tick: dora/timer/millis/10 + + - id: pytorch_kinematics + build: pip install -e ../../node-hub/dora-pytorch-kinematics + path: dora-pytorch-kinematics + inputs: + cmd_vel: gamepad/cmd_vel + outputs: + - cmd_vel + env: + MODEL_NAME: "z1_description" + END_EFFECTOR_LINK: "link05" + TRANSFORM: "0. 0. 0. 1. 0. 0. 0." diff --git a/node-hub/dora-pytorch-kinematics/dora_pytorch_kinematics/main.py b/node-hub/dora-pytorch-kinematics/dora_pytorch_kinematics/main.py index bd5d5ec5..63ae2b7e 100644 --- a/node-hub/dora-pytorch-kinematics/dora_pytorch_kinematics/main.py +++ b/node-hub/dora-pytorch-kinematics/dora_pytorch_kinematics/main.py @@ -11,6 +11,8 @@ import torch from dora import Node from pytorch_kinematics.transforms.rotation_conversions import matrix_to_euler_angles +POSITION_TOLERANCE = float(os.getenv("POSITION_TOLERANCE", "0.005")) +ROTATION_TOLERANCE = float(os.getenv("ROTATION_TOLERANCE", "0.05")) # in radians TRANSFORM = np.array(os.getenv("TRANSFORM", "0. 0. 0. 1. 0. 0. 0.").split(" ")).astype( np.float32, ) # wxyz format @@ -229,18 +231,21 @@ class RobotKinematics: # Instantiate and run the IK solver (core pytorch_kinematics objects/methods) ik_solver = pk.PseudoInverseIK( self.chain, - max_iterations=1_000, + max_iterations=100_000, retry_configs=q_init, joint_limits=torch.tensor(self.chain.get_joint_limits()), early_stopping_any_converged=True, - early_stopping_no_improvement="all", + # early_stopping_no_improvement="all", debug=False, lr=0.05, - rot_tolerance=1e-4, - pos_tolerance=1e-3, + rot_tolerance=ROTATION_TOLERANCE, + pos_tolerance=POSITION_TOLERANCE, ) solution_angles = ik_solver.solve(target_pose) - if solution_angles.err_rot > 1e-3 or solution_angles.err_pos > 1e-2: + if ( + solution_angles.err_rot > ROTATION_TOLERANCE + or solution_angles.err_pos > POSITION_TOLERANCE + ): print( f"IK did not converge: pos_err={solution_angles.err_pos}, rot_err={solution_angles.err_rot} for target {target_pose}", ) @@ -357,99 +362,96 @@ def main(): robot = RobotKinematics( urdf_path=model, urdf="", end_effector_link=end_effector_link ) - last_known_state = None + last_known_state = robot._default_q for event in node: if event["type"] == "INPUT": metadata = event["metadata"] if event["id"] == "cmd_vel": - if last_known_state is not None: - target_vel = event["value"].to_numpy() # expect 100ms - # Apply Forward Kinematics - target = robot.compute_fk(last_known_state) - target = ( - np.array(get_xyz_rpy_array_from_transform3d(target)) - + target_vel - ) - target = pa.array(target.ravel(), type=pa.float32()) - target = pk.Transform3d( - pos=target[:3], - rot=pk.transforms.euler_angles_to_matrix( - torch.tensor(target[3:6]), - convention="XYZ", - ), + target_vel = event["value"].to_numpy() # expect 100ms + # Apply Forward Kinematics + target = robot.compute_fk(last_known_state) + target = ( + np.array(get_xyz_rpy_array_from_transform3d(target).detach()) + + target_vel / 10 + ) + target = pa.array(target.ravel(), type=pa.float32()) + target = pk.Transform3d( + pos=target[:3], + rot=pk.transforms.euler_angles_to_matrix( + torch.tensor(target[3:6]), + convention="XYZ", + ), + ) + rob_target = ROB_TF.inverse().compose(target) + solution = robot.compute_ik(rob_target, last_known_state) + if solution is None: + print( + "No IK Solution for :", + target, + "skipping this frame.", ) - rob_target = ROB_TF.inverse().compose(target) - solution = robot.compute_ik(rob_target, last_known_state) - if solution is None: - print( - "No IK Solution for :", - target, - "skipping this frame.", - ) - continue - solution = solution.numpy().ravel() - metadata["encoding"] = "jointstate" - last_known_state = solution - solution = pa.array(last_known_state) - node.send_output(event["id"], solution, metadata=metadata) + continue + solution = solution.numpy().ravel() + metadata["encoding"] = "jointstate" + last_known_state = solution + solution = pa.array(last_known_state) + node.send_output(event["id"], solution, metadata=metadata) else: match metadata["encoding"]: case "xyzquat": # Apply Inverse Kinematics - if last_known_state is not None: - target = event["value"].to_numpy() - target = target.astype(np.float32) - target = pk.Transform3d( - pos=target[:3], - rot=torch.tensor(target[3:7]), - ) - rob_target = ROB_TF.inverse().compose(target) - solution = robot.compute_ik(rob_target, last_known_state) - metadata["encoding"] = "jointstate" - last_known_state = solution.numpy().ravel() - solution = pa.array(last_known_state) - node.send_output(event["id"], solution, metadata=metadata) + target = event["value"].to_numpy() + target = target.astype(np.float32) + target = pk.Transform3d( + pos=target[:3], + rot=torch.tensor(target[3:7]), + ) + rob_target = ROB_TF.inverse().compose(target) + solution = robot.compute_ik(rob_target, last_known_state) + metadata["encoding"] = "jointstate" + last_known_state = solution.numpy().ravel() + solution = pa.array(last_known_state) + node.send_output(event["id"], solution, metadata=metadata) case "xyzrpy": # Apply Inverse Kinematics - if last_known_state is not None: - target = event["value"].to_numpy() - target = target.astype(np.float32) - target = pk.Transform3d( - pos=target[:3], - rot=pk.transforms.euler_angles_to_matrix( - torch.tensor(target[3:6]), - convention="XYZ", - ), + target = event["value"].to_numpy() + target = target.astype(np.float32) + target = pk.Transform3d( + pos=target[:3], + rot=pk.transforms.euler_angles_to_matrix( + torch.tensor(target[3:6]), + convention="XYZ", + ), + ) + rob_target = ROB_TF.inverse().compose(target) + solution = robot.compute_ik(rob_target, last_known_state) + if solution is None: + print( + "No IK Solution for :", + target, + "skipping this frame.", ) - rob_target = ROB_TF.inverse().compose(target) - solution = robot.compute_ik(rob_target, last_known_state) - if solution is None: - print( - "No IK Solution for :", - target, - "skipping this frame.", - ) - continue - - solution = solution.numpy().ravel() - delta_angles = ( - solution - last_known_state[: len(solution)] - ) # match with dof - - valid = np.all( - (delta_angles >= -np.pi) & (delta_angles <= np.pi), + continue + + solution = solution.numpy().ravel() + delta_angles = ( + solution - last_known_state[: len(solution)] + ) # match with dof + + valid = np.all( + (delta_angles >= -np.pi) & (delta_angles <= np.pi), + ) + if not valid: + print( + "IK solution is not valid, as the rotation are too wide. skipping.", ) - if not valid: - print( - "IK solution is not valid, as the rotation are too wide. skipping.", - ) - continue - metadata["encoding"] = "jointstate" - last_known_state = solution - solution = pa.array(last_known_state) - node.send_output(event["id"], solution, metadata=metadata) + continue + metadata["encoding"] = "jointstate" + last_known_state = solution + solution = pa.array(last_known_state) + node.send_output(event["id"], solution, metadata=metadata) case "jointstate": target = event["value"].to_numpy() last_known_state = target diff --git a/node-hub/dora-rerun/Cargo.toml b/node-hub/dora-rerun/Cargo.toml index 84c075e9..d777d312 100644 --- a/node-hub/dora-rerun/Cargo.toml +++ b/node-hub/dora-rerun/Cargo.toml @@ -17,7 +17,7 @@ python = ["pyo3"] dora-node-api = { workspace = true, features = ["tracing"] } eyre = "0.6.8" tokio = { version = "1.24.2", features = ["rt"] } -rerun = { version = "0.23.1", features = ["web_viewer", "image"] } +rerun = { version = "0.23.3", features = ["web_viewer", "image"] } ndarray = "0.15.6" k = "0.32" pyo3 = { workspace = true, features = [ diff --git a/node-hub/dora-rerun/pyproject.toml b/node-hub/dora-rerun/pyproject.toml index 42ec1643..778822da 100644 --- a/node-hub/dora-rerun/pyproject.toml +++ b/node-hub/dora-rerun/pyproject.toml @@ -8,7 +8,7 @@ dynamic = ["version"] license = { text = "MIT" } requires-python = ">=3.10" -dependencies = ["rerun-loader-urdf", 'rerun_sdk==0.23.1'] +dependencies = ["rerun-loader-urdf", 'rerun_sdk>=0.23.1', "robot-descriptions"] scripts = { "dora-rerun" = "dora_rerun:py_main" } @@ -28,4 +28,5 @@ extend-select = [ ] [tool.uv.sources] -rerun-loader-urdf = { git = "https://github.com/haixuanTao/rerun-loader-python-example-urdf.git", branch = "patch-2" } +rerun-loader-urdf = { git = "https://github.com/dora-rs/rerun-loader-python-urdf.git" } +robot-descriptions = { git = "https://github.com/robot-descriptions/robot_descriptions.py.git" } diff --git a/node-hub/dora-rerun/src/urdf.rs b/node-hub/dora-rerun/src/urdf.rs index 0543b029..18d77313 100644 --- a/node-hub/dora-rerun/src/urdf.rs +++ b/node-hub/dora-rerun/src/urdf.rs @@ -1,4 +1,4 @@ -use std::{collections::HashMap, path::PathBuf}; +use std::{collections::HashMap, fmt::format, path::PathBuf}; use eyre::{Context, ContextCompat, Result}; use k::{nalgebra::Quaternion, Chain, Translation3, UnitQuaternion}; @@ -60,23 +60,39 @@ pub fn init_urdf(rec: &RecordingStream) -> Result>> { .collect::>(); let mut chains = HashMap::new(); for (key, urdf_path) in urdfs { - let path = key.replace("_urdf", ".urdf").replace("_URDF", ".urdf"); - let chain = k::Chain::::from_urdf_file(&urdf_path).context("Could not load URDF")?; + let urdf_path = if urdf_path.ends_with("_description") { + // Use robot description to get its path + let bash_cmd = format!("robot_descriptions pull {urdf_path}"); + let response = std::process::Command::new("bash") + .arg("-c") + .arg(bash_cmd) + .output() + .context("Failed to execute bash command. Are you sure `robot_descriptions` is installed?")? + .stdout; + let response_str = + String::from_utf8(response).context("Could not parse robot descriptions")?; + PathBuf::from(response_str.trim()) + } else { + // Use the path directly + PathBuf::from(urdf_path) + }; + let chain = k::Chain::::from_urdf_file(&urdf_path) + .context(format!("Could not load URDF {:#?}", urdf_path))?; + let path = key.replace("_urdf", ".urdf").replace("_URDF", ".urdf"); let transform = key.replace("_urdf", "_transform"); if PathBuf::from(&urdf_path).file_name() != PathBuf::from(&path).file_name() { return Err(eyre::eyre!( - "URDF path should be the same as the key without _urdf or _URDF. Got {} instead of {}", urdf_path, path + "URDF filename should be the same as the environment variable name and replacing the dot with a dash. Got {:#?} instead of {}", urdf_path, path )); } - if let Err(err) = rec.log_file_from_path(&urdf_path, None, false) { - println!("Could not log file: {}. Errored with {}", urdf_path, err); - println!("Make sure to install urdf loader with:"); - println!( - "pip install git+https://github.com/rerun-io/rerun-loader-python-example-urdf.git" - ) - }; + rec.log_file_from_path(&urdf_path, None, true) + .context(format!( + "Could not log URDF file {:#?} within rerun-urdf-loader", + urdf_path + ))?; + println!("Logging URDF file: {:#?}", urdf_path); // Get transform by replacing URDF_ with TRANSFORM_ if let Ok(transform) = std::env::var(transform) { diff --git a/node-hub/gamepad/gamepad/main.py b/node-hub/gamepad/gamepad/main.py index ad6385ea..e65f6d40 100644 --- a/node-hub/gamepad/gamepad/main.py +++ b/node-hub/gamepad/gamepad/main.py @@ -1,149 +1,151 @@ """Gamepad controller node for Dora.""" -from dora import Node -import pygame -import pyarrow as pa import json +import pyarrow as pa +import pygame +from dora import Node + + class Controller: """Controller mapping.""" def __init__(self): """Change this according to your controller mapping. Currently Logitech F710.""" self.axisNames = { - 'LEFT-X': 0, - 'LEFT-Y': 1, - 'RIGHT-X': 2, - 'RIGHT-Y': 3, + "LEFT-X": 0, + "LEFT-Y": 1, + "RIGHT-X": 2, + "RIGHT-Y": 3, } self.buttonNames = { - 'X': 0, - 'A': 1, - 'B': 2, - 'Y': 3, - 'LB': 4, - 'RB': 5, - 'LT': 6, - 'RT': 7, - 'BACK': 8, - 'START': 9, - 'LEFT-STICK': 10, - 'RIGHT-STICK': 11, + "X": 0, + "A": 1, + "B": 2, + "Y": 3, + "LB": 4, + "RB": 5, + "LT": 6, + "RT": 7, + "BACK": 8, + "START": 9, + "LEFT-STICK": 10, + "RIGHT-STICK": 11, } self.hatIndex = 0 # Index of the D-pad hat + def main(): node = Node("gamepad") - + pygame.init() pygame.joystick.init() - - if pygame.joystick.get_count() == 0: - print("No gamepad found! Please connect your controller.") - return - + + assert pygame.joystick.get_count(), ( + "No gamepad found! Please connect your controller." + ) + joystick = pygame.joystick.Joystick(0) joystick.init() controller = Controller() - + move_speed = 0.05 # Fixed increment for D-pad max_linear_z_speed = 0.1 # Maximum linear Z speed max_angular_speed = 0.8 # Maximum angular speed - + print(f"Detected controller: {joystick.get_name()}") print(f"Number of axes: {joystick.get_numaxes()}") print(f"Number of buttons: {joystick.get_numbuttons()}") print("Press Ctrl+C to exit") - try: - for event in node: - pygame.event.pump() - - # Get all controller states - axes = [joystick.get_axis(i) for i in range(joystick.get_numaxes())] - buttons = [joystick.get_button(i) for i in range(joystick.get_numbuttons())] - - # Get hat state (D-pad) - dpad_x, dpad_y = 0, 0 - if joystick.get_numhats() > 0: - dpad_x, dpad_y = joystick.get_hat(controller.hatIndex) - - # Create raw control state - raw_control = { - "axes": axes, - "buttons": buttons, - "hats": [[dpad_x, dpad_y]], - "mapping": { - "axes": controller.axisNames, - "buttons": controller.buttonNames - } - } - - # cmd_vel processing: - # 1. D-pad vertical for X axis - # 2. D-pad horizontal for Y axis - # 3. Right stick vertical for Z axis - # 4. Right stick horizontal for rotation around Z - # 5. Left stick vertical for rotation around X - # 6. Left stick horizontal for rotation around Y - - deadzone = 0.05 - - # Linear X velocity from D-pad vertical - linear_x = 0.0 - if dpad_y != 0: - linear_x = dpad_y * move_speed - - # Linear Y velocity from D-pad horizontal - linear_y = 0.0 - if dpad_x != 0: - linear_y = dpad_x * move_speed - - # Linear Z velocity from right stick vertical - right_y = -joystick.get_axis(controller.axisNames['RIGHT-Y']) - right_y = 0.0 if abs(right_y) < deadzone else right_y - linear_z = right_y * max_linear_z_speed - - # Angular Z velocity (rotation) from right stick horizontal - right_x = -joystick.get_axis(controller.axisNames['RIGHT-X']) - right_x = 0.0 if abs(right_x) < deadzone else right_x - angular_z = 0 * max_angular_speed # TODO: Make z non zero, but on my gamepad the value is never zero - - # Angular X velocity from left stick vertical - left_y = -joystick.get_axis(controller.axisNames['LEFT-Y']) - left_y = 0.0 if abs(left_y) < deadzone else left_y - angular_x = left_y * max_angular_speed - - # Angular Y velocity from left stick horizontal - left_x = -joystick.get_axis(controller.axisNames['LEFT-X']) - left_x = 0.0 if abs(left_x) < deadzone else left_x - angular_y = left_x * max_angular_speed - - cmd_vel = [linear_x, linear_y, linear_z, angular_x, angular_y, angular_z] - if any(cmd_vel): - node.send_output( - output_id="cmd_vel", - data=pa.array(cmd_vel, type=pa.float64()), - metadata={"type": "cmd_vel"} - ) - - node.send_output( - output_id="raw_control", - data=pa.array([json.dumps(raw_control)], type=pa.string()), - metadata={"type": "raw_control"} - ) - - except KeyboardInterrupt: - print("\nExiting...") - finally: - pygame.quit() - zero_cmd = [0.0] * 6 - node.send_output( - output_id="cmd_vel", - data=pa.array(zero_cmd, type=pa.float64()), - metadata={"type": "cmd_vel"} - ) + for event in node: + pygame.event.pump() + + # Get all controller states + axes = [joystick.get_axis(i) for i in range(joystick.get_numaxes())] + buttons = [joystick.get_button(i) for i in range(joystick.get_numbuttons())] + + # Get hat state (D-pad) + dpad_x, dpad_y = 0, 0 + if joystick.get_numhats() > 0: + dpad_x, dpad_y = joystick.get_hat(controller.hatIndex) + + # Create raw control state + raw_control = { + "axes": axes, + "buttons": buttons, + "hats": [[dpad_x, dpad_y]], + "mapping": { + "axes": controller.axisNames, + "buttons": controller.buttonNames, + }, + } + + # cmd_vel processing: + # 1. D-pad vertical for X axis + # 2. D-pad horizontal for Y axis + # 3. Right stick vertical for Z axis + # 4. Right stick horizontal for rotation around Z + # 5. Left stick vertical for rotation around X + # 6. Left stick horizontal for rotation around Y + + deadzone = 0.05 + + # Linear X velocity from D-pad vertical + linear_x = 0.0 + if dpad_y != 0: + linear_x = dpad_y * move_speed + + # Linear Y velocity from D-pad horizontal + linear_y = 0.0 + if dpad_x != 0: + linear_y = dpad_x * move_speed + + # Linear Z velocity from right stick vertical + right_y = -joystick.get_axis(controller.axisNames["RIGHT-Y"]) + right_y = 0.0 if abs(right_y) < deadzone else right_y + linear_z = right_y * max_linear_z_speed + + # Angular Z velocity (rotation) from right stick horizontal + right_x = -joystick.get_axis(controller.axisNames["RIGHT-X"]) + right_x = 0.0 if abs(right_x) < deadzone else right_x + angular_z = ( + right_x * max_angular_speed + ) # TODO: Make z non zero, but on my gamepad the value is never zero + + # Angular X velocity from left stick vertical + left_y = -joystick.get_axis(controller.axisNames["LEFT-Y"]) + left_y = 0.0 if abs(left_y) < deadzone else left_y + angular_x = left_y * max_angular_speed + + # Angular Y velocity from left stick horizontal + left_x = -joystick.get_axis(controller.axisNames["LEFT-X"]) + left_x = 0.0 if abs(left_x) < deadzone else left_x + angular_y = left_x * max_angular_speed + + cmd_vel = [linear_x, linear_y, linear_z, angular_x, angular_y, angular_z] + if any(cmd_vel): + node.send_output( + output_id="cmd_vel", + data=pa.array(cmd_vel, type=pa.float64()), + metadata={"type": "cmd_vel"}, + ) + + node.send_output( + output_id="raw_control", + data=pa.array([json.dumps(raw_control)], type=pa.string()), + metadata={"type": "raw_control"}, + ) + + pygame.quit() + zero_cmd = [0.0] * 6 + node.send_output( + output_id="cmd_vel", + data=pa.array(zero_cmd, type=pa.float64()), + metadata={"type": "cmd_vel"}, + ) + if __name__ == "__main__": - main() \ No newline at end of file + main()