Add Simulation Kinematics Scripts

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()

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()