Video recorder

def setup(self, state: MjState) -> Self:
    """Initializes the MuJoCo renderer for this model. Must be called before the simulation loop."""
    self._validate(state)

    try:
        self._renderer = mujoco.Renderer(
            model=state.model, height=self.height, width=self.width
        )
    except Exception as e:
        msg = "Failed to initialize the MuJoCo Renderer. If on a server, try setting 'export MUJOCO_GL=egl' in your terminal."
        logger.error(msg)
        raise RuntimeError(msg) from e

    # initialize the vopt with defaults
    mujoco.mjv_defaultOption(self._vopt)

    # whether or not to show debug graphics
    self._vopt.flags[mujoco.mjtVisFlag.mjVIS_PERTFORCE] = int(self.show_net_force)
    self._vopt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = int(self.show_contacts)
    return self

def is_due(self, state: MjState) -> bool:
    """Returns whether `capture_frame` would actually capture a frame for `state` right now, without any side effects. Lets callers skip expensive work (e.g. building `custom_traces`) that would otherwise go unused on the steps between frames."""
    if state.data.time < self._next_record_time:
        return False
    if not self.recording_trigger(state):
        return False
    if self.max_frames is not None and self._frame_count >= self.max_frames:
        return False
    return True

def capture_frame(
    self,
    state: MjState,
    custom_arrows: list[ArrowConfig],
    custom_lines: list[LineConfig],
    custom_traces: list[LineConfig],
):
    """Captures the current state as a video frame."""
    if state.data.time < self._next_record_time:
        return

    if not self.recording_trigger(state):
        return

    if self.max_frames is not None and self._frame_count >= self.max_frames:
        if not self._max_frames_warned:
            logger.warning(
                f"VideoRecorder for {self.path} reached max_frames={self.max_frames}; "
                "no further frames will be captured."
            )
            self._max_frames_warned = True
        return

    frame = self._render_frame(state, custom_arrows, custom_lines, custom_traces)

    if self.path.suffix.lower() in STREAMED_VIDEO_FORMAT:
        if self._encoder_proc is None:
            self._open_encoder()
        w, h = self._encode_size
        assert (
            self._encoder_proc is not None and self._encoder_proc.stdin is not None
        )
        self._encoder_proc.stdin.write(frame[:h, :w].tobytes())
    else:
        self._frames.append(frame)

    # increment the clock for the next frame
    self._frame_count += 1
    self._next_record_time += 1 / self.fps

def snapshot(
    self,
    state: MjState,
    path: Path,
    custom_arrows: list[ArrowConfig] | None = None,
    custom_lines: list[LineConfig] | None = None,
    custom_traces: list[LineConfig] | None = None,
):
    """Renders the current state and saves it as a single image to `path`, regardless of `recording_trigger` or `fps` timing."""
    from PIL import Image

    frame = self._render_frame(
        state, custom_arrows or [], custom_lines or [], custom_traces or []
    )
    Image.fromarray(frame).save(path)
    logger.info(f"Snapshot saved to {path}")

def save(self):
    """
    Finishes writing the video file.

    Supported formats:
    - `.mp4`: H.264 via ffmpeg; universally compatible.
    - `.webm`: VP9 via ffmpeg; smaller files and fully seekable.
    - `.gif`: via PIL; no audio, loops automatically, large file size, not seekable.

    `.mp4` and `.webm` are encoded incrementally: each frame is piped to ffmpeg as it's captured, so `save` only needs to close that pipe and wait for ffmpeg to finish. `.gif` (and any other format) buffers every frame in memory and is only encoded here.

    The output format is determined by the extension of `path`.
    """
    if self._frame_count == 0:
        return

    if self._encoder_proc is not None:
        self._finish_encoding()
    elif self.path.suffix.lower() == ".gif":
        from PIL import Image

        # convert arrays to PIL images
        pil_images = [Image.fromarray(frame) for frame in self._frames]

        # save gif
        pil_images[0].save(
            self.path,
            save_all=True,
            append_images=pil_images[1:],
            duration=int(1000 / self._output_fps),  # ms per frame
            loop=0,  # loop forever
        )
    else:
        import mediapy as media

        media.write_video(path=self.path, images=self._frames, fps=self._output_fps)
    logger.info(f"Video saved to {self.path}")

def close(self) -> None:
    """Releases the GL context held by the underlying MuJoCo renderer, and kills the ffmpeg encoder if `save` was never called. Call once recording is finished and `save` has been called."""
    if self._encoder_proc is not None and self._encoder_proc.poll() is None:
        proc = self._encoder_proc
        try:
            if proc.stdin is not None:
                proc.stdin.close()
        except BrokenPipeError:
            pass
        proc.kill()
        proc.wait()
    self._renderer.close()