Use uniform random for object pose.

isri-aist · Jan 25, 2025 · 1632c8c · 1632c8c
1 parent 5f1ed31
commit 1632c8c
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 1 deletion.
diff --git a/robo_manip_baselines/common/rollout/RolloutBase.py b/robo_manip_baselines/common/rollout/RolloutBase.py
@@ -169,6 +169,9 @@ def setup_args(self, parser=None, argv=None):
             argv = sys.argv
         self.args = parser.parse_args(argv[1:])
 
+        if self.args.world_random_scale is not None:
+            self.args.world_random_scale = np.array(self.args.world_random_scale)
+
     @abstractmethod
     def setup_policy(self):
         pass

diff --git a/robo_manip_baselines/envs/mujoco/ur5e/MujocoUR5eInsertEnv.py b/robo_manip_baselines/envs/mujoco/ur5e/MujocoUR5eInsertEnv.py
@@ -50,7 +50,9 @@ def modify_world(
 
         hole_pos = self.original_hole_pos + self.hole_pos_offsets[world_idx]
         if world_random_scale is not None:
-            hole_pos += np.random.normal(scale=world_random_scale, size=3)
+            hole_pos += np.random.uniform(
+                low=-1.0 * world_random_scale, high=world_random_scale, size=3
+            )
         self.model.body("hole").pos = hole_pos
 
         return world_idx