def reset_mocap2body_xpos()

in gym_xarm/tasks/mocap.py [0:0]

• 18 lines of code
• 11 McCabe index (conditional complexity)

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]