[bug-fix] Use correct memories for LSTM SAC (#5228)

* Use correct memories for LSTM SAC

* Add some comments

(cherry picked from commit 7077302)

Author: Ervin T

ml-agents/mlagents/trainers/sac/optimizer_torch.py

@@ -497,30 +497,17 @@ def update(self, batch: AgentBuffer, num_sequences: int) -> Dict[str, float]:
                 0, len(batch[BufferKey.CRITIC_MEMORY]), self.policy.sequence_length
             )
         ]
-        offset = 1 if self.policy.sequence_length > 1 else 0
-        next_value_memories_list = [
-            ModelUtils.list_to_tensor(
-                batch[BufferKey.CRITIC_MEMORY][i]
-            )  # only pass value part of memory to target network
-            for i in range(
-                offset, len(batch[BufferKey.CRITIC_MEMORY]), self.policy.sequence_length
-            )
-        ]
 
         if len(memories_list) > 0:
             memories = torch.stack(memories_list).unsqueeze(0)
             value_memories = torch.stack(value_memories_list).unsqueeze(0)
-            next_value_memories = torch.stack(next_value_memories_list).unsqueeze(0)
         else:
             memories = None
             value_memories = None
-            next_value_memories = None
 
         # Q and V network memories are 0'ed out, since we don't have them during inference.
         q_memories = (
-            torch.zeros_like(next_value_memories)
-            if next_value_memories is not None
-            else None
+            torch.zeros_like(value_memories) if value_memories is not None else None
         )
 
         # Copy normalizers from policy
@@ -568,6 +555,18 @@ def update(self, batch: AgentBuffer, num_sequences: int) -> Dict[str, float]:
             q1_stream, q2_stream = q1_out, q2_out
 
         with torch.no_grad():
+            # Since we didn't record the next value memories, evaluate one step in the critic to
+            # get them.
+            if value_memories is not None:
+                # Get the first observation in each sequence
+                just_first_obs = [
+                    _obs[:: self.policy.sequence_length] for _obs in current_obs
+                ]
+                _, next_value_memories = self._critic.critic_pass(
+                    just_first_obs, value_memories, sequence_length=1
+                )
+            else:
+                next_value_memories = None
             target_values, _ = self.target_network(
                 next_obs,
                 memories=next_value_memories,