mani-skill-nightly 2025.10.22.143__py3-none-any.whl → 2025.10.22.325__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 +32 -32
- 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/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/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 +11 -9
- 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.22.143.dist-info → mani_skill_nightly-2025.10.22.325.dist-info}/METADATA +1 -1
- {mani_skill_nightly-2025.10.22.143.dist-info → mani_skill_nightly-2025.10.22.325.dist-info}/RECORD +122 -122
- {mani_skill_nightly-2025.10.22.143.dist-info → mani_skill_nightly-2025.10.22.325.dist-info}/WHEEL +0 -0
- {mani_skill_nightly-2025.10.22.143.dist-info → mani_skill_nightly-2025.10.22.325.dist-info}/licenses/LICENSE +0 -0
- {mani_skill_nightly-2025.10.22.143.dist-info → mani_skill_nightly-2025.10.22.325.dist-info}/licenses/LICENSE-3RD-PARTY +0 -0
- {mani_skill_nightly-2025.10.22.143.dist-info → mani_skill_nightly-2025.10.22.325.dist-info}/top_level.txt +0 -0
|
@@ -15,7 +15,7 @@ in addition to initializing any task relevant data like a goal
|
|
|
15
15
|
See comments for how to make your own environment and what each required function should do
|
|
16
16
|
"""
|
|
17
17
|
|
|
18
|
-
from typing import Any,
|
|
18
|
+
from typing import Any, Union
|
|
19
19
|
|
|
20
20
|
import numpy as np
|
|
21
21
|
import sapien
|
|
@@ -191,7 +191,7 @@ class PushCubeEnv(BaseEnv):
|
|
|
191
191
|
"success": is_obj_placed,
|
|
192
192
|
}
|
|
193
193
|
|
|
194
|
-
def _get_obs_extra(self, info:
|
|
194
|
+
def _get_obs_extra(self, info: dict):
|
|
195
195
|
# some useful observation info for solving the task includes the pose of the tcp (tool center point) which is the point between the
|
|
196
196
|
# grippers of the robot
|
|
197
197
|
obs = dict(
|
|
@@ -206,7 +206,7 @@ class PushCubeEnv(BaseEnv):
|
|
|
206
206
|
)
|
|
207
207
|
return obs
|
|
208
208
|
|
|
209
|
-
def compute_dense_reward(self, obs: Any, action: Array, info:
|
|
209
|
+
def compute_dense_reward(self, obs: Any, action: Array, info: dict):
|
|
210
210
|
# We also create a pose marking where the robot should push the cube from that is easiest (pushing from behind the cube)
|
|
211
211
|
tcp_push_pose = Pose.create_from_pq(
|
|
212
212
|
p=self.obj.pose.p
|
|
@@ -226,13 +226,13 @@ class PushCubeEnv(BaseEnv):
|
|
|
226
226
|
)
|
|
227
227
|
place_reward = 1 - torch.tanh(5 * obj_to_goal_dist)
|
|
228
228
|
reward += place_reward * reached
|
|
229
|
-
|
|
229
|
+
|
|
230
230
|
# Compute a z reward to encourage the robot to keep the cube on the table
|
|
231
231
|
desired_obj_z = self.cube_half_size
|
|
232
232
|
current_obj_z = self.obj.pose.p[..., 2]
|
|
233
233
|
z_deviation = torch.abs(current_obj_z - desired_obj_z)
|
|
234
234
|
z_reward = 1 - torch.tanh(5 * z_deviation)
|
|
235
|
-
# We multiply the z reward by the place_reward and reached mask so that
|
|
235
|
+
# We multiply the z reward by the place_reward and reached mask so that
|
|
236
236
|
# we only add the z reward if the robot has reached the desired push pose
|
|
237
237
|
# and the z reward becomes more important as the robot gets closer to the goal.
|
|
238
238
|
reward += place_reward * z_reward * reached
|
|
@@ -241,7 +241,7 @@ class PushCubeEnv(BaseEnv):
|
|
|
241
241
|
reward[info["success"]] = 4
|
|
242
242
|
return reward
|
|
243
243
|
|
|
244
|
-
def compute_normalized_dense_reward(self, obs: Any, action: Array, info:
|
|
244
|
+
def compute_normalized_dense_reward(self, obs: Any, action: Array, info: dict):
|
|
245
245
|
# this should be equal to compute_dense_reward / max possible reward
|
|
246
246
|
max_reward = 4.0
|
|
247
247
|
return self.compute_dense_reward(obs=obs, action=action, info=info) / max_reward
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
from typing import Any
|
|
1
|
+
from typing import Any
|
|
2
2
|
|
|
3
3
|
import numpy as np
|
|
4
4
|
import sapien
|
|
@@ -489,7 +489,7 @@ class PushTEnv(BaseEnv):
|
|
|
489
489
|
|
|
490
490
|
return {"success": success}
|
|
491
491
|
|
|
492
|
-
def _get_obs_extra(self, info:
|
|
492
|
+
def _get_obs_extra(self, info: dict):
|
|
493
493
|
# ee position is super useful for pandastick robot
|
|
494
494
|
obs = dict(
|
|
495
495
|
tcp_pose=self.agent.tcp.pose.raw_pose,
|
|
@@ -502,7 +502,7 @@ class PushTEnv(BaseEnv):
|
|
|
502
502
|
)
|
|
503
503
|
return obs
|
|
504
504
|
|
|
505
|
-
def compute_dense_reward(self, obs: Any, action: Array, info:
|
|
505
|
+
def compute_dense_reward(self, obs: Any, action: Array, info: dict):
|
|
506
506
|
# reward for overlap of the tees
|
|
507
507
|
|
|
508
508
|
# legacy reward
|
|
@@ -533,6 +533,6 @@ class PushTEnv(BaseEnv):
|
|
|
533
533
|
reward[info["success"]] = 3
|
|
534
534
|
return reward
|
|
535
535
|
|
|
536
|
-
def compute_normalized_dense_reward(self, obs: Any, action: Array, info:
|
|
536
|
+
def compute_normalized_dense_reward(self, obs: Any, action: Array, info: dict):
|
|
537
537
|
max_reward = 3.0
|
|
538
538
|
return self.compute_dense_reward(obs=obs, action=action, info=info) / max_reward
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
from typing import Any
|
|
1
|
+
from typing import Any
|
|
2
2
|
|
|
3
3
|
import numpy as np
|
|
4
4
|
import sapien
|
|
@@ -135,7 +135,7 @@ class RollBallEnv(BaseEnv):
|
|
|
135
135
|
"success": is_obj_placed,
|
|
136
136
|
}
|
|
137
137
|
|
|
138
|
-
def _get_obs_extra(self, info:
|
|
138
|
+
def _get_obs_extra(self, info: dict):
|
|
139
139
|
|
|
140
140
|
obs = dict(
|
|
141
141
|
tcp_pose=self.agent.tcp.pose.raw_pose,
|
|
@@ -150,7 +150,7 @@ class RollBallEnv(BaseEnv):
|
|
|
150
150
|
)
|
|
151
151
|
return obs
|
|
152
152
|
|
|
153
|
-
def compute_dense_reward(self, obs: Any, action: Array, info:
|
|
153
|
+
def compute_dense_reward(self, obs: Any, action: Array, info: dict):
|
|
154
154
|
unit_vec = self.ball.pose.p - self.goal_region.pose.p
|
|
155
155
|
unit_vec = unit_vec / torch.linalg.norm(unit_vec, axis=1, keepdim=True)
|
|
156
156
|
tcp_hit_pose = Pose.create_from_pq(
|
|
@@ -176,6 +176,6 @@ class RollBallEnv(BaseEnv):
|
|
|
176
176
|
reward[info["success"]] = 30.0
|
|
177
177
|
return reward
|
|
178
178
|
|
|
179
|
-
def compute_normalized_dense_reward(self, obs: Any, action: Array, info:
|
|
179
|
+
def compute_normalized_dense_reward(self, obs: Any, action: Array, info: dict):
|
|
180
180
|
max_reward = 30.0
|
|
181
181
|
return self.compute_dense_reward(obs=obs, action=action, info=info) / max_reward
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
from typing import Any,
|
|
1
|
+
from typing import Any, Union
|
|
2
2
|
|
|
3
3
|
import numpy as np
|
|
4
4
|
import sapien
|
|
@@ -130,7 +130,7 @@ class StackCubeEnv(BaseEnv):
|
|
|
130
130
|
"success": success.bool(),
|
|
131
131
|
}
|
|
132
132
|
|
|
133
|
-
def _get_obs_extra(self, info:
|
|
133
|
+
def _get_obs_extra(self, info: dict):
|
|
134
134
|
obs = dict(tcp_pose=self.agent.tcp.pose.raw_pose)
|
|
135
135
|
if "state" in self.obs_mode:
|
|
136
136
|
obs.update(
|
|
@@ -142,7 +142,7 @@ class StackCubeEnv(BaseEnv):
|
|
|
142
142
|
)
|
|
143
143
|
return obs
|
|
144
144
|
|
|
145
|
-
def compute_dense_reward(self, obs: Any, action: torch.Tensor, info:
|
|
145
|
+
def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: dict):
|
|
146
146
|
# reaching reward
|
|
147
147
|
tcp_pose = self.agent.tcp.pose.p
|
|
148
148
|
cubeA_pos = self.cubeA.pose.p
|
|
@@ -181,6 +181,6 @@ class StackCubeEnv(BaseEnv):
|
|
|
181
181
|
return reward
|
|
182
182
|
|
|
183
183
|
def compute_normalized_dense_reward(
|
|
184
|
-
self, obs: Any, action: torch.Tensor, info:
|
|
184
|
+
self, obs: Any, action: torch.Tensor, info: dict
|
|
185
185
|
):
|
|
186
186
|
return self.compute_dense_reward(obs=obs, action=action, info=info) / 8
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
from typing import
|
|
1
|
+
from typing import Union
|
|
2
2
|
|
|
3
3
|
import numpy as np
|
|
4
4
|
import sapien
|
|
@@ -7,17 +7,18 @@ import torch
|
|
|
7
7
|
from mani_skill.agents.robots import Fetch, Panda
|
|
8
8
|
from mani_skill.envs.sapien_env import BaseEnv
|
|
9
9
|
from mani_skill.envs.utils import randomization
|
|
10
|
+
from mani_skill.sensors.camera import CameraConfig
|
|
11
|
+
from mani_skill.utils import common, sapien_utils
|
|
12
|
+
from mani_skill.utils.building import actors
|
|
10
13
|
from mani_skill.utils.geometry.rotation_conversions import (
|
|
11
14
|
euler_angles_to_matrix,
|
|
12
15
|
matrix_to_quaternion,
|
|
13
16
|
)
|
|
14
|
-
from mani_skill.
|
|
15
|
-
from mani_skill.utils import common, sapien_utils
|
|
16
|
-
from mani_skill.utils.building import actors
|
|
17
|
+
from mani_skill.utils.logging_utils import logger
|
|
17
18
|
from mani_skill.utils.registration import register_env
|
|
18
19
|
from mani_skill.utils.scene_builder.table import TableSceneBuilder
|
|
19
|
-
from mani_skill.utils.structs.pose import Pose
|
|
20
|
-
|
|
20
|
+
from mani_skill.utils.structs.pose import Pose
|
|
21
|
+
|
|
21
22
|
|
|
22
23
|
@register_env("StackPyramid-v1", max_episode_steps=250)
|
|
23
24
|
class StackPyramidEnv(BaseEnv):
|
|
@@ -40,7 +41,7 @@ class StackPyramidEnv(BaseEnv):
|
|
|
40
41
|
|
|
41
42
|
SUPPORTED_ROBOTS = ["panda_wristcam", "panda", "fetch"]
|
|
42
43
|
SUPPORTED_REWARD_MODES = ["none", "sparse"]
|
|
43
|
-
|
|
44
|
+
|
|
44
45
|
agent: Union[Panda, Fetch]
|
|
45
46
|
|
|
46
47
|
def __init__(
|
|
@@ -59,7 +60,6 @@ class StackPyramidEnv(BaseEnv):
|
|
|
59
60
|
pose = sapien_utils.look_at([0.6, 0.7, 0.6], [0.0, 0.0, 0.35])
|
|
60
61
|
return CameraConfig("render_camera", pose, 512, 512, 1, 0.01, 100)
|
|
61
62
|
|
|
62
|
-
|
|
63
63
|
def _load_scene(self, options: dict):
|
|
64
64
|
self.cube_half_size = common.to_tensor([0.02] * 3)
|
|
65
65
|
self.table_scene = TableSceneBuilder(
|
|
@@ -67,13 +67,25 @@ class StackPyramidEnv(BaseEnv):
|
|
|
67
67
|
)
|
|
68
68
|
self.table_scene.build()
|
|
69
69
|
self.cubeA = actors.build_cube(
|
|
70
|
-
self.scene,
|
|
70
|
+
self.scene,
|
|
71
|
+
half_size=0.02,
|
|
72
|
+
color=[1, 0, 0, 1],
|
|
73
|
+
name="cubeA",
|
|
74
|
+
initial_pose=sapien.Pose(p=[0, 0, 0.2]),
|
|
71
75
|
)
|
|
72
76
|
self.cubeB = actors.build_cube(
|
|
73
|
-
self.scene,
|
|
77
|
+
self.scene,
|
|
78
|
+
half_size=0.02,
|
|
79
|
+
color=[0, 1, 0, 1],
|
|
80
|
+
name="cubeB",
|
|
81
|
+
initial_pose=sapien.Pose(p=[1, 0, 0.2]),
|
|
74
82
|
)
|
|
75
83
|
self.cubeC = actors.build_cube(
|
|
76
|
-
self.scene,
|
|
84
|
+
self.scene,
|
|
85
|
+
half_size=0.02,
|
|
86
|
+
color=[0, 0, 1, 1],
|
|
87
|
+
name="cubeC",
|
|
88
|
+
initial_pose=sapien.Pose(p=[-1, 0, 0.2]),
|
|
77
89
|
)
|
|
78
90
|
|
|
79
91
|
def _initialize_episode(self, env_idx: torch.Tensor, options: dict):
|
|
@@ -85,7 +97,9 @@ class StackPyramidEnv(BaseEnv):
|
|
|
85
97
|
xyz[:, 2] = 0.02
|
|
86
98
|
xy = xyz[:, :2]
|
|
87
99
|
region = [[-0.1, -0.2], [0.1, 0.2]]
|
|
88
|
-
sampler = randomization.UniformPlacementSampler(
|
|
100
|
+
sampler = randomization.UniformPlacementSampler(
|
|
101
|
+
bounds=region, batch_size=b, device=self.device
|
|
102
|
+
)
|
|
89
103
|
radius = torch.linalg.norm(torch.tensor([0.02, 0.02]))
|
|
90
104
|
cubeA_xy = xy + sampler.sample(radius, 100)
|
|
91
105
|
cubeB_xy = xy + sampler.sample(radius, 100, verbose=False)
|
|
@@ -93,7 +107,7 @@ class StackPyramidEnv(BaseEnv):
|
|
|
93
107
|
|
|
94
108
|
# Cube A
|
|
95
109
|
xyz[:, :2] = cubeA_xy
|
|
96
|
-
|
|
110
|
+
|
|
97
111
|
qs = randomization.random_quaternions(
|
|
98
112
|
b,
|
|
99
113
|
lock_x=True,
|
|
@@ -112,7 +126,7 @@ class StackPyramidEnv(BaseEnv):
|
|
|
112
126
|
lock_z=False,
|
|
113
127
|
)
|
|
114
128
|
self.cubeB.set_pose(Pose.create_from_pq(p=xyz.clone(), q=qs))
|
|
115
|
-
|
|
129
|
+
|
|
116
130
|
# Cube C
|
|
117
131
|
xyz[:, :2] = cubeC_xy
|
|
118
132
|
qs = randomization.random_quaternions(
|
|
@@ -133,33 +147,39 @@ class StackPyramidEnv(BaseEnv):
|
|
|
133
147
|
offset_AC = pos_A - pos_C
|
|
134
148
|
|
|
135
149
|
def evaluate_cube_distance(offset, cube_a, cube_b, top_or_next):
|
|
136
|
-
xy_flag = (
|
|
137
|
-
|
|
138
|
-
|
|
139
|
-
|
|
150
|
+
xy_flag = (
|
|
151
|
+
torch.linalg.norm(offset[..., :2], axis=1)
|
|
152
|
+
<= torch.linalg.norm(2 * self.cube_half_size[:2]) + 0.005
|
|
153
|
+
)
|
|
140
154
|
z_flag = torch.abs(offset[..., 2]) > 0.02
|
|
141
155
|
if top_or_next == "top":
|
|
142
156
|
is_cubeA_on_cubeB = torch.logical_and(xy_flag, z_flag)
|
|
143
157
|
elif top_or_next == "next_to":
|
|
144
158
|
is_cubeA_on_cubeB = xy_flag
|
|
145
159
|
else:
|
|
146
|
-
return NotImplementedError(
|
|
147
|
-
|
|
160
|
+
return NotImplementedError(
|
|
161
|
+
f"Expect top_or_next to be either 'top' or 'next_to', got {top_or_next}"
|
|
162
|
+
)
|
|
163
|
+
|
|
148
164
|
is_cubeA_static = cube_a.is_static(lin_thresh=1e-2, ang_thresh=0.5)
|
|
149
165
|
is_cubeA_grasped = self.agent.is_grasping(cube_a)
|
|
150
166
|
|
|
151
|
-
success = is_cubeA_on_cubeB & is_cubeA_static & (~is_cubeA_grasped)
|
|
167
|
+
success = is_cubeA_on_cubeB & is_cubeA_static & (~is_cubeA_grasped)
|
|
152
168
|
return success.bool()
|
|
153
169
|
|
|
154
|
-
success_A_B = evaluate_cube_distance(
|
|
170
|
+
success_A_B = evaluate_cube_distance(
|
|
171
|
+
offset_AB, self.cubeA, self.cubeB, "next_to"
|
|
172
|
+
)
|
|
155
173
|
success_C_B = evaluate_cube_distance(offset_BC, self.cubeC, self.cubeB, "top")
|
|
156
174
|
success_C_A = evaluate_cube_distance(offset_AC, self.cubeC, self.cubeA, "top")
|
|
157
|
-
success = torch.logical_and(
|
|
175
|
+
success = torch.logical_and(
|
|
176
|
+
success_A_B, torch.logical_and(success_C_B, success_C_A)
|
|
177
|
+
)
|
|
158
178
|
return {
|
|
159
179
|
"success": success,
|
|
160
180
|
}
|
|
161
181
|
|
|
162
|
-
def _get_obs_extra(self, info:
|
|
182
|
+
def _get_obs_extra(self, info: dict):
|
|
163
183
|
obs = dict(tcp_pose=self.agent.tcp.pose.raw_pose)
|
|
164
184
|
if "state" in self.obs_mode:
|
|
165
185
|
obs.update(
|
|
@@ -174,4 +194,3 @@ class StackPyramidEnv(BaseEnv):
|
|
|
174
194
|
cubeA_to_cubeC_pos=self.cubeC.pose.p - self.cubeA.pose.p,
|
|
175
195
|
)
|
|
176
196
|
return obs
|
|
177
|
-
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
from typing import
|
|
1
|
+
from typing import Union
|
|
2
2
|
|
|
3
3
|
import numpy as np
|
|
4
4
|
import sapien
|
|
@@ -85,7 +85,7 @@ class TurnFaucetEnv(BaseEnv):
|
|
|
85
85
|
switch_link_ids = self._batched_episode_rng.randint(0, 2**31)
|
|
86
86
|
|
|
87
87
|
self._faucets = []
|
|
88
|
-
self._target_switch_links:
|
|
88
|
+
self._target_switch_links: list[Link] = []
|
|
89
89
|
self.model_offsets = []
|
|
90
90
|
for i, model_id in enumerate(model_ids):
|
|
91
91
|
# partnet-mobility is a dataset source and the ids are the ones we sampled
|
|
@@ -188,7 +188,7 @@ class TurnFaucetEnv(BaseEnv):
|
|
|
188
188
|
angle_dist = self.target_angle - self.current_angle
|
|
189
189
|
return dict(success=angle_dist < 0, angle_dist=angle_dist)
|
|
190
190
|
|
|
191
|
-
def _get_obs_extra(self, info:
|
|
191
|
+
def _get_obs_extra(self, info: dict):
|
|
192
192
|
obs = dict(
|
|
193
193
|
tcp_pose=self.agent.tcp.pose.raw_pose,
|
|
194
194
|
target_angle_diff=self.target_angle_diff,
|
|
@@ -231,7 +231,7 @@ class TurnFaucetEnv(BaseEnv):
|
|
|
231
231
|
# return reward
|
|
232
232
|
|
|
233
233
|
# def compute_normalized_dense_reward(
|
|
234
|
-
# self, obs: Any, action: torch.Tensor, info:
|
|
234
|
+
# self, obs: Any, action: torch.Tensor, info: dict
|
|
235
235
|
# ):
|
|
236
236
|
# max_reward = 10.0
|
|
237
237
|
# return self.compute_dense_reward(obs=obs, action=action, info=info) / max_reward
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
from typing import Any,
|
|
1
|
+
from typing import Any, Tuple
|
|
2
2
|
|
|
3
3
|
import numpy as np
|
|
4
4
|
import sapien
|
|
@@ -139,7 +139,7 @@ class TwoRobotPickCube(BaseEnv):
|
|
|
139
139
|
"is_right_arm_static": is_right_arm_static,
|
|
140
140
|
}
|
|
141
141
|
|
|
142
|
-
def _get_obs_extra(self, info:
|
|
142
|
+
def _get_obs_extra(self, info: dict):
|
|
143
143
|
obs = dict(
|
|
144
144
|
left_arm_tcp=self.left_agent.tcp.pose.raw_pose,
|
|
145
145
|
right_arm_tcp=self.right_agent.tcp.pose.raw_pose,
|
|
@@ -154,7 +154,7 @@ class TwoRobotPickCube(BaseEnv):
|
|
|
154
154
|
)
|
|
155
155
|
return obs
|
|
156
156
|
|
|
157
|
-
def compute_dense_reward(self, obs: Any, action: torch.Tensor, info:
|
|
157
|
+
def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: dict):
|
|
158
158
|
# Stage 1: Reach and push cube to be near other robot
|
|
159
159
|
tcp_to_obj_dist = torch.linalg.norm(
|
|
160
160
|
self.cube.pose.p - self.left_agent.tcp.pose.p, axis=1
|
|
@@ -257,6 +257,6 @@ class TwoRobotPickCube(BaseEnv):
|
|
|
257
257
|
return reward
|
|
258
258
|
|
|
259
259
|
def compute_normalized_dense_reward(
|
|
260
|
-
self, obs: Any, action: torch.Tensor, info:
|
|
260
|
+
self, obs: Any, action: torch.Tensor, info: dict
|
|
261
261
|
):
|
|
262
262
|
return self.compute_dense_reward(obs=obs, action=action, info=info) / 21
|
|
@@ -1,4 +1,4 @@
|
|
|
1
|
-
from typing import Any,
|
|
1
|
+
from typing import Any, Tuple
|
|
2
2
|
|
|
3
3
|
import numpy as np
|
|
4
4
|
import sapien
|
|
@@ -193,7 +193,7 @@ class TwoRobotStackCube(BaseEnv):
|
|
|
193
193
|
"success": success.bool(),
|
|
194
194
|
}
|
|
195
195
|
|
|
196
|
-
def _get_obs_extra(self, info:
|
|
196
|
+
def _get_obs_extra(self, info: dict):
|
|
197
197
|
obs = dict(
|
|
198
198
|
left_arm_tcp=self.left_agent.tcp.pose.raw_pose,
|
|
199
199
|
right_arm_tcp=self.right_agent.tcp.pose.raw_pose,
|
|
@@ -211,7 +211,7 @@ class TwoRobotStackCube(BaseEnv):
|
|
|
211
211
|
)
|
|
212
212
|
return obs
|
|
213
213
|
|
|
214
|
-
def compute_dense_reward(self, obs: Any, action: torch.Tensor, info:
|
|
214
|
+
def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: dict):
|
|
215
215
|
# Stage 1: Reach and grasp
|
|
216
216
|
# reaching reward for both robots to their respective cubes
|
|
217
217
|
cubeA_to_left_arm_tcp_dist = torch.linalg.norm(
|
|
@@ -291,6 +291,6 @@ class TwoRobotStackCube(BaseEnv):
|
|
|
291
291
|
return reward
|
|
292
292
|
|
|
293
293
|
def compute_normalized_dense_reward(
|
|
294
|
-
self, obs: Any, action: torch.Tensor, info:
|
|
294
|
+
self, obs: Any, action: torch.Tensor, info: dict
|
|
295
295
|
):
|
|
296
296
|
return self.compute_dense_reward(obs=obs, action=action, info=info) / 10
|
mani_skill/envs/template.py
CHANGED
|
@@ -21,7 +21,7 @@ mani_skill /envs/tasks/push_cube.py which is annotated with comments to explain
|
|
|
21
21
|
"""
|
|
22
22
|
|
|
23
23
|
|
|
24
|
-
from typing import Any,
|
|
24
|
+
from typing import Any, Union
|
|
25
25
|
|
|
26
26
|
import numpy as np
|
|
27
27
|
import sapien
|
|
@@ -169,21 +169,21 @@ class CustomEnv(BaseEnv):
|
|
|
169
169
|
"fail": torch.zeros(self.num_envs, device=self.device, dtype=bool),
|
|
170
170
|
}
|
|
171
171
|
|
|
172
|
-
def _get_obs_extra(self, info:
|
|
172
|
+
def _get_obs_extra(self, info: dict):
|
|
173
173
|
# should return an dict of additional observation data for your tasks
|
|
174
174
|
# this will be included as part of the observation in the "extra" key when obs_mode="state_dict" or any of the visual obs_modes
|
|
175
175
|
# and included as part of a flattened observation when obs_mode="state". Moreover, you have access to the info object
|
|
176
176
|
# which is generated by the `evaluate` function above
|
|
177
177
|
return dict()
|
|
178
178
|
|
|
179
|
-
def compute_dense_reward(self, obs: Any, action: torch.Tensor, info:
|
|
179
|
+
def compute_dense_reward(self, obs: Any, action: torch.Tensor, info: dict):
|
|
180
180
|
# you can optionally provide a dense reward function by returning a scalar value here. This is used when reward_mode="dense"
|
|
181
181
|
# note that as everything is batched, you must return a batch of of self.num_envs rewards as done in the example below.
|
|
182
182
|
# Moreover, you have access to the info object which is generated by the `evaluate` function above
|
|
183
183
|
return torch.zeros(self.num_envs, device=self.device)
|
|
184
184
|
|
|
185
185
|
def compute_normalized_dense_reward(
|
|
186
|
-
self, obs: Any, action: torch.Tensor, info:
|
|
186
|
+
self, obs: Any, action: torch.Tensor, info: dict
|
|
187
187
|
):
|
|
188
188
|
# this should be equal to compute_dense_reward / max possible reward
|
|
189
189
|
max_reward = 1.0
|
|
@@ -2,7 +2,6 @@
|
|
|
2
2
|
Functions that map a observation to a particular format, e.g. mapping the raw images to rgbd or pointcloud formats
|
|
3
3
|
"""
|
|
4
4
|
|
|
5
|
-
from typing import Dict
|
|
6
5
|
|
|
7
6
|
import numpy as np
|
|
8
7
|
import sapien.physx as physx
|
|
@@ -14,7 +13,7 @@ from mani_skill.sensors.camera import Camera
|
|
|
14
13
|
from mani_skill.utils import common
|
|
15
14
|
|
|
16
15
|
|
|
17
|
-
def sensor_data_to_pointcloud(observation:
|
|
16
|
+
def sensor_data_to_pointcloud(observation: dict, sensors: dict[str, BaseSensor]):
|
|
18
17
|
"""convert all camera data in sensor to pointcloud data"""
|
|
19
18
|
sensor_data = observation["sensor_data"]
|
|
20
19
|
camera_params = observation["sensor_param"]
|
|
@@ -29,7 +28,7 @@ def sensor_data_to_pointcloud(observation: Dict, sensors: Dict[str, BaseSensor])
|
|
|
29
28
|
# TODO: double check if the .clone()s are necessary
|
|
30
29
|
# Each pixel is (x, y, z, actor_id) in OpenGL camera space
|
|
31
30
|
# actor_id = 0 for the background
|
|
32
|
-
images:
|
|
31
|
+
images: dict[str, torch.Tensor]
|
|
33
32
|
position = images["position"].clone()
|
|
34
33
|
segmentation = images["segmentation"].clone()
|
|
35
34
|
position = position.float()
|
|
@@ -3,7 +3,7 @@ Code implementation for a batched random number generator. The goal is to enable
|
|
|
3
3
|
in CPU simulators and GPU simulators are the same
|
|
4
4
|
"""
|
|
5
5
|
|
|
6
|
-
from typing import
|
|
6
|
+
from typing import Union
|
|
7
7
|
|
|
8
8
|
import numpy as np
|
|
9
9
|
|
|
@@ -11,21 +11,21 @@ from mani_skill.utils import common
|
|
|
11
11
|
|
|
12
12
|
|
|
13
13
|
class BatchedRNG(np.random.RandomState):
|
|
14
|
-
def __init__(self, rngs:
|
|
14
|
+
def __init__(self, rngs: list):
|
|
15
15
|
self.rngs = rngs
|
|
16
16
|
self.batch_size = len(rngs)
|
|
17
17
|
|
|
18
18
|
@classmethod
|
|
19
|
-
def from_seeds(cls, seeds:
|
|
19
|
+
def from_seeds(cls, seeds: list[int], backend: str = "numpy:random_state"):
|
|
20
20
|
if backend == "numpy:random_state":
|
|
21
21
|
return cls(rngs=[np.random.RandomState(seed) for seed in seeds])
|
|
22
22
|
raise ValueError(f"Unknown batched RNG backend: {backend}")
|
|
23
23
|
|
|
24
24
|
@classmethod
|
|
25
|
-
def from_rngs(cls, rngs:
|
|
25
|
+
def from_rngs(cls, rngs: list):
|
|
26
26
|
return cls(rngs=rngs)
|
|
27
27
|
|
|
28
|
-
def __getitem__(self, idx: Union[int,
|
|
28
|
+
def __getitem__(self, idx: Union[int, list[int], np.ndarray]):
|
|
29
29
|
idx = common.to_numpy(idx)
|
|
30
30
|
if np.iterable(idx):
|
|
31
31
|
return BatchedRNG.from_rngs([self.rngs[i] for i in idx])
|
|
@@ -33,8 +33,8 @@ class BatchedRNG(np.random.RandomState):
|
|
|
33
33
|
|
|
34
34
|
def __setitem__(
|
|
35
35
|
self,
|
|
36
|
-
idx: Union[int,
|
|
37
|
-
value: Union[np.random.RandomState,
|
|
36
|
+
idx: Union[int, list[int], np.ndarray],
|
|
37
|
+
value: Union[np.random.RandomState, list[np.random.RandomState]],
|
|
38
38
|
):
|
|
39
39
|
idx = common.to_numpy(idx)
|
|
40
40
|
if np.iterable(idx):
|
|
@@ -2,7 +2,7 @@
|
|
|
2
2
|
Various sampling functions/classes for fast, vectorized sampling of e.g. object poses
|
|
3
3
|
"""
|
|
4
4
|
|
|
5
|
-
from typing import
|
|
5
|
+
from typing import Tuple
|
|
6
6
|
|
|
7
7
|
import torch
|
|
8
8
|
|
|
@@ -21,7 +21,7 @@ class UniformPlacementSampler:
|
|
|
21
21
|
|
|
22
22
|
def __init__(
|
|
23
23
|
self,
|
|
24
|
-
bounds: Tuple[
|
|
24
|
+
bounds: Tuple[list[float], list[float]],
|
|
25
25
|
batch_size: int,
|
|
26
26
|
device: Device = None,
|
|
27
27
|
) -> None:
|
|
@@ -65,7 +65,7 @@ class FrankaMoveBenchmarkEnv(BaseEnv):
|
|
|
65
65
|
qpos = self.agent.keyframes["rest"].qpos
|
|
66
66
|
qpos[0] = 0.5
|
|
67
67
|
self.agent.robot.set_qpos(qpos)
|
|
68
|
-
def _load_lighting(self, options:
|
|
68
|
+
def _load_lighting(self, options: dict):
|
|
69
69
|
self.scene.set_ambient_light(np.array([1,1,1])*0.1)
|
|
70
70
|
for i in range(self.num_envs):
|
|
71
71
|
self.scene.sub_scenes[i].set_environment_map(os.path.join(os.path.dirname(__file__), "kloofendal_28d_misty_puresky_1k.hdr"))
|
|
@@ -75,5 +75,5 @@ class FrankaMoveBenchmarkEnv(BaseEnv):
|
|
|
75
75
|
def evaluate(self):
|
|
76
76
|
return {}
|
|
77
77
|
|
|
78
|
-
def _get_obs_extra(self, info:
|
|
78
|
+
def _get_obs_extra(self, info: dict):
|
|
79
79
|
return dict()
|
|
@@ -79,7 +79,7 @@ class FrankaPickCubeBenchmarkEnv(BaseEnv):
|
|
|
79
79
|
qpos = self.agent.keyframes["rest"].qpos
|
|
80
80
|
self.agent.robot.set_qpos(qpos)
|
|
81
81
|
self.cube.set_pose(sapien.Pose(p=[0.6, 0, 0.02]))
|
|
82
|
-
def _load_lighting(self, options:
|
|
82
|
+
def _load_lighting(self, options: dict):
|
|
83
83
|
# self.scene.set_ambient_light(np.array([1,1,1])*0.05)
|
|
84
84
|
for i in range(self.num_envs):
|
|
85
85
|
self.scene.sub_scenes[i].set_environment_map(os.path.join(os.path.dirname(__file__), "kloofendal_28d_misty_puresky_1k.hdr"))
|
|
@@ -89,5 +89,5 @@ class FrankaPickCubeBenchmarkEnv(BaseEnv):
|
|
|
89
89
|
def evaluate(self):
|
|
90
90
|
return {}
|
|
91
91
|
|
|
92
|
-
def _get_obs_extra(self, info:
|
|
92
|
+
def _get_obs_extra(self, info: dict):
|
|
93
93
|
return dict()
|
|
@@ -31,7 +31,7 @@ class Profiler:
|
|
|
31
31
|
self.output_format = output_format
|
|
32
32
|
self.synchronize_torch = synchronize_torch
|
|
33
33
|
self.stats = dict()
|
|
34
|
-
|
|
34
|
+
|
|
35
35
|
# Initialize NVML
|
|
36
36
|
self._gpu_handle = None
|
|
37
37
|
try:
|
|
@@ -133,7 +133,7 @@ class Profiler:
|
|
|
133
133
|
return memory_usage
|
|
134
134
|
return None
|
|
135
135
|
def images_to_video(
|
|
136
|
-
images:
|
|
136
|
+
images: list[np.ndarray],
|
|
137
137
|
output_dir: str,
|
|
138
138
|
video_name: str,
|
|
139
139
|
fps: int = 10,
|
|
@@ -52,7 +52,7 @@ class Args:
|
|
|
52
52
|
quiet: bool = False
|
|
53
53
|
"""Disable verbose output."""
|
|
54
54
|
|
|
55
|
-
seed: Annotated[Optional[Union[int,
|
|
55
|
+
seed: Annotated[Optional[Union[int, list[int]]], tyro.conf.arg(aliases=["-s"])] = None
|
|
56
56
|
"""Seed(s) for random actions and simulator. Can be a single integer or a list of integers. Default is None (no seeds)"""
|
|
57
57
|
|
|
58
58
|
def main(args: Args):
|
mani_skill/render/shaders.py
CHANGED
|
@@ -1,5 +1,5 @@
|
|
|
1
1
|
from dataclasses import dataclass, field
|
|
2
|
-
from typing import Any, Callable
|
|
2
|
+
from typing import Any, Callable
|
|
3
3
|
|
|
4
4
|
import sapien
|
|
5
5
|
import torch
|
|
@@ -26,13 +26,13 @@ class ShaderConfig:
|
|
|
26
26
|
"""
|
|
27
27
|
|
|
28
28
|
shader_pack: str
|
|
29
|
-
texture_names:
|
|
29
|
+
texture_names: dict[str, list[str]] = field(default_factory=dict)
|
|
30
30
|
"""dictionary mapping shader texture names to the image modalities that are rendered. e.g. Color, Depth, Segmentation, etc."""
|
|
31
|
-
shader_pack_config:
|
|
31
|
+
shader_pack_config: dict[str, Any] = field(default_factory=dict)
|
|
32
32
|
"""configs for the shader pack. for e.g. the ray tracing shader you can configure the denoiser, samples per pixel, etc."""
|
|
33
33
|
|
|
34
|
-
texture_transforms:
|
|
35
|
-
str, Callable[[torch.Tensor],
|
|
34
|
+
texture_transforms: dict[
|
|
35
|
+
str, Callable[[torch.Tensor], dict[str, torch.Tensor]]
|
|
36
36
|
] = field(default_factory=dict)
|
|
37
37
|
"""texture transform functions that map each texture name to a function that converts the texture data into one or more standard image modalities. The return type should be a
|
|
38
38
|
dictionary with keys equal to the names of standard image modalities and values equal to the transformed data"""
|
|
@@ -1,5 +1,4 @@
|
|
|
1
1
|
from dataclasses import dataclass
|
|
2
|
-
from typing import Dict
|
|
3
2
|
|
|
4
3
|
import torch
|
|
5
4
|
|
|
@@ -36,7 +35,7 @@ class BaseSensor:
|
|
|
36
35
|
"""
|
|
37
36
|
raise NotImplementedError()
|
|
38
37
|
|
|
39
|
-
def get_params(self) ->
|
|
38
|
+
def get_params(self) -> dict:
|
|
40
39
|
"""
|
|
41
40
|
Get parameters for this sensor. Should return a dictionary with keys mapping to torch.Tensor values
|
|
42
41
|
"""
|
mani_skill/sensors/camera.py
CHANGED
|
@@ -2,7 +2,7 @@ from __future__ import annotations
|
|
|
2
2
|
|
|
3
3
|
import copy
|
|
4
4
|
from dataclasses import dataclass
|
|
5
|
-
from typing import TYPE_CHECKING,
|
|
5
|
+
from typing import TYPE_CHECKING, Optional, Union
|
|
6
6
|
|
|
7
7
|
import numpy as np
|
|
8
8
|
import sapien
|
|
@@ -68,7 +68,7 @@ class CameraConfig(BaseSensorConfig):
|
|
|
68
68
|
|
|
69
69
|
|
|
70
70
|
def update_camera_configs_from_dict(
|
|
71
|
-
camera_configs:
|
|
71
|
+
camera_configs: dict[str, CameraConfig], config_dict: dict[str, dict]
|
|
72
72
|
):
|
|
73
73
|
# Update CameraConfig to StereoDepthCameraConfig
|
|
74
74
|
if config_dict.pop("use_stereo_depth", False):
|
|
@@ -264,8 +264,8 @@ def normalize_depth(depth, min_depth=0, max_depth=None):
|
|
|
264
264
|
|
|
265
265
|
|
|
266
266
|
def camera_observations_to_images(
|
|
267
|
-
observations:
|
|
268
|
-
) ->
|
|
267
|
+
observations: dict[str, torch.Tensor], max_depth=None
|
|
268
|
+
) -> list[Array]:
|
|
269
269
|
"""Parse images from camera observations."""
|
|
270
270
|
images = dict()
|
|
271
271
|
for key in observations:
|