Use of MjSpec in environment #494

LilianLaporte · 2024-09-20T16:17:05Z

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

The text was updated successfully, but these errors were encountered:

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Use of MjSpec in environment #494

Use of MjSpec in environment #494

LilianLaporte commented Sep 20, 2024

Use of MjSpec in environment #494

Use of MjSpec in environment #494

Comments

LilianLaporte commented Sep 20, 2024

Example