[rllib, WIP, DO NOT MERGE] TRPO PyTorch implementation

python/ray/rllib/__init__.py

@@ -9,7 +9,7 @@
 
 def _register_all():
     for key in ["PPO", "ES", "DQN", "APEX", "A3C", "BC", "PG", "DDPG",
-                "DDPG2", "APEX_DDPG", "__fake", "__sigmoid_fake_data",
+                "DDPG2", "APEX_DDPG", "TRPO", "__fake", "__sigmoid_fake_data",
                 "__parameter_tuning"]:
         from ray.rllib.agent import get_agent_class
         register_trainable(key, get_agent_class(key))

python/ray/rllib/a3c/a3c.py

@@ -1,32 +1,29 @@
-from __future__ import absolute_import
-from __future__ import division
-from __future__ import print_function
+from __future__ import absolute_import, division, print_function
 
-import numpy as np
-import pickle
 import os
+import pickle
+
+import numpy as np
 
 import ray
+from ray.rllib.a3c.a3c_evaluator import (A3CEvaluator, GPURemoteA3CEvaluator,
+                                         RemoteA3CEvaluator,)
 from ray.rllib.agent import Agent
 from ray.rllib.optimizers import AsyncOptimizer
 from ray.rllib.utils import FilterManager
-from ray.rllib.a3c.a3c_evaluator import A3CEvaluator, RemoteA3CEvaluator, \
-    GPURemoteA3CEvaluator
 from ray.tune.result import TrainingResult
 from ray.tune.trial import Resources
 
 DEFAULT_CONFIG = {
     # Number of workers (excluding master)
     "num_workers": 4,
-    # Size of rollout batch
+    # Size of rollout
     "batch_size": 10,
-    # Use LSTM model - only applicable for image states
+    # Only applicable for image states
     "use_lstm": False,
-    # Use PyTorch as backend - no LSTM support
+    # No LSTM support if PyTorch used
     "use_pytorch": False,
-    # Which observation filter to apply to the observation
     "observation_filter": "NoFilter",
-    # Which reward filter to apply to the reward
     "reward_filter": "NoFilter",
     # Discount factor of MDP
     "gamma": 0.99,
@@ -36,9 +33,7 @@
     "grad_clip": 40.0,
     # Learning rate
     "lr": 0.0001,
-    # Value Function Loss coefficient
     "vf_loss_coeff": 0.5,
-    # Entropy coefficient
     "entropy_coeff": -0.01,
     # Whether to place workers on GPUs
     "use_gpu_for_workers": False,
@@ -84,15 +79,18 @@ def _init(self):
             self.config,
             self.logdir,
             start_sampler=False)
+
         if self.config["use_gpu_for_workers"]:
             remote_cls = GPURemoteA3CEvaluator
         else:
             remote_cls = RemoteA3CEvaluator
+
         self.remote_evaluators = [
             remote_cls.remote(self.registry, self.env_creator, self.config,
                               self.logdir)
             for i in range(self.config["num_workers"])
         ]
+
         self.optimizer = AsyncOptimizer(self.config["optimizer"],
                                         self.local_evaluator,
                                         self.remote_evaluators)
@@ -101,20 +99,23 @@ def _train(self):
         self.optimizer.step()
         FilterManager.synchronize(self.local_evaluator.filters,
                                   self.remote_evaluators)
-        res = self._fetch_metrics_from_remote_evaluators()
-        return res
+        result = self._fetch_metrics_from_remote_evaluators()
+        return result
 
     def _fetch_metrics_from_remote_evaluators(self):
         episode_rewards = []
         episode_lengths = []
+
         metric_lists = [
             a.get_completed_rollout_metrics.remote()
             for a in self.remote_evaluators
         ]
+
         for metrics in metric_lists:
             for episode in ray.get(metrics):
                 episode_lengths.append(episode.episode_length)
                 episode_rewards.append(episode.episode_reward)
