Develop add training area replicator (Unity-Technologies#5568)

* Added training area replicator to com.unity.ml-agents package. * Added num_areas to Unity RL Initialization proto. Added cli and config file support for num_areas. * Changed training area replicator to size grid automatically from number of areas. * Added tests for the training area replicator. * Added setup for tests for the training area replicator. * Added comments and updated create tutorial for training area replicator. * Updated CHANGELOG. * Fixed some failing tests. * Update com.unity.ml-agents/CHANGELOG.md Co-authored-by: Henry Peteet <[email protected]> * Update docs/Learning-Environment-Create-New.md Co-authored-by: Henry Peteet <[email protected]> * Update com.unity.ml-agents/Runtime/Areas/TrainingAreaReplicator.cs Co-authored-by: Henry Peteet <[email protected]> * Addressed CR comments. Co-authored-by: Miguel Alonso Jr <miguelalonsojr> Co-authored-by: Henry Peteet <[email protected]>
hubbardp · Oct 13, 2021 · 003e8ae · 003e8ae
1 parent f5e4a2b
commit 003e8ae
Show file tree

Hide file tree

Showing 24 changed files with 295 additions and 13 deletions.
diff --git a/com.unity.ml-agents/CHANGELOG.md b/com.unity.ml-agents/CHANGELOG.md
@@ -10,6 +10,7 @@ and this project adheres to
 ### Major Changes
 
 #### com.unity.ml-agents / com.unity.ml-agents.extensions (C#)
+- Added a new feature to replicate training areas dynamically during runtime. (#5568)
 
 #### ml-agents / ml-agents-envs / gym-unity (Python)
 

diff --git a/com.unity.ml-agents/Runtime/Academy.cs b/com.unity.ml-agents/Runtime/Academy.cs
@@ -183,6 +183,11 @@ public int InferenceSeed
             set { m_InferenceSeed = value; }
         }
 
+        // Number of training areas to instantiate
+        int m_NumAreas;
+
+        public int NumAreas => m_NumAreas;
+
         /// <summary>
         /// Returns the RLCapabilities of the python client that the unity process is connected to.
         /// </summary>
@@ -451,6 +456,7 @@ out var unityRlInitParameters
                         UnityEngine.Random.InitState(unityRlInitParameters.seed);
                         // We might have inference-only Agents, so set the seed for them too.
                         m_InferenceSeed = unityRlInitParameters.seed;
+                        m_NumAreas = unityRlInitParameters.numAreas;
                         TrainerCapabilities = unityRlInitParameters.TrainerCapabilities;
                         TrainerCapabilities.WarnOnPythonMissingBaseRLCapabilities();
                     }

