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

python/ray/autoscaler/updater.py

@@ -75,9 +75,9 @@ def run(self):
             self.provider.set_node_tags(self.node_id,
                                         {TAG_RAY_NODE_STATUS: "UpdateFailed"})
             if self.logfile is not None:
-                print("----- BEGIN REMOTE LOGS -----\n" + open(
-                    self.logfile.name).read() + "\n----- END REMOTE LOGS -----"
-                      )
+                print("----- BEGIN REMOTE LOGS -----\n" +
+                      open(self.logfile.name).read() +
+                      "\n----- END REMOTE LOGS -----")
             raise e
         self.provider.set_node_tags(
             self.node_id, {

python/ray/experimental/array/distributed/linalg.py

@@ -76,10 +76,9 @@ def tsqr(a):
                 lower = [a.shape[1], 0]
                 upper = [2 * a.shape[1], core.BLOCK_SIZE]
             ith_index //= 2
-            q_block_current = ra.dot.remote(q_block_current,
-                                            ra.subarray.remote(
-                                                q_tree[ith_index, j], lower,
-                                                upper))
+            q_block_current = ra.dot.remote(
+                q_block_current,
+                ra.subarray.remote(q_tree[ith_index, j], lower, upper))
         q_result.objectids[i] = q_block_current
     r = current_rs[0]
     return q_result, ray.get(r)
@@ -222,10 +221,10 @@ def qr(a):
         y_col_block = core.subblocks.remote(y_res, [], [i])
         q = core.subtract.remote(
             q,
-            core.dot.remote(y_col_block,
-                            core.dot.remote(
-                                Ts[i],
-                                core.dot.remote(
-                                    core.transpose.remote(y_col_block), q))))
+            core.dot.remote(
+                y_col_block,
+                core.dot.remote(
+                    Ts[i],
+                    core.dot.remote(core.transpose.remote(y_col_block), q))))
 
     return ray.get(q), r_res

python/ray/experimental/state.py

@@ -750,8 +750,8 @@ def micros_rel(ts):
                         "name":
                         "SubmitTask",
                         "args": {},
-                        "id": (parent_info["worker_id"] +
-                               str(micros(min(parent_times))))
+                        "id": (parent_info["worker_id"] + str(
+                            micros(min(parent_times))))
                     }
                     full_trace.append(parent)
 
@@ -825,8 +825,8 @@ def micros_rel(ts):
                         "name":
                         "SubmitTask",
                         "args": {},
-                        "id": (parent_info["worker_id"] +
-                               str(micros(min(parent_times))))
+                        "id": (parent_info["worker_id"] + str(
+                            micros(min(parent_times))))
                     }
                     full_trace.append(parent)
 

python/ray/experimental/ui.py

