From c79b3373bf4fa276fa05151494f4777c0d8f8a36 Mon Sep 17 00:00:00 2001 From: Peter Zhokhov Date: Fri, 5 Apr 2019 14:43:09 -0700 Subject: [PATCH] parse colon-separated env_id's --- baselines/common/cmd_util.py | 1 - baselines/run.py | 6 +++--- 2 files changed, 3 insertions(+), 4 deletions(-) diff --git a/baselines/common/cmd_util.py b/baselines/common/cmd_util.py index 8e0eced..47e96ff 100644 --- a/baselines/common/cmd_util.py +++ b/baselines/common/cmd_util.py @@ -150,7 +150,6 @@ def common_arg_parser(): parser.add_argument('--save_video_interval', help='Save video every x steps (0 = disabled)', default=0, type=int) parser.add_argument('--save_video_length', help='Length of recorded video. Default: 200', default=200, type=int) parser.add_argument('--play', default=False, action='store_true') - parser.add_argument('--extra_import', help='Extra module to import to access external environments', type=str, default=None) return parser def robotics_arg_parser(): diff --git a/baselines/run.py b/baselines/run.py index c438d1c..4f9ed15 100644 --- a/baselines/run.py +++ b/baselines/run.py @@ -1,4 +1,5 @@ import sys +import re import multiprocessing import os.path as osp import gym @@ -137,6 +138,8 @@ def get_env_type(args): if env_id in e: env_type = g break + if ':' in env_id: + env_type = re.sub(r':.*', '', env_id) assert env_type is not None, 'env_id {} is not recognized in env types'.format(env_id, _game_envs.keys()) return env_type, env_id @@ -197,9 +200,6 @@ def main(args): args, unknown_args = arg_parser.parse_known_args(args) extra_args = parse_cmdline_kwargs(unknown_args) - if args.extra_import is not None: - import_module(args.extra_import) - if MPI is None or MPI.COMM_WORLD.Get_rank() == 0: rank = 0 logger.configure()