mani-skill-nightly 2025.10.22.143__py3-none-any.whl → 2025.10.22.325__py3-none-any.whl

mani_skill/envs/sapien_env.py

@@ -2,7 +2,7 @@ import copy
 import gc
 import os
 from functools import cached_property
-from typing import Any, Dict, List, Optional, Sequence, Tuple, Union
+from typing import Any, Optional, Sequence, Tuple, Union
 
 import dacite
 import gymnasium as gym
@@ -80,7 +80,7 @@ class BaseEnv(gym.Env):
 
         viewer_camera_configs (dict): configurations of the viewer camera in the GUI to override any environment defaults. Similar usage as @sensor_configs.
 
-        robot_uids (Union[str, BaseAgent, List[Union[str, BaseAgent]]]): List of robots to instantiate and control in the environment.
+        robot_uids (Union[str, BaseAgent, list[Union[str, BaseAgent]]]): list of robots to instantiate and control in the environment.
 
         sim_config (Union[SimConfig, dict]): Configurations for simulation if used that override the environment defaults. If given
             a dictionary, it can just override specific attributes e.g. ``sim_config=dict(scene_config=dict(solver_iterations=25))``. If
@@ -118,7 +118,7 @@ class BaseEnv(gym.Env):
     """
 
     # fmt: off
-    SUPPORTED_ROBOTS: List[Union[str, Tuple[str]]] = None
+    SUPPORTED_ROBOTS: list[Union[str, Tuple[str]]] = None
     """Override this to enforce which robots or tuples of robots together are supported in the task. During env creation,
     setting robot_uids auto loads all desired robots into the scene, but not all tasks are designed to support some robot setups"""
     SUPPORTED_OBS_MODES = ("state", "state_dict", "none", "sensor_data", "any_textures", "pointcloud")
@@ -142,27 +142,27 @@ class BaseEnv(gym.Env):
     single_action_space: gym.Space
     """the unbatched action space of the environment"""
 
-    _sensors: Dict[str, BaseSensor]
+    _sensors: dict[str, BaseSensor]
     """all sensors configured in this environment"""
-    _sensor_configs: Dict[str, BaseSensorConfig]
+    _sensor_configs: dict[str, BaseSensorConfig]
     """all sensor configurations parsed from self._sensor_configs and agent._sensor_configs"""
-    _agent_sensor_configs: Dict[str, BaseSensorConfig]
+    _agent_sensor_configs: dict[str, BaseSensorConfig]
     """all agent sensor configs parsed from agent._sensor_configs"""
-    _human_render_cameras: Dict[str, Camera]
+    _human_render_cameras: dict[str, Camera]
     """cameras used for rendering the current environment retrievable via `env.render_rgb_array()`. These are not used to generate observations"""
-    _default_human_render_camera_configs: Dict[str, CameraConfig]
+    _default_human_render_camera_configs: dict[str, CameraConfig]
     """all camera configurations for cameras used for human render"""
-    _human_render_camera_configs: Dict[str, CameraConfig]
+    _human_render_camera_configs: dict[str, CameraConfig]
     """all camera configurations parsed from self._human_render_camera_configs"""
 
-    _hidden_objects: List[Union[Actor, Articulation]] = []
+    _hidden_objects: list[Union[Actor, Articulation]] = []
     """list of objects that are hidden during rendering when generating visual observations / running render_cameras()"""
 
     _main_rng: np.random.RandomState = None
     """main rng generator that generates episode seed sequences. For internal use only"""
     _batched_main_rng: BatchedRNG = None
     """the batched main RNG that generates episode seed sequences. For internal use only"""
-    _main_seed: List[int] = None
+    _main_seed: list[int] = None
     """main seed list for _main_rng and _batched_main_rng. _main_rng uses _main_seed[0]. For internal use only"""
     _episode_rng: np.random.RandomState = None
     """the numpy RNG that you can use to generate random numpy data. It is not recommended to use this. Instead use the _batched_episode_rng which helps ensure GPU and CPU simulation generate the same data with the same seeds."""
@@ -201,7 +201,7 @@ class BaseEnv(gym.Env):
         sensor_configs: Optional[dict] = dict(),
         human_render_camera_configs: Optional[dict] = dict(),
         viewer_camera_configs: Optional[dict] = dict(),
-        robot_uids: Union[str, BaseAgent, List[Union[str, BaseAgent]]] = None,
+        robot_uids: Union[str, BaseAgent, list[Union[str, BaseAgent]]] = None,
         sim_config: Union[SimConfig, dict] = dict(),
         reconfiguration_freq: Optional[int] = None,
         sim_backend: str = "auto",
@@ -439,7 +439,7 @@ class BaseEnv(gym.Env):
     def _default_sensor_configs(
         self,
     ) -> Union[
-        BaseSensorConfig, Sequence[BaseSensorConfig], Dict[str, BaseSensorConfig]
+        BaseSensorConfig, Sequence[BaseSensorConfig], dict[str, BaseSensorConfig]
     ]:
         """Add default (non-agent) sensors to the environment by returning sensor configurations. These can be overriden by the user at
         env creation time"""
@@ -448,7 +448,7 @@ class BaseEnv(gym.Env):
     def _default_human_render_camera_configs(
         self,
     ) -> Union[
-        CameraConfig, Sequence[CameraConfig], Dict[str, CameraConfig]
+        CameraConfig, Sequence[CameraConfig], dict[str, CameraConfig]
     ]:
         """Add default cameras for rendering when using render_mode='rgb_array'. These can be overriden by the user at env creation time """
         return []
@@ -498,7 +498,7 @@ class BaseEnv(gym.Env):
         """The current observation mode. This affects the observation returned by env.get_obs()"""
         return self._obs_mode
 
-    def get_obs(self, info: Optional[Dict] = None, unflattened: bool = False):
+    def get_obs(self, info: Optional[dict] = None, unflattened: bool = False):
         """
         Return the current observation of the environment. User may call this directly to get the current observation
         as opposed to taking a step with actions in the environment.
