use env.render in dummy_vec_env.render when num_envs == 1 (#74)

* use env.render in dummy_vec_env.render when num_envs == 1 * use shorter super() syntax per Alex's suggestion
2018-09-04 10:22:32 -07:00
parent 51cefc933b
commit b3f966aa02
1 changed files with 5 additions and 0 deletions
--- a/baselines/common/vec_env/dummy_vec_env.py
+++ b/baselines/common/vec_env/dummy_vec_env.py
@@ -66,3 +66,8 @@ class DummyVecEnv(VecEnv):
    def get_images(self):
        return [env.render(mode='rgb_array') for env in self.envs]
        
+    def render(self, mode='human'):
+        if self.num_envs == 1:
+            self.envs[0].render(mode=mode)
+        else:
+            super().render(mode=mode)