inference - fill 0s for done Agents (#3232)

* fill 0s for done agents

* docstrings

Author: Chris Elion

UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/GeneratorImpl.cs

@@ -101,20 +101,30 @@ public void Generate(TensorProxy tensorProxy, int batchSize, IEnumerable<AgentIn
             var agentIndex = 0;
             foreach (var info in infos)
             {
-                var tensorOffset = 0;
-                // Write each sensor consecutively to the tensor
-                foreach (var sensorIndex in m_SensorIndices)
+                if (info.agentInfo.done)
                 {
-                    var sensor = info.sensors[sensorIndex];
-                    m_WriteAdapter.SetTarget(tensorProxy, agentIndex, tensorOffset);
-                    var numWritten = sensor.Write(m_WriteAdapter);
-                    tensorOffset += numWritten;
+                    // If the agent is done, we might have a stale reference to the sensors
+                    // e.g. a dependent object might have been disposed.
+                    // To avoid this, just fill observation with zeroes instead of calling sensor.Write.
+                    TensorUtils.FillTensorBatch(tensorProxy, agentIndex, 0.0f);
+                }
+                else
+                {
+                    var tensorOffset = 0;
+                    // Write each sensor consecutively to the tensor
+                    foreach (var sensorIndex in m_SensorIndices)
+                    {
+                        var sensor = info.sensors[sensorIndex];
+                        m_WriteAdapter.SetTarget(tensorProxy, agentIndex, tensorOffset);
+                        var numWritten = sensor.Write(m_WriteAdapter);
+                        tensorOffset += numWritten;
+                    }
+                    Debug.AssertFormat(
+                        tensorOffset == vecObsSizeT,
+                        "mismatch between vector observation size ({0}) and number of observations written ({1})",
+                        vecObsSizeT, tensorOffset
+                    );
                 }
-                Debug.AssertFormat(
-                    tensorOffset == vecObsSizeT,
-                    "mismatch between vector observation size ({0}) and number of observations written ({1})",
-                    vecObsSizeT, tensorOffset
-                );
 
                 agentIndex++;
             }
@@ -356,8 +366,19 @@ public void Generate(TensorProxy tensorProxy, int batchSize, IEnumerable<AgentIn
             foreach (var infoSensorPair in infos)
             {
                 var sensor = infoSensorPair.sensors[m_SensorIndex];
-                m_WriteAdapter.SetTarget(tensorProxy, agentIndex, 0);
-                sensor.Write(m_WriteAdapter);
+                if (infoSensorPair.agentInfo.done)
+                {
+                    // If the agent is done, we might have a stale reference to the sensors
+                    // e.g. a dependent object might have been disposed.
+                    // To avoid this, just fill observation with zeroes instead of calling sensor.Write.
+                    TensorUtils.FillTensorBatch(tensorProxy, agentIndex, 0.0f);
+                }
+                else
+                {
+                    m_WriteAdapter.SetTarget(tensorProxy, agentIndex, 0);
+                    sensor.Write(m_WriteAdapter);
+
+                }
                 agentIndex++;
             }
         }

UnitySDK/Assets/ML-Agents/Scripts/InferenceBrain/TensorProxy.cs

@@ -90,6 +90,29 @@ public static TensorProxy TensorProxyFromBarracuda(Tensor src, string nameOverri
             };
         }
 
+        /// <summary>
+        /// Fill a specific batch of a TensorProxy with a given value
+        /// </summary>
+        /// <param name="tensorProxy"></param>
+        /// <param name="batch">The batch index to fill.</param>
+        /// <param name="fillValue"></param>
+        public static void FillTensorBatch(TensorProxy tensorProxy, int batch, float fillValue)
+        {
+            var height = tensorProxy.data.height;
+            var width = tensorProxy.data.width;
+            var channels = tensorProxy.data.channels;
+            for (var h = 0; h < height; h++)
+            {
+                for (var w = 0; w < width; w++)
+                {
+                    for (var c = 0; c < channels; c++)
+                    {
+                        tensorProxy.data[batch, h, w, c] = fillValue;
+                    }
+                }
+            }
+        }
+
         /// <summary>
         /// Fill a pre-allocated Tensor with random numbers
         /// </summary>