@@ -509,7 +509,7 @@ class BaseEnv(gym.Env):
         data in the info object by overriding the `self.evaluate` function.
 
         Args:
-            info (Dict): The info object of the environment. Generally should always be the result of `self.get_info()`.
+            info (dict): The info object of the environment. Generally should always be the result of `self.get_info()`.
                 If this is None (the default), this function will call `self.get_info()` itself
             unflattened (bool): Whether to return the observation without flattening even if the observation mode (`self.obs_mode`) asserts to return a flattened observation.
         """
@@ -543,7 +543,7 @@ class BaseEnv(gym.Env):
                 obs["state"] = common.flatten_state_dict(data, use_torch=True, device=self.device)
         return obs
 
-    def _get_obs_state_dict(self, info: Dict):
+    def _get_obs_state_dict(self, info: dict):
         """Get (ground-truth) state-based observations."""
         return dict(
             agent=self._get_obs_agent(),
@@ -555,7 +555,7 @@ class BaseEnv(gym.Env):
         Controller state is also included although most default controllers do not have any state."""
         return self.agent.get_proprioception()
 
-    def _get_obs_extra(self, info: Dict):
+    def _get_obs_extra(self, info: dict):
         """Get task-relevant extra observations. Usually defined on a task by task basis"""
         return dict()
 
@@ -564,11 +564,11 @@ class BaseEnv(gym.Env):
         for sensor in self._sensors.values():
             sensor.capture()
 
-    def get_sensor_images(self) -> Dict[str, Dict[str, torch.Tensor]]:
+    def get_sensor_images(self) -> dict[str, dict[str, torch.Tensor]]:
         """Get image (RGB) visualizations of what sensors currently sense. This function calls self._get_obs_sensor_data() internally which automatically hides objects and updates the render"""
         return self.scene.get_sensor_images(self._get_obs_sensor_data())
 
-    def get_sensor_params(self) -> Dict[str, Dict[str, torch.Tensor]]:
+    def get_sensor_params(self) -> dict[str, dict[str, torch.Tensor]]:
         """Get all sensor parameters."""
         params = dict()
         for name, sensor in self._sensors.items():
