From b3f966aa028c04d33c491268b021aac91b0b64c1 Mon Sep 17 00:00:00 2001 From: pzhokhov Date: Tue, 4 Sep 2018 10:22:32 -0700 Subject: [PATCH] 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 --- baselines/common/vec_env/dummy_vec_env.py | 5 +++++ 1 file changed, 5 insertions(+) diff --git a/baselines/common/vec_env/dummy_vec_env.py b/baselines/common/vec_env/dummy_vec_env.py index 3fbb8e1..09865b0 100644 --- 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)