[rllib] streaming minibatching for IMPALA (#3402)

* mb impala

* fix

* paropt

* update

* cpu warn

* on cpu

* fix mb

* doc

* docs

* comment

* larger num

* early release

* remove grad clip

* only check loader count in multi gpu mode

* revert bad multigpu changes

* num sgd iter

* comment

* reuse optimizer

* add test

* par load test

* loosen test

* Update run_multi_node_tests.sh

* fix local mode

* Update agent.py

python/ray/rllib/agents/agent.py

@@ -130,7 +130,7 @@ COMMON_CONFIG = {
     # Drop metric batches from unresponsive workers after this many seconds
     "collect_metrics_timeout": 180,
 
-    # === Offline Data Input / Output ===
+    # === Offline Data Input / Output (Experimental) ===
     # Specify how to generate experiences:
     #  - "sampler": generate experiences via online simulation (default)
     #  - a local directory or file glob expression (e.g., "/tmp/*.json")

python/ray/rllib/agents/impala/impala.py

@@ -18,10 +18,11 @@ OPTIMIZER_SHARED_CONFIGS = [
     "train_batch_size",
     "replay_buffer_num_slots",
     "replay_proportion",
-    "num_parallel_data_loaders",
-    "grad_clip",
+    "num_data_loader_buffers",
     "max_sample_requests_in_flight_per_worker",
     "broadcast_interval",
+    "num_sgd_iter",
+    "minibatch_buffer_size",
 ]
 
 # yapf: disable
@@ -33,6 +34,17 @@ DEFAULT_CONFIG = with_common_config({
     "vtrace_clip_pg_rho_threshold": 1.0,
 
     # System params.
+    #
+    # == Overview of data flow in IMPALA ==
+    # 1. Policy evaluation in parallel across `num_workers` actors produces
+    #    batches of size `sample_batch_size * num_envs_per_worker`.
+    # 2. If enabled, the replay buffer stores and produces batches of size
+    #    `sample_batch_size * num_envs_per_worker`.
+    # 3. If enabled, the minibatch ring buffer stores and replays batches of
+    #    size `train_batch_size` up to `num_sgd_iter` times per batch.
+    # 4. The learner thread executes data parallel SGD across `num_gpus` GPUs
+    #    on batches of size `train_batch_size`.
+    #
     "sample_batch_size": 50,
     "train_batch_size": 500,
     "min_iter_time_s": 10,
@@ -40,18 +52,23 @@ DEFAULT_CONFIG = with_common_config({
     # number of GPUs the learner should use.
     "num_gpus": 1,
     # set >1 to load data into GPUs in parallel. Increases GPU memory usage
-    # proportionally with the number of loaders.
-    "num_parallel_data_loaders": 1,
-    # level of queuing for sampling.
-    "max_sample_requests_in_flight_per_worker": 2,
-    # max number of workers to broadcast one set of weights to
-    "broadcast_interval": 1,
+    # proportionally with the number of buffers.
+    "num_data_loader_buffers": 1,
+    # how many train batches should be retained for minibatching. This conf
+    # only has an effect if `num_sgd_iter > 1`.
+    "minibatch_buffer_size": 1,
+    # number of passes to make over each train batch
+    "num_sgd_iter": 1,
     # set >0 to enable experience replay. Saved samples will be replayed with
     # a p:1 proportion to new data samples.
     "replay_proportion": 0.0,
     # number of sample batches to store for replay. The number of transitions
     # saved total will be (replay_buffer_num_slots * sample_batch_size).
     "replay_buffer_num_slots": 100,
+    # level of queuing for sampling.
+    "max_sample_requests_in_flight_per_worker": 2,
+    # max number of workers to broadcast one set of weights to
+    "broadcast_interval": 1,
 
     # Learning params.
     "grad_clip": 40.0,