@@ -624,7 +624,7 @@ class BaseEnv(gym.Env):
             torch.cuda.synchronize()
         return sensor_obs
 
-    def _get_obs_with_sensor_data(self, info: Dict, apply_texture_transforms: bool = True) -> dict:
+    def _get_obs_with_sensor_data(self, info: dict, apply_texture_transforms: bool = True) -> dict:
         """Get the observation with sensor data"""
         return dict(
             agent=self._get_obs_agent(),
@@ -645,7 +645,7 @@ class BaseEnv(gym.Env):
     def reward_mode(self):
         return self._reward_mode
 
-    def get_reward(self, obs: Any, action: torch.Tensor, info: Dict):
+    def get_reward(self, obs: Any, action: torch.Tensor, info: dict):
         """
         Compute the reward for environment at its current state. observation data, the most recent action, and the info dictionary (generated by the self.evaluate() function)
         are provided as inputs. By default the observation data will be in its most raw form, a dictionary (no flattening, wrappers etc.)
@@ -653,7 +653,7 @@ class BaseEnv(gym.Env):
         Args:
             obs (Any): The observation data.
             action (torch.Tensor): The most recent action.
-            info (Dict): The info dictionary.
+            info (dict): The info dictionary.
         """
         if self._reward_mode == "sparse":
             reward = self.compute_sparse_reward(obs=obs, action=action, info=info)
@@ -669,7 +669,7 @@ class BaseEnv(gym.Env):
             raise NotImplementedError(self._reward_mode)
         return reward
 
-    def compute_sparse_reward(self, obs: Any, action: torch.Tensor, info: Dict):
+    def compute_sparse_reward(self, obs: Any, action: torch.Tensor, info: dict):
         """
 
         Computes the sparse reward. By default this function tries to use the success/fail information in
@@ -678,7 +678,7 @@ class BaseEnv(gym.Env):
         Args:
             obs (Any): The observation data. By default the observation data will be in its most raw form, a dictionary (no flattening, wrappers etc.)
             action (torch.Tensor): The most recent action.
-            info (Dict): The info dictionary.
+            info (dict): The info dictionary.
         """
         if "success" in info:
             if "fail" in info:
@@ -695,19 +695,19 @@ class BaseEnv(gym.Env):
                 reward = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
         return reward
 
-    def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: Dict):
+    def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: dict):
         """
         Compute the dense reward.
 
         Args:
             obs (Any): The observation data. By default the observation data will be in its most raw form, a dictionary (no flattening, wrappers etc.)
             action (torch.Tensor): The most recent action.
-            info (Dict): The info dictionary.
+            info (dict): The info dictionary.
         """
         raise NotImplementedError()
 
     def compute_normalized_dense_reward(
-        self, obs: Any, action: torch.Tensor, info: Dict
+        self, obs: Any, action: torch.Tensor, info: dict
     ):
         """
         Compute the normalized dense reward.
@@ -715,7 +715,7 @@ class BaseEnv(gym.Env):
         Args:
             obs (Any): The observation data. By default the observation data will be in its most raw form, a dictionary (no flattening, wrappers etc.)
             action (torch.Tensor): The most recent action.
-            info (Dict): The info dictionary.
+            info (dict): The info dictionary.
         """
         raise NotImplementedError()
 
@@ -1035,7 +1035,7 @@ class BaseEnv(gym.Env):
     # Step
     # -------------------------------------------------------------------------- #
 
-    def step(self, action: Union[None, np.ndarray, torch.Tensor, Dict]):
+    def step(self, action: Union[None, np.ndarray, torch.Tensor, dict]):
         """
         Take a step through the environment with an action. Actions are automatically clipped to the action space.
 
@@ -1067,7 +1067,7 @@ class BaseEnv(gym.Env):
         )
 
     def _step_action(
-        self, action: Union[None, np.ndarray, torch.Tensor, Dict]
+        self, action: Union[None, np.ndarray, torch.Tensor, dict]
     ) -> Union[None, torch.Tensor]:
         set_action = False
         action_is_unbatched = False
@@ -1286,7 +1286,7 @@ class BaseEnv(gym.Env):
         """
         return common.flatten_state_dict(self.get_state_dict(), use_torch=True)
 
