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)
    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)
            r.update_scene(data, camera=cam)
            w.append_data(r.render())


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