Support multi-dimensional and compressed observations stacking (#4476)

Added stacking to multi-dimensional and compressed observations and added compressed channel mapping in communicator to support decompression.

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

Author: 3 people

com.unity.ml-agents/CHANGELOG.md

@@ -13,11 +13,19 @@ and this project adheres to
  - Added the Random Network Distillation (RND) intrinsic reward signal to the Pytorch
  trainers. To use RND, add a `rnd` section to the `reward_signals` section of your
  yaml configuration file. [More information here](https://github.com/Unity-Technologies/ml-agents/blob/master/docs/Training-Configuration-File.md#rnd-intrinsic-reward)
-
 ### Minor Changes
 #### com.unity.ml-agents (C#)
+ - Stacking for compressed observations is now supported. An addtional setting
+ option `Observation Stacks` is added in editor to sensor components that support
+ compressed observations. A new class `ISparseChannelSensor` with an
+ additional method `GetCompressedChannelMapping()`is added to generate a mapping
+ of the channels in compressed data to the actual channel after decompression,
+ for the python side to decompress correctly. (#4476)
 #### ml-agents / ml-agents-envs / gym-unity (Python)
-
+ - The Communication API was changed to 1.2.0 to indicate support for stacked
+ compressed observation. A new entry `compressed_channel_mapping` is added to the
+ proto to handle decompression correctly. Newer versions of the package that wish to
+ make use of this will also need a compatible version of the Python trainers. (#4476)
 ### Bug Fixes
 #### com.unity.ml-agents (C#)
 #### ml-agents / ml-agents-envs / gym-unity (Python)

com.unity.ml-agents/Editor/CameraSensorComponentEditor.cs

@@ -25,6 +25,7 @@ public override void OnInspectorGUI()
                 EditorGUILayout.PropertyField(so.FindProperty("m_Width"), true);
                 EditorGUILayout.PropertyField(so.FindProperty("m_Height"), true);
                 EditorGUILayout.PropertyField(so.FindProperty("m_Grayscale"), true);
+                EditorGUILayout.PropertyField(so.FindProperty("m_ObservationStacks"), true);
             }
             EditorGUI.EndDisabledGroup();
             EditorGUILayout.PropertyField(so.FindProperty("m_Compression"), true);

com.unity.ml-agents/Editor/RenderTextureSensorComponentEditor.cs

@@ -20,6 +20,7 @@ public override void OnInspectorGUI()
                 EditorGUILayout.PropertyField(so.FindProperty("m_RenderTexture"), true);
                 EditorGUILayout.PropertyField(so.FindProperty("m_SensorName"), true);
                 EditorGUILayout.PropertyField(so.FindProperty("m_Grayscale"), true);
+                EditorGUILayout.PropertyField(so.FindProperty("m_ObservationStacks"), true);
             }
             EditorGUI.EndDisabledGroup();
 

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

@@ -74,9 +74,13 @@ public class Academy : IDisposable
         ///         <term>1.1.0</term>
         ///         <description>Support concatenated PNGs for compressed observations.</description>
         ///     </item>
+        ///     <item>
+        ///         <term>1.2.0</term>
+        ///         <description>Support compression mapping for stacked compressed observations.</description>
+        ///     </item>
         /// </list>
         /// </remarks>
-        const string k_ApiVersion = "1.1.0";
+        const string k_ApiVersion = "1.2.0";
 
         /// <summary>
         /// Unity package version of com.unity.ml-agents.

com.unity.ml-agents/Runtime/Communicator/GrpcExtensions.cs

@@ -222,7 +222,8 @@ public static List<float[]> ToAgentActionList(this UnityRLInputProto.Types.ListA
         /// <summary>
         /// Static flag to make sure that we only fire the warning once.
         /// </summary>
-        private static bool s_HaveWarnedAboutTrainerCapabilities = false;
+        private static bool s_HaveWarnedTrainerCapabilitiesMultiPng = false;
+        private static bool s_HaveWarnedTrainerCapabilitiesMapping = false;
 
         /// <summary>
         /// Generate an ObservationProto for the sensor using the provided ObservationWriter.
@@ -243,10 +244,27 @@ public static ObservationProto GetObservationProto(this ISensor sensor, Observat
                 var trainerCanHandle = Academy.Instance.TrainerCapabilities == null || Academy.Instance.TrainerCapabilities.ConcatenatedPngObservations;
                 if (!trainerCanHandle)
                 {
-                    if (!s_HaveWarnedAboutTrainerCapabilities)
+                    if (!s_HaveWarnedTrainerCapabilitiesMultiPng)
                     {
                         Debug.LogWarning($"Attached trainer doesn't support multiple PNGs. Switching to uncompressed observations for sensor {sensor.GetName()}.");
-                        s_HaveWarnedAboutTrainerCapabilities = true;
+                        s_HaveWarnedTrainerCapabilitiesMultiPng = true;
+                    }
+                    compressionType = SensorCompressionType.None;
+                }
+            }
+            // Check capabilities if we need mapping for compressed observations
+            if (compressionType != SensorCompressionType.None && shape.Length == 3 && shape[2] > 3)
+            {
+                var trainerCanHandleMapping = Academy.Instance.TrainerCapabilities == null || Academy.Instance.TrainerCapabilities.CompressedChannelMapping;
+                var isTrivialMapping = IsTrivialMapping(sensor);
+                if (!trainerCanHandleMapping && !isTrivialMapping)
+                {
+                    if (!s_HaveWarnedTrainerCapabilitiesMapping)
+                    {
+                        Debug.LogWarning($"The sensor {sensor.GetName()} is using non-trivial mapping and " +
+                                "the attached trainer doesn't support compression mapping. " +
+                                "Switching to uncompressed observations.");
+                        s_HaveWarnedTrainerCapabilitiesMapping = true;
                     }
                     compressionType = SensorCompressionType.None;
                 }
@@ -283,12 +301,16 @@ public static ObservationProto GetObservationProto(this ISensor sensor, Observat
                         "return SensorCompressionType.None from GetCompressionType()."
                         );
                 }
-
                 observationProto = new ObservationProto
                 {
                     CompressedData = ByteString.CopyFrom(compressedObs),
                     CompressionType = (CompressionTypeProto)sensor.GetCompressionType(),
                 };
+                var compressibleSensor = sensor as ISparseChannelSensor;
+                if (compressibleSensor != null)
+                {
+                    observationProto.CompressedChannelMapping.AddRange(compressibleSensor.GetCompressedChannelMapping());
+                }
             }
             observationProto.Shape.AddRange(shape);
             return observationProto;
@@ -300,7 +322,8 @@ public static UnityRLCapabilities ToRLCapabilities(this UnityRLCapabilitiesProto
             return new UnityRLCapabilities
             {
                 BaseRLCapabilities = proto.BaseRLCapabilities,
-                ConcatenatedPngObservations = proto.ConcatenatedPngObservations
+                ConcatenatedPngObservations = proto.ConcatenatedPngObservations,
+                CompressedChannelMapping = proto.CompressedChannelMapping,
             };
         }
 
@@ -310,7 +333,36 @@ public static UnityRLCapabilitiesProto ToProto(this UnityRLCapabilities rlCaps)
             {
                 BaseRLCapabilities = rlCaps.BaseRLCapabilities,
                 ConcatenatedPngObservations = rlCaps.ConcatenatedPngObservations,
+                CompressedChannelMapping = rlCaps.CompressedChannelMapping,
             };
         }
+
+        internal static bool IsTrivialMapping(ISensor sensor)
+        {
+            var compressibleSensor = sensor as ISparseChannelSensor;
+            if (compressibleSensor is null)
+            {
+                return true;
+            }
+            var mapping = compressibleSensor.GetCompressedChannelMapping();
+            if (mapping == null)
+            {
+                return true;
+            }
+            // check if mapping equals zero mapping
+            if (mapping.Length == 3 && mapping.All(m => m == 0))
+            {
+                return true;
+            }
+            // check if mapping equals identity mapping
+            for (var i = 0; i < mapping.Length; i++)
+            {
+                if (mapping[i] != i)
+                {
+                    return false;
+                }
+            }
+            return true;
+        }
     }
 }