@@ -451,8 +451,8 @@ def task_completion_time_update(abs_earliest, abs_latest, abs_num_tasks,
         # Create the distribution to plot
         distr = []
         for task_id, data in tasks.items():
-            distr.append(
-                data["store_outputs_end"] - data["get_arguments_start"])
+            distr.append(data["store_outputs_end"] -
+                         data["get_arguments_start"])
 
         # Create a histogram from the distribution
         top, bin_edges = np.histogram(distr, bins="auto")
@@ -520,10 +520,10 @@ def compute_utilizations(abs_earliest,
         # Walk over each time bucket that this task intersects, adding the
         # amount of time that the task intersects within each bucket
         for bucket_idx in range(start_bucket, end_bucket + 1):
-            bucket_start_time = ((
-                earliest_time + bucket_idx) * bucket_time_length)
-            bucket_end_time = ((earliest_time +
-                                (bucket_idx + 1)) * bucket_time_length)
+            bucket_start_time = (
+                (earliest_time + bucket_idx) * bucket_time_length)
+            bucket_end_time = (
+                (earliest_time + (bucket_idx + 1)) * bucket_time_length)
 
             task_start_time_within_bucket = max(task_start_time,
                                                 bucket_start_time)

python/ray/log_monitor.py

@@ -39,8 +39,9 @@ def __init__(self, redis_ip_address, redis_port, node_ip_address):
     def update_log_filenames(self):
         """Get the most up-to-date list of log files to monitor from Redis."""
         num_current_log_files = len(self.log_files)
-        new_log_filenames = self.redis_client.lrange("LOG_FILENAMES:{}".format(
-            self.node_ip_address), num_current_log_files, -1)
+        new_log_filenames = self.redis_client.lrange(
+            "LOG_FILENAMES:{}".format(self.node_ip_address),
+            num_current_log_files, -1)
         for log_filename in new_log_filenames:
             print("Beginning to track file {}".format(log_filename))
             assert log_filename not in self.log_files

python/ray/monitor.py

@@ -189,10 +189,9 @@ def cleanup_object_table(self):
                 if manager in self.dead_plasma_managers:
                     # If the object was on a dead plasma manager, remove that
                     # location entry.
-                    ok = self.state._execute_command(object_id,
-                                                     "RAY.OBJECT_TABLE_REMOVE",
-                                                     object_id.id(),
-                                                     hex_to_binary(manager))
+                    ok = self.state._execute_command(
+                        object_id, "RAY.OBJECT_TABLE_REMOVE", object_id.id(),
+                        hex_to_binary(manager))
                     if ok != b"OK":
                         log.warn("Failed to remove object location for dead "
                                  "plasma manager.")
@@ -507,8 +506,8 @@ def run(self):
         log.debug("{} dead local schedulers, {} plasma managers total, {} "
                   "dead plasma managers".format(
                       len(self.dead_local_schedulers),
-                      (len(self.live_plasma_managers) +
-                       len(self.dead_plasma_managers)),
+                      (len(self.live_plasma_managers) + len(
+                          self.dead_plasma_managers)),
                       len(self.dead_plasma_managers)))
 
         # Handle messages from the subscription channels.

python/ray/rllib/a3c/shared_torch_policy.py

@@ -33,8 +33,8 @@ def compute(self, ob, *args):
         with self.lock:
             ob = Variable(torch.from_numpy(ob).float().unsqueeze(0))
             logits, values = self._model(ob)
-            samples = self._model.probs(logits).multinomial().squeeze()
-            values = values.squeeze(0)
+            samples = F.softmax(logits, dim=1).multinomial(num_samples=1).squeeze()
+            values = values.squeeze()
             return var_to_np(samples), {"vf_preds": var_to_np(values)}
 
     def compute_logits(self, ob, *args):
@@ -48,15 +48,17 @@ def value(self, ob, *args):
             ob = Variable(torch.from_numpy(ob).float().unsqueeze(0))
             res = self._model.hidden_layers(ob)
             res = self._model.value_branch(res)
-            res = res.squeeze(0)
+            res = res.squeeze()
             return var_to_np(res)
 
     def _evaluate(self, obs, actions):
         """Passes in multiple obs."""
         logits, values = self._model(obs)
-        log_probs = F.log_softmax(logits)
-        probs = self._model.probs(logits)
+        log_probs = F.log_softmax(logits,dim=1)
+        probs = F.softmax(logits, dim=1)
         action_log_probs = log_probs.gather(1, actions.view(-1, 1))
+        # TODO(alok): set distribution based on action space and use its
+        # .entropy() method to calculate automatically
         entropy = -(log_probs * probs).sum(-1).sum()
         return values, action_log_probs, entropy
 
@@ -67,7 +69,7 @@ def _backward(self, batch):
         states, acs, advs, rs, _ = convert_batch(batch)
         values, ac_logprobs, entropy = self._evaluate(states, acs)
         pi_err = -(advs * ac_logprobs).sum()
-        value_err = 0.5 * (values - rs).pow(2).sum()
+        value_err = F.mse_loss(values, rs)
 
         self.optimizer.zero_grad()
         overall_err = (pi_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 torch.autograd import Variable
 
@@ -26,6 +27,7 @@ def __init__(self, registry, ob_space, action_space, config,
         self.lock = Lock()
 
     def apply_gradients(self, grads):
+        grads = deepcopy(grads) # to prevent zero_grad from clearing grads since they may share memory with the policy's .grad tensors
         self.optimizer.zero_grad()
         for g, p in zip(grads, self._model.parameters()):
             p.grad = Variable(torch.from_numpy(g))

python/ray/rllib/models/pytorch/fcnet.py

@@ -32,12 +32,11 @@ def _init(self, inputs, num_outputs, options):
         self.hidden_layers = nn.Sequential(*layers)
 
         self.logits = SlimFC(
-            last_layer_size, num_outputs,
+            in_size=last_layer_size, out_size=num_outputs,
             initializer=normc_initializer(0.01),
             activation_fn=None)
-        self.probs = nn.Softmax()
         self.value_branch = SlimFC(
-            last_layer_size, 1,
+            in_size=last_layer_size, out_size=1,
             initializer=normc_initializer(1.0),
             activation_fn=None)
 
@@ -52,5 +51,5 @@ def forward(self, obs):
             value: value function for each state"""
         res = self.hidden_layers(obs)
         logits = self.logits(res)
-        value = self.value_branch(res)
+        value = self.value_branch(res).reshape(-1)
         return logits, value

python/ray/rllib/models/pytorch/misc.py

@@ -5,38 +5,32 @@
 
 import numpy as np
 import torch
-from torch.autograd import Variable
 
 
 def convert_batch(trajectory, has_features=False):
     """Convert trajectory from numpy to PT variable"""
-    states = Variable(torch.from_numpy(
-        trajectory["observations"]).float())
-    acs = Variable(torch.from_numpy(
-        trajectory["actions"]))
-    advs = Variable(torch.from_numpy(
-        trajectory["advantages"].copy()).float())
-    advs = advs.view(-1, 1)
-    rs = Variable(torch.from_numpy(
-        trajectory["value_targets"]).float())
-    rs = rs.view(-1, 1)
+    states = torch.from_numpy(trajectory["obs"]).float()
+    acs = torch.from_numpy(trajectory["actions"])
+    advs = torch.from_numpy(trajectory["advantages"].copy()).float().view(
+        -1, 1)
+    rs = torch.from_numpy(trajectory["rewards"]).float().view(-1)
     if has_features:
-        features = [Variable(torch.from_numpy(f))
-                    for f in trajectory["features"]]
+        features = [torch.from_numpy(f) for f in trajectory["features"]]
     else:
         features = trajectory["features"]
     return states, acs, advs, rs, features
 
 
 def var_to_np(var):
-    return var.data.numpy()[0]
+    return var.data.numpy()
 
 
 def normc_initializer(std=1.0):
     def initializer(tensor):
         tensor.data.normal_(0, 1)
         tensor.data *= std / torch.sqrt(
             tensor.data.pow(2).sum(1, keepdim=True))
+
     return initializer
 
 

python/ray/rllib/models/pytorch/model.py

@@ -53,14 +53,14 @@ def forward(self, x):
 class SlimFC(nn.Module):
     """Simple PyTorch of `linear` function"""
 
-    def __init__(self, in_size, size, initializer=None,
+    def __init__(self, in_size, out_size, initializer=None,
                  activation_fn=None, bias_init=0):
         super(SlimFC, self).__init__()
         layers = []
-        linear = nn.Linear(in_size, size)
+        linear = nn.Linear(in_size, out_size)
         if initializer:
             initializer(linear.weight)
-        nn.init.constant(linear.bias, bias_init)
+        nn.init.constant_(linear.bias, bias_init)
         layers.append(linear)
         if activation_fn:
             layers.append(activation_fn())

python/ray/rllib/models/pytorch/visionnet.py

@@ -41,7 +41,6 @@ def _init(self, inputs, num_outputs, options):
 
         self.logits = SlimFC(
             out_channels, num_outputs, initializer=nn.init.xavier_uniform)
-        self.probs = nn.Softmax()
         self.value_branch = SlimFC(
             out_channels, 1, initializer=normc_initializer())
 

python/ray/rllib/optimizers/sample_batch.py

@@ -36,8 +36,8 @@ def __init__(self, *args, **kwargs):
     @staticmethod
     def concat_samples(samples):
         out = {}
-        for k in samples[0].data.keys():
-            out[k] = np.concatenate([s.data[k] for s in samples])
+        for k in samples[0].keys():
+            out[k] = np.concatenate([s[k] for s in samples])
         return SampleBatch(out)
 
     def concat(self, other):
@@ -50,10 +50,10 @@ def concat(self, other):
             {"a": [1, 2, 3, 4, 5]}
         """
 
-        assert self.data.keys() == other.data.keys(), "must have same columns"
+        assert self.keys() == other.keys(), "must have same columns"
         out = {}
-        for k in self.data.keys():
-            out[k] = np.concatenate([self.data[k], other.data[k]])
+        for k in self.keys():
+            out[k] = np.concatenate([self[k], other[k]])
         return SampleBatch(out)
 
     def rows(self):
@@ -70,7 +70,7 @@ def rows(self):
 
         for i in range(self.count):
             row = {}
-            for k in self.data.keys():
+            for k in self.keys():
                 row[k] = self[k][i]
             yield row
 
@@ -85,19 +85,37 @@ def columns(self, keys):
 
         out = []
         for k in keys:
-            out.append(self.data[k])
+            out.append(self[k])
         return out
 
     def shuffle(self):
         permutation = np.random.permutation(self.count)
-        for key, val in self.data.items():
-            self.data[key] = val[permutation]
+        for key, val in self.items():
+            self[key] = val[permutation]
 
     def __getitem__(self, key):
         return self.data[key]
 
+    def __setitem__(self, key, item):
+        self.data[key] = item
+
     def __str__(self):
         return "SampleBatch({})".format(str(self.data))
 
     def __repr__(self):
         return "SampleBatch({})".format(str(self.data))
+
+    def keys(self):
+        return self.data.keys()
+
+    def items(self):
+        return self.data.items()
+
+    def __iter__(self):
+        return self.data.__iter__()
+
+    def __next__(self):
+        return self.data.__next__()
+
+    def __contains__(self, x):
+        return x in self.data

python/ray/rllib/test/test_optimizers.py

@@ -33,11 +33,11 @@ def testConcat(self):
         b2 = SampleBatch({"a": np.array([1]), "b": np.array([4])})
         b3 = SampleBatch({"a": np.array([1]), "b": np.array([5])})
         b12 = b1.concat(b2)
-        self.assertEqual(b12.data["a"].tolist(), [1, 2, 3, 1])
-        self.assertEqual(b12.data["b"].tolist(), [4, 5, 6, 4])
+        self.assertEqual(b12["a"].tolist(), [1, 2, 3, 1])
+        self.assertEqual(b12["b"].tolist(), [4, 5, 6, 4])
         b = SampleBatch.concat_samples([b1, b2, b3])
-        self.assertEqual(b.data["a"].tolist(), [1, 2, 3, 1, 1])
-        self.assertEqual(b.data["b"].tolist(), [4, 5, 6, 4, 5])
+        self.assertEqual(b["a"].tolist(), [1, 2, 3, 1, 1])
+        self.assertEqual(b["b"].tolist(), [4, 5, 6, 4, 5])
 
 
 if __name__ == '__main__':

python/ray/rllib/trpo/__init__.py

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