ML Agents in Unity

Most game characters follow rules that a developer wrote by hand: if the player is close, chase them, and if health is low, run away. For more complex behaviors such as walking, however, we can let the agents learn on their own, by trial and error. In this post, we train such agents with Unity’s ML-Agents toolkit.

Reinforcement Learning (RL)

ML-Agents is built on reinforcement learning (RL). Over and over, an agent makes observations about its surroundings, takes an action and gets a reward, like +1 for reaching the goal. Its “brain” is a small neural network called the policy, which turns observations into actions. Training is split into episodes, each one attempt at the task. Over thousands of episodes, the policy drifts toward whatever earned the most reward.

The Unity package runs the scene and the agents, while a Python package, mlagents, does the actual learning with PyTorch:

flowchart LR
    Agent[Agent in Unity] -->|observations and rewards| Trainer[mlagents-learn]
    Trainer -->|actions| Agent
    Trainer --> Model[ONNX model]

When training ends, the policy is saved as an .onnx file that Unity can run on its own, so the finished game doesn’t need Python.

Setup

ML-Agents only supports Python 3.10.1 to 3.10.12, so we create a separate environment with conda. The easiest way to set it up is via Miniconda. Once it’s installed, we create an environment for ML-Agents:

conda create -n mlagents python=3.10.12
conda activate mlagents

Next, we clone the ML-Agents repository, which also contains the example environments, and install both Python packages from it:

git clone https://github.com/Unity-Technologies/ml-agents.git
cd ml-agents
python -m pip install ./ml-agents-envs
python -m pip install ./ml-agents

The repository also includes a Unity project in its Project folder, with all the examples and the ML-Agents package already linked. We open it from Unity Hub with Add > Add project from disk. It was made for Unity 6.0, but it upgrades to 6.5 without any problems.

Testing the Examples

The examples live in Assets/ML-Agents/Examples, and each one ships with a pre-trained model. To train one yourself, start the trainer from the repository folder with the example’s configuration file:

mlagents-learn config/ppo/3DBall.yaml --run-id=3DBall --force

When it asks you to, press Play in Unity and the agents start learning. ppo stands for Proximal Policy Optimization, the default training algorithm, and --run-id names the folder in results/ where the trained model ends up.

Basic

Basic is the simplest example. The agent stands on a line and can step left, step right or stay put. A small goal 3 steps away is worth +0.1, and a large one 7 steps away in the other direction is worth +1. Every step costs 0.01, so the best strategy is to walk straight to the large goal, for a total of 0.93.

Start the training and run until the agent has figured it out. TensorBoard, which comes with ML-Agents, shows the training statistics as graphs at localhost:6006:

tensorboard --logdir results

The two most important graphs are Environment/Cumulative Reward, the average reward per episode, and Environment/Episode Length, the average number of steps per episode. As the agent learns, the reward rises and then levels off, and episodes usually get shorter. The Losses and Policy graphs show the trainer’s internals. For example, Entropy shows how random the actions still are and slowly drops.

At first, the agent wanders at random and usually ends up at the nearby small goal. Eventually, it learns to go for the large one instead, ending close to the best possible score. Once trained, we can drag Basic.onnx into the Model field of the agent’s Behavior Parameters so it can play without the trainer.

Building Our Own Environment

Now we build an environment from scratch, based on the RollerBall tutorial: a ball that learns to roll to a target.

The Scene

We start with a new Universal 3D project, as in Setting Up Unity, and install ML-Agents through the Package Manager with + > Add package by name and com.unity.ml-agents.

The scene is a 10 × 10 floor with four walls around it, a green cube as the target and a blue sphere with a Rigidbody as the agent.

The Agent

An agent is a script that inherits from Agent instead of MonoBehaviour. Here is the full RollerAgent.cs:

using UnityEngine;
using UnityEngine.InputSystem;
using Unity.MLAgents;
using Unity.MLAgents.Sensors;
using Unity.MLAgents.Actuators;

public class RollerAgent : Agent
{
    [SerializeField] private Transform target;
    [SerializeField] private float forceMultiplier = 5f;

    private Rigidbody rb;
    // A 7 x 7 area in the middle of the floor (Bounds takes the full size, not half)
    private Bounds spawnBounds = new Bounds(Vector3.zero, new Vector3(7f, 0f, 7f));

    public override void Initialize()
    {
        rb = GetComponent<Rigidbody>();
    }

    public override void OnEpisodeBegin()
    {
        // Don't carry speed over from the last episode
        rb.linearVelocity = Vector3.zero;
        rb.angularVelocity = Vector3.zero;

        // Local positions are relative to TrainingArea, so copies of the area work too
        transform.localPosition = GetRandomPosition(spawnBounds);
        target.localPosition = GetRandomPosition(spawnBounds);
    }

    public override void CollectObservations(VectorSensor sensor)
    {
        // 3 + 3 + 1 + 1 = 8 values, matching the observation size in the Inspector
        sensor.AddObservation(target.localPosition);
        sensor.AddObservation(transform.localPosition);
        sensor.AddObservation(rb.linearVelocity.x);
        sensor.AddObservation(rb.linearVelocity.z);
    }

