code format

mani_skill2/envs/pick_and_place/pick_cube.py

@@ -45,9 +45,7 @@ def _initialize_actors(self):
                 qs.append(q)
             qs = to_tensor(qs)
         # to set a batch of poses, use the Pose object or provide a raw tensor
-        obj_pose = Pose.create_from_pq(
-            p=xyz, q=qs
-        )
+        obj_pose = Pose.create_from_pq(p=xyz, q=qs)
 
         self.obj.set_pose(obj_pose)
 

mani_skill2/envs/pick_and_place/pick_single.py

@@ -79,7 +79,9 @@ def _check_assets(self):
         pass
 
     def _load_actors(self):
-        self.table_scene = TableSceneBuilder(env=self, robot_init_qpos_noise=self.robot_init_qpos_noise)
+        self.table_scene = TableSceneBuilder(
+            env=self, robot_init_qpos_noise=self.robot_init_qpos_noise
+        )
         self.table_scene.build()
         self._load_model()
         self.obj.set_linear_damping(0.1)
@@ -246,7 +248,7 @@ def evaluate(self, obs):
             obj_to_goal_pos=obj_to_goal_pos,
             is_obj_placed=is_obj_placed,
             is_robot_static=is_robot_static,
-            success=torch.logical_and(is_obj_placed, is_robot_static)
+            success=torch.logical_and(is_obj_placed, is_robot_static),
         )
 
     def compute_dense_reward(self, obs, action, info):
@@ -255,7 +257,6 @@ def compute_dense_reward(self, obs, action, info):
         # since the original reward can be unfriendly for RL,
         # even though MPC can solve many objects through the original reward.
 
-
         obj_pose = self.obj_pose
 
         # reaching reward
@@ -264,7 +265,8 @@ def compute_dense_reward(self, obs, action, info):
         reaching_reward = 1 - torch.tanh(
             3.0
             * torch.maximum(
-                tcp_to_obj_dist - torch.linalg.norm(self.model_bbox_size), torch.tensor(0.0)
+                tcp_to_obj_dist - torch.linalg.norm(self.model_bbox_size),
+                torch.tensor(0.0),
             )
         )
         reward = reaching_reward

mani_skill2/envs/tasks/__init__.py

@@ -1,3 +1,3 @@
 from .fmb.fmb import FMBEnv
 from .push_cube import PushCubeEnv
-from .push_object import PushObjectEnv
+from .push_object import PushObjectEnv

mani_skill2/envs/tasks/push_cube.py

@@ -96,7 +96,9 @@ def _initialize_actors(self):
 
         # here we write some randomization code that randomizes the x, y position of the cube we are pushing in the range [-0.1, -0.1] to [0.1, 0.1]
         xyz = torch.zeros((self.num_envs, 3), device=self.device)
-        xyz[..., :2] = torch.from_numpy(self._episode_rng.uniform(-0.1, 0.1, [self.num_envs, 2])).cuda()
+        xyz[..., :2] = torch.from_numpy(
+            self._episode_rng.uniform(-0.1, 0.1, [self.num_envs, 2])
+        ).cuda()
         xyz[..., 2] = self.cube_half_size
         q = [1, 0, 0, 0]
         # we can then create a pose object using Pose.create_from_pq to then set the cube pose with. Note that even though our quaternion
@@ -106,8 +108,12 @@ def _initialize_actors(self):
 
         # here we set the location of that red/white target (the goal region). In particular here, we set the position to be in front of the cube
         # and we further rotate 90 degrees on the y-axis to make the target object face up
