From 1259f6ab25c6f7261e33c4c3b92df869188f9260 Mon Sep 17 00:00:00 2001 From: Peter Zhokhov Date: Mon, 11 Mar 2019 17:44:03 -0700 Subject: [PATCH] check for environment being vectorized in the play logic in run.py --- baselines/run.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/baselines/run.py b/baselines/run.py index d1c0ac7..5a48b18 100644 --- a/baselines/run.py +++ b/baselines/run.py @@ -6,7 +6,7 @@ from collections import defaultdict import tensorflow as tf import numpy as np -from baselines.common.vec_env import VecFrameStack, VecNormalize +from baselines.common.vec_env import VecFrameStack, VecNormalize, VecEnv from baselines.common.vec_env.vec_video_recorder import VecVideoRecorder from baselines.common.cmd_util import common_arg_parser, parse_unknown_args, make_vec_env, make_env from baselines.common.tf_util import get_session @@ -228,7 +228,7 @@ def main(args): actions, _, _, _ = model.step(obs) obs, rew, done, _ = env.step(actions) - episode_rew += rew[0] + episode_rew += rew[0] if isinstance(env, VecEnv) else rew env.render() done = done.any() if isinstance(done, np.ndarray) else done if done: