Implementation:ARISE Initiative Robomimic TrainUtils save model
| Knowledge Sources | |
|---|---|
| Domains | Robotics, Training, Serialization |
| Last Updated | 2026-02-15 08:00 GMT |
Overview
Concrete tool for saving trained models with full metadata to PyTorch checkpoint files provided by the robomimic training utilities module.
Description
The save_model function serializes the algorithm state, configuration, environment metadata, shape metadata, and normalization statistics into a single .pth file using torch.save. It deep-copies metadata to avoid aliasing issues, converts normalization stats to lists for JSON compatibility, and validates that observation normalization is enabled in config when stats are provided.
Usage
Call this function at checkpointing intervals during training: periodically (every N epochs), on best rollout performance, and at training completion. The resulting .pth files are consumed by policy_from_checkpoint for evaluation.
Code Reference
Source Location
- Repository: robomimic
- File: robomimic/utils/train_utils.py
- Lines: L588-634
Signature
def save_model(model, config, env_meta, shape_meta, ckpt_path,
variable_state=None, obs_normalization_stats=None,
action_normalization_stats=None):
"""
Save model to a torch pth file.
Args:
model (Algo instance): model to save
config (BaseConfig instance): config to save
env_meta (dict): env metadata for this training run
shape_meta (dict): shape metadata for this training run
ckpt_path (str): writes model checkpoint to this path
variable_state (dict): internal variable state in main train loop
obs_normalization_stats (dict): observation normalization statistics
action_normalization_stats (dict): action normalization statistics
"""
Import
import robomimic.utils.train_utils as TrainUtils
# Call as:
TrainUtils.save_model(
model=model,
config=config,
env_meta=env_meta,
shape_meta=shape_meta,
ckpt_path="/path/to/checkpoint.pth",
)
I/O Contract
Inputs
| Name | Type | Required | Description |
|---|---|---|---|
| model | Algo | Yes | Algorithm instance to serialize |
| config | BaseConfig | Yes | Full experiment configuration |
| env_meta | dict | Yes | Environment construction metadata |
| shape_meta | dict | Yes | Observation/action shape metadata |
| ckpt_path | str | Yes | File path for the checkpoint |
| variable_state | dict | No | Training loop state for resuming |
| obs_normalization_stats | dict | No | Observation normalization mean/std |
| action_normalization_stats | dict | No | Action normalization mean/std |
Outputs
| Name | Type | Description |
|---|---|---|
| .pth file | File | Saved checkpoint containing keys: model, config, algo_name, env_metadata, shape_metadata, variable_state, obs_normalization_stats, action_normalization_stats |
Usage Examples
Save Checkpoint During Training
import robomimic.utils.train_utils as TrainUtils
import os
# Save periodic checkpoint
ckpt_path = os.path.join(output_dir, f"model_epoch_{epoch}.pth")
TrainUtils.save_model(
model=model,
config=config,
env_meta=env_meta,
shape_meta=shape_meta,
ckpt_path=ckpt_path,
variable_state=dict(epoch=epoch, best_return=best_return),
obs_normalization_stats=obs_normalization_stats,
)