Untitled

unknown
plain_text
a month ago
2.1 kB
Indexable
Never
import mediapy as media
import mujoco
import numpy as np

model_xml = """<mujoco model="kinematic_chain">
    <worldbody>
        <body name="body_0">
            <!-- Geometries will be added here -->
        </body>
    </worldbody>
</mujoco>"""


# Load the model
spec = mujoco.MjSpec()


spec.from_string(model_xml)
body_name = "world"

model = spec.compile()
data = mujoco.MjData(model)
height, width = (480, 640)
duration = 10
framerate = 30

frames = []
mujoco.mj_resetData(model, data)  # Reset state and time

with mujoco.Renderer(model, height, width) as renderer:

    add_body_time = 2  # seconds
    add_body = False

    start = data.time
    while data.time < duration:
        mujoco.mj_step(model, data)

        renderer.update_scene(data)
        if data.time - start > add_body_time and not add_body:
            print("Adding body...")
            add_body = True
            body = spec.find_body(body_name)
            kcgeom00 = body.add_geom()
            kcgeom00.type = mujoco.mjtGeom.mjGEOM_CAPSULE
            kcgeom00.size = [0.05, 0.0, 0]
            kcgeom00.fromto = [0, 0, 0, 0.5, 0, 0]
            kcgeom00.name = "Kc_geom00"
            kcgeom00.rgba = [0.1, 1, 0.1, 1]
            kcgeom00.contype = 0  # to avoid collision
            kcgeom00.conaffinity = 0  # to avoid collision
            kcgeom01 = body.add_geom()
            kcgeom01.type = mujoco.mjtGeom.mjGEOM_CAPSULE
            kcgeom01.size = [0.05, 0.0, 0]
            kcgeom01.fromto = [0.5, 0, 0, 0.5, 0.0, -0.5]
            kcgeom01.name = "Kc_geom01"
            kcgeom01.rgba = [0.1, 1, 0.1, 1]
            kcgeom01.contype = 0  # to avoid collision
            kcgeom01.conaffinity = 0  # to avoid collision

            model, data = spec.recompile(model, data)
            renderer.close()
            renderer = mujoco.Renderer(model, height=height, width=width)

        if len(frames) < data.time * framerate:
            frame = renderer.render()
            frames.append(frame)

# Display the video
media.show_video(frames, fps=framerate)