from rlzoo.common.env_wrappers import build_env
from rlzoo.common.utils import call_default_params
from rlzoo.algorithms import *
# choose an algorithm
AlgName = 'TD3'
# chose an environment
EnvName = 'Pendulum-v0'  
# select a corresponding environment type
EnvType = ['classic_control', 'atari', 'box2d', 'mujoco', 'robotics', 'dm_control', 'rlbench'][0] 
# build an environment with wrappers
env = build_env(EnvName, EnvType)  
# call default parameters for the algorithm and learning process
alg_params, learn_params = call_default_params(env, EnvType, AlgName)  
# instantiate the algorithm
alg = eval(AlgName+'(**alg_params)')
# start the training
alg.learn(env=env, mode='train', render=False, **learn_params)  
# test after training 
alg.learn(env=env, mode='test', render=True, **learn_params)  

# in the root folder of rlzoo package
cd rlzoo
python run_rlzoo.py

import gym
from rlzoo.common.utils import make_env, set_seed
from rlzoo.algorithms import AC
from rlzoo.common.value_networks import ValueNetwork
from rlzoo.common.policy_networks import StochasticPolicyNetwork

''' load environment '''
env = gym.make('CartPole-v0').unwrapped
obs_space = env.observation_space
act_space = env.action_space
# reproducible
seed = 2
set_seed(seed, env)

''' build networks for the algorithm '''
num_hidden_layer = 4 #number of hidden layers for the networks
hidden_dim = 64 # dimension of hidden layers for the networks
with tf.name_scope('AC'):
        with tf.name_scope('Critic'):
            	# choose the critic network, can be replaced with customized network
                critic = ValueNetwork(obs_space, hidden_dim_list=num_hidden_layer * [hidden_dim])
        with tf.name_scope('Actor'):
            	# choose the actor network, can be replaced with customized network
                actor = StochasticPolicyNetwork(obs_space, act_space, hidden_dim_list=num_hidden_layer * [hidden_dim], output_activation=tf.nn.tanh)
net_list = [actor, critic] # list of the networks

''' choose optimizers '''
a_lr, c_lr = 1e-4, 1e-2  # a_lr: learning rate of the actor; c_lr: learning rate of the critic
a_optimizer = tf.optimizers.Adam(a_lr)
c_optimizer = tf.optimizers.Adam(c_lr)
optimizers_list=[a_optimizer, c_optimizer]  # list of optimizers

# intialize the algorithm model, with algorithm parameters passed in
model = AC(net_list, optimizers_list)
''' 
full list of arguments for the algorithm
----------------------------------------
net_list: a list of networks (value and policy) used in the algorithm, from common functions or customization
optimizers_list: a list of optimizers for all networks and differentiable variables
gamma: discounted factor of reward
action_range: scale of action values
'''

# start the training process, with learning parameters passed in
model.learn(env, train_episodes=500,  max_steps=200,
            save_interval=50, mode='train', render=False)
''' 
full list of parameters for training
---------------------------------------
env: learning environment
train_episodes:  total number of episodes for training
test_episodes:  total number of episodes for testing
max_steps:  maximum number of steps for one episode
save_interval: time steps for saving the weights and plotting the results
mode: 'train' or 'test'
render:  if true, visualize the environment
'''

# test after training
model.learn(env, test_episodes=100, max_steps=200,  mode='test', render=True)

# in the root folder of rlzoo package
cd rlzoo
python algorithms/<ALGORITHM_NAME>/run_<ALGORITHM_NAME>.py 
# for example: run actor-critic
python algorithms/ac/run_ac.py

from rlzoo.common.env_wrappers import build_env
from rlzoo.common.policy_networks import *
from rlzoo.common.value_networks import *
from rlzoo.algorithms.dppo_clip_distributed.dppo_clip import DPPO_CLIP
from functools import partial

# Specify the training configurations
training_conf = {
    'total_step': int(1e7),  # overall training timesteps
    'traj_len': 200,         # length of the rollout trajectory
    'train_n_traj': 2,       # update the models after every certain number of trajectories for each learner 
    'save_interval': 10,     # saving the models after every certain number of updates
}

# Specify the environment and launch it
env_name, env_type = 'CartPole-v0', 'classic_control'
env_maker = partial(build_env, env_name, env_type)
temp_env = env_maker()
obs_shape, act_shape = temp_env.observation_space.shape, temp_env.action_space.shape

env_conf = {
    'env_name': env_name,
    'env_type': env_type,
    'env_maker': env_maker,
    'obs_shape': obs_shape,
    'act_shape': act_shape,
}


def build_network(observation_space, action_space, name='DPPO_CLIP'):
    """ build networks for the algorithm """
    hidden_dim = 256
    num_hidden_layer = 2
    critic = ValueNetwork(observation_space, [hidden_dim] * num_hidden_layer, name=name + '_value')

    actor = StochasticPolicyNetwork(observation_space, action_space,
                                    [hidden_dim] * num_hidden_layer,
                                    trainable=True,
                                    name=name + '_policy')
    return critic, actor


def build_opt(actor_lr=1e-4, critic_lr=2e-4):
    """ choose the optimizer for learning """
    import tensorflow as tf
    return [tf.optimizers.Adam(critic_lr), tf.optimizers.Adam(actor_lr)]


net_builder = partial(build_network, temp_env.observation_space, temp_env.action_space)
opt_builder = partial(build_opt, )

agent_conf = {
    'net_builder': net_builder,
    'opt_builder': opt_builder,
    'agent_generator': partial(DPPO_CLIP, net_builder, opt_builder),
}
del temp_env

from rlzoo.distributed.start_dis_role import main

print('Start Training.')
main(training_conf, env_conf, agent_conf)
print('Training Finished.')