    public override void OnActionReceived(ActionBuffers actions)
    {
        // Two continuous actions between -1 and 1: how hard to push along x and z
        float moveX = actions.ContinuousActions[0];
        float moveZ = actions.ContinuousActions[1];
        rb.AddForce(new Vector3(moveX, 0f, moveZ) * forceMultiplier);

        // Close enough to count as touching the cube
        if (Vector3.Distance(transform.localPosition, target.localPosition) < 1.42f)
        {
            SetReward(1f);
            EndEpisode();
        }
    }

    public override void Heuristic(in ActionBuffers actionsOut)
    {
        // Each key pair gives -1, 0 or 1, the same range as the policy's actions
        var continuousActions = actionsOut.ContinuousActions;
        Keyboard keyboard = Keyboard.current;
        continuousActions[0] = (keyboard.dKey.isPressed ? 1f : 0f) - (keyboard.aKey.isPressed ? 1f : 0f);
        continuousActions[1] = (keyboard.wKey.isPressed ? 1f : 0f) - (keyboard.sKey.isPressed ? 1f : 0f);
    }

    private Vector3 GetRandomPosition(Bounds bounds)
    {
        float randomX = Random.Range(bounds.min.x, bounds.max.x);
        float randomY = Random.Range(bounds.min.y, bounds.max.y);
        float randomZ = Random.Range(bounds.min.z, bounds.max.z);
        return new Vector3(randomX, randomY, randomZ);
    }
}

• Initialize runs once when the agent starts, like Start in a normal script, and gets the ball’s Rigidbody.
• OnEpisodeBegin stops the ball and places it and the target at random spots. Without the randomness, the agent would just memorize a single path.
• CollectObservations gives the policy 8 numbers: the target’s position, the agent’s position and its speed along the floor.
• OnActionReceived pushes the ball with the 2 numbers it gets from the policy. Getting close to the target earns +1 and ends the episode.
• Heuristic lets a human drive the ball with W, A, S and D, to test the environment before training.

In the Inspector, the agent’s Behavior Parameters get the Behavior Name RollerBall, an observation size of 8, 2 Continuous Actions and no Discrete Branches. A continuous action is a number between -1 and 1, useful when the agent needs to choose an amount, like how much force to apply. A discrete branch is a choice from a fixed list, like left, right or stay for the Basic agent.

A Decision Requester with a period of 10 lets the policy decide every 10 physics steps. Without one, the agent never asks for an action and just sits there. We also set Max Step to 1000, so episodes end even if it never finds the target.

Behavior Type decides who controls the agent. Default trains when mlagents-learn is running and otherwise uses the model, Heuristic Only hands control to the Heuristic method, which here means driving the ball yourself with the keyboard, and Inference Only always uses the model. It’s worth driving the ball with Heuristic Only first, then switching back to Default for training.

Training

The training settings go in a YAML file. We use these values:

behaviors:
  RollerBall:
    trainer_type: ppo
    hyperparameters:
      batch_size: 10
      buffer_size: 100
      learning_rate: 3.0e-4
      beta: 5.0e-4
      epsilon: 0.2
      lambd: 0.99
      num_epoch: 3
      learning_rate_schedule: linear
      beta_schedule: constant
      epsilon_schedule: linear
    network_settings:
      normalize: false
      hidden_units: 128
      num_layers: 2
    reward_signals:
      extrinsic:
        gamma: 0.99
        strength: 1.0
    max_steps: 500000
    time_horizon: 64
    summary_freq: 2000

• buffer_size and batch_size: the trainer collects 100 steps of experience, then updates the policy in chunks of 10. Small values mean frequent updates, which suits a simple task.
• learning_rate, beta and epsilon: how big each update is, how much the agent keeps trying random actions, and how far the policy may change in one update. The *_schedule settings either shrink a value over the run (linear) or keep it fixed (constant).
• network_settings: the policy is a network with 2 layers of 128 neurons.
• gamma: how much future rewards count compared to immediate ones.
• max_steps and summary_freq: training stops after 500,000 steps and logs statistics every 2,000.

Then it’s the same routine as before: start the trainer, then press Play.

mlagents-learn rollerball_config.yaml --run-id=RollerBall

During training, mlagents-learn runs Unity at 20 times normal speed, so episodes go by in a fraction of a second:

To stop early, press Ctrl+C in the terminal first, then Stop in Unity. For me, the full 500,000 steps took about 30 minutes:

Steps	Mean reward	Mean episode length
2,000	0.39	83.0
10,000	0.68	58.8
20,000	1.00	12.9
30,000	0.99	9.1
500,000	1.00	4.9

Episode length is counted in decisions, so a length of 100 means the ball ran out of time before reaching the cube. At first, the ball rolls around at random and finds the cube in fewer than half of the episodes. After about 20,000 steps, it reaches the cube every time, and from then on it only gets faster. Nothing improves after the first 50,000 steps or so, so a much smaller max_steps would have been enough.

Once training is done, the model ends up in results/RollerBall/RollerBall.onnx. Unity can only use files inside the project, so we first copy it into Assets. As with Basic, dragging it into the Model field of the Behavior Parameters lets the ball chase the target on its own.

Conclusion

In this post, we went from installing ML-Agents to training an example agent and a ball of our own that finds its target in about a second. Most of the agent’s behavior doesn’t come from its script but from the reward. The same ideas can be applied to create much more complex and interesting agents.