Skip to content

Commit ba75a8a

Browse files
wilsonchenghywilsonchenghy
committed
Add example scripts
1 parent a2b8fd1 commit ba75a8a

File tree

3 files changed

+482
-0
lines changed

3 files changed

+482
-0
lines changed

‎scripts/wato_additional_scripts/robot_arm_controllers/joint_space_test.py

Lines changed: 156 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,156 @@
1+
import argparse
2+
from isaaclab.app import AppLauncher
3+
4+
parser = argparse.ArgumentParser(description="Robot Arm with Joint Space Control")
5+
AppLauncher.add_app_launcher_args(parser)
6+
args_cli = parser.parse_args()
7+
8+
app_launcher = AppLauncher(args_cli)
9+
simulation_app = app_launcher.app
10+
11+
import torch
12+
13+
import isaaclab.sim as sim_utils
14+
from isaaclab.assets import AssetBaseCfg
15+
from isaaclab.controllers import DifferentialIKController, DifferentialIKControllerCfg
16+
from isaaclab.managers import SceneEntityCfg
17+
from isaaclab.markers import VisualizationMarkers
18+
from isaaclab.markers.config import FRAME_MARKER_CFG
19+
from isaaclab.scene import InteractiveScene, InteractiveSceneCfg
20+
from isaaclab.utils import configclass
21+
from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR
22+
from isaaclab.utils.math import subtract_frame_transforms
23+
24+
##
25+
# Pre-defined configs
26+
##
27+
from isaaclab_assets import UR10_CFG
28+
29+
30+
@configclass
31+
class TableTopSceneCfg(InteractiveSceneCfg):
32+
"""Configuration for a cart-pole scene."""
33+
34+
# ground plane
35+
ground = AssetBaseCfg(
36+
prim_path="/World/defaultGroundPlane",
37+
spawn=sim_utils.GroundPlaneCfg(),
38+
init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, -1.05)),
39+
)
40+
41+
# lights
42+
dome_light = AssetBaseCfg(
43+
prim_path="/World/Light",
44+
spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75))
45+
)
46+
47+
# mount
48+
table = AssetBaseCfg(
49+
prim_path="{ENV_REGEX_NS}/Table",
50+
spawn=sim_utils.UsdFileCfg(
51+
usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/Stand/stand_instanceable.usd", scale=(2.0, 2.0, 2.0)
52+
),
53+
)
54+
55+
# articulation
56+
robot = UR10_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
57+
58+
59+
def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene):
60+
61+
robot = scene["robot"]
62+
63+
# Create controller
64+
diff_ik_cfg = DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="dls")
65+
diff_ik_controller = DifferentialIKController(diff_ik_cfg, num_envs=scene.num_envs, device=sim.device)
66+
67+
# Markers
68+
frame_marker_cfg = FRAME_MARKER_CFG.copy()
69+
frame_marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1)
70+
ee_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_current"))
71+
72+
# Specify robot-specific parameters
73+
robot_entity_cfg = SceneEntityCfg("robot", joint_names=[".*"], body_names=["ee_link"])
74+
75+
# Resolving the scene entities
76+
robot_entity_cfg.resolve(scene)
77+
78+
# 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.
79+
if robot.is_fixed_base:
80+
ee_jacobi_idx = robot_entity_cfg.body_ids[0] - 1
81+
else:
82+
ee_jacobi_idx = robot_entity_cfg.body_ids[0]
83+
84+
# Define simulation stepping
85+
sim_dt = sim.get_physics_dt()
86+
count = 0
87+
joint_position_index = 0
88+
89+
while simulation_app.is_running():
90+
if count % 150 == 0:
91+
# reset time
92+
count = 0
93+
94+
# cycle between joint state
95+
joint_pos = robot.data.default_joint_pos.clone()
96+
joint_position_list = [
97+
[0.0, -0.712, 0.712, 0.0, 0.0, 0.0],
98+
[0.0, -1.712, 1.712, 0.0, 0.0, 0.0],
99+
[1.0, -1.712, 1.712, 0.0, 0.0, 0.0],
100+
[0.0, 1.712, 1.712, 0.0, 0.0, 0.0],
101+
]
102+
joint_position = torch.tensor(joint_position_list[joint_position_index], device=sim.device)
103+
joint_vel = robot.data.default_joint_vel.clone()
104+
105+
if joint_position_index >= len(joint_position_list) - 1:
106+
joint_position_index = 0
107+
else:
108+
joint_position_index += 1
109+
110+
robot.reset()
111+
112+
joint_pos_des = joint_position.unsqueeze(0)[:, robot_entity_cfg.joint_ids].clone()
113+
114+
# apply actions (Smooth movement)
115+
robot.set_joint_position_target(joint_pos_des, joint_ids=robot_entity_cfg.joint_ids)
116+
scene.write_data_to_sim()
117+
118+
# perform step
119+
sim.step()
120+
121+
# update sim-time
122+
count += 1
123+
124+
# update buffers
125+
scene.update(sim_dt)
126+
127+
# obtain quantities from simulation
128+
ee_pose_w = robot.data.body_state_w[:, robot_entity_cfg.body_ids[0], 0:7]
129+
# update marker positions
130+
ee_marker.visualize(ee_pose_w[:, 0:3], ee_pose_w[:, 3:7])
131+
132+
133+
def main():
134+
# Load kit helper
135+
sim_cfg = sim_utils.SimulationCfg(dt=0.01, device=args_cli.device)
136+
sim = sim_utils.SimulationContext(sim_cfg)
137+
138+
# Set main camera
139+
sim.set_camera_view([2.5, 2.5, 2.5], [0.0, 0.0, 0.0])
140+
141+
# Design scene
142+
scene_cfg = TableTopSceneCfg(num_envs=1, env_spacing=2.0)
143+
scene = InteractiveScene(scene_cfg)
144+
145+
# Play the simulator
146+
sim.reset()
147+
148+
print("[INFO]: Setup complete...")
149+
150+
# Run the simulator
151+
run_simulator(sim, scene)
152+
153+
154+
if __name__ == "__main__":
155+
main()
156+
simulation_app.close()