com.unity.ml-agents/Runtime/Communicator/UnityRLCapabilities.cs

@@ -6,15 +6,17 @@ internal class UnityRLCapabilities
     {
         public bool BaseRLCapabilities;
         public bool ConcatenatedPngObservations;
+        public bool CompressedChannelMapping;
 
         /// <summary>
         /// A class holding the capabilities flags for Reinforcement Learning across C# and the Trainer codebase.  This
         /// struct will be used to inform users if and when they are using C# / Trainer features that are mismatched.
         /// </summary>
-        public UnityRLCapabilities(bool baseRlCapabilities = true, bool concatenatedPngObservations = true)
+        public UnityRLCapabilities(bool baseRlCapabilities = true, bool concatenatedPngObservations = true, bool compressedChannelMapping = true)
         {
             BaseRLCapabilities = baseRlCapabilities;
             ConcatenatedPngObservations = concatenatedPngObservations;
+            CompressedChannelMapping = compressedChannelMapping;
         }
 
         /// <summary>

com.unity.ml-agents/Runtime/Grpc/CommunicatorObjects/Capabilities.cs

@@ -25,14 +25,15 @@ static CapabilitiesReflection() {
       byte[] descriptorData = global::System.Convert.FromBase64String(
           string.Concat(
             "CjVtbGFnZW50c19lbnZzL2NvbW11bmljYXRvcl9vYmplY3RzL2NhcGFiaWxp",
-            "dGllcy5wcm90bxIUY29tbXVuaWNhdG9yX29iamVjdHMiWwoYVW5pdHlSTENh",
+            "dGllcy5wcm90bxIUY29tbXVuaWNhdG9yX29iamVjdHMifQoYVW5pdHlSTENh",
             "cGFiaWxpdGllc1Byb3RvEhoKEmJhc2VSTENhcGFiaWxpdGllcxgBIAEoCBIj",
-            "Chtjb25jYXRlbmF0ZWRQbmdPYnNlcnZhdGlvbnMYAiABKAhCJaoCIlVuaXR5",
-            "Lk1MQWdlbnRzLkNvbW11bmljYXRvck9iamVjdHNiBnByb3RvMw=="));
+            "Chtjb25jYXRlbmF0ZWRQbmdPYnNlcnZhdGlvbnMYAiABKAgSIAoYY29tcHJl",
+            "c3NlZENoYW5uZWxNYXBwaW5nGAMgASgIQiWqAiJVbml0eS5NTEFnZW50cy5D",
+            "b21tdW5pY2F0b3JPYmplY3RzYgZwcm90bzM="));
       descriptor = pbr::FileDescriptor.FromGeneratedCode(descriptorData,
           new pbr::FileDescriptor[] { },
           new pbr::GeneratedClrTypeInfo(null, new pbr::GeneratedClrTypeInfo[] {
-            new pbr::GeneratedClrTypeInfo(typeof(global::Unity.MLAgents.CommunicatorObjects.UnityRLCapabilitiesProto), global::Unity.MLAgents.CommunicatorObjects.UnityRLCapabilitiesProto.Parser, new[]{ "BaseRLCapabilities", "ConcatenatedPngObservations" }, null, null, null)
+            new pbr::GeneratedClrTypeInfo(typeof(global::Unity.MLAgents.CommunicatorObjects.UnityRLCapabilitiesProto), global::Unity.MLAgents.CommunicatorObjects.UnityRLCapabilitiesProto.Parser, new[]{ "BaseRLCapabilities", "ConcatenatedPngObservations", "CompressedChannelMapping" }, null, null, null)
           }));
     }
     #endregion
