support fractional num_updates

Author: hnyu

alf/algorithms/algorithm.py

@@ -1517,8 +1517,8 @@ def _replay():
                         "No mini_batch_length is specified for off-policy training"
                     )
                     experience, batch_info = self._replay_buffer.get_batch(
-                        batch_size=(mini_batch_size *
-                                    config.num_updates_per_train_iter),
+                        batch_size=int(mini_batch_size *
+                                       config.num_updates_per_train_iter),
                         batch_length=config.mini_batch_length)
                     num_updates = 1
             return experience, batch_info, num_updates, mini_batch_size
@@ -1562,8 +1562,8 @@ def _replay():
 
             with record_time("time/offline_replay"):
                 offline_experience, offline_batch_info = self._offline_replay_buffer.get_batch(
-                    batch_size=(mini_batch_size *
-                                config.num_updates_per_train_iter),
+                    batch_size=int(mini_batch_size *
+                                   config.num_updates_per_train_iter),
                     batch_length=config.mini_batch_length)
             # train hybrid
             with record_time("time/offline_train"):
@@ -1641,14 +1641,26 @@ def _train_experience(self,
             torch.cuda.empty_cache()
 
         indices = None
+        # In the case where ``num_updates<1``, we will skip some mini-batches (only
+        # applicable if ``mini_batch_size<batch_size``).
+        training_fraction = min(1, num_updates)
+        training_every_n_batches = int(1. / training_fraction)
+        num_updates = int(max(1, np.ceil(num_updates)))
+        batches = 0
         for u in range(num_updates):
             if mini_batch_size < batch_size:
                 indices = torch.randperm(batch_size,
                                          device=experience.step_type.device)
             for b in range(0, batch_size, mini_batch_size):
+                if (b % (training_every_n_batches * mini_batch_size)) != 0:
+                    continue
 
-                is_last_mini_batch = (u == num_updates - 1
-                                      and b + mini_batch_size >= batch_size)
+                batches += 1
+
+                is_last_mini_batch = (
+                    u == num_updates - 1
+                    and b + mini_batch_size * training_every_n_batches
+                    >= batch_size)
                 do_summary = alf.summary.should_record_summaries() and (
                     is_last_mini_batch or update_counter_every_mini_batch)
 
@@ -1674,7 +1686,7 @@ def _train_experience(self,
                     # These are no longer used, release them to reduce memory usage.
                     del exp, train_info, loss_info, params
 
-        train_steps = batch_size * mini_batch_length * num_updates
+        train_steps = mini_batch_length * batches * mini_batch_size
         return train_steps
 
     def _prepare_experience_data(self,

alf/algorithms/config.py

@@ -280,8 +280,15 @@ def __init__(self,
             initial_collect_steps (int): if positive, number of steps each single
                 environment steps before perform first update. Only used
                 by ``OffPolicyAlgorithm``.
-            num_updates_per_train_iter (int): number of optimization steps for
-                one iteration. Only used by ``OffPolicyAlgorithm``.
+            num_updates_per_train_iter (float): number of optimization steps for
+                one iteration. Only used by ``OffPolicyAlgorithm``. If
+                ``config.whole_replay_buffer_training`` is False or doing offline training,
+                this is essentially a multiplier applied to ``config.mini_batch_size``.
+                Otherwise if this value is greater than 1, ``np.ceil(num_updates_per_train_iter)``
+                updates is performed; if this value is less than 1, it represents
+                the fraction of mini-batches from the entire buffer to train on. The latter
+                case is helpful if we want to shorten the training time for each
+                iteration.
             mini_batch_size (int): number of sequences for each minibatch. If None,
                 it's set to the replayer's ``batch_size``. Only used by
                 ``OffPolicyAlgorithm``.

alf/algorithms/distributed_off_policy_algorithm.py

@@ -16,6 +16,7 @@
 from typing import Callable
 import time
 import io
+import numpy as np
 import queue
 import random
 import threading
@@ -553,7 +554,8 @@ def _train_iter_off_policy(self):
             time.sleep(0.01)
 
         steps = super()._train_iter_off_policy()
-        self._total_updates += self._config.num_updates_per_train_iter
+        self._total_updates += int(
+            np.ceil(self._config.num_updates_per_train_iter))
 
         with record_time("time/trainer_send_params_to_unroller"):
             if (self._total_updates %

alf/utils/summary_utils.py

@@ -119,6 +119,8 @@ def summarize_variables(name_and_params, with_histogram=True):
         with_histogram (bool): If True, generate histogram.
     """
     for var_name, var in name_and_params:
+        if var.grad is None:
+            continue
         var_values = var
         if with_histogram and torch.all(torch.isfinite(var_values)):
             # Need to make sure all values are finite to avoid the histogram range