Systematically discovers joint offsets (qpos, qvel) for each object.

Notes:
- This code assumes unnamed joints but relies on joint order; if you prefer robustness, explicitly name each joint (e.g. <joint name="puck_joint" type="free"/>) and update the mj_name2id() calls accordingly.
- joint_qpos_addrs and joint_qvel_addrs are general lists that will support arbitrary numbers of bodies with free joints.
- For each object, the simulation handles push dynamics, orientation extraction, and saves a structured state and rendered GIF.

tray_sim.py

@@ -34,6 +34,29 @@ def run_tray_simulation(
     # Explicitly reset the simulation state
     model = mujoco.MjModel.from_xml_path(MODEL_PATH)
     data = mujoco.MjData(model)
+    mujoco.mj_resetData(model, data)
+
+    # Object naming convention: puck, obj0, obj1, ...
+    object_names = ["puck"] + [f"obj{i}" for i in range(num_objects - 1)]
+
+    # Validate all objects exist
+    for name in object_names:
+        if mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, name) == -1:
+            raise ValueError(f"Body '{name}' not found in model")
+
+    # Resolve joint offsets
+    joint_qpos_addrs = []
+    joint_qvel_addrs = []
+    for name in object_names:
+        joint_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_JOINT, f"{name}_joint")  # assumes <joint> has no name
+        if joint_id == -1:
+            # fallback: assume exactly one joint per object, unnamed
+            joint_id = sum(1 for i in range(model.njnt)
+                           if model.jnt_type[i] == mujoco.mjtJoint.mjJNT_FREE)
+            if len(joint_qpos_addrs) >= joint_id:
+                raise ValueError(f"Failed to resolve joint for body '{name}'")
+        joint_qpos_addrs.append(model.jnt_qposadr[joint_id])
+        joint_qvel_addrs.append(model.jnt_dofadr[joint_id])
 
     #tray_id = mujoco.mj_name2id(
     #    model,
@@ -48,37 +71,43 @@ def run_tray_simulation(
     # Compute DOF offsets
     
     # Verify the number of bodies (puck + tray)
-    expected_bodies = num_objects + 1  # Puck + tray
-    actual_bodies = sum(1 for i in range(model.nbody) if mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, i))
-    if actual_bodies < expected_bodies:
-        raise ValueError(f"Model has fewer bodies ({actual_bodies}) than expected ({expected_bodies})")
+    #expected_bodies = num_objects + 1  # Puck + tray
+    #actual_bodies = sum(1 for i in range(model.nbody) if mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, i))
+    #if actual_bodies < expected_bodies:
+    #    raise ValueError(f"Model has fewer bodies ({actual_bodies}) than expected ({expected_bodies})")
 
-    mujoco.mj_resetData(model, data)
 
     # Find puck body ID
-    puck_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "puck")
-    if puck_id == -1:
-        raise ValueError("Puck body not found in model")
+    #puck_id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, "puck")
+    #if puck_id == -1:
+    #    raise ValueError("Puck body not found in model")
 
-    print(f"Puck body id {puck_id} pos at start {data.xpos[puck_id]}")
+    #print(f"Puck body id {puck_id} pos at start {data.xpos[puck_id]}")
 
     # Puck (pusher) setup (indices for single object)
-    puck_idx = 0  # First body
-    puck_vel_idx = 0  # First body velocities
+    ##puck_idx = 0  # First body
+    ##puck_vel_idx = 0  # First body velocities
 
     # Verify qpos and qvel sizes
-    if puck_idx + 7 > model.nq:
-        raise ValueError(f"Puck position index {puck_idx} out of bounds for qpos size {model.nq}")
-    if puck_vel_idx + 6 > model.nv:
-        raise ValueError(f"Puck velocity index {puck_vel_idx} out of bounds for qvel size {model.nv}")
+    #if puck_idx + 7 > model.nq:
+    #    raise ValueError(f"Puck position index {puck_idx} out of bounds for qpos size {model.nq}")
+    #if puck_vel_idx + 6 > model.nv:
+    #    raise ValueError(f"Puck velocity index {puck_vel_idx} out of bounds for qvel size {model.nv}")
 
     # Puck position and orientation
+    #puck_pos = np.array([0.35, -0.35, 0.08])
+    #puck_quat = [1, 0, 0, 0]
+
+    # Set initial position and orientation for the puck
     puck_pos = np.array([0.35, -0.35, 0.08])
     puck_quat = [1, 0, 0, 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
+    
     #data.qpos[puck_idx:puck_idx+3] = puck_pos
     #data.qpos[puck_idx+3:puck_idx+7] = puck_quat
 
+    # Push velocity directed toward center
     # Compute velocity towards center
     target = np.array([0.0, 0.0, 0.05])  # Slightly outside tray corner
     direction = target - puck_pos
@@ -107,7 +136,6 @@ def run_tray_simulation(
             data.qvel[puck_vel_idx:puck_vel_idx+3] = push_velocity
 
         mujoco.mj_step(model, data)
-
         renderer.update_scene(data, camera=cam, scene_option=scene_option)
         img = renderer.render()
 
@@ -118,6 +146,7 @@ def run_tray_simulation(
 
     renderer.close()
 
+    # Orientation extraction
     euler_angles = []
     for i in range(num_objects):
         quat = data.qpos[i*7+3:i*7+7]  # qw, qx, qy, qz
@@ -129,15 +158,18 @@ def run_tray_simulation(
     # Fully structured snapshot of: Positions, 
     # Orientations (in both quaternion and Euler angle forms),
     # Linear and angular velocities
+    qpos = np.array([data.qpos[addr:addr+7] for addr in joint_qpos_addrs])
+    qvel = np.array([data.qvel[addr:addr+6] for addr in joint_qvel_addrs])
     physics_state = {
-        "positions_xyz": data.qpos[:num_objects * 7].reshape(num_objects, 7)[:, :3].tolist(),
-        "orientations_quat": data.qpos[:num_objects * 7].reshape(num_objects, 7)[:, 3:].tolist(),
+        "positions_xyz": qpos[:, :3].tolist(),
+        "orientations_quat": qpos[:, 3:].tolist(),
         "orientations_euler": euler_angles,
-        "velocities_linear": data.qvel[:num_objects * 6].reshape(num_objects, 6)[:, :3].tolist(),
-        "velocities_angular": data.qvel[:num_objects * 6].reshape(num_objects, 6)[:, 3:].tolist(),
+        "velocities_linear": qvel[:, :3].tolist(),
+        "velocities_angular": qvel[:, 3:].tolist(),
     }
     # Save JSON Snapshot with unique ID (optional, for LLM input/debugging)
     json_path = os.path.join(tempfile.gettempdir(), f"tray_sim_{unique_id}_state.json")
+
     with open(json_path, "w") as f:
         json.dump(physics_state, f, indent=2)
 
@@ -146,10 +178,9 @@ def run_tray_simulation(
     imageio.mimsave(gif_path, frames, fps=20)
 
     print(f"Bodies: {model.nbody}, qpos size: {model.nq}, qvel size: {model.nv}")
-
     for i in range(model.nbody):
         name = mujoco.mj_id2name(model, mujoco.mjtObj.mjOBJ_BODY, i)
         if name:
-            print(f"Body {i}: {name}")
+            print(f"Body {i}: {name}, xpos: {data.xpos[i]}")
 
     return gif_path, physics_state, json_path