mani-skill-nightly 2025.7.25.329__py3-none-any.whl → 2025.7.25.606__py3-none-any.whl

mani_skill/envs/sapien_env.py

@@ -37,7 +37,7 @@ from mani_skill.sensors.camera import (
     update_camera_configs_from_dict,
 )
 from mani_skill.sensors.depth_camera import StereoDepthCamera, StereoDepthCameraConfig
-from mani_skill.utils import common, gym_utils, sapien_utils
+from mani_skill.utils import common, gym_utils, sapien_utils, tree
 from mani_skill.utils.structs import Actor, Articulation
 from mani_skill.utils.structs.pose import Pose
 from mani_skill.utils.structs.types import Array, SimConfig
@@ -316,6 +316,8 @@ class BaseEnv(gym.Env):
         self._elapsed_steps = (
             torch.zeros(self.num_envs, device=self.device, dtype=torch.int32)
         )
+        self._last_obs = None
+        """the last observation returned by the environment"""
         obs, _ = self.reset(seed=[2022 + i for i in range(self.num_envs)], options=dict(reconfigure=True))
 
         self._init_raw_obs = common.to_cpu_tensor(obs)
@@ -850,7 +852,11 @@ class BaseEnv(gym.Env):
         options["reconfigure"] is True, will call self._reconfigure() which deletes the entire physx scene and reconstructs everything.
         Users building custom tasks generally do not need to override this function.
 
-        Returns the first observation and a info dictionary. The info dictionary is of type
+        If options["reset_to_env_states"] is given, we expect there to be options["reset_to_env_states"]["env_states"] and optionally options["reset_to_env_states"]["obs"], both with 
+        batch size equal to the number of environments being reset. "env_states" can be a dictionary or flat tensor and we skip calling the environment's _initialize_episode function which
+        generates the initial state on a normal reset. If "obs" is given we skip calling the environment's get_obs function which can save some compute/time.
+
+        Returns the observations and an info dictionary. The info dictionary is of type
 
 
         .. highlight:: python
@@ -917,12 +923,22 @@ class BaseEnv(gym.Env):
         if self.agent is not None:
             self.agent.reset()
 
-        if seed is not None or self._enhanced_determinism:
-            with torch.random.fork_rng():
-                torch.manual_seed(self._episode_seed[0])
-                self._initialize_episode(env_idx, options)
+        # we either reset to given env states or use the environment's defined _initialize_episode function to generate the initial state
+        reset_to_env_states_obs = None
+        if "reset_to_env_states" in options:
+            env_states = options["reset_to_env_states"]["env_states"]
+            reset_to_env_states_obs = options["reset_to_env_states"].get("obs", None)
+            if isinstance(env_states, dict):            
+                self.set_state_dict(env_states, env_idx)
+            else:
+                self.set_state(env_states, env_idx)
         else:
-            self._initialize_episode(env_idx, options)
+            if seed is not None or self._enhanced_determinism:
+                with torch.random.fork_rng():
+                    torch.manual_seed(self._episode_seed[0])
+                    self._initialize_episode(env_idx, options)
+            else:
+                self._initialize_episode(env_idx, options)
         # reset the reset mask back to all ones so any internal code in maniskill can continue to manipulate all scenes at once as usual
         self.scene._reset_mask = torch.ones(
             self.num_envs, dtype=bool, device=self.device
@@ -942,9 +958,13 @@ class BaseEnv(gym.Env):
                 self.agent.controller.reset()
 
         info = self.get_info()
-        obs = self.get_obs(info)
-
+        if reset_to_env_states_obs is None:
+            obs = self.get_obs(info)
+        else:
+            obs = self._last_obs
+            tree.replace(obs, env_idx, common.to_tensor(reset_to_env_states_obs, device=self.device))
         info["reconfigure"] = reconfigure
+        self._last_obs = obs
         return obs, info
 
     def _set_main_rng(self, seed):
@@ -1031,7 +1051,7 @@ class BaseEnv(gym.Env):
                 terminated = info["fail"].clone()
             else:
                 terminated = torch.zeros(self.num_envs, dtype=bool, device=self.device)
-
+        self._last_obs = obs
         return (
             obs,
             reward,

mani_skill/utils/structs/articulation.py

@@ -888,7 +888,7 @@ class Articulation(BaseStruct[physx.PhysxArticulation]):
             else:
                 gx, gy = self.get_joint_target_indices(joint_indices)
             self.px.cuda_articulation_target_qpos.torch()[
-                gx[self.scene._reset_mask], gy[self.scene._reset_mask]
+                gx[self.scene._reset_mask[self._scene_idxs]], gy[self.scene._reset_mask[self._scene_idxs]]
             ] = targets
         else:
             for i, joint in enumerate(joints):
@@ -911,7 +911,9 @@ class Articulation(BaseStruct[physx.PhysxArticulation]):
                 gx, gy = self.get_joint_target_indices(joints)
             else:
                 gx, gy = self.get_joint_target_indices(joint_indices)
-            self.px.cuda_articulation_target_qvel.torch()[gx, gy] = targets
+            self.px.cuda_articulation_target_qvel.torch()[
+                gx[self.scene._reset_mask[self._scene_idxs]], gy[self.scene._reset_mask[self._scene_idxs]]
+            ] = targets
         else:
             for i, joint in enumerate(joints):
                 joint.set_drive_velocity_target(targets[0, i])

mani_skill/utils/tree.py

@@ -0,0 +1,24 @@
+import torch
+
+
+# NOTE (stao): when tensordict is used we should replace all of this
+def slice(x, i):
+    if isinstance(x, dict):
+        return {k: slice(v, i) for k, v in x.items()}
+    else:
+        return x[i]
+
+
+def cat(x: list):
+    if isinstance(x[0], dict):
+        return {k: cat([d[k] for d in x]) for k in x[0].keys()}
+    else:
+        return torch.cat(x, dim=0)
+
+
+def replace(x, i, y):
+    if isinstance(x, dict):
+        for k, v in x.items():
+            replace(v, i, y[k])
+    else:
+        x[i] = y

mani_skill/utils/wrappers/__init__.py

@@ -1,3 +1,5 @@
+from .action_repeat import ActionRepeatWrapper
+from .cached_reset import CachedResetWrapper
 from .flatten import (
     FlattenActionSpaceWrapper,
     FlattenObservationWrapper,
@@ -6,4 +8,3 @@ from .flatten import (
 from .frame_stack import FrameStack
 from .gymnasium import CPUGymWrapper
 from .record import RecordEpisode
-from .action_repeat import ActionRepeatWrapper

mani_skill/utils/wrappers/cached_reset.py

@@ -0,0 +1,151 @@
+from dataclasses import asdict, dataclass
+from typing import List, Optional, Union
+
+import dacite
+import gymnasium as gym
+import torch
+
+from mani_skill.envs.sapien_env import BaseEnv
+from mani_skill.utils import common, tree
+from mani_skill.utils.structs.types import Device
+
+
+@dataclass
+class CachedResetsConfig:
+    num_resets: Optional[int] = None
+    """The number of reset states to cache. If none it will cache `num_envs` number of reset states."""
+    device: Optional[Device] = None
+    """The device to cache the reset states on. If none it will use the base environment's device."""
+    seed: Optional[int] = None
+    """The seed to use for generating the cached reset states."""
+
+    def dict(self):
+        return {k: v for k, v in asdict(self).items()}
+
+
+class CachedResetWrapper(gym.Wrapper):
+    """
+    Cached reset wrapper for ManiSkill3 environments. Caching resets allows you to skip slower parts of the reset function call and boost environment FPS as a result.
+
+    Args:
+        env: The environment to wrap.
+        reset_to_env_states: A dictionary with keys "env_states" and optionally "obs". "env_states" is a dictionary of environment states to reset to.
+            "obs" contains the corresponding observations generated at those env states. If reset_to_env_states is not provided, the wrapper will sample reset states
+            from the environment using the given seed.
+        config: A dictionary or a `CachedResetsConfig` object that contains the configuration for the cached resets.
+    """
+
+    def __init__(
+        self,
+        env: gym.Env,
+        reset_to_env_states: Optional[dict] = None,
+        config: Union[CachedResetsConfig, dict] = CachedResetsConfig(),
+    ):
+        super().__init__(env)
+        self.num_envs = self.base_env.num_envs
+        if isinstance(config, CachedResetsConfig):
+            config = config.dict()
+        self.cached_resets_config = dacite.from_dict(
+            data_class=CachedResetsConfig,
+            data=config,
+            config=dacite.Config(strict=True),
+        )
+        cached_data_device = self.cached_resets_config.device
+        if cached_data_device is None:
+            cached_data_device = self.base_env.device
+        self._num_cached_resets = 0
+        if reset_to_env_states is not None:
+            self._cached_resets_env_states = reset_to_env_states["env_states"]
+            self._cached_resets_obs_buffer = reset_to_env_states.get("obs", None)
+            self._num_cached_resets = len(self._cached_resets_env_states)
+        else:
+            if self.cached_resets_config.num_resets is None:
+                self.cached_resets_config.num_resets = 16384
+            self._cached_resets_env_states = []
+            self._cached_resets_obs_buffer = []
+            while self._num_cached_resets < self.cached_resets_config.num_resets:
+                obs, _ = self.env.reset(
+                    seed=self.cached_resets_config.seed,
+                    options=dict(
+                        env_idx=torch.arange(
+                            0,
+                            min(
+                                self.cached_resets_config.num_resets
+                                - self._num_cached_resets,
+                                self.num_envs,
+                            ),
+                            device=self.base_env.device,
+                        )
+                    ),
+                )
+                state = self.env.get_wrapper_attr("get_state_dict")()
+                if (
+                    self.cached_resets_config.num_resets - self._num_cached_resets
+                    < self.num_envs
+                ):
+                    obs = tree.slice(
+                        obs,
+                        slice(
+                            0,
+                            self.cached_resets_config.num_resets
+                            - self._num_cached_resets,
+                        ),
+                    )
+                    state = tree.slice(
+                        state,
+                        slice(
+                            0,
+                            self.cached_resets_config.num_resets
+                            - self._num_cached_resets,
+                        ),
+                    )
+                self._cached_resets_obs_buffer.append(
+                    common.to_tensor(obs, device=self.cached_resets_config.device)
+                )
+                self._cached_resets_env_states.append(
+                    common.to_tensor(state, device=self.cached_resets_config.device)
+                )
+                self._num_cached_resets += self.num_envs
+            self._cached_resets_env_states = tree.cat(self._cached_resets_env_states)
+            self._cached_resets_obs_buffer = tree.cat(self._cached_resets_obs_buffer)
+
+        self._cached_resets_env_states = common.to_tensor(
+            self._cached_resets_env_states, device=cached_data_device
+        )
+        if self._cached_resets_obs_buffer is not None:
+            self._cached_resets_obs_buffer = common.to_tensor(
+                self._cached_resets_obs_buffer, device=cached_data_device
+            )
+
+    @property
+    def base_env(self) -> BaseEnv:
+        return self.env.unwrapped
+
+    def reset(
+        self,
+        *args,
+        seed: Optional[Union[int, List[int]]] = None,
+        options: Optional[dict] = None,
+        **kwargs
+    ):
+        env_idx = None
+        if options is None:
+            options = dict()
+        if "env_idx" in options:
+            env_idx = options["env_idx"]
+        if self._cached_resets_env_states is not None:
+            sampled_ids = torch.randint(
+                0,
+                self._num_cached_resets,
+                size=(len(env_idx) if env_idx is not None else self.num_envs,),
+                device=self.base_env.device,
+            )
+            options["reset_to_env_states"] = dict(
+                env_states=tree.slice(self._cached_resets_env_states, sampled_ids),
+            )
+            if self._cached_resets_obs_buffer is not None:
+                options["reset_to_env_states"]["obs"] = tree.slice(
+                    self._cached_resets_obs_buffer, sampled_ids
+                )
+        obs, info = self.env.reset(seed=seed, options=options)
+        return obs, info

mani_skill/utils/wrappers/record.py

@@ -357,16 +357,15 @@ class RecordEpisode(gym.Wrapper):
         self,
         *args,
         seed: Optional[Union[int, List[int]]] = None,
-        options: Optional[dict] = dict(),
+        options: Optional[dict] = None,
         **kwargs,
     ):
-
         if self.save_on_reset:
             if self.save_video and self.num_envs == 1:
                 self.flush_video()
             # if doing a full reset then we flush all trajectories including incompleted ones
             if self._trajectory_buffer is not None:
-                if "env_idx" not in options:
+                if options is None or "env_idx" not in options:
                     self.flush_trajectory(env_idxs_to_flush=np.arange(self.num_envs))
                 else:
                     self.flush_trajectory(
@@ -415,7 +414,7 @@ class RecordEpisode(gym.Wrapper):
             if self.record_env_state:
                 first_step.state = common.to_numpy(common.batch(state_dict))
             env_idx = np.arange(self.num_envs)
-            if "env_idx" in options:
+            if options is not None and "env_idx" in options:
                 env_idx = common.to_numpy(options["env_idx"])
             if self._trajectory_buffer is None:
                 # Initialize trajectory buffer on the first episode based on given observation (which should be generated after all wrappers)

mani_skill/vector/wrappers/gymnasium.py

@@ -89,10 +89,10 @@ class ManiSkillVectorEnv(VectorEnv):
         self,
         *,
         seed: Optional[Union[int, List[int]]] = None,
-        options: Optional[dict] = dict(),
+        options: Optional[dict] = None,
     ):
         obs, info = self._env.reset(seed=seed, options=options)
-        if "env_idx" in options:
+        if options is not None and "env_idx" in options:
             env_idx = options["env_idx"]
             mask = torch.zeros(self.num_envs, dtype=bool, device=self.base_env.device)
             mask[env_idx] = True

{mani_skill_nightly-2025.7.25.329.dist-info → mani_skill_nightly-2025.7.25.606.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: mani-skill-nightly
-Version: 2025.7.25.329
+Version: 2025.7.25.606
 Summary: ManiSkill3: A Unified Benchmark for Generalizable Manipulation Skills
 Home-page: https://github.com/haosulab/ManiSkill
 Author: ManiSkill contributors

{mani_skill_nightly-2025.7.25.329.dist-info → mani_skill_nightly-2025.7.25.606.dist-info}/RECORD

@@ -539,7 +539,7 @@ mani_skill/assets/robots/xarm7/meshes/visual/link7.glb,sha256=aZatACOv20VJbi2tOE
 mani_skill/assets/robots/xarm7/meshes/visual/link_base.glb,sha256=vcy2lN1V72jIsSDRT0ZKVskR_0pVOXtDvBkxO2GENWs,467668
 mani_skill/envs/__init__.py,sha256=YPlttBErTcf9vSnkZ54EQ8vTABSfFFrBdUY0AkF4vmg,43
 mani_skill/envs/minimal_template.py,sha256=9THHWA1vkHatptc9g5Ojh-UBUKWQmLHVeq4fcaqv2aY,2200
-mani_skill/envs/sapien_env.py,sha256=VFP2hMAfaocOnOKl0xpuy95dJV_LhJBTqcNceBB2z-0,73858
+mani_skill/envs/sapien_env.py,sha256=-reQWgDhCQga2RqFovQlHrsAKxxnvpcYYtK1AqxPEyI,75338
 mani_skill/envs/scene.py,sha256=4ZAIJs61fwPPhfDvc3v845sj_Ftsd1sSYaju10KnXbQ,48465
 mani_skill/envs/sim2real_env.py,sha256=3mkQX4TonE2pUC5_Atmx0IYDH2_v6GSwOPJvQMEvCNY,19214
 mani_skill/envs/template.py,sha256=0wnwKjnGOF7RvTR5Gz4VopaUiFxnIioXwmb4nPVxAs8,11939
@@ -720,6 +720,7 @@ mani_skill/utils/logging_utils.py,sha256=Iomnw2VrJ56A6kvyIm9Zhgp94Qf-WOf9WuMAGNB
 mani_skill/utils/precompile_mpm.py,sha256=D497m_0mcNifeBEV7wmG5gHU_SstNpOoHIb-OQiqvUA,291
 mani_skill/utils/registration.py,sha256=u8ftfGvQP4qzlKNqLQjGi3cRF_-h6Rz-28xbLkW_YJ4,9718
 mani_skill/utils/sapien_utils.py,sha256=QMV0jRZO51KzIMB5CVW_Ne-4fPw0-mqM4a3yhNZaMYo,16430
+mani_skill/utils/tree.py,sha256=jFyqCqVkarR6AaTDissVo33nOfbPfkZp8CsqEQmCV9Q,515
 mani_skill/utils/assets/README.md,sha256=5kkmsIiV64ySEGO34HaAlpjXTyrGs1KTV5WnofK46G0,70
 mani_skill/utils/assets/__init__.py,sha256=gQVKwAczcImTXArSltBWKlSUUuguO12sZYO3Jh5KLso,159
 mani_skill/utils/assets/data.py,sha256=xEuibRoEPBDN_vEU-MM5UWf6VDb1omE6BfZKPvlMPdI,8807
@@ -801,7 +802,7 @@ mani_skill/utils/scene_builder/table/assets/table.glb,sha256=yw69itZDjBFg8JXZAr9
 mani_skill/utils/structs/README.md,sha256=qnYKimp_ZkgNcduURrYQxVTimNmq_usDMKoQ8VtMdCs,286
 mani_skill/utils/structs/__init__.py,sha256=BItR3Xe0z6xCrMHAEaH0AAAVyeonsQ3q-DJUyRUibAA,524
 mani_skill/utils/structs/actor.py,sha256=L0p6vkr8rGtJmF22xAq8Q7nhXKnDD5dahzODSAko0bg,17394
-mani_skill/utils/structs/articulation.py,sha256=DfNu3irWZ7LWoMCQAR4t0F8QRnojYhmPBnk2LNYoTtw,38341
+mani_skill/utils/structs/articulation.py,sha256=QvWQsHdgCUCiGHVfUzZOeEqJ_OQsbi7CrxkCgjy3wyM,38491
 mani_skill/utils/structs/articulation_joint.py,sha256=xDQkCAXM3XZ56YgFqLwH5Ec8aFqhR5BqMSvDYCS0bzw,12972
 mani_skill/utils/structs/base.py,sha256=meGQK5Y4KtHKLnp9VeOZS2gtwg9tE55whuEeqOguBaI,19465
 mani_skill/utils/structs/decorators.py,sha256=Lv6wQ989dOnreo2tB-qopDnkeBp_jsn1pmfUR-OY8VQ,535
@@ -815,20 +816,21 @@ mani_skill/utils/visualization/__init__.py,sha256=0QF97UR8d7poMHo6m52DsAUXAmUb3S
 mani_skill/utils/visualization/jupyter_utils.py,sha256=dXXUQz-rFTOV_Xq5yA6YE6cXg7DPw15YStw37NgB5Qc,1322
 mani_skill/utils/visualization/misc.py,sha256=KrDCef7F5GmGOdiBQ4qFUnmUTe-7-nNBz2DVBGFD8YU,5041
 mani_skill/utils/visualization/renderer.py,sha256=afFWwSQEeL-9c5CsBT1uug-zugGjOr1FDzmvd45-9dk,1646
-mani_skill/utils/wrappers/__init__.py,sha256=XCQdTP2RQ95UXrDjy1QxgcKEuevceGU_5KUGl636J6I,276
+mani_skill/utils/wrappers/__init__.py,sha256=QfSTw9RNtQEtEAzD1RLx0WXjNjIS7wVM2IIkSkLT6Xw,321
 mani_skill/utils/wrappers/action_repeat.py,sha256=RhCtzt3fYCtD-CClIOhAzdycGwVTXP_FG61yEf-QLqY,3542
+mani_skill/utils/wrappers/cached_reset.py,sha256=KV9Sd-mIK9NM_nes-7HtO3HAkkuflctVnUejSN77ecE,6272
 mani_skill/utils/wrappers/flatten.py,sha256=GuHJ3fCOdj9G_jm--XgG8k0p2G1eJx4LY1tesQQjnkg,4913
 mani_skill/utils/wrappers/frame_stack.py,sha256=pCp83HqXnFxbsKRYgwXreNBHnhD-yF0R2_7jdtGOTWQ,4213
 mani_skill/utils/wrappers/gymnasium.py,sha256=p0kl29kkedD2arIvGskClKhYDBAH97mZO4rTepz62jQ,4174
-mani_skill/utils/wrappers/record.py,sha256=757N7D6m6bH02BjIlLvFxGwbQDqDouiRv7F_6Prp57k,37322
+mani_skill/utils/wrappers/record.py,sha256=73g-dvnFrXcFy7t2BksgH8UqBUVBZG8BbCNPEbdtBSc,37362
 mani_skill/utils/wrappers/visual_encoders.py,sha256=ISLO5ceaRkINhvce92VuZMDMCU3I4F7cQWFW2aVP-14,2205
 mani_skill/vector/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 mani_skill/vector/wrappers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-mani_skill/vector/wrappers/gymnasium.py,sha256=aNPB-2oGDLep8qzdsuTSIlwGGO0OGQAQ193LefOGoTk,7434
+mani_skill/vector/wrappers/gymnasium.py,sha256=SbFD9inb_a8Hm36ckpdZumxVRfVcw5l1-nfyWGF0KfA,7456
 mani_skill/vector/wrappers/sb3.py,sha256=SlXdiEPqcNHYMhJCzA29kBU6zK7DKTe1nc0L6Z3QQtY,4722
 mani_skill/viewer/__init__.py,sha256=srvDBsk4LQU75K2VIttrhiQ68p_ro7PSDqQRls2PY5c,1722
-mani_skill_nightly-2025.7.25.329.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
-mani_skill_nightly-2025.7.25.329.dist-info/METADATA,sha256=ZORmjNSU7HXw-F69GCdOiTQnUMeQKOaCBbpJIfaMCP8,9271
-mani_skill_nightly-2025.7.25.329.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
-mani_skill_nightly-2025.7.25.329.dist-info/top_level.txt,sha256=bkBgOVl_MZMoQx2aRFsSFEYlZLxjWlip5vtJ39FB3jA,11
-mani_skill_nightly-2025.7.25.329.dist-info/RECORD,,
+mani_skill_nightly-2025.7.25.606.dist-info/LICENSE,sha256=xx0jnfkXJvxRnG63LTGOxlggYnIysveWIZ6H3PNdCrQ,11357
+mani_skill_nightly-2025.7.25.606.dist-info/METADATA,sha256=Fz68W0vbVEG-GCrKasNzdf8kz8wMbUYdJnzsWDJCKFY,9271
+mani_skill_nightly-2025.7.25.606.dist-info/WHEEL,sha256=tZoeGjtWxWRfdplE7E3d45VPlLNQnvbKiYnx7gwAy8A,92
+mani_skill_nightly-2025.7.25.606.dist-info/top_level.txt,sha256=bkBgOVl_MZMoQx2aRFsSFEYlZLxjWlip5vtJ39FB3jA,11
+mani_skill_nightly-2025.7.25.606.dist-info/RECORD,,