diff --git a/com.unity.ml-agents/Runtime/Areas.meta b/com.unity.ml-agents/Runtime/Areas.meta
diff --git a/com.unity.ml-agents/Runtime/Areas/TrainingAreaReplicator.cs b/com.unity.ml-agents/Runtime/Areas/TrainingAreaReplicator.cs
@@ -0,0 +1,88 @@
+using System;
+using Unity.Mathematics;
+using UnityEngine;
+
+namespace Unity.MLAgents.Areas
+{
+    /// <summary>
+    /// The Training Ares Replicator allows for a training area object group to be replicated dynamically during runtime.
+    /// </summary>
+    public class TrainingAreaReplicator : MonoBehaviour
+    {
+        public GameObject baseArea;
+        public int numAreas = 1;
+        public float separation = 10f;
+
+        int3 m_GridSize = new int3(1, 1, 1);
+        int m_areaCount = 0;
+        string m_TrainingAreaName;
+
+        public int3 GridSize => m_GridSize;
+        public string TrainingAreaName => m_TrainingAreaName;
+
+        public void Awake()
+        {
+            // Computes the Grid Size on Awake
+            ComputeGridSize();
+            // Sets the TrainingArea name to the name of the base area.
+            m_TrainingAreaName = baseArea.name;
+        }
+
+        public void OnEnable()
+        {
+            // Adds the training are replicas during OnEnable to ensure they are added before the Academy begins its work.
+            AddEnvironments();
+        }
+
+        /// <summary>
+        /// Computes the Grid Size for replicating the training area.
+        /// </summary>
+        void ComputeGridSize()
+        {
+            // check if running inference, if so, use the num areas set through the component,
+            // otherwise, pull it from the academy
+            if (Academy.Instance.Communicator != null)
+                numAreas = Academy.Instance.NumAreas;
+
+            var rootNumAreas = Mathf.Pow(numAreas, 1.0f / 3.0f);
+            m_GridSize.x = Mathf.CeilToInt(rootNumAreas);
+            m_GridSize.y = Mathf.CeilToInt(rootNumAreas);
+            var zSize = Mathf.CeilToInt((float)numAreas / (m_GridSize.x * m_GridSize.y));
+            m_GridSize.z = zSize == 0 ? 1 : zSize;
+        }
+
+        /// <summary>
+        /// Adds replicas of the training area to the scene.
+        /// </summary>
+        /// <exception cref="UnityAgentsException"></exception>
+        void AddEnvironments()
+        {
+            if (numAreas > m_GridSize.x * m_GridSize.y * m_GridSize.z)
+            {
+                throw new UnityAgentsException("The number of training areas that you have specified exceeds the size of the grid.");
+            }
+
+            for (int z = 0; z < m_GridSize.z; z++)
+            {
+                for (int y = 0; y < m_GridSize.y; y++)
+                {
+                    for (int x = 0; x < m_GridSize.x; x++)
+                    {
+                        if (m_areaCount == 0)
+                        {
+                            // Skip this first area since it already exists.
+                            m_areaCount = 1;
+                        }
+                        else if (m_areaCount < numAreas)
+                        {
+                            m_areaCount++;
+                            var area = Instantiate(baseArea, new Vector3(x * separation, y * separation, z * separation), Quaternion.identity);
+                            area.name = m_TrainingAreaName;
+                        }
+                    }
+                }
+            }
+        }
+    }
+}
+
diff --git a/com.unity.ml-agents/Runtime/Areas/TrainingAreaReplicator.cs.meta b/com.unity.ml-agents/Runtime/Areas/TrainingAreaReplicator.cs.meta
diff --git a/com.unity.ml-agents/Runtime/Communicator/GrpcExtensions.cs b/com.unity.ml-agents/Runtime/Communicator/GrpcExtensions.cs
@@ -298,6 +298,7 @@ public static UnityRLInitParameters ToUnityRLInitParameters(this UnityRLInitiali
             return new UnityRLInitParameters
             {
                 seed = inputProto.Seed,
+                numAreas = inputProto.NumAreas,
                 pythonLibraryVersion = inputProto.PackageVersion,
                 pythonCommunicationVersion = inputProto.CommunicationVersion,
                 TrainerCapabilities = inputProto.Capabilities.ToRLCapabilities()

diff --git a/com.unity.ml-agents/Runtime/Communicator/ICommunicator.cs b/com.unity.ml-agents/Runtime/Communicator/ICommunicator.cs
@@ -39,6 +39,11 @@ internal struct UnityRLInitParameters
         /// </summary>
         public int seed;
 
+        /// <summary>
+        /// The number of areas to replicate if Training Area Replication is used in the scene.
+        /// </summary>
+        public int numAreas;
+
         /// <summary>
         /// The library version of the python process.
         /// </summary>

diff --git a/com.unity.ml-agents/Runtime/Grpc/CommunicatorObjects/UnityRlInitializationInput.cs b/com.unity.ml-agents/Runtime/Grpc/CommunicatorObjects/UnityRlInitializationInput.cs
@@ -27,16 +27,16 @@ static UnityRlInitializationInputReflection() {
             "CkZtbGFnZW50c19lbnZzL2NvbW11bmljYXRvcl9vYmplY3RzL3VuaXR5X3Js",
             "X2luaXRpYWxpemF0aW9uX2lucHV0LnByb3RvEhRjb21tdW5pY2F0b3Jfb2Jq",
             "ZWN0cxo1bWxhZ2VudHNfZW52cy9jb21tdW5pY2F0b3Jfb2JqZWN0cy9jYXBh",
-            "YmlsaXRpZXMucHJvdG8irQEKH1VuaXR5UkxJbml0aWFsaXphdGlvbklucHV0",
+            "YmlsaXRpZXMucHJvdG8iwAEKH1VuaXR5UkxJbml0aWFsaXphdGlvbklucHV0",
             "UHJvdG8SDAoEc2VlZBgBIAEoBRIdChVjb21tdW5pY2F0aW9uX3ZlcnNpb24Y",
             "AiABKAkSFwoPcGFja2FnZV92ZXJzaW9uGAMgASgJEkQKDGNhcGFiaWxpdGll",
             "cxgEIAEoCzIuLmNvbW11bmljYXRvcl9vYmplY3RzLlVuaXR5UkxDYXBhYmls",
-            "aXRpZXNQcm90b0IlqgIiVW5pdHkuTUxBZ2VudHMuQ29tbXVuaWNhdG9yT2Jq",
-            "ZWN0c2IGcHJvdG8z"));
+            "aXRpZXNQcm90bxIRCgludW1fYXJlYXMYBSABKAVCJaoCIlVuaXR5Lk1MQWdl",
+            "bnRzLkNvbW11bmljYXRvck9iamVjdHNiBnByb3RvMw=="));
       descriptor = pbr::FileDescriptor.FromGeneratedCode(descriptorData,
           new pbr::FileDescriptor[] { global::Unity.MLAgents.CommunicatorObjects.CapabilitiesReflection.Descriptor, },
           new pbr::GeneratedClrTypeInfo(null, new pbr::GeneratedClrTypeInfo[] {
-            new pbr::GeneratedClrTypeInfo(typeof(global::Unity.MLAgents.CommunicatorObjects.UnityRLInitializationInputProto), global::Unity.MLAgents.CommunicatorObjects.UnityRLInitializationInputProto.Parser, new[]{ "Seed", "CommunicationVersion", "PackageVersion", "Capabilities" }, null, null, null)
+            new pbr::GeneratedClrTypeInfo(typeof(global::Unity.MLAgents.CommunicatorObjects.UnityRLInitializationInputProto), global::Unity.MLAgents.CommunicatorObjects.UnityRLInitializationInputProto.Parser, new[]{ "Seed", "CommunicationVersion", "PackageVersion", "Capabilities", "NumAreas" }, null, null, null)
           }));
     }
     #endregion
