Changing the API, not working as fast for now

Author: vincentpierre

Assets/ECS_MLAgents_v0/Core/AgentSystem.cs

@@ -5,6 +5,7 @@
 using Unity.Entities;
 using Unity.Jobs;
 using UnityEngine;
+using Unity.Burst;
 
 namespace ECS_MLAgents_v0.Core
 {
@@ -35,29 +36,27 @@ namespace ECS_MLAgents_v0.Core
      * data in batch and generate a new NativeArray<float> that will be used to populate the
      * Actuator data of all compatible Entities.
      */
-    public abstract class AgentSystem<TS, TA> : JobComponentSystem, IAgentSystem
+    public abstract class AgentSystem<TS, TA> : JobComponentSystem, IAgentSystem<TS, TA>
         where TS : struct, IComponentData
         where TA : struct, IComponentData
     {   
         private const int INITIAL_MEMORY_SIZE = 1024;
-        private const int SIZE_OF_FLOAT_IN_MEMORY = 4;
 
-        private int _sensorMemorySize = INITIAL_MEMORY_SIZE;
-        private int _actuatorMemorySize = INITIAL_MEMORY_SIZE;
+        private int _currentNAgents = INITIAL_MEMORY_SIZE;
 
         public IDecisionRequester DecisionRequester { get; set; }
         private int _phase;
 
-        public IAgentDecision Decision { get; set; }
+        public IAgentDecision<TS, TA> Decision { get; set; }
 
         private ComponentGroup _componentGroup;
         private int _sensorSize;
         private int _actuatorSize;
         // TODO : Make sure there is not extra cost for memory allocation here and when copying
-        private NativeArray<float> _sensorTensor =
-            new NativeArray<float>(INITIAL_MEMORY_SIZE, Allocator.Persistent);
-        private NativeArray<float> _actuatorTensor =
-            new NativeArray<float>(INITIAL_MEMORY_SIZE, Allocator.Persistent);
+        private NativeArray<TS> _sensorTensor =
+            new NativeArray<TS>(INITIAL_MEMORY_SIZE, Allocator.Persistent);
+        private NativeArray<TA> _actuatorTensor =
+            new NativeArray<TA>(INITIAL_MEMORY_SIZE, Allocator.Persistent);
 
         // TODO : Decide if we want to keep at all
         private Logger _logger;
@@ -138,17 +137,13 @@ protected override JobHandle OnUpdate(JobHandle inputDeps)
              * If there was more agents than allowed by the memory allocation of the sensor or
              * actuator, then the size is updated to the required size.
              */
-            if (nAgents * _sensorSize / SIZE_OF_FLOAT_IN_MEMORY > _sensorMemorySize)
+            if (nAgents > _currentNAgents)
             {
-                _sensorMemorySize = nAgents * _sensorSize / SIZE_OF_FLOAT_IN_MEMORY;
+                _currentNAgents = nAgents;
                 _sensorTensor.Dispose();
-                _sensorTensor = new NativeArray<float>(_sensorMemorySize, Allocator.Persistent);
-            }
-            if (nAgents * _actuatorSize / SIZE_OF_FLOAT_IN_MEMORY > _actuatorMemorySize)
-            {
-                _actuatorMemorySize = nAgents * _actuatorSize / SIZE_OF_FLOAT_IN_MEMORY;
                 _actuatorTensor.Dispose();
-                _actuatorTensor = new NativeArray<float>(_actuatorMemorySize, Allocator.Persistent);
+                _sensorTensor = new NativeArray<TS>(_currentNAgents, Allocator.Persistent);
+                _actuatorTensor = new NativeArray<TA>(_currentNAgents, Allocator.Persistent);
             }
 
             /*
@@ -167,48 +162,39 @@ protected override JobHandle OnUpdate(JobHandle inputDeps)
             {
                 Sensors = sensors,
                 SensorTensor = _sensorTensor,
-                SensorSize = _sensorSize
             };
             handle = copySensorsJob.Schedule(nAgents, 64, handle);
-            
+
             handle.Complete();
-            
-            /*
-             * The Decision is called here to populate the NativeArray<float> of Actuators.
-             */
-            handle = Decision.DecideBatch(ref _sensorTensor, 
-                ref _actuatorTensor, 
-                _sensorSize / SIZE_OF_FLOAT_IN_MEMORY, 
-                _actuatorSize / SIZE_OF_FLOAT_IN_MEMORY, 
-                nAgents,
-                handle);
-            
-            /*
+
+            Decision.BatchProcess(ref _sensorTensor, ref _actuatorTensor);
+
+                        /*
              * Copy the data from the actuator NativeArray<float> to the actuators of each entity.
              */
             var copyActuatorsJob = new CopyActuatorsJob
             {
                 ActuatorTensor = _actuatorTensor,
                 Actuators = actuators,
-                ActuatorSize = _actuatorSize
+
             }; 
 
             return copyActuatorsJob.Schedule(nAgents, 64, handle);
+
         }
 
         /*
          * This IJobParallelFor copied the Sensor data into a NativeArray<float>
          */
-//        [BurstCompile]
+       [BurstCompile]
         private struct CopySensorsJob : IJobParallelFor
         {
             [ReadOnly] public ComponentDataArray<TS> Sensors;
-            public NativeArray<float> SensorTensor;
-            [ReadOnly] public int SensorSize;
+            public NativeArray<TS> SensorTensor;
 
             public void Execute(int i)
             {
-                TensorUtility.CopyToNativeArray(Sensors[i], SensorTensor, i * SensorSize);
+                SensorTensor[i] = Sensors[i];
             }
         }
 
@@ -218,17 +204,12 @@ public void Execute(int i)
 //        [BurstCompile]
         private struct CopyActuatorsJob : IJobParallelFor
         {
-
             public ComponentDataArray<TA> Actuators;
-            public NativeArray<float> ActuatorTensor;
-            [ReadOnly] public int ActuatorSize;
+            public NativeArray<TA> ActuatorTensor;
 
             public void Execute(int i)
             {
-                var tmp = Actuators[i];
-                // TODO : Make sure there is no extra cost here
-                TensorUtility.CopyFromNativeArray(ActuatorTensor, out tmp, i * ActuatorSize);
-                Actuators[i] = tmp;
+                Actuators[i] = ActuatorTensor[i];
             }
         }
     }