‎scripts/wato_additional_scripts/robot_arm_controllers/task_space_test.py

Lines changed: 187 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,187 @@
1+
import argparse
2+
from isaaclab.app import AppLauncher
3+
4+
parser = argparse.ArgumentParser(description="Robot Arm with Task Space IK Control")
5+
parser.add_argument("--robot", type=str, default="franka_panda", help="Name of the robot.")
6+
AppLauncher.add_app_launcher_args(parser)
7+
args_cli = parser.parse_args()
8+
9+
app_launcher = AppLauncher(args_cli)
10+
simulation_app = app_launcher.app
11+
12+
import torch
13+
14+
import isaaclab.sim as sim_utils
15+
from isaaclab.assets import AssetBaseCfg
16+
from isaaclab.controllers import DifferentialIKController, DifferentialIKControllerCfg
17+
from isaaclab.managers import SceneEntityCfg
18+
from isaaclab.markers import VisualizationMarkers
19+
from isaaclab.markers.config import FRAME_MARKER_CFG
20+
from isaaclab.scene import InteractiveScene, InteractiveSceneCfg
21+
from isaaclab.utils import configclass
22+
from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR
23+
from isaaclab.utils.math import subtract_frame_transforms
24+
25+
##
26+
# Pre-defined configs
27+
##
28+
from isaaclab_assets import UR10_CFG, FRANKA_PANDA_HIGH_PD_CFG
29+
30+
31+
@configclass
32+
class TableTopSceneCfg(InteractiveSceneCfg):
33+
"""Configuration for a cart-pole scene."""
34+
35+
# ground plane
36+
ground = AssetBaseCfg(
37+
prim_path="/World/defaultGroundPlane",
38+
spawn=sim_utils.GroundPlaneCfg(),
39+
init_state=AssetBaseCfg.InitialStateCfg(pos=(0.0, 0.0, -1.05)),
40+
)
41+
42+
# lights
43+
dome_light = AssetBaseCfg(
44+
prim_path="/World/Light",
45+
spawn=sim_utils.DomeLightCfg(intensity=3000.0, color=(0.75, 0.75, 0.75))
46+
)
47+
48+
# mount
49+
table = AssetBaseCfg(
50+
prim_path="{ENV_REGEX_NS}/Table",
51+
spawn=sim_utils.UsdFileCfg(
52+
usd_path=f"{ISAAC_NUCLEUS_DIR}/Props/Mounts/Stand/stand_instanceable.usd", scale=(2.0, 2.0, 2.0)
53+
),
54+
)
55+
56+
# Cube
57+
cube = AssetBaseCfg(
58+
prim_path="/World/cube",
59+
spawn=sim_utils.CuboidCfg(
60+
size=[0.1, 0.1, 0.1]
61+
),
62+
init_state=AssetBaseCfg.InitialStateCfg(pos=(0.5, 0.0, 0.5)),
63+
)
64+
65+
# articulation
66+
if args_cli.robot == "franka_panda":
67+
robot = FRANKA_PANDA_HIGH_PD_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
68+
elif args_cli.robot == "ur10":
69+
robot = UR10_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
70+
else:
71+
raise ValueError(f"Robot {args_cli.robot} is not supported. Valid: franka_panda, ur10")
72+
# robot = UR10_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
73+
74+
75+
def run_simulator(sim: sim_utils.SimulationContext, scene: InteractiveScene):
76+
77+
robot = scene["robot"]
78+
79+
# Create controller
80+
diff_ik_cfg = DifferentialIKControllerCfg(command_type="pose", use_relative_mode=False, ik_method="dls")
81+
diff_ik_controller = DifferentialIKController(diff_ik_cfg, num_envs=scene.num_envs, device=sim.device)
82+
83+
# Markers
84+
frame_marker_cfg = FRAME_MARKER_CFG.copy()
85+
frame_marker_cfg.markers["frame"].scale = (0.1, 0.1, 0.1)
86+
ee_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_current"))
87+
goal_marker = VisualizationMarkers(frame_marker_cfg.replace(prim_path="/Visuals/ee_goal"))
88+
89+
# Specify robot-specific parameters
90+
if args_cli.robot == "franka_panda":
91+
robot_entity_cfg = SceneEntityCfg("robot", joint_names=["panda_joint.*"], body_names=["panda_hand"])
92+
elif args_cli.robot == "ur10":
93+
robot_entity_cfg = SceneEntityCfg("robot", joint_names=[".*"], body_names=["ee_link"])
94+
else:
95+
raise ValueError(f"Robot {args_cli.robot} is not supported. Valid: franka_panda, ur10")
96+
# robot_entity_cfg = SceneEntityCfg("robot", joint_names=[".*"], body_names=["ee_link"])
97+
98+
# Resolving the scene entities
99+
robot_entity_cfg.resolve(scene)
100+
101+
# 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.
102+
if robot.is_fixed_base:
103+
ee_jacobi_idx = robot_entity_cfg.body_ids[0] - 1
104+
else:
105+
ee_jacobi_idx = robot_entity_cfg.body_ids[0]
106+
107+
# Define simulation stepping
108+
sim_dt = sim.get_physics_dt()
109+
110+
# May have to set initial position first
111+
if args_cli.robot == "franka_panda":
112+
joint_position = robot.data.default_joint_pos.clone()
113+
joint_vel = robot.data.default_joint_vel.clone()
114+
robot.write_joint_state_to_sim(joint_position, joint_vel)
115+
else:
116+
joint_position_index = 0
117+
joint_position_list = [
118+
[0.0, 0.0, 0.0, 0.0, 0.0, 0.0],
119+
]
120+
joint_position = torch.tensor(joint_position_list[joint_position_index], device=sim.device)
121+
joint_vel = robot.data.default_joint_vel.clone()
122+
joint_pos_des = joint_position.unsqueeze(0)[:, robot_entity_cfg.joint_ids].clone()
123+
robot.write_joint_state_to_sim(joint_pos_des, joint_vel)
124+
125+
while simulation_app.is_running():
126+
# Get cube/target_point coordinates
127+
position, quaternion = scene["cube"].get_world_poses()
128+
ik_commands = torch.cat([position, quaternion], dim=1) # Quaternion is in (w, x, y, z)
129+
130+
diff_ik_controller.set_command(ik_commands)
131+
132+
# obtain quantities from simulation
133+
jacobian = robot.root_physx_view.get_jacobians()[:, ee_jacobi_idx, :, robot_entity_cfg.joint_ids]
134+
ee_pose_w = robot.data.body_state_w[:, robot_entity_cfg.body_ids[0], 0:7]
135+
root_pose_w = robot.data.root_state_w[:, 0:7]
136+
joint_pos = robot.data.joint_pos[:, robot_entity_cfg.joint_ids]
137+
138+
# compute frame in root frame
139+
ee_pos_b, ee_quat_b = subtract_frame_transforms(
140+
root_pose_w[:, 0:3], root_pose_w[:, 3:7], ee_pose_w[:, 0:3], ee_pose_w[:, 3:7]
141+
)
142+
143+
# compute the joint commands
144+
joint_pos_des = diff_ik_controller.compute(ee_pos_b, ee_quat_b, jacobian, joint_pos)
145+
146+
# apply actions (Smooth movement)
147+
robot.set_joint_position_target(joint_pos_des, joint_ids=robot_entity_cfg.joint_ids)
148+
scene.write_data_to_sim()
149+
150+
# perform step
151+
sim.step()
152+
153+
# update buffers
154+
scene.update(sim_dt)
155+
156+
# obtain quantities from simulation
157+
ee_pose_w = robot.data.body_state_w[:, robot_entity_cfg.body_ids[0], 0:7]
158+
159+
# update marker positions
160+
ee_marker.visualize(ee_pose_w[:, 0:3], ee_pose_w[:, 3:7])
161+
goal_marker.visualize(ik_commands[:, 0:3] + scene.env_origins, ik_commands[:, 3:7])
162+
163+
164+
def main():
165+
# Load kit helper
166+
sim_cfg = sim_utils.SimulationCfg(dt=0.01, device=args_cli.device)
167+
sim = sim_utils.SimulationContext(sim_cfg)
168+
169+
# Set main camera
170+
sim.set_camera_view([2.5, 2.5, 2.5], [0.0, 0.0, 0.0])
171+
172+
# Design scene
173+
scene_cfg = TableTopSceneCfg(num_envs=1, env_spacing=2.0)
174+
scene = InteractiveScene(scene_cfg)
175+
176+
# Play the simulator
177+
sim.reset()
178+
179+
print("[INFO]: Setup complete...")
180+
181+
# Run the simulator
182+
run_simulator(sim, scene)
183+
184+
185+
if __name__ == "__main__":
186+
main()
187+
simulation_app.close()

0 commit comments

Comments
 (0)