Package xarm

2024-03-29 13:59:09 +01:00
parent dfa6cc777f
commit 3713a3a87b
35 changed files with 1591 additions and 0 deletions
--- a/envs/sim_xarm/xarm/tasks/mocap.py
+++ b/envs/sim_xarm/xarm/tasks/mocap.py
@@ -0,0 +1,67 @@
+# import mujoco_py
+import mujoco
+import numpy as np
+
+
+def apply_action(model, model_names, data, action):
+    if model.nmocap > 0:
+        pos_action, gripper_action = np.split(action, (model.nmocap * 7,))
+        if data.ctrl is not None:
+            for i in range(gripper_action.shape[0]):
+                data.ctrl[i] = gripper_action[i]
+        pos_action = pos_action.reshape(model.nmocap, 7)
+        pos_delta, quat_delta = pos_action[:, :3], pos_action[:, 3:]
+        reset_mocap2body_xpos(model, model_names, data)
+        data.mocap_pos[:] = data.mocap_pos + pos_delta
+        data.mocap_quat[:] = data.mocap_quat + quat_delta
+
+
+def reset(model, data):
+    if model.nmocap > 0 and model.eq_data is not None:
+        for i in range(model.eq_data.shape[0]):
+            # if sim.model.eq_type[i] == mujoco_py.const.EQ_WELD:
+            if model.eq_type[i] == mujoco.mjtEq.mjEQ_WELD:
+                # model.eq_data[i, :] = np.array([0., 0., 0., 1., 0., 0., 0.])
+                model.eq_data[i, :] = np.array(
+                    [
+                        0.0,
+                        0.0,
+                        0.0,
+                        1.0,
+                        0.0,
+                        0.0,
+                        0.0,
+                        0.0,
+                        0.0,
+                        0.0,
+                        0.0,
+                    ]
+                )
+    # sim.forward()
+    mujoco.mj_forward(model, data)
+
+
+def reset_mocap2body_xpos(model, model_names, data):
+    if model.eq_type is None or model.eq_obj1id is None or model.eq_obj2id is None:
+        return
+
+    # For all weld constraints
+    for eq_type, obj1_id, obj2_id in zip(model.eq_type, model.eq_obj1id, model.eq_obj2id, strict=False):
+        # if eq_type != mujoco_py.const.EQ_WELD:
+        if eq_type != mujoco.mjtEq.mjEQ_WELD:
+            continue
+        # body2 = model.body_id2name(obj2_id)
+        body2 = model_names.body_id2name[obj2_id]
+        if body2 == "B0" or body2 == "B9" or body2 == "B1":
+            continue
+        mocap_id = model.body_mocapid[obj1_id]
+        if mocap_id != -1:
+            # obj1 is the mocap, obj2 is the welded body
+            body_idx = obj2_id
+        else:
+            # obj2 is the mocap, obj1 is the welded body
+            mocap_id = model.body_mocapid[obj2_id]
+            body_idx = obj1_id
+        assert mocap_id != -1
+        data.mocap_pos[mocap_id][:] = data.xpos[body_idx]
+        data.mocap_quat[mocap_id][:] = data.xquat[body_idx]