Spaces:
Sleeping
Sleeping
| import os | |
| import numpy as np | |
| from tray_sim import SIMULATOR_MODEL, run_tray_simulation, classify_stability | |
| def main(): | |
| sim = SIMULATOR_MODEL() | |
| sim.create_pusher('3.0 3.0 0.05') | |
| # Tray parameters | |
| physical_parameters_for_object_id_tray = { | |
| 'sliding_friction': 0.1, | |
| 'armature': 0.1, | |
| 'stiffness': 0.0, | |
| 'mass': 0.5, | |
| 'damping': 20 | |
| } | |
| sim.create_tray(object_physics=physical_parameters_for_object_id_tray) | |
| # Object parameters | |
| physical_parameters_for_object_id_1 = { # Bottle | |
| 'sliding_friction': 0.22, | |
| 'armature': 0.41, | |
| 'stiffness': 0.53, | |
| 'mass': 20.0, | |
| 'damping': 7.0 | |
| } | |
| physical_parameters_for_object_id_2 = { # Martini glass | |
| 'sliding_friction': 0.52, | |
| 'armature': 0.63, | |
| 'stiffness': 1.22, | |
| 'mass': 10.0, | |
| 'damping': 8.0 | |
| } | |
| physical_parameters_for_object_id_3 = { # Wine glass | |
| 'sliding_friction': 0.73, | |
| 'armature': 1.04, | |
| 'stiffness': 0.94, | |
| 'mass': 4.0, | |
| 'damping': 7.9 | |
| } | |
| # Create objects | |
| sim.create_object( | |
| object_id=1, | |
| object_name='bottle', | |
| object_location=('row_1', 'column_3'), | |
| object_color='orange', | |
| object_physics=physical_parameters_for_object_id_1 | |
| ) | |
| sim.create_object( | |
| object_id=2, | |
| object_name='martini_glass', | |
| object_location=('row_1', 'column_2'), | |
| object_color='orange', | |
| object_physics=physical_parameters_for_object_id_2 | |
| ) | |
| sim.create_object( | |
| object_id=3, | |
| object_name='wine_glass', | |
| object_location=('row_1', 'column_1'), | |
| object_color='orange', | |
| object_physics=physical_parameters_for_object_id_3 | |
| ) | |
| # Run simulation and collect trajectories | |
| model = sim.model | |
| data = sim.data | |
| trajectories = { | |
| 'bottle': [], | |
| 'martini_glass': [], | |
| 'wine_glass': [] | |
| } | |
| sample_steps = [0, 25, 50, 75, 100, 125, 150, 175, 199] # 9 points as in PAGE13 | |
| for t in range(200): | |
| if t == 50: # PUSH_START_STEP | |
| pusher_idx = 3 * 7 # 3 objects * 7 DOF | |
| pusher_vel_idx = 3 * 6 # 3 objects * 6 velocities | |
| direction = np.array([0.0, 0.0, 0.05]) - np.array([3.0, 3.0, 0.05]) | |
| direction[2] = 0.0 | |
| direction /= np.linalg.norm(direction) | |
| push_velocity = -4.8 * direction / np.sqrt(2) # Normalize to match (-4.8, -4.8) | |
| data.qvel[pusher_vel_idx:pusher_vel_idx+3] = push_velocity | |
| mujoco.mj_step(model, data) | |
| if t in sample_steps: | |
| for i, obj_name in enumerate(['bottle', 'martini_glass', 'wine_glass'], 1): | |
| pos = data.qpos[(i-1)*7:(i-1)*7+3] | |
| # Adjust z to match document's center of gravity | |
| if obj_name == 'bottle': | |
| pos[2] += 1.1 | |
| else: | |
| pos[2] += 0.5 | |
| trajectories[obj_name].append([round(p, 1) for p in pos]) | |
| sim.create_scene() | |
| sim_out = sim.run_simulation() # Ensure scene is finalized | |
| del sim | |
| # Format and save trajectories | |
| with open('object_traj_example_1.txt', 'w') as f: | |
| f.write("bottle_motion_trajectory (x, y, z) = [\n") | |
| for pos in trajectories['bottle']: | |
| f.write(f"({pos[0]}, {pos[1]}, {pos[2]}),\n") | |
| f.write("]\n") | |
| f.write("martini_glass_motion_trajectory (x, y, z) = [\n") | |
| for pos in trajectories['martini_glass']: | |
| f.write(f"({pos[0]}, {pos[1]}, {pos[2]}),\n") | |
| f.write("]\n") | |
| f.write("wine_glass_motion_trajectory (x, y, z) = [\n") | |
| for pos in trajectories['wine_glass']: | |
| f.write(f"({pos[0]}, {pos[1]}, {pos[2]}),\n") | |
| f.write("]\n") | |
| return trajectories | |
| if __name__ == "__main__": | |
| trajectories = main() | |
| print("Trajectories saved to object_traj_example_1.txt") |