@@ -75,6 +75,7 @@ public UnityRLInitializationInputProto(UnityRLInitializationInputProto other) :
       communicationVersion_ = other.communicationVersion_;
       packageVersion_ = other.packageVersion_;
       Capabilities = other.capabilities_ != null ? other.Capabilities.Clone() : null;
+      numAreas_ = other.numAreas_;
       _unknownFields = pb::UnknownFieldSet.Clone(other._unknownFields);
     }
 
@@ -136,6 +137,20 @@ public string PackageVersion {
       }
     }
 
+    /// <summary>Field number for the "num_areas" field.</summary>
+    public const int NumAreasFieldNumber = 5;
+    private int numAreas_;
+    /// <summary>
+    /// The number of training areas to instantiate
+    /// </summary>
+    [global::System.Diagnostics.DebuggerNonUserCodeAttribute]
+    public int NumAreas {
+      get { return numAreas_; }
+      set {
+        numAreas_ = value;
+      }
+    }
+
     [global::System.Diagnostics.DebuggerNonUserCodeAttribute]
     public override bool Equals(object other) {
       return Equals(other as UnityRLInitializationInputProto);
@@ -153,6 +168,7 @@ public bool Equals(UnityRLInitializationInputProto other) {
       if (CommunicationVersion != other.CommunicationVersion) return false;
       if (PackageVersion != other.PackageVersion) return false;
       if (!object.Equals(Capabilities, other.Capabilities)) return false;
+      if (NumAreas != other.NumAreas) return false;
       return Equals(_unknownFields, other._unknownFields);
     }
 
@@ -163,6 +179,7 @@ public override int GetHashCode() {
       if (CommunicationVersion.Length != 0) hash ^= CommunicationVersion.GetHashCode();
       if (PackageVersion.Length != 0) hash ^= PackageVersion.GetHashCode();
       if (capabilities_ != null) hash ^= Capabilities.GetHashCode();
+      if (NumAreas != 0) hash ^= NumAreas.GetHashCode();
       if (_unknownFields != null) {
         hash ^= _unknownFields.GetHashCode();
       }
@@ -192,6 +209,10 @@ public void WriteTo(pb::CodedOutputStream output) {
         output.WriteRawTag(34);
         output.WriteMessage(Capabilities);
       }
+      if (NumAreas != 0) {
+        output.WriteRawTag(40);
+        output.WriteInt32(NumAreas);
+      }
       if (_unknownFields != null) {
         _unknownFields.WriteTo(output);
       }
@@ -212,6 +233,9 @@ public int CalculateSize() {
       if (capabilities_ != null) {
         size += 1 + pb::CodedOutputStream.ComputeMessageSize(Capabilities);
       }
+      if (NumAreas != 0) {
+        size += 1 + pb::CodedOutputStream.ComputeInt32Size(NumAreas);
+      }
       if (_unknownFields != null) {
         size += _unknownFields.CalculateSize();
       }
@@ -238,6 +262,9 @@ public void MergeFrom(UnityRLInitializationInputProto other) {
         }
         Capabilities.MergeFrom(other.Capabilities);
       }
+      if (other.NumAreas != 0) {
+        NumAreas = other.NumAreas;
+      }
       _unknownFields = pb::UnknownFieldSet.MergeFrom(_unknownFields, other._unknownFields);
     }
 
@@ -268,6 +295,10 @@ public void MergeFrom(pb::CodedInputStream input) {
             input.ReadMessage(capabilities_);
             break;
           }
+          case 40: {
+            NumAreas = input.ReadInt32();
+            break;
+          }
         }
       }
     }

