check for environment being vectorized in the play logic in run.py

This commit is contained in:
Peter Zhokhov
2019-03-11 17:44:03 -07:00
parent 74101a9f24
commit 1259f6ab25

View File

@@ -6,7 +6,7 @@ from collections import defaultdict
import tensorflow as tf import tensorflow as tf
import numpy as np 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.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.cmd_util import common_arg_parser, parse_unknown_args, make_vec_env, make_env
from baselines.common.tf_util import get_session from baselines.common.tf_util import get_session
@@ -228,7 +228,7 @@ def main(args):
actions, _, _, _ = model.step(obs) actions, _, _, _ = model.step(obs)
obs, rew, done, _ = env.step(actions) obs, rew, done, _ = env.step(actions)
episode_rew += rew[0] episode_rew += rew[0] if isinstance(env, VecEnv) else rew
env.render() env.render()
done = done.any() if isinstance(done, np.ndarray) else done done = done.any() if isinstance(done, np.ndarray) else done
if done: if done: