mlagents-mcp

MCP server for controlling Unity ML-Agents training runs from Claude Code.

Launch, stop, resume, monitor, compare, and export ML-Agents training — all through natural conversation without leaving your editor.

Features

• Training control — start, stop, resume runs as background processes
• Instant checks — check training progress without blocking the conversation
• Live monitoring — read TensorBoard metrics, reward curves, and training logs in real time
• Run comparison — compare metrics across runs for hyperparameter tuning
• Config management — read and deep-merge update YAML training configs
• Model export — locate .onnx models and checkpoints
• Two training modes — Unity Editor (interactive) and built executable (headless batch)

Requirements

• Python 3.10+
• uv package manager
• Unity ML-Agents mlagents-learn available in PATH (or via conda env)
• Claude Code CLI

Installation

git clone https://github.com/limam-B/mlagents-mcp-server.git
cd mlagents-mcp-server
uv sync

Quick setup with Claude Code

# Register the MCP server:
claude mcp add --scope project --transport stdio mlagents-training \
  -- uv run --directory /path/to/mlagents-mcp-server mlagents-mcp

# Unregister (from all scopes to clean up stale configs):
claude mcp remove --scope local mlagents-training
claude mcp remove --scope user mlagents-training
claude mcp remove --scope project mlagents-training

# List registered servers:
claude mcp list

With environment variables

The server reads its configuration from environment variables. Add them to your .mcp.json (project-level) or pass them via the CLI:

{
  "mcpServers": {
    "mlagents-training": {
      "command": "uv",
      "args": ["run", "--directory", "/path/to/mlagents-mcp-server", "mlagents-mcp"],
      "env": {
        "MLAGENTS_PROJECT_ROOT": "/path/to/your/unity/project",
        "MLAGENTS_RESULTS_DIR": "results",
        "MLAGENTS_CONFIG_DIR": "config"
      }
    }
  }
}

With conda (if ML-Agents is installed in a conda env)

{
  "mcpServers": {
    "mlagents-training": {
      "command": "uv",
      "args": ["run", "--directory", "/path/to/mlagents-mcp-server", "mlagents-mcp"],
      "env": {
        "MLAGENTS_PROJECT_ROOT": "/path/to/your/unity/project",
        "MLAGENTS_RESULTS_DIR": "results",
        "MLAGENTS_CONFIG_DIR": "config",
        "MLAGENTS_CONDA_ENV": "mlagents",
        "MLAGENTS_CONDA_PATH": "/home/user/miniconda3"
      }
    }
  }
}

Environment variables

Variable	Default	Description
MLAGENTS_PROJECT_ROOT	.	Root directory of your Unity project
MLAGENTS_RESULTS_DIR	results	Training results directory (relative to project root)
MLAGENTS_CONFIG_DIR	config	Training config YAML directory (relative to project root)
MLAGENTS_CONDA_ENV	—	Conda environment name to activate before running mlagents-learn
MLAGENTS_CONDA_PATH	—	Path to conda installation (e.g. /home/user/miniconda3)

Tools (18)

Training control

Tool	Description
force_training	Launch a new training run (overwrites previous results). Blocks until ready by default.
stop_training	Gracefully stop a run (SIGINT, saves the model).
resume_training	Resume from checkpoint. Auto-reads config from previous run.
cleanup_processes	Kill orphaned mlagents-learn and Unity build processes not tracked by any active run.

Monitoring

Tool	Description
get_run_status	Status overview: reward trend, checkpoints, step progress.
get_metrics	Read TensorBoard scalars (reward, losses, learning rate, etc.).
get_training_logs	Tail live stdout/stderr from an active run.
list_runs	List all known runs with status filtering.

Comparison & export

Tool	Description
compare_runs	Compare a metric across multiple runs (min/max/final + trend).
export_model	Locate .onnx model files and checkpoints.

Configuration

Tool	Description
get_config	Read a YAML training config.
update_config	Deep-merge updates into a config (only specified keys change).

Wait & check

Tool	Default	Description
wait_for_first_metrics	Blocks	Blocks until first TensorBoard data point appears (~1-2 min).
wait_for_completion	Blocks	Blocks until training finishes. For automated run chaining (up to 4 hours).
check_step	Instant	Check if training reached a target step. Returns current progress.
check_reward	Instant	Check if mean reward reached a target. Returns current reward.
check_completion	Instant	Check if training finished. Returns current status and progress.
check_checkpoint	Instant	Check if new .onnx files appeared. Returns checkpoint list.

The check_* tools always return instantly — they never block the conversation. Use wait_for_completion when you want to block until a run finishes (e.g. to chain skill A → skill B automatically).

Two training modes

Editor mode (no env_path)

Training connects to the Unity Editor. force_training blocks until mlagents-learn prints "Listening on port... press Play", then you (or an AI agent) presses Play in Unity.

force_training(config_path="movement.yaml", run_id="Movement_v1")
# → blocks until "Listening on port 5004. Start training by pressing Play..."

Batch mode (with env_path)

Training launches a built executable directly — no Unity Editor needed. force_training blocks until the executable connects.

force_training(
    config_path="movement.yaml",
    run_id="Movement_v1",
    env_path="/path/to/Build.x86_64",
    num_envs=12,
    no_graphics=True,
)
# → blocks until "Connected to Unity environment"

Example workflow

A typical automated training session:

1. force_training(config, run_id, ...)     # launch, blocks until ready
2. wait_for_first_metrics(run_id)          # blocks until data flowing
3. wait_for_completion(run_id)             # blocks until training ends (hours)
4. export_model(run_id)                    # get checkpoint
5. update_config(next_skill, init_path=..) # chain checkpoint to next skill
6. [repeat from step 1 for next skill]

Steps 1-2 block briefly during startup. Step 3 blocks for the full duration (hours) — use this for automated chaining. For manual monitoring, use check_step/check_completion instead of step 3.

Development

# Install dev dependencies:
uv sync --group dev

# Lint:
uv run ruff check src/

# Format:
uv run ruff format src/

# Run the server directly (stdio):
uv run mlagents-mcp

Project structure

src/mlagents_mcp/
  server.py            # FastMCP app, all 18 tool definitions, entry point
  process_manager.py   # Subprocess launch/stop, log capture, port assignment
  metrics_reader.py    # TensorBoard event file parsing
  config_manager.py    # YAML config read/write/deep-merge
  run_registry.py      # Thread-safe run tracking + historical disk scan
  waiters.py           # Blocking wait logic for all wait_for_* tools
  types.py             # Shared dataclasses and enums

License

MIT