InputActuatorComponent to allow the generation of an action space from an InputActionAsset (#4881)

4 年前 · b0719885
--- a/.yamato/com.unity.ml-agents-test.yml
+++ b/.yamato/com.unity.ml-agents-test.yml
    assembly: Unity.ML-Agents
    minCoveragePct: 72
  - name: com.unity.ml-agents.extensions
-    assembly: Unity.ML-Agents.Extensions
+    assembly: Unity.ML-Agents.Extensions*
    minCoveragePct: 75
 ---

--- a/DevProject/Packages/manifest.json
+++ b/DevProject/Packages/manifest.json
  "dependencies": {
    "com.unity.2d.sprite": "1.0.0",
    "com.unity.2d.tilemap": "1.0.0",
-    "com.unity.ads": "3.4.9",
+    "com.unity.ads": "3.6.1",
-    "com.unity.ide.vscode": "1.2.1",
+    "com.unity.ide.vscode": "1.2.3",
+    "com.unity.inputsystem": "1.1.0-preview.3",
-    "com.unity.multiplayer-hlapi": "1.0.6",
+    "com.unity.multiplayer-hlapi": "1.0.8",
-    "com.unity.purchasing": "2.1.0",
-    "com.unity.test-framework": "1.1.16",
+    "com.unity.purchasing": "2.2.1",
+    "com.unity.test-framework": "1.1.20",
-    "com.unity.xr.legacyinputhelpers": "2.1.4",
+    "com.unity.xr.legacyinputhelpers": "2.1.7",
    "com.unity.modules.ai": "1.0.0",
    "com.unity.modules.androidjni": "1.0.0",
    "com.unity.modules.animation": "1.0.0",
  "registry": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates",
  "testables": [
    "com.unity.ml-agents",
-    "com.unity.ml-agents.extensions"
+    "com.unity.ml-agents.extensions",
+    "com.unity.inputsystem"
  ]
 }
--- a/DevProject/ProjectSettings/ProjectSettings.asset
+++ b/DevProject/ProjectSettings/ProjectSettings.asset
  androidBlitType: 0
  defaultIsNativeResolution: 1
  macRetinaSupport: 1
-  runInBackground: 0
+  runInBackground: 1
  captureSingleScreen: 0
  muteOtherAudioSources: 0
  Prepare IOS For Recording: 0
  xboxOneMonoLoggingLevel: 0
  xboxOneLoggingLevel: 1
  xboxOneDisableEsram: 0
+  xboxOneEnableTypeOptimization: 0
  xboxOnePresentImmediateThreshold: 0
  switchQueueCommandMemory: 1048576
  switchQueueControlMemory: 16384
  switchNVNOtherPoolsGranularity: 16777216
+  switchNVNMaxPublicTextureIDCount: 0
+  switchNVNMaxPublicSamplerIDCount: 0
+  stadiaPresentMode: 0
+  stadiaTargetFramerate: 0
+  vulkanEnableLateAcquireNextImage: 0
  m_SupportedAspectRatios:
    4:3: 1
    5:4: 1
  useHDRDisplay: 0
  D3DHDRBitDepth: 0
  m_ColorGamuts: 00000000
-  targetPixelDensity: 0
+  targetPixelDensity: 30
  resolutionScalingMode: 0
  androidSupportedAspectRatio: 1
  androidMaxAspectRatio: 2.1
  StripUnusedMeshComponents: 0
  VertexChannelCompressionMask: 4054
  iPhoneSdkVersion: 988
-  iOSTargetOSVersionString: 
+  iOSTargetOSVersionString: 10.0
-  tvOSTargetOSVersionString: 
+  tvOSTargetOSVersionString: 10.0
  uIPrerenderedIcon: 0
  uIRequiresPersistentWiFi: 0
  uIRequiresFullScreen: 1
-  iPhoneSplashScreen: {fileID: 0}
-  iPhoneHighResSplashScreen: {fileID: 0}
-  iPhoneTallHighResSplashScreen: {fileID: 0}
-  iPhone47inSplashScreen: {fileID: 0}
-  iPhone55inPortraitSplashScreen: {fileID: 0}
-  iPhone55inLandscapeSplashScreen: {fileID: 0}
-  iPhone58inPortraitSplashScreen: {fileID: 0}
-  iPhone58inLandscapeSplashScreen: {fileID: 0}
-  iPadPortraitSplashScreen: {fileID: 0}
-  iPadHighResPortraitSplashScreen: {fileID: 0}
-  iPadLandscapeSplashScreen: {fileID: 0}
-  iPadHighResLandscapeSplashScreen: {fileID: 0}
-  iPhone65inPortraitSplashScreen: {fileID: 0}
-  iPhone65inLandscapeSplashScreen: {fileID: 0}
-  iPhone61inPortraitSplashScreen: {fileID: 0}
-  iPhone61inLandscapeSplashScreen: {fileID: 0}
  appleTVSplashScreen: {fileID: 0}
  appleTVSplashScreen2x: {fileID: 0}
  tvOSSmallIconLayers: []
  metalEditorSupport: 1
  metalAPIValidation: 1
  iOSRenderExtraFrameOnPause: 0
+  iosCopyPluginsCodeInsteadOfSymlink: 0
  appleDeveloperTeamID: 
  iOSManualSigningProvisioningProfileID: 
  tvOSManualSigningProvisioningProfileID: 
  ps4ShareFilePath: 
  ps4ShareOverlayImagePath: 
  ps4PrivacyGuardImagePath: 
+  ps4ExtraSceSysFile: 
  ps4NPtitleDatPath: 
  ps4RemotePlayKeyAssignment: -1
  ps4RemotePlayKeyMappingDir: 
  ps4UseResolutionFallback: 0
  ps4ReprojectionSupport: 0
  ps4UseAudio3dBackend: 0
+  ps4UseLowGarlicFragmentationMode: 1
  ps4SocialScreenEnabled: 0
  ps4ScriptOptimizationLevel: 2
  ps4Audio3dVirtualSpeakerCount: 14
  ps4disableAutoHideSplash: 0
  ps4videoRecordingFeaturesUsed: 0
  ps4contentSearchFeaturesUsed: 0
+  ps4CompatibilityPS5: 0
+  ps4GPU800MHz: 1
  ps4attribEyeToEyeDistanceSettingVR: 0
  ps4IncludedModules: []
  ps4attribVROutputEnabled: 0
  additionalIl2CppArgs: 
  scriptingRuntimeVersion: 1
  gcIncremental: 0
+  assemblyVersionValidation: 1
  gcWBarrierValidation: 0
  apiCompatibilityLevelPerPlatform: {}
  m_RenderingPath: 1
  XboxOneCapability: []
  XboxOneGameRating: {}
  XboxOneIsContentPackage: 0
+  XboxOneEnhancedXboxCompatibilityMode: 0
  XboxOneEnableGPUVariability: 1
  XboxOneSockets: {}
  XboxOneSplashScreen: {fileID: 0}
  XboxOneOverrideIdentityName: 
+  XboxOneOverrideIdentityPublisher: 
  vrEditorSettings:
    daydream:
      daydreamIconForeground: {fileID: 0}
  projectName: 
  organizationId: 
  cloudEnabled: 0
-  enableNativePlatformBackendsForNewInputSystem: 0
-  disableOldInputManagerSupport: 0
+  enableNativePlatformBackendsForNewInputSystem: 1
+  disableOldInputManagerSupport: 1
  legacyClampBlendShapeWeights: 0
--- a/DevProject/ProjectSettings/ProjectVersion.txt
+++ b/DevProject/ProjectSettings/ProjectVersion.txt
-m_EditorVersion: 2019.4.7f1
-m_EditorVersionWithRevision: 2019.4.7f1 (e992b1a16e65)
+m_EditorVersion: 2019.4.19f1
+m_EditorVersionWithRevision: 2019.4.19f1 (ca5b14067cec)
--- a/com.unity.ml-agents.extensions/Documentation~/com.unity.ml-agents.extensions.md
+++ b/com.unity.ml-agents.extensions/Documentation~/com.unity.ml-agents.extensions.md
 * [Match-3 sensor and actuator](Match3.md)
 * [Grid-based sensor](Grid-Sensor.md)
 * Physics-based sensors
+ * [Input System Package Integration](InputActuatorComponent.md)

 ## Installation
 The ML-Agents Extensions package is not currently available in the Package Manager. There are two
 This version of the Unity ML-Agents Extensions package is compatible with the
 following versions of the Unity Editor:

- 2018.4 and later
+- If using the `InputActuatorComponent`
+    - 2019.4 or later
+    - install the `com.unity.inputsystem` package version `1.1.0-preview.3` or later.
+- Else 2018.4 and later
-none
+- For the `InputActuatorComponent`
+    - Limited implementation of `InputControls`
+    - No way to customize the action space of the `InputActuatorComponent`

 ## Need Help?
 The main [README](https://github.com/Unity-Technologies/ml-agents/tree/release_12_docs/README.md) contains links for contacting the team or getting support.
--- a/com.unity.ml-agents.extensions/Runtime/Unity.ML-Agents.Extensions.asmdef
+++ b/com.unity.ml-agents.extensions/Runtime/Unity.ML-Agents.Extensions.asmdef
    "name": "Unity.ML-Agents.Extensions",
    "references": [
        "Unity.Barracuda",
-        "Unity.ML-Agents"
+        "Unity.ML-Agents",
+        "Unity.ML-Agents.Extensions.Input"
    ],
    "includePlatforms": [],
    "excludePlatforms": []
--- a/com.unity.ml-agents/CHANGELOG.md
+++ b/com.unity.ml-agents/CHANGELOG.md
 - `InferenceDevice.Burst` was added, indicating that Agent's model will be run using Barracuda's Burst backend.
  This is the default for new Agents, but existing ones that use `InferenceDevice.CPU` should update to
  `InferenceDevice.Burst`. (#4925)
+- Add an InputActuatorComponent to allow the generation of Agent action spaces from an InputActionAsset.
+  Projects wanting to use this feature will need to add the
+  [Input System Package](https://docs.unity3d.com/Packages/com.unity.inputsystem@1.1/manual/index.html)
+  at version 1.1.0-preview.3 or later. (#4881)

 #### ml-agents / ml-agents-envs / gym-unity (Python)
 - Tensorboard now logs the Environment Reward as both a scalar and a histogram. (#4878)
--- a/com.unity.ml-agents/Runtime/Actuators/ActionSpec.cs
+++ b/com.unity.ml-agents/Runtime/Actuators/ActionSpec.cs
 using System;
+using System.Collections.Generic;
-using Unity.MLAgents.Policies;
 using UnityEngine;

 namespace Unity.MLAgents.Actuators
        public ActionSpec(int numContinuousActions = 0, int[] discreteBranchSizes = null)
        {
            m_NumContinuousActions = numContinuousActions;
-            BranchSizes = discreteBranchSizes;
+            BranchSizes = discreteBranchSizes ?? Array.Empty<int>();
        }

        /// <summary>
                    "ActionSpecs must be all continuous or all discrete."
                );
            }
+        }
+
+        /// <summary>
+        /// Combines a list of actions specs and allocates a new array of branch sizes if needed.
+        /// </summary>
+        /// <param name="specs">The list of action specs to combine.</param>
+        /// <returns>An ActionSpec which represents the aggregate of the ActionSpecs passed in.</returns>
+        public static ActionSpec Combine(params ActionSpec[] specs)
+        {
+            var numContinuous = 0;
+            var numDiscrete = 0;
+            for (var i = 0; i < specs.Length; i++)
+            {
+                var spec = specs[i];
+                numContinuous += spec.NumContinuousActions;
+                numDiscrete += spec.NumDiscreteActions;
+            }
+
+            if (numDiscrete <= 0)
+            {
+                return MakeContinuous(numContinuous);
+            }
+
+            var branchSizes = new int[numDiscrete];
+            var offset = 0;
+            for (var i = 0; i < specs.Length; i++)
+            {
+                var spec = specs[i];
+                if (spec.BranchSizes.Length == 0)
+                {
+                    continue;
+                }
+                var branchSizesLength = spec.BranchSizes.Length;
+                Array.Copy(spec.BranchSizes,
+                    0,
+                    branchSizes,
+                    offset,
+                    branchSizesLength);
+                offset += branchSizesLength;
+            }
+            return new ActionSpec(numContinuous, branchSizes);
        }
    }
 }
--- a/com.unity.ml-agents/Runtime/Actuators/ActuatorManager.cs
+++ b/com.unity.ml-agents/Runtime/Actuators/ActuatorManager.cs
 using System.Collections;
 using System.Collections.Generic;
 using UnityEngine;
+using UnityEngine.Profiling;

 namespace Unity.MLAgents.Actuators
 {
        /// actions for the IActuators in this list.</param>
        public void UpdateActions(ActionBuffers actions)
        {
+            Profiler.BeginSample("ActuatorManager.UpdateActions");
+            Profiler.EndSample();
        }

        static void UpdateActionArray<T>(ActionSegment<T> sourceActionBuffer, ActionSegment<T> destination)
        /// </summary>
        public void ApplyHeuristic(in ActionBuffers actionBuffersOut)
        {
+            Profiler.BeginSample("ActuatorManager.ApplyHeuristic");
            var continuousStart = 0;
            var discreteStart = 0;
            for (var i = 0; i < m_Actuators.Count; i++)
                continuousStart += numContinuousActions;
                discreteStart += numDiscreteActions;
            }
+            Profiler.EndSample();
        }

        /// <summary>
        /// </summary>
        public void ExecuteActions()
        {
+            Profiler.BeginSample("ActuatorManager.ExecuteActions");
            ReadyActuatorsForExecution();
            var continuousStart = 0;
            var discreteStart = 0;
                var numContinuousActions = actuator.ActionSpec.NumContinuousActions;
                var numDiscreteActions = actuator.ActionSpec.NumDiscreteActions;

+                if (numContinuousActions == 0 && numDiscreteActions == 0)
+                {
+                    continue;
+                }
+
                var continuousActions = ActionSegment<float>.Empty;
                if (numContinuousActions > 0)
                {
                continuousStart += numContinuousActions;
                discreteStart += numDiscreteActions;
            }
+            Profiler.EndSample();
        }

        /// <summary>
--- a/com.unity.ml-agents/Runtime/Actuators/VectorActuator.cs
+++ b/com.unity.ml-agents/Runtime/Actuators/VectorActuator.cs
+using UnityEngine.Profiling;
+
 namespace Unity.MLAgents.Actuators
 {
    /// <summary>
        /// <inheritdoc />
        public void OnActionReceived(ActionBuffers actionBuffers)
        {
-            ActionBuffers = actionBuffers;
-            m_ActionReceiver.OnActionReceived(ActionBuffers);
+            Profiler.BeginSample("VectorActuator.OnActionReceived");
+            m_ActionBuffers = actionBuffers;
+            m_ActionReceiver.OnActionReceived(m_ActionBuffers);
+            Profiler.EndSample();
+            Profiler.BeginSample("VectorActuator.Heuristic");
+            Profiler.EndSample();
        }

        /// <inheritdoc />
--- a/com.unity.ml-agents/Runtime/Agent.deprecated.cs
+++ b/com.unity.ml-agents/Runtime/Agent.deprecated.cs
 using System;
 using UnityEngine;
+using UnityEngine.Profiling;

 namespace Unity.MLAgents
 {
        [Obsolete("GetAction has been deprecated, please use GetStoredActionBuffers instead.")]
        public float[] GetAction()
        {
+            Profiler.BeginSample("Agent.GetAction.Deprecated");
            var actionSpec = m_PolicyFactory.BrainParameters.ActionSpec;
            // For continuous and discrete actions together, this shouldn't be called because we can only return one.
            if (actionSpec.NumContinuousActions > 0 && actionSpec.NumDiscreteActions > 0)
            {
                return storedAction.ContinuousActions.Array;
            }
-            else
-            {
-                return Array.ConvertAll(storedAction.DiscreteActions.Array, x => (float)x);
-            }
+            Profiler.EndSample();
+            return Array.ConvertAll(storedAction.DiscreteActions.Array, x => (float)x);
        }
    }
 }
--- a/com.unity.ml-agents/Runtime/AssemblyInfo.cs
+++ b/com.unity.ml-agents/Runtime/AssemblyInfo.cs
 [assembly: InternalsVisibleTo("Unity.ML-Agents.Editor.Tests")]
 [assembly: InternalsVisibleTo("Unity.ML-Agents.Editor")]
 [assembly: InternalsVisibleTo("Unity.ML-Agents.Extensions")]
+[assembly: InternalsVisibleTo("Unity.ML-Agents.Extensions.Input")]
--- a/com.unity.ml-agents/Runtime/Policies/BehaviorParameters.cs
+++ b/com.unity.ml-agents/Runtime/Policies/BehaviorParameters.cs
        [HideInInspector, SerializeField]
        BrainParameters m_BrainParameters = new BrainParameters();

+        public delegate void PolicyUpdated(bool isInHeuristicMode);
+
+        internal event PolicyUpdated OnPolicyUpdated;
+
        /// <summary>
        /// The associated <see cref="Policies.BrainParameters"/> for this behavior.
        /// </summary>
            get { return m_BehaviorName + "?team=" + TeamId; }
        }

