[rllib] Count actual sample batch size instead of configured batch size in A3C. (#2399)

This fixes a metrics accounting bug where the sample count is not reported correctly.

python/ray/rllib/evaluation/policy_evaluator.py

@@ -336,10 +336,11 @@ class PolicyEvaluator(EvaluatorInterface):
                 for pid, batch in samples.policy_batches.items():
                     grad_out[pid], info_out[pid] = (
                         self.policy_map[pid].compute_gradients(batch))
-            return grad_out, info_out
         else:
-            return self.policy_map[DEFAULT_POLICY_ID].compute_gradients(
-                samples)
+            grad_out, info_out = (
+                self.policy_map[DEFAULT_POLICY_ID].compute_gradients(samples))
+        info_out["batch_count"] = samples.count
+        return grad_out, info_out
 
     def apply_gradients(self, grads):
         if isinstance(grads, dict):

python/ray/rllib/optimizers/async_gradients_optimizer.py

@@ -14,12 +14,11 @@ class AsyncGradientsOptimizer(PolicyOptimizer):
     evaluators, sending updated weights back as needed. This pipelines the
     gradient computations on the remote workers.
     """
-    def _init(self, grads_per_step=100, batch_size=10):
+    def _init(self, grads_per_step=100):
         self.apply_timer = TimerStat()
         self.wait_timer = TimerStat()
         self.dispatch_timer = TimerStat()
         self.grads_per_step = grads_per_step
-        self.batch_size = batch_size
         if not self.remote_evaluators:
             raise ValueError(
                 "Async optimizer requires at least 1 remote evaluator")
@@ -40,11 +39,13 @@ class AsyncGradientsOptimizer(PolicyOptimizer):
         while gradient_queue:
             with self.wait_timer:
                 fut, e = gradient_queue.pop(0)
-                gradient, _ = ray.get(fut)
+                gradient, info = ray.get(fut)
 
             if gradient is not None:
                 with self.apply_timer:
                     self.local_evaluator.apply_gradients(gradient)
+                self.num_steps_sampled += info["batch_count"]
+                self.num_steps_trained += info["batch_count"]
 
             if num_gradients < self.grads_per_step:
                 with self.dispatch_timer:
@@ -53,9 +54,6 @@ class AsyncGradientsOptimizer(PolicyOptimizer):
                     gradient_queue.append((fut, e))
                     num_gradients += 1
 
-        self.num_steps_sampled += self.grads_per_step * self.batch_size
-        self.num_steps_trained += self.grads_per_step * self.batch_size
-
     def stats(self):
         return dict(PolicyOptimizer.stats(self), **{
             "wait_time_ms": round(1000 * self.wait_timer.mean, 3),

python/ray/rllib/test/mock_evaluator.py

@@ -28,7 +28,7 @@ class _MockEvaluator(object):
         return SampleBatch(samples_dict)
 
     def compute_gradients(self, samples):
-        return self._grad * samples.count, {}
+        return self._grad * samples.count, {"batch_count": samples.count}
 
     def apply_gradients(self, grads):
         self._weights += self._grad