-    def set_state_dict(self, state: Dict, env_idx: torch.Tensor = None):
+    def set_state_dict(self, state: dict, env_idx: torch.Tensor = None):
         """
         Set environment state with a state dictionary. Override to include task information (e.g., goal)
 

mani_skill/envs/scene.py

@@ -1,6 +1,6 @@
 from dataclasses import dataclass
 from functools import cached_property
-from typing import Any, Dict, List, Optional, Tuple, Union
+from typing import Any, Optional, Tuple, Union
 
 import numpy as np
 import sapien
@@ -33,8 +33,8 @@ if SAPIEN_RENDER_SYSTEM == "3.1":
 
 @dataclass
 class StateDictRegistry:
-    actors: Dict[str, Actor]
-    articulations: Dict[str, Articulation]
+    actors: dict[str, Actor]
+    articulations: dict[str, Articulation]
 
 
 class ManiSkillScene:
@@ -48,7 +48,7 @@ class ManiSkillScene:
 
     def __init__(
         self,
-        sub_scenes: Optional[List[sapien.Scene]] = None,
+        sub_scenes: Optional[list[sapien.Scene]] = None,
         sim_config: SimConfig = SimConfig(),
         debug_mode: bool = True,
         device: Device = None,
@@ -75,18 +75,18 @@ class ManiSkillScene:
         self.backend = backend  # references the backend object stored in BaseEnv class
 
         self.render_system_group: sapien.render.RenderSystemGroup = None
-        self.camera_groups: Dict[str, sapien.render.RenderCameraGroup] = dict()
+        self.camera_groups: dict[str, sapien.render.RenderCameraGroup] = dict()
 
-        self.actors: Dict[str, Actor] = dict()
-        self.articulations: Dict[str, Articulation] = dict()
+        self.actors: dict[str, Actor] = dict()
+        self.articulations: dict[str, Articulation] = dict()
 
-        self.actor_views: Dict[str, Actor] = dict()
+        self.actor_views: dict[str, Actor] = dict()
         """views of actors in any sub-scenes created by using Actor.merge and queryable as if it were a single Actor"""
-        self.articulation_views: Dict[str, Articulation] = dict()
+        self.articulation_views: dict[str, Articulation] = dict()
         """views of articulations in any sub-scenes created by using Articulation.merge and queryable as if it were a single Articulation"""
 
-        self.sensors: Dict[str, BaseSensor] = dict()
-        self.human_render_cameras: Dict[str, Camera] = dict()
+        self.sensors: dict[str, BaseSensor] = dict()
+        self.human_render_cameras: dict[str, Camera] = dict()
         self._sensors_initialized = False
         self._human_render_cameras_initialized = False
 
@@ -97,12 +97,12 @@ class ManiSkillScene:
         self._needs_fetch = False
         """Used internally to raise some errors ahead of time of when there may be undefined behaviors"""
 
-        self.pairwise_contact_queries: Dict[
+        self.pairwise_contact_queries: dict[
             str, physx.PhysxGpuContactPairImpulseQuery
         ] = dict()
         """dictionary mapping pairwise contact query keys to GPU contact queries. Used in GPU simulation only to cache queries as
         query creation will pause any GPU sim computation"""
-        self._pairwise_contact_query_unique_hashes: Dict[str, int] = dict()
+        self._pairwise_contact_query_unique_hashes: dict[str, int] = dict()
         """maps keys in self.pairwise_contact_queries to unique hashes dependent on the actual objects involved in the query.
         This is used to determine automatically when to rebuild contact queries as keys for self.pairwise_contact_queries are kept
         non-unique between episode resets in order to be easily rebuilt and deallocate old queries. This essentially acts as a way
@@ -203,7 +203,7 @@ class ManiSkillScene:
         height,
         near,
         far,
