Support for randomizer tag inheritance

com.unity.perception/Runtime/Randomization/Randomizers/RandomizerTagManager.cs

     /// </summary>
     public class RandomizerTagManager
     {
+        Dictionary<Type, List<Type>> m_TypeCache = new Dictionary<Type, List<Type>>();
+        Dictionary<GameObject, HashSet<Type>> m_ObjectTags = new Dictionary<GameObject, HashSet<Type>>();
+        /// <param name="returnSubclasses">Should this method retrieve all tags derived from the passed in tag also?</param>
-        public IEnumerable<GameObject> Query<T>() where T : RandomizerTag
+        public IEnumerable<GameObject> Query<T>(bool returnSubclasses = false) where T : RandomizerTag
+
+
-                yield return taggedObject;
+            {
+                if (returnSubclasses)
+                {
+                    yield return taggedObject.gameObject;
+                }
+                else
+                {
+                    if (m_ObjectTags.ContainsKey(taggedObject) && m_ObjectTags[taggedObject].Contains(type))
+                        yield return taggedObject.gameObject;
+                }
+            }
-            if (!m_TagMap.ContainsKey(tagType))
-                m_TagMap[tagType] = new HashSet<GameObject>();
-            m_TagMap[tagType].Add(obj);
+            // Add tag to the game object to tag map
+            if (!m_ObjectTags.ContainsKey(obj))
+                m_ObjectTags[obj] = new HashSet<Type>();
+
+            m_ObjectTags[obj].Add(tagType);
+
+            if (!m_TypeCache.ContainsKey(tagType))
+            {
+                var inheritedTags = new List<Type>();
+
+                while (tagType != null && tagType != typeof(RandomizerTag))
+                {
+                    inheritedTags.Add(tagType);
+                    tagType = tagType.BaseType;
+                }
+
+                m_TypeCache[tagType] = inheritedTags;
+            }
+
+            var tags = m_TypeCache[tagType];
+
+            foreach (var tag in tags)
+            {
+                if (!m_TagMap.ContainsKey(tag))
+                    m_TagMap[tag] = new HashSet<GameObject>();
+                m_TagMap[tag].Add(obj);
+            }
-            if (m_TagMap.ContainsKey(tagType))
-                m_TagMap[tagType].Remove(obj);
+            // Grab all of the tags from the inheritance tree, and remove
+            // the passed in object from all of them
+            if (m_TypeCache.ContainsKey(tagType))
+            {
+                var tags = m_TypeCache[tagType];
+
+                foreach (var tag in tags.Where(tag => m_TagMap.ContainsKey(tag)))
+                {
+                    m_TagMap[tag].Remove(obj);
+                }
+            }
+
+            // Remove entry from the object to tag map
+            if (m_ObjectTags.ContainsKey(obj))
+            {
+                m_ObjectTags[obj].Remove(tagType);
+
+                if (!m_ObjectTags[obj].Any())
+                    m_ObjectTags.Remove(obj);
+            }
         }
     }
 }

com.unity.perception/Tests/Runtime/Randomization/RandomizerTests/ExampleTag.cs

 namespace RandomizationTests.RandomizerTests
 {
     public class ExampleTag : RandomizerTag { }
+
+    public class ExampleTag2 : ExampleTag { }
 }

com.unity.perception/Tests/Runtime/Randomization/RandomizerTests/RandomizerTagTests.cs

             for (var i = 0; i < copyCount - 1; i++)
                 Object.Instantiate(gameObject);
 
+            var gameObject2 = new GameObject();
+            gameObject2.AddComponent<ExampleTag2>();
+            for (var i = 0; i < copyCount - 1; i++)
+                Object.Instantiate(gameObject2);
+
+
+            queriedObjects = m_Scenario.tagManager.Query<ExampleTag2>().ToArray();
+            Assert.AreEqual(queriedObjects.Length, copyCount);
+
+            queriedObjects = m_Scenario.tagManager.Query<ExampleTag>(true).ToArray();
+            Assert.AreEqual(queriedObjects.Length, copyCount * 2);
         }
     }
 }