Cleaned up scripts, added more comments

examples/acrobot_es.py

@@ -1,5 +1,4 @@
 import gym
-import numpy as np
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
@@ -10,8 +9,10 @@ import rltorch.memory as M
 import rltorch.env as E
 from rltorch.action_selector import StochasticSelector
 from tensorboardX import SummaryWriter
-import torch.multiprocessing as mp
 
+#
+## Networks
+#
 class Policy(nn.Module):
   def __init__(self, state_size, action_size):
     super(Policy, self).__init__()
@@ -32,37 +33,37 @@ class Policy(nn.Module):
     x = F.softmax(self.action_prob(x), dim = 1)
     return x
 
+#
+## Configuration
+#
 config = {}
 config['seed'] = 901
 config['environment_name'] = 'Acrobot-v1'
-config['memory_size'] = 2000
 config['total_training_episodes'] = 50
 config['total_evaluation_episodes'] = 5
-config['batch_size'] = 32
 config['learning_rate'] = 1e-1
-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
 
-
-
-def train(env, net, actor, config, logger = None, logwriter = None):
+#
+## Training Loop
+#
+def train(runner, net, config, logger = None, logwriter = None):
   finished = False
-  episode_num = 1
   while not finished:
-    rltorch.env.simulateEnvEps(env, actor, config, logger = logger, name = "Training")
-    episode_num += 1
+    runner.run()
     net.calc_gradients()
     net.step()
-    # When the episode number changes, log network paramters
     if logwriter is not None:
       net.log_named_parameters()
       logwriter.write(logger)
-    finished = episode_num > config['total_training_episodes']
+    finished = runner.episode_num > config['total_training_episodes']
  
+#
+## Loss function
+#
 def fitness(model):
   env = gym.make("Acrobot-v1")
   state = torch.from_numpy(env.reset()).float().unsqueeze(0)
@@ -75,9 +76,12 @@ def fitness(model):
     next_state, reward, done, _ = env.step(action)
     total_reward += reward
     state = torch.from_numpy(next_state).float().unsqueeze(0)
-  return total_reward
+  return -total_reward
 
 if __name__ == "__main__":
+  # Hide internal gym warnings
+  gym.logger.set_level(40)
+
   # Setting up the environment
   rltorch.set_seed(config['seed'])
   print("Setting up environment...", end = " ")
@@ -90,21 +94,21 @@ if __name__ == "__main__":
 
   # 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.ESNetwork(Policy(state_size, action_size), 
                       torch.optim.Adam, 100, fitness, config, device = device, name = "ES", logger = logger)
-  net.model.share_memory()
 
   # Actor takes a net and uses it to produce actions from given states
   actor = StochasticSelector(net, action_size, device = device)
 
-  print("Training...")
+  # Runner performs an episode of the environment
+  runner = rltorch.env.EnvironmentEpisodeSync(env, actor, config, name = "Training", logwriter = logwriter)
 
-  train(env, net, actor, config, logger = logger, logwriter = logwriter) 
+  print("Training...")
+  train(runner, net, config, logger = logger, logwriter = logwriter) 
 
   # For profiling...
   # import cProfile