-        fovy: Union[float, List, None] = None,
+        fovy: Union[float, list, None] = None,
         intrinsic: Union[Array, None] = None,
         mount: Union[Actor, Link, None] = None,
     ) -> RenderCamera:
@@ -225,7 +225,7 @@ class ManiSkillScene:
         height,
         near,
         far,
-        fovy: Union[float, List, None] = None,
+        fovy: Union[float, list, None] = None,
         intrinsic: Union[Array, None] = None,
         mount: Union[Actor, Link, None] = None,
     ) -> RenderCamera:
@@ -304,7 +304,7 @@ class ManiSkillScene:
         height,
         near,
         far,
-        fovy: Union[float, List, None] = None,
+        fovy: Union[float, list, None] = None,
         intrinsic: Union[Array, None] = None,
         mount: Union[Actor, Link, None] = None,
     ) -> RenderCamera:
@@ -583,7 +583,7 @@ class ManiSkillScene:
         shadow_near=0.1,
         shadow_far=10.0,
         shadow_map_size=2048,
-        scene_idxs: Optional[List[int]] = None,
+        scene_idxs: Optional[list[int]] = None,
     ):
         if scene_idxs is None:
             scene_idxs = list(range(len(self.sub_scenes)))
@@ -619,7 +619,7 @@ class ManiSkillScene:
         shadow_near=-10.0,
         shadow_far=10.0,
         shadow_map_size=2048,
-        scene_idxs: Optional[List[int]] = None,
+        scene_idxs: Optional[list[int]] = None,
     ):
         if scene_idxs is None:
             scene_idxs = list(range(len(self.sub_scenes)))
@@ -664,7 +664,7 @@ class ManiSkillScene:
         shadow_near=0.1,
         shadow_far=10.0,
         shadow_map_size=2048,
-        scene_idxs: Optional[List[int]] = None,
+        scene_idxs: Optional[list[int]] = None,
     ):
         if scene_idxs is None:
             scene_idxs = list(range(len(self.sub_scenes)))
@@ -873,7 +873,7 @@ class ManiSkillScene:
             del state_dict["articulations"]
         return state_dict
 
-    def set_sim_state(self, state: Dict, env_idx: torch.Tensor = None):
+    def set_sim_state(self, state: dict, env_idx: torch.Tensor = None):
         if env_idx is not None:
             prev_reset_mask = self._reset_mask.clone()
             # safe guard against setting the wrong states
@@ -908,7 +908,7 @@ class ManiSkillScene:
                 for scene in self.sub_scenes:
                     scene.update_render()
             self.px.gpu_init()
-        self.non_static_actors: List[Actor] = []
+        self.non_static_actors: list[Actor] = []
         # find non static actors, and set data indices that are now available after gpu_init was called
         for actor in self.actors.values():
             if actor.px_body_type == "static":
@@ -990,7 +990,7 @@ class ManiSkillScene:
     # ---------------------------------------------------------------------------- #
     def _get_all_render_bodies(
         self,
-    ) -> List[Tuple[sapien.render.RenderBodyComponent, int]]:
+    ) -> list[Tuple[sapien.render.RenderBodyComponent, int]]:
         all_render_bodies = []
         for actor in self.actors.values():
             if actor.px_body_type == "static":
@@ -1078,13 +1078,13 @@ class ManiSkillScene:
 
         self.render_system_group = sync_manager
 
-    def _gpu_setup_sensors(self, sensors: Dict[str, BaseSensor]):
+    def _gpu_setup_sensors(self, sensors: dict[str, BaseSensor]):
         if SAPIEN_RENDER_SYSTEM == "3.1":
             self._sapien_31_gpu_setup_sensors(sensors)
         else:
             self._sapien_gpu_setup_sensors(sensors)
 
-    def _sapien_gpu_setup_sensors(self, sensors: Dict[str, BaseSensor]):
+    def _sapien_gpu_setup_sensors(self, sensors: dict[str, BaseSensor]):
         for name, sensor in sensors.items():
             if isinstance(sensor, Camera):
                 try:
