Open
Description
Hello,
I would like to create a RL pipeline in order to find a policy to stack rebars in the most stable way depending on their shape. As I want to stack them, I need to spawn them at one by one. One way (and the easiest for me) would be to use the MjSpec class from mujoco and recompile the model every time I want to spawn a rebar. However, when I wanted to make the environment for the training, I realized that it was not trivial to use this MjSpec feature that I am using with the Physics class...
So my question would be to know if there is an easy way to do it or if I have to rewrite the Environment class?
Thank you
Example
Here is an example of my problem with a simple script.
xml file
<?xml version="1.0" ?>
<mujoco model="coil">
<extension>
<plugin plugin="mujoco.elasticity.cable"/>
</extension>
<option integrator="implicit" timestep="0.001"/>
<visual>
<map force="0.1" zfar="30"/>
<rgba haze="0.15 0.25 0.35 1"/>
<quality shadowsize="2048"/>
<global offwidth="800" offheight="800"/>
</visual>
<worldbody>
<body name="stirrup" pos="0 0 1.005" quat="1 0 0 0">
<freejoint/>
<inertial pos="0.28 0.09 0.005" mass="1" diaginertia="0.001 0.001 0.001"/>
<composite prefix="stirrup0" type="cable" vertex="
3.061617e-19 5.000000e-03 0.000000e+00
5.650000e-01 5.000000e-03 0.000000e+00
5.950000e-01 3.500000e-02 0.000000e+00
5.950000e-01 1.550000e-01 0.000000e+00
5.650000e-01 1.850000e-01 0.000000e+00
-9.184851e-19 1.850000e-01 0.000000e+00
" offset="0 0 0">
<plugin plugin="mujoco.elasticity.cable">
<config key="twist" value="5e8"/>
<config key="bend" value="15e8"/>
<config key="vmax" value="0"/>
</plugin>
<joint kind="main" damping="1" armature="100" stiffness="100" range="0 2"/>
<geom type="capsule" size=".005" rgba=".8 .2 .1 1"/>
</composite>
</body>
</worldbody>
<asset>
<texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512" height="3072"/>
<texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3" markrgb="0.8 0.8 0.8" width="300" height="300"/>
<material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
</asset>
<worldbody>
<light pos="0 0 3.5" dir="0 0 -1" directional="true"/>
<geom name="floor" size="0 0 0.05" type="plane" material="groundplane"/>
</worldbody>
</mujoco>
python code
import sys, os
import time
# Add the parent directory to sys.path
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
import mujoco
import mujoco.viewer
def display_video(frames, videoname='animation.mp4', framerate=30):
"""Display a list of frames as a video.
Args:
frames (list): list of frames to display.
videoname (str, optional): name of the video. Defaults to 'animation.mp4'.
framerate (int, optional): # frames per second. Defaults to 30.
"""
height, width, _ = frames[0].shape
dpi = 10
orig_backend = matplotlib.get_backend()
matplotlib.use('Agg') # Switch to headless 'Agg' to inhibit figure rendering.
fig, ax = plt.subplots(1, 1, figsize=(width / dpi, height / dpi), dpi=dpi)
matplotlib.use(orig_backend) # Switch back to the original backend.
ax.set_axis_off()
ax.set_aspect('equal')
ax.set_position([0, 0, 1, 1])
im = ax.imshow(frames[0])
def update(frame):
im.set_data(frame)
return [im]
interval = 1000/framerate
anim = animation.FuncAnimation(fig=fig, func=update, frames=frames,
interval=interval, blit=True, repeat=False)
# return HTML(anim.to_html5_video())
# Save as MP4
anim.save(videoname, writer='ffmpeg', fps=framerate)
plt.close(fig)
def add_body_in_spec(spec):
"""
Add a body to the mujoco model.
Args:
spec: mujoco.MjSpec, the mujoco model specification.
Returns:
spec: mujoco.MjSpec, the updated mujoco model specification.
Note: The body is a classic U-shape stirrup
TODO: make this process adaptve to every kind of stirrup.
"""
# plugin2_0 = spec.add_plugin()
# plugin2_0.instance_name = "compositestirrup0"
# plugin2_0.active = 1
body2 = spec.worldbody.add_body()
# body2 = spec.find_body("stirrup")
body2.name = "stirrup2"
body2.pos = [0.5, 0.1, 0.5]
# body2.quat = [ 0, 0, 0.7071068, 0.7071068 ]
body2.quat = [0.707107, 0, -0.707107, 0]
body2.mass = 1
body2.inertia = [0.001, 0.001, 0.001]
body2.ipos = [0.28, 0, 0.005]
joint = body2.add_joint()
joint.type = mujoco.mjtJoint.mjJNT_FREE
body2_0 = body2.add_body()
body2_0.name = "stirrup20B0"
body2_0.pos = [0, 0.005, 0]
joint2_0 = body2_0.add_joint()
joint2_0.name = "stirrup20J0"
joint2_0.type = mujoco.mjtJoint.mjJNT_BALL
joint2_0.group = 3
joint2_0.pos = [0, 0, 0]
joint2_0.armature = 100
joint2_0.damping = 1
joint2_0.stiffness = 100
joint2_0.range = [0, 2]
geom2_0 = body2_0.add_geom()
geom2_0.name = "stirrup2G0"
geom2_0.size = [0.005, 0.2825, 0]
geom2_0.pos = [0.2825, 0, 0]
geom2_0.quat = [0.707107, 0, -0.707107, 0]
# geom2_0.quat = [1, 0, 0, 0]
geom2_0.type = mujoco.mjtGeom.mjGEOM_CAPSULE
geom2_0.rgba = [0.8, 0.2, 0.1, 1]
site2_0 = body2_0.add_site()
site2_0.name = "stirrup2S0"
site2_0.pos = [0, 0, 0]
site2_0.group = 3
body2_1 = body2_0.add_body()
body2_1.name = "stirrup20B1"
body2_1.pos = [0.565, 0, 0]
body2_1.quat = [0.92388, 0, 0, 0.382683]
joint2_1 = body2_1.add_joint()
joint2_1.name = "stirrup20J1"
joint2_1.type = mujoco.mjtJoint.mjJNT_BALL
joint2_1.group = 3
joint2_1.pos = [0, 0, 0]
joint2_1.armature = 100
joint2_1.damping = 1
joint2_1.stiffness = 100
joint2_1.range = [0, 2]
geom2_1 = body2_1.add_geom()
geom2_1.name = "stirrup2G1"
geom2_1.size = [0.005, 0.0212132, 0]
geom2_1.pos = [0.0212132, 0, 0]
geom2_1.quat = [0.707107, 0, -0.707107, 0]
geom2_1.type = mujoco.mjtGeom.mjGEOM_CAPSULE
geom2_1.rgba = [0.8, 0.2, 0.1, 1]
site2_1 = body2_1.add_site()
site2_1.name = "stirrup2S1"
site2_1.pos = [0, 0, 0]
site2_1.group = 3
body2_2 = body2_1.add_body()
body2_2.name = "stirrup20B2"
body2_2.pos = [0.0424264, 0, 0]
body2_2.quat = [0.92388, 0, 0, 0.382683]
joint2_2 = body2_2.add_joint()
joint2_2.name = "stirrup20J2"
joint2_2.type = mujoco.mjtJoint.mjJNT_BALL
joint2_2.group = 3
joint2_2.pos = [0, 0, 0]
joint2_2.armature = 100
joint2_2.damping = 1
joint2_2.stiffness = 100
joint2_2.range = [0, 2]
geom2_2 = body2_2.add_geom()
geom2_2.name = "stirrup2G2"
geom2_2.size = [0.005, 0.06, 0]
geom2_2.pos = [0.06, 0, 0]
geom2_2.quat = [0.707107, 0, -0.707107, 0]
geom2_2.type = mujoco.mjtGeom.mjGEOM_CAPSULE
geom2_2.rgba = [0.8, 0.2, 0.1, 1]
site2_2 = body2_2.add_site()
site2_2.name = "stirrup2S2"
site2_2.pos = [0, 0, 0]
site2_2.group = 3
body2_3 = body2_2.add_body()
body2_3.name = "stirrup20B3"
body2_3.pos = [0.12, 0, 0]
body2_3.quat = [0.92388, 0, 0, 0.382683]
joint2_3 = body2_3.add_joint()
joint2_3.name = "stirrup20J3"
joint2_3.type = mujoco.mjtJoint.mjJNT_BALL
joint2_3.group = 3
joint2_3.pos = [0, 0, 0]
joint2_3.armature = 100
joint2_3.damping = 1
joint2_3.stiffness = 100
joint2_3.range = [0, 2]
geom2_3 = body2_3.add_geom()
geom2_3.name = "stirrup2G3"
geom2_3.size = [0.005, 0.0212132, 0]
geom2_3.pos = [0.0212132, 0, 0]
geom2_3.quat = [0.707107, 0, -0.707107, 0]
geom2_3.type = mujoco.mjtGeom.mjGEOM_CAPSULE
geom2_3.rgba = [0.8, 0.2, 0.1, 1]
site2_3 = body2_3.add_site()
site2_3.name = "stirrup2S3"
site2_3.pos = [0, 0, 0]
site2_3.group = 3
body2_4 = body2_3.add_body()
body2_4.name = "stirrup20B4"
body2_4.pos = [0.0424264, 0, 0]
body2_4.quat = [0.92388, 0, 0, 0.382683]
joint2_4 = body2_4.add_joint()
joint2_4.name = "stirrup20J4"
joint2_4.type = mujoco.mjtJoint.mjJNT_BALL
joint2_4.group = 3
joint2_4.pos = [0, 0, 0]
joint2_4.armature = 100
joint2_4.damping = 1
joint2_4.stiffness = 100
joint2_4.range = [0, 2]
geom2_4 = body2_4.add_geom()
geom2_4.name = "stirrup2G4"
geom2_4.size = [0.005, 0.2825, 0]
geom2_4.pos = [0.2825, 0, 0]
geom2_4.quat = [0.707107, 0, -0.707107, 0]
geom2_4.type = mujoco.mjtGeom.mjGEOM_CAPSULE
geom2_4.rgba = [0.8, 0.2, 0.1, 1]
site2_4 = body2_4.add_site()
site2_4.name = "stirrup2S4"
site2_4.pos = [0, 0, 0]
site2_4.group = 3
return spec
if __name__ == '__main__':
spec = mujoco.MjSpec()
spec.from_file("src/rebar_mujoco/environments/xml_files/mujoco.xml")
model = spec.compile()
data = mujoco.MjData(model)
framerate = 40 # (Hz)
timestep = 0.001 # (seconds)
duration = 5 # (seconds)
width, height = 500, 500
renderer = mujoco.Renderer(model, height=height, width=width)
video_name = "src/rebar_mujoco/tests/tmp/test_mjSpec.mp4"
scene_option = mujoco.MjvOption()
frames = []
frame_count = 0
camera = mujoco.MjvCamera()
mujoco.mjv_defaultFreeCamera(model, camera)
camera.distance = 2.5
camera.elevation = -20
# camera.azimuth = -130
camera.lookat = (0, 0, 0.5)
add_body_time = 2 # seconds
add_body = False
z_pos_vec = []
nb_step = 0
while nb_step*timestep < duration:
mujoco.mj_step(model, data)
if nb_step*timestep > add_body_time and not add_body:
print("Adding body")
add_body = True
spec = add_body_in_spec(spec)
## Recompile model and data while maintaining the state.
model, data = spec.recompile(model, data)
renderer.close()
renderer = mujoco.Renderer(model, height=height, width=width)
if frame_count < nb_step*timestep * framerate:
z_pos_vec.append(data.geom_xpos[-1][2])
renderer.update_scene(data, camera=camera, scene_option=scene_option)
frame = renderer.render()
frames.append(frame)
frame_count += 1
nb_step += 1
display_video(frames, video_name, framerate)
print(f"Video saved at {video_name}")
import matplotlib.pyplot as plt
plt.plot(z_pos_vec)
plt.show()Context:
- Operating system: Ubuntu 20.04
- MuJoCo version: 3.2.2
Metadata
Metadata
Assignees
Labels
No labels