[rllib] Fix stats collection and some docs bugs since the refactoring (#2361)

* fix

* fix pbt example

* fix

* fix

* single thread by default

* vec

* fix

* fix

python/ray/rllib/agents/a3c/a3c.py

@@ -9,7 +9,6 @@ import ray
 from ray.rllib.agents.agent import Agent, with_common_config
 from ray.rllib.optimizers import AsyncGradientsOptimizer
 from ray.rllib.utils import FilterManager
-from ray.rllib.evaluation.metrics import collect_metrics
 from ray.tune.trial import Resources
 
 DEFAULT_CONFIG = with_common_config({
@@ -98,12 +97,13 @@ class A3CAgent(Agent):
             self.config["optimizer"])
 
     def _train(self):
+        prev_steps = self.optimizer.num_steps_sampled
         self.optimizer.step()
         FilterManager.synchronize(
             self.local_evaluator.filters, self.remote_evaluators)
-        result = collect_metrics(self.local_evaluator, self.remote_evaluators)
+        result = self.optimizer.collect_metrics()
         result = result._replace(
-            info=self.optimizer.stats())
+            timesteps_this_iter=self.optimizer.num_steps_sampled - prev_steps)
         return result
 
     def _stop(self):

python/ray/rllib/agents/agent.py

@@ -39,8 +39,14 @@ COMMON_CONFIG = {
     "model": {},
     # Arguments to pass to the rllib optimizer
     "optimizer": {},
-    # Override default TF session args if non-empty
-    "tf_session_args": {},
+    # Configure TF for single-process operation by default
+    "tf_session_args": {
+        "intra_op_parallelism_threads": 1,
+        "inter_op_parallelism_threads": 1,
+        "gpu_options": {
+            "allow_growth": True,
+        },
+    },
     # Whether to LZ4 compress observations
     "compress_observations": False,
 

python/ray/rllib/agents/dqn/dqn.py

@@ -185,9 +185,17 @@ class DQNAgent(Agent):
             e.foreach_policy.remote(lambda p, _: p.set_epsilon(exp_val))
             exp_vals.append(exp_val)
 
-        result = collect_metrics(
-            self.local_evaluator, self.remote_evaluators)
+        if self.config["per_worker_exploration"]:
+            # Only collect metrics from the third of workers with lowest eps
+            result = collect_metrics(
+                self.local_evaluator,
+                self.remote_evaluators[-len(self.remote_evaluators) // 3:])
+        else:
+            result = collect_metrics(
+                self.local_evaluator, self.remote_evaluators)
+
         return result._replace(
+            timesteps_this_iter=self.global_timestep - start_timestep,
             info=dict({
                 "min_exploration": min(exp_vals),
                 "max_exploration": max(exp_vals),

python/ray/rllib/agents/pg/pg.py

@@ -4,7 +4,6 @@ from __future__ import print_function
 
 from ray.rllib.agents.agent import Agent, with_common_config
 from ray.rllib.agents.pg.pg_policy_graph import PGPolicyGraph
-from ray.rllib.evaluation.metrics import collect_metrics
 from ray.rllib.optimizers import SyncSamplesOptimizer
 from ray.tune.trial import Resources
 
@@ -49,6 +48,7 @@ class PGAgent(Agent):
             self.config["optimizer"])
 
     def _train(self):
+        prev_steps = self.optimizer.num_steps_sampled
         self.optimizer.step()
-        return collect_metrics(
-            self.optimizer.local_evaluator, self.optimizer.remote_evaluators)
+        return self.optimizer.collect_metrics()._replace(
+            timesteps_this_iter=self.optimizer.num_steps_sampled - prev_steps)

python/ray/rllib/agents/ppo/ppo.py

@@ -9,7 +9,6 @@ import pickle
 import ray
 from ray.rllib.agents import Agent, with_common_config
 from ray.rllib.agents.ppo.ppo_tf_policy import PPOTFPolicyGraph
-from ray.rllib.evaluation.metrics import collect_metrics
 from ray.rllib.utils import FilterManager
 from ray.rllib.optimizers.multi_gpu_optimizer import LocalMultiGPUOptimizer
 from ray.tune.trial import Resources
@@ -81,6 +80,8 @@ class PPOAgent(Agent):
              "timesteps_per_batch": self.config["timesteps_per_batch"]})
 
     def _train(self):
+        prev_steps = self.optimizer.num_steps_sampled
+
         def postprocess_samples(batch):
             # Divide by the maximum of value.std() and 1e-4
             # to guard against the case where all values are equal
@@ -92,6 +93,7 @@ class PPOAgent(Agent):
             if not self.config["use_gae"]:
                 batch.data["value_targets"] = dummy
                 batch.data["vf_preds"] = dummy
+
         extra_fetches = self.optimizer.step(postprocess_fn=postprocess_samples)
         kl = np.array(extra_fetches["kl"]).mean(axis=1)[-1]
         total_loss = np.array(extra_fetches["total_loss"]).mean(axis=1)[-1]
@@ -112,8 +114,10 @@ class PPOAgent(Agent):
 
         FilterManager.synchronize(
             self.local_evaluator.filters, self.remote_evaluators)
-        res = collect_metrics(self.local_evaluator, self.remote_evaluators)
-        res = res._replace(info=info)
+        res = self.optimizer.collect_metrics()
+        res = res._replace(
+            timesteps_this_iter=self.optimizer.num_steps_sampled - prev_steps,
+            info=dict(info, **res.info))
         return res
 
     def _stop(self):