+
         avg_reward = (np.mean(episode_rewards)
                       if episode_rewards else float('nan'))
         avg_length = (np.mean(episode_lengths)
@@ -137,21 +138,27 @@ def _stop(self):
     def _save(self, checkpoint_dir):
         checkpoint_path = os.path.join(checkpoint_dir,
                                        "checkpoint-{}".format(self.iteration))
+
         agent_state = ray.get(
             [a.save.remote() for a in self.remote_evaluators])
+
         extra_data = {
             "remote_state": agent_state,
             "local_state": self.local_evaluator.save()
         }
+
         pickle.dump(extra_data, open(checkpoint_path + ".extra_data", "wb"))
+
         return checkpoint_path
 
     def _restore(self, checkpoint_path):
         extra_data = pickle.load(open(checkpoint_path + ".extra_data", "rb"))
+
         ray.get([
             a.restore.remote(o)
             for a, o in zip(self.remote_evaluators, extra_data["remote_state"])
         ])
+
         self.local_evaluator.restore(extra_data["local_state"])
 
     def compute_action(self, observation):

python/ray/rllib/a3c/a3c_evaluator.py

@@ -26,23 +26,31 @@ class A3CEvaluator(PolicyEvaluator):
             rollouts.
         logdir: Directory for logging.
     """
-    def __init__(
-            self, registry, env_creator, config, logdir, start_sampler=True):
+
+    def __init__(self,
+                 registry,
+                 env_creator,
+                 config,
+                 logdir,
+                 start_sampler=True):
         env = ModelCatalog.get_preprocessor_as_wrapper(
             registry, env_creator(config["env_config"]), config["model"])
         self.env = env
         policy_cls = get_policy_cls(config)
         # TODO(rliaw): should change this to be just env.observation_space
-        self.policy = policy_cls(
-            registry, env.observation_space.shape, env.action_space, config)
+        self.policy = policy_cls(registry, env.observation_space.shape,
+                                 env.action_space, config)
         self.config = config
 
         # Technically not needed when not remote
-        self.obs_filter = get_filter(
-            config["observation_filter"], env.observation_space.shape)
+        self.obs_filter = get_filter(config["observation_filter"],
+                                     env.observation_space.shape)
         self.rew_filter = get_filter(config["reward_filter"], ())
-        self.filters = {"obs_filter": self.obs_filter,
-                        "rew_filter": self.rew_filter}
+        self.filters = {
+            "obs_filter": self.obs_filter,
+            "rew_filter": self.rew_filter
+        }
+
         self.sampler = AsyncSampler(env, self.policy, self.obs_filter,
                                     config["batch_size"])
         if start_sampler and self.sampler._async:
@@ -52,8 +60,11 @@ def __init__(
     def sample(self):
         rollout = self.sampler.get_data()
         samples = process_rollout(
-            rollout, self.rew_filter, gamma=self.config["gamma"],
-            lambda_=self.config["lambda"], use_gae=True)
+            rollout,
+            self.rew_filter,
+            gamma=self.config["gamma"],
+            lambda_=self.config["lambda"],
+            use_gae=True)
         return samples
 
     def get_completed_rollout_metrics(self):
@@ -79,9 +90,7 @@ def set_weights(self, params):
     def save(self):
         filters = self.get_filters(flush_after=True)
         weights = self.get_weights()
-        return pickle.dumps({
-            "filters": filters,
-            "weights": weights})
+        return pickle.dumps({"filters": filters, "weights": weights})
 
     def restore(self, objs):
         objs = pickle.loads(objs)

python/ray/rllib/a3c/shared_torch_policy.py

@@ -6,8 +6,8 @@
 import torch.nn.functional as F
 
 from ray.rllib.a3c.torchpolicy import TorchPolicy
-from ray.rllib.models.pytorch.misc import var_to_np, convert_batch
 from ray.rllib.models.catalog import ModelCatalog
+from ray.rllib.models.pytorch.misc import convert_batch, var_to_np
 
 
 class SharedTorchPolicy(TorchPolicy):
@@ -28,7 +28,7 @@ def _setup_graph(self, ob_space, ac_space):
             self._model.parameters(), lr=self.config["lr"])
 
     def compute(self, ob, *args):
-        """Should take in a SINGLE ob"""
+        """Should take in a SINGLE ob."""
         with self.lock:
             ob = torch.from_numpy(ob).float().unsqueeze(0)
             logits, values = self._model(ob)
@@ -64,16 +64,18 @@ def _evaluate(self, obs, actions):
         return values, action_log_probs, entropy
 
     def _backward(self, batch):
-        """Loss is encoded in here. Defining a new loss function
-        would start by rewriting this function"""
+        """Loss is encoded in here.
+
+        Defining a new loss function would start by rewriting this
+        function
+        """
 
         states, actions, advs, rs, _ = convert_batch(batch)
         values, action_log_probs, entropy = self._evaluate(states, actions)
         pi_err = -advs.dot(action_log_probs.reshape(-1))
         value_err = F.mse_loss(values.reshape(-1), rs)
 
         self.optimizer.zero_grad()
-
         overall_err = sum([
             pi_err,
             self.config["vf_loss_coeff"] * value_err,

python/ray/rllib/a3c/torchpolicy.py

@@ -2,6 +2,7 @@
 from __future__ import division
 from __future__ import print_function
 
+from copy import deepcopy
 import torch
 
 from ray.rllib.a3c.policy import Policy
@@ -30,6 +31,9 @@ def __init__(self,
         self.lock = Lock()
 
     def apply_gradients(self, grads):
+        # TODO(alok): see how A3C fills gradient buffers so that they don't get
+        # cleared by zero_grad
+        grads = deepcopy(grads)  # TODO rm
         self.optimizer.zero_grad()
         for g, p in zip(grads, self._model.parameters()):
             p.grad = torch.from_numpy(g)

python/ray/rllib/agent.py

@@ -261,6 +261,9 @@ def get_agent_class(alg):
     elif alg == "PG":
         from ray.rllib import pg
         return pg.PGAgent
+    elif alg == "TRPO":
+        from ray.rllib import trpo
+        return trpo.TRPOAgent
     elif alg == "script":
         from ray.tune import script_runner
         return script_runner.ScriptRunner

python/ray/rllib/trpo/__init__.py

@@ -0,0 +1,6 @@
+from ray.rllib.trpo.trpo import DEFAULT_CONFIG, TRPOAgent
+
+__all__ = [
+    'TRPOAgent',
+    'DEFAULT_CONFIG',
+]