Make robot initial pose seedable

robosuite/environments/robot_env.py

@@ -345,7 +345,7 @@ def _setup_observables(self):
             # Create sensor information
             sensors = []
             names = []
-            for (cam_name, cam_w, cam_h, cam_d, cam_segs) in zip(
+            for cam_name, cam_w, cam_h, cam_d, cam_segs in zip(
                 self.camera_names,
                 self.camera_widths,
                 self.camera_heights,
@@ -523,7 +523,7 @@ def _reset_internal(self):
 
         # Reset robot and update action space dimension along the way
         for robot in self.robots:
-            robot.reset(deterministic=self.deterministic_reset)
+            robot.reset(deterministic=self.deterministic_reset, rng=self.rng)
             self._action_dim += robot.action_dim
 
         # Update cameras if appropriate

robosuite/robots/fixed_base_robot.py

@@ -66,16 +66,17 @@ def load_model(self):
         # First, run the superclass method to load the relevant model
         super().load_model()
 
-    def reset(self, deterministic=False):
+    def reset(self, deterministic=False, rng=None):
         """
         Sets initial pose of arm and grippers. Overrides gripper joint configuration if we're using a
         deterministic reset (e.g.: hard reset from xml file)
 
         Args:
             deterministic (bool): If true, will not randomize initializations within the sim
+            rng (numpy.random._generator.Generator): Seeded random number generator
         """
         # First, run the superclass method to reset the position and controller
-        super().reset(deterministic)
+        super().reset(deterministic, rng)
 
     def setup_references(self):
         """

robosuite/robots/legged_robot.py

@@ -117,16 +117,17 @@ def load_model(self):
         # First, run the superclass method to load the relevant model
         super().load_model()
 
-    def reset(self, deterministic=False):
+    def reset(self, deterministic=False, rng=None):
         """
         Sets initial pose of arm and grippers. Overrides gripper joint configuration if we're using a
         deterministic reset (e.g.: hard reset from xml file)
 
         Args:
             deterministic (bool): If true, will not randomize initializations within the sim
+            rng (numpy.random._generator.Generator): Seeded random number generator
         """
         # First, run the superclass method to reset the position and controller
-        super().reset(deterministic)
+        super().reset(deterministic, rng)
 
         # Set initial q pos of the legged base
         if isinstance(self.robot_model.base, LegBaseModel):

robosuite/robots/mobile_robot.py

@@ -192,16 +192,17 @@ def load_model(self):
         # First, run the superclass method to load the relevant model
         super().load_model()
 
-    def reset(self, deterministic=False):
+    def reset(self, deterministic=False, rng=None):
         """
         Sets initial pose of arm and grippers. Overrides gripper joint configuration if we're using a
         deterministic reset (e.g.: hard reset from xml file)
 
         Args:
             deterministic (bool): If true, will not randomize initializations within the sim
+            rng (numpy.random._generator.Generator): Seeded random number generator
         """
         # First, run the superclass method to reset the position and controller
-        super().reset(deterministic)
+        super().reset(deterministic, rng)
 
     def setup_references(self):
         """

robosuite/robots/robot.py

@@ -231,13 +231,14 @@ def reset_sim(self, sim: MjSim):
         """
         self.sim = sim
 
-    def reset(self, deterministic=False):
+    def reset(self, deterministic=False, rng=None):
         """
         Sets initial pose of arm and grippers. Overrides robot joint configuration if we're using a
         deterministic reset (e.g.: hard reset from xml file)
 
         Args:
             deterministic (bool): If true, will not randomize initializations within the sim
+            rng (numpy.random._generator.Generator): Seeded random number generator
 
         Raises:
             ValueError: [Invalid noise type]
@@ -246,9 +247,15 @@ def reset(self, deterministic=False):
         if not deterministic:
             # Determine noise
             if self.initialization_noise["type"] == "gaussian":
-                noise = np.random.randn(len(self.init_qpos)) * self.initialization_noise["magnitude"]
+                if rng is None:
+                    noise = np.random.randn(len(self.init_qpos)) * self.initialization_noise["magnitude"]
+                else:
+                    noise = rng.standard_normal(len(self.init_qpos)) * self.initialization_noise["magnitude"]
             elif self.initialization_noise["type"] == "uniform":
-                noise = np.random.uniform(-1.0, 1.0, len(self.init_qpos)) * self.initialization_noise["magnitude"]
+                if rng is None:
+                    noise = np.random.uniform(-1.0, 1.0, len(self.init_qpos)) * self.initialization_noise["magnitude"]
+                else:
+                    noise = rng.uniform(-1.0, 1.0, len(self.init_qpos)) * self.initialization_noise["magnitude"]
             else:
                 raise ValueError("Error: Invalid noise type specified. Options are 'gaussian' or 'uniform'.")
             init_qpos += noise

robosuite/robots/wheeled_robot.py

@@ -73,16 +73,17 @@ def load_model(self):
         # First, run the superclass method to load the relevant model
         super().load_model()
 
-    def reset(self, deterministic=False):
+    def reset(self, deterministic=False, rng=None):
         """
         Sets initial pose of arm and grippers. Overrides gripper joint configuration if we're using a
         deterministic reset (e.g.: hard reset from xml file)
 
         Args:
             deterministic (bool): If true, will not randomize initializations within the sim
+            rng (numpy.random._generator.Generator): Seeded random number generator
         """
         # First, run the superclass method to reset the position and controller
-        super().reset(deterministic)
+        super().reset(deterministic, rng)
 
     def setup_references(self):
         """