@@ -1126,8 +1126,8 @@ class ManiSkillScene:
                 )
 
     def get_sensor_images(
-        self, obs: Dict[str, Any]
-    ) -> Dict[str, Dict[str, torch.Tensor]]:
+        self, obs: dict[str, Any]
+    ) -> dict[str, dict[str, torch.Tensor]]:
         """Get raw sensor data as images for visualization purposes."""
         sensor_data = dict()
         for name, sensor in self.sensors.items():
@@ -1136,7 +1136,7 @@ class ManiSkillScene:
 
     def get_human_render_camera_images(
         self, camera_name: str = None
-    ) -> Dict[str, torch.Tensor]:
+    ) -> dict[str, torch.Tensor]:
         image_data = dict()
         if self.gpu_sim_enabled:
             if self.parallel_in_single_scene:

mani_skill/envs/scenes/base_env.py

@@ -1,4 +1,4 @@
-from typing import Any, Dict, Union
+from typing import Any, Union
 
 import numpy as np
 import sapien as sapien
@@ -142,11 +142,11 @@ class SceneManipulationEnv(BaseEnv):
     def evaluate(self) -> dict:
         return dict()
 
-    def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: Dict):
+    def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: dict):
         return 0
 
     def compute_normalized_dense_reward(
-        self, obs: Any, action: torch.Tensor, info: Dict
+        self, obs: Any, action: torch.Tensor, info: dict
     ):
         return self.compute_dense_reward(obs=obs, action=action, info=info) / 1
 

mani_skill/envs/sim2real_env.py

@@ -1,5 +1,5 @@
 import time
-from typing import Any, Callable, Dict, List, Optional
+from typing import Any, Callable, Optional
 
 import gymnasium as gym
 import numpy as np
@@ -37,7 +37,7 @@ class Sim2RealEnv(gym.Env):
                     self.agent.reset(qpos=self.base_sim_env.agent.robot.qpos.cpu().flatten())
                     input("Press enter if the environment is reset")
 
-        sensor_data_preprocessing_function (Optional[Callable[[Dict], Dict]]): The function to call to process the sensor data returned by the BaseRealAgent.get_sensor_data function.
+        sensor_data_preprocessing_function (Optional[Callable[[dict], dict]]): The function to call to process the sensor data returned by the BaseRealAgent.get_sensor_data function.
             By default this is None and we use a default processing function which does the following for each sensor type:
             - Camera: Perform a center crop of the real sensor image (rgb or depth) to have the same aspect ratio as the simulation sensor image. Then resize the image to the simulation sensor image shape using cv2.resize
 
@@ -56,7 +56,7 @@ class Sim2RealEnv(gym.Env):
         real_reset_function: Optional[
             Callable[["Sim2RealEnv", Optional[int], Optional[dict]], None]
         ] = None,
-        sensor_data_preprocessing_function: Optional[Callable[[Dict], Dict]] = None,
+        sensor_data_preprocessing_function: Optional[Callable[[dict], dict]] = None,
         render_mode: Optional[str] = "sensors",
         skip_data_checks: bool = False,
         control_freq: Optional[int] = None,
@@ -122,7 +122,7 @@ class Sim2RealEnv(gym.Env):
                 return self
 
         cur_env = self.sim_env
-        wrappers: List[gym.Wrapper] = []
+        wrappers: list[gym.Wrapper] = []
         while isinstance(cur_env, gym.Wrapper):
             wrappers.append(cur_env)
             cur_env = cur_env.env
@@ -232,7 +232,7 @@ class Sim2RealEnv(gym.Env):
         # using the original user implemented sim env's _get_obs_agent function in case they modify it e.g. to remove qvel values as they might be too noisy
         return self.base_sim_env.__class__._get_obs_agent(self)
 
-    def _get_obs_extra(self, info: Dict):
+    def _get_obs_extra(self, info: dict):
         # using the original user implemented sim env's _get_obs_extra function in case they modify it e.g. to include engineered features like the tcp_pose of the robot
         try:
             return self.base_sim_env.__class__._get_obs_extra(self, info)
