Spaces:
Sleeping
Sleeping
File size: 3,970 Bytes
5e35a81 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 |
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") |