Implementation:Haosulab ManiSkill CachedResetWrapper

Knowledge Sources	Haosulab_ManiSkill
Domains	Robotics, Simulation, Environment Wrapper
Last Updated	2026-02-15 08:00 GMT

Overview

Concrete tool for caching environment reset states to accelerate repeated resets by skipping expensive reconfiguration.

Description

The CachedResetWrapper is a Gymnasium wrapper that caches a set of environment reset states and randomly samples from them on each reset() call, bypassing the slower full reset procedure. This can significantly boost environment FPS during training.

CachedResetsConfig:

• num_resets -- Number of states to cache (defaults to num_envs).
• device -- Device for cached data storage (defaults to env device).
• seed -- Seed for generating cached states.

Initialization:

• If reset_to_env_states is provided, uses those pre-computed states directly.
• Otherwise, runs the environment reset multiple times to collect states and observations.
• Handles cases where num_resets is not a multiple of num_envs by slicing excess data.
• Stores cached states and observations on the configured device.

Reset behavior:

• Randomly samples num_envs indices from the cached states.
• Passes the sampled states and optionally observations to the underlying environment's reset via options["reset_to_env_states"].
• Supports partial resets via env_idx in options.

Usage

Wrap a ManiSkill environment during training to speed up resets. Particularly effective for GPU-parallel environments where reset overhead can be significant.

Code Reference

Source Location

• Repository: Haosulab_ManiSkill
• File: mani_skill/utils/wrappers/cached_reset.py

Signature

@dataclass
class CachedResetsConfig:
    num_resets: Optional[int] = None
    device: Optional[Device] = None
    seed: Optional[int] = None

class CachedResetWrapper(gym.Wrapper):
    def __init__(
        self,
        env: gym.Env,
        reset_to_env_states: Optional[dict] = None,
        config: Union[CachedResetsConfig, dict] = CachedResetsConfig(),
    ): ...
    def reset(self, *, seed=None, options=None) -> tuple: ...

Import

from mani_skill.utils.wrappers.cached_reset import CachedResetWrapper, CachedResetsConfig

I/O Contract

Inputs

Name	Type	Required	Description
env	gym.Env	Yes	The environment to wrap
reset_to_env_states	dict	No	Pre-computed states {"env_states": ..., "obs": ...}
config	CachedResetsConfig or dict	No	Cache configuration

Outputs

Name	Type	Description
obs	dict	Observation from the cached reset state
info	dict	Info dictionary from reset

Usage Examples

Basic Usage

import gymnasium as gym
from mani_skill.utils.wrappers.cached_reset import CachedResetWrapper, CachedResetsConfig

env = gym.make("PickCube-v1", num_envs=16)
env = CachedResetWrapper(
    env,
    config=CachedResetsConfig(num_resets=100, seed=42),
)

# Each reset samples from 100 cached states
obs, info = env.reset()

Related Pages

• Environment:Haosulab_ManiSkill_Python_SAPIEN_Core
• Haosulab_ManiSkill_CommonUtils