Make num_updates_per_train_iter scheduleable

Author: QuantuMope

alf/algorithms/algorithm.py

@@ -1503,6 +1503,10 @@ def train_from_replay_buffer(self, update_global_counter=False):
                 ``config.update_counter_every_mini_batch=True``.
         """
         config: TrainerConfig = self._config
+        num_updates_per_train_iter = int(config.num_updates_per_train_iter())
+
+        print(num_updates_per_train_iter)
+        print("-----")
 
         # returns 0 if haven't started training yet, when ``_replay_buffer`` is
         # not None and the number of samples in the buffer is less than
@@ -1535,7 +1539,7 @@ def _replay(batch_multiplier):
                 if config.whole_replay_buffer_training:
                     experience, batch_info = self._replay_buffer.gather_all(
                         ignore_earliest_frames=True)
-                    num_updates = config.num_updates_per_train_iter
+                    num_updates = num_updates_per_train_iter
                 else:
                     assert config.mini_batch_length is not None, (
                         "No mini_batch_length is specified for off-policy training"
@@ -1550,7 +1554,7 @@ def _replay(batch_multiplier):
             if (config.sample_mini_batch_per_update
                     and not config.whole_replay_buffer_training):
                 train_steps = 0
-                for _ in range(config.num_updates_per_train_iter):
+                for _ in range(num_updates_per_train_iter):
                     experience, batch_info, num_updates, mini_batch_size = _replay(
                         1)
                     with record_time("time/train"):
@@ -1567,7 +1571,7 @@ def _replay(batch_multiplier):
                 return train_steps
             else:
                 experience, batch_info, num_updates, mini_batch_size = _replay(
-                    config.num_updates_per_train_iter)
+                    num_updates_per_train_iter)
                 with record_time("time/train"):
                     return self._train_experience(
                         experience,
@@ -1597,7 +1601,7 @@ def _replay(batch_multiplier):
             if (config.sample_mini_batch_per_update
                     and not config.whole_replay_buffer_training):
                 train_steps = 0
-                for _ in range(config.num_updates_per_train_iter):
+                for _ in range(num_updates_per_train_iter):
                     if self._RL_train:
                         experience, batch_info, num_updates, mini_batch_size = _replay(
                             1)
@@ -1623,7 +1627,7 @@ def _replay(batch_multiplier):
             else:
                 if self._RL_train:
                     experience, batch_info, num_updates, mini_batch_size = _replay(
-                        config.num_updates_per_train_iter)
+                        num_updates_per_train_iter)
                 else:
                     experience = None
                     batch_info = None
@@ -1633,7 +1637,7 @@ def _replay(batch_multiplier):
                 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),
+                                    num_updates_per_train_iter),
                         batch_length=config.mini_batch_length)
                 # train hybrid
                 with record_time("time/offline_train"):

alf/algorithms/config.py

@@ -285,8 +285,8 @@ 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 (int|Scheduler): number of optimization
+                steps for one iteration. Only used by ``OffPolicyAlgorithm``.
             sample_mini_batch_per_update (bool): If True and
                 ``whole_replay_buffer_training`` is False, sample one minibatch
                 per update instead of sampling ``num_updates_per_train_iter``
@@ -440,7 +440,8 @@ def __init__(self,
         self.summarize_action_distributions = summarize_action_distributions
         self.summarize_output = summarize_output
         self.initial_collect_steps = initial_collect_steps
-        self.num_updates_per_train_iter = num_updates_per_train_iter
+        self.num_updates_per_train_iter = as_scheduler(
+            num_updates_per_train_iter)
         self.sample_mini_batch_per_update = sample_mini_batch_per_update
         self.mini_batch_length = mini_batch_length
         self.mini_batch_size = mini_batch_size

alf/algorithms/distributed_off_policy_algorithm.py

@@ -560,7 +560,7 @@ 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(self._config.num_updates_per_train_iter())
 
         with record_time("time/trainer_send_params_to_unroller"):
             if (self._total_updates %

alf/algorithms/muzero_representation_learner.py

@@ -1027,7 +1027,8 @@ def __init__(self,
             updated.clear_replay_buffer = False
             updated.mini_batch_length = training_options.mini_batch_length
             updated.mini_batch_size = training_options.mini_batch_size
-            updated.num_updates_per_train_iter = training_options.num_updates_per_train_iter
+            updated.num_updates_per_train_iter = as_scheduler(
+                training_options.num_updates_per_train_iter)
             updated.replay_buffer_length = training_options.replay_buffer_length
             updated.initial_collect_steps = training_options.initial_collect_steps
             updated.priority_replay = training_options.priority_replay

alf/algorithms/ppg/ppg_aux_algorithm.py

@@ -24,6 +24,7 @@
 from alf.data_structures import TimeStep, AlgStep, LossInfo
 from alf.experience_replayers.replay_buffer import ReplayBuffer
 from alf.utils import dist_utils
+from alf.utils.schedulers import as_scheduler
 from alf.tensor_specs import TensorSpec
 
 # Data structure to store the options for PPG's auxiliary phase
@@ -109,7 +110,8 @@ def __init__(self,
         updated_config.mini_batch_length = (aux_options.mini_batch_length
                                             or config.unroll_length)
         updated_config.mini_batch_size = aux_options.mini_batch_size
-        updated_config.num_updates_per_train_iter = aux_options.num_updates_per_train_iter
+        updated_config.num_updates_per_train_iter = as_scheduler(
+            aux_options.num_updates_per_train_iter)
 
         # Since we are going to store already-transformed experience in the
         # replay buffer, the aux algorithm shall not inherit the data

alf/algorithms/rlpd_algorithm.py

@@ -186,6 +186,9 @@ def __init__(self,
         else:
             total_utd = alf.config_util.get_config_value(
                 "num_updates_per_train_iter")
+            if callable(total_utd):
+                total_utd = total_utd()
+            total_utd = int(total_utd)
             if critic_utd is not None:
                 assert critic_utd < total_utd, (
                     "critic_utd should be less than num_updates_per_train_iter"