fix rgb_array rendering

SatireY · Mar 7, 2022 · 568f9d9 · 568f9d9
1 parent 97f1268
commit 568f9d9
Showing 1 changed file with 3 additions and 1 deletion.
diff --git a/rlbench/gym/rlbench_env.py b/rlbench/gym/rlbench_env.py
@@ -94,7 +94,9 @@ def render(self, mode='human') -> Union[None, np.ndarray]:
                 'You passed in mode %s, but expected %s.' % (
                     mode, self._render_mode))
         if mode == 'rgb_array':
-            return self._gym_cam.capture_rgb()
+            frame = self._gym_cam.capture_rgb()
+            frame = np.clip((frame * 255.).astype(np.uint8), 0, 255)
+            return frame
 
     def reset(self) -> Dict[str, np.ndarray]:
         descriptions, obs = self.task.reset()