Support multi-step td

Author: Haichao Zhang

alf/algorithms/sac_algorithm.py

@@ -28,6 +28,8 @@
 from alf.algorithms.off_policy_algorithm import OffPolicyAlgorithm
 from alf.algorithms.one_step_loss import OneStepTDLoss
 from alf.algorithms.rl_algorithm import RLAlgorithm
+from alf.experience_replayers.replay_buffer import ReplayBuffer
+from alf.nest.utils import convert_device
 from alf.data_structures import TimeStep, Experience, LossInfo, namedtuple
 from alf.data_structures import AlgStep, StepType
 from alf.nest import nest
@@ -1073,14 +1075,21 @@ def _calc_critic_loss(self, info: SacInfo):
 
     def preprocess_experience(self, time_step: TimeStep, rollout_info: SacInfo,
                               batch_info):
-
-        _, mini_batch_length = time_step.step_type.shape
-        discounted_return = batch_info.discounted_return
-        discounted_return = discounted_return.unsqueeze(1).expand(
-            -1, mini_batch_length)
-        rollout_info = alf.nest.set_field(
-            rollout_info, 'discounted_return',
-            discounted_return if self._use_mc_return else ())
+        if self._use_mc_return:
+            assert batch_info != ()
+            replay_buffer: ReplayBuffer = batch_info.replay_buffer
+            mini_batch_length = time_step.step_type.shape[1]
+
+            with alf.device(replay_buffer.device):
+                # [B, 1]
+                positions = convert_device(batch_info.positions).unsqueeze(-1)
+                # [B, 1]
+                env_ids = convert_device(batch_info.env_ids).unsqueeze(-1)
+                # [B, T]
+                positions = positions + torch.arange(mini_batch_length)
+                discounted_return = replay_buffer.get_discounted_return(env_ids=env_ids, positions=positions)
+                discounted_return = convert_device(discounted_return)
+            rollout_info = rollout_info._replace(discounted_return=discounted_return)
         return time_step, rollout_info
 
     def _trainable_attributes_to_ignore(self):

alf/examples/cart_pole_mc_return.png

@@ -20,10 +20,14 @@
 
 default_return = -1000
 use_mc_return = True
+mini_batch_length = 2 # set to a value > 2 for multi-step learning
 
 alf.config("ReplayBuffer",
            keep_episodic_info=True,
            record_episodic_return=True,
            default_return=default_return)
 alf.config("TDLoss", default_return=default_return)
 alf.config("SacAlgorithm", use_mc_return=use_mc_return)
+
+alf.config('TrainerConfig',
+           mini_batch_length=mini_batch_length)

alf/examples/sac_cart_pole_mc_return_conf.py

@@ -839,6 +839,25 @@ def get_field(self, field_name, env_ids, positions):
         result = alf.nest.map_structure(lambda x: x[indices], field)
         return convert_device(result)
 
+    def get_discounted_return(self, env_ids, positions):
+        """Get discounted return from the replay buffer by ``env_ids`` and ``positions``.
+
+        Args:
+            env_ids (Tensor): 1-D int64 Tensor.
+            positions (Tensor): 1-D int64 Tensor with same shape as ``env_ids``.
+                These positions should be obtained from the BatchInfo returned
+                by ``get_batch()``.
+        Returns:
+            Tensor: with the same shape as broadcasted shape of env_ids and positions
+        """
+        current_pos = self._current_pos[env_ids]
+        assert torch.all(positions < current_pos), "Invalid positions"
+        assert torch.all(positions >= current_pos -
+                         self._max_length), "Invalid positions"
+        indices = (env_ids, self.circular(positions))
+        result = self._episodic_discounted_return[indices]
+        return convert_device(result)
+
     @property
     def total_size(self):
         """Total size from all environments."""