-        target_region_xyz = xyz + torch.tensor([0.1 + self.goal_radius, 0, 0], device=self.device)
-        target_region_xyz[..., 2] = 1e-3 # set a little bit above 0 so the target is sitting on the table
+        target_region_xyz = xyz + torch.tensor(
+            [0.1 + self.goal_radius, 0, 0], device=self.device
+        )
+        target_region_xyz[
+            ..., 2
+        ] = 1e-3  # set a little bit above 0 so the target is sitting on the table
         self.goal_region.set_pose(
             Pose.create_from_pq(
                 p=target_region_xyz,
@@ -116,7 +122,7 @@ def _initialize_actors(self):
         )
 
     def _get_obs_extra(self):
-        # some useful observation info for solving the task includes the pose of the tcp (tool center point) which is the point between the 
+        # some useful observation info for solving the task includes the pose of the tcp (tool center point) which is the point between the
         # grippers of the robot
         obs = OrderedDict(
             tcp_pose=vectorize_pose(self.agent.tcp.pose),
@@ -131,29 +137,30 @@ def _get_obs_extra(self):
         return obs
 
     def evaluate(self, obs: Any):
-        # success is achieved when the cube's xy position on the table is within the 
+        # success is achieved when the cube's xy position on the table is within the
         # goal region's area (a circle centered at the goal region's xy position)
         is_obj_placed = (
             torch.linalg.norm(
                 self.obj.pose.p[..., :2] - self.goal_region.pose.p[..., :2], axis=1
             )
             < self.goal_radius
         )
-        
+
         return {
             "success": is_obj_placed,
         }
 
     def compute_dense_reward(self, obs: Any, action: Array, info: Dict):
         # We also create a pose marking where the robot should push the cube from that is easiest (pushing from behind the cube)
         tcp_push_pose = Pose.create_from_pq(
-            p=self.obj.pose.p + torch.tensor([-self.cube_half_size - 0.005, 0, 0], device=self.device)
+            p=self.obj.pose.p
+            + torch.tensor([-self.cube_half_size - 0.005, 0, 0], device=self.device)
         )
         tcp_to_push_pose = tcp_push_pose.p - self.agent.tcp.pose.p
         tcp_to_push_pose_dist = torch.linalg.norm(tcp_to_push_pose, axis=1)
         reaching_reward = 1 - torch.tanh(5 * tcp_to_push_pose_dist)
         reward = reaching_reward
-        
+
         # compute a placement reward to encourage robot to move the cube to the center of the goal region
         # we further multiply the place_reward by a mask reached so we only add the place reward if the robot has reached the desired push pose
         # This reward design helps train RL agents faster by staging the reward out.

mani_skill2/envs/tasks/push_object.py

@@ -9,7 +9,13 @@
 from mani_skill2.agents.robots.xmate3.xmate3 import Xmate3Robotiq
 from mani_skill2.envs.sapien_env import BaseEnv
 from mani_skill2.sensors.camera import CameraConfig
-from mani_skill2.utils.building.actors import build_actor_ycb, build_cube, build_red_white_target, model_dbs, _load_ycb_dataset
+from mani_skill2.utils.building.actors import (
+    _load_ycb_dataset,
+    build_actor_ycb,
+    build_cube,
+    build_red_white_target,
+    model_dbs,
+)
 from mani_skill2.utils.registration import register_env
 from mani_skill2.utils.sapien_utils import (  # import various useful utilities for working with sapien
     look_at,
@@ -53,9 +59,14 @@ def _load_actors(self):
         )
         self.table_scene.build()
 
-
-        self.obj, height = build_actor_ycb(self.model_id, self._scene, name=self.model_id, body_type="dynamic", add_collision=True)
-        self.obj_z = height/2
+        self.obj, height = build_actor_ycb(
+            self.model_id,
+            self._scene,
+            name=self.model_id,
+            body_type="dynamic",
+            add_collision=True,
+        )
+        self.obj_z = height / 2
 
         self.goal_region = build_red_white_target(
             self._scene,
@@ -83,7 +94,9 @@ def _initialize_actors(self):
 
         # here we write some randomization code that randomizes the x, y position of the cube we are pushing in the range [-0.1, -0.1] to [0.1, 0.1]
         xyz = torch.zeros((self.num_envs, 3), device=self.device)
-        xyz[..., :2] = torch.from_numpy(self._episode_rng.uniform(-0.1, 0.1, [self.num_envs, 2])).cuda()
+        xyz[..., :2] = torch.from_numpy(
+            self._episode_rng.uniform(-0.1, 0.1, [self.num_envs, 2])
+        ).cuda()
         xyz[..., 2] = self.obj_z
         q = [1, 0, 0, 0]
         # we can then create a pose object using Pose.create_from_pq to then set the cube pose with. Note that even though our quaternion
@@ -94,8 +107,12 @@ def _initialize_actors(self):
         # here we set the location of that red/white target (the goal region). In particular here, we set the position to be in front of the cube
         # and we further rotate 90 degrees on the y-axis to make the target object face up
         xyz = torch.zeros((self.num_envs, 3), device=self.device)
-        target_region_xyz = xyz + torch.tensor([0.1 + self.goal_radius, 0, 0], device=self.device)
-        target_region_xyz[..., 2] = 1e-3 # set a little bit above 0 so the target is sitting on the table
+        target_region_xyz = xyz + torch.tensor(
+            [0.1 + self.goal_radius, 0, 0], device=self.device
+        )
+        target_region_xyz[
+            ..., 2
+        ] = 1e-3  # set a little bit above 0 so the target is sitting on the table
         self.goal_region.set_pose(
             Pose.create_from_pq(
                 p=target_region_xyz,
@@ -104,7 +121,7 @@ def _initialize_actors(self):
         )
 
     def _get_obs_extra(self):
-        # some useful observation info for solving the task includes the pose of the tcp (tool center point) which is the point between the 
+        # some useful observation info for solving the task includes the pose of the tcp (tool center point) which is the point between the
         # grippers of the robot
         obs = OrderedDict(
             tcp_pose=vectorize_pose(self.agent.tcp.pose),
@@ -119,15 +136,15 @@ def _get_obs_extra(self):
         return obs
 
     def evaluate(self, obs: Any):
-        # success is achieved when the cube's xy position on the table is within the 
+        # success is achieved when the cube's xy position on the table is within the
         # goal region's area (a circle centered at the goal region's xy position)
         # is_obj_placed = (
         #     torch.linalg.norm(
         #         self.obj.pose.p[..., :2] - self.goal_region.pose.p[..., :2], axis=1
         #     )
         #     < self.goal_radius
         # )
-        
+
         return {
             "success": torch.zeros(self.num_envs, device=self.device),
         }
@@ -136,13 +153,14 @@ def compute_dense_reward(self, obs: Any, action: Array, info: Dict):
         return torch.zeros(self.num_envs, device=self.device)
         # We also create a pose marking where the robot should push the cube from that is easiest (pushing from behind the cube)
         tcp_push_pose = Pose.create_from_pq(
-            p=self.obj.pose.p + torch.tensor([-self.cube_half_size - 0.005, 0, 0], device=self.device)
+            p=self.obj.pose.p
+            + torch.tensor([-self.cube_half_size - 0.005, 0, 0], device=self.device)
         )
         tcp_to_push_pose = tcp_push_pose.p - self.agent.tcp.pose.p
         tcp_to_push_pose_dist = torch.linalg.norm(tcp_to_push_pose, axis=1)
         reaching_reward = 1 - torch.tanh(5 * tcp_to_push_pose_dist)
         reward = reaching_reward
-        
+
         # compute a placement reward to encourage robot to move the cube to the center of the goal region
         # we further multiply the place_reward by a mask reached so we only add the place reward if the robot has reached the desired push pose
         # This reward design helps train RL agents faster by staging the reward out.

mani_skill2/utils/building/actors.py

@@ -265,7 +265,7 @@ def build_actor_ycb(
     name: str,
     root_dir=ASSET_DIR / "mani_skill2_ycb",
     body_type: str = "dynamic",
-    add_collision: bool = True
+    add_collision: bool = True,
 ):
     if "YCB" not in model_dbs:
         _load_ycb_dataset()
@@ -288,7 +288,7 @@ def build_actor_ycb(
             material=physical_material,
             density=density,
             decomposition="coacd",
-    )
+        )
 
     visual_file = str(model_dir / "textured.obj")
     builder.add_visual_from_file(filename=visual_file, scale=[scale] * 3)