Added a single process environment runner. Also added an example for using such class.

This commit is contained in:
Brandon Rozek 2019-02-16 18:15:45 -05:00
parent 736e73a1f7
commit 14ba64d525
2 changed files with 179 additions and 0 deletions

View file

@ -0,0 +1,135 @@
import gym
import numpy as np
import torch
import torch.nn as nn
import torch.nn.functional as F
import rltorch
import rltorch.network as rn
import rltorch.memory as M
import rltorch.env as E
from rltorch.action_selector import ArgMaxSelector
from tensorboardX import SummaryWriter
import torch.multiprocessing as mp
class Value(nn.Module):
def __init__(self, state_size, action_size):
super(Value, self).__init__()
self.state_size = state_size
self.action_size = action_size
self.fc1 = rn.NoisyLinear(state_size, 255)
self.fc_norm = nn.LayerNorm(255)
self.value_fc = rn.NoisyLinear(255, 255)
self.value_fc_norm = nn.LayerNorm(255)
self.value = rn.NoisyLinear(255, 1)
self.advantage_fc = rn.NoisyLinear(255, 255)
self.advantage_fc_norm = nn.LayerNorm(255)
self.advantage = rn.NoisyLinear(255, action_size)
def forward(self, x):
x = F.relu(self.fc_norm(self.fc1(x)))
state_value = F.relu(self.value_fc_norm(self.value_fc(x)))
state_value = self.value(state_value)
advantage = F.relu(self.advantage_fc_norm(self.advantage_fc(x)))
advantage = self.advantage(advantage)
x = state_value + advantage - advantage.mean()
return x
config = {}
config['seed'] = 901
config['environment_name'] = 'Acrobot-v1'
config['memory_size'] = 2000
config['total_training_episodes'] = 5
config['total_evaluation_episodes'] = 2
config['batch_size'] = 32
config['learning_rate'] = 1e-3
config['target_sync_tau'] = 1e-1
config['discount_rate'] = 0.99
config['replay_skip'] = 0
# How many episodes between printing out the episode stats
config['print_stat_n_eps'] = 1
config['disable_cuda'] = False
# Prioritized vs Random Sampling
# 0 - Random sampling
# 1 - Only the highest prioirities
config['prioritized_replay_sampling_priority'] = 0.6
# How important are the weights for the loss?
# 0 - Treat all losses equally
# 1 - Lower the importance of high losses
# Should ideally start from 0 and move your way to 1 to prevent overfitting
config['prioritized_replay_weight_importance'] = rltorch.scheduler.ExponentialScheduler(initial_value = 0.4, end_value = 1, iterations = 5000)
def train(runner, agent, config, logger = None, logwriter = None):
finished = False
last_episode_num = 1
while not finished:
runner.run(config['replay_skip'] + 1)
agent.learn()
if logwriter is not None:
if last_episode_num < runner.episode_num:
last_episode_num = runner.episode_num
agent.net.log_named_parameters()
logwriter.write(logger)
finished = runner.episode_num > config['total_training_episodes']
if __name__ == "__main__":
torch.multiprocessing.set_sharing_strategy('file_system') # To not hit file descriptor memory limit
# Setting up the environment
rltorch.set_seed(config['seed'])
print("Setting up environment...", end = " ")
env = E.TorchWrap(gym.make(config['environment_name']))
env.seed(config['seed'])
print("Done.")
state_size = env.observation_space.shape[0]
action_size = env.action_space.n
# Logging
logger = rltorch.log.Logger()
# logwriter = rltorch.log.LogWriter(logger, SummaryWriter())
logwriter = rltorch.log.LogWriter(SummaryWriter())
# Setting up the networks
device = torch.device("cuda:0" if torch.cuda.is_available() and not config['disable_cuda'] else "cpu")
net = rn.Network(Value(state_size, action_size),
torch.optim.Adam, config, device = device, name = "DQN", logger = logger)
target_net = rn.TargetNetwork(net, device = device)
net.model.share_memory()
target_net.model.share_memory()
# Actor takes a net and uses it to produce actions from given states
actor = ArgMaxSelector(net, action_size, device = device)
# Memory stores experiences for later training
memory = M.PrioritizedReplayMemory(capacity = config['memory_size'], alpha = config['prioritized_replay_sampling_priority'])
# memory = M.ReplayMemory(capacity = config['memory_size'])
# Runner performs a certain number of steps in the environment
runner = rltorch.env.EnvironmentRunSync(env, actor, config, name = "Training", memory = memory, logwriter = logwriter)
# Agent is what performs the training
agent = rltorch.agents.DQNAgent(net, memory, config, target_net = target_net, logger = logger)
print("Training...")
train(runner, agent, config, logger = logger, logwriter = logwriter)
# For profiling...
# import cProfile
# cProfile.run('train(runner, agent, config, logger = logger, logwriter = logwriter )')
# python -m torch.utils.bottleneck /path/to/source/script.py [args] is also a good solution...
print("Training Finished.")
print("Evaluating...")
rltorch.env.simulateEnvEps(env, actor, config, total_episodes = config['total_evaluation_episodes'], logger = logger, name = "Evaluation")
print("Evaulations Done.")
logwriter.close() # We don't need to write anything out to disk anymore

View file

@ -1,3 +1,6 @@
from copy import deepcopy
import rltorch
def simulateEnvEps(env, actor, config, total_episodes = 1, memory = None, logger = None, name = ""):
for episode in range(total_episodes):
state = env.reset()
@ -19,3 +22,44 @@ def simulateEnvEps(env, actor, config, total_episodes = 1, memory = None, logger
if logger is not None:
logger.append(name + '/EpisodeReward', episode_reward)
class EnvironmentRunSync():
def __init__(self, env, actor, config, memory = None, logwriter = None, name = ""):
self.env = env
self.name = name
self.actor = actor
self.config = deepcopy(config)
self.logwriter = logwriter
self.memory = memory
self.episode_num = 1
self.episode_reward = 0
self.last_state = env.reset()
def run(self, iterations):
state = self.last_state
logger = rltorch.log.Logger() if self.logwriter is not None else None
for _ in range(iterations):
action = self.actor.act(state)
next_state, reward, done, _ = self.env.step(action)
self.episode_reward += reward
if self.memory is not None:
self.memory.append(state, action, reward, next_state, done)
state = next_state
if done:
if self.episode_num % self.config['print_stat_n_eps'] == 0:
print("episode: {}/{}, score: {}"
.format(self.episode_num, self.config['total_training_episodes'], self.episode_reward))
if self.logwriter is not None:
logger.append(self.name + '/EpisodeReward', self.episode_reward)
self.episode_reward = 0
state = self.env.reset()
self.episode_num += 1
if self.logwriter is not None:
self.logwriter.write(logger)
self.last_state = state