Repository: ray-project/ray

Issue:
AsyncSampler auto-concat feature causes non-requested batch size increase
<!--
General questions should be asked on the mailing list ray-dev@googlegroups.com.

Before submitting an issue, please fill out the following form.
-->

### System information
- **OS Platform and Distribution (e.g., Linux Ubuntu 16.04)**: Linux Ubuntu 16.04
- **Ray installed from (source or binary)**: source
- **Ray version**: 0.6.0
- **Python version**: 3.6.7
- **Exact command to reproduce**: rllib train --run PPO --env CartPole-v0 --config '{"train_batch_size":200,"sample_batch_size":200,"num_workers":1,"sample_async":true,"observation_filter":"NoFilter"}'

<!--
You can obtain the Ray version with

python -c "import ray; print(ray.__version__)"
-->

### Describe the problem
<!-- Describe the problem clearly here. -->
When working with the AsyncSampler, the auto-concat feature in the 'get_data()' routine causes sample sizes which exceed the 'batch_steps' passed to PolicyEvaluator, when using batch_mode="truncate_episodes". 
In this mode the doc says the batch size will be exactly 'batch_steps * num_envs', however due to auto-concat feature, assuming the ENV is faster than the training, we will get the entire queue as additional data (in the provided example command, we get 1200 train steps each iteration instead of the requested 200, due to the AsyncSampler queue size of 5).

In case of batch_mode="complete_episodes", the auto-concat feature has no effect (Since each rollout always ends with 'done')

The comment suggests this is to improve A3C performance, which it probably does for fast envs, but that would always be true by increasing the batch size. 
In many algorithms it's important to have a fixed batch size using 'truncate_episodes', while still using the AsyncSampler (For example for real-time environments which continue running during train time)

Maybe make the auto-concat an optional configuration? Maybe as part of a new 'batch_mode' option (batch_mode="max_available"?)

### Source code / logs
<!-- Include any logs or source code that would be helpful to diagnose the problem. If including tracebacks, please include the full traceback. Large logs and files should be attached. Try to provide a reproducible test case that is the bare minimum necessary to generate the problem. -->



Discussion:

Resolving patch:
diff --git a/python/ray/rllib/evaluation/sampler.py b/python/ray/rllib/evaluation/sampler.py
--- a/python/ray/rllib/evaluation/sampler.py
+++ b/python/ray/rllib/evaluation/sampler.py
@@ -10,8 +10,7 @@
 import threading
 
 from ray.rllib.evaluation.episode import MultiAgentEpisode, _flatten_action
-from ray.rllib.evaluation.sample_batch import MultiAgentSampleBatchBuilder, \
-    MultiAgentBatch
+from ray.rllib.evaluation.sample_batch import MultiAgentSampleBatchBuilder
 from ray.rllib.evaluation.tf_policy_graph import TFPolicyGraph
 from ray.rllib.env.async_vector_env import AsyncVectorEnv
 from ray.rllib.env.atari_wrappers import get_wrapper_by_cls, MonitorEnv
@@ -164,20 +163,6 @@ def get_data(self):
         if isinstance(rollout, BaseException):
             raise rollout
 
-        # We can't auto-concat rollouts in these modes
-        if self.async_vector_env.num_envs > 1 or \
-                isinstance(rollout, MultiAgentBatch):
-            return rollout
-
-        # Auto-concat rollouts; TODO(ekl) is this important for A3C perf?
-        while not rollout["dones"][-1]:
-            try:
-                part = self.queue.get_nowait()
-                if isinstance(part, BaseException):
-                    raise rollout
-                rollout = rollout.concat(part)
-            except queue.Empty:
-                break
         return rollout
 
     def get_metrics(self):


Test patch:
diff --git a/python/ray/rllib/test/test_policy_evaluator.py b/python/ray/rllib/test/test_policy_evaluator.py
--- a/python/ray/rllib/test/test_policy_evaluator.py
+++ b/python/ray/rllib/test/test_policy_evaluator.py
@@ -263,18 +263,6 @@ def testAsync(self):
             self.assertIn(key, batch)
         self.assertGreater(batch["advantages"][0], 1)
 
-    def testAutoConcat(self):
-        ev = PolicyEvaluator(
-            env_creator=lambda _: MockEnv(episode_length=40),
-            policy_graph=MockPolicyGraph,
-            sample_async=True,
-            batch_steps=10,
-            batch_mode="truncate_episodes",
-            observation_filter="ConcurrentMeanStdFilter")
-        time.sleep(2)
-        batch = ev.sample()
-        self.assertEqual(batch.count, 40)  # auto-concat up to 5 episodes
-
     def testAutoVectorization(self):
         ev = PolicyEvaluator(
             env_creator=lambda cfg: MockEnv(episode_length=20, config=cfg),