Team manager prototype (#4850)

* remove group id

* very rough sketch for TeamManager interface

* add team manager id to proto

* team manager for hallway

* add manager to hallway

* send and process team manager id

* remove print

* small cleanup

Co-authored-by: Chris Elion <[email protected]>

Author: Ruo-Ping Dong

Project/Assets/ML-Agents/Examples/Hallway/Scripts/HallwayCollabAgent.cs

@@ -15,6 +15,17 @@ public class HallwayCollabAgent : HallwayAgent
 
     [HideInInspector]
     public int selection = 0;
+
+    public override void Initialize()
+    {
+        base.Initialize();
+        if (isSpotter)
+        {
+            var teamManager = new HallwayTeamManager();
+            SetTeamManager(teamManager);
+            teammate.SetTeamManager(teamManager);
+        }
+    }
     public override void OnEpisodeBegin()
     {
         m_Message = -1;

Project/Assets/ML-Agents/Examples/Hallway/Scripts/HallwayTeamManager.cs

@@ -0,0 +1,25 @@
+using System.Collections.Generic;
+using Unity.MLAgents;
+using Unity.MLAgents.Extensions.Teams;
+using Unity.MLAgents.Sensors;
+
+public class HallwayTeamManager : BaseTeamManager
+{
+    List<Agent> m_AgentList = new List<Agent> { };
+
+
+    public override void RegisterAgent(Agent agent)
+    {
+        m_AgentList.Add(agent);
+    }
+
+    public override void OnAgentDone(Agent agent, Agent.DoneReason doneReason, List<ISensor> sensors)
+    {
+        agent.SendDoneToTrainer();
+    }
+
+    public override void AddTeamReward(float reward)
+    {
+
+    }
+}

Project/Assets/ML-Agents/Examples/Hallway/Scripts/HallwayTeamManager.cs.meta

@@ -0,0 +1,35 @@
+using System.Collections.Generic;
+using Unity.MLAgents;
+using Unity.MLAgents.Sensors;
+
+namespace Unity.MLAgents.Extensions.Teams
+{
+    public class BaseTeamManager : ITeamManager
+    {
+        readonly string m_Id = System.Guid.NewGuid().ToString();
+
+        public virtual void RegisterAgent(Agent agent)
+        {
+            throw new System.NotImplementedException();
+        }
+
+        public virtual void OnAgentDone(Agent agent, Agent.DoneReason doneReason, List<ISensor> sensors)
+        {
+            // Possible implementation - save reference to Agent's IPolicy so that we can repeatedly
+            // call IPolicy.RequestDecision on behalf of the Agent after it's dead
+            // If so, we'll need dummy sensor impls with the same shape as the originals.
+            throw new System.NotImplementedException();
+        }
+
+        public virtual void AddTeamReward(float reward)
+        {
+
+        }
+
+        public string GetId()
+        {
+            return m_Id;
+        }
+
+    }
+}

Project/Assets/ML-Agents/Examples/Hallway/TFModels/HallwayCollab.onnx.meta

@@ -26,7 +26,6 @@ internal class BehaviorParametersEditor : UnityEditor.Editor
         const string k_InferenceDeviceName = "m_InferenceDevice";
         const string k_BehaviorTypeName = "m_BehaviorType";
         const string k_TeamIdName = "TeamId";
-        const string k_GroupIdName = "GroupId";
         const string k_UseChildSensorsName = "m_UseChildSensors";
         const string k_ObservableAttributeHandlingName = "m_ObservableAttributeHandling";
 
@@ -68,7 +67,6 @@ public override void OnInspectorGUI()
             }
             needPolicyUpdate = needPolicyUpdate || EditorGUI.EndChangeCheck();
 
-            EditorGUILayout.PropertyField(so.FindProperty(k_GroupIdName));
             EditorGUILayout.PropertyField(so.FindProperty(k_TeamIdName));
             EditorGUI.BeginDisabledGroup(!EditorUtilities.CanUpdateModelProperties());
             {

com.unity.ml-agents.extensions/Runtime/Teams.meta

@@ -50,6 +50,11 @@ internal struct AgentInfo
         /// </summary>
         public int episodeId;
 
+        /// <summary>
+        /// Team Manager identifier.
+        /// </summary>
+        public string teamManagerId;
+
         public void ClearActions()
         {
             storedActions.Clear();
@@ -312,6 +317,8 @@ internal struct AgentParameters
         /// </summary>
         float[] m_LegacyActionCache;
 
+        private ITeamManager m_TeamManager;
+
         /// <summary>
         /// Called when the attached [GameObject] becomes enabled and active.
         /// [GameObject]: https://docs.unity3d.com/Manual/GameObjects.html
@@ -443,6 +450,11 @@ public void LazyInitialize()
                 new int[m_ActuatorManager.NumDiscreteActions]
             );
 
+            if (m_TeamManager != null)
+            {
+                m_Info.teamManagerId = m_TeamManager.GetId();
+            }
+
             // The first time the Academy resets, all Agents in the scene will be
             // forced to reset through the <see cref="AgentForceReset"/> event.
             // To avoid the Agent resetting twice, the Agents will not begin their
@@ -459,7 +471,7 @@ public void LazyInitialize()
         /// <summary>
         /// The reason that the Agent has been set to "done".
         /// </summary>
-        enum DoneReason
+        public enum DoneReason
         {
             /// <summary>
             /// The episode was ended manually by calling <see cref="EndEpisode"/>.
@@ -535,9 +547,17 @@ void NotifyAgentDone(DoneReason doneReason)
                 }
             }
             // Request the last decision with no callbacks
-            // We request a decision so Python knows the Agent is done immediately
-            m_Brain?.RequestDecision(m_Info, sensors);
-            ResetSensors();
+            if (m_TeamManager != null)
+            {
+                // Send final observations to TeamManager if it exists.
+                // The TeamManager is responsible to keeping track of the Agent after it's
+                // done, including propagating any "posthumous" rewards.
+                m_TeamManager.OnAgentDone(this, doneReason, sensors);
+            }
+            else
+            {
+                SendDoneToTrainer();
+            }
 
             // We also have to write any to any DemonstationStores so that they get the "done" flag.
             foreach (var demoWriter in DemonstrationWriters)
@@ -560,6 +580,13 @@ void NotifyAgentDone(DoneReason doneReason)
             m_Info.storedActions.Clear();
         }
 
+        public void SendDoneToTrainer()
+        {
+            // We request a decision so Python knows the Agent is done immediately
+            m_Brain?.RequestDecision(m_Info, sensors);
+            ResetSensors();
+        }
+
         /// <summary>
         /// Updates the Model assigned to this Agent instance.
         /// </summary>
@@ -1344,5 +1371,12 @@ void DecideAction()
             m_Info.CopyActions(actions);
             m_ActuatorManager.UpdateActions(actions);
         }
+
+        public void SetTeamManager(ITeamManager teamManager)
+        {
+            m_TeamManager = teamManager;
+            m_Info.teamManagerId = teamManager?.GetId();
+            teamManager?.RegisterAgent(this);
+        }
     }
 }

com.unity.ml-agents.extensions/Runtime/Teams/BaseTeamManager.cs

@@ -67,6 +67,11 @@ public static AgentInfoProto ToAgentInfoProto(this AgentInfo ai)
                 agentInfoProto.ActionMask.AddRange(ai.discreteActionMasks);
             }
 
+            if (ai.teamManagerId != null)
+            {
+                agentInfoProto.TeamManagerId = ai.teamManagerId;
+            }
+
             return agentInfoProto;
         }