reference/generated/torchrl.checkpoint.Checkpoint

Run in Google Colab

Colab

Download Notebook

Notebook

View on GitHub

GitHub

Checkpoint

class torchrl.checkpoint.Checkpoint(*, format: Literal['directory', 'archive'] = 'directory', strict: Literal['error', 'warn', 'ignore'] = 'error', archive_compression: Literal['stored', 'deflate'] = 'stored', save_components: Collection[str] | None = None, **components: Any)

Standard TorchRL checkpoint coordinator.

Components are bound to the checkpoint and may be independently selected for each save or restore. The default directory format supports direct, lazy access to component payloads; format="archive" stores the same manifest and component tree in one ZIP file.

Parameters:

• format – Default output format, "directory" or "archive".

• strict – Default restoration behavior for missing or incompatible requested components.

• archive_compression – "stored" avoids recompressing payloads; "deflate" enables ZIP deflate compression.

• save_components – Optional default component selection for saves. This is useful for excluding large components from scheduled saves.

• **components – Named objects or JSON-compatible values to register.

Examples

>>> import tempfile
>>> import torch
>>> from torchrl.checkpoint import Checkpoint
>>> source = torch.nn.Linear(2, 1)
>>> target = torch.nn.Linear(2, 1)
>>> with tempfile.TemporaryDirectory() as tmpdir:
...     Checkpoint(policy=source).save(f"{tmpdir}/checkpoint")
...     result = Checkpoint(policy=target).load(f"{tmpdir}/checkpoint")
>>> result.loaded
{'policy'}

property components: Mapping[str, Any]

Read-only mapping view of registered component values.

classmethod is_checkpoint(path: str | Path) → bool

Return whether path contains a TorchRL checkpoint manifest.

load(path: str | Path, *, components: Collection[str] | None = None, component_options: Mapping[str, CheckpointOptions] | None = None, map_location: Any = None, tensor_load_kwargs: Mapping[str, Any] | None = None, strict: Literal['error', 'warn', 'ignore'] | None = None) → CheckpointLoadResult

Restore selected components from a local checkpoint.

Parameters:

• path – Directory or archive checkpoint.

• components – Registered names to restore. None selects all registered components.

• component_options – Per-operation option overrides.

• map_location – Device mapping used while reading tensor payloads.

• tensor_load_kwargs – Additional keyword arguments passed to torch.load() for state-dict components explicitly saved with the torch payload format. weights_only defaults to True.

• strict – Per-operation strictness override.

Returns:

A structured component load report.

classmethod manifest(path: str | Path) → dict[str, Any]

Read and validate a checkpoint manifest without loading payloads.

classmethod read_component(path: str | ~pathlib.Path, name: str, *, map_location: ~typing.Any = None, default: ~typing.Any = <object object>) → Any

Return one component’s stored payload without a live object.

State-dict components return the decoded state dict and JSON components return the stored value, so a checkpoint can be inspected before the objects it belongs to exist, for example to read the saved logger state or configuration. Tensors are copied out of the checkpoint. Components stored through dump/load or a custom adapter require a live object and are rejected.

Parameters:

• path – Directory or archive checkpoint.

• name – Manifest component name.

• map_location – Device mapping used while reading tensor payloads.

• default – Value returned instead of raising when the component is absent.

Returns:

The state dict or JSON value stored for name.

Raises:

• KeyError – If the checkpoint has no component named name and no default is given.

• CheckpointError – If the component cannot be read without an object.

register(name: str, component: Any, *, adapter: CheckpointAdapter | None = None, options: CheckpointOptions | None = None) → Checkpoint

Register a named component and return self.

Parameters:

• name – Stable manifest component name.

• component – Live object or JSON-compatible value.

• adapter – Optional explicit serialization adapter.

• options – Persistent component method arguments.

Returns:

This checkpoint.

register_adapter(component_type: type, adapter: CheckpointAdapter) → Checkpoint

Register an adapter for a type on this checkpoint instance.

classmethod register_migration(from_version: int, migration: Callable[[dict[str, Any]], dict[str, Any]]) → None

Register one manifest migration from from_version to the next.

Migrations must return a new manifest whose format_version is exactly from_version + 1. Payload migrations remain the responsibility of the component adapter identified by the migrated manifest.

save(path: str | Path, *, components: Collection[str] | None = None, component_options: Mapping[str, CheckpointOptions] | None = None, format: Literal['directory', 'archive'] | None = None, metadata: Mapping[str, Any] | None = None) → Path

Save selected registered components atomically.

Parameters:

• path – Local checkpoint destination.

• components – Component names to save. None uses save_components from construction, or saves all when no default selection was configured.

• component_options – Per-operation option overrides.

• format – Per-operation container override.

• metadata – JSON-compatible manifest metadata.

Returns:

Expanded destination path.