@@ -259,7 +259,7 @@ class Sim2RealEnv(gym.Env):
         return data
 
     def _get_obs_with_sensor_data(
-        self, info: Dict, apply_texture_transforms: bool = True
+        self, info: dict, apply_texture_transforms: bool = True
     ) -> dict:
         """Get the observation with sensor data"""
         return self.base_sim_env.__class__._get_obs_with_sensor_data(
@@ -298,7 +298,7 @@ class Sim2RealEnv(gym.Env):
     def get_reward(self, obs, action, info):
         return self.base_sim_env.__class__.get_reward(self, obs, action, info)
 
-    def compute_sparse_reward(self, obs: Any, action: torch.Tensor, info: Dict):
+    def compute_sparse_reward(self, obs: Any, action: torch.Tensor, info: dict):
         """
         Computes the sparse reward. By default this function tries to use the success/fail information in
         returned by the evaluate function and gives +1 if success, -1 if fail, 0 otherwise"""
@@ -306,11 +306,11 @@ class Sim2RealEnv(gym.Env):
             self, obs, action, info
         )
 
-    def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: Dict):
+    def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: dict):
         raise NotImplementedError()
 
     def compute_normalized_dense_reward(
-        self, obs: Any, action: torch.Tensor, info: Dict
+        self, obs: Any, action: torch.Tensor, info: dict
     ):
         raise NotImplementedError()
 
@@ -347,7 +347,7 @@ class Sim2RealEnv(gym.Env):
         self.agent.stop()
 
     def preprocess_sensor_data(
-        self, sensor_data: Dict, sensor_names: Optional[List[str]] = None
+        self, sensor_data: dict, sensor_names: Optional[list[str]] = None
     ):
         import cv2
 

mani_skill/envs/tasks/control/ant.py

@@ -1,5 +1,5 @@
 import os
-from typing import Any, Dict, Optional, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import sapien
@@ -152,7 +152,7 @@ class AntEnv(BaseEnv):
             link.name for link in self.active_links if "foot" in link.name
         ]
 
-    def _initialize_episode(self, env_idx: torch.Tensor, options: Dict):
+    def _initialize_episode(self, env_idx: torch.Tensor, options: dict):
         with torch.device(self.device):
             b = len(env_idx)
             # set agent root pose - torso now centered at dummy root at (0,0,0)
@@ -235,7 +235,7 @@ class AntEnv(BaseEnv):
         )
 
     # cache re-used computation
-    def evaluate(self) -> Dict:
+    def evaluate(self) -> dict:
         link_angvels, link_linvels, cmass_linvel = self.get_vels
         return dict(
             link_angvels=link_angvels,
@@ -243,7 +243,7 @@ class AntEnv(BaseEnv):
             cmass_linvel=cmass_linvel,
         )
 
-    def _get_obs_extra(self, info: Dict):
+    def _get_obs_extra(self, info: dict):
         obs = super()._get_obs_extra(info)
         if self.obs_mode_struct.use_state:
             obs.update(
@@ -282,14 +282,14 @@ class AntEnv(BaseEnv):
             .view(-1)
         )
 
-    def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: Dict):
+    def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: dict):
         small_control = (4 + self.control_rew(action)) / 5
         return (
             small_control * self.move_x_rew(info, self.move_speed) * self.standing_rew()
         )
 
     def compute_normalized_dense_reward(
-        self, obs: Any, action: torch.Tensor, info: Dict
+        self, obs: Any, action: torch.Tensor, info: dict
     ):
         return self.compute_dense_reward(obs, action, info)
 

mani_skill/envs/tasks/control/cartpole.py

@@ -1,7 +1,7 @@
 """Adapted from https://github.com/google-deepmind/dm_control/blob/main/dm_control/suite/cartpole.py"""
 
 import os
-from typing import Any, Dict, Optional, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import sapien
@@ -132,7 +132,7 @@ class CartpoleEnv(BaseEnv):
     def evaluate(self):
         return dict()
 
