make scene builder api consistent

mani_skill2/envs/tasks/dexterity/rotate_single_object_in_hand.py

@@ -143,7 +143,7 @@ def _initialize_actors(self, env_idx: torch.Tensor):
         with torch.device(self.device):
             b = len(env_idx)
             # Initialize object pose
-            self.table_scene.initialize()
+            self.table_scene.initialize(env_idx)
             pose = self.obj.pose
             new_pos = torch.randn((b, 3)) * self.obj_init_pos_noise
             # hand_init_height is robot hand position while the 0.03 is a margin to ensure

mani_skill2/envs/tasks/dexterity/rotate_valve.py

@@ -1,5 +1,5 @@
 from collections import OrderedDict
-from typing import Union, Dict, Any, List
+from typing import Any, Dict, List, Union
 
 import numpy as np
 import torch
@@ -11,8 +11,7 @@
 from mani_skill2.utils.building.articulations import build_robel_valve
 from mani_skill2.utils.geometry.rotation_conversions import axis_angle_to_quaternion
 from mani_skill2.utils.registration import register_env
-from mani_skill2.utils.sapien_utils import get_obj_by_name
-from mani_skill2.utils.sapien_utils import look_at
+from mani_skill2.utils.sapien_utils import get_obj_by_name, look_at
 from mani_skill2.utils.scene_builder.table.table_scene_builder import TableSceneBuilder
 from mani_skill2.utils.structs.articulation import Articulation
 from mani_skill2.utils.structs.pose import Pose, vectorize_pose
@@ -138,7 +137,7 @@ def _load_articulations(self):
     def _initialize_actors(self, env_idx: torch.Tensor):
         with torch.device(self.device):
             b = len(env_idx)
-            self.table_scene.initialize()
+            self.table_scene.initialize(env_idx)
 
             # Initialize task related information
             if self.difficulty_level <= 3:

mani_skill2/envs/tasks/fmb/fmb.py

@@ -125,7 +125,7 @@ def _load_actors(self):
         self.bridge_grasp = builder.build_kinematic(name="bridge_grasp")
 
     def _initialize_actors(self, env_idx: torch.Tensor):
-        self.table_scene.initialize()
+        self.table_scene.initialize(env_idx)
         offset_pose = sapien.Pose(p=[0.02, -0.115, 0], q=euler2quat(0, 0, np.pi / 2))
         self.board.set_pose(
             sapien.Pose(p=np.array([0.115, 0.115, 0.034444])) * offset_pose

mani_skill2/envs/tasks/pick_cube.py

@@ -86,7 +86,7 @@ def _load_actors(self):
     def _initialize_actors(self, env_idx: torch.Tensor):
         with torch.device(self.device):
             b = len(env_idx)
-            self.table_scene.initialize()
+            self.table_scene.initialize(env_idx)
             xyz = torch.zeros((b, 3))
             xyz[:, :2] = torch.rand((b, 2)) * 0.2 - 0.1
             xyz[:, 2] = self.cube_half_size

mani_skill2/envs/tasks/pick_single_ycb.py

@@ -144,7 +144,7 @@ def _load_actors(self):
     def _initialize_actors(self, env_idx: torch.Tensor):
         with torch.device(self.device):
             b = len(env_idx)
-            self.table_scene.initialize()
+            self.table_scene.initialize(env_idx)
             xyz = torch.zeros((b, 3))
             xyz[:, :2] = torch.rand((b, 2)) * 0.2 - 0.1
             for i in range(b):

mani_skill2/envs/tasks/push_cube.py

@@ -132,7 +132,7 @@ def _initialize_actors(self, env_idx: torch.Tensor):
             b = len(env_idx)
             # when using scene builders, you must always call .initialize on them so they can set the correct poses of objects in the prebuilt scene
             # note that the table scene is built such that z=0 is the surface of the table.
-            self.table_scene.initialize()
+            self.table_scene.initialize(env_idx)
 
             # 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((b, 3))

mani_skill2/envs/tasks/stack_cube.py

@@ -78,7 +78,7 @@ def _load_actors(self):
     def _initialize_actors(self, env_idx: torch.Tensor):
         with torch.device(self.device):
             b = len(env_idx)
-            self.table_scene.initialize()
+            self.table_scene.initialize(env_idx)
 
             xyz = torch.zeros((b, 3))
             xyz[:, 2] = 0.02

mani_skill2/envs/tasks/two_robot_pick_cube.py

@@ -92,7 +92,7 @@ def _load_actors(self):
     def _initialize_actors(self, env_idx: torch.Tensor):
         with torch.device(self.device):
             b = len(env_idx)
-            self.table_scene.initialize()
+            self.table_scene.initialize(env_idx)
             xyz = torch.zeros((b, 3))
             xyz[:, 0] = torch.rand((b,)) * 0.1 - 0.05
             # ensure cube is spawned on the left side of the table

mani_skill2/envs/tasks/two_robot_stack_cube.py

@@ -98,7 +98,7 @@ def _load_actors(self):
     def _initialize_actors(self, env_idx: torch.Tensor):
         with torch.device(self.device):
             b = len(env_idx)
-            self.table_scene.initialize()
+            self.table_scene.initialize(env_idx)
             # the table scene initializes two robots. the first one self.agents[0] is on the left and the second one is on the right
 
             torch.zeros((b, 3))