render_worker.py

"""Render a qpos trajectory to MP4. Run as a subprocess by rollout.write_video
with MUJOCO_GL already set, because MuJoCo fixes its GL backend at import.

    python render_worker.py model.mjb qpos.npy out.mp4 fps width height camera
"""
import sys

import numpy as np


def main(mjb, qpos_path, out, fps, width, height, camera):
    import imageio.v2 as imageio
    import mujoco

    model = mujoco.MjModel.from_binary_path(mjb)
    data = mujoco.MjData(model)
    qpos = np.load(qpos_path)

    # "follow" is a free camera kept above and behind the robot. The scene's own
    # "track" camera sits 0.35 m above the centre of mass, which ends up inside
    # the hillside on sloped terrain and renders the ground black.
    follow = camera == "follow"
    if follow:
        cam = mujoco.MjvCamera()
        cam.azimuth, cam.elevation, cam.distance = 140.0, -15.0, 6.0
    else:
        cam = camera if mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_CAMERA, camera) >= 0 else -1

    with mujoco.Renderer(model, height=int(height), width=int(width)) as r, \
            imageio.get_writer(out, fps=float(fps), codec="libx264", quality=7,
                               macro_block_size=None) as w:
        for q in qpos:
            data.qpos[:] = q
            mujoco.mj_forward(model, data)
            if follow:
                cam.lookat[:] = data.qpos[:3]
            r.update_scene(data, camera=cam)
            w.append_data(r.render())


if __name__ == "__main__":
    main(*sys.argv[1:8])