-    def _get_obs_extra(self, info: Dict):
+    def _get_obs_extra(self, info: dict):
         obs = dict(
             velocity=self.agent.robot.links_map["pole_1"].linear_velocity,
             angular_velocity=self.agent.robot.links_map["pole_1"].angular_velocity,
@@ -143,7 +143,7 @@ class CartpoleEnv(BaseEnv):
     def pole_angle_cosine(self):
         return torch.cos(self.agent.robot.joints_map["hinge_1"].qpos)
 
-    def compute_dense_reward(self, obs: Any, action: Array, info: Dict):
+    def compute_dense_reward(self, obs: Any, action: Array, info: dict):
         cart_pos = self.agent.robot.links_map["cart"].pose.p[
             :, 0
         ]  # (B, ), we only care about x position
@@ -169,7 +169,7 @@ class CartpoleEnv(BaseEnv):
         reward = upright * centered * small_control * small_velocity
         return reward
 
-    def compute_normalized_dense_reward(self, obs: Any, action: Array, info: Dict):
+    def compute_normalized_dense_reward(self, obs: Any, action: Array, info: dict):
         # this should be equal to compute_dense_reward / max possible reward
         max_reward = 1.0
         return self.compute_dense_reward(obs=obs, action=action, info=info) / max_reward

mani_skill/envs/tasks/control/hopper.py

@@ -1,7 +1,7 @@
 """Adapted from https://github.com/google-deepmind/dm_control/blob/main/dm_control/suite/hopper.py"""
 
 import os
-from typing import Any, Dict, Optional, Union
+from typing import Any, Optional, Union
 
 import numpy as np
 import sapien
@@ -163,7 +163,7 @@ class HopperEnv(BaseEnv):
         self.planar_scene = PlanarSceneBuilder(env=self)
         self.planar_scene.build()
 
-    def _initialize_episode(self, env_idx: torch.Tensor, options: Dict):
+    def _initialize_episode(self, env_idx: torch.Tensor, options: dict):
         with torch.device(self.device):
             b = len(env_idx)
             # qpos sampled same as dm_control, but ensure no self intersection explicitly here
@@ -206,7 +206,7 @@ class HopperEnv(BaseEnv):
         return torch.log1p(force_mag)
 
     # dm_control also includes foot pressures as state obs space
-    def _get_obs_state_dict(self, info: Dict):
+    def _get_obs_state_dict(self, info: dict):
         return dict(
             agent=self._get_obs_agent(),
             toe_touch=self.touch("foot_toe"),
@@ -231,11 +231,11 @@ class HopperStandEnv(HopperEnv):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-    def compute_dense_reward(self, obs: Any, action: Array, info: Dict):
+    def compute_dense_reward(self, obs: Any, action: Array, info: dict):
         standing = rewards.tolerance(self.height, lower=_STAND_HEIGHT, upper=2.0)
         return standing.view(-1)
 
-    def compute_normalized_dense_reward(self, obs: Any, action: Array, info: Dict):
+    def compute_normalized_dense_reward(self, obs: Any, action: Array, info: dict):
         # this should be equal to compute_dense_reward / max possible reward
         max_reward = 1.0
         return self.compute_dense_reward(obs=obs, action=action, info=info) / max_reward
@@ -258,7 +258,7 @@ class HopperHopEnv(HopperEnv):
     def __init__(self, *args, **kwargs):
         super().__init__(*args, **kwargs)
 
-    def compute_dense_reward(self, obs: Any, action: Array, info: Dict):
+    def compute_dense_reward(self, obs: Any, action: Array, info: dict):
         standing = rewards.tolerance(self.height, lower=_STAND_HEIGHT, upper=2.0)
         hopping = rewards.tolerance(
             self.subtreelinvelx,
@@ -271,6 +271,6 @@ class HopperHopEnv(HopperEnv):
 
         return standing.view(-1) * hopping.view(-1)
 
-    def compute_normalized_dense_reward(self, obs: Any, action: Array, info: Dict):
+    def compute_normalized_dense_reward(self, obs: Any, action: Array, info: dict):
         max_reward = 1.0
         return self.compute_dense_reward(obs=obs, action=action, info=info) / max_reward