mani-skill-nightly 2025.10.21.2011__py3-none-any.whl → 2025.10.22.157__py3-none-any.whl
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Potentially problematic release.
This version of mani-skill-nightly might be problematic. Click here for more details.
- mani_skill/agents/base_agent.py +20 -14
- mani_skill/agents/base_real_agent.py +6 -6
- mani_skill/agents/controllers/base_controller.py +6 -6
- mani_skill/agents/controllers/pd_joint_pos.py +2 -2
- mani_skill/agents/controllers/utils/kinematics.py +27 -12
- mani_skill/agents/multi_agent.py +5 -5
- mani_skill/agents/registration.py +3 -4
- mani_skill/agents/robots/allegro_hand/allegro.py +1 -2
- mani_skill/agents/robots/allegro_hand/allegro_touch.py +3 -3
- mani_skill/agents/robots/dclaw/dclaw.py +2 -3
- mani_skill/agents/robots/fetch/fetch.py +2 -2
- mani_skill/agents/robots/floating_ability_hand/floating_ability_hand.py +10 -13
- mani_skill/agents/robots/floating_robotiq_2f_85_gripper/floating_robotiq_2f_85_gripper.py +2 -2
- mani_skill/agents/robots/lerobot/manipulator.py +4 -4
- mani_skill/agents/robots/panda/panda_stick.py +2 -2
- mani_skill/agents/robots/trifingerpro/trifingerpro.py +1 -2
- mani_skill/agents/robots/xarm/xarm7_ability.py +2 -2
- mani_skill/agents/utils.py +2 -2
- mani_skill/envs/minimal_template.py +4 -4
- mani_skill/envs/sapien_env.py +36 -33
- mani_skill/envs/scene.py +27 -27
- mani_skill/envs/scenes/base_env.py +3 -3
- mani_skill/envs/sim2real_env.py +10 -10
- mani_skill/envs/tasks/control/ant.py +6 -6
- mani_skill/envs/tasks/control/cartpole.py +4 -4
- mani_skill/envs/tasks/control/hopper.py +7 -7
- mani_skill/envs/tasks/control/humanoid.py +20 -20
- mani_skill/envs/tasks/dexterity/insert_flower.py +41 -23
- mani_skill/envs/tasks/dexterity/rotate_single_object_in_hand.py +6 -6
- mani_skill/envs/tasks/dexterity/rotate_valve.py +5 -5
- mani_skill/envs/tasks/digital_twins/base_env.py +4 -4
- mani_skill/envs/tasks/digital_twins/bridge_dataset_eval/base_env.py +22 -12
- mani_skill/envs/tasks/digital_twins/so100_arm/grasp_cube.py +4 -4
- mani_skill/envs/tasks/drawing/draw.py +1 -3
- mani_skill/envs/tasks/drawing/draw_svg.py +6 -8
- mani_skill/envs/tasks/drawing/draw_triangle.py +1 -2
- mani_skill/envs/tasks/empty_env.py +1 -3
- mani_skill/envs/tasks/fmb/fmb.py +1 -2
- mani_skill/envs/tasks/humanoid/humanoid_pick_place.py +7 -7
- mani_skill/envs/tasks/humanoid/humanoid_stand.py +5 -5
- mani_skill/envs/tasks/humanoid/transport_box.py +4 -4
- mani_skill/envs/tasks/mobile_manipulation/open_cabinet_drawer.py +8 -8
- mani_skill/envs/tasks/mobile_manipulation/robocasa/kitchen.py +2 -3
- mani_skill/envs/tasks/quadruped/quadruped_reach.py +5 -5
- mani_skill/envs/tasks/quadruped/quadruped_spin.py +5 -5
- mani_skill/envs/tasks/rotate_cube.py +4 -4
- mani_skill/envs/tasks/tabletop/assembling_kits.py +2 -2
- mani_skill/envs/tasks/tabletop/lift_peg_upright.py +4 -4
- mani_skill/envs/tasks/tabletop/peg_insertion_side.py +4 -4
- mani_skill/envs/tasks/tabletop/pick_clutter_ycb.py +4 -4
- mani_skill/envs/tasks/tabletop/pick_cube.py +4 -4
- mani_skill/envs/tasks/tabletop/pick_single_ycb.py +5 -5
- mani_skill/envs/tasks/tabletop/place_sphere.py +4 -4
- mani_skill/envs/tasks/tabletop/plug_charger.py +2 -2
- mani_skill/envs/tasks/tabletop/poke_cube.py +4 -4
- mani_skill/envs/tasks/tabletop/pull_cube.py +5 -5
- mani_skill/envs/tasks/tabletop/pull_cube_tool.py +4 -4
- mani_skill/envs/tasks/tabletop/push_cube.py +6 -6
- mani_skill/envs/tasks/tabletop/push_t.py +4 -4
- mani_skill/envs/tasks/tabletop/roll_ball.py +4 -4
- mani_skill/envs/tasks/tabletop/stack_cube.py +4 -4
- mani_skill/envs/tasks/tabletop/stack_pyramid.py +44 -25
- mani_skill/envs/tasks/tabletop/turn_faucet.py +4 -4
- mani_skill/envs/tasks/tabletop/two_robot_pick_cube.py +4 -4
- mani_skill/envs/tasks/tabletop/two_robot_stack_cube.py +4 -4
- mani_skill/envs/template.py +4 -4
- mani_skill/envs/utils/observations/observations.py +2 -3
- mani_skill/envs/utils/randomization/batched_rng.py +7 -7
- mani_skill/envs/utils/randomization/samplers.py +2 -2
- mani_skill/examples/benchmarking/envs/maniskill/franka_move.py +2 -2
- mani_skill/examples/benchmarking/envs/maniskill/franka_pick_cube.py +2 -2
- mani_skill/examples/benchmarking/profiling.py +2 -2
- mani_skill/examples/demo_random_action.py +1 -1
- mani_skill/render/shaders.py +5 -5
- mani_skill/sensors/base_sensor.py +1 -2
- mani_skill/sensors/camera.py +4 -4
- mani_skill/trajectory/replay_trajectory.py +0 -1
- mani_skill/utils/assets/data.py +3 -3
- mani_skill/utils/building/_mjcf_loader.py +11 -11
- mani_skill/utils/building/actor_builder.py +4 -4
- mani_skill/utils/building/articulation_builder.py +3 -3
- mani_skill/utils/building/mjcf_loader.py +6 -6
- mani_skill/utils/building/urdf_loader.py +6 -6
- mani_skill/utils/common.py +2 -2
- mani_skill/utils/geometry/bounding_cylinder.py +4 -4
- mani_skill/utils/geometry/geometry.py +1 -3
- mani_skill/utils/geometry/trimesh_utils.py +1 -3
- mani_skill/utils/gym_utils.py +2 -4
- mani_skill/utils/registration.py +6 -6
- mani_skill/utils/sapien_utils.py +21 -21
- mani_skill/utils/scene_builder/ai2thor/constants.py +1 -2
- mani_skill/utils/scene_builder/ai2thor/scene_builder.py +9 -9
- mani_skill/utils/scene_builder/control/planar/scene_builder.py +2 -4
- mani_skill/utils/scene_builder/kitchen_counter/scene_builder.py +1 -2
- mani_skill/utils/scene_builder/registration.py +1 -2
- mani_skill/utils/scene_builder/replicacad/rearrange/scene_builder.py +16 -16
- mani_skill/utils/scene_builder/replicacad/scene_builder.py +15 -15
- mani_skill/utils/scene_builder/robocasa/fixtures/windows.py +2 -4
- mani_skill/utils/scene_builder/robocasa/scene_builder.py +5 -5
- mani_skill/utils/scene_builder/scene_builder.py +15 -15
- mani_skill/utils/scene_builder/table/scene_builder.py +1 -2
- mani_skill/utils/structs/actor.py +6 -6
- mani_skill/utils/structs/articulation.py +32 -30
- mani_skill/utils/structs/articulation_joint.py +6 -6
- mani_skill/utils/structs/base.py +14 -9
- mani_skill/utils/structs/drive.py +2 -2
- mani_skill/utils/structs/link.py +10 -8
- mani_skill/utils/structs/pose.py +3 -3
- mani_skill/utils/structs/render_camera.py +4 -4
- mani_skill/utils/structs/types.py +3 -1
- mani_skill/utils/visualization/jupyter_utils.py +1 -3
- mani_skill/utils/visualization/misc.py +5 -5
- mani_skill/utils/wrappers/cached_reset.py +5 -3
- mani_skill/utils/wrappers/flatten.py +1 -2
- mani_skill/utils/wrappers/record.py +10 -8
- mani_skill/utils/wrappers/visual_encoders.py +2 -2
- mani_skill/vector/wrappers/gymnasium.py +23 -13
- mani_skill/vector/wrappers/sb3.py +5 -5
- {mani_skill_nightly-2025.10.21.2011.dist-info → mani_skill_nightly-2025.10.22.157.dist-info}/METADATA +1 -1
- {mani_skill_nightly-2025.10.21.2011.dist-info → mani_skill_nightly-2025.10.22.157.dist-info}/RECORD +124 -124
- {mani_skill_nightly-2025.10.21.2011.dist-info → mani_skill_nightly-2025.10.22.157.dist-info}/WHEEL +0 -0
- {mani_skill_nightly-2025.10.21.2011.dist-info → mani_skill_nightly-2025.10.22.157.dist-info}/licenses/LICENSE +0 -0
- {mani_skill_nightly-2025.10.21.2011.dist-info → mani_skill_nightly-2025.10.22.157.dist-info}/licenses/LICENSE-3RD-PARTY +0 -0
- {mani_skill_nightly-2025.10.21.2011.dist-info → mani_skill_nightly-2025.10.22.157.dist-info}/top_level.txt +0 -0
|
@@ -3,7 +3,7 @@ from __future__ import annotations
|
|
|
3
3
|
import typing
|
|
4
4
|
from dataclasses import dataclass
|
|
5
5
|
from functools import cached_property
|
|
6
|
-
from typing import TYPE_CHECKING,
|
|
6
|
+
from typing import TYPE_CHECKING, Optional
|
|
7
7
|
|
|
8
8
|
import numpy as np
|
|
9
9
|
import sapien.physx as physx
|
|
@@ -39,7 +39,7 @@ class ArticulationJoint(BaseStruct[physx.PhysxArticulationJoint]):
|
|
|
39
39
|
parent_link: Optional[Link] = None
|
|
40
40
|
name: str = None
|
|
41
41
|
|
|
42
|
-
_physx_articulations:
|
|
42
|
+
_physx_articulations: list[physx.PhysxArticulation] = None
|
|
43
43
|
|
|
44
44
|
def __str__(self):
|
|
45
45
|
return f"<{self.name}: struct of type {self.__class__}; managing {self._num_objs} {self._objs[0].__class__} objects>"
|
|
@@ -53,8 +53,8 @@ class ArticulationJoint(BaseStruct[physx.PhysxArticulationJoint]):
|
|
|
53
53
|
@classmethod
|
|
54
54
|
def create(
|
|
55
55
|
cls,
|
|
56
|
-
physx_joints:
|
|
57
|
-
physx_articulations:
|
|
56
|
+
physx_joints: list[physx.PhysxArticulationJoint],
|
|
57
|
+
physx_articulations: list[physx.PhysxArticulation],
|
|
58
58
|
scene: ManiSkillScene,
|
|
59
59
|
scene_idxs: torch.Tensor,
|
|
60
60
|
joint_index: torch.Tensor,
|
|
@@ -232,7 +232,7 @@ class ArticulationJoint(BaseStruct[physx.PhysxArticulationJoint]):
|
|
|
232
232
|
return torch.tensor([obj.dof for obj in self._objs])
|
|
233
233
|
|
|
234
234
|
@property
|
|
235
|
-
def drive_mode(self) ->
|
|
235
|
+
def drive_mode(self) -> list[typing.Literal["force", "acceleration"]]:
|
|
236
236
|
"""
|
|
237
237
|
:type: typing.Literal['force', 'acceleration']
|
|
238
238
|
"""
|
|
@@ -375,7 +375,7 @@ class ArticulationJoint(BaseStruct[physx.PhysxArticulationJoint]):
|
|
|
375
375
|
@property
|
|
376
376
|
def type(
|
|
377
377
|
self,
|
|
378
|
-
) ->
|
|
378
|
+
) -> list[
|
|
379
379
|
typing.Literal["fixed", "revolute", "revolute_unwrapped", "prismatic", "free"]
|
|
380
380
|
]:
|
|
381
381
|
return [obj.type for obj in self._objs]
|
mani_skill/utils/structs/base.py
CHANGED
|
@@ -2,7 +2,7 @@ from __future__ import annotations
|
|
|
2
2
|
|
|
3
3
|
from dataclasses import dataclass
|
|
4
4
|
from functools import cached_property
|
|
5
|
-
from typing import TYPE_CHECKING, Generic,
|
|
5
|
+
from typing import TYPE_CHECKING, Generic, TypeVar
|
|
6
6
|
|
|
7
7
|
import numpy as np
|
|
8
8
|
import sapien.physx as physx
|
|
@@ -24,7 +24,7 @@ class BaseStruct(Generic[T]):
|
|
|
24
24
|
Base class of all structs that manage sapien objects on CPU/GPU
|
|
25
25
|
"""
|
|
26
26
|
|
|
27
|
-
_objs:
|
|
27
|
+
_objs: list[T]
|
|
28
28
|
"""list of objects of type T managed by this dataclass. This should not be modified after initialization. The struct hash is dependent on the hash of this list."""
|
|
29
29
|
_scene_idxs: torch.Tensor
|
|
30
30
|
"""a list of indexes parallel to `self._objs` indicating which sub-scene each of those objects are actually in by index"""
|
|
@@ -67,7 +67,7 @@ class BaseStruct(Generic[T]):
|
|
|
67
67
|
|
|
68
68
|
@dataclass
|
|
69
69
|
class PhysxRigidBaseComponentStruct(BaseStruct[T], Generic[T]):
|
|
70
|
-
_bodies:
|
|
70
|
+
_bodies: list[physx.PhysxRigidBaseComponent]
|
|
71
71
|
|
|
72
72
|
# ---------------------------------------------------------------------------- #
|
|
73
73
|
# API from physx.PhysxRigidBaseComponent
|
|
@@ -91,7 +91,7 @@ class PhysxRigidBaseComponentStruct(BaseStruct[T], Generic[T]):
|
|
|
91
91
|
|
|
92
92
|
@dataclass
|
|
93
93
|
class PhysxRigidBodyComponentStruct(PhysxRigidBaseComponentStruct[T], Generic[T]):
|
|
94
|
-
_bodies:
|
|
94
|
+
_bodies: list[physx.PhysxRigidBodyComponent]
|
|
95
95
|
_body_data_name: str
|
|
96
96
|
_body_data_index_internal: slice = None
|
|
97
97
|
|
|
@@ -210,7 +210,7 @@ class PhysxRigidBodyComponentStruct(PhysxRigidBaseComponentStruct[T], Generic[T]
|
|
|
210
210
|
else:
|
|
211
211
|
return torch.tensor(
|
|
212
212
|
np.array([body.angular_velocity for body in self._bodies]),
|
|
213
|
-
device=self.device
|
|
213
|
+
device=self.device,
|
|
214
214
|
)
|
|
215
215
|
|
|
216
216
|
@property
|
|
@@ -265,7 +265,9 @@ class PhysxRigidBodyComponentStruct(PhysxRigidBaseComponentStruct[T], Generic[T]
|
|
|
265
265
|
# for link entities, namely 7:10 was angular velocity and 10:13 was linear velocity. SAPIEN 3.0.0 and above fixes this
|
|
266
266
|
return self._body_data[self._body_data_index, 7:10]
|
|
267
267
|
else:
|
|
268
|
-
return torch.from_numpy(self._bodies[0].linear_velocity[None, :]).to(
|
|
268
|
+
return torch.from_numpy(self._bodies[0].linear_velocity[None, :]).to(
|
|
269
|
+
self.device
|
|
270
|
+
)
|
|
269
271
|
|
|
270
272
|
@property
|
|
271
273
|
def mass(self) -> torch.Tensor:
|
|
@@ -298,7 +300,7 @@ class PhysxRigidBodyComponentStruct(PhysxRigidBaseComponentStruct[T], Generic[T]
|
|
|
298
300
|
|
|
299
301
|
@dataclass
|
|
300
302
|
class PhysxRigidDynamicComponentStruct(PhysxRigidBodyComponentStruct[T], Generic[T]):
|
|
301
|
-
_bodies:
|
|
303
|
+
_bodies: list[physx.PhysxRigidDynamicComponent]
|
|
302
304
|
|
|
303
305
|
def get_angular_velocity(self) -> torch.Tensor:
|
|
304
306
|
return self.angular_velocity
|
|
@@ -357,7 +359,9 @@ class PhysxRigidDynamicComponentStruct(PhysxRigidBodyComponentStruct[T], Generic
|
|
|
357
359
|
if self.scene.gpu_sim_enabled:
|
|
358
360
|
return self._body_data[self._body_data_index, 10:13]
|
|
359
361
|
else:
|
|
360
|
-
return torch.from_numpy(self._bodies[0].angular_velocity[None, :]).to(
|
|
362
|
+
return torch.from_numpy(self._bodies[0].angular_velocity[None, :]).to(
|
|
363
|
+
self.device
|
|
364
|
+
)
|
|
361
365
|
|
|
362
366
|
@angular_velocity.setter
|
|
363
367
|
def angular_velocity(self, arg1: Array):
|
|
@@ -426,7 +430,8 @@ class PhysxRigidDynamicComponentStruct(PhysxRigidBodyComponentStruct[T], Generic
|
|
|
426
430
|
return self._body_data[self._body_data_index, 7:10]
|
|
427
431
|
else:
|
|
428
432
|
return torch.tensor(
|
|
429
|
-
np.array([body.linear_velocity for body in self._bodies]),
|
|
433
|
+
np.array([body.linear_velocity for body in self._bodies]),
|
|
434
|
+
device=self.device,
|
|
430
435
|
)
|
|
431
436
|
|
|
432
437
|
@linear_velocity.setter
|
|
@@ -2,7 +2,7 @@ from __future__ import annotations
|
|
|
2
2
|
|
|
3
3
|
import typing
|
|
4
4
|
from dataclasses import dataclass
|
|
5
|
-
from typing import TYPE_CHECKING,
|
|
5
|
+
from typing import TYPE_CHECKING, Sequence, Union
|
|
6
6
|
|
|
7
7
|
import sapien
|
|
8
8
|
import sapien.physx as physx
|
|
@@ -38,7 +38,7 @@ class Drive(PhysxJointComponentStruct[physx.PhysxDriveComponent]):
|
|
|
38
38
|
pose1: Union[sapien.Pose, Pose] = None,
|
|
39
39
|
scene_idxs: torch.Tensor = None,
|
|
40
40
|
):
|
|
41
|
-
physx_drives:
|
|
41
|
+
physx_drives: list[physx.PhysxDriveComponent] = []
|
|
42
42
|
assert bodies1 is not None
|
|
43
43
|
if bodies0 is None:
|
|
44
44
|
bodies0 = [None] * len(bodies1)
|
mani_skill/utils/structs/link.py
CHANGED
|
@@ -2,7 +2,7 @@ from __future__ import annotations
|
|
|
2
2
|
|
|
3
3
|
from dataclasses import dataclass, field
|
|
4
4
|
from functools import cached_property
|
|
5
|
-
from typing import TYPE_CHECKING, Callable,
|
|
5
|
+
from typing import TYPE_CHECKING, Callable, Union
|
|
6
6
|
|
|
7
7
|
import sapien
|
|
8
8
|
import sapien.physx as physx
|
|
@@ -38,7 +38,7 @@ class Link(PhysxRigidBodyComponentStruct[physx.PhysxArticulationLinkComponent]):
|
|
|
38
38
|
joint: ArticulationJoint = None
|
|
39
39
|
"""the joint of which this link is a child of. If this is a view/merged link then this joint is also a view/merged joint"""
|
|
40
40
|
|
|
41
|
-
meshes:
|
|
41
|
+
meshes: dict[str, list[trimesh.Trimesh]] = field(default_factory=dict)
|
|
42
42
|
"""
|
|
43
43
|
map from user-defined mesh groups (e.g. "handle" meshes for cabinets) to a list of trimesh.Trimesh objects corresponding to each physx link object managed here
|
|
44
44
|
"""
|
|
@@ -58,7 +58,7 @@ class Link(PhysxRigidBodyComponentStruct[physx.PhysxArticulationLinkComponent]):
|
|
|
58
58
|
@classmethod
|
|
59
59
|
def create(
|
|
60
60
|
cls,
|
|
61
|
-
physx_links:
|
|
61
|
+
physx_links: list[physx.PhysxArticulationLinkComponent],
|
|
62
62
|
scene: ManiSkillScene,
|
|
63
63
|
scene_idxs: torch.Tensor,
|
|
64
64
|
):
|
|
@@ -75,7 +75,7 @@ class Link(PhysxRigidBodyComponentStruct[physx.PhysxArticulationLinkComponent]):
|
|
|
75
75
|
)
|
|
76
76
|
|
|
77
77
|
@classmethod
|
|
78
|
-
def merge(cls, links:
|
|
78
|
+
def merge(cls, links: list["Link"], name: str = None):
|
|
79
79
|
objs = []
|
|
80
80
|
joint_objs = []
|
|
81
81
|
merged_joint_indexes = []
|
|
@@ -126,7 +126,7 @@ class Link(PhysxRigidBodyComponentStruct[physx.PhysxArticulationLinkComponent]):
|
|
|
126
126
|
"""
|
|
127
127
|
Returns each managed link objects render shape list (a list of lists)
|
|
128
128
|
"""
|
|
129
|
-
all_render_shapes:
|
|
129
|
+
all_render_shapes: list[list[sapien.render.RenderShape]] = []
|
|
130
130
|
for obj in self._objs:
|
|
131
131
|
rb_comp = obj.entity.find_component_by_type(
|
|
132
132
|
sapien.render.RenderBodyComponent
|
|
@@ -137,7 +137,7 @@ class Link(PhysxRigidBodyComponentStruct[physx.PhysxArticulationLinkComponent]):
|
|
|
137
137
|
|
|
138
138
|
def get_visual_meshes(
|
|
139
139
|
self, to_world_frame: bool = True, first_only: bool = False
|
|
140
|
-
) ->
|
|
140
|
+
) -> list[trimesh.Trimesh]:
|
|
141
141
|
"""
|
|
142
142
|
Returns the visual mesh of each managed link object. Note results of this are not cached or optimized at the moment
|
|
143
143
|
so this function can be slow if called too often
|
|
@@ -183,7 +183,7 @@ class Link(PhysxRigidBodyComponentStruct[physx.PhysxArticulationLinkComponent]):
|
|
|
183
183
|
filter: Callable[
|
|
184
184
|
[physx.PhysxArticulationLinkComponent, sapien.render.RenderShape], bool
|
|
185
185
|
],
|
|
186
|
-
) ->
|
|
186
|
+
) -> list[trimesh.primitives.Box]:
|
|
187
187
|
# First we need to pre-compute the bounding box of the link at 0. This will be slow the first time
|
|
188
188
|
bboxes = []
|
|
189
189
|
for link, link_render_shapes in zip(self._objs, self.render_shapes):
|
|
@@ -244,7 +244,9 @@ class Link(PhysxRigidBodyComponentStruct[physx.PhysxArticulationLinkComponent]):
|
|
|
244
244
|
raw_pose = new_pose
|
|
245
245
|
return Pose.create(raw_pose)
|
|
246
246
|
else:
|
|
247
|
-
return Pose.create(
|
|
247
|
+
return Pose.create(
|
|
248
|
+
[obj.entity_pose for obj in self._objs], device=self.device
|
|
249
|
+
)
|
|
248
250
|
|
|
249
251
|
@pose.setter
|
|
250
252
|
def pose(self, arg1: Union[Pose, sapien.Pose, Array]) -> None:
|
mani_skill/utils/structs/pose.py
CHANGED
|
@@ -1,5 +1,5 @@
|
|
|
1
1
|
from dataclasses import dataclass
|
|
2
|
-
from typing import
|
|
2
|
+
from typing import Optional, Union
|
|
3
3
|
|
|
4
4
|
import numpy as np
|
|
5
5
|
import sapien
|
|
@@ -21,7 +21,7 @@ def add_batch_dim(x):
|
|
|
21
21
|
return x
|
|
22
22
|
|
|
23
23
|
|
|
24
|
-
def to_batched_tensor(x: Union[
|
|
24
|
+
def to_batched_tensor(x: Union[list, Array], device: Optional[Device] = None):
|
|
25
25
|
if x is None:
|
|
26
26
|
return None
|
|
27
27
|
return add_batch_dim(common.to_tensor(x, device=device))
|
|
@@ -121,7 +121,7 @@ class Pose:
|
|
|
121
121
|
@classmethod
|
|
122
122
|
def create(
|
|
123
123
|
cls,
|
|
124
|
-
pose: Union[torch.Tensor, sapien.Pose,
|
|
124
|
+
pose: Union[torch.Tensor, sapien.Pose, list[sapien.Pose], "Pose"],
|
|
125
125
|
device: Optional[Device] = None,
|
|
126
126
|
) -> "Pose":
|
|
127
127
|
"""Creates a Pose object from a given ``pose``, which can be a torch tensor, sapien.Pose, list of sapien.Pose, or Pose"""
|
|
@@ -1,5 +1,5 @@
|
|
|
1
1
|
from dataclasses import dataclass
|
|
2
|
-
from typing import TYPE_CHECKING, Any,
|
|
2
|
+
from typing import TYPE_CHECKING, Any, Union
|
|
3
3
|
|
|
4
4
|
from mani_skill.utils.geometry.rotation_conversions import quaternion_to_matrix
|
|
5
5
|
|
|
@@ -30,7 +30,7 @@ class RenderCamera:
|
|
|
30
30
|
Wrapper around sapien.render.RenderCameraComponent
|
|
31
31
|
"""
|
|
32
32
|
|
|
33
|
-
_render_cameras:
|
|
33
|
+
_render_cameras: list[sapien.render.RenderCameraComponent]
|
|
34
34
|
name: str
|
|
35
35
|
# NOTE (stao): I cannot seem to use ManiSkillScene as a type here, it complains it is undefined despite using TYPE_CHECKING variable. Without typchecking there is a ciruclar import error
|
|
36
36
|
scene: Any
|
|
@@ -47,7 +47,7 @@ class RenderCamera:
|
|
|
47
47
|
@classmethod
|
|
48
48
|
def create(
|
|
49
49
|
cls,
|
|
50
|
-
render_cameras:
|
|
50
|
+
render_cameras: list[sapien.render.RenderCameraComponent],
|
|
51
51
|
scene: Any,
|
|
52
52
|
mount: Union[Actor, Link] = None,
|
|
53
53
|
):
|
|
@@ -157,7 +157,7 @@ class RenderCamera:
|
|
|
157
157
|
def get_near(self) -> float:
|
|
158
158
|
return self._render_cameras[0].get_near()
|
|
159
159
|
|
|
160
|
-
def get_picture(self, names: Union[str,
|
|
160
|
+
def get_picture(self, names: Union[str, list[str]]) -> list[torch.Tensor]:
|
|
161
161
|
if isinstance(names, str):
|
|
162
162
|
names = [names]
|
|
163
163
|
if self.scene.gpu_sim_enabled and not self.scene.parallel_in_single_scene:
|
|
@@ -36,7 +36,9 @@ class GPUMemoryConfig:
|
|
|
36
36
|
|
|
37
37
|
@dataclass
|
|
38
38
|
class SceneConfig:
|
|
39
|
-
gravity: np.ndarray = field(
|
|
39
|
+
gravity: Union[np.ndarray, list[float]] = field(
|
|
40
|
+
default_factory=lambda: np.array([0, 0, -9.81])
|
|
41
|
+
)
|
|
40
42
|
bounce_threshold: float = 2.0
|
|
41
43
|
sleep_threshold: float = 0.005
|
|
42
44
|
contact_offset: float = 0.02
|
|
@@ -1,12 +1,10 @@
|
|
|
1
|
-
from typing import List
|
|
2
|
-
|
|
3
1
|
import matplotlib.pyplot as plt
|
|
4
2
|
import numpy as np
|
|
5
3
|
from IPython.display import HTML, display
|
|
6
4
|
from matplotlib import animation
|
|
7
5
|
|
|
8
6
|
|
|
9
|
-
def display_images(images:
|
|
7
|
+
def display_images(images: list[np.ndarray], dpi=100.0, format="html5_video", **kwargs):
|
|
10
8
|
"""Display images as an animation in jupyter notebook.
|
|
11
9
|
|
|
12
10
|
Args:
|
|
@@ -1,5 +1,5 @@
|
|
|
1
1
|
import os
|
|
2
|
-
from typing import
|
|
2
|
+
from typing import Optional
|
|
3
3
|
|
|
4
4
|
import imageio
|
|
5
5
|
import numpy as np
|
|
@@ -11,7 +11,7 @@ from mani_skill.utils.structs.types import Array
|
|
|
11
11
|
|
|
12
12
|
|
|
13
13
|
def images_to_video(
|
|
14
|
-
images:
|
|
14
|
+
images: list[Array],
|
|
15
15
|
output_dir: str,
|
|
16
16
|
video_name: str,
|
|
17
17
|
fps: int = 10,
|
|
@@ -51,7 +51,7 @@ def images_to_video(
|
|
|
51
51
|
writer.close()
|
|
52
52
|
|
|
53
53
|
|
|
54
|
-
def tile_images(images:
|
|
54
|
+
def tile_images(images: list[Array], nrows=1) -> Array:
|
|
55
55
|
"""
|
|
56
56
|
Tile multiple images to a single image comprised of nrows and an appropriate number of columns to fit all the images.
|
|
57
57
|
The images can also be batched (e.g. of shape (B, H, W, C)), but give images must all have the same batch size.
|
|
@@ -118,7 +118,7 @@ def tile_images(images: List[Array], nrows=1) -> Array:
|
|
|
118
118
|
TEXT_FONT = None
|
|
119
119
|
|
|
120
120
|
|
|
121
|
-
def put_text_on_image(image: np.ndarray, lines:
|
|
121
|
+
def put_text_on_image(image: np.ndarray, lines: list[str]):
|
|
122
122
|
global TEXT_FONT
|
|
123
123
|
assert image.dtype == np.uint8, image.dtype
|
|
124
124
|
image = image.copy()
|
|
@@ -139,7 +139,7 @@ def put_text_on_image(image: np.ndarray, lines: List[str]):
|
|
|
139
139
|
return np.array(image)
|
|
140
140
|
|
|
141
141
|
|
|
142
|
-
def put_info_on_image(image, info:
|
|
142
|
+
def put_info_on_image(image, info: dict[str, float], extras=None, overlay=True):
|
|
143
143
|
lines = [
|
|
144
144
|
f"{k}: {v:.3f}" if isinstance(v, float) else f"{k}: {v}"
|
|
145
145
|
for k, v in info.items()
|
|
@@ -1,5 +1,5 @@
|
|
|
1
1
|
from dataclasses import asdict, dataclass
|
|
2
|
-
from typing import
|
|
2
|
+
from typing import Optional, Union
|
|
3
3
|
|
|
4
4
|
import dacite
|
|
5
5
|
import gymnasium as gym
|
|
@@ -66,7 +66,9 @@ class CachedResetWrapper(gym.Wrapper):
|
|
|
66
66
|
self._cached_resets_obs_buffer = []
|
|
67
67
|
while self._num_cached_resets < self.cached_resets_config.num_resets:
|
|
68
68
|
obs, _ = self.env.reset(
|
|
69
|
-
seed=self.cached_resets_config.seed
|
|
69
|
+
seed=self.cached_resets_config.seed
|
|
70
|
+
if self._num_cached_resets == 0
|
|
71
|
+
else None,
|
|
70
72
|
options=dict(
|
|
71
73
|
env_idx=torch.arange(
|
|
72
74
|
0,
|
|
@@ -125,7 +127,7 @@ class CachedResetWrapper(gym.Wrapper):
|
|
|
125
127
|
def reset(
|
|
126
128
|
self,
|
|
127
129
|
*args,
|
|
128
|
-
seed: Optional[Union[int,
|
|
130
|
+
seed: Optional[Union[int, list[int]]] = None,
|
|
129
131
|
options: Optional[dict] = None,
|
|
130
132
|
**kwargs
|
|
131
133
|
):
|
|
@@ -1,5 +1,4 @@
|
|
|
1
1
|
import copy
|
|
2
|
-
from typing import Dict
|
|
3
2
|
|
|
4
3
|
import gymnasium as gym
|
|
5
4
|
import gymnasium.spaces.utils
|
|
@@ -42,7 +41,7 @@ class FlattenRGBDObservationWrapper(gym.ObservationWrapper):
|
|
|
42
41
|
new_obs = self.observation(self.base_env._init_raw_obs)
|
|
43
42
|
self.base_env.update_obs_space(new_obs)
|
|
44
43
|
|
|
45
|
-
def observation(self, observation:
|
|
44
|
+
def observation(self, observation: dict):
|
|
46
45
|
sensor_data = observation.pop("sensor_data")
|
|
47
46
|
del observation["sensor_param"]
|
|
48
47
|
rgb_images = []
|
|
@@ -2,7 +2,7 @@ import copy
|
|
|
2
2
|
import time
|
|
3
3
|
from dataclasses import dataclass
|
|
4
4
|
from pathlib import Path
|
|
5
|
-
from typing import Callable,
|
|
5
|
+
from typing import Callable, Optional, Union
|
|
6
6
|
|
|
7
7
|
import gymnasium as gym
|
|
8
8
|
import h5py
|
|
@@ -119,21 +119,21 @@ class RecordEpisode(gym.Wrapper):
|
|
|
119
119
|
|
|
120
120
|
Each JSON file contains:
|
|
121
121
|
|
|
122
|
-
- `env_info` (
|
|
122
|
+
- `env_info` (dict): task (also known as environment) information, which can be used to initialize the task
|
|
123
123
|
- `env_id` (str): task id
|
|
124
124
|
- `max_episode_steps` (int)
|
|
125
|
-
- `env_kwargs` (
|
|
126
|
-
- `episodes` (
|
|
125
|
+
- `env_kwargs` (dict): keyword arguments to initialize the task. **Essential to recreate the environment.**
|
|
126
|
+
- `episodes` (list[dict]): episode information
|
|
127
127
|
- `source_type` (Optional[str]): a simple category string describing what process generated the trajectory data. ManiSkill official datasets will usually write one of "human", "motionplanning", or "rl" at the moment.
|
|
128
128
|
- `source_desc` (Optional[str]): a longer explanation of how the data was generated.
|
|
129
129
|
|
|
130
130
|
The episode information (the element of `episodes`) includes:
|
|
131
131
|
|
|
132
132
|
- `episode_id` (int): a unique id to index the episode
|
|
133
|
-
- `reset_kwargs` (
|
|
133
|
+
- `reset_kwargs` (dict): keyword arguments to reset the task. **Essential to reproduce the trajectory.**
|
|
134
134
|
- `control_mode` (str): control mode used for the episode.
|
|
135
135
|
- `elapsed_steps` (int): trajectory length
|
|
136
|
-
- `info` (
|
|
136
|
+
- `info` (dict): information at the end of the episode.
|
|
137
137
|
|
|
138
138
|
With just the meta data, you can reproduce the task the same way it was created when the trajectories were collected as so:
|
|
139
139
|
|
|
@@ -356,7 +356,7 @@ class RecordEpisode(gym.Wrapper):
|
|
|
356
356
|
def reset(
|
|
357
357
|
self,
|
|
358
358
|
*args,
|
|
359
|
-
seed: Optional[Union[int,
|
|
359
|
+
seed: Optional[Union[int, list[int]]] = None,
|
|
360
360
|
options: Optional[dict] = None,
|
|
361
361
|
save=True,
|
|
362
362
|
**kwargs,
|
|
@@ -367,7 +367,9 @@ class RecordEpisode(gym.Wrapper):
|
|
|
367
367
|
# if doing a full reset then we flush all trajectories including incompleted ones
|
|
368
368
|
if self._trajectory_buffer is not None:
|
|
369
369
|
if options is None or "env_idx" not in options:
|
|
370
|
-
self.flush_trajectory(
|
|
370
|
+
self.flush_trajectory(
|
|
371
|
+
env_idxs_to_flush=np.arange(self.num_envs), save=save
|
|
372
|
+
)
|
|
371
373
|
else:
|
|
372
374
|
self.flush_trajectory(
|
|
373
375
|
env_idxs_to_flush=common.to_numpy(options["env_idx"], save=save)
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
from typing import
|
|
1
|
+
from typing import Literal
|
|
2
2
|
|
|
3
3
|
import gymnasium as gym
|
|
4
4
|
import torch
|
|
@@ -38,7 +38,7 @@ class VisualEncoderWrapper(gym.ObservationWrapper):
|
|
|
38
38
|
super().__init__(env)
|
|
39
39
|
|
|
40
40
|
@torch.no_grad()
|
|
41
|
-
def observation(self, obs:
|
|
41
|
+
def observation(self, obs: dict):
|
|
42
42
|
vec_img_embeddings_list = []
|
|
43
43
|
image_obs = obs.pop("sensor_data")
|
|
44
44
|
del obs["sensor_param"]
|
|
@@ -1,6 +1,6 @@
|
|
|
1
1
|
from __future__ import annotations
|
|
2
2
|
|
|
3
|
-
from typing import TYPE_CHECKING,
|
|
3
|
+
from typing import TYPE_CHECKING, Optional, Tuple, Union
|
|
4
4
|
|
|
5
5
|
import gymnasium as gym
|
|
6
6
|
import torch
|
|
@@ -36,7 +36,7 @@ class ManiSkillVectorEnv(VectorEnv):
|
|
|
36
36
|
def __init__(
|
|
37
37
|
self,
|
|
38
38
|
env: Union[Env, str],
|
|
39
|
-
num_envs: int =
|
|
39
|
+
num_envs: int = 1,
|
|
40
40
|
auto_reset: bool = True,
|
|
41
41
|
ignore_terminations: bool = False,
|
|
42
42
|
record_metrics: bool = False,
|
|
@@ -78,7 +78,7 @@ class ManiSkillVectorEnv(VectorEnv):
|
|
|
78
78
|
|
|
79
79
|
@property
|
|
80
80
|
def base_env(self) -> BaseEnv:
|
|
81
|
-
return self._env.unwrapped
|
|
81
|
+
return self._env.unwrapped # type: ignore
|
|
82
82
|
|
|
83
83
|
@property
|
|
84
84
|
def unwrapped(self):
|
|
@@ -87,13 +87,15 @@ class ManiSkillVectorEnv(VectorEnv):
|
|
|
87
87
|
def reset(
|
|
88
88
|
self,
|
|
89
89
|
*,
|
|
90
|
-
seed: Optional[Union[int,
|
|
90
|
+
seed: Optional[Union[int, list[int]]] = None,
|
|
91
91
|
options: Optional[dict] = None,
|
|
92
92
|
):
|
|
93
|
-
obs, info = self._env.reset(seed=seed, options=options)
|
|
93
|
+
obs, info = self._env.reset(seed=seed, options=options) # type: ignore
|
|
94
94
|
if options is not None and "env_idx" in options:
|
|
95
95
|
env_idx = options["env_idx"]
|
|
96
|
-
mask = torch.zeros(
|
|
96
|
+
mask = torch.zeros(
|
|
97
|
+
self.num_envs, dtype=torch.bool, device=self.base_env.device
|
|
98
|
+
)
|
|
97
99
|
mask[env_idx] = True
|
|
98
100
|
if self.record_metrics:
|
|
99
101
|
self.success_once[mask] = False
|
|
@@ -106,11 +108,11 @@ class ManiSkillVectorEnv(VectorEnv):
|
|
|
106
108
|
self.returns[:] = 0
|
|
107
109
|
return obs, info
|
|
108
110
|
|
|
109
|
-
def step(
|
|
110
|
-
self, actions: Union[Array,
|
|
111
|
-
) -> Tuple[Array, Array, Array, Array,
|
|
111
|
+
def step( # pyright: ignore[reportIncompatibleMethodOverride]
|
|
112
|
+
self, actions: Union[Array, dict]
|
|
113
|
+
) -> Tuple[Array, Array, Array, Array, dict]:
|
|
112
114
|
obs, rew, terminations, truncations, infos = self._env.step(actions)
|
|
113
|
-
|
|
115
|
+
episode_info: Optional[dict] = None
|
|
114
116
|
if self.record_metrics:
|
|
115
117
|
episode_info = dict()
|
|
116
118
|
self.returns += rew
|
|
@@ -131,7 +133,7 @@ class ManiSkillVectorEnv(VectorEnv):
|
|
|
131
133
|
|
|
132
134
|
if self.ignore_terminations:
|
|
133
135
|
terminations[:] = False
|
|
134
|
-
if
|
|
136
|
+
if episode_info:
|
|
135
137
|
if "success" in infos:
|
|
136
138
|
episode_info["success_at_end"] = infos["success"].clone()
|
|
137
139
|
if "fail" in infos:
|
|
@@ -139,7 +141,9 @@ class ManiSkillVectorEnv(VectorEnv):
|
|
|
139
141
|
if self.record_metrics:
|
|
140
142
|
infos["episode"] = episode_info
|
|
141
143
|
|
|
142
|
-
dones = torch.logical_or(
|
|
144
|
+
dones = torch.logical_or(
|
|
145
|
+
terminations, truncations # pyright: ignore[reportArgumentType]
|
|
146
|
+
)
|
|
143
147
|
|
|
144
148
|
if dones.any() and self.auto_reset:
|
|
145
149
|
final_obs = torch_clone_dict(obs)
|
|
@@ -155,7 +159,13 @@ class ManiSkillVectorEnv(VectorEnv):
|
|
|
155
159
|
infos["_final_observation"] = dones
|
|
156
160
|
infos["_elapsed_steps"] = dones
|
|
157
161
|
# NOTE (stao): Unlike gymnasium, the code here does not add masks for every key in the info object.
|
|
158
|
-
return
|
|
162
|
+
return (
|
|
163
|
+
obs,
|
|
164
|
+
rew,
|
|
165
|
+
terminations,
|
|
166
|
+
truncations,
|
|
167
|
+
infos,
|
|
168
|
+
) # pyright: ignore[reportReturnType]
|
|
159
169
|
|
|
160
170
|
def close_extras(self, **kwargs):
|
|
161
171
|
self._env.close()
|
|
@@ -1,5 +1,5 @@
|
|
|
1
1
|
import time
|
|
2
|
-
from typing import Any,
|
|
2
|
+
from typing import Any, Optional, Type, Union
|
|
3
3
|
|
|
4
4
|
import gymnasium as gym
|
|
5
5
|
import numpy as np
|
|
@@ -52,7 +52,7 @@ class ManiSkillSB3VectorEnv(SB3VecEnv):
|
|
|
52
52
|
def base_env(self) -> BaseEnv:
|
|
53
53
|
return self._env.unwrapped
|
|
54
54
|
|
|
55
|
-
def seed(self, seed: Optional[int] = None) ->
|
|
55
|
+
def seed(self, seed: Optional[int] = None) -> list[Union[None, int]]:
|
|
56
56
|
self._last_seed = seed
|
|
57
57
|
|
|
58
58
|
def reset(self) -> VecEnvObs:
|
|
@@ -113,7 +113,7 @@ class ManiSkillSB3VectorEnv(SB3VecEnv):
|
|
|
113
113
|
def close(self) -> None:
|
|
114
114
|
return self._env.close()
|
|
115
115
|
|
|
116
|
-
def get_attr(self, attr_name: str, indices: VecEnvIndices = None) ->
|
|
116
|
+
def get_attr(self, attr_name: str, indices: VecEnvIndices = None) -> list[Any]:
|
|
117
117
|
return self._env.get_attr(attr_name, indices)
|
|
118
118
|
|
|
119
119
|
def set_attr(
|
|
@@ -127,12 +127,12 @@ class ManiSkillSB3VectorEnv(SB3VecEnv):
|
|
|
127
127
|
*method_args,
|
|
128
128
|
indices: VecEnvIndices = None,
|
|
129
129
|
**method_kwargs
|
|
130
|
-
) ->
|
|
130
|
+
) -> list[Any]:
|
|
131
131
|
return self._env.env_method(
|
|
132
132
|
method_name, *method_args, indices=indices, **method_kwargs
|
|
133
133
|
)
|
|
134
134
|
|
|
135
135
|
def env_is_wrapped(
|
|
136
136
|
self, wrapper_class: Type[gym.Wrapper], indices: VecEnvIndices = None
|
|
137
|
-
) ->
|
|
137
|
+
) -> list[bool]:
|
|
138
138
|
return [False] * self.num_envs
|