@@ -71,6 +72,7 @@ public UnityRLCapabilitiesProto() {
     public UnityRLCapabilitiesProto(UnityRLCapabilitiesProto other) : this() {
       baseRLCapabilities_ = other.baseRLCapabilities_;
       concatenatedPngObservations_ = other.concatenatedPngObservations_;
+      compressedChannelMapping_ = other.compressedChannelMapping_;
       _unknownFields = pb::UnknownFieldSet.Clone(other._unknownFields);
     }
 
@@ -107,6 +109,20 @@ public bool ConcatenatedPngObservations {
       }
     }
 
+    /// <summary>Field number for the "compressedChannelMapping" field.</summary>
+    public const int CompressedChannelMappingFieldNumber = 3;
+    private bool compressedChannelMapping_;
+    /// <summary>
+    /// compression mapping for stacking compressed observations.
+    /// </summary>
+    [global::System.Diagnostics.DebuggerNonUserCodeAttribute]
+    public bool CompressedChannelMapping {
+      get { return compressedChannelMapping_; }
+      set {
+        compressedChannelMapping_ = value;
+      }
+    }
+
     [global::System.Diagnostics.DebuggerNonUserCodeAttribute]
     public override bool Equals(object other) {
       return Equals(other as UnityRLCapabilitiesProto);
@@ -122,6 +138,7 @@ public bool Equals(UnityRLCapabilitiesProto other) {
       }
       if (BaseRLCapabilities != other.BaseRLCapabilities) return false;
       if (ConcatenatedPngObservations != other.ConcatenatedPngObservations) return false;
+      if (CompressedChannelMapping != other.CompressedChannelMapping) return false;
       return Equals(_unknownFields, other._unknownFields);
     }
 
