Controlling Unity ML-agent environment Scenes for different experiments with python

Suppose, our environment has multiple scenes to support different experiments in the same environment. Based on different configurations we want to load different scenes. One simple approach to handle this requirement is to Build different executables for different scenes.

But I wanted to keep one single executable and control it programmatically from python script. So far, I couldn’t find a direct way to do this. Instead I used another API called ‘Side Channels’ provided by Unity ML-agent. Details about Side Channels can be found here. Basically side channels provide a communication channel. When the simulation/training is running, we can pass data between unity environment and python script. We can also create custom channels if we need more control following this.

Let’s start with implementations in unity-

Unity Side implementations

In unity, add a new script (Right click -> Create -> C# Script), ang give it a name. I name it EnvironmentController. Attach the script to any object of the environment, or add it to any empty object, so that the script is executed when the game/environment runs.

Let’s just first define an enum with the scene/experiments-

enum Experiment
{
    Main = 1,
    PaperRod = 2
}

Now we use the Awake() method of the MonoBehaviour script to get access to the Float Properties channel using the Academy instance (a Singleton class) like below. We use the Awake() method to fetch the scene before the script starts running.

private float scene;
public void Awake()
{
    var floatChannel = Academy.Instance.FloatProperties;
    scene = floatChannel.GetPropertyWithDefault("sceneToLoad", (float) Experiment.Main);
}

From the floatChannel we can retrieve the data that was passed through python API. This works like a key-value dictionary with a default value. We are retrieving the scene no. with ‘sceneToLoad’ key, and setting a default to our Main scene enum value, in case there is no value passed from python API.

Now that we have the scene no, we use the Start() method to load the scene using Unity’s SceneManager.

private void Start()
{
    if (scene == (float) Experiment.PaperRod && !SceneManager.GetActiveScene().name.Equals("Paper_Rod_Experiment"))
    {
        Academy.Instance.Dispose();
        SceneManager.LoadScene("Paper_Rod_Experiment", LoadSceneMode.Single);
    }
    else if (scene == (float)Experiment.Main && !SceneManager.GetActiveScene().name.Equals("MainScene"))
    {
        Academy.Instance.Dispose();
        SceneManager.LoadScene("MainScene", LoadSceneMode.Single);
    }
}

Here we are just checking if the current scene is not the one we want to change with SceneManager.GetActiveScene(). If it is a different scene we are using the SceneManager.LoadScene method to load the scene with the exact scene names.

Note: It is important to call ‘Academy.Instance.Dispose()’. This ensures that the previous scene is cleared and the agent ‘Brain’ is reloaded and connected to the agent. Without this the agent doesn’t receive any call to AgentAction method.

That’s it. Easy way to change the scene from python API. Now, we just need to export the unity app. Be sure to include all your scenes when exporting through the build setting window (File->Build Settings). To add scenes to the exported app/binary, add all the scenes (Scenes need to be added in the hierarchy window.

Then build the app binary.

Python Side Implementation

Similar to Unity first we will define some enum for the scenes in python side also-

import enum

class Experiments(enum.Enum):
    Main = 1
    PaperRod = 2

To pass the controlling data through python APIs and FloatPropertiesChannel we use

channel = FloatPropertiesChannel()
channel.set_property("sceneToLoad", float(Experiments.Main.value))

We initialize a FloatPropertiesChannel object. Then we set the scene value with key ‘sceneToLoad’ exactly the same as we used in unity side.

Now we will pass this channel in the environment creation step-

env = UnityEnvironment(base_port = 5005, file_name=env_name, \
                  side_channels = [engine_configuration_channel, channel])

We pass our channel as the side_channels argument. To learn about the details on how to start the environment with ML-agent follow this documentations.

Here is the output of two enum (Main, PaperRod) values passed as FloatPropertiesChannel.

Hope this helps someone.