Allow exclude arm from randomization.

SatireY · Mar 28, 2022 · 5ea7e98 · 5ea7e98
1 parent c2bd28c
commit 5ea7e98
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 5 deletions.
diff --git a/rlbench/sim2real/domain_randomization.py b/rlbench/sim2real/domain_randomization.py
@@ -43,14 +43,17 @@ class RandomizationConfig(object):
 
     def __init__(self,
                  whitelist: List[str]=None,
-                 blacklist: List[str]=None):
+                 blacklist: List[str]=None,
+                 randomize_arm: bool = True):
         self.whitelist = whitelist
         self.blacklist = [] if blacklist is None else blacklist
+        self.randomize_arm = randomize_arm
 
     def should_randomize(self, obj_name: str):
         return ((self.whitelist is None and len(self.blacklist) == 0) or
                 (self.whitelist is not None and obj_name in self.whitelist) or
-                (obj_name not in self.blacklist))
+                (obj_name not in self.blacklist)) and (
+                self.randomize_arm or 'panda' not in obj_name.lower())
 
 
 class DynamicsRandomizationConfig(RandomizationConfig):
@@ -61,9 +64,10 @@ class VisualRandomizationConfig(RandomizationConfig):
 
     def __init__(self,
                  image_directory: str,
-                 whitelist: List[str]=None,
-                 blacklist: List[str]=None):
-        super().__init__(whitelist, blacklist)
+                 whitelist: List[str] = None,
+                 blacklist: List[str] = None,
+                 randomize_arm: bool = True):
+        super().__init__(whitelist, blacklist, randomize_arm)
         self._image_directory = image_directory
         if not os.path.exists(image_directory):
             raise NotADirectoryError(

diff --git a/rlbench/sim2real/domain_randomization_scene.py b/rlbench/sim2real/domain_randomization_scene.py
@@ -15,6 +15,7 @@
     'mapping_mode': TextureMappingMode.PLANE,
     'repeat_along_u': True,
     'repeat_along_v': True,
+    'uv_scaling': [4., 4.]
 }
 
 
@@ -47,6 +48,7 @@ def __init__(self,
         self._scene_objects += self.robot.gripper.get_visuals()
         if self._visual_rand_config is not None:
             # Make the floor plane renderable (to cover old floor)
+            self._scene_objects[0].set_position([0, 0, 0.01])
             self._scene_objects[0].set_renderable(True)
 
     def _should_randomize_episode(self, index: int):