๐พ Checkpoints#
Genesis provides state save/load functionality for training resumption and episode resets.
Basic Save/Load#
import genesis as gs
scene = gs.Scene()
robot = scene.add_entity(gs.morphs.MJCF(file="franka.xml"))
scene.build()
# Simulate
for _ in range(100):
scene.step()
# Save checkpoint
scene.save_checkpoint("checkpoint.pkl")
# Load in new scene
scene2 = gs.Scene()
robot2 = scene2.add_entity(gs.morphs.MJCF(file="franka.xml"))
scene2.build()
scene2.load_checkpoint("checkpoint.pkl")
State Objects#
# Get current state (in-memory)
state = scene.get_state()
# Reset to initial state
scene.reset()
# Reset to custom state
scene.reset(state=state)
RL Episode Resets#
scene.build(n_envs=N)
# Snapshot initial state
init_state = scene.get_state()
for episode in range(num_episodes):
scene.reset(state=init_state)
for step in range(episode_length):
scene.step()
obs, reward, done = get_observations()
# Reset environments where episode ended
if done.any():
done_envs = torch.where(done)[0].tolist()
scene.reset(state=init_state, envs_idx=done_envs)
Selective Environment Reset#
scene.build(n_envs=16)
# Reset all environments
scene.reset()
# Reset specific environments
scene.reset(envs_idx=[0, 2, 5])
# Reset with custom state for specific envs
scene.reset(state=init_state, envs_idx=[1, 3, 7])
State Contents#
The SimState object contains:
Solver |
State Variables |
|---|---|
Rigid |
|
MPM |
|
SPH |
|
PBD |
|
FEM |
|
Checkpoint File Format#
Checkpoints are pickled dictionaries:
{
"timestamp": time.time(),
"step_index": scene.t,
"arrays": { # Numpy arrays keyed by solver/field
"RigidSolver.qpos": np.array(...),
"MPMSolver.pos": np.array(...),
...
}
}
Serialization for Transfer#
# Make state serializable (detach from graph)
state = scene.get_state()
state_serializable = state.serializable()
# Now safe to pickle
import pickle
with open("state.pkl", "wb") as f:
pickle.dump(state_serializable, f)
Important Notes#
Checkpoints require compatible scene configuration (same entities, solver options)
32-bit precision may lose ~2e-6 accuracy between save/load
Use
envs_idxparameter for efficient partial resetsscene.tstores the simulation step count