Merge branch 'develop-hybrid-actions-singleton' into develop-hybrid-actions-csharp

.github/ISSUE_TEMPLATE/bug_report.md

 - Unity Version: [e.g. Unity 2020.1f1]
 - OS + version: [e.g. Windows 10]
 - _ML-Agents version_: (e.g. ML-Agents v0.8, or latest `develop` branch from source)
-- _TensorFlow version_: (you can run `pip3 show tensorflow` to get this)
+- _Torch version_: (you can run `pip3 show torch` to get this)
 - _Environment_: (which example environment you used to reproduce the error)
 
 **NOTE:** We are unable to help reproduce bugs with custom environments.  Please attempt to reproduce your issue with one of the example environments, or provide a minimal patch to one of the environments needed to reproduce the issue.

.github/workflows/pytest.yml

         python -m pip install --progress-bar=off -r test_requirements.txt -c ${{ matrix.pip_constraints }}
         python -m pip install --progress-bar=off -e ./gym-unity -c ${{ matrix.pip_constraints }}
     - name: Save python dependencies
-      run: pip freeze > pip_versions-${{ matrix.python-version }}.txt
+      run: |
+        pip freeze > pip_versions-${{ matrix.python-version }}.txt
+        cat pip_versions-${{ matrix.python-version }}.txt
     - name: Run pytest
       run: pytest --cov=ml-agents --cov=ml-agents-envs --cov=gym-unity --cov-report html --junitxml=junit/test-results-${{ matrix.python-version }}.xml -p no:warnings
     - name: Upload pytest test results

.yamato/test_versions.metafile

 # List of editor versions for standalone-build-test and its dependencies.
+# csharp_backcompat_version is used in training-int-tests to determine the
+# older package version to run the backwards compat tests against.
+    csharp_backcompat_version: 1.0.0
+    csharp_backcompat_version: 1.0.0
-# Waiting on a barracuda fix, see https://jira.unity3d.com/browse/MLA-1464
-#  - version: 2020.2
+    csharp_backcompat_version: 1.0.0
+  - version: 2020.2
+    # 2020.2 moved the AssetImporters namespace
+    # but we didn't handle this until 1.2.0
+    csharp_backcompat_version: 1.2.0

.yamato/training-int-tests.yml

     # If we make a breaking change to the communication protocol, these will need
     # to be disabled until the next release.
     - python -u -m ml-agents.tests.yamato.training_int_tests --python=0.16.0
-    - python -u -m ml-agents.tests.yamato.training_int_tests --csharp=1.0.0
+    - python -u -m ml-agents.tests.yamato.training_int_tests --csharp={{ editor.csharp_backcompat_version }}
   dependencies:
     - .yamato/standalone-build-test.yml#test_mac_standalone_{{ editor.version }}
   triggers:

Project/Assets/ML-Agents/Examples/SharedAssets/Scripts/ModelOverrider.cs

             var bp = m_Agent.GetComponent<BehaviorParameters>();
             var behaviorName = bp.BehaviorName;
 
