Fix padding in optimizer value estimate

ml-agents/mlagents/trainers/optimizer/torch_optimizer.py

 from mlagents.torch_utils import torch
 import numpy as np
 import math
+from collections import defaultdict
 
 from mlagents.trainers.buffer import AgentBuffer, AgentBufferField
 from mlagents.trainers.trajectory import ObsUtil
         # Compute the number of elements in this padded seq.
         leftover = num_experiences % self.policy.sequence_length
 
-        # Compute values for the potentially truncated initial sequence
-        seq_obs = []
-
-        first_seq_len = leftover if leftover > 0 else self.policy.sequence_length
-        for _obs in tensor_obs:
-            first_seq_obs = _obs[0:first_seq_len]
-            seq_obs.append(first_seq_obs)
-
-        # For the first sequence, the initial memory should be the one at the
-        # beginning of this trajectory.
-        for _ in range(first_seq_len):
-            all_next_memories.append(ModelUtils.to_numpy(initial_memory.squeeze()))
-
-        init_values, _mem = self.critic.critic_pass(
-            seq_obs, initial_memory, sequence_length=first_seq_len
-        )
-        all_values = {
-            signal_name: [init_values[signal_name]]
-            for signal_name in init_values.keys()
-        }
-
+        all_values: Dict[str, List[np.ndarray]] = defaultdict(list)
+        _mem = initial_memory
-            1, math.ceil((num_experiences) / (self.policy.sequence_length))
+            0, math.floor((num_experiences) / (self.policy.sequence_length))
-            start = seq_num * self.policy.sequence_length - (
-                self.policy.sequence_length - leftover
-            )
-            end = (seq_num + 1) * self.policy.sequence_length - (
-                self.policy.sequence_length - leftover
-            )
+            start = seq_num * self.policy.sequence_length
+            end = (seq_num + 1) * self.policy.sequence_length
+
             for _obs in tensor_obs:
                 seq_obs.append(_obs[start:end])
             values, _mem = self.critic.critic_pass(
                 all_values[signal_name].append(_val)
+
+        # Compute values for the potentially truncated last sequence
+        seq_obs = []
+
+        last_seq_len = leftover
+        if last_seq_len > 0:
+            for _obs in tensor_obs:
+                last_seq_obs = _obs[0:last_seq_len]
+                seq_obs.append(last_seq_obs)
+
+            # For the first sequence, the initial memory should be the one at the
+            # beginning of this trajectory.
+            for _ in range(last_seq_len):
+                all_next_memories.append(ModelUtils.to_numpy(_mem.squeeze()))
+
+            last_values, _mem = self.critic.critic_pass(
+                seq_obs, _mem, sequence_length=last_seq_len
+            )
+            for signal_name, _val in last_values.items():
+                all_values[signal_name].append(_val)
+
         # Create one tensor per reward signal
         all_value_tensors = {
             signal_name: torch.cat(value_list, dim=0)