[rllib] Add magic methods for rollouts (#2024)

python/ray/rllib/utils/sampler.py

@@ -56,9 +56,30 @@ class PartialRollout(object):
             terminal (bool): if rollout has terminated."""
         return self.data["dones"][-1]
 
+    def __getitem__(self, key):
+        return self.data[key]
 
-CompletedRollout = namedtuple(
-    "CompletedRollout", ["episode_length", "episode_reward"])
+    def __setitem__(self, key, item):
+        self.data[key] = item
+
+    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
+
+
+CompletedRollout = namedtuple("CompletedRollout",
+                              ["episode_length", "episode_reward"])
 
 
 class SyncSampler(object):
@@ -71,16 +92,15 @@ class SyncSampler(object):
     thread."""
     async = False
 
-    def __init__(self, env, policy, obs_filter,
-                 num_local_steps, horizon=None):
+    def __init__(self, env, policy, obs_filter, num_local_steps, horizon=None):
         self.num_local_steps = num_local_steps
         self.horizon = horizon
         self.env = env
         self.policy = policy
         self._obs_filter = obs_filter
-        self.rollout_provider = _env_runner(
-            self.env, self.policy, self.num_local_steps, self.horizon,
-            self._obs_filter)
+        self.rollout_provider = _env_runner(self.env, self.policy,
+                                            self.num_local_steps, self.horizon,
+                                            self._obs_filter)
         self.metrics_queue = queue.Queue()
 
     def get_data(self):
@@ -108,10 +128,10 @@ class AsyncSampler(threading.Thread):
     accumulate and the gradient can be calculated on up to 5 batches."""
     async = True
 
-    def __init__(self, env, policy, obs_filter,
-                 num_local_steps, horizon=None):
-        assert getattr(obs_filter, "is_concurrent", False), (
-            "Observation Filter must support concurrent updates.")
+    def __init__(self, env, policy, obs_filter, num_local_steps, horizon=None):
+        assert getattr(
+            obs_filter, "is_concurrent",
+            False), ("Observation Filter must support concurrent updates.")
         threading.Thread.__init__(self)
         self.queue = queue.Queue(5)
         self.metrics_queue = queue.Queue()
@@ -132,9 +152,9 @@ class AsyncSampler(threading.Thread):
             raise e
 
     def _run(self):
-        rollout_provider = _env_runner(
-            self.env, self.policy, self.num_local_steps,
-            self.horizon, self._obs_filter)
+        rollout_provider = _env_runner(self.env, self.policy,
+                                       self.num_local_steps, self.horizon,
+                                       self._obs_filter)
         while True:
             # The timeout variable exists because apparently, if one worker
             # dies, the other workers won't die with it, unless the timeout is
@@ -232,13 +252,14 @@ def _env_runner(env, policy, num_local_steps, horizon, obs_filter):
                 action = np.concatenate(action, axis=0).flatten()
 
             # Collect the experience.
-            rollout.add(obs=last_observation,
-                        actions=action,
-                        rewards=reward,
-                        dones=terminal,
-                        features=last_features,
-                        new_obs=observation,
-                        **pi_info)
+            rollout.add(
+                obs=last_observation,
+                actions=action,
+                rewards=reward,
+                dones=terminal,
+                features=last_features,
+                new_obs=observation,
+                **pi_info)
 
             last_observation = observation
             last_features = features
@@ -247,8 +268,8 @@ def _env_runner(env, policy, num_local_steps, horizon, obs_filter):
                 terminal_end = True
                 yield CompletedRollout(length, rewards)
 
-                if (length >= horizon or
-                        not env.metadata.get("semantics.autoreset")):
+                if (length >= horizon
+                        or not env.metadata.get("semantics.autoreset")):
                     last_observation = obs_filter(env.reset())
                     if hasattr(policy, "get_initial_features"):
                         last_features = policy.get_initial_features()