text: """Runner script for single and multi-agent reinforcement learning experiments.
  
  This script performs an RL experiment using the PPO algorithm. Choice of
  hyperparameters can be seen and adjusted from the code below.
  
  Usage
  	python train.py EXP_CONFIG
  """
  import argparse
  import json
  import os
  import sys
  from time import strftime
  from copy import deepcopy
  
  from flow.core.util import ensure_dir
  from flow.utils.registry import env_constructor
  from flow.utils.rllib import FlowParamsEncoder, get_flow_params
  from flow.utils.registry import make_create_env
  
  import numpy as np
  import wandb
  
  os.environ.setdefault("TUNE_GLOBAL_CHECKPOINT_S", str(sys.maxsize))
  
  # Callbacks
  # Custom state can be stored for the episode in the info["episode"].user_data dict
  # Custom scalar metrics reported by saving values to the info["episode"].custom_metrics dict
  def on_episode_start(info):
  	episode = info["episode"]
  	episode.user_data["true_reward"] = []
  
  def on_episode_step(info):
  	episode = info["episode"]
  	env = info["env"]
  
  	kernel = env.vector_env.envs[0].true_reward
  	# vel = np.array([
  	# 		kernel.vehicle.get_speed(veh_id)
  	# 		for veh_id in kernel.vehicle.get_ids()
  	# 	])
  
  	# # reward average velocity
  	# eta_2 = 4.
  	# true_reward = eta_2 * np.mean(vel) / 20
  
  	# # punish accelerations (should lead to reduced stop-and-go waves)
  	# eta = 4  # 0.25
  	# mean_actions = np.mean(np.abs(np.array(episode.last_action_for())))
  	# accel_threshold = 0
  
  	# if mean_actions > accel_threshold:
  	# 	true_reward += eta * (accel_threshold - mean_actions)
  	episode.user_data["true_reward"].append(true_reward)
  
  
  def on_episode_step_multi(info):
  	episode = info["episode"]
  	env = info["env"]
  
  	true_reward = env.envs[0].true_reward
  	# vel = np.array([
  	# 		kernel.vehicle.get_speed(veh_id)
  	# 		for veh_id in kernel.vehicle.get_ids()
  	# 	])
  
  	# # reward average velocity
  	# eta_2 = 4.
  	# true_reward = eta_2 * np.mean(vel) / 20
  
  	# # punish accelerations (should lead to reduced stop-and-go waves)
  	# eta = 4  # 0.25
  	# mean_actions = np.mean(np.abs(np.array(episode.last_action_for())))
  	# accel_threshold = 0
  
  	# if mean_actions > accel_threshold:
  	# 	true_reward += eta * (accel_threshold - mean_actions)
  	episode.user_data["true_reward"].append(true_reward)
  
  def on_episode_end(info):
  	episode = info["episode"]
  	mean_rew = np.sum(episode.user_data["true_reward"])
  	print(mean_rew)
  	episode.custom_metrics["true_reward"] = mean_rew
  
  def on_train_result(info):
  	pass
  
  def on_postprocess_traj(info):
  	pass
  
  
  def parse_args(args):
  	"""Parse training options user can specify in command line.
  
  	Returns
  	-------
  	argparse.Namespace
  		the output parser object
  	"""
  	parser = argparse.ArgumentParser(
  		formatter_class=argparse.RawDescriptionHelpFormatter,
  		description="Parse argument used when running a Flow simulation.",
  		epilog="python train.py EXP_CONFIG")
  
  	# required input parameters
  	parser.add_argument(
  		'exp_config', type=str,
  		help='Name of the experiment configuration file, as located in '
  			 'exp_configs/rl/singleagent or exp_configs/rl/multiagent.')
  	parser.add_argument(
  		'--multi', action='store_true', help='Run multiagent experiment')
  	parser.add_argument(
  		'--test', action='store_true', help='No wandb')
  
  	# optional input parameters
  	parser.add_argument(
  		'--rl_trainer', type=str, default="rllib",
  		help='the RL trainer to use. either rllib or Stable-Baselines')
  
  	parser.add_argument(
  		'--num_cpus', type=int, default=1,
  		help='How many CPUs to use')
  	parser.add_argument(
  		'--num_steps', type=int, default=5000,
  		help='How many total steps to perform learning over')
  	parser.add_argument(
  		'--rollout_size', type=int, default=20,
  		help='How many steps are in a training batch.')
  	parser.add_argument(
  		'--horizon', type=int, default=400,
  		help='Number of steps in each episode.')
  	parser.add_argument(
  		'--checkpoint', type=int, default=50,
  		help='How frequently to checkpoint model.')
  	parser.add_argument(
  		'--checkpoint_path', type=str, default=None,
  		help='Directory with checkpoint to restore training from.')
  
  	parser.add_argument('-p', '--metric', type=str, default='reward')
  	parser.add_argument('-g', '--goal', type=str, default='maximize')
  	parser.add_argument('-n', '--name', type=str, default='untitled')
  	parser.add_argument('-gpu', '--gpuid', type=int, default=0)
  	parser.add_argument('--sweep', action='store_true', help='run hparam sweep')
  	parser.add_argument('--replicas', type=bool, default=False, 
  		help='number of different hparam settings to try')
  
  	parser.add_argument('--optim', type=str, default='adam')
  	parser.add_argument('--batch_size', type=int, default=128)
  	parser.add_argument('--lr', type=float, default=5e-5)
  	parser.add_argument('--epochs', type=int, default=1000)
  	parser.add_argument('--d_model', type=int, default=128)
  	parser.add_argument('--dropout', type=float, default=0.2)
  
  	return parser.parse_known_args(args)[0]
  
  
  def run_model_stablebaseline(flow_params,
  							 num_cpus=1,
  							 rollout_size=50,
  							 num_steps=50):
  	"""Run the model for num_steps if provided.
  
  	Parameters
  	----------
  	flow_params : dict
  		flow-specific parameters
  	num_cpus : int
  		number of CPUs used during training
  	rollout_size : int
  		length of a single rollout
  	num_steps : int
  		total number of training steps
  	The total rollout length is rollout_size.
  
  	Returns
  	-------
  	stable_baselines.*
  		the trained model
  	"""
  	from stable_baselines.common.vec_env import DummyVecEnv, SubprocVecEnv
  	from stable_baselines import PPO2
  
  	if num_cpus == 1:
  		constructor = env_constructor(params=flow_params, version=0)()
  		# The algorithms require a vectorized environment to run
  		env = DummyVecEnv([lambda: constructor])
  	else:
  		env = SubprocVecEnv([env_constructor(params=flow_params, version=i)
  							 for i in range(num_cpus)])
  
  	train_model = PPO2('MlpPolicy', env, verbose=1, n_steps=rollout_size)
  	train_model.learn(total_timesteps=num_steps)
  	return train_model
  
  
  def setup_exps_rllib(flow_params,
  					 n_cpus,
  					 n_rollouts,
  					 policy_graphs=None,
  					 policy_mapping_fn=None,
  					 policies_to_train=None):
  	"""Return the relevant components of an RLlib experiment.
  
  	Parameters
  	----------
  	flow_params : dict
  		flow-specific parameters (see flow/utils/registry.py)
  	n_cpus : int
  		number of CPUs to run the experiment over
  	n_rollouts : int
  		number of rollouts per training iteration
  	policy_graphs : dict, optional
  		TODO
  	policy_mapping_fn : function, optional
  		TODO
  	policies_to_train : list of str, optional
  		TODO
  
  	Returns
  	-------
  	str
  		name of the training algorithm
  	str
  		name of the gym environment to be trained
  	dict
  		training configuration parameters
  	"""
  	from ray import tune
  	from ray.tune.registry import register_env
  	try:
  		from ray.rllib.agents.agent import get_agent_class
  	except ImportError:
  		from ray.rllib.agents.registry import get_agent_class
  	
  	cfg = wandb.config
  	horizon = flow_params['env'].horizon
  	print(f"Horizon: {horizon}")
  	alg_run = "PPO"
  
  	agent_cls = get_agent_class(alg_run)
  	config = deepcopy(agent_cls._default_config)
  	
  	config["seed"] = 17
  	
  	config["num_workers"] = n_cpus - 1
  	config["train_batch_size"] = horizon * n_rollouts
  	config["sgd_minibatch_size"] = min(16 * 1024, config["train_batch_size"])
  	config["gamma"] = 0.999  # discount rate
  	fcnet = [int(sys.argv[4])] * int(sys.argv[5]) #[32, 32, 32]
  	config["model"].update({"fcnet_hiddens": fcnet})
  	config["use_gae"] = True
  	config["lambda"] = 0.97
  	config["kl_target"] = 0.02
  	config["vf_clip_param"] = 10000
  	config["num_sgd_iter"] = 10
  	config["horizon"] = horizon
  	
  	# save the flow params for replay
  	flow_json = json.dumps(
  		flow_params, cls=FlowParamsEncoder, sort_keys=True, indent=4)
  	config['env_config']['flow_params'] = flow_json
  	config['env_config']['run'] = alg_run
  
  	# multiagent configuration
  	if policy_graphs is not None:
  		print("policy_graphs", policy_graphs)
  		config['multiagent'].update({'policies': policy_graphs})
  	if policy_mapping_fn is not None:
  		config['multiagent'].update(
  			{'policy_mapping_fn': tune.function(policy_mapping_fn)})
  	if policies_to_train is not None:
  		config['multiagent'].update({'policies_to_train': policies_to_train})
  
  	config['callbacks'] = {
  					"on_episode_start": on_episode_start,
  					"on_episode_step": on_episode_step_multi if cfg.multi else on_episode_step,
  					"on_episode_end": on_episode_end,
  				}
  
  	create_env, gym_name = make_create_env(params=flow_params)
  
  	# Register as rllib env
  	register_env(gym_name, create_env)
  	return alg_run, gym_name, config
  
  
  def train_rllib(submodule, flags):
  	"""Train policies using the PPO algorithm in RLlib."""
  	import ray
  	from ray.tune import run_experiments
  
  	flow_params = submodule.flow_params
  	flow_params["exp_tag"] = sys.argv[2]
  	flow_params["env"].additional_params["eta"] = float(sys.argv[3])
  	flow_params["env"].horizon = flags.horizon
  	n_cpus = int(sys.argv[6])
  	n_rollouts = flags.rollout_size
  	policy_graphs = getattr(submodule, "POLICY_GRAPHS", None)
  	policy_mapping_fn = getattr(submodule, "policy_mapping_fn", None)
  	policies_to_train = getattr(submodule, "policies_to_train", None)
  
  	alg_run, gym_name, config = setup_exps_rllib(
  		flow_params, n_cpus, n_rollouts,
  		policy_graphs, policy_mapping_fn, policies_to_train)
  
  	ray.init(address=os.environ["ip_head"])
  	print(f"Epochs: {flags.num_steps}")
  	exp_config = {
  		"run": alg_run,
  		"env": gym_name,
  		"config": {
  			**config
  		},
  		"checkpoint_freq": flags.checkpoint,
  		"checkpoint_at_end": True,
  		"max_failures": 999,
  		"stop": {
  			"training_iteration": flags.num_steps,
  		},
  	}
  
  	if flags.checkpoint_path is not None:
  		exp_config['restore'] = flags.checkpoint_path
  	run_experiments({flow_params["exp_tag"]: exp_config})
  
  
  def train_h_baselines(env_name, args, multiagent):
  	"""Train policies using SAC and TD3 with h-baselines."""
  	from hbaselines.algorithms import OffPolicyRLAlgorithm
  	from hbaselines.utils.train import parse_options, get_hyperparameters
  
  	# Get the command-line arguments that are relevant here
  	args = parse_options(description="", example_usage="", args=args)
  
  	# the base directory that the logged data will be stored in
  	base_dir = "training_data"
  
  	for i in range(args.n_training):
  		# value of the next seed
  		seed = args.seed + i
  
  		# The time when the current experiment started.
  		now = strftime("%Y-%m-%d-%H:%M:%S")
  
  		# Create a save directory folder (if it doesn't exist).
  		dir_name = os.path.join(base_dir, '{}/{}'.format(args.env_name, now))
  		ensure_dir(dir_name)
  
  		# Get the policy class.
  		if args.alg == "TD3":
  			if multiagent:
  				from hbaselines.multi_fcnet.td3 import MultiFeedForwardPolicy
  				policy = MultiFeedForwardPolicy
  			else:
  				from hbaselines.fcnet.td3 import FeedForwardPolicy
  				policy = FeedForwardPolicy
  		elif args.alg == "SAC":
  			if multiagent:
  				from hbaselines.multi_fcnet.sac import MultiFeedForwardPolicy
  				policy = MultiFeedForwardPolicy
  			else:
  				from hbaselines.fcnet.sac import FeedForwardPolicy
  				policy = FeedForwardPolicy
  		else:
  			raise ValueError("Unknown algorithm: {}".format(args.alg))
  
  		# Get the hyperparameters.
  		hp = get_hyperparameters(args, policy)
  
  		# Add the seed for logging purposes.
  		params_with_extra = hp.copy()
  		params_with_extra['seed'] = seed
  		params_with_extra['env_name'] = args.env_name
  		params_with_extra['policy_name'] = policy.__name__
  		params_with_extra['algorithm'] = args.alg
  		params_with_extra['date/time'] = now
  
  		# Add the hyperparameters to the folder.
  		with open(os.path.join(dir_name, 'hyperparameters.json'), 'w') as f:
  			json.dump(params_with_extra, f, sort_keys=True, indent=4)
  
  		# Create the algorithm object.
  		alg = OffPolicyRLAlgorithm(
  			policy=policy,
  			env="flow:{}".format(env_name),
  			eval_env="flow:{}".format(env_name) if args.evaluate else None,
  			**hp
  		)
  
  		# Perform training.
  		alg.learn(
  			total_steps=args.total_steps,
  			log_dir=dir_name,
  			log_interval=args.log_interval,
  			eval_interval=args.eval_interval,
  			save_interval=args.save_interval,
  			initial_exploration_steps=args.initial_exploration_steps,
  			seed=seed,
  		)
  
  
  def train_stable_baselines(submodule, flags):
  	"""Train policies using the PPO algorithm in stable-baselines."""
  	from stable_baselines.common.vec_env import DummyVecEnv
  	from stable_baselines import PPO2
  
  	flow_params = submodule.flow_params
  	# Path to the saved files
  	exp_tag = flow_params['exp_tag']
  	result_name = '{}/{}'.format(exp_tag, strftime("%Y-%m-%d-%H:%M:%S"))
  
  	# Perform training.
  	print('Beginning training.')
  	model = run_model_stablebaseline(
  		flow_params, flags.num_cpus, flags.rollout_size, flags.num_steps)
  
  	# Save the model to a desired folder and then delete it to demonstrate
  	# loading.
  	print('Saving the trained model!')
  	path = os.path.realpath(os.path.expanduser('~/baseline_results'))
  	ensure_dir(path)
  	save_path = os.path.join(path, result_name)
  	model.save(save_path)
  
  	# dump the flow params
  	with open(os.path.join(path, result_name) + '.json', 'w') as outfile:
  		json.dump(flow_params, outfile,
  				  cls=FlowParamsEncoder, sort_keys=True, indent=4)
  
  	# Replay the result by loading the model
  	print('Loading the trained model and testing it out!')
  	model = PPO2.load(save_path)
  	flow_params = get_flow_params(os.path.join(path, result_name) + '.json')
  	flow_params['sim'].render = True
  	env = env_constructor(params=flow_params, version=0)()
  	# The algorithms require a vectorized environment to run
  	eval_env = DummyVecEnv([lambda: env])
  	obs = eval_env.reset()
  	reward = 0
  	for _ in range(flow_params['env'].horizon):
  		action, _states = model.predict(obs)
  		obs, rewards, dones, info = eval_env.step(action)
  		reward += rewards
  	print('the final reward is {}'.format(reward))
  
  
  def main(args):
  	"""Perform the training operations."""
  	# Parse script-level arguments (not including package arguments).
  	flags = parse_args(args)
  
  	# Import the sub-module containing the specified exp_config and determine
  	# whether the environment is single agent or multi-agent.
  	if flags.multi:
  		assert flags.rl_trainer.lower() in ["rllib", "h-baselines"], \
  			"Currently, multiagent experiments are only supported through "\
  			"RLlib. Try running this experiment using RLlib: " \
  			"'python train.py EXP_CONFIG'"
  
  	if flags.sweep:
  		flags.__dict__.update(SWEEP_CONFIG)
  		config = flags.__dict__
  		sweep_id = wandb.sweep(config, entity="aypan17", project="value-learning", group="traffic", sync_tensorboard=True)
  		wandb.agent(sweep_id, train, count=flags.replicas)
  	else:
  		flags.__dict__.update(DEFAULT_CONFIG)
  		config = flags.__dict__
  		wandb.init(entity="aypan17", project="value-learning", group="traffic", config=config, sync_tensorboard=True)
  		train()
  
  def train():
  	config = wandb.config
  
  	# Import relevant information from the exp_config script.
  	if config.multi:
  		module = __import__(
  			"flow_cfg.exp_configs.rl.multiagent", fromlist=[config.exp_config])
  	else:
  		module = __import__(
  			"flow_cfg.exp_configs.rl.singleagent", fromlist=[config.exp_config])
  
  	submodule = getattr(module, config.exp_config)
  
  	# Perform the training operation.
  	if config.rl_trainer.lower() == "rllib":
  		train_rllib(submodule, config)
  	elif config.rl_trainer.lower() == "stable-baselines":
  		train_stable_baselines(submodule, config)
  	elif config.rl_trainer.lower() == "h-baselines":
  		train_h_baselines(config.exp_config, args, config.multi)
  	else:
  		raise ValueError("rl_trainer should be either 'rllib', 'h-baselines', "
  						 "or 'stable-baselines'.")
  
  SWEEP_CONFIG = {
  			'method': 'random', #'grid'
  			'parameters': {
  				'epochs': {
  					'values': [2, 5, 10]
  				},
  				'batch_size': {
  					'values': [256, 128, 64, 32]
  				},
  				'dropout': {
  					'values': [0.3, 0.4, 0.5]
  				},
  				'lr': {
  					'values': [1e-2, 1e-3, 1e-4, 3e-4, 3e-5, 1e-5]
  				},
  				'd_model':{
  					'values':[128,256,512]
  				},
  				'optimizer': {
  					'values': ['adam', 'sgd']
  				},
  			}
  		}
  
  DEFAULT_CONFIG = {}
  
  if __name__ == "__main__":
  	print(sys.argv)
  	np.seterr(all='raise')	
  	main(sys.argv[1:])
  