Assets/ECS_MLAgents_v0/Core/ExternalDecision.cs

@@ -6,10 +6,13 @@
 using Unity.Collections.LowLevel.Unsafe;
 using UnityEngine;
 using UnityEngine.Profiling;
+using Unity.Entities;
 
 
 namespace ECS_MLAgents_v0.Core{
-    public class ExternalDecision : IAgentDecision
+    public class ExternalDecision<TS, TA> : IAgentDecision<TS, TA> 
+        where TS : struct, IComponentData
+        where TA : struct, IComponentData 
     {
 
         // [ Unity Ready (1) , nAgents (4) , sensorSize (4) , actuatorSize (4) , Data
@@ -27,6 +30,14 @@ public class ExternalDecision : IAgentDecision
 
 
         private float[] actuatorData = new float[0];
+
+
+
+        System.Type _sensorType;
+        System.Type _actuatorType;
+
+        private int _sensorSize;
+        private int _actuatorSize;
 
         // This is a temporary test file
         // TODO : Replace with a file creation system
@@ -47,36 +58,35 @@ public ExternalDecision()
             Debug.Log("Is Ready to Communicate");
         }
 
-        public JobHandle DecideBatch(
-            ref NativeArray<float> sensor, 
-            ref NativeArray<float> actuator, 
-            int sensorSize,
-            int actuatorSize, 
-            int nAgents, 
-            JobHandle handle)
+         public void BatchProcess(ref NativeArray<TS> sensors, ref NativeArray<TA> actuators )
         {
             Profiler.BeginSample("Communicating");
-            if (sensor.Length > 4 * 50000)
+
+            VerifySensor(typeof(TS));
+            VerifyActuator(typeof(TA));
+
+            int batch = sensors.Length;
+            if (batch != actuators.Length)
             {
-                throw new Exception("TOO much data to send");
+                throw new Exception("Error in the length of the sensors and actuators");
             }
-            
-            if (actuator.Length > 4 * 50000)
+
+            if (batch > 50000)
             {
                 throw new Exception("TOO much data to send");
             }
 
-            if (actuatorData.Length < actuator.Length)
+            if (actuatorData.Length < _actuatorSize* batch)
             {
-                actuatorData = new float[actuator.Length];
+                actuatorData = new float[_actuatorSize * batch];
             }
 
 
-            accessor.Write(NUMBER_AGENTS_POSITION, nAgents);
-            accessor.Write(SENSOR_SIZE_POSITION, sensorSize);
-            accessor.Write(ACTUATOR_SIZE_POSITION, actuatorSize);
+            accessor.Write(NUMBER_AGENTS_POSITION, batch);
+            accessor.Write(SENSOR_SIZE_POSITION, _sensorSize);
+            accessor.Write(ACTUATOR_SIZE_POSITION, _actuatorSize);
 
-            accessor.WriteArray(SENSOR_DATA_POSITION, sensor.ToArray(), 0, sensor.Length);
+            accessor.WriteArray(SENSOR_DATA_POSITION, sensors.ToArray(), 0, batch);
 
             accessor.Write(PYTHON_READY_POSITION, false);
 
@@ -96,11 +106,36 @@ public JobHandle DecideBatch(
                 }
             }
 
-            accessor.ReadArray(ACTUATOR_DATA_POSITION, actuatorData, 0, actuator.Length);
-            actuator.CopyFrom(actuatorData);
+            accessor.ReadArray(ACTUATOR_DATA_POSITION, actuatorData, 0, batch * _actuatorSize);
+
+            // actuator.CopyFrom(actuatorData);
+
+            var tmpA = new NativeArray<float>(batch * _actuatorSize, Allocator.Persistent);
+            tmpA.CopyFrom(actuatorData);
+            for(var i = 0; i< batch; i++){
+                var act = new TA();
+                TensorUtility.CopyFromNativeArray(tmpA, out act, i * _sensorSize * 4);
+                actuators[i] = act;
+            }
+            tmpA.Dispose();
+
 
             Profiler.BeginSample("Communicating");
-            return handle;
+        }
+
+        private void VerifySensor(System.Type t){
+            if (! t.Equals(_sensorType)){
+                TensorUtility.DebugCheckStructure(t);
+                _sensorSize = System.Runtime.InteropServices.Marshal.SizeOf(t) / 4;
+                _sensorType = t;
+            }
+        }
+        private void VerifyActuator(System.Type t){
+            if (! t.Equals(_actuatorType)){
+                TensorUtility.DebugCheckStructure(t);
+                _actuatorSize = System.Runtime.InteropServices.Marshal.SizeOf(t) / 4;
+                _actuatorType = t;
+            }
         }
     }
 }

