pend_npg_GAE.py

import torch
import torch.nn as nn
import numpy as np
import matplotlib.pyplot as plt
from actor import Actor
from algorithms import GAE, ActorOnlyMC, NPG, get_returns
from utils import sample_memory
import gym
from experiment_class import Experiment, mult_seed_exp
from time import time
import os

class Critic(nn.Module):
    def __init__(self, num_inputs, num_hidden):
        super(Critic, self).__init__()
        self.fc1 = nn.Linear(num_inputs, num_hidden)
        self.fc2 = nn.Linear(num_hidden, num_hidden)
        self.fc3 = nn.Linear(num_hidden, 1)
        
        # init weights
        gains = [np.sqrt(2), np.sqrt(2), 1]
        layers = [self.fc1, self.fc2, self.fc3]
        for i in range(len(layers)):
            nn.init.xavier_uniform_(layers[i].weight, gain=gains[i])
            layers[i].bias.data.fill_(0.01)
        
    def forward(self, x):
        x = nn.functional.relu(self.fc1(x))
        x = nn.functional.relu(self.fc2(x))
        v = self.fc3(x)
        return v



class ContActor(Actor):
    '''
    Actor for continous action space
    '''

    def __init__(self, num_input, num_hidden, **kwargs):
            super(ContActor, self).__init__(dist_type=torch.distributions.normal.Normal)
            self.fc1 = nn.Linear(num_input, num_hidden)
            self.fc2 = nn.Linear(num_hidden, num_hidden)
            self.fc3 = nn.Linear(num_hidden, 2)
            
            # init weights
            gains = [np.sqrt(2), np.sqrt(2), 1]
            layers = [self.fc1, self.fc2, self.fc3]
            for i in range(len(layers)):
                nn.init.xavier_uniform_(layers[i].weight, gain=gains[i])
                layers[i].bias.data.fill_(0.01)
            
            
            #self.log_var = nn.parameter.Parameter(data=torch.Tensor([0]), requires_grad=True)

    # override the parent class' forward method
    def forward(self, states):
        '''
        Computes parameters of gaussian distributions over actions given a tensor of states
        Input:
            states - pytorch tensor of shape (batch, *|S|)
        Output:
            params - pytorch tensor of shape (batch, *|params|)
        '''

        x = nn.functional.relu(self.fc1(states))
        x = nn.functional.relu(self.fc2(x))
        params = self.fc3(x)

        params[:,1] = torch.exp(params[:,1]) # make sure variance is positive
        #var = torch.exp(self.log_var)
        #params = torch.cat((params, torch.ones_like(params) * var), dim=1)

        return params

# run a single experiment

####
# dont change these params:
ac_kwargs = {'num_hidden':64}
critic_kwargs = {'num_hidden':64}
critic_optim_kwargs = {'lr':3e-4}
target_alg_kwargs = {'batch_size':16, 'epochs':3}
num_iters = 200
seeds = [0,1,2,3,4,5,6,7,8,9] # subject to change
log_dir = './pendulum/npg/GAE/'
try:
    os.mkdir(log_dir)
except:
    pass
try:
    os.mkdir(log_dir+'plots/')
except:
    pass

####

####
# do change these params:
target_alg_kwargs['gamma'] = 0.8
target_alg_kwargs['lamda'] = 0.97
ep_per_iter = 5
ac_alg_kwargs = {'lr':0.1}
lr_list = [0.1]
####

experiment_parameters =   {'seed':42, 
                           'env_str':'Pendulum-v0',
                           'ac':ContActor,
                           'ac_kwargs':ac_kwargs,
                           'ac_alg':NPG,
                           'ac_alg_kwargs':ac_alg_kwargs,
                           'target_alg':GAE,
                           'target_alg_kwargs':target_alg_kwargs,
                           'critic':Critic,
                           'critic_kwargs':critic_kwargs,
                           'critic_optim':torch.optim.Adam,
                           'critic_optim_kwargs':critic_optim_kwargs,
                           'num_iters':num_iters,
                           'ep_per_iter':ep_per_iter,
                           'log_file':'./pendulum/npg/GAE/single.npz'}




search_time = time()
max_return = -np.inf
lr_return_list = []
lr_return_std_list = []

for n in range(len(lr_list)):

    trial_time = time()

    # draw a sample from the loguniform distribution
    lr =  lr_list[n]

    print('\nNow running wiht lr = {}'.format(lr))
    
    # set up experiment
    ac_alg_kwargs['lr'] = lr
    experiment_parameters['ac_alg_kwargs'] = ac_alg_kwargs
    trial_log_dir = log_dir + 'lr_{0:1.4f}_gam_{1:1.2f}_lam_{2:1.2f}/'.format(lr, target_alg_kwargs['gamma'], target_alg_kwargs['lamda'])
    try:
        os.mkdir(trial_log_dir)
        os.mkdir(trial_log_dir+'plots/')
    except:
        pass

    #experiment = Experiment(**experiment_parameters)

    # run mutliple experiments with different seeds
    mult_exp = mult_seed_exp(experiment_parameters, seeds, trial_log_dir, save_all=False)
    mult_exp.run()

    try:
        mult_exp.plot(trial_log_dir+'plots/')
    except:
        pass

    # get results
    mean_results, std_results = mult_exp.get_mean_results()

    print('Mean return at termination with lr {0:1.5f} was {1:1.2f}'.format(lr, mean_results[-1,0]))

    lr_return_list.append(mean_results[-1,0])
    lr_return_std_list.append(std_results[-1,0])

    if mean_results[1,0]>max_return:
        print('Updating max return run..')
        max_return = mean_results[-1,0]

        mult_exp.plot(log_dir+ 'plots/best_run_')
    #mult_exp.plot(plot_path=log_dir+'plots/')

    print('This trial took {0:1.2f} seconds'.format((time()-trial_time)))

print('Hyperparameter search finished. Time: {0:1.0f} minutes and {1:1.2f} seconds'.format((time()-search_time)//60, (time()-search_time)))

print('Plotting Return - KL diagram')

plt.figure()
plt.errorbar(lr_list, lr_return_list, yerr=lr_return_std_list, fmt='o', capsize=6)
plt.xlabel('lr')
plt.ylabel('Log Final Return')
#plt.yscale('log')
plt.savefig(log_dir + 'plots/lr_return_plot.pdf')

# save results
np.savez_compressed(log_dir + 'kl_return_summary.npz', lr_list=lr_list, lr_return_list=lr_return_list, lr_return_std_list=lr_return_std_list)

#def run_loguniform_hp_search(low, high, num_trials, keyword, params):
#    '''
#    Runs a loguniform hyperparameter search of the param 'keyword'
#    '''