-            var nnModel = GetModelForBehaviorName(behaviorName);
+            NNModel nnModel = null;
+            try
+            {
+                nnModel = GetModelForBehaviorName(behaviorName);
+            }
+            catch (Exception e)
+            {
+                overrideError = $"Exception calling GetModelForBehaviorName: {e}";
+            }
+
-                overrideError =
-                    $"Didn't find a model for behaviorName {behaviorName}. Make " +
-                    $"sure the behaviorName is set correctly in the commandline " +
-                    $"and that the model file exists";
+                if (string.IsNullOrEmpty(overrideError))
+                {
+                    overrideError =
+                        $"Didn't find a model for behaviorName {behaviorName}. Make " +
+                        "sure the behaviorName is set correctly in the commandline " +
+                        "and that the model file exists";
+                }
             }
             else
             {

README.md

 
 **The Unity Machine Learning Agents Toolkit** (ML-Agents) is an open-source
 project that enables games and simulations to serve as environments for
-training intelligent agents. Agents can be trained using reinforcement learning,
-imitation learning, neuroevolution, or other machine learning methods through a
-simple-to-use Python API. We also provide implementations (based on TensorFlow)
+training intelligent agents. We provide implementations (based on PyTorch)
-train intelligent agents for 2D, 3D and VR/AR games. These trained agents can be
+train intelligent agents for 2D, 3D and VR/AR games. Researchers can also use the
+provided simple-to-use Python API to train Agents using reinforcement learning,
+imitation learning, neuroevolution, or any other methods. These trained agents can be
 used for multiple purposes, including controlling NPC behavior (in a variety of
 settings such as multi-agent and adversarial), automated testing of game builds
 and evaluating different game design decisions pre-release. The ML-Agents

com.unity.ml-agents/CHANGELOG.md

 ### Major Changes
 #### com.unity.ml-agents (C#)
 #### ml-agents / ml-agents-envs / gym-unity (Python)
+ - PyTorch trainers are now the default. See the
+ [installation docs](https://github.com/Unity-Technologies/ml-agents/blob/master/docs/Installation.md) for
+ more information on installing PyTorch. For the time being, TensorFlow is still available;
+ you can use the TensorFlow backend by adding `--tensorflow` to the CLI, or
+ adding `framework: tensorflow` in the configuration YAML. (#4517)
+- The Barracuda dependency was upgraded to 1.1.2 (#4571)
+- The `action_probs` node is no longer listed as an output in TensorFlow models (#4613).
+- `Agent.CollectObservations()` and `Agent.EndEpisode()` will now throw an exception
+if they are called recursively (for example, if they call `Agent.EndEpisode()`).
+Previously, this would result in an infinite loop and cause the editor to hang. (#4573)
+- Fixed an issue where runs could not be resumed when using TensorFlow and Ghost Training. (#4593)
 
 
 ## [1.5.0-preview] - 2020-10-14

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

         // Flag used to keep track of the first time the Academy is reset.
         bool m_HadFirstReset;
 
-        // Whether the Academy is in the middle of a step. This is used to detect and Academy
-        // step called by user code that is also called by the Academy.
-        bool m_IsStepping;
+        // Detect an Academy step called by user code that is also called by the Academy.
+        private RecursionChecker m_StepRecursionChecker = new RecursionChecker("EnvironmentStep");
 
         // Random seed used for inference.
         int m_InferenceSeed;
         /// </summary>
         public void EnvironmentStep()
         {
-            // Check whether we're already in the middle of a step.
-            // This shouldn't happen generally, but could happen if user code (e.g. CollectObservations)
-            // that is called by EnvironmentStep() also calls EnvironmentStep(). This would result
-            // in an infinite loop and/or stack overflow, so stop it before it happens.
-            if (m_IsStepping)
-            {
-                throw new UnityAgentsException(
-                    "Academy.EnvironmentStep() called recursively. " +
-                    "This might happen if you call EnvironmentStep() from custom code such as " +
-                    "CollectObservations() or OnActionReceived()."
-                );
-            }
-
-            m_IsStepping = true;
-
-            try
+            using (m_StepRecursionChecker.Start())
             {
                 if (!m_HadFirstReset)
                 {
                 {
                     AgentAct?.Invoke();
                 }
-            }
-            finally
-            {
-                // Reset m_IsStepping when we're done (or if an exception occurred).
-                m_IsStepping = false;
             }
         }
 

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

         /// </summary>
         internal VectorSensor collectObservationsSensor;
 
+        private RecursionChecker m_CollectObservationsChecker = new RecursionChecker("CollectObservations");
+        private RecursionChecker m_OnEpisodeBeginChecker = new RecursionChecker("OnEpisodeBegin");
+
         /// <summary>
         /// List of IActuators that this Agent will delegate actions to if any exist.
         /// </summary>
             // episode when initializing until after the Academy had its first reset.
             if (Academy.Instance.TotalStepCount != 0)
             {
-                OnEpisodeBegin();
+                using (m_OnEpisodeBeginChecker.Start())
+                {
+                    OnEpisodeBegin();
+                }
             }
         }
 
             {
                 // Make sure the latest observations are being passed to training.
                 collectObservationsSensor.Reset();
-                CollectObservations(collectObservationsSensor);
+                using (m_CollectObservationsChecker.Start())
+                {
+                    CollectObservations(collectObservationsSensor);
+                }
             }
             // Request the last decision with no callbacks
             // We request a decision so Python knows the Agent is done immediately
             UpdateSensors();
             using (TimerStack.Instance.Scoped("CollectObservations"))
             {
-                CollectObservations(collectObservationsSensor);
+                using (m_CollectObservationsChecker.Start())
+                {
+                    CollectObservations(collectObservationsSensor);
+                }
             }
             using (TimerStack.Instance.Scoped("CollectDiscreteActionMasks"))
             {
         {
             ResetData();
             m_StepCount = 0;
-            OnEpisodeBegin();
+            using (m_OnEpisodeBeginChecker.Start())
+            {
+                OnEpisodeBegin();
+            }
+
         }
 
         /// <summary>

com.unity.ml-agents/Tests/Editor/MLAgentsEditModeTest.cs

             }
         }
     }
+
+    [TestFixture]
+    public class AgentRecursionTests
+    {
+        [SetUp]
+        public void SetUp()
+        {
+            if (Academy.IsInitialized)
+            {
+                Academy.Instance.Dispose();
+            }
+        }
+
+        class CollectObsEndEpisodeAgent : Agent
+        {
+            public override void CollectObservations(VectorSensor sensor)
+            {
+                // NEVER DO THIS IN REAL CODE!
+                EndEpisode();
+            }
+        }
+
+        class OnEpisodeBeginEndEpisodeAgent : Agent
+        {
+            public override void OnEpisodeBegin()
+            {
+                // NEVER DO THIS IN REAL CODE!
+                EndEpisode();
+            }
+        }
+
+        void TestRecursiveThrows<T>() where T : Agent
+        {
+            var gameObj = new GameObject();
+            var agent = gameObj.AddComponent<T>();
+            agent.LazyInitialize();
+            agent.RequestDecision();
+
+            Assert.Throws<UnityAgentsException>(() =>
+            {
+                Academy.Instance.EnvironmentStep();
+            });
+        }
+
+        [Test]
+        public void TestRecursiveCollectObsEndEpisodeThrows()
+        {
+            TestRecursiveThrows<CollectObsEndEpisodeAgent>();
+        }
+
+        [Test]
+        public void TestRecursiveOnEpisodeBeginEndEpisodeThrows()
+        {
+            TestRecursiveThrows<OnEpisodeBeginEndEpisodeAgent>();
+        }
+    }
 }

com.unity.ml-agents/package.json

   "unity": "2018.4",
   "description": "Use state-of-the-art machine learning to create intelligent character behaviors in any Unity environment (games, robotics, film, etc.).",
   "dependencies": {
-    "com.unity.barracuda": "1.1.1-preview",
+    "com.unity.barracuda": "1.1.2-preview",
     "com.unity.modules.imageconversion": "1.0.0",
     "com.unity.modules.jsonserialize": "1.0.0",
     "com.unity.modules.physics": "1.0.0",

docs/Background-Machine-Learning.md

 one where the number of observations an agent perceives and the number of
 actions they can take are large). Many of the algorithms we provide in ML-Agents
 use some form of deep learning, built on top of the open-source library,
-[TensorFlow](Background-TensorFlow.md).
+[PyTorch](Background-PyTorch.md).

docs/Getting-Started.md

 
 ## Running a pre-trained model
 
-We include pre-trained models for our agents (`.nn` files) and we use the
+We include pre-trained models for our agents (`.onnx` files) and we use the
 [Unity Inference Engine](Unity-Inference-Engine.md) to run these models inside
 Unity. In this section, we will use the pre-trained model for the 3D Ball
 example.
 
 ## Training a new model with Reinforcement Learning
 
-While we provide pre-trained `.nn` files for the agents in this environment, any
+While we provide pre-trained models for the agents in this environment, any
 environment you make yourself will require training agents from scratch to
 generate a new model file. In this section we will demonstrate how to use the
 reinforcement learning algorithms that are part of the ML-Agents Python package
 use it with compatible Agents (the Agents that generated the model). **Note:**
 Do not just close the Unity Window once the `Saved Model` message appears.
 Either wait for the training process to close the window or press `Ctrl+C` at
-the command-line prompt. If you close the window manually, the `.nn` file
+the command-line prompt. If you close the window manually, the `.onnx` file
 containing the trained model is not exported into the ml-agents folder.
 
 If you've quit the training early using `Ctrl+C` and want to resume training,
 mlagents-learn config/ppo/3DBall.yaml --run-id=first3DBallRun --resume
 ```
 
-Your trained model will be at `results/<run-identifier>/<behavior_name>.nn` where
+Your trained model will be at `results/<run-identifier>/<behavior_name>.onnx` where
 `<behavior_name>` is the name of the `Behavior Name` of the agents corresponding
 to the model. This file corresponds to your model's latest checkpoint. You can
 now embed this trained model into your Agents by following the steps below,
    `Project/Assets/ML-Agents/Examples/3DBall/TFModels/`.
 1. Open the Unity Editor, and select the **3DBall** scene as described above.
 1. Select the **3DBall** prefab Agent object.
-1. Drag the `<behavior_name>.nn` file from the Project window of the Editor to
+1. Drag the `<behavior_name>.onnx` file from the Project window of the Editor to
    the **Model** placeholder in the **Ball3DAgent** inspector window.
 1. Press the **Play** button at the top of the Editor.
 

docs/Installation.md

 [instructions](https://packaging.python.org/guides/installing-using-linux-tools/#installing-pip-setuptools-wheel-with-linux-package-managers)
 on installing it.
 
-Although we do not provide support for Anaconda installation on Windows, the
-previous
-[Windows Anaconda Installation (Deprecated) guide](Installation-Anaconda-Windows.md)
-is still available.
-
 ### Clone the ML-Agents Toolkit Repository (Optional)
 
 Now that you have installed Unity and Python, you can now install the Unity and
 dependencies for each project and are supported on Mac / Windows / Linux. We
 offer a dedicated [guide on Virtual Environments](Using-Virtual-Environment.md).
 
+#### (Windows) Installing PyTorch
+
+On Windows, you'll have to install the PyTorch package separately prior to
+installing ML-Agents. Activate your virtual environment and run from the command line:
+
+```sh
+pip3 install torch==1.7.0 -f https://download.pytorch.org/whl/torch_stable.html
+```
+
+Note that on Windows, you may also need Microsoft's
+[Visual C++ Redistributable](https://support.microsoft.com/en-us/help/2977003/the-latest-supported-visual-c-downloads)
+if you don't have it already. See the [PyTorch installation guide](https://pytorch.org/get-started/locally/)
+for more installation options and versions.
+
+#### Installing `mlagents`
+
 To install the `mlagents` Python package, activate your virtual environment and
 run from the command line:
 
 
 By installing the `mlagents` package, the dependencies listed in the
 [setup.py file](../ml-agents/setup.py) are also installed. These include
-[TensorFlow](Background-TensorFlow.md) (Requires a CPU w/ AVX support).
+[PyTorch](Background-PyTorch.md) (Requires a CPU w/ AVX support).
 
 #### Advanced: Local Installation for Development
 
 the repository's root directory, run:
 
 ```sh
+pip3 install torch -f https://download.pytorch.org/whl/torch_stable.html
 pip3 install -e ./ml-agents-envs
 pip3 install -e ./ml-agents
 ```

docs/Learning-Environment-Executable.md

 ```
 
 You can press Ctrl+C to stop the training, and your trained model will be at
-`results/<run-identifier>/<behavior_name>.nn`, which corresponds to your model's
+`results/<run-identifier>/<behavior_name>.onnx`, which corresponds to your model's
 latest checkpoint. (**Note:** There is a known bug on Windows that causes the
 saving of the model to fail when you early terminate the training, it's
 recommended to wait until Step has reached the max_steps parameter you set in
    `Project/Assets/ML-Agents/Examples/3DBall/TFModels/`.
 1. Open the Unity Editor, and select the **3DBall** scene as described above.
 1. Select the **3DBall** prefab from the Project window and select **Agent**.
-1. Drag the `<behavior_name>.nn` file from the Project window of the Editor to
+1. Drag the `<behavior_name>.onnx` file from the Project window of the Editor to
    the **Model** placeholder in the **Ball3DAgent** inspector window.
 1. Press the **Play** button at the top of the Editor.

docs/ML-Agents-Overview.md

 for training intelligent agents. Agents can be trained using reinforcement
 learning, imitation learning, neuroevolution, or other machine learning methods
 through a simple-to-use Python API. We also provide implementations (based on
-TensorFlow) of state-of-the-art algorithms to enable game developers and
+PyTorch) of state-of-the-art algorithms to enable game developers and
 hobbyists to easily train intelligent agents for 2D, 3D and VR/AR games. These
 trained agents can be used for multiple purposes, including controlling NPC
 behavior (in a variety of settings such as multi-agent and adversarial),
 that include overviews and helpful resources on the
 [Unity Engine](Background-Unity.md),
 [machine learning](Background-Machine-Learning.md) and
-[TensorFlow](Background-TensorFlow.md). We **strongly** recommend browsing the
+[PyTorch](Background-PyTorch.md). We **strongly** recommend browsing the
-machine learning concepts or have not previously heard of TensorFlow.
+machine learning concepts or have not previously heard of PyTorch.
 
 The remainder of this page contains a deep dive into ML-Agents, its key
 components, different training modes and scenarios. By the end of it, you should
 
 ### Custom Training and Inference
 
-In the previous mode, the Agents were used for training to generate a TensorFlow
+In the previous mode, the Agents were used for training to generate a PyTorch
 model that the Agents can later use. However, any user of the ML-Agents Toolkit
 can leverage their own algorithms for training. In this case, the behaviors of
 all the Agents in the scene will be controlled within Python. You can even turn

docs/Readme.md

 - [ML-Agents Toolkit Overview](ML-Agents-Overview.md)
   - [Background: Unity](Background-Unity.md)
   - [Background: Machine Learning](Background-Machine-Learning.md)
-  - [Background: TensorFlow](Background-TensorFlow.md)
+  - [Background: PyTorch](Background-PyTorch.md)
 - [Example Environments](Learning-Environment-Examples.md)
 
 ## Creating Learning Environments

docs/Training-Configuration-File.md

 | `time_horizon`           | (default = `64`) How many steps of experience to collect per-agent before adding it to the experience buffer. When this limit is reached before the end of an episode, a value estimate is used to predict the overall expected reward from the agent's current state. As such, this parameter trades off between a less biased, but higher variance estimate (long time horizon) and more biased, but less varied estimate (short time horizon). In cases where there are frequent rewards within an episode, or episodes are prohibitively large, a smaller number can be more ideal. This number should be large enough to capture all the important behavior within a sequence of an agent's actions. <br><br> Typical range: `32` - `2048` |
 | `max_steps`              | (default = `500000`) Total number of steps (i.e., observation collected and action taken) that must be taken in the environment (or across all environments if using multiple in parallel) before ending the training process. If you have multiple agents with the same behavior name within your environment, all steps taken by those agents will contribute to the same `max_steps` count. <br><br>Typical range: `5e5` - `1e7`                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                                          |
 | `keep_checkpoints`         | (default = `5`) The maximum number of model checkpoints to keep. Checkpoints are saved after the number of steps specified by the checkpoint_interval option. Once the maximum number of checkpoints has been reached, the oldest checkpoint is deleted when saving a new checkpoint. |
-| `checkpoint_interval`         | (default = `500000`) The number of experiences collected between each checkpoint by the trainer. A maximum of `keep_checkpoints` checkpoints are saved before old ones are deleted. Each checkpoint saves the `.nn` (and `.onnx` if applicable) files in `results/` folder.|
+| `checkpoint_interval`         | (default = `500000`) The number of experiences collected between each checkpoint by the trainer. A maximum of `keep_checkpoints` checkpoints are saved before old ones are deleted. Each checkpoint saves the `.onnx` (and `.nn` if using TensorFlow) files in `results/` folder.|
 | `init_path`              | (default = None) Initialize trainer from a previously saved model. Note that the prior run should have used the same trainer configurations as the current run, and have been saved with the same version of ML-Agents. <br><br>You should provide the full path to the folder where the checkpoints were saved, e.g. `./models/{run-id}/{behavior_name}`. This option is provided in case you want to initialize different behaviors from different runs; in most cases, it is sufficient to use the `--initialize-from` CLI parameter to initialize all models from the same run.                                                                                                                                  |
 | `threaded`               | (default = `true`) By default, model updates can happen while the environment is being stepped. This violates the [on-policy](https://spinningup.openai.com/en/latest/user/algorithms.html#the-on-policy-algorithms) assumption of PPO slightly in exchange for a training speedup. To maintain the strict on-policyness of PPO, you can disable parallel updates by setting `threaded` to `false`. There is usually no reason to turn `threaded` off for SAC.                                                                                                                                                                                                                                                       |
 | `hyperparameters -> learning_rate`          | (default = `3e-4`) Initial learning rate for gradient descent. Corresponds to the strength of each gradient descent update step. This should typically be decreased if training is unstable, and the reward does not consistently increase. <br><br>Typical range: `1e-5` - `1e-3`                                                                                                                                                                                                                                                                                                                                                                                                                                                                |

docs/Training-ML-Agents.md

     - [Curriculum Learning](#curriculum)
       - [Training with a Curriculum](#training-with-a-curriculum)
   - [Training Using Concurrent Unity Instances](#training-using-concurrent-unity-instances)
-  - [Using PyTorch (Experimental)](#using-pytorch-experimental)
 
 For a broad overview of reinforcement learning, imitation learning and all the
 training scenarios, methods and options within the ML-Agents Toolkit, see
    values. See [Using TensorBoard](Using-Tensorboard.md) for more details on how
    to visualize the training metrics.
 1. Models: these contain the model checkpoints that
-   are updated throughout training and the final model file (`.nn`). This final
+   are updated throughout training and the final model file (`.onnx`). This final
    model file is generated once either when training completes or is
    interrupted.
 1. Timers file (under `results/<run-identifier>/run_logs`): this contains aggregated
 - **Result Variation Using Concurrent Unity Instances** - If you keep all the
   hyperparameters the same, but change `--num-envs=<n>`, the results and model
   would likely change.
-
-### Using PyTorch (Experimental)
-
-ML-Agents, by default, uses TensorFlow as its backend, but experimental support
-for PyTorch has been added. To use PyTorch, the `torch` Python package must
-be installed, and PyTorch must be enabled for your trainer.
-
-#### Installing PyTorch
-
-If you've already installed ML-Agents, follow the
-[official PyTorch install instructions](https://pytorch.org/get-started/locally/) for
-your platform and configuration. Note that on Windows, you may also need Microsoft's
-[Visual C++ Redistributable](https://support.microsoft.com/en-us/help/2977003/the-latest-supported-visual-c-downloads) if you don't have it already.
-
-If you're installing or upgrading ML-Agents on Linux or Mac, you can also run
-`pip3 install mlagents[torch]` instead of `pip3 install mlagents`
-during [installation](Installation.md). On Windows, install ML-Agents first and then
-separately install PyTorch.
-
-#### Enabling PyTorch
-
-PyTorch can be enabled in one of two ways. First, by adding `--torch` to the
-`mlagents-learn` command. This will make all behaviors train with PyTorch.
-
-Second, by changing the `framework` option for your agent behavior in the
-configuration YAML as below. This will use PyTorch just for that behavior.
-
-```yaml
-behaviors:
-  YourAgentBehavior:
-    framework: pytorch
-```

docs/Training-on-Amazon-Web-Service.md

     # Download and install the latest Nvidia driver for ubuntu
     # Please refer to http://download.nvidia.com/XFree86/Linux-#x86_64/latest.txt
     $ wget http://download.nvidia.com/XFree86/Linux-x86_64/390.87/NVIDIA-Linux-x86_64-390.87.run
-    $ sudo /bin/bash ./NVIDIA-Linux-x86_64-390.67.run --accept-license --no-questions --ui=none
+    $ sudo /bin/bash ./NVIDIA-Linux-x86_64-390.87.run --accept-license --no-questions --ui=none
 
     # Disable Nouveau as it will clash with the Nvidia driver
     $ sudo echo 'blacklist nouveau'  | sudo tee -a /etc/modprobe.d/blacklist.conf

docs/Unity-Inference-Engine.md

   [industry-standard open format](https://onnx.ai/about.html) produced by the
   [tf2onnx package](https://github.com/onnx/tensorflow-onnx).
 
-Export to ONNX is currently considered beta. To enable it, make sure
-`tf2onnx>=1.5.5` is installed in pip. tf2onnx does not currently support
-tensorflow 2.0.0 or later, or earlier than 1.12.0.
+Export to ONNX is used if using PyTorch (the default). To enable it
+while using TensorFlow, make sure `tf2onnx>=1.6.1` is installed in pip.
 
 ## Using the Unity Inference Engine
 

gym-unity/gym_unity/envs/__init__.py

         self._previous_decision_step = decision_steps
 
         # Set action spaces
-        if self.group_spec.is_action_discrete():
-            branches = self.group_spec.discrete_action_branches
-            if self.group_spec.action_size == 1:
+        if self.group_spec.action_spec.is_discrete():
+            self.action_size = self.group_spec.action_spec.discrete_size
+            branches = self.group_spec.action_spec.discrete_branches
+            if self.group_spec.action_spec.discrete_size == 1:
                 self._action_space = spaces.Discrete(branches[0])
             else:
                 if flatten_branched:
                     self._action_space = spaces.MultiDiscrete(branches)
 
-        else:
+        elif self.group_spec.action_spec.is_continuous():
-            high = np.array([1] * self.group_spec.action_shape)
+
+            self.action_size = self.group_spec.action_spec.continuous_size
+            high = np.array([1] * self.group_spec.action_spec.continuous_size)
+        else:
+            raise UnityGymException(
+                "The gym wrapper does not provide explicit support for both discrete "
+                "and continuous actions."
+            )
 
         # Set observations space
         list_spaces: List[gym.Space] = []
             # Translate action into list
             action = self._flattener.lookup_action(action)
 
-        spec = self.group_spec
-        action = np.array(action).reshape((1, spec.action_size))
+        action = np.array(action).reshape((1, self.action_size))
         self._env.set_actions(self.name, action)
 
         self._env.step()

gym-unity/gym_unity/tests/test_gym.py

 from gym_unity.envs import UnityToGymWrapper
 from mlagents_envs.base_env import (
     BehaviorSpec,
-    ActionType,
+    ActionSpec,
     DecisionSteps,
     TerminalSteps,
     BehaviorMapping,
     Creates a mock BrainParameters object with parameters.
     """
     # Avoid using mutable object as default param
-    act_type = ActionType.DISCRETE
-        act_type = ActionType.CONTINUOUS
+        action_spec = ActionSpec.create_continuous(vector_action_space_size)
+        action_spec = ActionSpec.create_discrete(vector_action_space_size)
-    return BehaviorSpec(obs_shapes, act_type, vector_action_space_size)
+    return BehaviorSpec(obs_shapes, action_spec)
 
 
 def create_mock_vector_steps(specs, num_agents=1, number_visual_observations=0):

ml-agents-envs/mlagents_envs/base_env.py

     NamedTuple,
     Tuple,
     Optional,
-    Union,
     Dict,
     Iterator,
     Any,
-from enum import Enum
+
+from mlagents_envs.exception import UnityActionException
 
 AgentId = int
 BehaviorName = str
         )
 
 
-class ActionType(Enum):
-    DISCRETE = 0
-    CONTINUOUS = 1
+class ActionTuple:
+    """
+    An object whose fields correspond to actions of different types.
+    Continuous and discrete actions are numpy arrays of type float32 and
+    int32, respectively and are type checked on construction.
+    Dimensions are of (n_agents, continuous_size) and (n_agents, discrete_size),
+    respectively.
+    """
+    def __init__(
+        self,
+        continuous: Optional[np.ndarray] = None,
+        discrete: Optional[np.ndarray] = None,
+    ):
+        if continuous is not None and continuous.dtype != np.float32:
+            continuous = continuous.astype(np.float32, copy=False)
+        self._continuous = continuous
+        if discrete is not None and discrete.dtype != np.int32:
+            discrete = discrete.astype(np.int32, copy=False)
+        self._discrete = discrete
-class BehaviorSpec(NamedTuple):
+    @property
+    def continuous(self) -> np.ndarray:
+        return self._continuous
+
+    @property
+    def discrete(self) -> np.ndarray:
+        return self._discrete
+
+
+class ActionSpec(NamedTuple):
-    A NamedTuple to containing information about the observations and actions
-    spaces for a group of Agents under the same behavior.
-     - observation_shapes is a List of Tuples of int : Each Tuple corresponds
-     to an observation's dimensions. The shape tuples have the same ordering as
-     the ordering of the DecisionSteps and TerminalSteps.
-     - action_type is the type of data of the action. it can be discrete or
-     continuous. If discrete, the action tensors are expected to be int32. If
-     continuous, the actions are expected to be float32.
-     - action_shape is:
-       - An int in continuous action space corresponding to the number of
-     floats that constitute the action.
-       - A Tuple of int in discrete action space where each int corresponds to
-       the number of discrete actions available to the agent.
+    A NamedTuple containing utility functions and information about the action spaces
+    for a group of Agents under the same behavior.
+    - num_continuous_actions is an int corresponding to the number of floats which
+    constitute the action.
+    - discrete_branch_sizes is a Tuple of int where each int corresponds to
+    the number of discrete actions available to the agent on an independent action branch.
-    observation_shapes: List[Tuple]
-    action_type: ActionType
-    action_shape: Union[int, Tuple[int, ...]]
+    continuous_size: int
+    discrete_branches: Tuple[int, ...]
+
+    def __eq__(self, other):
+        return (
+            self.continuous_size == other.continuous_size
+            and self.discrete_branches == other.discrete_branches
+        )
+
+    def __str__(self):
+        return f"Continuous: {self.continuous_size}, Discrete: {self.discrete_branches}"
-    def is_action_discrete(self) -> bool:
+    # For backwards compatibility
+    def is_discrete(self) -> bool:
-        return self.action_type == ActionType.DISCRETE
+        return self.discrete_size > 0 and self.continuous_size == 0
-    def is_action_continuous(self) -> bool:
+    # For backwards compatibility
+    def is_continuous(self) -> bool:
-        return self.action_type == ActionType.CONTINUOUS
+        return self.discrete_size == 0 and self.continuous_size > 0
-    def action_size(self) -> int:
-        """
-        Returns the dimension of the action.
-         - In the continuous case, will return the number of continuous actions.
-         - In the (multi-)discrete case, will return the number of action.
-         branches.
-        """
-        if self.action_type == ActionType.DISCRETE:
-            return len(self.action_shape)  # type: ignore
-        else:
-            return self.action_shape  # type: ignore
-
-    @property
-    def discrete_action_branches(self) -> Optional[Tuple[int, ...]]:
+    def discrete_size(self) -> int:
-        Returns a Tuple of int corresponding to the number of possible actions
-        for each branch (only for discrete actions). Will return None in
-        for continuous actions.
+        Returns a an int corresponding to the number of discrete branches.
-        if self.action_type == ActionType.DISCRETE:
-            return self.action_shape  # type: ignore
-        else:
-            return None
+        return len(self.discrete_branches)
-    def create_empty_action(self, n_agents: int) -> np.ndarray:
+    def empty_action(self, n_agents: int) -> ActionTuple:
-        Generates a numpy array corresponding to an empty action (all zeros)
+        Generates ActionTuple corresponding to an empty action (all zeros)
-        if self.action_type == ActionType.DISCRETE:
-            return np.zeros((n_agents, self.action_size), dtype=np.int32)
-        else:
-            return np.zeros((n_agents, self.action_size), dtype=np.float32)
+        continuous = np.zeros((n_agents, self.continuous_size), dtype=np.float32)
+        discrete = np.zeros((n_agents, self.discrete_size), dtype=np.int32)
+        return ActionTuple(continuous, discrete)
-    def create_random_action(self, n_agents: int) -> np.ndarray:
+    def random_action(self, n_agents: int) -> ActionTuple:
-        Generates a numpy array corresponding to a random action (either discrete
+        Generates ActionTuple corresponding to a random action (either discrete
-        :param generator: The random number generator used for creating random action
-        if self.is_action_continuous():
-            action = np.random.uniform(
-                low=-1.0, high=1.0, size=(n_agents, self.action_size)
-            ).astype(np.float32)
-            return action
-        elif self.is_action_discrete():
-            branch_size = self.discrete_action_branches
-            action = np.column_stack(
+        continuous = np.random.uniform(
+            low=-1.0, high=1.0, size=(n_agents, self.continuous_size)
+        )
+        discrete = np.zeros((n_agents, self.discrete_size), dtype=np.int32)
+        if self.discrete_size > 0:
+            discrete = np.column_stack(
-                        branch_size[i],  # type: ignore
+                        self.discrete_branches[i],  # type: ignore
-                    for i in range(self.action_size)
+                    for i in range(self.discrete_size)
-            return action
+        return ActionTuple(continuous, discrete)
+
+    def _validate_action(
+        self, actions: ActionTuple, n_agents: int, name: str
+    ) -> ActionTuple:
+        """
+        Validates that action has the correct action dim
+        for the correct number of agents and ensures the type.
+        """
+        _expected_shape = (n_agents, self.continuous_size)
+        if self.continuous_size > 0 and actions.continuous.shape != _expected_shape:
+            raise UnityActionException(
+                f"The behavior {name} needs a continuous input of dimension "
+                f"{_expected_shape} for (<number of agents>, <action size>) but "
+                f"received input of dimension {actions.continuous.shape}"
+            )
+        _expected_shape = (n_agents, self.discrete_size)
+        if self.discrete_size > 0 and actions.discrete.shape != _expected_shape:
+            raise UnityActionException(
+                f"The behavior {name} needs a discrete input of dimension "
+                f"{_expected_shape} for (<number of agents>, <action size>) but "
+                f"received input of dimension {actions.discrete.shape}"
+            )
+        return actions
+
+    @staticmethod
+    def create_continuous(continuous_size: int) -> "ActionSpec":
+        """
+        Creates an ActionSpec that is homogenously continuous
+        """
+        return ActionSpec(continuous_size, ())
+
+    @staticmethod
+    def create_discrete(discrete_branches: Tuple[int]) -> "ActionSpec":
+        """
+        Creates an ActionSpec that is homogenously discrete
+        """
+        return ActionSpec(0, discrete_branches)
+
+
+class BehaviorSpec(NamedTuple):
+    """
+    A NamedTuple containing information about the observation and action
+    spaces for a group of Agents under the same behavior.
+    - observation_shapes is a List of Tuples of int : Each Tuple corresponds
+    to an observation's dimensions. The shape tuples have the same ordering as
+    the ordering of the DecisionSteps and TerminalSteps.
+    - action_spec is an ActionSpec NamedTuple
+    """
+
+    observation_shapes: List[Tuple]
+    action_spec: ActionSpec
 
 
 class BehaviorMapping(Mapping):
         """
 
     @abstractmethod
-    def set_actions(self, behavior_name: BehaviorName, action: np.ndarray) -> None:
+    def set_actions(self, behavior_name: BehaviorName, action: ActionTuple) -> None:
-        :param action: A two dimensional np.ndarray corresponding to the action
-        (either int or float)
+        :param action: ActionTuple tuple of continuous and/or discrete action.
+        Actions are np.arrays with dimensions  (n_agents, continuous_size) and
+        (n_agents, discrete_size), respectively.
-        self, behavior_name: BehaviorName, agent_id: AgentId, action: np.ndarray
+        self, behavior_name: BehaviorName, agent_id: AgentId, action: ActionTuple
     ) -> None:
         """
         Sets the action for one of the agents in the simulation for the next
-        :param action: A one dimensional np.ndarray corresponding to the action
-        (either int or float)
+        :param action: ActionTuple tuple of continuous and/or discrete action
+        Actions are np.arrays with dimensions  (1, continuous_size) and
+        (1, discrete_size), respectively. Note, this initial dimensions of 1 is because
+        this action is meant for a single agent.
         """
 
     @abstractmethod

ml-agents-envs/mlagents_envs/environment.py

     DecisionSteps,
     TerminalSteps,
     BehaviorSpec,
+    ActionTuple,
     BehaviorName,
     AgentId,
     BehaviorMapping,
 
         self._env_state: Dict[str, Tuple[DecisionSteps, TerminalSteps]] = {}
         self._env_specs: Dict[str, BehaviorSpec] = {}
-        self._env_actions: Dict[str, np.ndarray] = {}
+        self._env_actions: Dict[str, ActionTuple] = {}
         self._is_first_message = True
         self._update_behavior_specs(aca_output)
 
                     n_agents = len(self._env_state[group_name][0])
                 self._env_actions[group_name] = self._env_specs[
                     group_name
-                ].create_empty_action(n_agents)
+                ].action_spec.empty_action(n_agents)
         step_input = self._generate_step_input(self._env_actions)
         with hierarchical_timer("communicator.exchange"):
             outputs = self._communicator.exchange(step_input)
                 f"agent group in the environment"
             )
 
-    def set_actions(self, behavior_name: BehaviorName, action: np.ndarray) -> None:
+    def set_actions(self, behavior_name: BehaviorName, action: ActionTuple) -> None:
-        spec = self._env_specs[behavior_name]
-        expected_type = np.float32 if spec.is_action_continuous() else np.int32
-        expected_shape = (len(self._env_state[behavior_name][0]), spec.action_size)
-        if action.shape != expected_shape:
-            raise UnityActionException(
-                f"The behavior {behavior_name} needs an input of dimension "
-                f"{expected_shape} for (<number of agents>, <action size>) but "
-                f"received input of dimension {action.shape}"
-            )
-        if action.dtype != expected_type:
-            action = action.astype(expected_type)
+        action_spec = self._env_specs[behavior_name].action_spec
+        num_agents = len(self._env_state[behavior_name][0])
+        action = action_spec._validate_action(action, num_agents, behavior_name)
-        self, behavior_name: BehaviorName, agent_id: AgentId, action: np.ndarray
+        self, behavior_name: BehaviorName, agent_id: AgentId, action: ActionTuple
-        spec = self._env_specs[behavior_name]
-        expected_shape = (spec.action_size,)
-        if action.shape != expected_shape:
-            raise UnityActionException(
-                f"The Agent {agent_id} with BehaviorName {behavior_name} needs "
-                f"an input of dimension {expected_shape} but received input of "
-                f"dimension {action.shape}"
-            )
-        expected_type = np.float32 if spec.is_action_continuous() else np.int32
-        if action.dtype != expected_type:
-            action = action.astype(expected_type)
-
+        action_spec = self._env_specs[behavior_name].action_spec
+        num_agents = len(self._env_state[behavior_name][0])
+        action = action_spec._validate_action(action, num_agents, behavior_name)
-            self._env_actions[behavior_name] = spec.create_empty_action(
-                len(self._env_state[behavior_name][0])
-            )
+            self._env_actions[behavior_name] = action_spec.empty_action(num_agents)
         try:
             index = np.where(self._env_state[behavior_name][0].agent_id == agent_id)[0][
                 0
                     agent_id
                 )
             ) from ie
-        self._env_actions[behavior_name][index] = action
+        if action_spec.continuous_size > 0:
+            self._env_actions[behavior_name].continuous[index] = action.continuous[0, :]
+        if action_spec.discrete_size > 0:
+            self._env_actions[behavior_name].discrete[index] = action.discrete[0, :]
 
     def get_steps(
         self, behavior_name: BehaviorName
 
     @timed
     def _generate_step_input(
-        self, vector_action: Dict[str, np.ndarray]
+        self, vector_action: Dict[str, ActionTuple]
     ) -> UnityInputProto:
         rl_in = UnityRLInputProto()
         for b in vector_action:
             for i in range(n_agents):
-                action = AgentActionProto(vector_actions=vector_action[b][i])
+                # TODO: extend to AgentBuffers
+                if vector_action[b].continuous is not None:
+                    _act = vector_action[b].continuous[i]
+                else:
+                    _act = vector_action[b].discrete[i]
+                action = AgentActionProto(vector_actions=_act)
                 rl_in.agent_actions[b].value.extend([action])
                 rl_in.command = STEP
         rl_in.side_channel = bytes(

ml-agents-envs/mlagents_envs/rpc_utils.py

 from mlagents_envs.base_env import (
+    ActionSpec,
-    ActionType,
     DecisionSteps,
     TerminalSteps,
 )
 from mlagents_envs.communicator_objects.brain_parameters_pb2 import BrainParametersProto
 import numpy as np
 import io
-from typing import cast, List, Tuple, Union, Collection, Optional, Iterable
+from typing import cast, List, Tuple, Collection, Optional, Iterable
 from PIL import Image
 
 
     :return: BehaviorSpec object.
     """
     observation_shape = [tuple(obs.shape) for obs in agent_info.observations]
-    action_type = (
-        ActionType.DISCRETE
-        if brain_param_proto.vector_action_space_type_deprecated == 0
-        else ActionType.CONTINUOUS
+    action_spec_proto = brain_param_proto.action_spec
+    action_spec = ActionSpec(
+        action_spec_proto.num_continuous_actions,
+        tuple(branch for branch in action_spec_proto.discrete_branch_sizes),
-    if action_type == ActionType.CONTINUOUS:
-        action_shape: Union[
-            int, Tuple[int, ...]
-        ] = brain_param_proto.vector_action_size_deprecated[0]
-    else:
-        action_shape = tuple(brain_param_proto.vector_action_size_deprecated)
-    return BehaviorSpec(observation_shape, action_type, action_shape)
+    return BehaviorSpec(observation_shape, action_spec)
 
 
 class OffsetBytesIO:
         [agent_info.id for agent_info in terminal_agent_info_list], dtype=np.int32
     )
     action_mask = None
-    if behavior_spec.is_action_discrete():
+    if behavior_spec.action_spec.discrete_size > 0:
-            a_size = np.sum(behavior_spec.discrete_action_branches)
+            a_size = np.sum(behavior_spec.action_spec.discrete_branches)
             mask_matrix = np.ones((n_agents, a_size), dtype=np.bool)
             for agent_index, agent_info in enumerate(decision_agent_info_list):
                 if agent_info.action_mask is not None:
                             for k in range(a_size)
                         ]
             action_mask = (1 - mask_matrix).astype(np.bool)
-            indices = _generate_split_indices(behavior_spec.discrete_action_branches)
+            indices = _generate_split_indices(
+                behavior_spec.action_spec.discrete_branches
+            )
             action_mask = np.split(action_mask, indices, axis=1)
     return (
         DecisionSteps(

ml-agents-envs/mlagents_envs/tests/test_envs.py

 from unittest import mock
 import pytest
 
-import numpy as np
-
-from mlagents_envs.base_env import DecisionSteps, TerminalSteps
+from mlagents_envs.base_env import DecisionSteps, TerminalSteps, ActionTuple
 from mlagents_envs.exception import UnityEnvironmentException, UnityActionException
 from mlagents_envs.mock_communicator import MockCommunicator
 
     env.step()
     decision_steps, terminal_steps = env.get_steps("RealFakeBrain")
     n_agents = len(decision_steps)
-    env.set_actions(
-        "RealFakeBrain", np.zeros((n_agents, spec.action_size), dtype=np.float32)
-    )
+    env.set_actions("RealFakeBrain", spec.action_spec.empty_action(n_agents))
-        env.set_actions(
-            "RealFakeBrain",
-            np.zeros((n_agents - 1, spec.action_size), dtype=np.float32),
-        )
+        env.set_actions("RealFakeBrain", spec.action_spec.empty_action(n_agents - 1))
-    env.set_actions(
-        "RealFakeBrain", -1 * np.ones((n_agents, spec.action_size), dtype=np.float32)
-    )
+    _empty_act = spec.action_spec.empty_action(n_agents)
+    next_action = ActionTuple(_empty_act.continuous - 1, _empty_act.discrete - 1)
+    env.set_actions("RealFakeBrain", next_action)
     env.step()
 
     env.close()

ml-agents-envs/mlagents_envs/tests/test_rpc_utils.py

 from mlagents_envs.communicator_objects.agent_action_pb2 import AgentActionProto
 from mlagents_envs.base_env import (
     BehaviorSpec,
-    ActionType,
+    ActionSpec,
     DecisionSteps,
     TerminalSteps,
 )
 def test_batched_step_result_from_proto():
     n_agents = 10
     shapes = [(3,), (4,)]
-    spec = BehaviorSpec(shapes, ActionType.CONTINUOUS, 3)
+    spec = BehaviorSpec(shapes, ActionSpec.create_continuous(3))
     ap_list = generate_list_agent_proto(n_agents, shapes)
     decision_steps, terminal_steps = steps_from_proto(ap_list, spec)
     for agent_id in range(n_agents):
 def test_action_masking_discrete():
     n_agents = 10
     shapes = [(3,), (4,)]
-    behavior_spec = BehaviorSpec(shapes, ActionType.DISCRETE, (7, 3))
+    behavior_spec = BehaviorSpec(shapes, ActionSpec.create_discrete((7, 3)))
     ap_list = generate_list_agent_proto(n_agents, shapes)
     decision_steps, terminal_steps = steps_from_proto(ap_list, behavior_spec)
     masks = decision_steps.action_mask
 def test_action_masking_discrete_1():
     n_agents = 10
     shapes = [(3,), (4,)]
-    behavior_spec = BehaviorSpec(shapes, ActionType.DISCRETE, (10,))
+    behavior_spec = BehaviorSpec(shapes, ActionSpec.create_discrete((10,)))
     ap_list = generate_list_agent_proto(n_agents, shapes)
     decision_steps, terminal_steps = steps_from_proto(ap_list, behavior_spec)
     masks = decision_steps.action_mask
 def test_action_masking_discrete_2():
     n_agents = 10
     shapes = [(3,), (4,)]
-    behavior_spec = BehaviorSpec(shapes, ActionType.DISCRETE, (2, 2, 6))
+    behavior_spec = BehaviorSpec(shapes, ActionSpec.create_discrete((2, 2, 6)))
     ap_list = generate_list_agent_proto(n_agents, shapes)
     decision_steps, terminal_steps = steps_from_proto(ap_list, behavior_spec)
     masks = decision_steps.action_mask
 def test_action_masking_continuous():
     n_agents = 10
     shapes = [(3,), (4,)]
-    behavior_spec = BehaviorSpec(shapes, ActionType.CONTINUOUS, 10)
+    behavior_spec = BehaviorSpec(shapes, ActionSpec.create_continuous(10))
     ap_list = generate_list_agent_proto(n_agents, shapes)
     decision_steps, terminal_steps = steps_from_proto(ap_list, behavior_spec)
     masks = decision_steps.action_mask
     bp.vector_action_size_deprecated.extend([5, 4])
     bp.vector_action_space_type_deprecated = 0
     behavior_spec = behavior_spec_from_proto(bp, agent_proto)
-    assert behavior_spec.is_action_discrete()
-    assert not behavior_spec.is_action_continuous()
+    assert behavior_spec.action_spec.is_discrete()
+    assert not behavior_spec.action_spec.is_continuous()
-    assert behavior_spec.discrete_action_branches == (5, 4)
-    assert behavior_spec.action_size == 2
+    assert behavior_spec.action_spec.discrete_branches == (5, 4)
+    assert behavior_spec.action_spec.discrete_size == 2
-    assert not behavior_spec.is_action_discrete()
-    assert behavior_spec.is_action_continuous()
-    assert behavior_spec.action_size == 6
+    assert not behavior_spec.action_spec.is_discrete()
+    assert behavior_spec.action_spec.is_continuous()
+    assert behavior_spec.action_spec.continuous_size == 6
-    behavior_spec = BehaviorSpec(shapes, ActionType.CONTINUOUS, 3)
+    behavior_spec = BehaviorSpec(shapes, ActionSpec.create_continuous(3))
     ap_list = generate_list_agent_proto(n_agents, shapes, infinite_rewards=True)
     with pytest.raises(RuntimeError):
         steps_from_proto(ap_list, behavior_spec)
     n_agents = 10
     shapes = [(3,), (4,)]
-    behavior_spec = BehaviorSpec(shapes, ActionType.CONTINUOUS, 3)
+    behavior_spec = BehaviorSpec(shapes, ActionSpec.create_continuous(3))
     ap_list = generate_list_agent_proto(n_agents, shapes, nan_observations=True)
     with pytest.raises(RuntimeError):
         steps_from_proto(ap_list, behavior_spec)

ml-agents-envs/mlagents_envs/tests/test_steps.py

 from mlagents_envs.base_env import (
     DecisionSteps,
     TerminalSteps,
-    ActionType,
+    ActionSpec,
     BehaviorSpec,
 )
 
 
 def test_empty_decision_steps():
     specs = BehaviorSpec(
-        observation_shapes=[(3, 2), (5,)],
-        action_type=ActionType.CONTINUOUS,
-        action_shape=3,
+        observation_shapes=[(3, 2), (5,)], action_spec=ActionSpec.create_continuous(3)
     )
     ds = DecisionSteps.empty(specs)
     assert len(ds.obs) == 2
 
 def test_empty_terminal_steps():
     specs = BehaviorSpec(
-        observation_shapes=[(3, 2), (5,)],
-        action_type=ActionType.CONTINUOUS,
-        action_shape=3,
+        observation_shapes=[(3, 2), (5,)], action_spec=ActionSpec.create_continuous(3)
     )
     ts = TerminalSteps.empty(specs)
     assert len(ts.obs) == 2
 
 def test_specs():
-    specs = BehaviorSpec(
-        observation_shapes=[(3, 2), (5,)],
-        action_type=ActionType.CONTINUOUS,
-        action_shape=3,
-    )
-    assert specs.discrete_action_branches is None
-    assert specs.action_size == 3
-    assert specs.create_empty_action(5).shape == (5, 3)
-    assert specs.create_empty_action(5).dtype == np.float32
+    specs = ActionSpec.create_continuous(3)
+    assert specs.discrete_branches == ()
+    assert specs.discrete_size == 0
+    assert specs.continuous_size == 3
+    assert specs.empty_action(5).continuous.shape == (5, 3)
+    assert specs.empty_action(5).continuous.dtype == np.float32
-    specs = BehaviorSpec(
-        observation_shapes=[(3, 2), (5,)],
-        action_type=ActionType.DISCRETE,
-        action_shape=(3,),
-    )
-    assert specs.discrete_action_branches == (3,)
-    assert specs.action_size == 1
-    assert specs.create_empty_action(5).shape == (5, 1)
-    assert specs.create_empty_action(5).dtype == np.int32
+    specs = ActionSpec.create_discrete((3,))
+    assert specs.discrete_branches == (3,)
+    assert specs.discrete_size == 1
+    assert specs.continuous_size == 0
+    assert specs.empty_action(5).discrete.shape == (5, 1)
+    assert specs.empty_action(5).discrete.dtype == np.int32
+
+    specs = ActionSpec(3, (3,))
+    assert specs.continuous_size == 3
+    assert specs.discrete_branches == (3,)
+    assert specs.discrete_size == 1
+    assert specs.empty_action(5).continuous.shape == (5, 3)
+    assert specs.empty_action(5).continuous.dtype == np.float32
+    assert specs.empty_action(5).discrete.shape == (5, 1)
+    assert specs.empty_action(5).discrete.dtype == np.int32
-    specs = BehaviorSpec(
-        observation_shapes=[(5,)],
-        action_type=ActionType.CONTINUOUS,
-        action_shape=action_len,
-    )
-    zero_action = specs.create_empty_action(4)
+    specs = ActionSpec.create_continuous(action_len)
+    zero_action = specs.empty_action(4).continuous
-    random_action = specs.create_random_action(4)
+    print(specs.random_action(4))
+    random_action = specs.random_action(4).continuous
+    print(random_action)
     assert random_action.dtype == np.float32
     assert random_action.shape == (4, action_len)
     assert np.min(random_action) >= -1
     action_shape = (10, 20, 30)
-    specs = BehaviorSpec(
-        observation_shapes=[(5,)],
-        action_type=ActionType.DISCRETE,
-        action_shape=action_shape,
-    )
-    zero_action = specs.create_empty_action(4)
+    specs = ActionSpec.create_discrete(action_shape)
+    zero_action = specs.empty_action(4).discrete
-    random_action = specs.create_random_action(4)
+    random_action = specs.random_action(4).discrete
     assert random_action.dtype == np.int32
     assert random_action.shape == (4, len(action_shape))
     assert np.min(random_action) >= 0

ml-agents/mlagents/tf_utils/__init__.py

 from mlagents.tf_utils.tf import tf as tf  # noqa
 from mlagents.tf_utils.tf import set_warnings_enabled  # noqa
 from mlagents.tf_utils.tf import generate_session_config  # noqa
+from mlagents.tf_utils.tf import is_available  # noqa

ml-agents/mlagents/tf_utils/tf.py

 # This should be the only place that we import tensorflow directly.
 # Everywhere else is caught by the banned-modules setting for flake8
-import tensorflow as tf  # noqa I201
+
+try:
+    import tensorflow as tf  # noqa I201
-# LooseVersion handles things "1.2.3a" or "4.5.6-rc7" fairly sensibly.
-_is_tensorflow2 = LooseVersion(tf.__version__) >= LooseVersion("2.0.0")
+    # LooseVersion handles things "1.2.3a" or "4.5.6-rc7" fairly sensibly.
+    _is_tensorflow2 = LooseVersion(tf.__version__) >= LooseVersion("2.0.0")
-if _is_tensorflow2:
-    import tensorflow.compat.v1 as tf
+    if _is_tensorflow2:
+        import tensorflow.compat.v1 as tf
-    tf.disable_v2_behavior()
-    tf_logging = tf.logging
-else:
-    try:
-        # Newer versions of tf 1.x will complain that tf.logging is deprecated
-        tf_logging = tf.compat.v1.logging
-    except AttributeError:
-        # Fall back to the safe import, even if it might generate a warning or two.
+        tf.disable_v2_behavior()
+    else:
+        try:
+            # Newer versions of tf 1.x will complain that tf.logging is deprecated
+            tf_logging = tf.compat.v1.logging
+        except AttributeError:
+            # Fall back to the safe import, even if it might generate a warning or two.
+            tf_logging = tf.logging
+except ImportError:
+    tf = None
+
+
+def is_available():
+    """