• Search by Module
• Search by Words
• Search Projects

• Java
• Python
• JavaScript
• TypeScript
• C++
• Scala
• Blog

• lirpg-master
  • LICENSE
  • setup.py
  • data
  • baselines
    • gail
      • statistics.py
      • result
        • gail-result.md
      • run_mujoco.py
      • adversary.py
      • dataset
        • mujoco_dset.py
        • __init__.py
      • mlp_policy.py
      • __init__.py
      • README.md
      • trpo_mpi.py
      • behavior_clone.py
      • gail-eval.py
    • her
      • util.py
      • ddpg.py
      • rollout.py
      • replay_buffer.py
      • actor_critic.py
      • normalizer.py
      • __init__.py
      • README.md
      • experiment
        • config.py
        • play.py
        • plot.py
        • train.py
        • __init__.py
      • her.py
    • deepq
      • simple.py
      • replay_buffer.py
      • build_graph.py
      • models.py
      • experiments
        • run_atari.py
        • train_cartpole.py
        • enjoy_mountaincar.py
        • enjoy_cartpole.py
        • enjoy_pong.py
        • train_mountaincar.py
        • custom_cartpole.py
        • __init__.py
      • __init__.py
      • README.md
      • utils.py
    • ddpg
      • ddpg.py
      • training.py
      • models.py
      • noise.py
      • memory.py
      • __init__.py
      • main.py
      • README.md
    • acer
      • run_atari.py
      • buffer.py
      • policies.py
      • acer_simple.py
      • __init__.py
      • README.md
    • a2c
      • run_atari.py
      • policies.py
      • __init__.py
      • README.md
      • utils.py
      • a2c.py
    • ppo2
      • run_atari.py
      • ppo2.py
      • run_mujoco.py
      • policies.py
      • __init__.py
      • README.md
    • ppo1
      • pposgd_simple.py
      • run_atari.py
      • run_mujoco.py
      • cnn_policy.py
      • mlp_policy.py
      • __init__.py
      • README.md
    • common
      • console_util.py
      • mpi_adam.py
      • mpi_moments.py
      • misc_util.py
      • mpi_fork.py
      • distributions.py
      • schedules.py
      • math_util.py
      • vec_env
        • subproc_vec_env.py
        • dummy_vec_env.py
        • vec_normalize.py
        • vec_frame_stack.py
        • __init__.py
      • cmd_util.py
      • atari_wrappers.py
      • cg.py
      • __init__.py
      • tests
        • test_tf_util.py
        • test_schedules.py
        • test_segment_tree.py
      • mpi_running_mean_std.py
      • dataset.py
      • tf_util.py
      • running_mean_std.py
      • segment_tree.py
    • logger.py
    • __init__.py
    • results_plotter.py
    • trpo_mpi
      • run_atari.py
      • run_mujoco.py
      • nosharing_cnn_policy.py
      • __init__.py
      • README.md
      • trpo_mpi.py
    • acktr
      • run_atari.py
      • acktr_cont.py
      • run_mujoco.py
      • running_stat.py
      • value_functions.py
      • policies.py
      • filters.py
      • acktr_disc.py
      • __init__.py
      • kfac_utils.py
      • README.md
      • utils.py
      • kfac.py
    • bench
      • monitor.py
      • __init__.py
      • benchmarks.py
  • README.md
  • .gitignore

"""
Helpers for scripts like run_atari.py.
"""

import os
import gym
from gym.wrappers import FlattenDictWrapper
from baselines import logger
from baselines.bench import Monitor
from baselines.common import set_global_seeds
from baselines.common.atari_wrappers import make_atari, wrap_deepmind
from baselines.common.vec_env.subproc_vec_env import SubprocVecEnv
from mpi4py import MPI

def make_atari_env(env_id, num_env, seed, wrapper_kwargs=None, start_index=0):
    """
    Create a wrapped, monitored SubprocVecEnv for Atari.
    """
    if wrapper_kwargs is None: wrapper_kwargs = {}
    def make_env(rank): # pylint: disable=C0111
        def _thunk():
            env = make_atari(env_id)
            env.seed(seed + rank)
            env = Monitor(env, logger.get_dir() and os.path.join(logger.get_dir(), str(rank)))
            return wrap_deepmind(env, **wrapper_kwargs)
        return _thunk
    set_global_seeds(seed)
    return SubprocVecEnv([make_env(i + start_index) for i in range(num_env)])

def make_mujoco_env(env_id, seed):
    """
    Create a wrapped, monitored gym.Env for MuJoCo.
    """
    set_global_seeds(seed)
    env = gym.make(env_id)
    env = Monitor(env, logger.get_dir())
    env.seed(seed)
    return env

def make_robotics_env(env_id, seed, rank=0):
    """
    Create a wrapped, monitored gym.Env for MuJoCo.
    """
    set_global_seeds(seed)
    env = gym.make(env_id)
    env = FlattenDictWrapper(env, ['observation', 'desired_goal'])
    env = Monitor(
        env, logger.get_dir() and os.path.join(logger.get_dir(), str(rank)),
        info_keywords=('is_success',))
    env.seed(seed)
    return env

def arg_parser():
    """
    Create an empty argparse.ArgumentParser.
    """
    import argparse
    return argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)

def atari_arg_parser():
    """
    Create an argparse.ArgumentParser for run_atari.py.
    """
    parser = arg_parser()
    parser.add_argument('--env', help='environment ID', default='BreakoutNoFrameskip-v4')
    parser.add_argument('--seed', help='RNG seed', type=int, default=0)
    parser.add_argument('--num-timesteps', type=int, default=int(10e6))
    return parser

def mujoco_arg_parser():
    """
    Create an argparse.ArgumentParser for run_mujoco.py.
    """
    parser = arg_parser()
    parser.add_argument('--env', help='environment ID', type=str, default='Reacher-v2')
    parser.add_argument('--seed', help='RNG seed', type=int, default=0)
    parser.add_argument('--num-timesteps', type=int, default=int(1e6))
    return parser

def robotics_arg_parser():
    """
    Create an argparse.ArgumentParser for run_mujoco.py.
    """
    parser = arg_parser()
    parser.add_argument('--env', help='environment ID', type=str, default='FetchReach-v0')
    parser.add_argument('--seed', help='RNG seed', type=int, default=0)
    parser.add_argument('--num-timesteps', type=int, default=int(1e6))
    return parser