code: import sys
  import multiprocessing
  import os.path as osp
  import gym
  from collections import defaultdict
  import tensorflow as tf
  import numpy as np
  
  from baselines.common.vec_env.vec_video_recorder import VecVideoRecorder
  from baselines.common.vec_env.vec_frame_stack import VecFrameStack
  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 import logger
  from importlib import import_module
  
  from baselines.common.vec_env.vec_normalize import VecNormalize
  
  try:
      from mpi4py import MPI
  except ImportError:
      print("can't import MPI ")
      MPI = None
  
  try:
      import pybullet_envs
  except ImportError:
      pybullet_envs = None
  
  try:
      import roboschool
  except ImportError:
      roboschool = None
  
  _game_envs = defaultdict(set)
  for env in gym.envs.registry.all():
      # TODO: solve this with regexes
      env_type = env._entry_point.split(':')[0].split('.')[-1]
      _game_envs[env_type].add(env.id)
  
  # reading benchmark names directly from retro requires
  # importing retro here, and for some reason that crashes tensorflow
  # in ubuntu
  _game_envs['retro'] = {
      'BubbleBobble-Nes',
      'SuperMarioBros-Nes',
      'TwinBee3PokoPokoDaimaou-Nes',
      'SpaceHarrier-Nes',
      'SonicTheHedgehog-Genesis',
      'Vectorman-Genesis',
      'FinalFight-Snes',
      'SpaceInvaders-Snes',
  }
  
  
  def train(args, extra_args):
  
      env_type, env_id = get_env_type(args.env)
      print("In the train function with env_type {} env_id {}".format(env_type , env_id))
      # print('env_type: {}'.format(env_type))
  
      total_timesteps = int(args.num_timesteps)
      seed = args.seed
  
      learn = get_learn_function(args.alg)
      alg_kwargs = get_learn_function_defaults(args.alg, env_type)
      alg_kwargs.update(extra_args)
  
  
  
      print("Now called build_env env function with arg :: ",args)
      env = build_env(args)
      if args.save_video_interval != 0:
          env = VecVideoRecorder(env, osp.join(logger.Logger.CURRENT.dir, "videos"), record_video_trigger=lambda x: x % args.save_video_interval == 0, video_length=args.save_video_length)
  
      if args.network:
          alg_kwargs['network'] = args.network
      else:
          if alg_kwargs.get('network') is None:
              alg_kwargs['network'] = get_default_network(env_type)
  
      print('Training {} on {}:{} with arguments \n{}'.format(args.alg, env_type, env_id, alg_kwargs))
  
      model = learn(
          env=env,
          seed=seed,
          total_timesteps=total_timesteps,
          **alg_kwargs
      )
  
      return model, env
  
  
  def random_agent_ob_mean_std(env, nsteps=10000):
      ob = np.asarray(env.reset())
      if MPI.COMM_WORLD.Get_rank() == 0:
          obs = [ob]
          print("::: Entered the random action for 1000 steps ::: ")
          for _ in range(nsteps):
              ac = env.action_space.sample() # random action 
              ob, _, done, _ = env.step(ac)
              if done:
                  ob = env.reset()
              obs.append(np.asarray(ob))
          mean = np.mean(obs, 0).astype(np.float32)
          std = np.std(obs, 0).mean().astype(np.float32)
      else:
          mean = np.empty(shape=ob.shape, dtype=np.float32)
          std = np.empty(shape=(), dtype=np.float32)
      MPI.COMM_WORLD.Bcast(mean, root=0)
      MPI.COMM_WORLD.Bcast(std, root=0)
      return mean, std
  
  def build_env(args):
  
      ncpu = multiprocessing.cpu_count()
      if sys.platform == 'darwin': ncpu //= 2
      nenv = args.num_env or ncpu
      alg = args.alg
      seed = args.seed
  
      env_type, env_id = get_env_type(args.env)
      print("In the build_env function with alg :: ",alg)
      if env_type in {'atari', 'retro'}:
          if alg == 'deepq':
              env = make_env(env_id, env_type, seed=seed, wrapper_kwargs={'frame_stack': True})
          elif alg == 'trpo_mpi':
              env = make_env(env_id, env_type, seed=seed)
          else:
              frame_stack_size = 4
              print("make_vec_env arguments env_id {} , env_type {} , nenv {} ,seed {} , gamestate {} reward_scale {}".format(
                  env_id , env_type , nenv , seed , args.gamestate , args.reward_scale))
              
              #>
              # print("Called environment for mean and std")
              # env = make_vec_env(env_id, env_type, 1, seed, gamestate=args.gamestate, reward_scale=args.reward_scale)
              # # env = VecFrameStack(env, frame_stack_size) ## No need for frame stacking while calculation of mean and std
              # ob_mean, ob_std = random_agent_ob_mean_std(env)
              # print(" environment complete with mean {} and std {}".format(ob_mean , ob_std))
              # del env 
              #>
  
              env = make_vec_env(env_id, env_type, nenv, seed, gamestate=args.gamestate, reward_scale=args.reward_scale)
              
              # print("Received env from make_vec_env  type env {} and env ".format(
                  # type(env) , env))
              print("ob_space {} and ac_space {} ".format(env.observation_space, env.action_space))
              env = VecFrameStack(env, frame_stack_size)
  
  
              print("After Frame stacking env would become " )
  
      else:
         config = tf.ConfigProto(allow_soft_placement=True,
                                 intra_op_parallelism_threads=1,
                                 inter_op_parallelism_threads=1)
         config.gpu_options.allow_growth = True
         get_session(config=config)
  
         env = make_vec_env(env_id, env_type, args.num_env or 1, seed, reward_scale=args.reward_scale)
  
         if env_type == 'mujoco':
             env = VecNormalize(env)
  
      return env #, ob_mean, ob_std
  
  
  def get_env_type(env_id):
      if env_id in _game_envs.keys():
          env_type = env_id
          env_id = [g for g in _game_envs[env_type]][0]
      else:
          env_type = None
          for g, e in _game_envs.items():
              if env_id in e:
                  env_type = g
                  break
          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
  
  
  def get_default_network(env_type):
      if env_type in {'atari', 'retro'}:
          return 'cnn'
      else:
          return 'mlp'
  
  def get_alg_module(alg, submodule=None):
      submodule = submodule or alg
      try:
          # first try to import the alg module from baselines
          alg_module = import_module('.'.join(['baselines', alg, submodule]))
      except ImportError:
          # then from rl_algs
          alg_module = import_module('.'.join(['rl_' + 'algs', alg, submodule]))
  
      return alg_module
  
  
  def get_learn_function(alg):
      return get_alg_module(alg).learn
  
  
  def get_learn_function_defaults(alg, env_type):
      try:
          alg_defaults = get_alg_module(alg, 'defaults')
          kwargs = getattr(alg_defaults, env_type)()
      except (ImportError, AttributeError):
          kwargs = {}
      return kwargs
  
  
  def parse_cmdline_kwargs(args):
      '''
      convert a list of '='-spaced command-line arguments to a dictionary, evaluating python objects when possible
      '''
      def parse(v):
  
          assert isinstance(v, str)
          try:
              return eval(v)
          except (NameError, SyntaxError):
              return v
  
      return {k: parse(v) for k,v in parse_unknown_args(args).items()}
  
  
  
  def main(args):
      # configure logger, disable logging in child MPI processes (with rank > 0)
  
      arg_parser = common_arg_parser()
      args, unknown_args = arg_parser.parse_known_args(args)
      extra_args = parse_cmdline_kwargs(unknown_args)
  
      if MPI is None or MPI.COMM_WORLD.Get_rank() == 0:
          rank = 0
          logger.configure()
      else:
          logger.configure(format_strs=[])
          rank = MPI.COMM_WORLD.Get_rank()
  
      print("Called the trained function")
      model, env = train(args, extra_args)
  
      env.close()
  
      if args.save_path is not None and rank == 0:
          save_path = osp.expanduser(args.save_path)
          model.save(save_path)
  
      if args.play:
          logger.log("Running trained model")
          env ,ob_mean , ob_std = build_env(args)
          obs = env.reset()
          def initialize_placeholders(nlstm=128,**kwargs):
              return np.zeros((args.num_env or 1, 2*nlstm)), np.zeros((1))
          state, dones = initialize_placeholders(**extra_args)
          while True:
              actions, _, state, _ = model.step(obs,S=state, M=dones)
              obs, _, done, _ = env.step(actions)
              env.render()
              done = done.any() if isinstance(done, np.ndarray) else done
  
              if done:
                  obs = env.reset()
  
          env.close()
  
      return model
  
  if __name__ == '__main__':
      main(sys.argv)
  
apis: numpy.empty numpy.zeros numpy.asarray numpy.std tensorflow.ConfigProto numpy.mean