Merge pull request #4825 from Unity-Technologies/sensor-types

[WIP] Observation Types

Author: awjuliani

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

@@ -411,6 +411,17 @@ public static ObservationProto GetObservationProto(this ISensor sensor, Observat
                 }
             }
             observationProto.Shape.AddRange(shape);
+
+            // Add the observation type, if any, to the observationProto
+            var typeSensor = sensor as ITypedSensor;
+            if (typeSensor != null)
+            {
+                observationProto.ObservationType = (ObservationTypeProto)typeSensor.GetObservationType();
+            }
+            else
+            {
+                observationProto.ObservationType = ObservationTypeProto.Default;
+            }
             return observationProto;
         }
 

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

@@ -25,20 +25,23 @@ static ObservationReflection() {
       byte[] descriptorData = global::System.Convert.FromBase64String(
           string.Concat(
             "CjRtbGFnZW50c19lbnZzL2NvbW11bmljYXRvcl9vYmplY3RzL29ic2VydmF0",
-            "aW9uLnByb3RvEhRjb21tdW5pY2F0b3Jfb2JqZWN0cyK7AgoQT2JzZXJ2YXRp",
+            "aW9uLnByb3RvEhRjb21tdW5pY2F0b3Jfb2JqZWN0cyKBAwoQT2JzZXJ2YXRp",
             "b25Qcm90bxINCgVzaGFwZRgBIAMoBRJEChBjb21wcmVzc2lvbl90eXBlGAIg",
             "ASgOMiouY29tbXVuaWNhdG9yX29iamVjdHMuQ29tcHJlc3Npb25UeXBlUHJv",
             "dG8SGQoPY29tcHJlc3NlZF9kYXRhGAMgASgMSAASRgoKZmxvYXRfZGF0YRgE",
             "IAEoCzIwLmNvbW11bmljYXRvcl9vYmplY3RzLk9ic2VydmF0aW9uUHJvdG8u",
             "RmxvYXREYXRhSAASIgoaY29tcHJlc3NlZF9jaGFubmVsX21hcHBpbmcYBSAD",
-            "KAUSHAoUZGltZW5zaW9uX3Byb3BlcnRpZXMYBiADKAUaGQoJRmxvYXREYXRh",
-            "EgwKBGRhdGEYASADKAJCEgoQb2JzZXJ2YXRpb25fZGF0YSopChRDb21wcmVz",
-            "c2lvblR5cGVQcm90bxIICgROT05FEAASBwoDUE5HEAFCJaoCIlVuaXR5Lk1M",
-            "QWdlbnRzLkNvbW11bmljYXRvck9iamVjdHNiBnByb3RvMw=="));
+            "KAUSHAoUZGltZW5zaW9uX3Byb3BlcnRpZXMYBiADKAUSRAoQb2JzZXJ2YXRp",
+            "b25fdHlwZRgHIAEoDjIqLmNvbW11bmljYXRvcl9vYmplY3RzLk9ic2VydmF0",
+            "aW9uVHlwZVByb3RvGhkKCUZsb2F0RGF0YRIMCgRkYXRhGAEgAygCQhIKEG9i",
+            "c2VydmF0aW9uX2RhdGEqKQoUQ29tcHJlc3Npb25UeXBlUHJvdG8SCAoETk9O",
+            "RRAAEgcKA1BORxABKkYKFE9ic2VydmF0aW9uVHlwZVByb3RvEgsKB0RFRkFV",
+            "TFQQABIICgRHT0FMEAESCgoGUkVXQVJEEAISCwoHTUVTU0FHRRADQiWqAiJV",
+            "bml0eS5NTEFnZW50cy5Db21tdW5pY2F0b3JPYmplY3RzYgZwcm90bzM="));
       descriptor = pbr::FileDescriptor.FromGeneratedCode(descriptorData,
           new pbr::FileDescriptor[] { },
-          new pbr::GeneratedClrTypeInfo(new[] {typeof(global::Unity.MLAgents.CommunicatorObjects.CompressionTypeProto), }, new pbr::GeneratedClrTypeInfo[] {
-            new pbr::GeneratedClrTypeInfo(typeof(global::Unity.MLAgents.CommunicatorObjects.ObservationProto), global::Unity.MLAgents.CommunicatorObjects.ObservationProto.Parser, new[]{ "Shape", "CompressionType", "CompressedData", "FloatData", "CompressedChannelMapping", "DimensionProperties" }, new[]{ "ObservationData" }, null, new pbr::GeneratedClrTypeInfo[] { new pbr::GeneratedClrTypeInfo(typeof(global::Unity.MLAgents.CommunicatorObjects.ObservationProto.Types.FloatData), global::Unity.MLAgents.CommunicatorObjects.ObservationProto.Types.FloatData.Parser, new[]{ "Data" }, null, null, null)})
+          new pbr::GeneratedClrTypeInfo(new[] {typeof(global::Unity.MLAgents.CommunicatorObjects.CompressionTypeProto), typeof(global::Unity.MLAgents.CommunicatorObjects.ObservationTypeProto), }, new pbr::GeneratedClrTypeInfo[] {
+            new pbr::GeneratedClrTypeInfo(typeof(global::Unity.MLAgents.CommunicatorObjects.ObservationProto), global::Unity.MLAgents.CommunicatorObjects.ObservationProto.Parser, new[]{ "Shape", "CompressionType", "CompressedData", "FloatData", "CompressedChannelMapping", "DimensionProperties", "ObservationType" }, new[]{ "ObservationData" }, null, new pbr::GeneratedClrTypeInfo[] { new pbr::GeneratedClrTypeInfo(typeof(global::Unity.MLAgents.CommunicatorObjects.ObservationProto.Types.FloatData), global::Unity.MLAgents.CommunicatorObjects.ObservationProto.Types.FloatData.Parser, new[]{ "Data" }, null, null, null)})
           }));
     }
     #endregion
@@ -50,6 +53,13 @@ internal enum CompressionTypeProto {
     [pbr::OriginalName("PNG")] Png = 1,
   }
 
+  internal enum ObservationTypeProto {
+    [pbr::OriginalName("DEFAULT")] Default = 0,
+    [pbr::OriginalName("GOAL")] Goal = 1,
+    [pbr::OriginalName("REWARD")] Reward = 2,
+    [pbr::OriginalName("MESSAGE")] Message = 3,
+  }
+
   #endregion
 
   #region Messages
@@ -82,6 +92,7 @@ public ObservationProto(ObservationProto other) : this() {
       compressionType_ = other.compressionType_;
       compressedChannelMapping_ = other.compressedChannelMapping_.Clone();
       dimensionProperties_ = other.dimensionProperties_.Clone();
+      observationType_ = other.observationType_;
       switch (other.ObservationDataCase) {
         case ObservationDataOneofCase.CompressedData:
           CompressedData = other.CompressedData;
@@ -162,6 +173,17 @@ public ObservationProto Clone() {
       get { return dimensionProperties_; }
     }
 
+    /// <summary>Field number for the "observation_type" field.</summary>
+    public const int ObservationTypeFieldNumber = 7;
+    private global::Unity.MLAgents.CommunicatorObjects.ObservationTypeProto observationType_ = 0;
+    [global::System.Diagnostics.DebuggerNonUserCodeAttribute]
+    public global::Unity.MLAgents.CommunicatorObjects.ObservationTypeProto ObservationType {
+      get { return observationType_; }
+      set {
+        observationType_ = value;
+      }
+    }
+
     private object observationData_;
     /// <summary>Enum of possible cases for the "observation_data" oneof.</summary>
     public enum ObservationDataOneofCase {
@@ -200,6 +222,7 @@ public bool Equals(ObservationProto other) {
       if (!object.Equals(FloatData, other.FloatData)) return false;
       if(!compressedChannelMapping_.Equals(other.compressedChannelMapping_)) return false;
       if(!dimensionProperties_.Equals(other.dimensionProperties_)) return false;
+      if (ObservationType != other.ObservationType) return false;
       if (ObservationDataCase != other.ObservationDataCase) return false;
       return Equals(_unknownFields, other._unknownFields);
     }
@@ -213,6 +236,7 @@ public override int GetHashCode() {
       if (observationDataCase_ == ObservationDataOneofCase.FloatData) hash ^= FloatData.GetHashCode();
       hash ^= compressedChannelMapping_.GetHashCode();
       hash ^= dimensionProperties_.GetHashCode();
+      if (ObservationType != 0) hash ^= ObservationType.GetHashCode();
       hash ^= (int) observationDataCase_;
       if (_unknownFields != null) {
         hash ^= _unknownFields.GetHashCode();
@@ -242,6 +266,10 @@ public void WriteTo(pb::CodedOutputStream output) {
       }
       compressedChannelMapping_.WriteTo(output, _repeated_compressedChannelMapping_codec);
       dimensionProperties_.WriteTo(output, _repeated_dimensionProperties_codec);
+      if (ObservationType != 0) {
+        output.WriteRawTag(56);
+        output.WriteEnum((int) ObservationType);
+      }
       if (_unknownFields != null) {
         _unknownFields.WriteTo(output);
       }
@@ -262,6 +290,9 @@ public int CalculateSize() {
       }
       size += compressedChannelMapping_.CalculateSize(_repeated_compressedChannelMapping_codec);
       size += dimensionProperties_.CalculateSize(_repeated_dimensionProperties_codec);
+      if (ObservationType != 0) {
+        size += 1 + pb::CodedOutputStream.ComputeEnumSize((int) ObservationType);
+      }
       if (_unknownFields != null) {
         size += _unknownFields.CalculateSize();
       }
@@ -279,6 +310,9 @@ public void MergeFrom(ObservationProto other) {
       }
       compressedChannelMapping_.Add(other.compressedChannelMapping_);
       dimensionProperties_.Add(other.dimensionProperties_);
+      if (other.ObservationType != 0) {
+        ObservationType = other.ObservationType;
+      }
       switch (other.ObservationDataCase) {
         case ObservationDataOneofCase.CompressedData:
           CompressedData = other.CompressedData;
@@ -334,6 +368,10 @@ public void MergeFrom(pb::CodedInputStream input) {
             dimensionProperties_.AddEntriesFrom(input, _repeated_dimensionProperties_codec);
             break;
           }
+          case 56: {
+            observationType_ = (global::Unity.MLAgents.CommunicatorObjects.ObservationTypeProto) input.ReadEnum();
+            break;
+          }
         }
       }
     }

com.unity.ml-agents/Runtime/Sensors/ITypedSensor.cs

@@ -0,0 +1,31 @@
+namespace Unity.MLAgents.Sensors
+{
+
+    /// <summary>
+    /// The ObservationType enum of the Sensor.
+    /// </summary>
+    internal enum ObservationType
+    {
+        // Collected observations are generic.
+        Default = 0,
+        // Collected observations contain goal information.
+        Goal = 1,
+        // Collected observations contain reward information.
+        Reward = 2,
+        // Collected observations are messages from other agents.
+        Message = 3,
+    }
+
+
+    /// <summary>
+    /// Sensor interface for sensors with variable types.
+    /// </summary>
+    internal interface ITypedSensor
+    {
+        /// <summary>
+        /// Returns the ObservationType enum corresponding to the type of the sensor.
+        /// </summary>
+        /// <returns>The ObservationType enum</returns>
+        ObservationType GetObservationType();
+    }
+}

com.unity.ml-agents/Runtime/Sensors/ITypedSensor.cs.meta

@@ -227,13 +227,14 @@ A `TerminalStep` has the following fields:
 
 A `BehaviorSpec` has the following fields :
 
-- `sensor_specs` is a List of `SensorSpec` objects : Each `SensorSpec`
+- `observation_specs` is a List of `ObservationSpec` objects : Each `ObservationSpec`
   corresponds to an observation's properties: `shape` is a tuple of ints that
   corresponds to the shape of the observation (without the number of agents dimension).
   `dimension_property` is a tuple of flags containing extra information about how the
-  data should be processed in the corresponding dimension. Note that the `SensorSpec`
-  have the same ordering as the ordering of observations in the DecisionSteps,
-  DecisionStep, TerminalSteps and TerminalStep.
+  data should be processed in the corresponding dimension. `observation_type` is an enum
+  corresponding to what type of observation is generating the data (i.e., default, goal,
+  etc). Note that the `ObservationSpec` have the same ordering as the ordering of observations
+  in the DecisionSteps, DecisionStep, TerminalSteps and TerminalStep.
 - `action_spec` is an `ActionSpec` namedtuple that defines the number and types
   of actions for the Agent.
 

docs/Python-API.md

@@ -229,16 +229,16 @@ def _preprocess_single(self, single_visual_obs: np.ndarray) -> np.ndarray:
 
     def _get_n_vis_obs(self) -> int:
         result = 0
-        for sen_spec in self.group_spec.sensor_specs:
-            if len(sen_spec.shape) == 3:
+        for obs_spec in self.group_spec.observation_specs:
+            if len(obs_spec.shape) == 3:
                 result += 1
         return result
 
     def _get_vis_obs_shape(self) -> List[Tuple]:
         result: List[Tuple] = []
-        for sen_spec in self.group_spec.sensor_specs:
-            if len(sen_spec.shape) == 3:
-                result.append(sen_spec.shape)
+        for obs_spec in self.group_spec.observation_specs:
+            if len(obs_spec.shape) == 3:
+                result.append(obs_spec.shape)
         return result
 
     def _get_vis_obs_list(
@@ -261,9 +261,9 @@ def _get_vector_obs(
 
     def _get_vec_obs_size(self) -> int:
         result = 0
-        for sen_spec in self.group_spec.sensor_specs:
-            if len(sen_spec.shape) == 1:
-                result += sen_spec.shape[0]
+        for obs_spec in self.group_spec.observation_specs:
+            if len(obs_spec.shape) == 1:
+                result += obs_spec.shape[0]
         return result
 
     def render(self, mode="rgb_array"):

gym-unity/gym_unity/envs/__init__.py

@@ -11,7 +11,7 @@
     TerminalSteps,
     BehaviorMapping,
 )
-from mlagents.trainers.tests.dummy_config import create_sensor_specs_with_shapes
+from mlagents.trainers.tests.dummy_config import create_observation_specs_with_shapes
 
 
 def test_gym_wrapper():
@@ -227,8 +227,8 @@ def create_mock_group_spec(
     obs_shapes = [(vector_observation_space_size,)]
     for _ in range(number_visual_observations):
         obs_shapes += [(8, 8, 3)]
-    sen_spec = create_sensor_specs_with_shapes(obs_shapes)
-    return BehaviorSpec(sen_spec, action_spec)
+    obs_spec = create_observation_specs_with_shapes(obs_shapes)
+    return BehaviorSpec(obs_spec, action_spec)
 
 
 def create_mock_vector_steps(specs, num_agents=1, number_visual_observations=0):

gym-unity/gym_unity/tests/test_gym.py

@@ -28,7 +28,7 @@
     Any,
     Mapping as MappingType,
 )
-from enum import IntFlag
+from enum import IntFlag, Enum
 import numpy as np
 
 from mlagents_envs.exception import UnityActionException
@@ -137,7 +137,7 @@ def empty(spec: "BehaviorSpec") -> "DecisionSteps":
         :param spec: The BehaviorSpec for the DecisionSteps
         """
         obs: List[np.ndarray] = []
-        for sen_spec in spec.sensor_specs:
+        for sen_spec in spec.observation_specs:
             obs += [np.zeros((0,) + sen_spec.shape, dtype=np.float32)]
         return DecisionSteps(
             obs=obs,
@@ -235,7 +235,7 @@ def empty(spec: "BehaviorSpec") -> "TerminalSteps":
         :param spec: The BehaviorSpec for the TerminalSteps
         """
         obs: List[np.ndarray] = []
-        for sen_spec in spec.sensor_specs:
+        for sen_spec in spec.observation_specs:
             obs += [np.zeros((0,) + sen_spec.shape, dtype=np.float32)]
         return TerminalSteps(
             obs=obs,
@@ -458,31 +458,49 @@ class DimensionProperty(IntFlag):
     VARIABLE_SIZE = 4
 
 
-class SensorSpec(NamedTuple):
+class ObservationType(Enum):
+    """
+    An Enum which defines the type of information carried in the observation
+    of the agent.
+    """
+
+    # Observation information is generic.
+    DEFAULT = 0
+    # Observation contains goal information for current task.
+    GOAL = 1
+    # Observation contains reward information for current task.
+    REWARD = 2
+    # Observation contains a message from another agent.
+    MESSAGE = 3
+
+
+class ObservationSpec(NamedTuple):
     """
     A NamedTuple containing information about the observation of Agents.
     - shape is a Tuple of int : It corresponds to the shape of
     an observation's dimensions.
     - dimension_property is a Tuple of DimensionProperties flag, one flag for each
     dimension.
+    - observation_type is an enum of ObservationType.
     """
 
     shape: Tuple[int, ...]
     dimension_property: Tuple[DimensionProperty, ...]
+    observation_type: ObservationType
 
 
 class BehaviorSpec(NamedTuple):
     """
     A NamedTuple containing information about the observation and action
     spaces for a group of Agents under the same behavior.
-    - sensor_specs is a List of SensorSpec NamedTuple containing
+    - observation_specs is a List of ObservationSpec NamedTuple containing
     information about the information of the Agent's observations such as their shapes.
     The order of the SensorSpec is the same as the order of the observations of an
     agent.
     - action_spec is an ActionSpec NamedTuple.
     """
 
-    sensor_specs: List[SensorSpec]
+    observation_specs: List[ObservationSpec]
     action_spec: ActionSpec