Question regarding the init state for a kinematically coupled entity

[gokulmk-12]

Hi,

Thanks for introducing such an exciting project in mujoco, I am really loving it.

I am working on a system of 2 quadrupeds, connected kinematically with a cuboidal payload of certain fixed mass. In mujoco xml, i use the equality connect constraint between sites, defined on the quadruped and the payload. I define this in mjlab, by starting with a simple payload.xml, then using mjspec to add 2 quadruped bodies and including two equality constraints. Now to generate the entitycfg for the RL task, i have to define the init_state, but I am kinda stuck there. How do i define init state for such a coupled system, which is acting as a single entity ?

Specifically in this code, how do i incorporate the states for the other quadruped and the payload ?

INIT_STATE = EntityCfg.InitialStateCfg(
    pos=(0.0, 0.0, 0.32),
    joint_pos={
        ".*thigh_joint": 0.9,
        ".*calf_joint": -1.8,
        ".*R_hip_joint": 0.1,
        ".*L_hip_joint": -0.1,
    },
    joint_vel={".*": 0.0},
)

Thanks in advance !!!

[kevinzakka]

Hey, thanks for using mjlab!

The way to do this is with 3 separate entities (2 quadrupeds + 1 payload), each with its own EntityCfg and InitialStateCfg. Then use the scene-level spec_fn to add the equality connect constraints between them.

When entities are added to a scene, their element names get prefixed with the entity name. So if your scene config is:

SceneCfg(
  entities={
    "quad1": quad_cfg,
    "quad2": quad_cfg,
    "payload": payload_cfg,
  },
  spec_fn=add_constraints,
)

You can reference cross-entity sites in the spec_fn:

def add_constraints(spec: mujoco.MjSpec) -> None:
  spec.add_equality(
    type=mujoco.mjtEq.mjEQ_CONNECT,
    objtype=mujoco.mjtObj.mjOBJ_SITE,
    name1="quad1/attach_site",
    name2="payload/left_site",
  )
  spec.add_equality(
    type=mujoco.mjtEq.mjEQ_CONNECT,
    objtype=mujoco.mjtObj.mjOBJ_SITE,
    name1="quad2/attach_site",
    name2="payload/right_site",
  )

Each entity gets its own InitialStateCfg, so you set pos and joint_pos independently for each quadruped and the payload. With site-based connect constraints, MuJoCo will snap the sites together even if they don't perfectly overlap in the initial configuration, so the positions don't need to be exact.

[gokulmk-12]

Hey, thanks for the idea, it worked. I was having a hard time, editing MjSpec to get this done, but this approach is great and simple.