diff --git a/com.unity.ml-agents/Runtime/Unity.ML-Agents.asmdef b/com.unity.ml-agents/Runtime/Unity.ML-Agents.asmdef
@@ -2,7 +2,8 @@
     "name": "Unity.ML-Agents",
     "references": [
         "Unity.Barracuda",
-        "Unity.ML-Agents.CommunicatorObjects"
+        "Unity.ML-Agents.CommunicatorObjects",
+        "Unity.Mathematics"
     ],
     "optionalUnityReferences": [],
     "includePlatforms": [],

diff --git a/com.unity.ml-agents/Tests/Editor/Areas.meta b/com.unity.ml-agents/Tests/Editor/Areas.meta
diff --git a/com.unity.ml-agents/Tests/Editor/Areas/TrainingAreaReplicatorTests.cs b/com.unity.ml-agents/Tests/Editor/Areas/TrainingAreaReplicatorTests.cs
@@ -0,0 +1,61 @@
+using System.Linq;
+using NUnit.Framework;
+using Unity.Mathematics;
+using Unity.MLAgents.Areas;
+using UnityEngine;
+
+namespace Unity.MLAgents.Tests.Areas
+{
+    [TestFixture]
+    public class TrainingAreaReplicatorTests
+    {
+        private TrainingAreaReplicator m_Replicator;
+
+        [SetUp]
+        public void Setup()
+        {
+            var gameObject = new GameObject();
+            var trainingArea = new GameObject();
+            trainingArea.name = "MyTrainingArea";
+            m_Replicator = gameObject.AddComponent<TrainingAreaReplicator>();
+            m_Replicator.baseArea = trainingArea;
+        }
+
+        private static object[] NumAreasCases =
+        {
+            new object[] {1},
+            new object[] {2},
+            new object[] {5},
+            new object[] {7},
+            new object[] {8},
+            new object[] {64},
+            new object[] {63},
+        };
+
+        [TestCaseSource(nameof(NumAreasCases))]
+        public void TestComputeGridSize(int numAreas)
+        {
+            m_Replicator.numAreas = numAreas;
+            m_Replicator.Awake();
+            m_Replicator.OnEnable();
+            var m_CorrectGridSize = int3.zero;
+            var m_RootNumAreas = Mathf.Pow(numAreas, 1.0f / 3.0f);
+            m_CorrectGridSize.x = Mathf.CeilToInt(m_RootNumAreas);
+            m_CorrectGridSize.y = Mathf.CeilToInt(m_RootNumAreas);
+            m_CorrectGridSize.z = Mathf.CeilToInt((float)numAreas / (m_CorrectGridSize.x * m_CorrectGridSize.y));
+            Assert.GreaterOrEqual(m_Replicator.GridSize.x * m_Replicator.GridSize.y * m_Replicator.GridSize.z, m_Replicator.numAreas);
+            Assert.AreEqual(m_CorrectGridSize, m_Replicator.GridSize);
+        }
+
+        [Test]
+        public void TestAddEnvironments()
+        {
+            m_Replicator.numAreas = 10;
+            m_Replicator.Awake();
+            m_Replicator.OnEnable();
+            var trainingAreas = Resources.FindObjectsOfTypeAll<GameObject>().Where(obj => obj.name == m_Replicator.TrainingAreaName);
+            Assert.AreEqual(10, trainingAreas.Count());
+
+        }
+    }
+}
diff --git a/com.unity.ml-agents/Tests/Editor/Areas/TrainingAreaReplicatorTests.cs.meta b/com.unity.ml-agents/Tests/Editor/Areas/TrainingAreaReplicatorTests.cs.meta
diff --git a/com.unity.ml-agents/Tests/Editor/Unity.ML-Agents.Editor.Tests.asmdef b/com.unity.ml-agents/Tests/Editor/Unity.ML-Agents.Editor.Tests.asmdef
@@ -4,6 +4,7 @@
         "Unity.ML-Agents.Editor",
         "Unity.ML-Agents",
         "Unity.Barracuda",
