POCA trainer

com.unity.ml-agents/Runtime/SimpleMultiAgentGroup.cs

@@ -7,7 +7,7 @@ namespace Unity.MLAgents
     /// <summary>
     /// A basic class implementation of MultiAgentGroup.
     /// </summary>
-    internal class SimpleMultiAgentGroup : IMultiAgentGroup, IDisposable
+    public class SimpleMultiAgentGroup : IMultiAgentGroup, IDisposable
     {
         readonly int m_Id = MultiAgentGroupIdCounter.GetGroupId();
         HashSet<Agent> m_Agents = new HashSet<Agent>();

ml-agents/mlagents/trainers/buffer.py

@@ -35,6 +35,7 @@ class BufferKey(enum.Enum):
     MASKS = "masks"
     MEMORY = "memory"
     CRITIC_MEMORY = "critic_memory"
+    BASELINE_MEMORY = "poca_baseline_memory"
     PREV_ACTION = "prev_action"
 
     ADVANTAGES = "advantages"
@@ -63,6 +64,7 @@ class RewardSignalKeyPrefix(enum.Enum):
     VALUE_ESTIMATES = "value_estimates"
     RETURNS = "returns"
     ADVANTAGE = "advantage"
+    BASELINES = "baselines"
 
 
 AgentBufferKey = Union[
@@ -87,6 +89,10 @@ def returns_key(name: str) -> AgentBufferKey:
     def advantage_key(name: str) -> AgentBufferKey:
         return RewardSignalKeyPrefix.ADVANTAGE, name
 
+    @staticmethod
+    def baseline_estimates_key(name: str) -> AgentBufferKey:
+        return RewardSignalKeyPrefix.BASELINES, name
+
 
 class AgentBufferField(list):
     """

ml-agents/mlagents/trainers/ghost/trainer.py

@@ -192,7 +192,9 @@ def _process_trajectory(self, trajectory: Trajectory) -> None:
         """
         if trajectory.done_reached:
             # Assumption is that final reward is >0/0/<0 for win/draw/loss
-            final_reward = trajectory.steps[-1].reward
+            final_reward = (
+                trajectory.steps[-1].reward + trajectory.steps[-1].group_reward
+            )
             result = 0.5
             if final_reward > 0:
                 result = 1.0

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

@@ -10,7 +10,11 @@
 
 from mlagents.trainers.policy.torch_policy import TorchPolicy
 from mlagents.trainers.optimizer import Optimizer
-from mlagents.trainers.settings import TrainerSettings
+from mlagents.trainers.settings import (
+    TrainerSettings,
+    RewardSignalSettings,
+    RewardSignalType,
+)
 from mlagents.trainers.torch.utils import ModelUtils
 
 
@@ -44,7 +48,9 @@ def critic(self):
     def update(self, batch: AgentBuffer, num_sequences: int) -> Dict[str, float]:
         pass
 
-    def create_reward_signals(self, reward_signal_configs):
+    def create_reward_signals(
+        self, reward_signal_configs: Dict[RewardSignalType, RewardSignalSettings]
+    ) -> None:
         """
         Create reward signals
         :param reward_signal_configs: Reward signal config.
@@ -56,7 +62,7 @@ def create_reward_signals(self, reward_signal_configs):
             )
 
     def _evaluate_by_sequence(
-        self, tensor_obs: List[torch.Tensor], initial_memory: np.ndarray
+        self, tensor_obs: List[torch.Tensor], initial_memory: torch.Tensor
     ) -> Tuple[Dict[str, torch.Tensor], AgentBufferField, torch.Tensor]:
         """
         Evaluate a trajectory sequence-by-sequence, assembling the result. This enables us to get the
@@ -78,10 +84,8 @@ def _evaluate_by_sequence(
         # Compute values for the potentially truncated initial sequence
         seq_obs = []
 
-        first_seq_len = self.policy.sequence_length
+        first_seq_len = leftover if leftover > 0 else self.policy.sequence_length
         for _obs in tensor_obs:
-            if leftover > 0:
-                first_seq_len = leftover
             first_seq_obs = _obs[0:first_seq_len]
             seq_obs.append(first_seq_obs)
 
@@ -106,13 +110,13 @@ def _evaluate_by_sequence(
             seq_obs = []
             for _ in range(self.policy.sequence_length):
                 all_next_memories.append(ModelUtils.to_numpy(_mem.squeeze()))
+            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
+            )
             for _obs in tensor_obs:
-                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
-                )
                 seq_obs.append(_obs[start:end])
             values, _mem = self.critic.critic_pass(
                 seq_obs, _mem, sequence_length=self.policy.sequence_length