Access attributes programmatically using the MuJoCo Python API

tray_sim.py

@@ -44,6 +44,7 @@ def run_tray_simulation(
         if mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, name) == -1:
             raise ValueError(f"Body '{name}' not found in model")
 
+    # Validate all objects exist and get joint addresses
     # Resolve joint offsets
     joint_qpos_addrs = []
     joint_qvel_addrs = []
@@ -59,6 +60,12 @@ def run_tray_simulation(
         joint_qpos_addrs.append(model.jnt_qposadr[joint_addr])
         joint_qvel_addrs.append(model.jnt_dofadr[joint_addr])
 
+    # Extract puck's initial position and quaternion from the model
+    puck_body_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "puck")
+    puck_pos = model.body_pos[puck_body_id].copy()  # Extract initial position
+    puck_quat = model.body_quat[puck_body_id].copy()  # Extract initial quaternion
+    puck_vel_idx = joint_qvel_addrs[0]
+
     #tray_id = mujoco.mj_name2id(
     #    model,
     #    mujoco.mjtObj.mjOBJ_BODY,
@@ -72,10 +79,6 @@ def run_tray_simulation(
     # Compute DOF offsets
 
     # Set initial position and orientation for the puck
-    puck_pos = np.array([0.0, 0.0, 0.05])
-    puck_quat = [1, 0, 0, 0]
-    puck_vel_idx = joint_qvel_addrs[0]
-
     data.qpos[joint_qpos_addrs[0]:joint_qpos_addrs[0]+3] = puck_pos
     data.qpos[joint_qpos_addrs[0]+3:joint_qpos_addrs[0]+7] = puck_quat