+        "Unity.Mathematics",
         "Unity.ML-Agents.CommunicatorObjects",
         "Unity.ML-Agents.Runtime.Utils.Tests",
         "Unity.ML-Agents.Runtime.Sensor.Tests"

diff --git a/docs/Learning-Environment-Create-New.md b/docs/Learning-Environment-Create-New.md
@@ -475,6 +475,14 @@ RollerBall environment:
 1. You can now instantiate copies of the TrainingArea prefab. Drag them into
    your scene, positioning them so that they do not overlap.
 
+Alternatively, you can use the `TrainingAreaReplicator` to replicate training areas. Use the following steps:
+
+1. Create a new empty Game Object in the scene.
+2. Click on the new object and add a TrainingAreaReplicator component to the empty Game Object through the inspector.
+3. Drag the training area to `Base Area` in the Training Area Replicator.
+4. Specify the number of areas to replicate and the separation between areas.
+5. Hit play and the areas will be replicated automatically!
+
 ## Optional: Training Using Concurrent Unity Instances
 Another level of parallelization comes by training using
 [concurrent Unity instances](ML-Agents-Overview.md#additional-features).

diff --git a/docs/Python-API-Documentation.md b/docs/Python-API-Documentation.md
@@ -623,7 +623,7 @@ class UnityEnvironment(BaseEnv)
 #### \_\_init\_\_
 
 ```python
- | __init__(file_name: Optional[str] = None, worker_id: int = 0, base_port: Optional[int] = None, seed: int = 0, no_graphics: bool = False, timeout_wait: int = 60, additional_args: Optional[List[str]] = None, side_channels: Optional[List[SideChannel]] = None, log_folder: Optional[str] = None)
+ | __init__(file_name: Optional[str] = None, worker_id: int = 0, base_port: Optional[int] = None, seed: int = 0, no_graphics: bool = False, timeout_wait: int = 60, additional_args: Optional[List[str]] = None, side_channels: Optional[List[SideChannel]] = None, log_folder: Optional[str] = None, num_areas: int = 1)
 ```
 
 Starts a new unity environment and establishes a connection with the environment.