Import upstream snapshot c19f713c415a699a79d71cd96aa13c3104a05047

Upstream: https://github.com/michaelgillett/mjlab
Upstream-Commit: c19f713c415a699a79d71cd96aa13c3104a05047
Upstream-Branch: main

scripts/demos/differential_ik.py

@@ -0,0 +1,194 @@
+"""Interactive IK control demo.
+
+Drag the 3D transform control in the viser viewer to move the YAM end-effector.
+
+Run with:
+  MJLAB_WARP_QUIET=1 uv run scripts/demos/differential_ik.py
+"""
+
+from __future__ import annotations
+
+import time
+from types import SimpleNamespace
+
+import torch
+import viser
+
+from mjlab.asset_zoo.robots.i2rt_yam.yam_constants import get_yam_robot_cfg
+from mjlab.entity import Entity, EntityCfg
+from mjlab.envs.mdp.actions import DifferentialIKAction, DifferentialIKActionCfg
+from mjlab.sim.sim import MujocoCfg, Simulation, SimulationCfg
+from mjlab.utils.lab_api.math import quat_from_matrix
+from mjlab.viewer.viser import ViserMujocoScene
+
+DEMO_INIT_STATE = EntityCfg.InitialStateCfg(
+  pos=(0.0, 0.0, 0.01),
+  joint_pos={
+    "joint2": 0.6,
+    "joint3": 0.6,
+    "joint4": 0.0,
+    "left_finger": 0.037,
+    "right_finger": -0.037,
+  },
+  joint_vel={".*": 0.0},
+)
+
+IK_ITERATIONS = 10
+
+
+def main() -> None:
+  device = "cuda:0" if torch.cuda.is_available() else "cpu"
+
+  robot_cfg = get_yam_robot_cfg()
+  robot_cfg.init_state = DEMO_INIT_STATE
+  entity = Entity(robot_cfg)
+  model = entity.compile()
+  sim_cfg = SimulationCfg(mujoco=MujocoCfg(gravity=(0, 0, -9.81)))
+  sim = Simulation(num_envs=1, cfg=sim_cfg, model=model, device=device)
+  entity.initialize(model, sim.model, sim.data, device)
+  entity.write_joint_position_to_sim(entity.data.default_joint_pos, joint_ids=None)
+  sim.forward()
+
+  env = SimpleNamespace(num_envs=1, device=device, scene={"robot": entity}, sim=sim)
+  ik_cfg = DifferentialIKActionCfg(
+    entity_name="robot",
+    actuator_names=("joint.*",),
+    frame_name="grasp_site",
+    frame_type="site",
+    posture_weight=0.02,
+    joint_limit_weight=1e-1,
+    damping=1e-1,
+    use_relative_mode=False,
+  )
+  ik_action: DifferentialIKAction = ik_cfg.build(env)  # type: ignore[arg-type]
+  joint_ids = ik_action._joint_ids
+
+  grip_ids, _ = entity.find_joints("left_finger")
+  grip_joint_ids = torch.tensor(grip_ids, device=device, dtype=torch.long)
+  grip_open = torch.tensor([[0.037]], device=device)
+
+  server = viser.ViserServer(label="IK Control Demo")
+  scene = ViserMujocoScene(server, sim.mj_model, num_envs=1)
+  scene.create_scene_gui(
+    camera_distance=0.1,
+    camera_azimuth=135.0,
+    camera_elevation=30.0,
+  )
+
+  site_id = ik_action._frame_id
+  pos = sim.data.site_xpos[0, site_id].cpu().numpy()
+  xmat = sim.data.site_xmat[0, site_id]
+  quat = quat_from_matrix(xmat).cpu().numpy()
+
+  transform_ctrl = server.scene.add_transform_controls(
+    "/ik_target",
+    position=(float(pos[0]), float(pos[1]), float(pos[2])),
+    wxyz=(float(quat[0]), float(quat[1]), float(quat[2]), float(quat[3])),
+    scale=0.12,
+  )
+
+  needs_reset = [False]
+
+  with server.gui.add_folder("IK Control"):
+    reset_button = server.gui.add_button("Reset")
+    reset_button.on_click(lambda _: needs_reset.__setitem__(0, True))
+    iterations_slider = server.gui.add_slider(
+      "IK Iterations",
+      min=1,
+      max=50,
+      step=1,
+      initial_value=IK_ITERATIONS,
+    )
+
+  with server.gui.add_folder("IK Weights"):
+    damping_slider = server.gui.add_slider(
+      "Damping (λ)",
+      min=1e-2,
+      max=1.0,
+      step=1e-3,
+      initial_value=ik_cfg.damping,
+    )
+    pos_w_slider = server.gui.add_slider(
+      "Position Weight",
+      min=0.0,
+      max=10.0,
+      step=0.1,
+      initial_value=ik_cfg.position_weight,
+    )
+    ori_w_slider = server.gui.add_slider(
+      "Orientation Weight",
+      min=0.0,
+      max=10.0,
+      step=0.1,
+      initial_value=ik_cfg.orientation_weight,
+    )
+    jlim_w_slider = server.gui.add_slider(
+      "Joint Limit Weight",
+      min=0.0,
+      max=1.0,
+      step=0.01,
+      initial_value=ik_cfg.joint_limit_weight,
+    )
+    posture_w_slider = server.gui.add_slider(
+      "Posture Weight",
+      min=0.0,
+      max=1.0,
+      step=0.01,
+      initial_value=ik_cfg.posture_weight,
+    )
+
+  print("=" * 60)
+  print("IK Control Demo")
+  print("  Open the viser URL printed above")
+  print("  Drag the 3D transform control to move the end-effector")
+  print("=" * 60)
+
+  target_action = torch.zeros(1, 7, device=device)
+
+  def _reset() -> None:
+    entity.write_joint_position_to_sim(entity.data.default_joint_pos, joint_ids=None)
+    sim.forward()
+    ik_action.reset()
+    p = sim.data.site_xpos[0, site_id].cpu().numpy()
+    q = quat_from_matrix(sim.data.site_xmat[0, site_id]).cpu().numpy()
+    transform_ctrl.position = (float(p[0]), float(p[1]), float(p[2]))
+    transform_ctrl.wxyz = (float(q[0]), float(q[1]), float(q[2]), float(q[3]))
+
+  try:
+    while True:
+      if needs_reset[0]:
+        needs_reset[0] = False
+        _reset()
+
+      ik_cfg.damping = max(damping_slider.value, 1e-2)
+      ik_cfg.position_weight = max(pos_w_slider.value, 0.0)
+      ik_cfg.orientation_weight = max(ori_w_slider.value, 0.0)
+      ik_cfg.joint_limit_weight = max(jlim_w_slider.value, 0.0)
+      ik_cfg.posture_weight = max(posture_w_slider.value, 0.0)
+
+      p = transform_ctrl.position
+      w = transform_ctrl.wxyz
+      target_action[0, :3] = torch.tensor([p[0], p[1], p[2]], device=device)
+      target_action[0, 3:] = torch.tensor([w[0], w[1], w[2], w[3]], device=device)
+      ik_action.process_actions(target_action)
+
+      n_iter = int(iterations_slider.value)
+      for _ in range(n_iter):
+        dq = ik_action.compute_dq()
+        q = entity.data.joint_pos[:, joint_ids] + dq
+        entity.write_joint_position_to_sim(q, joint_ids=joint_ids)
+        entity.write_joint_position_to_sim(grip_open, joint_ids=grip_joint_ids)
+        sim.forward()
+
+      scene.update(sim.data)
+      if scene.needs_update:
+        scene.refresh_visualization()
+
+      time.sleep(1.0 / 30.0)
+  except KeyboardInterrupt:
+    print("\nShutting down...")
+    server.stop()
+
+
+if __name__ == "__main__":
+  main()