68 lines
2.4 KiB
Python
68 lines
2.4 KiB
Python
|
# 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]
|