OakEmu-RL v0.1.0

Appearance

BackgroundCheckpointWriter ​

python
BackgroundCheckpointWriter(checkpoint_manager: CheckpointManager, metrics_tracker: MetricsTracker, onnx_exporter: OnnxExporter, get_global_step: Callable[[], int], *, save_onnx_with_checkpoints: bool, checkpoint_onnx_filename_template: str, game_name: str, experiment_name: str, save_movies_with_checkpoints: bool = False, checkpoint_movie_maximum_steps: int = 5000, checkpoint_movie_extra_end_frames: int = 0, checkpoint_movie_filename_template: str = '{game}_checkpoint_{checkpoint}.mp4', checkpoint_movie_command: str = 'oakemu') -> None

Saves checkpoints on a single background worker thread.

The checkpoint number and directory are reserved synchronously so callers can read CheckpointManager.latest_checkpoint_number immediately; the file writes, compression, and optional ONNX export then run on the background worker. Files passed to save must therefore be immutable snapshots of the captured state. A failure from the previous background write is re-raised on the next call to save or wait_for_pending_write.

Initialize the writer.

Parameters:

  • checkpoint_manager (CheckpointManager) – Reserves and writes checkpoint directories.
  • metrics_tracker (MetricsTracker) – Logs the checkpoint-number metric.
  • onnx_exporter (OnnxExporter) – Exports the optional ONNX model alongside a checkpoint.
  • get_global_step (Callable[[], int]) – Returns the current global step, for the checkpoint-number metric.
  • save_onnx_with_checkpoints (bool) – Whether to export ONNX alongside every checkpoint.
  • checkpoint_onnx_filename_template (str) – Template for the ONNX file's path within the checkpoint directory.
  • game_name (str) – The experiment's lowercased game name, for template substitution.
  • experiment_name (str) – The experiment's name, for template substitution.
  • save_movies_with_checkpoints (bool) – Record a movie of each exported ONNX model into its checkpoint directory.
  • checkpoint_movie_maximum_steps (int) – Step cap for checkpoint movies.
  • checkpoint_movie_extra_end_frames (int) – Frames recorded after the episode ends.
  • checkpoint_movie_filename_template (str) – Template for the movie's path within the checkpoint directory.
  • checkpoint_movie_command (str) – The OakEmu tool executable used to record movies.

Source

Methods ​

save ​

python
save(files: list[CheckpointFile], *, background: bool = True) -> Path

Save a checkpoint.

Parameters:

  • files (list[CheckpointFile]) – The checkpoint files to save.
  • background (bool) – Whether to write the files on the background worker. Final shutdown saves must write synchronously because the worker may already be torn down while threads are being joined.

Returns:

  • Path – The reserved checkpoint directory.

Source

shutdown ​

python
shutdown(*, wait: bool) -> None

Shut down the background worker.

Parameters:

  • wait (bool) – Whether to block until the worker has finished.

Source

wait_for_pending_write ​

python
wait_for_pending_write() -> None

Block until the in-flight background checkpoint write completes.

Re-raises any exception from the background write on the calling thread.

Source