+        void Awake()
+        {
+            OnPolicyUpdated += mode => { };
+        }
+
        internal IPolicy GeneratePolicy(ActionSpec actionSpec, ActuatorManager actuatorManager)
        {
            switch (m_BehaviorType)
            }
        }

+        /// <summary>
+        /// Query the behavior parameters in order to see if the Agent is running in Heuristic Mode.
+        /// </summary>
+        /// <returns>true if the Agent is running in Heuristic mode.</returns>
+        public bool IsInHeuristicMode()
+        {
+            if (BehaviorType == BehaviorType.HeuristicOnly)
+            {
+                return true;
+            }
+
+            return BehaviorType == BehaviorType.Default &&
+                ReferenceEquals(Model, null) &&
+                (!Academy.IsInitialized ||
+                    Academy.IsInitialized &&
+                    !Academy.Instance.IsCommunicatorOn);
+        }
+
        internal void UpdateAgentPolicy()
        {
            var agent = GetComponent<Agent>();
            }
            agent.ReloadPolicy();
+            OnPolicyUpdated?.Invoke(IsInHeuristicMode());
        }
    }
 }
--- a/com.unity.ml-agents/Tests/Editor/BehaviorParameterTests.cs
+++ b/com.unity.ml-agents/Tests/Editor/BehaviorParameterTests.cs
 using NUnit.Framework;
+using Unity.Barracuda;
+using UnityEditor;
+using UnityEngine.TestTools;

 namespace Unity.MLAgents.Tests
 {
+        const string k_continuousONNXPath = "Packages/com.unity.ml-agents/Tests/Editor/TestModels/continuous2vis8vec2action.onnx";
        public void Heuristic(in ActionBuffers actionsOut)
        {
            // No-op
            {
                bp.GeneratePolicy(actionSpec, new ActuatorManager());
            });
+        }
+
+        [Test]
+        public void TestIsInHeuristicMode()
+        {
+            var gameObj = new GameObject();
+            var bp = gameObj.AddComponent<BehaviorParameters>();
+            bp.Model = null;
+            gameObj.AddComponent<Agent>();
+            bp.BehaviorType = BehaviorType.HeuristicOnly;
+            Assert.IsTrue(bp.IsInHeuristicMode());
+
+            bp.BehaviorType = BehaviorType.Default;
+            Assert.IsTrue(bp.IsInHeuristicMode());
+
+            bp.Model = ScriptableObject.CreateInstance<NNModel>();
+            Assert.IsFalse(bp.IsInHeuristicMode());
+        }
+
+        [Test]
+        public void TestPolicyUpdateEventFired()
+        {
+            var gameObj = new GameObject();
+            var bp = gameObj.AddComponent<BehaviorParameters>();
+            gameObj.AddComponent<Agent>().LazyInitialize();
+            bp.OnPolicyUpdated += delegate (bool isInHeuristicMode) { Debug.Log($"OnPolicyChanged:{isInHeuristicMode}"); };
+            bp.BehaviorType = BehaviorType.HeuristicOnly;
+            LogAssert.Expect(LogType.Log, $"OnPolicyChanged:{true}");
+
+            bp.BehaviorType = BehaviorType.Default;
+            LogAssert.Expect(LogType.Log, $"OnPolicyChanged:{true}");
+
+            Assert.Throws<UnityAgentsException>(() =>
+            {
+                bp.BehaviorType = BehaviorType.InferenceOnly;
+            });
+
+            bp.Model = AssetDatabase.LoadAssetAtPath<NNModel>(k_continuousONNXPath);
+            LogAssert.Expect(LogType.Log, $"OnPolicyChanged:{false}");
+
+            bp.BehaviorType = BehaviorType.HeuristicOnly;
+            LogAssert.Expect(LogType.Log, $"OnPolicyChanged:{true}");
        }
    }
 }