@@ -130,6 +147,7 @@ public override int GetHashCode() {
       int hash = 1;
       if (BaseRLCapabilities != false) hash ^= BaseRLCapabilities.GetHashCode();
       if (ConcatenatedPngObservations != false) hash ^= ConcatenatedPngObservations.GetHashCode();
+      if (CompressedChannelMapping != false) hash ^= CompressedChannelMapping.GetHashCode();
       if (_unknownFields != null) {
         hash ^= _unknownFields.GetHashCode();
       }
@@ -151,6 +169,10 @@ public void WriteTo(pb::CodedOutputStream output) {
         output.WriteRawTag(16);
         output.WriteBool(ConcatenatedPngObservations);
       }
+      if (CompressedChannelMapping != false) {
+        output.WriteRawTag(24);
+        output.WriteBool(CompressedChannelMapping);
+      }
       if (_unknownFields != null) {
         _unknownFields.WriteTo(output);
       }
@@ -165,6 +187,9 @@ public int CalculateSize() {
       if (ConcatenatedPngObservations != false) {
         size += 1 + 1;
       }
+      if (CompressedChannelMapping != false) {
+        size += 1 + 1;
+      }
       if (_unknownFields != null) {
         size += _unknownFields.CalculateSize();
       }
@@ -182,6 +207,9 @@ public void MergeFrom(UnityRLCapabilitiesProto other) {
       if (other.ConcatenatedPngObservations != false) {
         ConcatenatedPngObservations = other.ConcatenatedPngObservations;
       }
+      if (other.CompressedChannelMapping != false) {
+        CompressedChannelMapping = other.CompressedChannelMapping;
+      }
       _unknownFields = pb::UnknownFieldSet.MergeFrom(_unknownFields, other._unknownFields);
     }
 
@@ -201,6 +229,10 @@ public void MergeFrom(pb::CodedInputStream input) {
             ConcatenatedPngObservations = input.ReadBool();
             break;
           }
+          case 24: {
+            CompressedChannelMapping = input.ReadBool();
+            break;
+          }
         }
       }
     }