Assets/ECS_MLAgents_v0/Core/IAgentDecision.cs

@@ -1,34 +1,26 @@
 using Unity.Collections;
-using Unity.Jobs;
+using Unity.Entities;
 
 namespace ECS_MLAgents_v0.Core
 {
     /*
      * The Interface to define a Decision process by which a bach of agent updates its actuator
      * based on the information present in the sensor.
      */
-    public interface IAgentDecision
+    public interface IAgentDecision<TS, TA> 
+        where TS : struct, IComponentData
+        where TA : struct, IComponentData 
     {
         /// <summary>
-        /// DecideBatch updates the aggregated actuators of the agents present in the batch from
-        /// the aggregated actuators. 
+        /// DecideBatch updates the actuators of the agents present in the batch from
+        /// the data present in the sensors. 
         /// </summary>
-        /// <param name="sensor">The aggregated data for the sensor information present in the
-        /// batch. The sensor data is linearly arranged.</param>
-        /// <param name="actuator">The aggregated data for the actuator information present in the
-        /// batch. The sensor data is linearly arranged.</param>
-        /// <param name="sensorSize">The number of float values present in a sensor for one agent
-        /// </param>
-        /// <param name="actuatorSize">The number of float values present in an actuator
-        /// for one agent</param>
-        /// <param name="nAgents">The number of agents present in the batch</param>
-        /// <param name="handle">The JobHandle for the input dependencies.</param>
-        /// <returns>The Job Handle for the output dependencies.</returns>
-        JobHandle DecideBatch(ref NativeArray<float> sensor, 
-            ref NativeArray<float> actuator, 
-            int sensorSize, 
-            int actuatorSize,
-            int nAgents,
-            JobHandle handle);
+        /// <param name="sensors">The aggregated data for the sensor information present in the
+        /// batch. T.</param>
+        /// <param name="actuators">The aggregated data for the actuator information present in the
+        /// batch. </param>
+        void BatchProcess(ref NativeArray<TS> sensors, ref NativeArray<TA> actuators);
+
     }
+
 }

Assets/ECS_MLAgents_v0/Core/IAgentSystem.cs

@@ -2,7 +2,9 @@
 
 namespace ECS_MLAgents_v0.Core
 {
-    public interface IAgentSystem
+    public interface IAgentSystem<TS, TA>
+        where TS : struct , IComponentData
+        where TA:struct, IComponentData
     {
         /// <summary>
         /// If true, the AgentSystem will perform on the agents
@@ -12,7 +14,7 @@ public interface IAgentSystem
         /// <summary>
         /// The IAgentDecision that will be used to update the Actuators of compatible Entities.
         /// </summary>
-        IAgentDecision Decision { get; set; }
+        IAgentDecision<TS, TA> Decision { get; set; } 
 
         /// <summary>
         /// This method defines what are the required ComponentType that are needed on an Entity