--- a/DevProject/DevProject.sln.DotSettings
+++ b/DevProject/DevProject.sln.DotSettings
+<wpf:ResourceDictionary xml:space="preserve" xmlns:x="http://schemas.microsoft.com/winfx/2006/xaml" xmlns:s="clr-namespace:System;assembly=mscorlib" xmlns:ss="urn:shemas-jetbrains-com:settings-storage-xaml" xmlns:wpf="http://schemas.microsoft.com/winfx/2006/xaml/presentation">
+	<s:String x:Key="/Default/CodeStyle/Naming/CSharpNaming/Abbreviations/=RL/@EntryIndexedValue">RL</s:String></wpf:ResourceDictionary>
--- a/DevProject/Packages/packages-lock.json
+++ b/DevProject/Packages/packages-lock.json
+{
+  "dependencies": {
+    "com.unity.2d.sprite": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.2d.tilemap": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.ads": {
+      "version": "3.6.1",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {
+        "com.unity.ugui": "1.0.0"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.analytics": {
+      "version": "3.3.5",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {
+        "com.unity.ugui": "1.0.0"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.barracuda": {
+      "version": "1.3.0-preview",
+      "depth": 1,
+      "source": "registry",
+      "dependencies": {
+        "com.unity.burst": "1.3.4",
+        "com.unity.modules.jsonserialize": "1.0.0",
+        "com.unity.modules.imageconversion": "1.0.0"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.burst": {
+      "version": "1.3.4",
+      "depth": 2,
+      "source": "registry",
+      "dependencies": {
+        "com.unity.mathematics": "1.2.1"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.coding": {
+      "version": "0.1.0-preview.13",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {
+        "nuget.moq": "1.0.0"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.collab-proxy": {
+      "version": "1.2.16",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {},
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.ext.nunit": {
+      "version": "1.0.6",
+      "depth": 1,
+      "source": "registry",
+      "dependencies": {},
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.ide.rider": {
+      "version": "1.1.4",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {
+        "com.unity.test-framework": "1.1.1"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.ide.vscode": {
+      "version": "1.2.3",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {},
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.inputsystem": {
+      "version": "1.1.0-preview.3",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {},
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.mathematics": {
+      "version": "1.2.1",
+      "depth": 3,
+      "source": "registry",
+      "dependencies": {},
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.ml-agents": {
+      "version": "file:../../com.unity.ml-agents",
+      "depth": 0,
+      "source": "local",
+      "dependencies": {
+        "com.unity.barracuda": "1.3.0-preview",
+        "com.unity.modules.imageconversion": "1.0.0",
+        "com.unity.modules.jsonserialize": "1.0.0",
+        "com.unity.modules.physics": "1.0.0",
+        "com.unity.modules.physics2d": "1.0.0",
+        "com.unity.modules.unityanalytics": "1.0.0"
+      }
+    },
+    "com.unity.ml-agents.extensions": {
+      "version": "file:../../com.unity.ml-agents.extensions",
+      "depth": 0,
+      "source": "local",
+      "dependencies": {
+        "com.unity.ml-agents": "1.7.2-preview"
+      }
+    },
+    "com.unity.multiplayer-hlapi": {
+      "version": "1.0.8",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {
+        "nuget.mono-cecil": "0.1.6-preview"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.nuget.mono-cecil": {
+      "version": "0.1.6-preview.2",
+      "depth": 1,
+      "source": "registry",
+      "dependencies": {
+        "nuget.mono-cecil": "0.1.6-preview"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.nuget.newtonsoft-json": {
+      "version": "2.0.0-preview",
+      "depth": 1,
+      "source": "registry",
+      "dependencies": {},
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.package-manager-doctools": {
+      "version": "1.7.0-preview",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {
+        "com.unity.package-validation-suite": "0.10.0-preview",
+        "com.unity.nuget.newtonsoft-json": "2.0.0-preview"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.package-validation-suite": {
+      "version": "0.19.0-preview",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {
+        "com.unity.nuget.mono-cecil": "0.1.6-preview.2"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.purchasing": {
+      "version": "2.2.1",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {
+        "com.unity.ugui": "1.0.0"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.settings-manager": {
+      "version": "1.0.1",
+      "depth": 1,
+      "source": "registry",
+      "dependencies": {},
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.test-framework": {
+      "version": "1.1.20",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {
+        "com.unity.ext.nunit": "1.0.6",
+        "com.unity.modules.imgui": "1.0.0",
+        "com.unity.modules.jsonserialize": "1.0.0"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.test-framework.performance": {
+      "version": "2.2.0-preview",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {
+        "com.unity.test-framework": "1.1.0",
+        "com.unity.nuget.newtonsoft-json": "2.0.0-preview"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.testtools.codecoverage": {
+      "version": "1.0.0-pre.3",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {
+        "com.unity.test-framework": "1.0.16",
+        "com.unity.settings-manager": "1.0.1"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.textmeshpro": {
+      "version": "2.0.1",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {
+        "com.unity.ugui": "1.0.0"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.timeline": {
+      "version": "1.2.12",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {},
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.ugui": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.ui": "1.0.0",
+        "com.unity.modules.imgui": "1.0.0"
+      }
+    },
+    "com.unity.xr.legacyinputhelpers": {
+      "version": "2.1.7",
+      "depth": 0,
+      "source": "registry",
+      "dependencies": {
+        "com.unity.modules.vr": "1.0.0",
+        "com.unity.modules.xr": "1.0.0"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "nuget.castle-core": {
+      "version": "1.0.1",
+      "depth": 2,
+      "source": "registry",
+      "dependencies": {},
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "nuget.mono-cecil": {
+      "version": "0.1.6-preview",
+      "depth": 1,
+      "source": "registry",
+      "dependencies": {},
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "nuget.moq": {
+      "version": "1.0.0",
+      "depth": 1,
+      "source": "registry",
+      "dependencies": {
+        "nuget.castle-core": "1.0.1"
+      },
+      "url": "https://artifactory.prd.cds.internal.unity3d.com/artifactory/api/npm/upm-candidates"
+    },
+    "com.unity.modules.ai": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.androidjni": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.animation": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.assetbundle": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.audio": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.cloth": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.physics": "1.0.0"
+      }
+    },
+    "com.unity.modules.director": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.audio": "1.0.0",
+        "com.unity.modules.animation": "1.0.0"
+      }
+    },
+    "com.unity.modules.imageconversion": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.imgui": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.jsonserialize": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.particlesystem": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.physics": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.physics2d": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.screencapture": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.imageconversion": "1.0.0"
+      }
+    },
+    "com.unity.modules.subsystems": {
+      "version": "1.0.0",
+      "depth": 1,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.jsonserialize": "1.0.0"
+      }
+    },
+    "com.unity.modules.terrain": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.terrainphysics": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.physics": "1.0.0",
+        "com.unity.modules.terrain": "1.0.0"
+      }
+    },
+    "com.unity.modules.tilemap": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.physics2d": "1.0.0"
+      }
+    },
+    "com.unity.modules.ui": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.uielements": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.imgui": "1.0.0",
+        "com.unity.modules.jsonserialize": "1.0.0"
+      }
+    },
+    "com.unity.modules.umbra": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.unityanalytics": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.unitywebrequest": "1.0.0",
+        "com.unity.modules.jsonserialize": "1.0.0"
+      }
+    },
+    "com.unity.modules.unitywebrequest": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.unitywebrequestassetbundle": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.assetbundle": "1.0.0",
+        "com.unity.modules.unitywebrequest": "1.0.0"
+      }
+    },
+    "com.unity.modules.unitywebrequestaudio": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.unitywebrequest": "1.0.0",
+        "com.unity.modules.audio": "1.0.0"
+      }
+    },
+    "com.unity.modules.unitywebrequesttexture": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.unitywebrequest": "1.0.0",
+        "com.unity.modules.imageconversion": "1.0.0"
+      }
+    },
+    "com.unity.modules.unitywebrequestwww": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.unitywebrequest": "1.0.0",
+        "com.unity.modules.unitywebrequestassetbundle": "1.0.0",
+        "com.unity.modules.unitywebrequestaudio": "1.0.0",
+        "com.unity.modules.audio": "1.0.0",
+        "com.unity.modules.assetbundle": "1.0.0",
+        "com.unity.modules.imageconversion": "1.0.0"
+      }
+    },
+    "com.unity.modules.vehicles": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.physics": "1.0.0"
+      }
+    },
+    "com.unity.modules.video": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.audio": "1.0.0",
+        "com.unity.modules.ui": "1.0.0",
+        "com.unity.modules.unitywebrequest": "1.0.0"
+      }
+    },
+    "com.unity.modules.vr": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.jsonserialize": "1.0.0",
+        "com.unity.modules.physics": "1.0.0",
+        "com.unity.modules.xr": "1.0.0"
+      }
+    },
+    "com.unity.modules.wind": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {}
+    },
+    "com.unity.modules.xr": {
+      "version": "1.0.0",
+      "depth": 0,
+      "source": "builtin",
+      "dependencies": {
+        "com.unity.modules.physics": "1.0.0",
+        "com.unity.modules.jsonserialize": "1.0.0",
+        "com.unity.modules.subsystems": "1.0.0"
+      }
+    }
+  }
+}
--- a/com.unity.ml-agents.extensions/Documentation~/InputActuatorComponent.md
+++ b/com.unity.ml-agents.extensions/Documentation~/InputActuatorComponent.md
+# Integration of the Input System Package with ML-Agents
+
+## Overview
+One area we are always trying to improve is getting developers up and running with ML-Agents.  With this in mind,
+we have implemented an `InputActuatorComponent`.  This component integrates with the
+[Input System Package](https://docs.unity3d.com/Packages/com.unity.inputsystem@1.1/manual/QuickStartGuide.html)
+to set up an action space for your `Agent` based on an `InputActionAsset` that is referenced by the
+`IInputActionAssetProvider` interface, or the `PlayerInput` component that may be living on your player controlled
+`Agent`.  This means that if you have code outside of your agent that handles input, you will not need to implement
+the Heuristic function in agent as well.  The `InputActuatorComponent` will handle this for you.  You can now train and
+run inference on `Agents` with an action space defined by an `InputActionAsset`.
+
+This implementation includes:
+
+* C# `InputActuatorComponent` you can attach to your Agent.
+* Implement the `IInputActionAssetProvider` in the `Componenet` where you handle player input.
+* An example environment where the input handling code is not in the Heuristic function of the Agent subclass.
+
+### Feedback
+We have only implemented a small subset of `InputControl` types that we thought would cover a large portion of what
+most developers would use.  Please let us know if you want more control types implemented by posting in the [ML-Agents
+forum.](https://forum.unity.com/forums/ml-agents.453/)
+
+We would also like your feedback on the workflow of integrating this into your games.  If you run
+into workflow issues please let us know in the ML-Agents forums, or if you've discovered a bug,
+please file a bug on our GitHub page.
+
+## Getting started
+The C# code for the `InputActuatorComponent` exists inside of the extensions package (com.unity.ml-agents.extensions).  A good first step would be to familiarize with the extensions package by reading the document [here](com.unity.ml-agents.extensions.md).  The second step would be to take a look at how we have implemented the C# code in the example Input Integration scene (located under  ML-Agents-Input-Example/Assets/ML-Agents/Examples/PushBlock/).  Once you have some familiarity, then the next step would be to add the InputActuatorComponent to your player Agent.  The example we have implemented uses C# Events to send information from the Input System.
+
+Additionally, see below for additional technical specifications on the C# code for the InputActuatorComponent.
+
+## Technical specifications for the InputActuatorComponent
+
+### `IInputActionsAssetProvider` Interface
+The `InputActuatorComponent` searches for a `Component` that implements
+`IInputActionAssetProvider` on the `GameObject` they both are attached to.  It is important to note
+that if multiple `Components` on your `GameObject` need to access an `InputActionAsset` to handle events,
+they will need to share the same instance of the `InputActionAsset` that is returned from the
+`IInputActionAssetProvider`.
+
+### `InputActuatorComponent` class
+The `InputActuatorComponent` is the bridge between ML-Agents and the Input System.. It allows ML-Agents to
+* create an `ActionSpec` for your Agent based on an `InputActionAsset` that comes from an
+`IInputActionAssetProvider`.
+* send simulated input from a training process or a neural network
+* let developers keep their input handling code in one place
+
+This is accomplished by adding the `InputActuatorComponenet` to an Agent which already has the PlayerInput component attached.
+
+### Setting up a scene using the `InputActuatorComponent`
+1. Add the `com.unity.inputsystem` version 1.1.0-preview.3 or later to your project via the Package Manager window.
+2. If you have already setup an InputActionAsset skip to Step 3, otherwise follow these sub steps:
+    1. Create an InputActionAsset to allow your Agent to be controlled by the Input System.
+    2. Handle the events from the Input System where you normally would (i.e. a script external to your Agent class).
+3. Add the InputSystemActuatorComponent to the GameObject that has the `PlayerInput` and `Agent` components attached.
+
--- a/com.unity.ml-agents.extensions/Editor/Input.meta
+++ b/com.unity.ml-agents.extensions/Editor/Input.meta
+fileFormatVersion: 2
+guid: 1f773a20e85042999e87f5d3c7b55281
+timeCreated: 1613637190
--- a/com.unity.ml-agents.extensions/Runtime/Input.meta
+++ b/com.unity.ml-agents.extensions/Runtime/Input.meta
+fileFormatVersion: 2
+guid: 1694e881b9ec420ba1c201f0612392d6
+timeCreated: 1610754907
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input.meta
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input.meta
+fileFormatVersion: 2
+guid: 27f8e1ce37d7485f814ce50a37101203
+timeCreated: 1612908869
--- a/com.unity.ml-agents/Tests/Editor/Actuators/ActionSpecTests.cs
+++ b/com.unity.ml-agents/Tests/Editor/Actuators/ActionSpecTests.cs
+using System;
+using System.Collections.Generic;
+using System.Linq;
+using NUnit.Framework;
+using Unity.MLAgents.Actuators;
+using UnityEngine.TestTools.Constraints;
+using Is = UnityEngine.TestTools.Constraints.Is;
+
+namespace Unity.MLAgents.Tests.Actuators
+{
+    [TestFixture]
+    public class ActionSpecTests
+    {
+        [Test]
+        public void ActionSpecCombineTest()
+        {
+            var as0 = new ActionSpec(3, new[] { 3, 2, 1 });
+            var as1 = new ActionSpec(1, new[] { 35, 122, 1, 3, 8, 3 });
+
+            var as0NumCon = 3;
+            var as0NumDis = as0.NumDiscreteActions;
+            var as1NumCon = 1;
+            var as1NumDis = as1.NumDiscreteActions;
+            var branchSizes = new List<int>();
+            branchSizes.AddRange(as0.BranchSizes);
+            branchSizes.AddRange(as1.BranchSizes);
+
+            var asc = ActionSpec.Combine(as0, as1);
+
+            Assert.AreEqual(as0NumCon + as1NumCon, asc.NumContinuousActions);
+            Assert.AreEqual(as0NumDis + as1NumDis, asc.NumDiscreteActions);
+            Assert.IsTrue(branchSizes.ToArray().SequenceEqual(asc.BranchSizes));
+
+            as0 = new ActionSpec(3);
+            as1 = new ActionSpec(1);
+            asc = ActionSpec.Combine(as0, as1);
+            Assert.IsEmpty(asc.BranchSizes);
+        }
+    }
+}
--- a/com.unity.ml-agents/Tests/Editor/Actuators/ActionSpecTests.cs.meta
+++ b/com.unity.ml-agents/Tests/Editor/Actuators/ActionSpecTests.cs.meta
+fileFormatVersion: 2
+guid: 99d76ec04c944b75bc6b85abfff4ac4e
+timeCreated: 1613680505
--- a/DevProject/ProjectSettings/Packages/com.unity.testtools.codecoverage/Settings.json
+++ b/DevProject/ProjectSettings/Packages/com.unity.testtools.codecoverage/Settings.json
+{
+    "m_Name": "Settings",
+    "m_Path": "ProjectSettings/Packages/com.unity.testtools.codecoverage/Settings.json",
+    "m_Dictionary": {
+        "m_DictionaryValues": []
+    }
+}
--- a/com.unity.ml-agents.extensions/Editor/Input/InputActuatorComponentEditor.cs
+++ b/com.unity.ml-agents.extensions/Editor/Input/InputActuatorComponentEditor.cs
+#if MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
+using Unity.MLAgents.Extensions.Input;
+using UnityEditor;
+
+namespace Unity.MLAgents.Extensions.Editor.Input
+{
+    [CustomEditor(typeof(InputActuatorComponent))]
+    internal class InputActuatorComponentEditor : UnityEditor.Editor
+    {
+        const string k_ActionSpecName = "m_ActionSpec";
+
+        public override void OnInspectorGUI()
+        {
+            var so = serializedObject;
+            so.Update();
+            InputActuatorComponent o = so.targetObject as InputActuatorComponent;
+            _ = o.ActionSpec;
+            EditorGUI.indentLevel++;
+            EditorGUI.BeginDisabledGroup(true);
+            EditorGUILayout.PropertyField(so.FindProperty(k_ActionSpecName));
+            EditorGUI.EndDisabledGroup();
+            EditorGUI.indentLevel--;
+        }
+    }
+}
+#endif // MLA_INPUT_SYSTEM && UNITY_2019_OR_NEWER
--- a/com.unity.ml-agents.extensions/Editor/Input/InputActuatorComponentEditor.cs.meta
+++ b/com.unity.ml-agents.extensions/Editor/Input/InputActuatorComponentEditor.cs.meta
+fileFormatVersion: 2
+guid: 3435eeef4d1645be8c0f770b68f2ba19
+timeCreated: 1613637202
--- a/com.unity.ml-agents.extensions/Editor/Input/Unity.ML-Agents.Extensions.Editor.Input.asmdef
+++ b/com.unity.ml-agents.extensions/Editor/Input/Unity.ML-Agents.Extensions.Editor.Input.asmdef
+{
+    "name": "Unity.ML-Agents.Extensions.Editor.Input",
+    "references": [
+        "Unity.ML-Agents",
+        "Unity.ML-Agents.Extensions.Input",
+        "Unity.ML-Agents.Editor",
+        "Unity.InputSystem"
+    ],
+    "includePlatforms": [
+        "Editor"
+    ],
+    "excludePlatforms": [],
+    "allowUnsafeCode": false,
+    "overrideReferences": false,
+    "precompiledReferences": [],
+    "autoReferenced": true,
+    "defineConstraints": [],
+    "versionDefines": [
+        {
+            "name": "com.unity.inputsystem",
+            "expression": "1.1.0-preview",
+            "define": "MLA_INPUT_SYSTEM"
+        }
+    ],
+    "noEngineReferences": false
+}
--- a/com.unity.ml-agents.extensions/Editor/Input/Unity.ML-Agents.Extensions.Editor.Input.asmdef.meta
+++ b/com.unity.ml-agents.extensions/Editor/Input/Unity.ML-Agents.Extensions.Editor.Input.asmdef.meta
+fileFormatVersion: 2
+guid: 4851f2d02f9f1423a8593f60b1a9cd7e
+AssemblyDefinitionImporter:
+  externalObjects: {}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 
--- a/com.unity.ml-agents.extensions/Runtime/Input/Adaptors.meta
+++ b/com.unity.ml-agents.extensions/Runtime/Input/Adaptors.meta
+fileFormatVersion: 2
+guid: de3fc3f4fd664e3ab579a102f7fabc88
+timeCreated: 1612204931
--- a/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/ButtonInputActionAdaptor.cs
+++ b/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/ButtonInputActionAdaptor.cs
+#if MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
+using Unity.MLAgents.Actuators;
+using UnityEngine;
+using UnityEngine.InputSystem;
+using UnityEngine.InputSystem.LowLevel;
+
+namespace Unity.MLAgents.Extensions.Input
+{
+    /// <summary>
+    /// Class that translates data between the a <see cref="UnityEngine.InputSystem.Controls.ButtonControl"/> and
+    /// the ML-Agents <see cref="ActionBuffers"/> object.
+    /// </summary>
+    public class ButtonInputActionAdaptor : IRLActionInputAdaptor
+    {
+        /// <summary>
+        /// TODO this method needs to be more nuanced depending the types of controls that can back it.  i.e. TriggerControls
+        /// are continuous buttons, etc.
+        /// Currently returns an <see cref="ActionSpec"/> with 1 branch of size 2.  One value for not pressed, and one
+        /// for pressed.
+        /// </summary>
+        /// <param name="action">The action associated with this adaptor to help determine the action space.</param>
+        /// <returns></returns>
+        public ActionSpec GetActionSpecForInputAction(InputAction action)
+        {
+            return ActionSpec.MakeDiscrete(2);
+        }
+
+        /// TODO again this might need to be more nuanced for things like continuous buttons.
+        /// <inheritdoc cref="IRLActionInputAdaptor.QueueInputEventForAction"/>
+        public void QueueInputEventForAction(InputAction action, InputControl control, ActionSpec actionSpec, in ActionBuffers actionBuffers)
+        {
+            var val = actionBuffers.DiscreteActions[0];
+            InputSystem.QueueDeltaStateEvent(control, (byte)val);
+        }
+
+        /// <inheritdoc cref="IRLActionInputAdaptor.WriteToHeuristic"/>>
+        public void WriteToHeuristic(InputAction action, in ActionBuffers actionBuffers)
+        {
+            var discreteActions = actionBuffers.DiscreteActions;
+            var val = action.ReadValue<float>();
+            discreteActions[0] = (int)val;
+        }
+    }
+}
+#endif // MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/ButtonInputActionAdaptor.cs.meta
+++ b/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/ButtonInputActionAdaptor.cs.meta
+fileFormatVersion: 2
+guid: f3a0fe3f0bd446958c729d6f71e8d00b
+timeCreated: 1612373241
--- a/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/DoubleInputActionAdaptor.cs
+++ b/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/DoubleInputActionAdaptor.cs
+#if MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
+using Unity.MLAgents.Actuators;
+using UnityEngine.InputSystem;
+using UnityEngine.InputSystem.LowLevel;
+
+namespace Unity.MLAgents.Extensions.Input
+{
+    /// <summary>
+    /// Translates data from a <see cref="UnityEngine.InputSystem.Controls.DoubleControl"/>.
+    /// </summary>
+    public class DoubleInputActionAdaptor : IRLActionInputAdaptor
+    {
+        /// <inheritdoc cref="IRLActionInputAdaptor.GetActionSpecForInputAction"/>
+        public ActionSpec GetActionSpecForInputAction(InputAction action)
+        {
+            return ActionSpec.MakeContinuous(1);
+        }
+
+        /// <inheritdoc cref="IRLActionInputAdaptor.QueueInputEventForAction"/>
+        public void QueueInputEventForAction(InputAction action, InputControl control, ActionSpec actionSpec, in ActionBuffers actionBuffers)
+        {
+            var val = actionBuffers.ContinuousActions[0];
+            InputSystem.QueueDeltaStateEvent(control,(double)val);
+        }
+
+        /// <inheritdoc cref="IRLActionInputAdaptor.WriteToHeuristic"/>
+        public void WriteToHeuristic(InputAction action, in ActionBuffers actionBuffers)
+        {
+            var actions = actionBuffers.ContinuousActions;
+            var val = (float)action.ReadValue<double>();
+            actions[0] = val;
+        }
+    }
+}
+#endif // MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/DoubleInputActionAdaptor.cs.meta
+++ b/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/DoubleInputActionAdaptor.cs.meta
+fileFormatVersion: 2
+guid: 14ec2823be0a4b3bbee5490ed4840e9c
+timeCreated: 1612574422
--- a/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/FloatInputActionAdaptor.cs
+++ b/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/FloatInputActionAdaptor.cs
+#if MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
+using Unity.MLAgents.Actuators;
+using UnityEngine.InputSystem;
+using UnityEngine.InputSystem.LowLevel;
+
+namespace Unity.MLAgents.Extensions.Input
+{
+    /// <summary>
+    /// Translates data from any control that extends from <see cref="InputControl{Single}"/>.
+    /// </summary>
+    public class FloatInputActionAdaptor : IRLActionInputAdaptor
+    {
+        /// <inheritdoc cref="IRLActionInputAdaptor.GetActionSpecForInputAction"/>
+        public ActionSpec GetActionSpecForInputAction(InputAction action)
+        {
+            return ActionSpec.MakeContinuous(1);
+        }
+
+        /// <inheritdoc cref="IRLActionInputAdaptor.QueueInputEventForAction"/>
+        public void QueueInputEventForAction(InputAction action, InputControl control, ActionSpec actionSpec, in ActionBuffers actionBuffers)
+        {
+            var val = actionBuffers.ContinuousActions[0];
+            InputSystem.QueueDeltaStateEvent(control, val);
+        }
+
+        /// <inheritdoc cref="IRLActionInputAdaptor.WriteToHeuristic"/>
+        public void WriteToHeuristic(InputAction action, in ActionBuffers actionBuffers)
+        {
+            var actions = actionBuffers.ContinuousActions;
+            var val = action.ReadValue<float>();
+            actions[0] = val;
+        }
+    }
+}
+#endif // MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/FloatInputActionAdaptor.cs.meta
+++ b/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/FloatInputActionAdaptor.cs.meta
+fileFormatVersion: 2
+guid: 6c98cc3fdaec4664aae128a05cfe6560
+timeCreated: 1612573580
--- a/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/IntegerInputActionAdaptor.cs
+++ b/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/IntegerInputActionAdaptor.cs
+#if MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
+using Unity.MLAgents.Actuators;
+using UnityEngine.InputSystem;
+using UnityEngine.InputSystem.LowLevel;
+
+namespace Unity.MLAgents.Extensions.Input
+{
+    /// <summary>
+    /// Translates data from a <see cref="UnityEngine.InputSystem.Controls.IntegerControl"/>.
+    /// </summary>
+    public class IntegerInputActionAdaptor : IRLActionInputAdaptor
+    {
+        // TODO need to figure out how we can infer the branch size from here.
+        /// <inheritdoc cref="IRLActionInputAdaptor.GetActionSpecForInputAction"/>
+        public ActionSpec GetActionSpecForInputAction(InputAction action)
+        {
+            return ActionSpec.MakeDiscrete(2);
+        }
+
+        /// <inheritdoc cref="IRLActionInputAdaptor.QueueInputEventForAction"/>
+        public void QueueInputEventForAction(InputAction action, InputControl control, ActionSpec actionSpec, in ActionBuffers actionBuffers)
+        {
+            var val = actionBuffers.DiscreteActions[0];
+            InputSystem.QueueDeltaStateEvent(control, val);
+        }
+
+        /// <inheritdoc cref="IRLActionInputAdaptor.WriteToHeuristic"/>
+        public void WriteToHeuristic(InputAction action, in ActionBuffers actionBuffers)
+        {
+            var actions = actionBuffers.DiscreteActions;
+            var val = action.ReadValue<int>();
+            actions[0] = val;
+        }
+    }
+}
+#endif // MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/IntegerInputActionAdaptor.cs.meta
+++ b/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/IntegerInputActionAdaptor.cs.meta
+fileFormatVersion: 2
+guid: 592012f9a30847a29c618f0cf8addfdf
+timeCreated: 1612572952
--- a/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/Vector2InputActionAdaptor.cs
+++ b/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/Vector2InputActionAdaptor.cs
+#if MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
+using System;
+using Unity.MLAgents.Actuators;
+using UnityEngine;
+using UnityEngine.InputSystem;
+using UnityEngine.InputSystem.LowLevel;
+
+namespace Unity.MLAgents.Extensions.Input
+{
+    /// <summary>
+    /// Translates data from any control that extends from <see cref="InputControl{Vector2}"/>.
+    /// </summary>
+    public class Vector2InputActionAdaptor : IRLActionInputAdaptor
+    {
+        /// <inheritdoc cref="IRLActionInputAdaptor.GetActionSpecForInputAction"/>
+        public ActionSpec GetActionSpecForInputAction(InputAction action)
+        {
+            // TODO create the action spec based on what controls back the action
+            return ActionSpec.MakeContinuous(2);
+        }
+
+        /// <inheritdoc cref="IRLActionInputAdaptor.QueueInputEventForAction"/>
+        public void QueueInputEventForAction(InputAction action,
+            InputControl control,
+            ActionSpec actionSpec,
+            in ActionBuffers actionBuffers)
+        {
+            var x = actionBuffers.ContinuousActions[0];
+            var y = actionBuffers.ContinuousActions[1];
+            InputSystem.QueueDeltaStateEvent(control, new Vector2(x, y));
+        }
+
+        /// <inheritdoc cref="IRLActionInputAdaptor.WriteToHeuristic"/>
+        public void WriteToHeuristic(InputAction action, in ActionBuffers actionBuffers)
+        {
+            var value = action.ReadValue<Vector2>();
+            var continuousActions = actionBuffers.ContinuousActions;
+            continuousActions[0] = value.x;
+            continuousActions[1] = value.y;
+        }
+
+    }
+}
+#endif // MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/Vector2InputActionAdaptor.cs.meta
+++ b/com.unity.ml-agents.extensions/Runtime/Input/Adaptors/Vector2InputActionAdaptor.cs.meta
+fileFormatVersion: 2
+guid: aa311fde3dac44f3b7998bb7fee77225
+timeCreated: 1611356491
--- a/com.unity.ml-agents.extensions/Runtime/Input/AssemblyInfo.cs
+++ b/com.unity.ml-agents.extensions/Runtime/Input/AssemblyInfo.cs
+using System.Runtime.CompilerServices;
+
+[assembly: InternalsVisibleTo("Unity.ML-Agents.Extensions.Input.Tests.Runtime")]
--- a/com.unity.ml-agents.extensions/Runtime/Input/AssemblyInfo.cs.meta
+++ b/com.unity.ml-agents.extensions/Runtime/Input/AssemblyInfo.cs.meta
+fileFormatVersion: 2
+guid: 989e62db4b694586bf2c832ce13e2d50
+timeCreated: 1612916938
--- a/com.unity.ml-agents.extensions/Runtime/Input/IInputActionAssetProvider.cs
+++ b/com.unity.ml-agents.extensions/Runtime/Input/IInputActionAssetProvider.cs
+#if MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
+using UnityEngine.InputSystem;
+
+namespace Unity.MLAgents.Extensions.Input
+{
+    /// <summary>
+    /// Implement this interface if you are listening to C# events from the generated C# class from the
+    /// <see cref="InputActionAsset"/>. This interface works with the <see cref="InputActuatorComponent"/> in order
+    /// to allow ML-Agents to simulate input actions based on the instance of the <see cref="InputActionAsset"/>
+    /// used to listen to events.  If you implement this interface the <see cref="InputActuatorComponent"/> will use
+    /// what is returned from <see cref="GetInputActionAsset"/> as the asset to base it's simulated input for.
+    /// Otherwise, the <see cref="InputActuatorComponent"/> will look for the <see cref="PlayerInput"/> component
+    /// and use the asset from there.  If you have multiple components handling PlayerInput on the same GameObject
+    /// they will need to share the same instance of the <see cref="InputActionAsset"/> in order to get the simulated
+    /// input.
+    /// </summary>
+    public interface IInputActionAssetProvider
+    {
+        /// <summary>
+        /// Returns the <see cref="InputActionAsset"/> instance being from the generated C# class of the
+        /// <see cref="InputActionAsset"/> in order to correctly fire events when simulating input from ML-Agents.
+        /// </summary>
+        /// <returns>The instance of the <see cref="InputActionAsset"/> you are listening for events on.</returns>
+        (InputActionAsset, IInputActionCollection2) GetInputActionAsset();
+    }
+}
+#endif // MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Runtime/Input/IInputActionAssetProvider.cs.meta
+++ b/com.unity.ml-agents.extensions/Runtime/Input/IInputActionAssetProvider.cs.meta
+fileFormatVersion: 2
+guid: 3ce1e3218e46e48a4ab64f6f631553a4
+MonoImporter:
+  externalObjects: {}
+  serializedVersion: 2
+  defaultReferences: []
+  executionOrder: 0
+  icon: {instanceID: 0}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 
--- a/com.unity.ml-agents.extensions/Runtime/Input/IRLActionInputAdaptor.cs
+++ b/com.unity.ml-agents.extensions/Runtime/Input/IRLActionInputAdaptor.cs
+#if MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
+using System;
+using Unity.MLAgents.Actuators;
+using UnityEngine.InputSystem;
+
+namespace Unity.MLAgents.Extensions.Input
+{
+    /// <summary>
+    /// Implement this interface in order to customize how information is translated <see cref="InputControl"/>s
+    /// and <see cref="ActionBuffers"/>.
+    /// </summary>
+    public interface IRLActionInputAdaptor
+    {
+        /// <summary>
+        /// Generate an <see cref="ActionSpec"/> for a given action which determines how data is translated between
+        /// the <see cref="InputSystem"/> and ML-Agents.
+        /// </summary>
+        /// <param name="action">The <see cref="InputAction"/> to based the <see cref="ActionSpec"/> from.</param>
+        /// <returns>An <see cref="ActionSpec"/> instance based off the information in the <see cref="InputAction"/>.</returns>
+        ActionSpec GetActionSpecForInputAction(InputAction action);
+
+        /// <summary>
+        /// Translates data from the <see cref="ActionBuffers"/> object to the <see cref="InputSystem"/>.
+        /// </summary>
+        /// <param name="action">The action associated with this adaptor.</param>
+        /// <param name="control">The control which will write the event to the <see cref="InputSystem"/>.</param>
+        /// <param name="actionSpec">The <see cref="ActionSpec"/> associated with this action and adaptor pair.</param>
+        /// <param name="actionBuffers">The <see cref="ActionBuffers"/> object to read from.</param>
+        void QueueInputEventForAction(InputAction action, InputControl control, ActionSpec actionSpec, in ActionBuffers actionBuffers);
+
+        /// <summary>
+        /// Writes data from the <paramref name="action"/> to the <paramref name="actionBuffers"/>.
+        /// </summary>
+        /// <param name="action">The <paramref name="action"/> to read data from.</param>
+        /// <param name="actionBuffers">The <paramref name="actionBuffers"/> object to write data to.</param>
+        void WriteToHeuristic(InputAction action, in ActionBuffers actionBuffers);
+    }
+}
+#endif // MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Runtime/Input/IRLActionInputAdaptor.cs.meta
+++ b/com.unity.ml-agents.extensions/Runtime/Input/IRLActionInputAdaptor.cs.meta
+fileFormatVersion: 2
+guid: 010af7dd365b48849a9d498b84a4be94
+timeCreated: 1611350962
--- a/com.unity.ml-agents.extensions/Runtime/Input/InputActionActuator.cs
+++ b/com.unity.ml-agents.extensions/Runtime/Input/InputActionActuator.cs
+#if MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
+
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Policies;
+using UnityEngine;
+using UnityEngine.InputSystem;
+using UnityEngine.Profiling;
+
+namespace Unity.MLAgents.Extensions.Input
+{
+    /// <summary>
+    /// This implementation of <see cref="IActuator"/> will send events from the ML-Agents training process, or from
+    /// neural networks to the <see cref="InputSystem"/> via the <see cref="IRLActionInputAdaptor"/> interface.  If an
+    /// <see cref="Agent"/>'s <see cref="BehaviorParameters"/> indicate that the Agent is running in Heuristic Mode,
+    /// this Actuator will write actions from the <see cref="InputSystem"/> to the <see cref="ActionBuffers"/> object.
+    /// </summary>
+    public class InputActionActuator : IActuator, IHeuristicProvider
+    {
+        readonly BehaviorParameters m_BehaviorParameters;
+        readonly InputAction m_Action;
+        readonly IRLActionInputAdaptor m_InputAdaptor;
+        InputDevice m_Device;
+        InputControl m_Control;
+
+        /// <summary>
+        /// Construct an <see cref="InputActionActuator"/> with the <see cref="BehaviorParameters"/> of the
+        /// <see cref="Agent"/> component, the relevant <see cref="InputAction"/>, and the relevant
+        /// <see cref="IRLActionInputAdaptor"/> to convert between ml-agents &lt;--&gt; <see cref="InputSystem"/>.
+        /// </summary>
+        /// <param name="inputDevice">The input device this action is bound to.</param>
+        /// <param name="behaviorParameters">Used to determine if the <see cref="Agent"/> is running in
+        ///     heuristic mode.</param>
+        /// <param name="action">The <see cref="InputAction"/> this <see cref="IActuator"/> we read/write data to/from
+        ///     via the <see cref="IRLActionInputAdaptor"/>.</param>
+        /// <param name="adaptor">The <see cref="IRLActionInputAdaptor"/> that will convert data between ML-Agents
+        ///     and the <see cref="InputSystem"/>.</param>
+        public InputActionActuator(InputDevice inputDevice, BehaviorParameters behaviorParameters,
+            InputAction action,
+            IRLActionInputAdaptor adaptor)
+        {
+            m_BehaviorParameters = behaviorParameters;
+            Name = $"InputActionActuator-{action.name}";
+            m_Action = action;
+            m_InputAdaptor = adaptor;
+            ActionSpec = adaptor.GetActionSpecForInputAction(m_Action);
+            m_Device = inputDevice;
+            m_Control = m_Device?.GetChildControl(m_Action.name);
+        }
+
+        /// <inheritdoc cref="IActionReceiver.OnActionReceived"/>
+        public void OnActionReceived(ActionBuffers actionBuffers)
+        {
+            Profiler.BeginSample("InputActionActuator.OnActionReceived");
+            if (!m_BehaviorParameters.IsInHeuristicMode())
+            {
+                m_InputAdaptor.QueueInputEventForAction(m_Action, m_Control, ActionSpec, actionBuffers);
+            }
+            Profiler.EndSample();
+        }
+
+        /// <inheritdoc cref="IActionReceiver.WriteDiscreteActionMask"/>
+        public void WriteDiscreteActionMask(IDiscreteActionMask actionMask)
+        {
+            // TODO configure mask from editor UI?
+        }
+
+        /// <inheritdoc cref="IActuator.ActionSpec"/>
+        public ActionSpec ActionSpec { get; }
+
+        /// <inheritdoc cref="IActuator.Name"/>
+        public string Name { get; }
+
+        /// <inheritdoc cref="IActuator.ResetData"/>
+        public void ResetData()
+        {
+            // do nothing for now
+        }
+
+        /// <inheritdoc cref="IHeuristicProvider.Heuristic"/>
+        public void Heuristic(in ActionBuffers actionBuffersOut)
+        {
+            Profiler.BeginSample("InputActionActuator.Heuristic");
+            m_InputAdaptor.WriteToHeuristic(m_Action, actionBuffersOut);
+            Profiler.EndSample();
+        }
+    }
+}
+
+#endif // MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Runtime/Input/InputActionActuator.cs.meta
+++ b/com.unity.ml-agents.extensions/Runtime/Input/InputActionActuator.cs.meta
+fileFormatVersion: 2
+guid: d4dfb5125cc2461088f18e27b6c21842
+timeCreated: 1612305510
--- a/com.unity.ml-agents.extensions/Runtime/Input/InputActuatorComponent.cs
+++ b/com.unity.ml-agents.extensions/Runtime/Input/InputActuatorComponent.cs
+#if MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
+using System;
+using System.Collections.Generic;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Policies;
+using UnityEngine;
+using UnityEngine.Assertions;
+using UnityEngine.InputSystem;
+using UnityEngine.InputSystem.Controls;
+using UnityEngine.InputSystem.Layouts;
+using UnityEngine.InputSystem.Utilities;
+#if UNITY_EDITOR
+using UnityEditor;
+#endif
+
+namespace Unity.MLAgents.Extensions.Input
+{
+    /// <summary>
+    /// Component class that handles the parsing of the <see cref="InputActionAsset"/> and translates that into
+    /// <see cref="InputActionActuator"/>s.
+    /// </summary>
+    [RequireComponent(typeof(PlayerInput), typeof(IInputActionAssetProvider))]
+    public class InputActuatorComponent : ActuatorComponent
+    {
+        InputActionAsset m_InputAsset;
+        IInputActionCollection2 m_AssetCollection;
+        PlayerInput m_PlayerInput;
+        BehaviorParameters m_BehaviorParameters;
+        IActuator[] m_Actuators;
+        InputDevice m_Device;
+
+        /// <summary>
+        /// Mapping of <see cref="InputControl"/> types to types of <see cref="IRLActionInputAdaptor"/> concrete classes.
+        /// </summary>
+        public static Dictionary<Type, Type> controlTypeToAdaptorType = new Dictionary<Type, Type>
+        {
+            { typeof(Vector2Control), typeof(Vector2InputActionAdaptor) },
+            { typeof(ButtonControl), typeof(ButtonInputActionAdaptor) },
+            { typeof(IntegerControl), typeof(IntegerInputActionAdaptor) },
+            { typeof(AxisControl), typeof(FloatInputActionAdaptor) },
+            { typeof(DoubleControl), typeof(DoubleInputActionAdaptor) }
+        };
+
+        string m_LayoutName;
+        [SerializeField]
+        ActionSpec m_ActionSpec;
+        InputControlScheme m_ControlScheme;
+
+        public const string mlAgentsLayoutFormat = "MLAT";
+        public const string mlAgentsLayoutName = "MLAgentsLayout";
+        public const string mlAgentsControlSchemeName = "ml-agents";
+
+        /// <inheritdoc cref="IActuator.ActionSpec"/>
+        public override ActionSpec ActionSpec
+        {
+            get
+            {
+#if UNITY_EDITOR
+                FindNeededComponents();
+                var actuators = CreateActuatorsFromMap(m_InputAsset.FindActionMap(m_PlayerInput.defaultActionMap), m_BehaviorParameters, null);
+                m_ActionSpec = CombineActuatorActionSpecs(actuators);
+#endif
+                return m_ActionSpec;
+            }
+        }
+
+        void OnDisable()
+        {
+            CleanupActionAsset();
+        }
+
+        /// <summary>
+        /// This method is where the <see cref="InputActionAsset"/> gets parsed and translated into
+        /// <see cref="InputActionActuator"/>s that communicate with the <see cref="InputSystem"/> via a
+        /// virtual <see cref="InputDevice"/>.
+        /// <remarks>
+        /// The flow of this method is as follows:
+        /// <list type="number">
+        ///     <item>
+        ///     <description>Ensure that our custom <see cref="InputBindingComposite{TValue}"/>s are registered with
+        ///     the InputSystem.</description>
+        ///     </item>
+        ///     <item>
+        ///     <description>Look for the components that are needed by this class in order to retrieve the
+        ///     <see cref="InputActionAsset"/>.  It first looks for <see cref="IInputActionAssetProvider"/>, if that
+        ///     is not found, it will get the asset from the <see cref="PlayerInput"/> component.</description>
+        ///     </item>
+        ///     <item>
+        ///     <description>Create the list <see cref="InputActionActuator"/>s, one for each action in the default
+        ///     <see cref="InputActionMap"/> as set by the <see cref="PlayerInput"/> component.  Within the method
+        ///     where the actuators are being created, an <see cref="InputControlLayout"/> is also being built based
+        ///     on the number and types of <see cref="InputAction"/>s.  This will be used to create a virtual
+        ///     <see cref="InputDevice"/> with a <see cref="InputControlLayout"/> that is specific to the
+        ///     <see cref="InputActionMap"/> specified by <see cref="PlayerInput"/></description>
+        ///     </item>
+        ///     <item>
+        ///     <description>Create our device based on the layout that was generated and registered during
+        ///     actuator creation.</description>
+        ///     </item>
+        ///     <item>
+        ///     <description>Create an ml-agents control scheme and add it to the <see cref="InputActionAsset"/> so
+        ///     our virtual devices can be used.</description>
+        ///     </item>
+        ///     <item>
+        ///     <description>Add our virtual <see cref="InputDevice"/> to the input system.</description>
+        ///     </item>
+        /// </list>
+        /// </remarks>
+        /// </summary>
+        /// <returns>A list of </returns>
+        public override IActuator[] CreateActuators()
+        {
+            FindNeededComponents();
+            var collection = m_AssetCollection ?? m_InputAsset;
+            collection.Disable();
+            var inputActionMap = m_InputAsset.FindActionMap(m_PlayerInput.defaultActionMap);
+
+            RegisterLayoutBuilder(inputActionMap, m_LayoutName);
+            m_Device = InputSystem.AddDevice(m_LayoutName);
+
+            m_Actuators = CreateActuatorsFromMap(inputActionMap, m_BehaviorParameters, m_Device);
+
+            UpdateDeviceBinding(m_BehaviorParameters.IsInHeuristicMode());
+            inputActionMap.Enable();
+
+            m_ActionSpec = CombineActuatorActionSpecs(m_Actuators);
+            collection.Enable();
+            return m_Actuators;
+        }
+
+        static ActionSpec CombineActuatorActionSpecs(IActuator[] actuators)
+        {
+            var specs = new ActionSpec[actuators.Length];
+            for (var i = 0; i < actuators.Length; i++)
+            {
+                specs[i] = actuators[i].ActionSpec;
+            }
+            return ActionSpec.Combine(specs);
+        }
+
+        internal static IActuator[] CreateActuatorsFromMap(InputActionMap inputActionMap,
+            BehaviorParameters behaviorParameters,
+            InputDevice inputDevice)
+        {
+            var actuators = new IActuator[inputActionMap.actions.Count];
+            for (var i = 0; i < inputActionMap.actions.Count; i++)
+            {
+                var action = inputActionMap.actions[i];
+                var actionLayout = InputSystem.LoadLayout(action.expectedControlType);
+                var adaptor = (IRLActionInputAdaptor)Activator.CreateInstance(controlTypeToAdaptorType[actionLayout.type]);
+                actuators[i] = new InputActionActuator(inputDevice, behaviorParameters, action, adaptor);
+
+                // Reasonably, the input system starts adding numbers after the first none numbered name
+                // is added.  So for device ID of 0, we use the empty string in the path.
+                var path = $"{inputDevice?.path}{InputControlPath.Separator}{action.name}";
+                action.AddBinding(path,
+                    action.interactions,
+                    action.processors,
+                    mlAgentsControlSchemeName);
+            }
+            return actuators;
+        }
+
+        /// <summary>
+        /// Set up bindings based on whether or not the BehaviorParameters are working in Heuristic mode or not.
+        /// If we are working in Heuristic mode, we want the input system to handle everything.  If not, we
+        /// want the neural network to send input from virtual devices.
+        /// </summary>
+        /// <param name="isInHeuristicMode">true if the Agent connected to this GameObject is working in
+        /// Heuristic mode.</param>
+        /// <seealso cref="BehaviorParameters.IsInHeuristicMode"/>
+        internal void UpdateDeviceBinding(bool isInHeuristicMode)
+        {
+            if (ReferenceEquals(m_Device, null))
+            {
+                return;
+            }
+            var collection = m_AssetCollection ?? m_InputAsset;
+            m_ControlScheme = CreateControlScheme(m_Device, isInHeuristicMode, m_InputAsset);
+            if (m_InputAsset.FindControlSchemeIndex(m_ControlScheme.name) != -1)
+            {
+                m_InputAsset.RemoveControlScheme(m_ControlScheme.name);
+            }
+
+            if (!isInHeuristicMode)
+            {
+                var inputActionMap = m_InputAsset.FindActionMap(m_PlayerInput.defaultActionMap);
+                m_InputAsset.AddControlScheme(m_ControlScheme);
+                collection.bindingMask = InputBinding.MaskByGroup(m_ControlScheme.bindingGroup);
+                collection.devices = new ReadOnlyArray<InputDevice>(new[] { m_Device });
+                inputActionMap.bindingMask = collection.bindingMask;
+                inputActionMap.devices = collection.devices;
+            }
+            else
+            {
+                var inputActionMap = m_InputAsset.FindActionMap(m_PlayerInput.defaultActionMap);
+                collection.bindingMask = null;
+                collection.devices = InputSystem.devices;
+                inputActionMap.devices = InputSystem.devices;
+                inputActionMap.bindingMask = null;
+            }
+            collection.Enable();
+        }
+
+        /// <summary>
+        /// This method creates a control scheme and adds it to the <see cref="InputActionAsset"/> passed in so
+        /// we can add our device to in order for it to be discovered by the <see cref="InputSystem"/>.
+        /// </summary>
+        /// <param name="device">The virtual device to add to our custom control scheme.</param>
+        /// <param name="isInHeuristicMode">if we are in heuristic mode, we need to add other other device requirements.</param>
+        /// <param name="asset">The InputActionAsset to get the device requirements from</param>
+        internal static InputControlScheme CreateControlScheme(InputControl device,
+            bool isInHeuristicMode,
+            InputActionAsset asset)
+        {
+            var deviceRequirements = new List<InputControlScheme.DeviceRequirement>
+            {
+                new InputControlScheme.DeviceRequirement
+                {
+                    controlPath = InputBinding.Separator + mlAgentsLayoutName
+                }
+            };
+
+            if (isInHeuristicMode)
+            {
+                for (var i = 0; i < asset.controlSchemes.Count; i++)
+                {
+                    var scheme = asset.controlSchemes[i];
+                    for (var ii = 0; ii < scheme.deviceRequirements.Count; ii++)
+                    {
+                        deviceRequirements.Add(scheme.deviceRequirements[ii]);
+                    }
+                }
+            }
+
+            var inputControlScheme = new InputControlScheme(
+                 mlAgentsControlSchemeName,
+                deviceRequirements);
+
+            return inputControlScheme;
+        }
+
+#pragma warning disable 672
+        /// <inheritdoc cref="ActuatorComponent.CreateActuator"/>
+        public override IActuator CreateActuator() { return null; }
+#pragma warning restore 672
+
+        /// <summary>
+        ///
+        /// </summary>
+        /// <param name="defaultMap"></param>
+        /// <param name="layoutName"></param>
+        /// <returns></returns>
+        internal static void RegisterLayoutBuilder(InputActionMap defaultMap, string layoutName)
+        {
+            if (InputSystem.LoadLayout(layoutName) == null)
+            {
+                InputSystem.RegisterLayoutBuilder(() =>
+                {
+                    // TODO does this need to change based on the action map we use?
+                    var builder = new InputControlLayout.Builder()
+                        .WithName(layoutName)
+                        .WithFormat(mlAgentsLayoutFormat);
+                    for(var i = 0; i < defaultMap.actions.Count; i++)
+                    {
+                        var action = defaultMap.actions[i];
+                        builder.AddControl(action.name)
+                            .WithLayout(action.expectedControlType);
+                    }
+                    return builder.Build();
+
+                }, layoutName);
+            }
+        }
+
+        internal void FindNeededComponents()
+        {
+            if (m_InputAsset == null)
+            {
+                var assetProvider = GetComponent<IInputActionAssetProvider>();
+                Assert.IsNotNull(assetProvider);
+                (m_InputAsset, m_AssetCollection) = assetProvider.GetInputActionAsset();
+                Assert.IsNotNull(m_InputAsset, "An InputActionAsset could not be found on IInputActionAssetProvider or PlayerInput.");
+            }
+            if (m_PlayerInput == null)
+            {
+                m_PlayerInput = GetComponent<PlayerInput>();
+                Assert.IsNotNull(m_PlayerInput, "PlayerInput component could not be found on this GameObject.");
+            }
+
+            if (m_BehaviorParameters == null)
+            {
+                m_BehaviorParameters = GetComponent<BehaviorParameters>();
+                Assert.IsNotNull(m_BehaviorParameters, "BehaviorParameters were not on the current GameObject.");
+                m_BehaviorParameters.OnPolicyUpdated += UpdateDeviceBinding;
+                m_LayoutName = mlAgentsLayoutName + m_BehaviorParameters.BehaviorName;
+            }
+        }
+
+        internal void CleanupActionAsset()
+        {
+            InputSystem.RemoveLayout(mlAgentsLayoutName);
+            if (!ReferenceEquals(m_Device, null))
+            {
+                InputSystem.RemoveDevice(m_Device);
+            }
+
+            if (!ReferenceEquals(m_InputAsset, null)
+                && m_InputAsset.FindControlSchemeIndex(mlAgentsControlSchemeName) != -1)
+            {
+                m_InputAsset.RemoveControlScheme(mlAgentsControlSchemeName);
+            }
+
+            if (m_Actuators != null)
+            {
+                Array.Clear(m_Actuators, 0, m_Actuators.Length);
+            }
+
+            if (!ReferenceEquals(m_BehaviorParameters, null))
+            {
+                m_BehaviorParameters.OnPolicyUpdated -= UpdateDeviceBinding;
+            }
+
+            m_InputAsset = null;
+            m_PlayerInput = null;
+            m_BehaviorParameters = null;
+            m_Device = null;
+        }
+    }
+}
+#endif // MLA_INPUT_SYSTEM && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Runtime/Input/InputActuatorComponent.cs.meta
+++ b/com.unity.ml-agents.extensions/Runtime/Input/InputActuatorComponent.cs.meta
+fileFormatVersion: 2
+guid: 33005b124d7f841a191249baf2bacb2a
+MonoImporter:
+  externalObjects: {}
+  serializedVersion: 2
+  defaultReferences: []
+  executionOrder: 0
+  icon: {instanceID: 0}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 
--- a/com.unity.ml-agents.extensions/Runtime/Input/Unity.ML-Agents.Extensions.Input.asmdef
+++ b/com.unity.ml-agents.extensions/Runtime/Input/Unity.ML-Agents.Extensions.Input.asmdef
+{
+    "name": "Unity.ML-Agents.Extensions.Input",
+    "references": [
+        "Unity.ML-Agents",
+        "Unity.Barracuda",
+        "Unity.InputSystem"
+    ],
+    "includePlatforms": [],
+    "excludePlatforms": [],
+    "allowUnsafeCode": false,
+    "overrideReferences": false,
+    "precompiledReferences": [],
+    "autoReferenced": true,
+    "defineConstraints": [],
+    "versionDefines": [
+        {
+            "name": "com.unity.inputsystem",
+            "expression": "1.1.0-preview",
+            "define": "MLA_INPUT_SYSTEM"
+        }
+    ],
+    "noEngineReferences": false
+}
--- a/com.unity.ml-agents.extensions/Runtime/Input/Unity.ML-Agents.Extensions.Input.asmdef.meta
+++ b/com.unity.ml-agents.extensions/Runtime/Input/Unity.ML-Agents.Extensions.Input.asmdef.meta
+fileFormatVersion: 2
+guid: a7b0d999fb2a7493a85c4c7017412530
+AssemblyDefinitionImporter:
+  externalObjects: {}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: 
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors.meta
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors.meta
+fileFormatVersion: 2
+guid: 9b59c16f3e30469ca26156b96351164a
+timeCreated: 1613081382
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/ButtonInputActionAdaptorTests.cs
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/ButtonInputActionAdaptorTests.cs
+#if MLA_INPUT_TESTS && UNITY_2019_4_OR_NEWER
+using NUnit.Framework;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Extensions.Input;
+using UnityEngine;
+using UnityEngine.InputSystem;
+
+namespace Unity.MLAgents.Extensions.Tests.Runtime.Input
+{
+    public class ButtonInputActionAdaptorTests : InputTestFixture
+    {
+        ButtonInputActionAdaptor m_Adaptor;
+        InputDevice m_Device;
+        InputControl<float> m_Control;
+        InputAction m_Action;
+
+        public override void Setup()
+        {
+            base.Setup();
+            const string kLayout = @"
+                {
+                    ""name"" : ""TestDevice"",
+                    ""extend"" : ""HID"",
+                    ""controls"" : [
+                        { ""name"" : ""button"", ""layout"" : ""Button"" }
+                    ]
+                }";
+            InputSystem.RegisterLayout(kLayout);
+            m_Device = InputSystem.AddDevice("TestDevice");
+            m_Control = (InputControl<float>)m_Device["button"];
+            m_Action = new InputAction("action", InputActionType.Button, "/TestDevice/button", null, null, "Button");
+            m_Action.Enable();
+            m_Adaptor = new ButtonInputActionAdaptor();
+        }
+
+        public override void TearDown()
+        {
+            base.TearDown();
+            m_Adaptor = null;
+        }
+
+        [Test]
+        public void TestGenerateActionSpec()
+        {
+            var actionSpec = m_Adaptor.GetActionSpecForInputAction(new InputAction());
+            Assert.IsTrue(actionSpec.NumDiscreteActions == 1);
+            Assert.IsTrue(actionSpec.BranchSizes[0] == 2);
+        }
+
+        [Test]
+        public void TestQueueEvent()
+        {
+            var actionBuffers = new ActionBuffers(ActionSegment<float>.Empty, new ActionSegment<int>(new[] { 1 }));
+            m_Adaptor.QueueInputEventForAction(m_Action, m_Control, new ActionSpec(), actionBuffers);
+            InputSystem.Update();
+            var val = m_Action.ReadValue<float>();
+            Assert.IsTrue(Mathf.Approximately(1f, val));
+        }
+
+        [Test]
+        public void TestWriteToHeuristic()
+        {
+            var actionBuffers = new ActionBuffers(ActionSegment<float>.Empty, new ActionSegment<int>(new[] { 1 }));
+            m_Adaptor.QueueInputEventForAction(m_Action, m_Control, new ActionSpec(), actionBuffers);
+            InputSystem.Update();
+            var buffer = new ActionBuffers(ActionSegment<float>.Empty, new ActionSegment<int>(new[] { 1 }));
+            m_Adaptor.WriteToHeuristic(m_Action, buffer);
+            Assert.IsTrue(buffer.DiscreteActions[0] == 1);
+        }
+    }
+}
+#endif // MLA_INPUT_TESTS && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/ButtonInputActionAdaptorTests.cs.meta
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/ButtonInputActionAdaptorTests.cs.meta
+fileFormatVersion: 2
+guid: f53bd7ea7b154c9a992b48c9fbc34e47
+timeCreated: 1613081399
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/DoubleInputActionAdaptorTests.cs
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/DoubleInputActionAdaptorTests.cs
+#if MLA_INPUT_TESTS && UNITY_2019_4_OR_NEWER
+using NUnit.Framework;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Extensions.Input;
+using UnityEngine;
+using UnityEngine.InputSystem;
+
+namespace Unity.MLAgents.Extensions.Tests.Runtime.Input
+{
+    public class DoubleInputActionAdaptorTests : InputTestFixture
+    {
+        DoubleInputActionAdaptor m_Adaptor;
+        InputDevice m_Device;
+        InputControl<double> m_Control;
+        InputAction m_Action;
+
+        public override void Setup()
+        {
+            base.Setup();
+            const string kLayout = @"
+                {
+                    ""name"" : ""TestDevice"",
+                    ""extend"" : ""HID"",
+                    ""controls"" : [
+                        { ""name"" : ""button"", ""layout"" : ""Double"" }
+                    ]
+                }";
+            InputSystem.RegisterLayout(kLayout);
+            m_Device = InputSystem.AddDevice("TestDevice");
+            m_Control = (InputControl<double>)m_Device["button"];
+            m_Action = new InputAction("action", InputActionType.Value, "/TestDevice/button", null, null, "double");
+            m_Action.Enable();
+            m_Adaptor = new DoubleInputActionAdaptor();
+        }
+
+        public override void TearDown()
+        {
+            base.TearDown();
+            m_Adaptor = null;
+        }
+
+        [Test]
+        public void TestGenerateActionSpec()
+        {
+            var actionSpec = m_Adaptor.GetActionSpecForInputAction(new InputAction());
+            Assert.IsTrue(actionSpec.NumContinuousActions == 1);
+        }
+
+        [Test]
+        public void TestQueueEvent()
+        {
+            var actionBuffers = new ActionBuffers(new ActionSegment<float>(new[] { 1f }), ActionSegment<int>.Empty);
+            m_Adaptor.QueueInputEventForAction(m_Action, m_Control, new ActionSpec(), actionBuffers);
+            InputSystem.Update();
+            Assert.IsTrue(Mathf.Approximately(1f, (float)m_Action.ReadValue<double>()));
+        }
+
+        [Test]
+        public void TestWriteToHeuristic()
+        {
+            var actionBuffers = new ActionBuffers(new ActionSegment<float>(new[] { 1f }), ActionSegment<int>.Empty);
+            m_Adaptor.QueueInputEventForAction(m_Action, m_Control, new ActionSpec(), actionBuffers);
+            InputSystem.Update();
+            var buffer = new ActionBuffers(new ActionSegment<float>(new[] { 1f }), ActionSegment<int>.Empty);
+            m_Adaptor.WriteToHeuristic(m_Action, buffer);
+            Assert.IsTrue(Mathf.Approximately(buffer.ContinuousActions[0], 1f));
+        }
+    }
+}
+#endif // MLA_INPUT_TESTS && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/DoubleInputActionAdaptorTests.cs.meta
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/DoubleInputActionAdaptorTests.cs.meta
+fileFormatVersion: 2
+guid: ae2012c098aa4b57ab09dbf72bcd4efe
+timeCreated: 1613074622
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/FloatInputActionAdapatorTests.cs
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/FloatInputActionAdapatorTests.cs
+#if MLA_INPUT_TESTS && UNITY_2019_4_OR_NEWER
+using NUnit.Framework;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Extensions.Input;
+using UnityEngine;
+using UnityEngine.InputSystem;
+
+namespace Unity.MLAgents.Extensions.Tests.Runtime.Input
+{
+    public class FloatInputActionAdaptorTests : InputTestFixture
+    {
+        FloatInputActionAdaptor m_Adaptor;
+        InputDevice m_Device;
+        InputControl<float> m_Control;
+        InputAction m_Action;
+
+        public override void Setup()
+        {
+            base.Setup();
+            const string kLayout = @"
+                {
+                    ""name"" : ""TestDevice"",
+                    ""extend"" : ""HID"",
+                    ""controls"" : [
+                        { ""name"" : ""button"", ""layout"" : ""Axis"" }
+                    ]
+                }";
+            InputSystem.RegisterLayout(kLayout);
+            m_Device = InputSystem.AddDevice("TestDevice");
+            m_Control = (InputControl<float>)m_Device["button"];
+            m_Action = new InputAction("action", InputActionType.Value, "/TestDevice/button", null, null, "Axis");
+            m_Action.Enable();
+            m_Adaptor = new FloatInputActionAdaptor();
+        }
+
+        public override void TearDown()
+        {
+            base.TearDown();
+            m_Adaptor = null;
+        }
+
+        [Test]
+        public void TestGenerateActionSpec()
+        {
+            var actionSpec = m_Adaptor.GetActionSpecForInputAction(new InputAction());
+            Assert.IsTrue(actionSpec.NumContinuousActions == 1);
+        }
+
+        [Test]
+        public void TestQueueEvent()
+        {
+            var actionBuffers = new ActionBuffers(new ActionSegment<float>(new[] { 1f }), ActionSegment<int>.Empty);
+            m_Adaptor.QueueInputEventForAction(m_Action, m_Control, new ActionSpec(), actionBuffers);
+            InputSystem.Update();
+            var val = m_Action.ReadValue<float>();
+            Assert.IsTrue(Mathf.Approximately(1f, val));
+        }
+
+        [Test]
+        public void TestWriteToHeuristic()
+        {
+            var actionBuffers = new ActionBuffers(new ActionSegment<float>(new[] { 1f }), ActionSegment<int>.Empty);
+            m_Adaptor.QueueInputEventForAction(m_Action, m_Control, new ActionSpec(), actionBuffers);
+            InputSystem.Update();
+            var buffer = new ActionBuffers(new ActionSegment<float>(new[] { 1f }), ActionSegment<int>.Empty);
+            m_Adaptor.WriteToHeuristic(m_Action, buffer);
+            Assert.IsTrue(Mathf.Approximately(1f, buffer.ContinuousActions[0]));
+        }
+    }
+}
+#endif // MLA_INPUT_TESTS && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/FloatInputActionAdapatorTests.cs.meta
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/FloatInputActionAdapatorTests.cs.meta
+fileFormatVersion: 2
+guid: a85fbb8b3e154eccadbdab826c3c5cae
+timeCreated: 1613083958
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/IntegerInputActionAdaptorTests.cs
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/IntegerInputActionAdaptorTests.cs
+#if MLA_INPUT_TESTS && UNITY_2019_4_OR_NEWER
+using System;
+using NUnit.Framework;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Extensions.Input;
+using UnityEngine;
+using UnityEngine.InputSystem;
+
+namespace Unity.MLAgents.Extensions.Tests.Runtime.Input
+{
+    public class IntegerInputActionAdaptorTests : InputTestFixture
+    {
+        IntegerInputActionAdaptor m_Adaptor;
+        InputDevice m_Device;
+        InputControl<int> m_Control;
+        InputAction m_Action;
+
+        public override void Setup()
+        {
+            base.Setup();
+            const string kLayout = @"
+                {
+                    ""name"" : ""TestDevice"",
+                    ""extend"" : ""HID"",
+                    ""controls"" : [
+                        { ""name"" : ""button"", ""layout"" : ""integer"" }
+                    ]
+                }";
+            InputSystem.RegisterLayout(kLayout);
+            m_Device = InputSystem.AddDevice("TestDevice");
+            m_Control = (InputControl<int>)m_Device["button"];
+            m_Action = new InputAction("action", InputActionType.Value, "/TestDevice/button", null, null, "int");
+            m_Action.Enable();
+            m_Adaptor = new IntegerInputActionAdaptor();
+        }
+
+        public override void TearDown()
+        {
+            base.TearDown();
+            m_Adaptor = null;
+        }
+
+        [Test]
+        public void TestGenerateActionSpec()
+        {
+            var actionSpec = m_Adaptor.GetActionSpecForInputAction(new InputAction());
+            Assert.IsTrue(actionSpec.NumDiscreteActions == 1);
+            Assert.IsTrue(actionSpec.SumOfDiscreteBranchSizes == 2);
+        }
+
+        [Test]
+        public void TestQueueEvent()
+        {
+            var actionBuffers = new ActionBuffers(ActionSegment<float>.Empty, new ActionSegment<int>(new[] { 1 }));
+            m_Adaptor.QueueInputEventForAction(m_Action, m_Control, new ActionSpec(), actionBuffers);
+            InputSystem.Update();
+            var val = m_Action.ReadValue<int>();
+            Assert.IsTrue(val == 1);
+        }
+
+        [Test]
+        public void TestWriteToHeuristic()
+        {
+            var actionBuffers = new ActionBuffers(ActionSegment<float>.Empty, new ActionSegment<int>(new[] { 1 }));
+            m_Adaptor.QueueInputEventForAction(m_Action, m_Control, new ActionSpec(), actionBuffers);
+            InputSystem.Update();
+            var buffer = new ActionBuffers(ActionSegment<float>.Empty, new ActionSegment<int>(new int[1]));
+            m_Adaptor.WriteToHeuristic(m_Action, buffer);
+            Assert.IsTrue(buffer.DiscreteActions[0] == 1);
+        }
+    }
+}
+#endif // MLA_INPUT_TESTS && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/IntegerInputActionAdaptorTests.cs.meta
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/IntegerInputActionAdaptorTests.cs.meta
+fileFormatVersion: 2
+guid: a8631ede70c049a7836284217ddfdc94
+timeCreated: 1613082414
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/Vector2InputActionAdaptorTests.cs
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/Vector2InputActionAdaptorTests.cs
+#if MLA_INPUT_TESTS && UNITY_2019_4_OR_NEWER
+using NUnit.Framework;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Extensions.Input;
+using UnityEngine;
+using UnityEngine.InputSystem;
+
+namespace Unity.MLAgents.Extensions.Tests.Runtime.Input
+{
+    public class Vector2InputActionAdaptorTests : InputTestFixture
+    {
+        Vector2InputActionAdaptor m_Adaptor;
+        InputDevice m_Device;
+        InputControl<Vector2> m_Control;
+        InputAction m_Action;
+
+        public override void Setup()
+        {
+            base.Setup();
+            const string kLayout = @"
+                {
+                    ""name"" : ""TestDevice"",
+                    ""extend"" : ""HID"",
+                    ""controls"" : [
+                        { ""name"" : ""button"", ""layout"" : ""Vector2"" }
+                    ]
+                }";
+            InputSystem.RegisterLayout(kLayout);
+            m_Device = InputSystem.AddDevice("TestDevice");
+            m_Control = (InputControl<Vector2>)m_Device["button"];
+            m_Action = new InputAction("action", InputActionType.Value, "/TestDevice/button", null, null, "Vector2");
+            m_Action.Enable();
+            m_Adaptor = new Vector2InputActionAdaptor();
+        }
+
+        public override void TearDown()
+        {
+            base.TearDown();
+            m_Adaptor = null;
+        }
+
+        [Test]
+        public void TestGenerateActionSpec()
+        {
+            var actionSpec = m_Adaptor.GetActionSpecForInputAction(new InputAction());
+            Assert.IsTrue(actionSpec.NumContinuousActions == 2);
+        }
+
+        [Test]
+        public void TestQueueEvent()
+        {
+            var actionBuffers = new ActionBuffers(new ActionSegment<float>(new[] { 0f, 1f }), ActionSegment<int>.Empty);
+            m_Adaptor.QueueInputEventForAction(m_Action, m_Control, new ActionSpec(), actionBuffers);
+            InputSystem.Update();
+            var val = m_Action.ReadValue<Vector2>();
+            Assert.IsTrue(Mathf.Approximately(0f, val.x));
+            Assert.IsTrue(Mathf.Approximately(1f, val.y));
+        }
+
+        [Test]
+        public void TestWriteToHeuristic()
+        {
+            var actionBuffers = new ActionBuffers(new ActionSegment<float>(new[] { 0f, 1f }), ActionSegment<int>.Empty);
+            m_Adaptor.QueueInputEventForAction(m_Action, m_Control, new ActionSpec(), actionBuffers);
+            InputSystem.Update();
+            var buffer = new ActionBuffers(new ActionSegment<float>(new float[2]), ActionSegment<int>.Empty);
+            m_Adaptor.WriteToHeuristic(m_Action, buffer);
+            Assert.IsTrue(Mathf.Approximately(buffer.ContinuousActions[0], 0f));
+            Assert.IsTrue(Mathf.Approximately(buffer.ContinuousActions[1], 1f));
+        }
+    }
+}
+#endif // MLA_INPUT_TESTS && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/Vector2InputActionAdaptorTests.cs.meta
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/Adaptors/Vector2InputActionAdaptorTests.cs.meta
+fileFormatVersion: 2
+guid: 135c4a0f33174ea090da18b0d67bd169
+timeCreated: 1613082080
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/InputActionActuatorTests.cs
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/InputActionActuatorTests.cs
+#if MLA_INPUT_TESTS && UNITY_2019_4_OR_NEWER
+using NUnit.Framework;
+using Unity.Barracuda;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Extensions.Input;
+using Unity.MLAgents.Policies;
+using UnityEngine;
+using UnityEngine.InputSystem;
+
+namespace Unity.MLAgents.Extensions.Tests.Runtime.Input
+{
+    class TestAdaptor : IRLActionInputAdaptor
+    {
+        public bool eventQueued;
+        public bool writtenToHeuristic;
+
+        public ActionSpec GetActionSpecForInputAction(InputAction action)
+        {
+            return ActionSpec.MakeContinuous(1);
+        }
+
+        public void QueueInputEventForAction(InputAction action, InputControl control, ActionSpec actionSpec, in ActionBuffers actionBuffers)
+        {
+            eventQueued = true;
+        }
+
+        public void WriteToHeuristic(InputAction action, in ActionBuffers actionBuffers)
+        {
+            writtenToHeuristic = true;
+        }
+
+        public void Reset()
+        {
+            eventQueued = false;
+            writtenToHeuristic = false;
+        }
+    }
+
+    [TestFixture]
+    public class InputActionActuatorTests
+    {
+        BehaviorParameters m_BehaviorParameters;
+        InputActionActuator m_Actuator;
+        TestAdaptor m_Adaptor;
+
+        [SetUp]
+        public void Setup()
+        {
+            var go = new GameObject();
+            m_BehaviorParameters = go.AddComponent<BehaviorParameters>();
+            var action = new InputAction("action");
+            m_Adaptor = new TestAdaptor();
+            m_Actuator = new InputActionActuator(null, m_BehaviorParameters, action, m_Adaptor);
+        }
+
+        [Test]
+        public void TestOnActionReceived()
+        {
+            m_BehaviorParameters.BehaviorType = BehaviorType.HeuristicOnly;
+            m_Actuator.OnActionReceived(new ActionBuffers());
+            m_Actuator.Heuristic(new ActionBuffers());
+            Assert.IsFalse(m_Adaptor.eventQueued);
+            Assert.IsTrue(m_Adaptor.writtenToHeuristic);
+            m_Adaptor.Reset();
+
+            m_BehaviorParameters.BehaviorType = BehaviorType.Default;
+            m_Actuator.OnActionReceived(new ActionBuffers());
+            Assert.IsFalse(m_Adaptor.eventQueued);
+            m_Adaptor.Reset();
+
+            m_BehaviorParameters.Model = ScriptableObject.CreateInstance<NNModel>();
+            m_Actuator.OnActionReceived(new ActionBuffers());
+            Assert.IsTrue(m_Adaptor.eventQueued);
+            m_Adaptor.Reset();
+
+            Assert.AreEqual(m_Actuator.Name, "InputActionActuator-action");
+            m_Actuator.ResetData();
+            m_Actuator.WriteDiscreteActionMask(null);
+        }
+    }
+}
+#endif // MLA_INPUT_TESTS && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/InputActionActuatorTests.cs.meta
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/InputActionActuatorTests.cs.meta
+fileFormatVersion: 2
+guid: c5adebe0f22e48938e7730469f034c70
+timeCreated: 1613070186
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/InputActuatorComponentTests.cs
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/InputActuatorComponentTests.cs
+#if MLA_INPUT_TESTS
+using System;
+using System.Linq;
+using NUnit.Framework;
+using Unity.MLAgents.Actuators;
+using Unity.MLAgents.Extensions.Input;
+using Unity.MLAgents.Policies;
+using UnityEngine;
+using UnityEngine.InputSystem;
+
+namespace Unity.MLAgents.Extensions.Tests.Runtime.Input
+{
+    class TestProvider : MonoBehaviour, IInputActionAssetProvider
+    {
+        public InputActionAsset asset;
+        public IInputActionCollection2 collection;
+
+        public (InputActionAsset, IInputActionCollection2) GetInputActionAsset()
+        {
+            return (asset, collection);
+        }
+    }
+    public class InputActuatorComponentTests : InputTestFixture
+    {
+        InputActionAsset m_Asset;
+        GameObject m_GameObject;
+        PlayerInput m_PlayerInput;
+        BehaviorParameters m_BehaviorParameters;
+        InputActuatorComponent m_ActuatorComponent;
+        TestPushBlockActions m_Actions;
+        TestProvider m_Provider;
+
+        public override void Setup()
+        {
+            base.Setup();
+            m_Actions = new TestPushBlockActions();
+            m_Asset = m_Actions.asset;
+            m_GameObject = new GameObject();
+            m_PlayerInput = m_GameObject.AddComponent<PlayerInput>();
+            m_Provider = m_GameObject.AddComponent<TestProvider>();
+            m_Provider.asset = m_Asset;
+            m_Provider.collection = m_Actions;
+            m_ActuatorComponent = m_GameObject.AddComponent<InputActuatorComponent>();
+            m_BehaviorParameters = m_GameObject.AddComponent<BehaviorParameters>();
+            m_BehaviorParameters.BehaviorName = "InputTest";
+            m_BehaviorParameters.BehaviorType = BehaviorType.Default;
+        }
+
+        public override void TearDown()
+        {
+            m_ActuatorComponent.CleanupActionAsset();
+            base.TearDown();
+        }
+
+        [Test]
+        public void InputActuatorComponentTestCreateActuators()
+        {
+            // Use the Assert class to test conditions.
+            m_PlayerInput.actions = m_Asset;
+            m_PlayerInput.defaultActionMap = m_Asset.actionMaps[0].name;
+            var actuators = m_ActuatorComponent.CreateActuators();
+            Assert.IsTrue(actuators.Length == 2);
+            Assert.IsTrue(actuators[0].ActionSpec.Equals(ActionSpec.MakeContinuous(2)));
+            Assert.IsTrue(actuators[1].ActionSpec.NumDiscreteActions == 1);
+
+            var actuatorComponentActionSpec = m_ActuatorComponent.ActionSpec;
+            Assert.IsTrue(actuatorComponentActionSpec.BranchSizes.SequenceEqual(new[] {2}));
+            Assert.IsTrue(actuatorComponentActionSpec.NumContinuousActions == 2);
+        }
+
+        [Test]
+        public void InputActuatorComponentTestGenerateActuatorsFromAsset()
+        {
+            // Use the Assert class to test conditions.
+            m_PlayerInput.actions = m_Asset;
+            m_PlayerInput.defaultActionMap = m_Asset.actionMaps[0].name;
+            var inputActionMap = m_Asset.FindActionMap(m_PlayerInput.defaultActionMap);
+            InputActuatorComponent.RegisterLayoutBuilder(
+                inputActionMap,
+                "TestLayout");
+
+            var device = InputSystem.AddDevice("TestLayout");
+
+            var actuators = InputActuatorComponent.CreateActuatorsFromMap(inputActionMap, m_BehaviorParameters, device);
+            Assert.IsTrue(actuators.Length == 2);
+            Assert.IsTrue(actuators[0].ActionSpec.Equals(ActionSpec.MakeContinuous(2)));
+            Assert.IsTrue(actuators[1].ActionSpec.NumDiscreteActions == 1);
+        }
+
+        [Test]
+        public void InputActuatorComponentTestCreateDevice()
+        {
+            // Use the Assert class to test conditions.
+            m_PlayerInput.actions = m_Asset;
+            m_PlayerInput.defaultActionMap = m_Asset.actionMaps[0].name;
+
+            // need to call this to load the layout in the input system
+            InputActuatorComponent.RegisterLayoutBuilder(
+                m_Asset.FindActionMap(m_PlayerInput.defaultActionMap),
+                "TestLayout");
+
+            InputSystem.LoadLayout("TestLayout");
+            var device = InputSystem.AddDevice("TestLayout");
+            Assert.AreEqual("TestLayout", device.layout);
+            Assert.IsTrue(device.children.Count == 2);
+            Assert.AreEqual(device.children[0].path, $"{device.path}{InputControlPath.Separator}movement");
+            Assert.AreEqual(device.children[1].path, $"{device.path}{InputControlPath.Separator}jump");
+            Assert.NotNull(InputSystem.LoadLayout("TestLayout"));
+        }
+    }
+}
+#endif // MLA_INPUT_TESTS
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/InputActuatorComponentTests.cs.meta
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/InputActuatorComponentTests.cs.meta
+fileFormatVersion: 2
+guid: 1894abc7110a4ab39719618db7eb55a9
+timeCreated: 1612910162
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/TestPushBlockActions.cs
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/TestPushBlockActions.cs
+#if MLA_INPUT_TESTS && UNITY_2019_4_OR_NEWER
+//------------------------------------------------------------------------------
+// <auto-generated>
+//     This code was auto-generated by com.unity.inputsystem:InputActionCodeGenerator
+//     version 1.1.0
+//     from Assets/ML-Agents/Examples/PushBlock/TestPushBlockActions.inputactions
+//
+//     Changes to this file may cause incorrect behavior and will be lost if
+//     the code is regenerated.
+// </auto-generated>
+//------------------------------------------------------------------------------
+
+using System;
+using System.Collections;
+using System.Collections.Generic;
+using UnityEngine.InputSystem;
+using UnityEngine.InputSystem.Utilities;
+
+public partial class TestPushBlockActions : IInputActionCollection2, IDisposable
+{
+    public InputActionAsset asset { get; }
+    public TestPushBlockActions()
+    {
+        asset = InputActionAsset.FromJson(@"{
+    ""name"": ""TestPushBlockActions"",
+    ""maps"": [
+        {
+            ""name"": ""Movement"",
+            ""id"": ""03a2e5d4-ae81-47f1-a575-0779fb7da538"",
+            ""actions"": [
+                {
+                    ""name"": ""movement"",
+                    ""type"": ""Value"",
+                    ""id"": ""5f47cbc6-de46-4d33-93e2-2b1af4f5996d"",
+                    ""expectedControlType"": ""Vector2"",
+                    ""processors"": """",
+                    ""interactions"": """"
+                },
+                {
+                    ""name"": ""jump"",
+                    ""type"": ""Value"",
+                    ""id"": ""ca5eb833-5dfb-4b7c-880d-6118bd5d1378"",
+                    ""expectedControlType"": ""Integer"",
+                    ""processors"": """",
+                    ""interactions"": """"
+                }
+            ],
+            ""bindings"": [
+                {
+                    ""name"": ""gamepad_move"",
+                    ""id"": ""477500ef-6d32-4b84-b9f8-158f18bcb906"",
+                    ""path"": ""2DVector"",
+                    ""interactions"": """",
+                    ""processors"": """",
+                    ""groups"": """",
+                    ""action"": ""movement"",
+                    ""isComposite"": true,
+                    ""isPartOfComposite"": false
+                },
+                {
+                    ""name"": ""up"",
+                    ""id"": ""6d2537b8-2266-4a50-8575-fb0fe310daa5"",
+                    ""path"": ""<Gamepad>/dpad/up"",
+                    ""interactions"": """",
+                    ""processors"": """",
+                    ""groups"": ""Keyboard"",
+                    ""action"": ""movement"",
+                    ""isComposite"": false,
+                    ""isPartOfComposite"": true
+                },
+                {
+                    ""name"": ""down"",
+                    ""id"": ""50584c83-beb6-4e90-a453-a635c03a761e"",
+                    ""path"": ""<Gamepad>/dpad/down"",
+                    ""interactions"": """",
+                    ""processors"": """",
+                    ""groups"": ""Keyboard"",
+                    ""action"": ""movement"",
+                    ""isComposite"": false,
+                    ""isPartOfComposite"": true
+                },
+                {
+                    ""name"": ""left"",
+                    ""id"": ""44408b8f-27e7-4c6d-b078-7536ba020d1a"",
+                    ""path"": ""<Gamepad>/dpad/left"",
+                    ""interactions"": """",
+                    ""processors"": """",
+                    ""groups"": ""Keyboard"",
+                    ""action"": ""movement"",
+                    ""isComposite"": false,
+                    ""isPartOfComposite"": true
+                },
+                {
+                    ""name"": ""right"",
+                    ""id"": ""f5681423-d3e3-41a5-b85e-0a7642c774aa"",
+                    ""path"": ""<Gamepad>/dpad/right"",
+                    ""interactions"": """",
+                    ""processors"": """",
+                    ""groups"": ""Keyboard"",
+                    ""action"": ""movement"",
+                    ""isComposite"": false,
+                    ""isPartOfComposite"": true
+                },
+                {
+                    ""name"": ""keyboard_move"",
+                    ""id"": ""6bcba4bf-5ce0-4005-9e6a-0de2487211b0"",
+                    ""path"": ""2DVector"",
+                    ""interactions"": """",
+                    ""processors"": """",
+                    ""groups"": """",
+                    ""action"": ""movement"",
+                    ""isComposite"": true,
+                    ""isPartOfComposite"": false
+                },
+                {
+                    ""name"": ""up"",
+                    ""id"": ""63da699e-b354-4e63-b0f8-26fb92abea41"",
+                    ""path"": ""<Keyboard>/w"",
+                    ""interactions"": """",
+                    ""processors"": """",
+                    ""groups"": ""Keyboard"",
+                    ""action"": ""movement"",
+                    ""isComposite"": false,
+                    ""isPartOfComposite"": true
+                },
+                {
+                    ""name"": ""down"",
+                    ""id"": ""39409748-9002-4aff-9a09-cdc05b9708ad"",
+                    ""path"": ""<Keyboard>/s"",
+                    ""interactions"": """",
+                    ""processors"": """",
+                    ""groups"": ""Keyboard"",
+                    ""action"": ""movement"",
+                    ""isComposite"": false,
+                    ""isPartOfComposite"": true
+                },
+                {
+                    ""name"": ""left"",
+                    ""id"": ""0afe45fc-dc45-4310-9c73-7dc3c503addf"",
+                    ""path"": ""<Keyboard>/a"",
+                    ""interactions"": """",
+                    ""processors"": """",
+                    ""groups"": ""Keyboard"",
+                    ""action"": ""movement"",
+                    ""isComposite"": false,
+                    ""isPartOfComposite"": true
+                },
+                {
+                    ""name"": ""right"",
+                    ""id"": ""69fe0335-9e0c-495d-a90d-4b0fcbfd2b34"",
+                    ""path"": ""<Keyboard>/d"",
+                    ""interactions"": """",
+                    ""processors"": """",
+                    ""groups"": ""Keyboard"",
+                    ""action"": ""movement"",
+                    ""isComposite"": false,
+                    ""isPartOfComposite"": true
+                },
+                {
+                    ""name"": """",
+                    ""id"": ""ab696218-63cd-4eb8-9fe1-48a68e32e92f"",
+                    ""path"": ""<Keyboard>/space"",
+                    ""interactions"": """",
+                    ""processors"": """",
+                    ""groups"": ""Keyboard"",
+                    ""action"": ""jump"",
+                    ""isComposite"": false,
+                    ""isPartOfComposite"": false
+                },
+                {
+                    ""name"": """",
+                    ""id"": ""7adcb138-5175-4cc4-addc-d2b02cb5f0de"",
+                    ""path"": ""<Gamepad>/buttonSouth"",
+                    ""interactions"": """",
+                    ""processors"": """",
+                    ""groups"": ""Keyboard"",
+                    ""action"": ""jump"",
+                    ""isComposite"": false,
+                    ""isPartOfComposite"": false
+                }
+            ]
+        }
+    ],
+    ""controlSchemes"": [
+        {
+            ""name"": ""Keyboard"",
+            ""bindingGroup"": ""Keyboard"",
+            ""devices"": [
+                {
+                    ""devicePath"": ""<Keyboard>"",
+                    ""isOptional"": true,
+                    ""isOR"": false
+                },
+                {
+                    ""devicePath"": ""<Gamepad>"",
+                    ""isOptional"": true,
+                    ""isOR"": false
+                }
+            ]
+        }
+    ]
+}");
+        // Movement
+        m_Movement = asset.FindActionMap("Movement", throwIfNotFound: true);
+        m_Movement_movement = m_Movement.FindAction("movement", throwIfNotFound: true);
+        m_Movement_jump = m_Movement.FindAction("jump", throwIfNotFound: true);
+    }
+
+    public void Dispose()
+    {
+        UnityEngine.Object.Destroy(asset);
+    }
+
+    public InputBinding? bindingMask
+    {
+        get => asset.bindingMask;
+        set => asset.bindingMask = value;
+    }
+
+    public ReadOnlyArray<InputDevice>? devices
+    {
+        get => asset.devices;
+        set => asset.devices = value;
+    }
+
+    public ReadOnlyArray<InputControlScheme> controlSchemes => asset.controlSchemes;
+
+    public bool Contains(InputAction action)
+    {
+        return asset.Contains(action);
+    }
+
+    public IEnumerator<InputAction> GetEnumerator()
+    {
+        return asset.GetEnumerator();
+    }
+
+    IEnumerator IEnumerable.GetEnumerator()
+    {
+        return GetEnumerator();
+    }
+
+    public void Enable()
+    {
+        asset.Enable();
+    }
+
+    public void Disable()
+    {
+        asset.Disable();
+    }
+    public IEnumerable<InputBinding> bindings => asset.bindings;
+
+    public InputAction FindAction(string actionNameOrId, bool throwIfNotFound = false)
+    {
+        return asset.FindAction(actionNameOrId, throwIfNotFound);
+    }
+    public int FindBinding(InputBinding bindingMask, out InputAction action)
+    {
+        return asset.FindBinding(bindingMask, out action);
+    }
+
+    // Movement
+    private readonly InputActionMap m_Movement;
+    private IMovementActions m_MovementActionsCallbackInterface;
+    private readonly InputAction m_Movement_movement;
+    private readonly InputAction m_Movement_jump;
+    public struct MovementActions
+    {
+        private TestPushBlockActions m_Wrapper;
+        public MovementActions(TestPushBlockActions wrapper) { m_Wrapper = wrapper; }
+        public InputAction @movement => m_Wrapper.m_Movement_movement;
+        public InputAction @jump => m_Wrapper.m_Movement_jump;
+        public InputActionMap Get() { return m_Wrapper.m_Movement; }
+        public void Enable() { Get().Enable(); }
+        public void Disable() { Get().Disable(); }
+        public bool enabled => Get().enabled;
+        public static implicit operator InputActionMap(MovementActions set) { return set.Get(); }
+        public void SetCallbacks(IMovementActions instance)
+        {
+            if (m_Wrapper.m_MovementActionsCallbackInterface != null)
+            {
+                @movement.started -= m_Wrapper.m_MovementActionsCallbackInterface.OnMovement;
+                @movement.performed -= m_Wrapper.m_MovementActionsCallbackInterface.OnMovement;
+                @movement.canceled -= m_Wrapper.m_MovementActionsCallbackInterface.OnMovement;
+                @jump.started -= m_Wrapper.m_MovementActionsCallbackInterface.OnJump;
+                @jump.performed -= m_Wrapper.m_MovementActionsCallbackInterface.OnJump;
+                @jump.canceled -= m_Wrapper.m_MovementActionsCallbackInterface.OnJump;
+            }
+            m_Wrapper.m_MovementActionsCallbackInterface = instance;
+            if (instance != null)
+            {
+                @movement.started += instance.OnMovement;
+                @movement.performed += instance.OnMovement;
+                @movement.canceled += instance.OnMovement;
+                @jump.started += instance.OnJump;
+                @jump.performed += instance.OnJump;
+                @jump.canceled += instance.OnJump;
+            }
+        }
+    }
+    public MovementActions @Movement => new MovementActions(this);
+    private int m_KeyboardSchemeIndex = -1;
+    public InputControlScheme KeyboardScheme
+    {
+        get
+        {
+            if (m_KeyboardSchemeIndex == -1) m_KeyboardSchemeIndex = asset.FindControlSchemeIndex("Keyboard");
+            return asset.controlSchemes[m_KeyboardSchemeIndex];
+        }
+    }
+    public interface IMovementActions
+    {
+        void OnMovement(InputAction.CallbackContext context);
+        void OnJump(InputAction.CallbackContext context);
+    }
+}
+#endif // MLA_INPUT_TESTS && UNITY_2019_4_OR_NEWER
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/TestPushBlockActions.cs.meta
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/TestPushBlockActions.cs.meta
+fileFormatVersion: 2
+guid: 87c358a50faa48c1bba3e5f48a5cb75d
+timeCreated: 1613172830
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/Unity.ML-Agents.Extensions.Input.Tests.Runtime.asmdef
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/Unity.ML-Agents.Extensions.Input.Tests.Runtime.asmdef
+{
+    "name": "Unity.ML-Agents.Extensions.Input.Tests.Runtime",
+    "references": [
+        "Unity.ML-Agents",
+        "Unity.ML-Agents.Extensions.Input",
+        "Unity.InputSystem.TestFramework",
+        "UnityEngine.TestRunner",
+        "UnityEditor.TestRunner",
+        "Unity.InputSystem",
+        "Unity.Barracuda"
+    ],
+    "includePlatforms": [],
+    "excludePlatforms": [],
+    "allowUnsafeCode": false,
+    "overrideReferences": true,
+    "precompiledReferences": [
+        "nunit.framework.dll"
+    ],
+    "autoReferenced": true,
+    "defineConstraints": [
+        "UNITY_INCLUDE_TESTS"
+    ],
+    "versionDefines": [
+        {
+            "name": "com.unity.inputsystem",
+            "expression": "1.1.0-preview",
+            "define": "MLA_INPUT_TESTS"
+        }
+    ],
+    "noEngineReferences": false
+}
--- a/com.unity.ml-agents.extensions/Tests/Runtime/Input/Unity.ML-Agents.Extensions.Input.Tests.Runtime.asmdef.meta
+++ b/com.unity.ml-agents.extensions/Tests/Runtime/Input/Unity.ML-Agents.Extensions.Input.Tests.Runtime.asmdef.meta
+fileFormatVersion: 2
+guid: ca257bdcba9544f71baf0c291c36b05a
+AssemblyDefinitionImporter:
+  externalObjects: {}
+  userData: 
+  assetBundleName: 
+  assetBundleVariant: