From 9a3958d1e9c56f6be5d8b12d096da942aeae38c5 Mon Sep 17 00:00:00 2001 From: wilsonchenghy Date: Sat, 19 Jul 2025 13:30:04 -0400 Subject: [PATCH] Add Simulation Kinematics Scripts --- .../robot_arm_controllers/joint_space_test.py | 156 +++++++++++++++ .../robot_arm_controllers/task_space_test.py | 187 ++++++++++++++++++ .../shadow_hand_joint_space_test.py | 139 +++++++++++++ 3 files changed, 482 insertions(+) create mode 100644 autonomy/simulation/Wato_additional_scripts/robot_arm_controllers/joint_space_test.py create mode 100644 autonomy/simulation/Wato_additional_scripts/robot_arm_controllers/task_space_test.py create mode 100644 autonomy/simulation/Wato_additional_scripts/shadow_hand/shadow_hand_joint_space_test.py diff --git a/autonomy/simulation/Wato_additional_scripts/robot_arm_controllers/joint_space_test.py b/autonomy/simulation/Wato_additional_scripts/robot_arm_controllers/joint_space_test.py new file mode 100644 index 0000000..06bc1b8 --- /dev/null +++ b/autonomy/simulation/Wato_additional_scripts/robot_arm_controllers/joint_space_test.py @@ -0,0 +1,156 @@ +import argparse +from isaaclab.app import AppLauncher + +parser = argparse.ArgumentParser(description="Robot Arm with Joint Space Control") +AppLauncher.add_app_launcher_args(parser) +args_cli = parser.parse_args() + +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +import torch + +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg +from isaaclab.controllers import DifferentialIKController, DifferentialIKControllerCfg +from isaaclab.managers import SceneEntityCfg +from isaaclab.markers import VisualizationMarkers +from isaaclab.markers.config import FRAME_MARKER_CFG +from isaaclab.scene import InteractiveScene, InteractiveSceneCfg +from isaaclab.utils import configclass +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.math import subtract_frame_transforms + +## +# Pre-defined configs +## +from isaaclab_assets import UR10_CFG + + +@configclass +class TableTopSceneCfg(InteractiveSceneCfg): + """Configuration for a cart-pole scene.""" + + # ground plane + ground = AssetBaseCfg( + prim_path="/World/defaultGroundPlane", + spawn=sim_utils.GroundPlaneCfg(), + init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, -1.05)), + ) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", + spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)) + ) + + # mount + table = AssetBaseCfg( + prim_path="{ENV_REGEX_NS}/Table", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/Stand/stand_instanceable.usd", scale=(2.0, 2.0, 2.0) + ), + ) + + # articulation + robot = UR10_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + +def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): + + robot = scene["robot"] + + # Create controller + diff_ik_cfg = DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="dls") + diff_ik_controller = DifferentialIKController(diff_ik_cfg, num_envs=scene.num_envs, device=sim.device) + + # Markers + frame_marker_cfg = FRAME_MARKER_CFG.copy() + frame_marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1) + ee_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_current")) + + # Specify robot-specific parameters + robot_entity_cfg = SceneEntityCfg("robot", joint_names=[".*"], body_names=["ee_link"]) + + # Resolving the scene entities + robot_entity_cfg.resolve(scene) + + # Obtain the frame index of the end-effector ; For a fixed base robot, the frame index is one less than the body index. This is because ; the root body is not included in the returned Jacobians. + if robot.is_fixed_base: + ee_jacobi_idx = robot_entity_cfg.body_ids[0] - 1 + else: + ee_jacobi_idx = robot_entity_cfg.body_ids[0] + + # Define simulation stepping + sim_dt = sim.get_physics_dt() + count = 0 + joint_position_index = 0 + + while simulation_app.is_running(): + if count % 150 == 0: + # reset time + count = 0 + + # cycle between joint state + joint_pos = robot.data.default_joint_pos.clone() + joint_position_list = [ + [0.0, -0.712, 0.712, 0.0, 0.0, 0.0], + [0.0, -1.712, 1.712, 0.0, 0.0, 0.0], + [1.0, -1.712, 1.712, 0.0, 0.0, 0.0], + [0.0, 1.712, 1.712, 0.0, 0.0, 0.0], + ] + joint_position = torch.tensor(joint_position_list[joint_position_index], device=sim.device) + joint_vel = robot.data.default_joint_vel.clone() + + if joint_position_index >= len(joint_position_list) - 1: + joint_position_index = 0 + else: + joint_position_index += 1 + + robot.reset() + + joint_pos_des = joint_position.unsqueeze(0)[:, robot_entity_cfg.joint_ids].clone() + + # apply actions (Smooth movement) + robot.set_joint_position_target(joint_pos_des, joint_ids=robot_entity_cfg.joint_ids) + scene.write_data_to_sim() + + # perform step + sim.step() + + # update sim-time + count += 1 + + # update buffers + scene.update(sim_dt) + + # obtain quantities from simulation + ee_pose_w = robot.data.body_state_w[:, robot_entity_cfg.body_ids[0], 0:7] + # update marker positions + ee_marker.visualize(ee_pose_w[:, 0:3], ee_pose_w[:, 3:7]) + + +def main(): + # Load kit helper + sim_cfg = sim_utils.SimulationCfg(dt=0.01, device=args_cli.device) + sim = sim_utils.SimulationContext(sim_cfg) + + # Set main camera + sim.set_camera_view([2.5, 2.5, 2.5], [0.0, 0.0, 0.0]) + + # Design scene + scene_cfg = TableTopSceneCfg(num_envs=1, env_spacing=2.0) + scene = InteractiveScene(scene_cfg) + + # Play the simulator + sim.reset() + + print("[INFO]: Setup complete...") + + # Run the simulator + run_simulator(sim, scene) + + +if __name__ == "__main__": + main() + simulation_app.close() \ No newline at end of file diff --git a/autonomy/simulation/Wato_additional_scripts/robot_arm_controllers/task_space_test.py b/autonomy/simulation/Wato_additional_scripts/robot_arm_controllers/task_space_test.py new file mode 100644 index 0000000..79161cd --- /dev/null +++ b/autonomy/simulation/Wato_additional_scripts/robot_arm_controllers/task_space_test.py @@ -0,0 +1,187 @@ +import argparse +from isaaclab.app import AppLauncher + +parser = argparse.ArgumentParser(description="Robot Arm with Task Space IK Control") +parser.add_argument("--robot", type=str, default="franka_panda", help="Name of the robot.") +AppLauncher.add_app_launcher_args(parser) +args_cli = parser.parse_args() + +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +import torch + +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg +from isaaclab.controllers import DifferentialIKController, DifferentialIKControllerCfg +from isaaclab.managers import SceneEntityCfg +from isaaclab.markers import VisualizationMarkers +from isaaclab.markers.config import FRAME_MARKER_CFG +from isaaclab.scene import InteractiveScene, InteractiveSceneCfg +from isaaclab.utils import configclass +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.math import subtract_frame_transforms + +## +# Pre-defined configs +## +from isaaclab_assets import UR10_CFG, FRANKA_PANDA_HIGH_PD_CFG + + +@configclass +class TableTopSceneCfg(InteractiveSceneCfg): + """Configuration for a cart-pole scene.""" + + # ground plane + ground = AssetBaseCfg( + prim_path="/World/defaultGroundPlane", + spawn=sim_utils.GroundPlaneCfg(), + init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, -1.05)), + ) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", + spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)) + ) + + # mount + table = AssetBaseCfg( + prim_path="{ENV_REGEX_NS}/Table", + spawn=sim_utils.UsdFileCfg( + usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/Stand/stand_instanceable.usd", scale=(2.0, 2.0, 2.0) + ), + ) + + # Cube + cube = AssetBaseCfg( + prim_path="/World/cube", + spawn=sim_utils.CuboidCfg( + size=[0.1, 0.1, 0.1] + ), + init_state=AssetBaseCfg.InitialStateCfg(pos=(0.5, 0.0, 0.5)), + ) + + # articulation + if args_cli.robot == "franka_panda": + robot = FRANKA_PANDA_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + elif args_cli.robot == "ur10": + robot = UR10_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + else: + raise ValueError(f"Robot {args_cli.robot} is not supported. Valid: franka_panda, ur10") + # robot = UR10_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + +def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): + + robot = scene["robot"] + + # Create controller + diff_ik_cfg = DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="dls") + diff_ik_controller = DifferentialIKController(diff_ik_cfg, num_envs=scene.num_envs, device=sim.device) + + # Markers + frame_marker_cfg = FRAME_MARKER_CFG.copy() + frame_marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1) + ee_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_current")) + goal_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_goal")) + + # Specify robot-specific parameters + if args_cli.robot == "franka_panda": + robot_entity_cfg = SceneEntityCfg("robot", joint_names=["panda_joint.*"], body_names=["panda_hand"]) + elif args_cli.robot == "ur10": + robot_entity_cfg = SceneEntityCfg("robot", joint_names=[".*"], body_names=["ee_link"]) + else: + raise ValueError(f"Robot {args_cli.robot} is not supported. Valid: franka_panda, ur10") + # robot_entity_cfg = SceneEntityCfg("robot", joint_names=[".*"], body_names=["ee_link"]) + + # Resolving the scene entities + robot_entity_cfg.resolve(scene) + + # Obtain the frame index of the end-effector ; For a fixed base robot, the frame index is one less than the body index. This is because ; the root body is not included in the returned Jacobians. + if robot.is_fixed_base: + ee_jacobi_idx = robot_entity_cfg.body_ids[0] - 1 + else: + ee_jacobi_idx = robot_entity_cfg.body_ids[0] + + # Define simulation stepping + sim_dt = sim.get_physics_dt() + + # May have to set initial position first + if args_cli.robot == "franka_panda": + joint_position = robot.data.default_joint_pos.clone() + joint_vel = robot.data.default_joint_vel.clone() + robot.write_joint_state_to_sim(joint_position, joint_vel) + else: + joint_position_index = 0 + joint_position_list = [ + [0.0, 0.0, 0.0, 0.0, 0.0, 0.0], + ] + joint_position = torch.tensor(joint_position_list[joint_position_index], device=sim.device) + joint_vel = robot.data.default_joint_vel.clone() + joint_pos_des = joint_position.unsqueeze(0)[:, robot_entity_cfg.joint_ids].clone() + robot.write_joint_state_to_sim(joint_pos_des, joint_vel) + + while simulation_app.is_running(): + # Get cube/target_point coordinates + position, quaternion = scene["cube"].get_world_poses() + ik_commands = torch.cat([position, quaternion], dim=1) # Quaternion is in (w, x, y, z) + + diff_ik_controller.set_command(ik_commands) + + # obtain quantities from simulation + jacobian = robot.root_physx_view.get_jacobians()[:, ee_jacobi_idx, :, robot_entity_cfg.joint_ids] + ee_pose_w = robot.data.body_state_w[:, robot_entity_cfg.body_ids[0], 0:7] + root_pose_w = robot.data.root_state_w[:, 0:7] + joint_pos = robot.data.joint_pos[:, robot_entity_cfg.joint_ids] + + # compute frame in root frame + ee_pos_b, ee_quat_b = subtract_frame_transforms( + root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7] + ) + + # compute the joint commands + joint_pos_des = diff_ik_controller.compute(ee_pos_b, ee_quat_b, jacobian, joint_pos) + + # apply actions (Smooth movement) + robot.set_joint_position_target(joint_pos_des, joint_ids=robot_entity_cfg.joint_ids) + scene.write_data_to_sim() + + # perform step + sim.step() + + # update buffers + scene.update(sim_dt) + + # obtain quantities from simulation + ee_pose_w = robot.data.body_state_w[:, robot_entity_cfg.body_ids[0], 0:7] + + # update marker positions + ee_marker.visualize(ee_pose_w[:, 0:3], ee_pose_w[:, 3:7]) + goal_marker.visualize(ik_commands[:, 0:3] + scene.env_origins, ik_commands[:, 3:7]) + + +def main(): + # Load kit helper + sim_cfg = sim_utils.SimulationCfg(dt=0.01, device=args_cli.device) + sim = sim_utils.SimulationContext(sim_cfg) + + # Set main camera + sim.set_camera_view([2.5, 2.5, 2.5], [0.0, 0.0, 0.0]) + + # Design scene + scene_cfg = TableTopSceneCfg(num_envs=1, env_spacing=2.0) + scene = InteractiveScene(scene_cfg) + + # Play the simulator + sim.reset() + + print("[INFO]: Setup complete...") + + # Run the simulator + run_simulator(sim, scene) + + +if __name__ == "__main__": + main() + simulation_app.close() \ No newline at end of file diff --git a/autonomy/simulation/Wato_additional_scripts/shadow_hand/shadow_hand_joint_space_test.py b/autonomy/simulation/Wato_additional_scripts/shadow_hand/shadow_hand_joint_space_test.py new file mode 100644 index 0000000..200207c --- /dev/null +++ b/autonomy/simulation/Wato_additional_scripts/shadow_hand/shadow_hand_joint_space_test.py @@ -0,0 +1,139 @@ +import argparse +from isaaclab.app import AppLauncher + +parser = argparse.ArgumentParser(description="Robot Arm with Joint Space Control") +AppLauncher.add_app_launcher_args(parser) +args_cli = parser.parse_args() + +app_launcher = AppLauncher(args_cli) +simulation_app = app_launcher.app + +import torch + +import isaaclab.sim as sim_utils +from isaaclab.assets import AssetBaseCfg +from isaaclab.controllers import DifferentialIKController, DifferentialIKControllerCfg +from isaaclab.managers import SceneEntityCfg +from isaaclab.markers import VisualizationMarkers +from isaaclab.markers.config import FRAME_MARKER_CFG +from isaaclab.scene import InteractiveScene, InteractiveSceneCfg +from isaaclab.utils import configclass +from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR +from isaaclab.utils.math import subtract_frame_transforms + +from isaaclab_assets import SHADOW_HAND_CFG + +# Shadow hand has 24 DOF - found from print(robot_entity_cfg.joint_ids) +# Print joint position - robot.data.default_joint_pos + + +@configclass +class TableTopSceneCfg(InteractiveSceneCfg): + """Configuration for a cart-pole scene.""" + + # ground plane + ground = AssetBaseCfg( + prim_path="/World/defaultGroundPlane", + spawn=sim_utils.GroundPlaneCfg(), + init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, -1.05)), + ) + + # lights + dome_light = AssetBaseCfg( + prim_path="/World/Light", + spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75)) + ) + + # mount + # table = AssetBaseCfg( + # prim_path="{ENV_REGEX_NS}/Table", + # spawn=sim_utils.UsdFileCfg( + # usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/Stand/stand_instanceable.usd", scale=(2.0, 2.0, 2.0) + # ), + # ) + + # articulation + robot = SHADOW_HAND_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot") + + +def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene): + + robot = scene["robot"] + + # Specify robot-specific parameters + robot_entity_cfg = SceneEntityCfg("robot", joint_names=[".*"], body_names=[".*"]) + + # Resolving the scene entities + robot_entity_cfg.resolve(scene) + + # Obtain the frame index of the end-effector ; For a fixed base robot, the frame index is one less than the body index. This is because ; the root body is not included in the returned Jacobians. + if robot.is_fixed_base: + ee_jacobi_idx = robot_entity_cfg.body_ids[0] - 1 + else: + ee_jacobi_idx = robot_entity_cfg.body_ids[0] + + # Define simulation stepping + sim_dt = sim.get_physics_dt() + count = 0 + joint_position_index = 0 + + while simulation_app.is_running(): + if count % 150 == 0: + # reset time + count = 0 + + # cycle between joint state + joint_position_list = [ + [0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], + [1., 1., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0.], + [0., 0., 1., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 0., 1., 1., 0., 0., 0., 0., 0.] + ] + joint_position = torch.tensor(joint_position_list[joint_position_index], device=sim.device) + + if joint_position_index >= len(joint_position_list) - 1: + joint_position_index = 0 + else: + joint_position_index += 1 + + robot.reset() + + joint_pos_des = joint_position.unsqueeze(0)[:, robot_entity_cfg.joint_ids].clone() + + # apply actions (Smooth movement) + robot.set_joint_position_target(joint_pos_des, joint_ids=robot_entity_cfg.joint_ids) + scene.write_data_to_sim() + + # perform step + sim.step() + + # update sim-time + count += 1 + + # update buffers + scene.update(sim_dt) + + +def main(): + # Load kit helper + sim_cfg = sim_utils.SimulationCfg(dt=0.01, device=args_cli.device) + sim = sim_utils.SimulationContext(sim_cfg) + + # Set main camera + sim.set_camera_view([2.5, 2.5, 2.5], [0.0, 0.0, 0.0]) + + # Design scene + scene_cfg = TableTopSceneCfg(num_envs=1, env_spacing=2.0) + scene = InteractiveScene(scene_cfg) + + # Play the simulator + sim.reset() + + print("[INFO]: Setup complete...") + + # Run the simulator + run_simulator(sim, scene) + + +if __name__ == "__main__": + main() + simulation_app.close() \ No newline at end of file