feat(*): Add self-play RL training pipeline with PPO trainer, in-game GDScript policy inference, and bot opponent support in Match mode

.claude/agents/godot.md

@@ -0,0 +1,15 @@
+---
+name: godot
+description: Runs all godot-mcp tool interactions (reading/editing scenes, running the project, in-game verification, screenshots, debug output) on Haiku to keep verbose Godot tool outputs out of the main model's context and reduce cost. Give it a concrete, self-contained task and tell it exactly what to report back.
+tools: mcp__godot-mcp__*, Read, Grep, Glob, Bash
+model: haiku
+---
+
+You operate the Godot project at `Game/` (repo root: the parent directory) via the godot-mcp tools. Prefer godot-mcp tools over shell commands or manual file parsing for anything involving scenes, nodes, scripts, the input map, or the running game.
+
+Ground rules:
+- The Godot executable is at `/Applications/Godot.app/Contents/MacOS/Godot` if you need the CLI (e.g. headless runs); otherwise use the MCP tools.
+- `get_debug_output` and `stop_project` can return enormous logs — never dump them into your report. Grep/filter for the relevant lines and quote only those.
+- When verifying gameplay, follow the existing patterns: interact via input actions (`game_key_press` with action names from `project.godot`), inspect state with `game_get_property`/`game_get_scene_tree`, and check `game_get_errors`.
+- Always stop a running project before finishing unless told otherwise.
+- Report back exactly what was asked: the outcome, key evidence (short quotes, property values), and any errors — not raw tool output.

.gitignore

@@ -1 +1,8 @@
-.DS_Store
+.DS_Store
+
+# RL training artifacts (training/ code is committed; outputs are not)
+training/.venv/
+training/logs/
+training/checkpoints/
+training/smoke_run.log
+training/__pycache__/

CLAUDE.md

@@ -47,3 +47,13 @@ The structure was deliberately chosen so an RL-trained AI opponent and, later, m
 - **HUD / telemetry pattern**: `Ship` emits flight data via signals only when values change past thresholds (`_last_*` fields, `*_THRESHOLD` constants). `HUDController` (`scripts/HUDController.gd` on `scenes/HUD.tscn`, instanced by each mode's scene) discovers the ship, camera rig, and game mode via groups (`"ship"`, `"ship_camera"`, `"game"`), connects to signals, and only updates labels — no polling. Follow this discovery-by-group + signal-push pattern for new instruments or cross-node communication, not hardcoded `get_node` paths or per-frame polling.
 - **Input actions** are defined in `Game/project.godot` under `[input]` (`move_forward`, `turn_left`, `turbo`, `reset_ball`, etc.) and read only by `PlayerShipController` (plus mode-level `_unhandled_input` for `reset_ball`/`ui_cancel`) — add new controls there rather than hardcoding key checks.
 - Physics engine is Jolt (`Game/project.godot`, `[physics] 3d/physics_engine="Jolt Physics"`).
+
+## Reinforcement learning / AI bots
+
+See `TRAINING.md` for the full workflow (training, exporting, evaluating, difficulty tiers). Architecture summary:
+
+- `scenes/training.tscn` (`scripts/training_mode.gd`, extends `GameMode`) is the headless self-play environment: two ships driven by `RLShipController`s, with `ShipAIController` (extends the vendored plugin's `AIController3D`) as the only class touching godot_rl types. The plugin is vendored (not a submodule) at `Game/addons/godot_rl_agents` — see its `VENDORED.md`; its C#/ONNX files are unused.
+- `scripts/ship_observations.gd` is the shared observation builder used by both training and in-game inference — never fork or diverge these two paths. Team 1's observations are mirrored (180° about Y) so one policy plays both sides.
+- In-game bots: `scripts/ai_ship_controller.gd` (a `ShipController`) runs the exported policy JSON via `scripts/policy_network.gd` (pure-GDScript MLP) — no .NET/ONNX/Python at runtime. Models live in `Game/bots/`; Match mode's `bot_model_path`/`bot_reaction_ticks`/`bot_action_noise` exports configure the opponent.
+- Python side lives in `training/` (venv, not committed): `train.py` (SB3 PPO, launches parallel headless Godot instances from source), `export_policy.py` (checkpoint → JSON with parity check), `evaluate.py` (head-to-head eval, appends `training/eval_history.json`).
+- The flattened action space is Box(7): thrust xyz, rotation xyz, turbo (>0 = on) — this is `ShipAction` verbatim; change either only deliberately and together.

Game/addons/godot_rl_agents/LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 Edward Beeching
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE. 

Game/addons/godot_rl_agents/VENDORED.md

@@ -0,0 +1,3 @@
+Vendored from https://github.com/edbeeching/godot_rl_agents_plugin
+commit 998c357a0cd09b37f40a36d70c7867fc9f682338 (2026-06-13), MIT license (see LICENSE).
+The onnx/csharp files are unused (require the .NET Godot build); in-game inference uses scripts/policy_network.gd instead.

Game/addons/godot_rl_agents/controller/ai_controller_2d.gd

@@ -0,0 +1,136 @@
+extends Node2D
+class_name AIController2D
+
+enum ControlModes {
+	INHERIT_FROM_SYNC, ## Inherit setting from sync node
+	HUMAN, ## Test the environment manually
+	TRAINING, ## Train a model
+	ONNX_INFERENCE, ## Load a pretrained model using an .onnx file
+	RECORD_EXPERT_DEMOS ## Record observations and actions for expert demonstrations
+}
+@export var control_mode: ControlModes = ControlModes.INHERIT_FROM_SYNC
+## The path to a trained .onnx model file to use for inference (overrides the path set in sync node).
+@export var onnx_model_path := ""
+## Once the number of steps has passed, the flag 'needs_reset' will be set to 'true' for this instance.
+@export var reset_after := 1000
+
+@export_group("Record expert demos mode options")
+## Path where the demos will be saved. The file can later be used for imitation learning.
+@export var expert_demo_save_path: String
+## The action that erases the last recorded episode from the currently recorded data.
+@export var remove_last_episode_key: InputEvent
+## Action will be repeated for n frames. Will introduce control lag if larger than 1.
+## Can be used to ensure that action_repeat on inference and training matches
+## the recorded demonstrations.
+@export var action_repeat: int = 1
+
+@export_group("Multi-policy mode options")
+## Allows you to set certain agents to use different policies.
+## Changing has no effect with default SB3 training. Works with Rllib example.
+## Tutorial: https://github.com/edbeeching/godot_rl_agents/blob/main/docs/TRAINING_MULTIPLE_POLICIES.md
+@export var policy_name: String = "shared_policy"
+
+var onnx_model: ONNXModel
+
+var heuristic := "human"
+var done := false
+var reward := 0.0
+var n_steps := 0
+var needs_reset := false
+
+var _player: Node2D
+
+
+func _ready():
+	add_to_group("AGENT")
+
+
+func init(player: Node2D):
+	_player = player
+
+
+#region Methods that need implementing using the "extend script" option in Godot
+func get_obs() -> Dictionary:
+	assert(false, "the get_obs method is not implemented when extending from ai_controller")
+	return {"obs": []}
+
+
+func get_reward() -> float:
+	assert(false, "the get_reward method is not implemented when extending from ai_controller")
+	return 0.0
+
+
+func get_action_space() -> Dictionary:
+	assert(
+		false, "the get_action_space method is not implemented when extending from ai_controller"
+	)
+	return {
+		"example_actions_continous": {"size": 2, "action_type": "continuous"},
+		"example_actions_discrete": {"size": 2, "action_type": "discrete"},
+	}
+
+
+func set_action(action) -> void:
+	assert(false, "the set_action method is not implemented when extending from ai_controller")
+
+
+#endregion
+
+
+#region Methods that sometimes need implementing using the "extend script" option in Godot
+# Only needed if you are recording expert demos with this AIController
+func get_action() -> Array:
+	assert(
+		false,
+		"the get_action method is not implemented in extended AIController but demo_recorder is used"
+	)
+	return []
+
+
+# For providing additional info (e.g. `is_success` for SB3 training)
+func get_info() -> Dictionary:
+	return {}
+
+
+#endregion
+
+
+func _physics_process(delta):
+	n_steps += 1
+	if n_steps > reset_after:
+		needs_reset = true
+
+
+func get_obs_space():
+	# may need overriding if the obs space is complex
+	var obs = get_obs()
+	return {
+		"obs": {"size": [len(obs["obs"])], "space": "box"},
+	}
+
+
+func reset():
+	n_steps = 0
+	needs_reset = false
+
+
+func reset_if_done():
+	if done:
+		reset()
+
+
+func set_heuristic(h):
+	# sets the heuristic from "human" or "model" nothing to change here
+	heuristic = h
+
+
+func get_done():
+	return done
+
+
+func set_done_false():
+	done = false
+
+
+func zero_reward():
+	reward = 0.0

Game/addons/godot_rl_agents/controller/ai_controller_2d.gd.uid

@@ -0,0 +1 @@
+uid://dgfbh3y7s07g7

Game/addons/godot_rl_agents/controller/ai_controller_3d.gd

@@ -0,0 +1,136 @@
+extends Node3D
+class_name AIController3D
+
+enum ControlModes {
+	INHERIT_FROM_SYNC, ## Inherit setting from sync node
+	HUMAN, ## Test the environment manually
+	TRAINING, ## Train a model
+	ONNX_INFERENCE, ## Load a pretrained model using an .onnx file
+	RECORD_EXPERT_DEMOS ## Record observations and actions for expert demonstrations
+}
+@export var control_mode: ControlModes = ControlModes.INHERIT_FROM_SYNC
+## The path to a trained .onnx model file to use for inference (overrides the path set in sync node).
+@export var onnx_model_path := ""
+## Once the number of steps has passed, the flag 'needs_reset' will be set to 'true' for this instance.
+@export var reset_after := 1000
+
+@export_group("Record expert demos mode options")
+## Path where the demos will be saved. The file can later be used for imitation learning.
+@export var expert_demo_save_path: String
+## The action that erases the last recorded episode from the currently recorded data.
+@export var remove_last_episode_key: InputEvent
+## Action will be repeated for n frames. Will introduce control lag if larger than 1.
+## Can be used to ensure that action_repeat on inference and training matches
+## the recorded demonstrations.
+@export var action_repeat: int = 1
+
+@export_group("Multi-policy mode options")
+## Allows you to set certain agents to use different policies.
+## Changing has no effect with default SB3 training. Works with Rllib example.
+## Tutorial: https://github.com/edbeeching/godot_rl_agents/blob/main/docs/TRAINING_MULTIPLE_POLICIES.md
+@export var policy_name: String = "shared_policy"
+
+var onnx_model: ONNXModel
+
+var heuristic := "human"
+var done := false
+var reward := 0.0
+var n_steps := 0
+var needs_reset := false
+
+var _player: Node3D
+
+
+func _ready():
+	add_to_group("AGENT")
+
+
+func init(player: Node3D):
+	_player = player
+
+
+#region Methods that need implementing using the "extend script" option in Godot
+func get_obs() -> Dictionary:
+	assert(false, "the get_obs method is not implemented when extending from ai_controller")
+	return {"obs": []}
+
+
+func get_reward() -> float:
+	assert(false, "the get_reward method is not implemented when extending from ai_controller")
+	return 0.0
+
+
+func get_action_space() -> Dictionary:
+	assert(
+		false, "the get_action_space method is not implemented when extending from ai_controller"
+	)
+	return {
+		"example_actions_continous": {"size": 2, "action_type": "continuous"},
+		"example_actions_discrete": {"size": 2, "action_type": "discrete"},
+	}
+
+
+func set_action(action) -> void:
+	assert(false, "the set_action method is not implemented when extending from ai_controller")
+
+
+#endregion
+
+
+#region Methods that sometimes need implementing using the "extend script" option in Godot
+# Only needed if you are recording expert demos with this AIController
+func get_action() -> Array:
+	assert(
+		false,
+		"the get_action method is not implemented in extended AIController but demo_recorder is used"
+	)
+	return []
+
+
+# For providing additional info (e.g. `is_success` for SB3 training)
+func get_info() -> Dictionary:
+	return {}
+
+
+#endregion
+
+
+func _physics_process(delta):
+	n_steps += 1
+	if n_steps > reset_after:
+		needs_reset = true
+
+
+func get_obs_space():
+	# may need overriding if the obs space is complex
+	var obs = get_obs()
+	return {
+		"obs": {"size": [len(obs["obs"])], "space": "box"},
+	}
+
+
+func reset():
+	n_steps = 0
+	needs_reset = false
+
+
+func reset_if_done():
+	if done:
+		reset()
+
+
+func set_heuristic(h):
+	# sets the heuristic from "human" or "model" nothing to change here
+	heuristic = h
+
+
+func get_done():
+	return done
+
+
+func set_done_false():
+	done = false
+
+
+func zero_reward():
+	reward = 0.0

Game/addons/godot_rl_agents/controller/ai_controller_3d.gd.uid

@@ -0,0 +1 @@
+uid://qm8wpg7ccydk

Game/addons/godot_rl_agents/godot_rl_agents.gd

@@ -0,0 +1,16 @@
+@tool
+extends EditorPlugin
+
+
+func _enter_tree():
+	# Initialization of the plugin goes here.
+	# Add the new type with a name, a parent type, a script and an icon.
+	add_custom_type("Sync", "Node", preload("sync.gd"), preload("icon.png"))
+	#add_custom_type("RaycastSensor2D2", "Node", preload("raycast_sensor_2d.gd"), preload("icon.png"))
+
+
+func _exit_tree():
+	# Clean-up of the plugin goes here.
+	# Always remember to remove it from the engine when deactivated.
+	remove_custom_type("Sync")
+	#remove_custom_type("RaycastSensor2D2")

Game/addons/godot_rl_agents/godot_rl_agents.gd.uid

@@ -0,0 +1 @@
+uid://bdnkavfpf6ge6

Game/addons/godot_rl_agents/icon.png.import

@@ -0,0 +1,40 @@
+[remap]
+
+importer="texture"
+type="CompressedTexture2D"
+uid="uid://bxy3je5atsh68"
+path="res://.godot/imported/icon.png-45a871b53434e556222f5901d598ab34.ctex"
+metadata={
+"vram_texture": false
+}
+
+[deps]
+
+source_file="res://addons/godot_rl_agents/icon.png"
+dest_files=["res://.godot/imported/icon.png-45a871b53434e556222f5901d598ab34.ctex"]
+
+[params]
+
+compress/mode=0
+compress/high_quality=false
+compress/lossy_quality=0.7
+compress/uastc_level=0
+compress/rdo_quality_loss=0.0
+compress/hdr_compression=1
+compress/normal_map=0
+compress/channel_pack=0
+mipmaps/generate=false
+mipmaps/limit=-1
+roughness/mode=0
+roughness/src_normal=""
+process/channel_remap/red=0
+process/channel_remap/green=1
+process/channel_remap/blue=2
+process/channel_remap/alpha=3
+process/fix_alpha_border=true
+process/premult_alpha=false
+process/normal_map_invert_y=false
+process/hdr_as_srgb=false
+process/hdr_clamp_exposure=false
+process/size_limit=0
+detect_3d/compress_to=1

Game/addons/godot_rl_agents/onnx/csharp/ONNXInference.cs

@@ -0,0 +1,115 @@
+using Godot;
+using Microsoft.ML.OnnxRuntime;
+using Microsoft.ML.OnnxRuntime.Tensors;
+using System.Collections.Generic;
+using System.Linq;
+
+namespace GodotONNX
+{
+	/// <include file='docs/ONNXInference.xml' path='docs/members[@name="ONNXInference"]/ONNXInference/*'/>
+	public partial class ONNXInference : GodotObject
+	{
+
+		private InferenceSession session;
+		/// <summary>
+		/// Path to the ONNX model. Use Initialize to change it. 
+		/// </summary>
+		private string modelPath;
+		private int batchSize;
+
+		private SessionOptions SessionOpt;
+
+		/// <summary>
+		/// init function
+		/// </summary>
+		/// <param name="Path"></param>
+		/// <param name="BatchSize"></param>
+		/// <returns>Returns the output size of the model</returns>
+		public int Initialize(string Path, int BatchSize)
+		{
+			modelPath = Path;
+			batchSize = BatchSize;
+			SessionOpt = SessionConfigurator.MakeConfiguredSessionOptions();
+			session = LoadModel(modelPath);
+			return session.OutputMetadata["output"].Dimensions[1];
+		}
+
+
+		/// <include file='docs/ONNXInference.xml' path='docs/members[@name="ONNXInference"]/Run/*'/>
+		public Godot.Collections.Dictionary<string, Godot.Collections.Array<float>> RunInference(Godot.Collections.Dictionary<string, Godot.Collections.Array<float>> obs, int state_ins)
+		{
+			//Current model: Any (Godot Rl Agents)
+			//Expects a tensor of shape [batch_size, input_size] type float for any output of the agents observation dictionary and a tensor of shape [batch_size] type float named state_ins
+
+			var modelInputsList = new List<NamedOnnxValue>
+			{
+				NamedOnnxValue.CreateFromTensor("state_ins", new DenseTensor<float>(new float[] { state_ins }, new int[] { batchSize }))
+			};
+			foreach (var key in obs.Keys)
+			{
+				var subObs = obs[key];
+				// Fill the input tensors for each key of the observation
+				// create span of observation from specific inputSize
+				var obsData = new float[subObs.Count]; //There's probably a better way to do this
+				for (int i = 0; i < subObs.Count; i++)
+				{
+					obsData[i] = subObs[i];
+				}
+				modelInputsList.Add(
+					NamedOnnxValue.CreateFromTensor(key, new DenseTensor<float>(obsData, new int[] { batchSize, subObs.Count }))
+				);
+			}
+
+			IReadOnlyCollection<string> outputNames = new List<string> { "output", "state_outs" }; //ONNX is sensible to these names, as well as the input names
+
+			IDisposableReadOnlyCollection<DisposableNamedOnnxValue> results;
+			//We do not use "using" here so we get a better exception explaination later
+			try
+			{
+				results = session.Run(modelInputsList, outputNames);
+			}
+			catch (OnnxRuntimeException e)
+			{
+				//This error usually means that the model is not compatible with the input, beacause of the input shape (size)
+				GD.Print("Error at inference: ", e);
+				return null;
+			}
+			//Can't convert IEnumerable<float> to Variant, so we have to convert it to an array or something
+			Godot.Collections.Dictionary<string, Godot.Collections.Array<float>> output = new Godot.Collections.Dictionary<string, Godot.Collections.Array<float>>();
+			DisposableNamedOnnxValue output1 = results.First();
+			DisposableNamedOnnxValue output2 = results.Last();
+			Godot.Collections.Array<float> output1Array = new Godot.Collections.Array<float>();
+			Godot.Collections.Array<float> output2Array = new Godot.Collections.Array<float>();
+
+			foreach (float f in output1.AsEnumerable<float>())
+			{
+				output1Array.Add(f);
+			}
+
+			foreach (float f in output2.AsEnumerable<float>())
+			{
+				output2Array.Add(f);
+			}
+
+			output.Add(output1.Name, output1Array);
+			output.Add(output2.Name, output2Array);
+
+			//Output is a dictionary of arrays, ex: { "output" : [0.1, 0.2, 0.3, 0.4, ...], "state_outs" : [0.5, ...]}
+			results.Dispose();
+			return output;
+		}
+		/// <include file='docs/ONNXInference.xml' path='docs/members[@name="ONNXInference"]/Load/*'/>
+		public InferenceSession LoadModel(string Path)
+		{
+			using Godot.FileAccess file = FileAccess.Open(Path, Godot.FileAccess.ModeFlags.Read);
+			byte[] model = file.GetBuffer((int)file.GetLength());
+			//file.Close(); file.Dispose(); //Close the file, then dispose the reference.
+			return new InferenceSession(model, SessionOpt); //Load the model
+		}
+		public void FreeDisposables()
+		{
+			session.Dispose();
+			SessionOpt.Dispose();
+		}
+	}
+}

Game/addons/godot_rl_agents/onnx/csharp/SessionConfigurator.cs

@@ -0,0 +1,131 @@
+using Godot;
+using Microsoft.ML.OnnxRuntime;
+
+namespace GodotONNX
+{
+	/// <include file='docs/SessionConfigurator.xml' path='docs/members[@name="SessionConfigurator"]/SessionConfigurator/*'/>
+
+	public static class SessionConfigurator
+	{
+		public enum ComputeName
+		{
+			CUDA,
+			ROCm,
+			DirectML,
+			CoreML,
+			CPU
+		}
+
+        /// <include file='docs/SessionConfigurator.xml' path='docs/members[@name="SessionConfigurator"]/GetSessionOptions/*'/>
+        public static SessionOptions MakeConfiguredSessionOptions()
+        {
+            SessionOptions sessionOptions = new();
+            SetOptions(sessionOptions);
+            return sessionOptions;
+        }
+
+        private static void SetOptions(SessionOptions sessionOptions)
+        {
+            sessionOptions.LogSeverityLevel = OrtLoggingLevel.ORT_LOGGING_LEVEL_WARNING;
+            ApplySystemSpecificOptions(sessionOptions);
+        }
+
+        /// <include file='docs/SessionConfigurator.xml' path='docs/members[@name="SessionConfigurator"]/SystemCheck/*'/>
+        static public void ApplySystemSpecificOptions(SessionOptions sessionOptions)
+        {
+            //Most code for this function is verbose only, the only reason it exists is to track
+            //implementation progress of the different compute APIs.
+
+            //December 2022: CUDA is not working. 
+
+            string OSName = OS.GetName(); //Get OS Name
+
+            //ComputeName ComputeAPI = ComputeCheck(); //Get Compute API
+            //                                         //TODO: Get CPU architecture
+
+            //Linux can use OpenVINO (C#) on x64 and ROCm on x86 (GDNative/C++)
+            //Windows can use OpenVINO (C#) on x64
+            //TODO: try TensorRT instead of CUDA
+            //TODO: Use OpenVINO for Intel Graphics
+
+            // Temporarily using CPU on all platforms to avoid errors detected with DML
+            ComputeName ComputeAPI = ComputeName.CPU;
+
+            //match OS and Compute API
+            GD.Print($"OS: {OSName} Compute API: {ComputeAPI}");
+
+            // CPU is set by default without appending necessary
+            // sessionOptions.AppendExecutionProvider_CPU(0);
+
+            /*
+            switch (OSName)
+            {
+                case "Windows": //Can use CUDA, DirectML
+                    if (ComputeAPI is ComputeName.CUDA)
+                    {
+                        //CUDA 
+                        //sessionOptions.AppendExecutionProvider_CUDA(0);
+                        //sessionOptions.AppendExecutionProvider_DML(0);
+                    }
+                    else if (ComputeAPI is ComputeName.DirectML)
+                    {
+                        //DirectML
+                        //sessionOptions.AppendExecutionProvider_DML(0);
+                    }
+                    break;
+                case "X11": //Can use CUDA, ROCm
+                    if (ComputeAPI is ComputeName.CUDA)
+                    {
+                        //CUDA
+                        //sessionOptions.AppendExecutionProvider_CUDA(0);
+                    }
+                    if (ComputeAPI is ComputeName.ROCm)
+                    {
+                        //ROCm, only works on x86 
+                        //Research indicates that this has to be compiled as a GDNative plugin
+                        //GD.Print("ROCm not supported yet, using CPU.");
+                        //sessionOptions.AppendExecutionProvider_CPU(0);
+                    }
+                    break;
+                case "macOS": //Can use CoreML
+                    if (ComputeAPI is ComputeName.CoreML)
+                    { //CoreML
+                      //TODO: Needs testing
+                        //sessionOptions.AppendExecutionProvider_CoreML(0);
+                        //CoreML on ARM64, out of the box, on x64 needs .tar file from GitHub
+                    }
+                    break;
+                default:
+                    GD.Print("OS not Supported.");
+                    break;
+            }
+            */
+        }
+
+
+        /// <include file='docs/SessionConfigurator.xml' path='docs/members[@name="SessionConfigurator"]/ComputeCheck/*'/>
+        public static ComputeName ComputeCheck()
+		{
+			string adapterName = Godot.RenderingServer.GetVideoAdapterName();
+			//string adapterVendor = Godot.RenderingServer.GetVideoAdapterVendor();
+			adapterName = adapterName.ToUpper(new System.Globalization.CultureInfo(""));
+			//TODO: GPU vendors for MacOS, what do they even use these days?
+		  
+			if (adapterName.Contains("INTEL"))
+			{
+				return ComputeName.DirectML;
+			}
+			if (adapterName.Contains("AMD") || adapterName.Contains("RADEON"))
+			{
+				return ComputeName.DirectML;
+			}
+			if (adapterName.Contains("NVIDIA"))
+			{
+				return ComputeName.CUDA;
+			}
+
+			GD.Print("Graphics Card not recognized."); //Should use CPU
+			return ComputeName.CPU;
+		}
+	}
+}

Game/addons/godot_rl_agents/onnx/csharp/docs/ONNXInference.xml

@@ -0,0 +1,31 @@
+<docs>
+	<members name="ONNXInference">
+		<ONNXInference>
+			<summary>
+       		The main <c>ONNXInference</c> Class that handles the inference process.
+        	</summary>
+		</ONNXInference>
+   		<Initialize>
+        	<summary>
+			Starts the inference process. 
+			</summary>
+			<param name="Path">Path to the ONNX model, expects a path inside resources.</param>
+			<param name="BatchSize">How many observations will the model recieve.</param>
+		</Initialize>
+		<Run>
+			<summary>
+			Runs the given input through the model and returns the output.
+			</summary>
+			<param name="obs">Dictionary containing all observations.</param>
+			<param name="state_ins">How many different agents are creating these observations.</param>
+			<returns>A Dictionary of arrays, containing instructions based on the observations.</returns>
+		</Run>
+		<Load>
+			<summary>
+			Loads the given model into the inference process, using the best Execution provider available.
+			</summary>
+			<param name="Path">Path to the ONNX model, expects a path inside resources.</param>
+			<returns>InferenceSession ready to run.</returns>
+		</Load>
+	</members>
+</docs>

Game/addons/godot_rl_agents/onnx/csharp/docs/SessionConfigurator.xml

@@ -0,0 +1,29 @@
+<docs>
+	<members name="SessionConfigurator">
+		<SessionConfigurator>
+			<summary>
+       		The main <c>SessionConfigurator</c> Class that handles the execution options and providers for the inference process.
+        	</summary>
+		</SessionConfigurator>
+  		<GetSessionOptions>
+    	    <summary>
+    	    Creates a SessionOptions with all available execution providers.
+    	    </summary>
+			<returns>SessionOptions with all available execution providers.</returns>
+    	</GetSessionOptions>
+    	<SystemCheck>
+        	<summary>
+        	Appends any execution provider available in the current system.  
+        	</summary>
+			<remarks>
+				This function is mainly verbose for tracking implementation progress of different compute APIs. 
+			</remarks>
+    	</SystemCheck>
+    	<ComputeCheck>
+        	<summary>
+        		Checks for available GPUs.
+        	</summary>
+			<returns>An integer identifier for each compute platform.</returns>
+    	</ComputeCheck>
+	</members>
+</docs>

Game/addons/godot_rl_agents/onnx/wrapper/ONNX_wrapper.gd

@@ -0,0 +1,51 @@
+extends Resource
+class_name ONNXModel
+var inferencer_script = load("res://addons/godot_rl_agents/onnx/csharp/ONNXInference.cs")
+
+var inferencer = null
+
+## How many action values the model outputs
+var action_output_size: int
+
+## Used to differentiate models
+## that only output continuous action mean (e.g. sb3, cleanrl export)
+## versus models that output mean and logstd (e.g. rllib export)
+var action_means_only: bool
+
+## Whether action_means_value has been set already for this model
+var action_means_only_set: bool
+
+# Must provide the path to the model and the batch size
+func _init(model_path, batch_size):
+	inferencer = inferencer_script.new()
+	action_output_size = inferencer.Initialize(model_path, batch_size)
+
+# This function is the one that will be called from the game,
+# requires the observations as an Dictionary and the state_ins as an int
+# returns a Dictionary containing the action the model takes.
+func run_inference(obs: Dictionary, state_ins: int) -> Dictionary:
+	if inferencer == null:
+		printerr("Inferencer not initialized")
+		return {}
+	return inferencer.RunInference(obs, state_ins)
+
+
+func _notification(what):
+	if what == NOTIFICATION_PREDELETE:
+		inferencer.FreeDisposables()
+		inferencer.free()
+
+# Check whether agent uses a continuous actions model with only action means or not
+func set_action_means_only(agent_action_space):
+	action_means_only_set = true
+	var continuous_only: bool = true
+	var continuous_actions: int
+	for action in agent_action_space:
+		if not agent_action_space[action]["action_type"] == "continuous":
+			continuous_only = false
+			break
+		else:
+			continuous_actions += agent_action_space[action]["size"]
+	if continuous_only:
+		if continuous_actions == action_output_size:
+			action_means_only = true

Game/addons/godot_rl_agents/onnx/wrapper/ONNX_wrapper.gd.uid

@@ -0,0 +1 @@
+uid://c35ckkxpe764s

Game/addons/godot_rl_agents/plugin.cfg

@@ -0,0 +1,7 @@
+[plugin]
+
+name="GodotRLAgents"
+description="Custom nodes for the godot rl agents toolkit	"
+author="Edward Beeching"
+version="0.1"
+script="godot_rl_agents.gd"

Game/addons/godot_rl_agents/rewards/ApproachNodeReward2D.gd

@@ -0,0 +1,30 @@
+extends RewardFunction2D
+class_name ApproachNodeReward2D
+
+## Calculates the reward for approaching node
+## a reward is only added when the agent reaches a new
+## best distance to the target object.
+
+## Best distance reward will be calculated for this object
+@export var target_node: Node2D
+
+## Scales the reward, 1.0 means the reward is equal to 
+## how much closer the agent is than the previous best.
+@export_range(0.0, 1.0, 0.0001, "or_greater") var reward_scale: float = 1.0
+
+var _best_distance
+
+
+func get_reward() -> float:
+	var reward := 0.0
+	var current_distance := global_position.distance_to(target_node.global_position)
+	if not _best_distance:
+		_best_distance = current_distance
+	if current_distance < _best_distance:
+		reward = (_best_distance - current_distance) * reward_scale
+		_best_distance = current_distance
+	return reward
+
+
+func reset():
+	_best_distance = null

Game/addons/godot_rl_agents/rewards/ApproachNodeReward2D.gd.uid

@@ -0,0 +1 @@
+uid://2jcows6svnje

Game/addons/godot_rl_agents/rewards/ApproachNodeReward3D.gd

@@ -0,0 +1,30 @@
+extends RewardFunction3D
+class_name ApproachNodeReward3D
+
+## Calculates the reward for approaching node
+## a reward is only added when the agent reaches a new
+## best distance to the target object.
+
+## Best distance reward will be calculated for this object
+@export var target_node: Node3D
+
+## Scales the reward, 1.0 means the reward is equal to 
+## how much closer the agent is than the previous best.
+@export_range(0.0, 1.0, 0.0001, "or_greater") var reward_scale: float = 1.0
+
+var _best_distance
+
+
+func get_reward() -> float:
+	var reward := 0.0
+	var current_distance := global_position.distance_to(target_node.global_position)
+	if not _best_distance:
+		_best_distance = current_distance
+	if current_distance < _best_distance:
+		reward = (_best_distance - current_distance) * reward_scale
+		_best_distance = current_distance
+	return reward
+
+
+func reset():
+	_best_distance = null

Game/addons/godot_rl_agents/rewards/ApproachNodeReward3D.gd.uid

@@ -0,0 +1 @@
+uid://bpgeoecqatvwi

Game/addons/godot_rl_agents/rewards/RewardFunction2D.gd

@@ -0,0 +1,10 @@
+extends Node2D
+class_name RewardFunction2D
+
+
+func get_reward():
+	return 0.0
+
+
+func reset():
+	return

Game/addons/godot_rl_agents/rewards/RewardFunction2D.gd.uid

@@ -0,0 +1 @@
+uid://52jl48u122l8

Game/addons/godot_rl_agents/rewards/RewardFunction3D.gd

@@ -0,0 +1,10 @@
+extends Node3D
+class_name RewardFunction3D
+
+
+func get_reward():
+	return 0.0
+
+
+func reset():
+	return

Game/addons/godot_rl_agents/rewards/RewardFunction3D.gd.uid

@@ -0,0 +1 @@
+uid://bwqq2ytgnepid

Game/addons/godot_rl_agents/sensors/sensors_2d/ExampleRaycastSensor2D.tscn

@@ -0,0 +1,48 @@
+[gd_scene load_steps=5 format=3 uid="uid://ddeq7mn1ealyc"]
+
+[ext_resource type="Script" path="res://addons/godot_rl_agents/sensors/sensors_2d/RaycastSensor2D.gd" id="1"]
+
+[sub_resource type="GDScript" id="2"]
+script/source = "extends Node2D
+
+
+
+func _physics_process(delta: float) -> void:
+    print(\"step start\")
+    
+"
+
+[sub_resource type="GDScript" id="1"]
+script/source = "extends RayCast2D
+
+var steps = 1
+
+func _physics_process(delta: float) -> void:
+    print(\"processing raycast\")
+    steps += 1
+    if steps % 2:
+        force_raycast_update()
+
+    print(is_colliding())
+"
+
+[sub_resource type="CircleShape2D" id="3"]
+
+[node name="ExampleRaycastSensor2D" type="Node2D"]
+script = SubResource("2")
+
+[node name="ExampleAgent" type="Node2D" parent="."]
+position = Vector2(573, 314)
+rotation = 0.286234
+
+[node name="RaycastSensor2D" type="Node2D" parent="ExampleAgent"]
+script = ExtResource("1")
+
+[node name="TestRayCast2D" type="RayCast2D" parent="."]
+script = SubResource("1")
+
+[node name="StaticBody2D" type="StaticBody2D" parent="."]
+position = Vector2(1, 52)
+
+[node name="CollisionShape2D" type="CollisionShape2D" parent="StaticBody2D"]
+shape = SubResource("3")

Game/addons/godot_rl_agents/sensors/sensors_2d/GridSensor2D.gd

@@ -0,0 +1,235 @@
+@tool
+extends ISensor2D
+class_name GridSensor2D
+
+@export var debug_view := false:
+	get:
+		return debug_view
+	set(value):
+		debug_view = value
+		_update()
+
+@export_flags_2d_physics var detection_mask := 0:
+	get:
+		return detection_mask
+	set(value):
+		detection_mask = value
+		_update()
+
+@export var collide_with_areas := false:
+	get:
+		return collide_with_areas
+	set(value):
+		collide_with_areas = value
+		_update()
+
+@export var collide_with_bodies := true:
+	get:
+		return collide_with_bodies
+	set(value):
+		collide_with_bodies = value
+		_update()
+
+@export_range(1, 200, 0.1) var cell_width := 20.0:
+	get:
+		return cell_width
+	set(value):
+		cell_width = value
+		_update()
+
+@export_range(1, 200, 0.1) var cell_height := 20.0:
+	get:
+		return cell_height
+	set(value):
+		cell_height = value
+		_update()
+
+@export_range(1, 21, 2, "or_greater") var grid_size_x := 3:
+	get:
+		return grid_size_x
+	set(value):
+		grid_size_x = value
+		_update()
+
+@export_range(1, 21, 2, "or_greater") var grid_size_y := 3:
+	get:
+		return grid_size_y
+	set(value):
+		grid_size_y = value
+		_update()
+
+var _obs_buffer: PackedFloat64Array
+var _rectangle_shape: RectangleShape2D
+var _collision_mapping: Dictionary
+var _n_layers_per_cell: int
+
+var _highlighted_cell_color: Color
+var _standard_cell_color: Color
+
+
+func get_observation():
+	return _obs_buffer
+
+
+func _update():
+	if Engine.is_editor_hint():
+		if is_node_ready():
+			_spawn_nodes()
+
+
+func _ready() -> void:
+	_set_colors()
+
+	if Engine.is_editor_hint():
+		if get_child_count() == 0:
+			_spawn_nodes()
+	else:
+		_spawn_nodes()
+
+
+func _set_colors() -> void:
+	_standard_cell_color = Color(100.0 / 255.0, 100.0 / 255.0, 100.0 / 255.0, 100.0 / 255.0)
+	_highlighted_cell_color = Color(255.0 / 255.0, 100.0 / 255.0, 100.0 / 255.0, 100.0 / 255.0)
+
+
+func _get_collision_mapping() -> Dictionary:
+	# defines which layer is mapped to which cell obs index
+	var total_bits = 0
+	var collision_mapping = {}
+	for i in 32:
+		var bit_mask = 2 ** i
+		if (detection_mask & bit_mask) > 0:
+			collision_mapping[i] = total_bits
+			total_bits += 1
+
+	return collision_mapping
+
+
+func _spawn_nodes():
+	for cell in get_children():
+		cell.name = "_%s" % cell.name  # Otherwise naming below will fail
+		cell.queue_free()
+
+	_collision_mapping = _get_collision_mapping()
+	#prints("collision_mapping", _collision_mapping, len(_collision_mapping))
+	# allocate memory for the observations
+	_n_layers_per_cell = len(_collision_mapping)
+	_obs_buffer = PackedFloat64Array()
+	_obs_buffer.resize(grid_size_x * grid_size_y * _n_layers_per_cell)
+	_obs_buffer.fill(0)
+	#prints(len(_obs_buffer), _obs_buffer )
+
+	_rectangle_shape = RectangleShape2D.new()
+	_rectangle_shape.set_size(Vector2(cell_width, cell_height))
+
+	var shift := Vector2(
+		-(grid_size_x / 2) * cell_width,
+		-(grid_size_y / 2) * cell_height,
+	)
+
+	for i in grid_size_x:
+		for j in grid_size_y:
+			var cell_position = Vector2(i * cell_width, j * cell_height) + shift
+			_create_cell(i, j, cell_position)
+
+
+func _create_cell(i: int, j: int, position: Vector2):
+	var cell := Area2D.new()
+	cell.position = position
+	cell.name = "GridCell %s %s" % [i, j]
+	cell.modulate = _standard_cell_color
+
+	if collide_with_areas:
+		cell.area_entered.connect(_on_cell_area_entered.bind(i, j))
+		cell.area_exited.connect(_on_cell_area_exited.bind(i, j))
+
+	if collide_with_bodies:
+		cell.body_entered.connect(_on_cell_body_entered.bind(i, j))
+		cell.body_exited.connect(_on_cell_body_exited.bind(i, j))
+
+	cell.collision_layer = 0
+	cell.collision_mask = detection_mask
+	cell.monitorable = true
+	add_child(cell)
+	cell.set_owner(get_tree().edited_scene_root)
+
+	var col_shape := CollisionShape2D.new()
+	col_shape.shape = _rectangle_shape
+	col_shape.name = "CollisionShape2D"
+	cell.add_child(col_shape)
+	col_shape.set_owner(get_tree().edited_scene_root)
+
+	if debug_view:
+		var quad = MeshInstance2D.new()
+		quad.name = "MeshInstance2D"
+		var quad_mesh = QuadMesh.new()
+
+		quad_mesh.set_size(Vector2(cell_width, cell_height))
+
+		quad.mesh = quad_mesh
+		cell.add_child(quad)
+		quad.set_owner(get_tree().edited_scene_root)
+
+
+func _update_obs(cell_i: int, cell_j: int, collision_layer: int, entered: bool):
+	for key in _collision_mapping:
+		var bit_mask = 2 ** key
+		if (collision_layer & bit_mask) > 0:
+			var collison_map_index = _collision_mapping[key]
+
+			var obs_index = (
+				(cell_i * grid_size_y * _n_layers_per_cell)
+				+ (cell_j * _n_layers_per_cell)
+				+ collison_map_index
+			)
+			#prints(obs_index, cell_i, cell_j)
+			if entered:
+				_obs_buffer[obs_index] += 1
+			else:
+				_obs_buffer[obs_index] -= 1
+
+
+func _toggle_cell(cell_i: int, cell_j: int):
+	var cell = get_node_or_null("GridCell %s %s" % [cell_i, cell_j])
+
+	if cell == null:
+		print("cell not found, returning")
+
+	var n_hits = 0
+	var start_index = (cell_i * grid_size_y * _n_layers_per_cell) + (cell_j * _n_layers_per_cell)
+	for i in _n_layers_per_cell:
+		n_hits += _obs_buffer[start_index + i]
+
+	if n_hits > 0:
+		cell.modulate = _highlighted_cell_color
+	else:
+		cell.modulate = _standard_cell_color
+
+
+func _on_cell_area_entered(area: Area2D, cell_i: int, cell_j: int):
+	#prints("_on_cell_area_entered", cell_i, cell_j)
+	_update_obs(cell_i, cell_j, area.collision_layer, true)
+	if debug_view:
+		_toggle_cell(cell_i, cell_j)
+	#print(_obs_buffer)
+
+
+func _on_cell_area_exited(area: Area2D, cell_i: int, cell_j: int):
+	#prints("_on_cell_area_exited", cell_i, cell_j)
+	_update_obs(cell_i, cell_j, area.collision_layer, false)
+	if debug_view:
+		_toggle_cell(cell_i, cell_j)
+
+
+func _on_cell_body_entered(body: Node2D, cell_i: int, cell_j: int):
+	#prints("_on_cell_body_entered", cell_i, cell_j)
+	_update_obs(cell_i, cell_j, body.collision_layer, true)
+	if debug_view:
+		_toggle_cell(cell_i, cell_j)
+
+
+func _on_cell_body_exited(body: Node2D, cell_i: int, cell_j: int):
+	#prints("_on_cell_body_exited", cell_i, cell_j)
+	_update_obs(cell_i, cell_j, body.collision_layer, false)
+	if debug_view:
+		_toggle_cell(cell_i, cell_j)

Game/addons/godot_rl_agents/sensors/sensors_2d/GridSensor2D.gd.uid

@@ -0,0 +1 @@
+uid://ctfloulhiht0v

Game/addons/godot_rl_agents/sensors/sensors_2d/ISensor2D.gd

@@ -0,0 +1,25 @@
+extends Node2D
+class_name ISensor2D
+
+var _obs: Array = []
+var _active := false
+
+
+func get_observation():
+	pass
+
+
+func activate():
+	_active = true
+
+
+func deactivate():
+	_active = false
+
+
+func _update_observation():
+	pass
+
+
+func reset():
+	pass

Game/addons/godot_rl_agents/sensors/sensors_2d/ISensor2D.gd.uid

@@ -0,0 +1 @@
+uid://dwhx5ltv86b4e

Game/addons/godot_rl_agents/sensors/sensors_2d/PositionSensor2D.gd

@@ -0,0 +1,65 @@
+extends ISensor2D
+class_name PositionSensor2D
+
+@export var objects_to_observe: Array[Node2D]
+
+## Whether to include relative x position in obs
+@export var include_x := true
+## Whether to include relative y position in obs
+@export var include_y := true
+
+## Max distance, values in obs will be normalized,
+## 0 will represent the closest distance possible, and 1 the farthest.
+## Do not use a much larger value than needed, as it would make the obs
+## very small after normalization.
+@export_range(0.01, 20_000) var max_distance := 1.0
+
+@export var use_separate_direction: bool = false
+
+@export var debug_lines: bool = true
+@export var debug_color: Color = Color.GREEN
+
+@onready var line: Line2D
+
+
+func _ready() -> void:
+	if debug_lines:
+		line = Line2D.new()
+		add_child(line)
+		line.width = 1
+		line.default_color = debug_color
+
+func get_observation():
+	var observations: Array[float]
+
+	if debug_lines:
+		line.clear_points()
+
+	for obj in objects_to_observe:
+		var relative_position := Vector2.ZERO
+
+		## If object has been removed, keep the zeroed position
+		if is_instance_valid(obj): relative_position = to_local(obj.global_position)
+
+		if debug_lines:
+			line.add_point(Vector2.ZERO)
+			line.add_point(relative_position)
+
+		var direction := Vector2.ZERO 
+		var distance := 0.0
+		if use_separate_direction:
+			direction = relative_position.normalized()
+			distance = min(relative_position.length() / max_distance, 1.0)
+			if include_x:
+				observations.append(direction.x)
+			if include_y:
+				observations.append(direction.y)
+			observations.append(distance)
+		else:
+			relative_position = relative_position.limit_length(max_distance) / max_distance
+			if include_x:
+				observations.append(relative_position.x)
+			if include_y:
+				observations.append(relative_position.y)
+
+	return observations

Game/addons/godot_rl_agents/sensors/sensors_2d/PositionSensor2D.gd.uid

@@ -0,0 +1 @@
+uid://dvny3m75wta1e

Game/addons/godot_rl_agents/sensors/sensors_2d/RGBCameraSensor2D.gd

@@ -0,0 +1,77 @@
+extends Node2D
+class_name RGBCameraSensor2D
+var camera_pixels = null
+
+@export var camera_zoom_factor := Vector2(0.1, 0.1)
+@onready var camera := $SubViewport/Camera
+@onready var preview_window := $Control
+@onready var camera_texture := $Control/CameraTexture as Sprite2D
+@onready var processed_texture := $Control/ProcessedTexture as Sprite2D
+@onready var sub_viewport := $SubViewport as SubViewport
+@onready var displayed_image: ImageTexture
+
+@export var render_image_resolution := Vector2i(36, 36)
+## Display size does not affect rendered or sent image resolution.
+## Scale is relative to either render image or downscale image resolution
+## depending on which mode is set.
+@export var displayed_image_scale_factor := Vector2i(8, 8)
+
+@export_group("Downscale image options")
+## Enable to downscale the rendered image before sending the obs.
+@export var downscale_image: bool = false
+## If downscale_image is true, will display the downscaled image instead of rendered image.
+@export var display_downscaled_image: bool = true
+## This is the resolution of the image that will be sent after downscaling
+@export var resized_image_resolution := Vector2i(36, 36)
+
+
+func _ready():
+	DisplayServer.register_additional_output(self)
+
+	camera.zoom = camera_zoom_factor
+
+	var preview_size: Vector2
+
+	sub_viewport.world_2d = get_tree().get_root().get_world_2d()
+	sub_viewport.size = render_image_resolution
+	camera_texture.scale = displayed_image_scale_factor
+
+	if downscale_image and display_downscaled_image:
+		camera_texture.visible = false
+		processed_texture.scale = displayed_image_scale_factor
+		preview_size = displayed_image_scale_factor * resized_image_resolution
+	else:
+		processed_texture.visible = false
+		preview_size = displayed_image_scale_factor * render_image_resolution
+
+	preview_window.size = preview_size
+
+
+func get_camera_pixel_encoding():
+	var image := camera_texture.get_texture().get_image() as Image
+
+	if downscale_image:
+		image.resize(
+			resized_image_resolution.x, resized_image_resolution.y, Image.INTERPOLATE_NEAREST
+		)
+		if display_downscaled_image:
+			if not processed_texture.texture:
+				displayed_image = ImageTexture.create_from_image(image)
+				processed_texture.texture = displayed_image
+			else:
+				displayed_image.update(image)
+
+	return image.get_data().hex_encode()
+
+
+func get_camera_shape() -> Array:
+	var size = resized_image_resolution if downscale_image else render_image_resolution
+
+	assert(
+		size.x >= 36 and size.y >= 36,
+		"Camera sensor sent image resolution must be 36x36 or larger."
+	)
+	if sub_viewport.transparent_bg:
+		return [4, size.y, size.x]
+	else:
+		return [3, size.y, size.x]

Game/addons/godot_rl_agents/sensors/sensors_2d/RGBCameraSensor2D.gd.uid

@@ -0,0 +1 @@
+uid://c8f8s63q3yyu1

Game/addons/godot_rl_agents/sensors/sensors_2d/RGBCameraSensor2D.tscn

@@ -0,0 +1,36 @@
+[gd_scene load_steps=3 format=3 uid="uid://bav1cl8uwc45c"]
+
+[ext_resource type="Script" path="res://addons/godot_rl_agents/sensors/sensors_2d/RGBCameraSensor2D.gd" id="1_txpo2"]
+
+[sub_resource type="ViewportTexture" id="ViewportTexture_jks1s"]
+viewport_path = NodePath("SubViewport")
+
+[node name="RGBCameraSensor2D" type="Node2D"]
+script = ExtResource("1_txpo2")
+displayed_image_scale_factor = Vector2(3, 3)
+
+[node name="RemoteTransform" type="RemoteTransform2D" parent="."]
+remote_path = NodePath("../SubViewport/Camera")
+
+[node name="SubViewport" type="SubViewport" parent="."]
+canvas_item_default_texture_filter = 0
+size = Vector2i(36, 36)
+render_target_update_mode = 4
+
+[node name="Camera" type="Camera2D" parent="SubViewport"]
+position_smoothing_speed = 2.0
+
+[node name="Control" type="Window" parent="."]
+canvas_item_default_texture_filter = 0
+title = "CameraSensor"
+position = Vector2i(20, 40)
+size = Vector2i(64, 64)
+theme_override_font_sizes/title_font_size = 12
+metadata/_edit_use_anchors_ = true
+
+[node name="CameraTexture" type="Sprite2D" parent="Control"]
+texture = SubResource("ViewportTexture_jks1s")
+centered = false
+
+[node name="ProcessedTexture" type="Sprite2D" parent="Control"]
+centered = false

Game/addons/godot_rl_agents/sensors/sensors_2d/RaycastSensor2D.gd

@@ -0,0 +1,123 @@
+@tool
+extends ISensor2D
+class_name RaycastSensor2D
+
+@export_flags_2d_physics var collision_mask := 1:
+	get:
+		return collision_mask
+	set(value):
+		collision_mask = value
+		_update()
+
+@export var collide_with_areas := false:
+	get:
+		return collide_with_areas
+	set(value):
+		collide_with_areas = value
+		_update()
+
+@export var collide_with_bodies := true:
+	get:
+		return collide_with_bodies
+	set(value):
+		collide_with_bodies = value
+		_update()
+
+@export var n_rays := 16.0:
+	get:
+		return n_rays
+	set(value):
+		n_rays = value
+		_update()
+
+@export_range(5, 3000, 5.0) var ray_length := 200:
+	get:
+		return ray_length
+	set(value):
+		ray_length = value
+		_update()
+@export_range(5, 360, 5.0) var cone_width := 360.0:
+	get:
+		return cone_width
+	set(value):
+		cone_width = value
+		_update()
+
+@export var debug_draw := false:
+	get:
+		return debug_draw
+	set(value):
+		debug_draw = value
+		_update()
+
+var _angles = []
+var rays := []
+
+
+func _update():
+	if Engine.is_editor_hint():
+		if is_node_ready():
+			_spawn_nodes()
+
+func _ready() -> void:
+	if Engine.is_editor_hint():
+		if get_child_count() == 0:
+			_spawn_nodes()
+	else:
+		_spawn_nodes()
+
+
+func _spawn_nodes():
+	for ray in get_children():
+		ray.queue_free()
+	rays = []
+
+	_angles = []
+	var step = cone_width / (n_rays)
+	var start = step / 2 - cone_width / 2
+
+	for i in n_rays:
+		var angle = start + i * step
+		var ray = RayCast2D.new()
+		ray.set_target_position(
+			Vector2(ray_length * cos(deg_to_rad(angle)), ray_length * sin(deg_to_rad(angle)))
+		)
+		if debug_draw:
+			ray.enabled = true
+		else:
+			ray.enabled = false
+		ray.collide_with_areas = collide_with_areas
+		ray.collide_with_bodies = collide_with_bodies
+		ray.collision_mask = collision_mask
+		add_child(ray)
+		ray.set_owner(get_tree().edited_scene_root)
+		ray.set_name("node_" + str(i))
+		rays.append(ray)
+
+		_angles.append(start + i * step)
+
+
+func get_observation() -> Array:
+	return self.calculate_raycasts()
+
+
+func calculate_raycasts() -> Array:
+	var result = []
+	for ray in rays:
+		if not debug_draw:
+			ray.enabled = true
+		ray.force_raycast_update()
+		var distance = _get_raycast_distance(ray)
+		result.append(distance)
+		if not debug_draw:
+			ray.enabled = false
+	return result
+
+
+func _get_raycast_distance(ray: RayCast2D) -> float:
+	if !ray.is_colliding():
+		return 0.0
+
+	var distance = (global_position - ray.get_collision_point()).length()
+	distance = clamp(distance, 0.0, ray_length)
+	return (ray_length - distance) / ray_length

Game/addons/godot_rl_agents/sensors/sensors_2d/RaycastSensor2D.gd.uid

@@ -0,0 +1 @@
+uid://c0eh0jjfthdgw

Game/addons/godot_rl_agents/sensors/sensors_2d/RaycastSensor2D.tscn

@@ -0,0 +1,7 @@
+[gd_scene load_steps=2 format=3 uid="uid://drvfihk5esgmv"]
+
+[ext_resource type="Script" path="res://addons/godot_rl_agents/sensors/sensors_2d/RaycastSensor2D.gd" id="1"]
+
+[node name="RaycastSensor2D" type="Node2D"]
+script = ExtResource("1")
+n_rays = 17.0

Game/addons/godot_rl_agents/sensors/sensors_3d/ExampleRaycastSensor3D.tscn

@@ -0,0 +1,6 @@
+[gd_scene format=3 uid="uid://biu787qh4woik"]
+
+[node name="ExampleRaycastSensor3D" type="Node3D"]
+
+[node name="Camera3D" type="Camera3D" parent="."]
+transform = Transform3D(1, 0, 0, 0, 1, 0, 0, 0, 1, 0.804183, 0, 2.70146)

Game/addons/godot_rl_agents/sensors/sensors_3d/GridSensor3D.gd

@@ -0,0 +1,258 @@
+@tool
+extends ISensor3D
+class_name GridSensor3D
+
+@export var debug_view := false:
+	get:
+		return debug_view
+	set(value):
+		debug_view = value
+		_update()
+
+@export_flags_3d_physics var detection_mask := 0:
+	get:
+		return detection_mask
+	set(value):
+		detection_mask = value
+		_update()
+
+@export var collide_with_areas := false:
+	get:
+		return collide_with_areas
+	set(value):
+		collide_with_areas = value
+		_update()
+
+@export var collide_with_bodies := false:
+	# NOTE! The sensor will not detect StaticBody3D, add an area to static bodies to detect them
+	get:
+		return collide_with_bodies
+	set(value):
+		collide_with_bodies = value
+		_update()
+
+@export_range(0.1, 2, 0.1) var cell_width := 1.0:
+	get:
+		return cell_width
+	set(value):
+		cell_width = value
+		_update()
+
+@export_range(0.1, 2, 0.1) var cell_height := 1.0:
+	get:
+		return cell_height
+	set(value):
+		cell_height = value
+		_update()
+
+@export_range(1, 21, 1, "or_greater") var grid_size_x := 3:
+	get:
+		return grid_size_x
+	set(value):
+		grid_size_x = value
+		_update()
+
+@export_range(1, 21, 1, "or_greater") var grid_size_z := 3:
+	get:
+		return grid_size_z
+	set(value):
+		grid_size_z = value
+		_update()
+
+var _obs_buffer: PackedFloat64Array
+var _box_shape: BoxShape3D
+var _collision_mapping: Dictionary
+var _n_layers_per_cell: int
+
+var _highlighted_box_material: StandardMaterial3D
+var _standard_box_material: StandardMaterial3D
+
+
+func get_observation():
+	return _obs_buffer
+
+
+func reset():
+	_obs_buffer.fill(0)
+
+
+func _update():
+	if Engine.is_editor_hint():
+		if is_node_ready():
+			_spawn_nodes()
+
+
+func _ready() -> void:
+	_make_materials()
+
+	if Engine.is_editor_hint():
+		if get_child_count() == 0:
+			_spawn_nodes()
+	else:
+		_spawn_nodes()
+
+
+func _make_materials() -> void:
+	if _highlighted_box_material != null and _standard_box_material != null:
+		return
+
+	_standard_box_material = StandardMaterial3D.new()
+	_standard_box_material.set_transparency(1)  # ALPHA
+	_standard_box_material.albedo_color = Color(
+		100.0 / 255.0, 100.0 / 255.0, 100.0 / 255.0, 100.0 / 255.0
+	)
+
+	_highlighted_box_material = StandardMaterial3D.new()
+	_highlighted_box_material.set_transparency(1)  # ALPHA
+	_highlighted_box_material.albedo_color = Color(
+		255.0 / 255.0, 100.0 / 255.0, 100.0 / 255.0, 100.0 / 255.0
+	)
+
+
+func _get_collision_mapping() -> Dictionary:
+	# defines which layer is mapped to which cell obs index
+	var total_bits = 0
+	var collision_mapping = {}
+	for i in 32:
+		var bit_mask = 2 ** i
+		if (detection_mask & bit_mask) > 0:
+			collision_mapping[i] = total_bits
+			total_bits += 1
+
+	return collision_mapping
+
+
+func _spawn_nodes():
+	for cell in get_children():
+		cell.name = "_%s" % cell.name  # Otherwise naming below will fail
+		cell.queue_free()
+
+	_collision_mapping = _get_collision_mapping()
+	#prints("collision_mapping", _collision_mapping, len(_collision_mapping))
+	# allocate memory for the observations
+	_n_layers_per_cell = len(_collision_mapping)
+	_obs_buffer = PackedFloat64Array()
+	_obs_buffer.resize(grid_size_x * grid_size_z * _n_layers_per_cell)
+	_obs_buffer.fill(0)
+	#prints(len(_obs_buffer), _obs_buffer )
+
+	_box_shape = BoxShape3D.new()
+	_box_shape.set_size(Vector3(cell_width, cell_height, cell_width))
+
+	var shift := Vector3(
+		-(grid_size_x / 2) * cell_width,
+		0,
+		-(grid_size_z / 2) * cell_width,
+	)
+
+	for i in grid_size_x:
+		for j in grid_size_z:
+			var cell_position = Vector3(i * cell_width, 0.0, j * cell_width) + shift
+			_create_cell(i, j, cell_position)
+
+
+func _create_cell(i: int, j: int, position: Vector3):
+	var cell := Area3D.new()
+	cell.position = position
+	cell.name = "GridCell %s %s" % [i, j]
+
+	if collide_with_areas:
+		cell.area_entered.connect(_on_cell_area_entered.bind(i, j))
+		cell.area_exited.connect(_on_cell_area_exited.bind(i, j))
+
+	if collide_with_bodies:
+		cell.body_entered.connect(_on_cell_body_entered.bind(i, j))
+		cell.body_exited.connect(_on_cell_body_exited.bind(i, j))
+
+#		cell.body_shape_entered.connect(_on_cell_body_shape_entered.bind(i, j))
+#		cell.body_shape_exited.connect(_on_cell_body_shape_exited.bind(i, j))
+
+	cell.collision_layer = 0
+	cell.collision_mask = detection_mask
+	cell.monitorable = true
+	cell.input_ray_pickable = false
+	add_child(cell)
+	cell.set_owner(get_tree().edited_scene_root)
+
+	var col_shape := CollisionShape3D.new()
+	col_shape.shape = _box_shape
+	col_shape.name = "CollisionShape3D"
+	cell.add_child(col_shape)
+	col_shape.set_owner(get_tree().edited_scene_root)
+
+	if debug_view:
+		var box = MeshInstance3D.new()
+		box.name = "MeshInstance3D"
+		var box_mesh = BoxMesh.new()
+
+		box_mesh.set_size(Vector3(cell_width, cell_height, cell_width))
+		box_mesh.material = _standard_box_material
+
+		box.mesh = box_mesh
+		cell.add_child(box)
+		box.set_owner(get_tree().edited_scene_root)
+
+
+func _update_obs(cell_i: int, cell_j: int, collision_layer: int, entered: bool):
+	for key in _collision_mapping:
+		var bit_mask = 2 ** key
+		if (collision_layer & bit_mask) > 0:
+			var collison_map_index = _collision_mapping[key]
+
+			var obs_index = (
+				(cell_i * grid_size_z * _n_layers_per_cell)
+				+ (cell_j * _n_layers_per_cell)
+				+ collison_map_index
+			)
+			#prints(obs_index, cell_i, cell_j)
+			if entered:
+				_obs_buffer[obs_index] += 1
+			else:
+				_obs_buffer[obs_index] -= 1
+
+
+func _toggle_cell(cell_i: int, cell_j: int):
+	var cell = get_node_or_null("GridCell %s %s" % [cell_i, cell_j])
+
+	if cell == null:
+		print("cell not found, returning")
+
+	var n_hits = 0
+	var start_index = (cell_i * grid_size_z * _n_layers_per_cell) + (cell_j * _n_layers_per_cell)
+	for i in _n_layers_per_cell:
+		n_hits += _obs_buffer[start_index + i]
+
+	var cell_mesh = cell.get_node_or_null("MeshInstance3D")
+	if n_hits > 0:
+		cell_mesh.mesh.material = _highlighted_box_material
+	else:
+		cell_mesh.mesh.material = _standard_box_material
+
+
+func _on_cell_area_entered(area: Area3D, cell_i: int, cell_j: int):
+	#prints("_on_cell_area_entered", cell_i, cell_j)
+	_update_obs(cell_i, cell_j, area.collision_layer, true)
+	if debug_view:
+		_toggle_cell(cell_i, cell_j)
+	#print(_obs_buffer)
+
+
+func _on_cell_area_exited(area: Area3D, cell_i: int, cell_j: int):
+	#prints("_on_cell_area_exited", cell_i, cell_j)
+	_update_obs(cell_i, cell_j, area.collision_layer, false)
+	if debug_view:
+		_toggle_cell(cell_i, cell_j)
+
+
+func _on_cell_body_entered(body: Node3D, cell_i: int, cell_j: int):
+	#prints("_on_cell_body_entered", cell_i, cell_j)
+	_update_obs(cell_i, cell_j, body.collision_layer, true)
+	if debug_view:
+		_toggle_cell(cell_i, cell_j)
+
+
+func _on_cell_body_exited(body: Node3D, cell_i: int, cell_j: int):
+	#prints("_on_cell_body_exited", cell_i, cell_j)
+	_update_obs(cell_i, cell_j, body.collision_layer, false)
+	if debug_view:
+		_toggle_cell(cell_i, cell_j)

Game/addons/godot_rl_agents/sensors/sensors_3d/GridSensor3D.gd.uid

@@ -0,0 +1 @@
+uid://me8mehqmblq8

Game/addons/godot_rl_agents/sensors/sensors_3d/ISensor3D.gd

@@ -0,0 +1,25 @@
+extends Node3D
+class_name ISensor3D
+
+var _obs: Array = []
+var _active := false
+
+
+func get_observation():
+	pass
+
+
+func activate():
+	_active = true
+
+
+func deactivate():
+	_active = false
+
+
+func _update_observation():
+	pass
+
+
+func reset():
+	pass

Game/addons/godot_rl_agents/sensors/sensors_3d/ISensor3D.gd.uid

@@ -0,0 +1 @@
+uid://b6dvaob0xndoh

Game/addons/godot_rl_agents/sensors/sensors_3d/PositionSensor3D.gd

@@ -0,0 +1,79 @@
+extends ISensor3D
+class_name PositionSensor3D
+
+@export var objects_to_observe: Array[Node3D]
+
+## Whether to include relative x position in obs
+@export var include_x := true
+## Whether to include relative y position in obs
+@export var include_y := true
+## Whether to include relative z position in obs
+@export var include_z := true
+
+## Max distance, values in obs will be normalized,
+## 0 will represent the closest distance possible, and 1 the farthest.
+## Do not use a much larger value than needed, as it would make the obs
+## very small after normalization.
+@export_range(0.01, 2_500) var max_distance := 1.0
+
+@export var use_separate_direction: bool = false
+
+@export var debug_lines: bool = true
+@export var debug_color: Color = Color.GREEN
+
+@onready var mesh: ImmediateMesh
+
+
+func _ready() -> void:
+	if debug_lines:
+		var debug_mesh = MeshInstance3D.new()
+		add_child(debug_mesh)
+		var line_material := StandardMaterial3D.new()
+		line_material.albedo_color = debug_color
+		debug_mesh.material_override = line_material
+		debug_mesh.mesh = ImmediateMesh.new()
+		mesh = debug_mesh.mesh
+
+
+func get_observation():
+	var observations: Array[float]
+
+	if debug_lines:
+		mesh.clear_surfaces()
+		mesh.surface_begin(Mesh.PRIMITIVE_LINES)
+		mesh.surface_set_color(debug_color)
+
+	for obj in objects_to_observe:
+		var relative_position := Vector3.ZERO
+
+		## If object has been removed, keep the zeroed position
+		if is_instance_valid(obj): relative_position = to_local(obj.global_position)
+
+		if debug_lines:
+			mesh.surface_add_vertex(Vector3.ZERO)
+			mesh.surface_add_vertex(relative_position)
+
+		var direction := Vector3.ZERO 
+		var distance := 0.0
+		if use_separate_direction:
+			direction = relative_position.normalized()
+			distance = min(relative_position.length() / max_distance, 1.0)
+			if include_x:
+				observations.append(direction.x)
+			if include_y:
+				observations.append(direction.y)
+			if include_z:
+				observations.append(direction.z)
+			observations.append(distance)
+		else:
+			relative_position = relative_position.limit_length(max_distance) / max_distance
+			if include_x:
+				observations.append(relative_position.x)
+			if include_y:
+				observations.append(relative_position.y)
+			if include_z:
+				observations.append(relative_position.z)
+
+	if debug_lines:
+		mesh.surface_end()
+	return observations

Game/addons/godot_rl_agents/sensors/sensors_3d/PositionSensor3D.gd.uid

@@ -0,0 +1 @@
+uid://cew2a213sw1q

Game/addons/godot_rl_agents/sensors/sensors_3d/RGBCameraSensor3D.gd

@@ -0,0 +1,63 @@
+extends Node3D
+class_name RGBCameraSensor3D
+var camera_pixels = null
+
+@onready var camera_texture := $Control/CameraTexture as Sprite2D
+@onready var processed_texture := $Control/ProcessedTexture as Sprite2D
+@onready var sub_viewport := $SubViewport as SubViewport
+@onready var displayed_image: ImageTexture
+
+@export var render_image_resolution := Vector2i(36, 36)
+## Display size does not affect rendered or sent image resolution.
+## Scale is relative to either render image or downscale image resolution
+## depending on which mode is set.
+@export var displayed_image_scale_factor := Vector2i(8, 8)
+
+@export_group("Downscale image options")
+## Enable to downscale the rendered image before sending the obs.
+@export var downscale_image: bool = false
+## If downscale_image is true, will display the downscaled image instead of rendered image.
+@export var display_downscaled_image: bool = true
+## This is the resolution of the image that will be sent after downscaling
+@export var resized_image_resolution := Vector2i(36, 36)
+
+
+func _ready():
+	sub_viewport.size = render_image_resolution
+	camera_texture.scale = displayed_image_scale_factor
+
+	if downscale_image and display_downscaled_image:
+		camera_texture.visible = false
+		processed_texture.scale = displayed_image_scale_factor
+	else:
+		processed_texture.visible = false
+
+
+func get_camera_pixel_encoding():
+	var image := camera_texture.get_texture().get_image() as Image
+
+	if downscale_image:
+		image.resize(
+			resized_image_resolution.x, resized_image_resolution.y, Image.INTERPOLATE_NEAREST
+		)
+		if display_downscaled_image:
+			if not processed_texture.texture:
+				displayed_image = ImageTexture.create_from_image(image)
+				processed_texture.texture = displayed_image
+			else:
+				displayed_image.update(image)
+
+	return image.get_data().hex_encode()
+
+
+func get_camera_shape() -> Array:
+	var size = resized_image_resolution if downscale_image else render_image_resolution
+
+	assert(
+		size.x >= 36 and size.y >= 36,
+		"Camera sensor sent image resolution must be 36x36 or larger."
+	)
+	if sub_viewport.transparent_bg:
+		return [4, size.y, size.x]
+	else:
+		return [3, size.y, size.x]

Game/addons/godot_rl_agents/sensors/sensors_3d/RGBCameraSensor3D.gd.uid

@@ -0,0 +1 @@
+uid://6e38006xhqf0

Game/addons/godot_rl_agents/sensors/sensors_3d/RGBCameraSensor3D.tscn

@@ -0,0 +1,35 @@
+[gd_scene load_steps=3 format=3 uid="uid://baaywi3arsl2m"]
+
+[ext_resource type="Script" path="res://addons/godot_rl_agents/sensors/sensors_3d/RGBCameraSensor3D.gd" id="1"]
+
+[sub_resource type="ViewportTexture" id="ViewportTexture_y72s3"]
+viewport_path = NodePath("SubViewport")
+
+[node name="RGBCameraSensor3D" type="Node3D"]
+script = ExtResource("1")
+
+[node name="RemoteTransform" type="RemoteTransform3D" parent="."]
+remote_path = NodePath("../SubViewport/Camera")
+
+[node name="SubViewport" type="SubViewport" parent="."]
+size = Vector2i(36, 36)
+render_target_update_mode = 3
+
+[node name="Camera" type="Camera3D" parent="SubViewport"]
+near = 0.5
+
+[node name="Control" type="Control" parent="."]
+layout_mode = 3
+anchors_preset = 15
+anchor_right = 1.0
+anchor_bottom = 1.0
+grow_horizontal = 2
+grow_vertical = 2
+metadata/_edit_use_anchors_ = true
+
+[node name="CameraTexture" type="Sprite2D" parent="Control"]
+texture = SubResource("ViewportTexture_y72s3")
+centered = false
+
+[node name="ProcessedTexture" type="Sprite2D" parent="Control"]
+centered = false

Game/addons/godot_rl_agents/sensors/sensors_3d/RaycastSensor3D.gd

@@ -0,0 +1,197 @@
+@tool
+extends ISensor3D
+class_name RayCastSensor3D
+@export_flags_3d_physics var collision_mask = 1:
+	get:
+		return collision_mask
+	set(value):
+		collision_mask = value
+		_update()
+@export_flags_3d_physics var boolean_class_mask = 1:
+	get:
+		return boolean_class_mask
+	set(value):
+		boolean_class_mask = value
+		_update()
+
+@export var n_rays_width := 6.0:
+	get:
+		return n_rays_width
+	set(value):
+		n_rays_width = value
+		_update()
+
+@export var n_rays_height := 6.0:
+	get:
+		return n_rays_height
+	set(value):
+		n_rays_height = value
+		_update()
+
+@export var ray_length := 10.0:
+	get:
+		return ray_length
+	set(value):
+		ray_length = value
+		_update()
+
+@export var cone_width := 60.0:
+	get:
+		return cone_width
+	set(value):
+		cone_width = value
+		_update()
+
+@export var cone_height := 60.0:
+	get:
+		return cone_height
+	set(value):
+		cone_height = value
+		_update()
+
+@export var collide_with_areas := false:
+	get:
+		return collide_with_areas
+	set(value):
+		collide_with_areas = value
+		_update()
+
+@export var collide_with_bodies := true:
+	get:
+		return collide_with_bodies
+	set(value):
+		collide_with_bodies = value
+		_update()
+
+@export var class_sensor := false
+
+@export var debug_draw := false:
+	get:
+		return debug_draw
+	set(value):
+		debug_draw = value
+		_update()
+
+var rays := []
+var geo = null
+
+
+func _update():
+	if Engine.is_editor_hint():
+		if is_node_ready():
+			_spawn_nodes()
+
+
+func _ready() -> void:
+	if Engine.is_editor_hint():
+		if get_child_count() == 0:
+			_spawn_nodes()
+	else:
+		_spawn_nodes()
+
+
+func _spawn_nodes():
+	print("spawning nodes")
+	for ray in get_children():
+		ray.queue_free()
+	if geo:
+		geo.clear()
+	#$Lines.remove_points()
+	rays = []
+
+	var horizontal_step = cone_width / (n_rays_width)
+	var vertical_step = cone_height / (n_rays_height)
+
+	var horizontal_start = horizontal_step / 2 - cone_width / 2
+	var vertical_start = vertical_step / 2 - cone_height / 2
+
+	var points = []
+
+	for i in n_rays_width:
+		for j in n_rays_height:
+			var angle_w = horizontal_start + i * horizontal_step
+			var angle_h = vertical_start + j * vertical_step
+			#angle_h = 0.0
+			var ray = RayCast3D.new()
+			var cast_to = to_spherical_coords(ray_length, angle_w, angle_h)
+			ray.set_target_position(cast_to)
+
+			points.append(cast_to)
+
+			if debug_draw:
+				ray.enabled = true
+			else:
+				ray.enabled = false
+			ray.collide_with_bodies = collide_with_bodies
+			ray.collide_with_areas = collide_with_areas
+			ray.collision_mask = collision_mask
+			add_child(ray)
+			ray.set_owner(get_tree().edited_scene_root)
+			ray.set_name("node_" + str(i) + " " + str(j))
+			rays.append(ray)
+			ray.force_raycast_update()
+
+
+#    if Engine.editor_hint:
+#        _create_debug_lines(points)
+
+
+func _create_debug_lines(points):
+	if not geo:
+		geo = ImmediateMesh.new()
+		add_child(geo)
+
+	geo.clear()
+	geo.begin(Mesh.PRIMITIVE_LINES)
+	for point in points:
+		geo.set_color(Color.AQUA)
+		geo.add_vertex(Vector3.ZERO)
+		geo.add_vertex(point)
+	geo.end()
+
+
+func display():
+	if geo:
+		geo.display()
+
+
+func to_spherical_coords(r, inc, azimuth) -> Vector3:
+	return Vector3(
+		r * sin(deg_to_rad(inc)) * cos(deg_to_rad(azimuth)),
+		r * sin(deg_to_rad(azimuth)),
+		r * cos(deg_to_rad(inc)) * cos(deg_to_rad(azimuth))
+	)
+
+
+func get_observation() -> Array:
+	return self.calculate_raycasts()
+
+
+func calculate_raycasts() -> Array:
+	var result = []
+	for ray in rays:
+		if not debug_draw:
+			ray.set_enabled(true)
+		ray.force_raycast_update()
+		var distance = _get_raycast_distance(ray)
+
+		result.append(distance)
+		if class_sensor:
+			var hit_class: float = 0
+			if ray.get_collider():
+				var hit_collision_layer = ray.get_collider().collision_layer
+				hit_collision_layer = hit_collision_layer & collision_mask
+				hit_class = (hit_collision_layer & boolean_class_mask) > 0
+			result.append(float(hit_class))
+		if not debug_draw:
+			ray.set_enabled(false)
+	return result
+
+
+func _get_raycast_distance(ray: RayCast3D) -> float:
+	if !ray.is_colliding():
+		return 0.0
+
+	var distance = (global_transform.origin - ray.get_collision_point()).length()
+	distance = clamp(distance, 0.0, ray_length)
+	return (ray_length - distance) / ray_length

Game/addons/godot_rl_agents/sensors/sensors_3d/RaycastSensor3D.gd.uid

@@ -0,0 +1 @@
+uid://glop37kjpwci

Game/addons/godot_rl_agents/sensors/sensors_3d/RaycastSensor3D.tscn

@@ -0,0 +1,27 @@
+[gd_scene load_steps=2 format=3 uid="uid://b803cbh1fmy66"]
+
+[ext_resource type="Script" path="res://addons/godot_rl_agents/sensors/sensors_3d/RaycastSensor3D.gd" id="1"]
+
+[node name="RaycastSensor3D" type="Node3D"]
+script = ExtResource("1")
+n_rays_width = 4.0
+n_rays_height = 2.0
+ray_length = 11.0
+
+[node name="node_1 0" type="RayCast3D" parent="."]
+target_position = Vector3(-1.38686, -2.84701, 10.5343)
+
+[node name="node_1 1" type="RayCast3D" parent="."]
+target_position = Vector3(-1.38686, 2.84701, 10.5343)
+
+[node name="node_2 0" type="RayCast3D" parent="."]
+target_position = Vector3(1.38686, -2.84701, 10.5343)
+
+[node name="node_2 1" type="RayCast3D" parent="."]
+target_position = Vector3(1.38686, 2.84701, 10.5343)
+
+[node name="node_3 0" type="RayCast3D" parent="."]
+target_position = Vector3(4.06608, -2.84701, 9.81639)
+
+[node name="node_3 1" type="RayCast3D" parent="."]
+target_position = Vector3(4.06608, 2.84701, 9.81639)

Game/addons/godot_rl_agents/sync.gd

@@ -0,0 +1,621 @@
+extends Node
+class_name Sync
+
+# --fixed-fps 2000 --disable-render-loop
+
+enum ControlModes {
+	HUMAN, ## Test the environment manually
+	TRAINING, ## Train a model
+	ONNX_INFERENCE ## Load a pretrained model using an .onnx file
+}
+@export var control_mode: ControlModes = ControlModes.TRAINING
+## Action will be repeated for n frames (Godot physics steps).
+@export_range(1, 10, 1, "or_greater") var action_repeat := 8
+## Speeds up the physics in the environment to enable faster training.
+@export_range(0, 10, 0.1, "or_greater") var speed_up := 1.0
+## The path to a trained .onnx model file to use for inference (only needed for the 'Onnx Inference' control mode).
+@export var onnx_model_path := ""
+## Whether the inference will be deterministic (NOTE: Only applies to discrete actions in onnx inference mode)
+@export var deterministic_inference := true
+
+# Onnx model stored for each requested path
+var onnx_models: Dictionary
+
+@onready var start_time = Time.get_ticks_msec()
+
+const MAJOR_VERSION := "0"
+const MINOR_VERSION := "7"
+const DEFAULT_PORT := "11008"
+const DEFAULT_SEED := "1"
+var stream: StreamPeerTCP = null
+var connected = false
+var message_center
+var should_connect = true
+
+var all_agents: Array
+var agents_training: Array
+## Policy name of each agent, for use with multi-policy multi-agent RL cases
+var agents_training_policy_names: Array[String] = ["shared_policy"]
+var agents_inference: Array
+var agents_heuristic: Array
+
+## For recording expert demos
+var agent_demo_record: Node
+## File path for writing recorded trajectories
+var expert_demo_save_path: String
+## Stores recorded trajectories
+var demo_trajectories: Array
+## A trajectory includes obs: Array, acts: Array, terminal (set in Python env instead)
+var current_demo_trajectory: Array
+
+var need_to_send_obs = false
+var args = null
+var initialized = false
+var just_reset = false
+var onnx_model = null
+var n_action_steps = 0
+
+var _action_space_training: Array[Dictionary] = []
+var _action_space_inference: Array[Dictionary] = []
+var _obs_space_training: Array[Dictionary] = []
+
+
+# Called when the node enters the scene tree for the first time.
+func _ready():
+	await get_parent().ready
+	get_tree().set_pause(true)
+	_initialize()
+	await get_tree().create_timer(1.0).timeout
+	get_tree().set_pause(false)
+
+
+func _initialize():
+	_get_agents()
+	args = _get_args()
+	Engine.physics_ticks_per_second = _get_speedup() * 60  # Replace with function body.
+	Engine.time_scale = _get_speedup() * 1.0
+	prints(
+		"physics ticks",
+		Engine.physics_ticks_per_second,
+		Engine.time_scale,
+		_get_speedup(),
+		speed_up
+	)
+
+	_set_heuristic("human", all_agents)
+
+	_initialize_training_agents()
+	_initialize_inference_agents()
+	_initialize_demo_recording()
+
+	_set_seed()
+	_set_action_repeat()
+	initialized = true
+
+
+func _initialize_training_agents():
+	if agents_training.size() > 0:
+		_obs_space_training.resize(agents_training.size())
+		_action_space_training.resize(agents_training.size())
+		for agent_idx in range(0, agents_training.size()):
+			_obs_space_training[agent_idx] = agents_training[agent_idx].get_obs_space()
+			_action_space_training[agent_idx] = agents_training[agent_idx].get_action_space()
+		connected = connect_to_server()
+		if connected:
+			_set_heuristic("model", agents_training)
+			_handshake()
+			_send_env_info()
+		else:
+			push_warning(
+				"Couldn't connect to Python server, using human controls instead. ",
+				"Did you start the training server using e.g. `gdrl` from the console?"
+			)
+
+
+func _initialize_inference_agents():
+	if agents_inference.size() > 0:
+		if control_mode == ControlModes.ONNX_INFERENCE:
+			assert(
+				FileAccess.file_exists(onnx_model_path),
+				"Onnx Model Path set on Sync node does not exist: %s" % onnx_model_path
+			)
+			onnx_models[onnx_model_path] = ONNXModel.new(onnx_model_path, 1)
+
+		for agent in agents_inference:
+			var action_space = agent.get_action_space()
+			_action_space_inference.append(action_space)
+
+			var agent_onnx_model: ONNXModel
+			if agent.onnx_model_path.is_empty():
+				assert(
+					onnx_models.has(onnx_model_path),
+					(
+						"Node %s has no onnx model path set " % agent.get_path()
+						+ "and sync node's control mode is not set to OnnxInference. "
+						+ "Either add the path to the AIController, "
+						+ "or if you want to use the path set on sync node instead, "
+						+ "set control mode to OnnxInference."
+					)
+				)
+				prints(
+					"Info: AIController %s" % agent.get_path(),
+					"has no onnx model path set.",
+					"Using path set on the sync node instead."
+				)
+				agent_onnx_model = onnx_models[onnx_model_path]
+			else:
+				if not onnx_models.has(agent.onnx_model_path):
+					assert(
+						FileAccess.file_exists(agent.onnx_model_path),
+						(
+							"Onnx Model Path set on %s node does not exist: %s"
+							% [agent.get_path(), agent.onnx_model_path]
+						)
+					)
+					onnx_models[agent.onnx_model_path] = ONNXModel.new(agent.onnx_model_path, 1)
+				agent_onnx_model = onnx_models[agent.onnx_model_path]
+
+			agent.onnx_model = agent_onnx_model
+			if not agent_onnx_model.action_means_only_set:
+				agent_onnx_model.set_action_means_only(action_space)
+
+		_set_heuristic("model", agents_inference)
+
+
+func _initialize_demo_recording():
+	if agent_demo_record:
+		expert_demo_save_path = agent_demo_record.expert_demo_save_path
+		assert(
+			not expert_demo_save_path.is_empty(),
+			"Expert demo save path set in %s is empty." % agent_demo_record.get_path()
+		)
+
+		InputMap.add_action("RemoveLastDemoEpisode")
+		InputMap.action_add_event(
+			"RemoveLastDemoEpisode", agent_demo_record.remove_last_episode_key
+		)
+		current_demo_trajectory.resize(2)
+		current_demo_trajectory[0] = []
+		current_demo_trajectory[1] = []
+		agent_demo_record.heuristic = "demo_record"
+
+
+func _physics_process(_delta):
+	# two modes, human control, agent control
+	# pause tree, send obs, get actions, set actions, unpause tree
+
+	_demo_record_process()
+
+	if n_action_steps % action_repeat != 0:
+		n_action_steps += 1
+		return
+
+	n_action_steps += 1
+
+	_training_process()
+	_inference_process()
+	_heuristic_process()
+
+
+func _training_process():
+	if connected:
+		get_tree().set_pause(true)
+
+		var obs = _get_obs_from_agents(agents_training)
+		var info = _get_info_from_agents(agents_training)
+
+		if just_reset:
+			just_reset = false
+
+			var reply = {"type": "reset", "obs": obs, "info": info}
+			_send_dict_as_json_message(reply)
+			# this should go straight to getting the action and setting it checked the agent, no need to perform one phyics tick
+			get_tree().set_pause(false)
+			return
+
+		if need_to_send_obs:
+			need_to_send_obs = false
+			var reward = _get_reward_from_agents()
+			var done = _get_done_from_agents()
+			#_reset_agents_if_done() # this ensures the new observation is from the next env instance : NEEDS REFACTOR
+
+			var reply = {"type": "step", "obs": obs, "reward": reward, "done": done, "info": info}
+			_send_dict_as_json_message(reply)
+
+		var handled = handle_message()
+
+
+func _inference_process():
+	if agents_inference.size() > 0:
+		var obs: Array = _get_obs_from_agents(agents_inference)
+		var actions = []
+
+		for agent_id in range(0, agents_inference.size()):
+			var model: ONNXModel = agents_inference[agent_id].onnx_model
+			var action = model.run_inference(obs[agent_id], 1.0)
+			var action_dict = _extract_action_dict(
+				action["output"], _action_space_inference[agent_id], model.action_means_only
+			)
+			actions.append(action_dict)
+
+		_set_agent_actions(actions, agents_inference)
+		_reset_agents_if_done(agents_inference)
+		get_tree().set_pause(false)
+
+
+func _demo_record_process():
+	if not agent_demo_record:
+		return
+
+	if Input.is_action_just_pressed("RemoveLastDemoEpisode"):
+		print("[Sync script][Demo recorder] Removing last recorded episode.")
+		demo_trajectories.remove_at(demo_trajectories.size() - 1)
+		print("Remaining episode count: %d" % demo_trajectories.size())
+
+	if n_action_steps % agent_demo_record.action_repeat != 0:
+		return
+
+	var obs_dict: Dictionary = agent_demo_record.get_obs()
+
+	# Get the current obs from the agent
+	assert(
+		obs_dict.has("obs"),
+		"Demo recorder needs an 'obs' key in get_obs() returned dictionary to record obs from."
+	)
+	current_demo_trajectory[0].append(obs_dict.obs)
+
+	# Get the action applied for the current obs from the agent
+	agent_demo_record.set_action()
+	var acts = agent_demo_record.get_action()
+
+	var terminal = agent_demo_record.get_done()
+	# Record actions only for non-terminal states
+	if terminal:
+		agent_demo_record.set_done_false()
+	else:
+		current_demo_trajectory[1].append(acts)
+
+	if terminal:
+		#current_demo_trajectory[2].append(true)
+		demo_trajectories.append(current_demo_trajectory.duplicate(true))
+		print("[Sync script][Demo recorder] Recorded episode count: %d" % demo_trajectories.size())
+		current_demo_trajectory[0].clear()
+		current_demo_trajectory[1].clear()
+
+
+func _heuristic_process():
+	for agent in agents_heuristic:
+		_reset_agents_if_done(agents_heuristic)
+
+
+func _extract_action_dict(action_array: Array, action_space: Dictionary, action_means_only: bool):
+	var index = 0
+	var result = {}
+	for key in action_space.keys():
+		var size = action_space[key]["size"]
+		var action_type = action_space[key]["action_type"]
+
+		if action_type == "discrete":
+			var largest_logit: float = -INF  # Value of the largest logit for this action in the actions array
+			var largest_logit_idx: int  # Index of the largest logit for this action in the actions array
+			for logit_idx in range(0, size):
+				var logit_value = action_array[index + logit_idx]
+				if logit_value > largest_logit:
+					largest_logit = logit_value
+					largest_logit_idx = logit_idx
+			if deterministic_inference:
+				result[key] = largest_logit_idx  # Index of the largest logit is the discrete action value
+			else:
+				var exp_logit_sum: float  # Sum of exp of each logit
+				var exp_logits: Array[float]
+
+				for logit_idx in range(0, size):
+					# Normalize using the max logit to add stability in case a logit would be huge after exp
+					exp_logits.append(exp(action_array[index + logit_idx] - largest_logit))
+					exp_logit_sum += exp_logits[logit_idx]
+
+				# Choose a random number, will be used to select an action
+				var random_value = randf_range(0, exp_logit_sum)
+
+				# Select the first index at which the sum is larger than the random number
+				var sum: float
+				for exp_logit_idx in exp_logits.size():
+					sum += exp_logits[exp_logit_idx]
+					if sum > random_value:
+						result[key] = exp_logit_idx
+						break
+			index += size
+		elif action_type == "continuous":
+			# For continous actions, we only take the action mean values
+			result[key] = clamp_array(action_array.slice(index, index + size), -1.0, 1.0)
+			if action_means_only:
+				index += size  # model only outputs action means, so we move index by size
+			else:
+				index += size * 2  # model outputs logstd after action mean, we skip the logstd part
+
+		else:
+			assert(
+				false,
+				(
+					'Only "discrete" and "continuous" action types supported. Found: %s action type set.'
+					% action_type
+				)
+			)
+
+	return result
+
+
+## For AIControllers that inherit mode from sync, sets the correct mode.
+func _set_agent_mode(agent: Node):
+	var agent_inherits_mode: bool = agent.control_mode == agent.ControlModes.INHERIT_FROM_SYNC
+
+	if agent_inherits_mode:
+		match control_mode:
+			ControlModes.HUMAN:
+				agent.control_mode = agent.ControlModes.HUMAN
+			ControlModes.TRAINING:
+				agent.control_mode = agent.ControlModes.TRAINING
+			ControlModes.ONNX_INFERENCE:
+				agent.control_mode = agent.ControlModes.ONNX_INFERENCE
+
+
+func _get_agents():
+	all_agents = get_tree().get_nodes_in_group("AGENT")
+	for agent in all_agents:
+		_set_agent_mode(agent)
+
+		if agent.control_mode == agent.ControlModes.TRAINING:
+			agents_training.append(agent)
+		elif agent.control_mode == agent.ControlModes.ONNX_INFERENCE:
+			agents_inference.append(agent)
+		elif agent.control_mode == agent.ControlModes.HUMAN:
+			agents_heuristic.append(agent)
+		elif agent.control_mode == agent.ControlModes.RECORD_EXPERT_DEMOS:
+			assert(
+				not agent_demo_record,
+				"Currently only a single AIController can be used for recording expert demos."
+			)
+			agent_demo_record = agent
+
+	var training_agent_count = agents_training.size()
+	agents_training_policy_names.resize(training_agent_count)
+	for i in range(0, training_agent_count):
+		agents_training_policy_names[i] = agents_training[i].policy_name
+
+
+func _set_heuristic(heuristic, agents: Array):
+	for agent in agents:
+		agent.set_heuristic(heuristic)
+
+
+func _handshake():
+	print("performing handshake")
+
+	var json_dict = _get_dict_json_message()
+	assert(json_dict["type"] == "handshake")
+	var major_version = json_dict["major_version"]
+	var minor_version = json_dict["minor_version"]
+	if major_version != MAJOR_VERSION:
+		print("WARNING: major verison mismatch ", major_version, " ", MAJOR_VERSION)
+	if minor_version != MINOR_VERSION:
+		print("WARNING: minor verison mismatch ", minor_version, " ", MINOR_VERSION)
+
+	print("handshake complete")
+
+
+func _get_dict_json_message():
+	# returns a dictionary from of the most recent message
+	# this is not waiting
+	while stream.get_available_bytes() == 0:
+		stream.poll()
+		if stream.get_status() != 2:
+			print("server disconnected status, closing")
+			get_tree().quit()
+			return null
+
+		OS.delay_usec(10)
+
+	var message = stream.get_string()
+	var json_data = JSON.parse_string(message)
+
+	return json_data
+
+
+func _send_dict_as_json_message(dict):
+	stream.put_string(JSON.stringify(dict, "", false))
+
+
+func _send_env_info():
+	var json_dict = _get_dict_json_message()
+	assert(json_dict["type"] == "env_info")
+
+	var message = {
+		"type": "env_info",
+		"observation_space": _obs_space_training,
+		"action_space": _action_space_training,
+		"n_agents": len(agents_training),
+		"agent_policy_names": agents_training_policy_names
+	}
+	_send_dict_as_json_message(message)
+
+
+func connect_to_server():
+	print("Waiting for one second to allow server to start")
+	OS.delay_msec(1000)
+	print("trying to connect to server")
+	stream = StreamPeerTCP.new()
+
+	# "localhost" was not working on windows VM, had to use the IP
+	var ip = "127.0.0.1"
+	var port = _get_port()
+	var connect = stream.connect_to_host(ip, port)
+	stream.set_no_delay(true)  # TODO check if this improves performance or not
+	stream.poll()
+	# Fetch the status until it is either connected (2) or failed to connect (3)
+	while stream.get_status() < 2:
+		stream.poll()
+	return stream.get_status() == 2
+
+
+func _get_args():
+	print("getting command line arguments")
+	var arguments = {}
+	for argument in OS.get_cmdline_args():
+		print(argument)
+		if argument.find("=") > -1:
+			var key_value = argument.split("=")
+			arguments[key_value[0].lstrip("--")] = key_value[1]
+		else:
+			# Options without an argument will be present in the dictionary,
+			# with the value set to an empty string.
+			arguments[argument.lstrip("--")] = ""
+
+	return arguments
+
+
+func _get_speedup():
+	print(args)
+	return args.get("speedup", str(speed_up)).to_float()
+
+
+func _get_port():
+	return args.get("port", DEFAULT_PORT).to_int()
+
+
+func _set_seed():
+	var _seed = args.get("env_seed", DEFAULT_SEED).to_int()
+	seed(_seed)
+
+
+func _set_action_repeat():
+	action_repeat = args.get("action_repeat", str(action_repeat)).to_int()
+
+
+func disconnect_from_server():
+	stream.disconnect_from_host()
+
+
+func handle_message() -> bool:
+	# get json message: reset, step, close
+	var message = _get_dict_json_message()
+	if message["type"] == "close":
+		print("received close message, closing game")
+		get_tree().quit()
+		get_tree().set_pause(false)
+		return true
+
+	if message["type"] == "reset":
+		print("resetting all agents")
+		_reset_agents()
+		just_reset = true
+		get_tree().set_pause(false)
+		#print("resetting forcing draw")
+#        RenderingServer.force_draw()
+#        var obs = _get_obs_from_agents()
+#        print("obs ", obs)
+#        var reply = {
+#            "type": "reset",
+#            "obs": obs
+#        }
+#        _send_dict_as_json_message(reply)
+		return true
+
+	if message["type"] == "call":
+		var method = message["method"]
+		var returns = _call_method_on_agents(method)
+		var reply = {"type": "call", "returns": returns}
+		print("calling method from Python")
+		_send_dict_as_json_message(reply)
+		return handle_message()
+
+	if message["type"] == "action":
+		var action = message["action"]
+		_set_agent_actions(action, agents_training)
+		need_to_send_obs = true
+		get_tree().set_pause(false)
+		return true
+
+	print("message was not handled")
+	return false
+
+
+func _call_method_on_agents(method):
+	var returns = []
+	for agent in all_agents:
+		returns.append(agent.call(method))
+
+	return returns
+
+
+func _reset_agents_if_done(agents = all_agents):
+	for agent in agents:
+		if agent.get_done():
+			agent.set_done_false()
+
+
+func _reset_agents(agents = all_agents):
+	for agent in agents:
+		agent.needs_reset = true
+		#agent.reset()
+
+
+func _get_obs_from_agents(agents: Array = all_agents):
+	var obs = []
+	for agent in agents:
+		obs.append(agent.get_obs())
+	return obs
+
+
+func _get_reward_from_agents(agents: Array = agents_training):
+	var rewards = []
+	for agent in agents:
+		rewards.append(agent.get_reward())
+		agent.zero_reward()
+	return rewards
+
+
+func _get_info_from_agents(agents: Array = all_agents):
+	var info = []
+	for agent in agents:
+		info.append(agent.get_info())
+	return info
+
+
+func _get_done_from_agents(agents: Array = agents_training):
+	var dones = []
+	for agent in agents:
+		var done = agent.get_done()
+		if done:
+			agent.set_done_false()
+		dones.append(done)
+	return dones
+
+
+func _set_agent_actions(actions, agents: Array = all_agents):
+	for i in range(len(actions)):
+		agents[i].set_action(actions[i])
+
+
+func clamp_array(arr: Array, min: float, max: float):
+	var output: Array = []
+	for a in arr:
+		output.append(clamp(a, min, max))
+	return output
+
+
+## Save recorded export demos on window exit (Close game window instead of "Stop" button in Godot Editor)
+func _notification(what):
+	if demo_trajectories.size() == 0 or expert_demo_save_path.is_empty():
+		return
+
+	if what == NOTIFICATION_PREDELETE:
+		var json_string = JSON.stringify(demo_trajectories, "", false)
+		var file = FileAccess.open(expert_demo_save_path, FileAccess.WRITE)
+
+		if not file:
+			var error: Error = FileAccess.get_open_error()
+			assert(not error, "There was an error opening the file: %d" % error)
+
+		file.store_line(json_string)
+		var error = file.get_error()
+		assert(not error, "There was an error after trying to write to the file: %d" % error)

Game/addons/godot_rl_agents/sync.gd.uid

@@ -0,0 +1 @@
+uid://b4vk6sql5v1jl

Game/project.godot

@@ -21,6 +21,10 @@ run/main_scene="uid://bcq14356s3e2i"
 config/features=PackedStringArray("4.7", "Forward Plus")
 config/icon="res://icon.svg"
 
+[editor_plugins]
+
+enabled=PackedStringArray("res://addons/godot_rl_agents/plugin.cfg")
+
 [display]
 
 window/size/viewport_width=1920
@@ -114,3 +118,4 @@ roll_right={
 
 
 
+

Game/scenes/match.tscn

@@ -6,6 +6,7 @@
 
 [node name="Match" type="Node3D"]
 script = ExtResource("1_m")
+bot_model_path = "res://bots/rookie.json"
 
 [node name="Arena" parent="." instance=ExtResource("2_m")]
 

Game/scenes/training.tscn

@@ -0,0 +1,15 @@
+[gd_scene load_steps=4 format=3]
+
+[ext_resource type="Script" path="res://scripts/training_mode.gd" id="1_tr"]
+[ext_resource type="PackedScene" path="res://scenes/arena_01.tscn" id="2_tr"]
+[ext_resource type="Script" path="res://addons/godot_rl_agents/sync.gd" id="3_tr"]
+
+[node name="Training" type="Node3D"]
+script = ExtResource("1_tr")
+
+[node name="Arena" parent="." instance=ExtResource("2_tr")]
+
+[node name="Sync" type="Node" parent="."]
+script = ExtResource("3_tr")
+action_repeat = 8
+speed_up = 8.0

Game/scripts/ai_ship_controller.gd

@@ -0,0 +1,90 @@
+class_name AIShipController
+extends ShipController
+
+# Drives a ship from a trained self-play policy (see TRAINING.md). Builds the
+# same canonical observation as training (ShipObservations) and runs the
+# policy MLP in GDScript (PolicyNetwork) — the shipped bot has no Python,
+# .NET, or network dependency.
+#
+# Difficulty is (model, reaction_ticks, action_noise): weaker checkpoints make
+# easier bots outright, and the two knobs handicap a given model further —
+# slower reactions and noisier execution. Models live in res://bots/.
+
+@export_file("*.json") var model_path: String = ""
+# Decide a new action every N physics ticks, holding the last one between
+# decisions. 8 matches the training action_repeat; larger = slower reactions.
+@export_range(1, 60) var reaction_ticks: int = 8
+# Uniform noise magnitude added to each action axis (0 = play at full skill).
+@export_range(0.0, 1.0) var action_noise: float = 0.0
+
+var _policy: PolicyNetwork
+var _action := ShipAction.new()
+var _ticks_until_decision := 0
+
+var _ship: Ship
+var _opponent: Ship
+var _ball: RigidBody3D
+var _attack_goal_position: Vector3
+var _scene_refs_ready := false
+
+
+func _ready():
+	if not model_path.is_empty():
+		_policy = PolicyNetwork.load_from_file(model_path)
+
+
+func get_action() -> ShipAction:
+	if _policy == null:
+		return _action  # unloaded model: behaves like the inert placeholder
+	if not _scene_refs_ready and not _discover_scene_refs():
+		return _action
+
+	_ticks_until_decision -= 1
+	if _ticks_until_decision <= 0:
+		_ticks_until_decision = reaction_ticks
+		_decide()
+	return _action
+
+
+func _decide() -> void:
+	var obs := ShipObservations.build(_ship, _opponent, _ball, _attack_goal_position)
+	var out := _policy.forward(obs)
+	# Output layout matches the flattened training action space (Box(7)):
+	# thrust xyz, rotation xyz, turbo (> 0 means on).
+	_action.thrust = Vector3(
+		_axis(out[0]),
+		_axis(out[1]),
+		_axis(out[2])
+	)
+	_action.rotation = Vector3(
+		_axis(out[3]),
+		_axis(out[4]),
+		_axis(out[5])
+	)
+	_action.turbo = out[6] > 0.0
+
+
+func _axis(value: float) -> float:
+	if action_noise > 0.0:
+		value += randf_range(-action_noise, action_noise)
+	return clampf(value, -1.0, 1.0)
+
+
+# Find ship/ball/opponent/goal once everything is spawned. ShipAction axes
+# are body-frame so only observations need team context (ShipObservations).
+func _discover_scene_refs() -> bool:
+	_ship = get_parent() as Ship
+	if _ship == null or not is_inside_tree():
+		return false
+	_ball = get_tree().get_first_node_in_group("ball")
+	for node in get_tree().get_nodes_in_group("ship"):
+		if node != _ship:
+			_opponent = node
+			break
+	for goal in get_tree().get_nodes_in_group("goal"):
+		if goal.team == 1 - _ship.team:
+			_attack_goal_position = goal.global_position
+	if _ball == null:
+		return false
+	_scene_refs_ready = true
+	return true

Game/scripts/ai_ship_controller.gd.uid

@@ -0,0 +1 @@
+uid://4emuhiolrkb2

Game/scripts/match_mode.gd

@@ -1,15 +1,22 @@
 extends GameMode
 
 # Timed match: two teams, score tracking, kickoff resets after each goal.
-# The opponent ship is currently inert (base ShipController, zero action) —
-# it becomes the AI opponent once an AIShipController exists (see TODO.md),
-# and additional player ships once multiplayer lands.
+# The opponent is a trained AI bot when a policy model is configured
+# (see TRAINING.md for training and promoting models into res://bots/),
+# otherwise an inert placeholder ship.
 
 signal timer_updated(minutes: int, seconds: int)
 signal score_changed(score: Dictionary)
 
 @export var match_length_seconds := 150.0
 
+@export_group("AI opponent")
+# Trained policy for the opponent; empty = inert placeholder ship.
+@export_file("*.json") var bot_model_path: String = ""
+# Difficulty handicaps, applied on top of the model (see AIShipController).
+@export_range(1, 60) var bot_reaction_ticks: int = 8
+@export_range(0.0, 1.0) var bot_action_noise: float = 0.0
+
 var score := {0: 0, 1: 0}
 var match_timer: Timer
 
@@ -18,7 +25,7 @@ func _start() -> void:
 	spawn_ball()
 	var player_ship := spawn_ship(0, 0, PlayerShipController.new())
 	spawn_camera_rig(player_ship)
-	spawn_ship(1, 0, ShipController.new())  # inert placeholder opponent
+	spawn_ship(1, 0, _make_opponent_controller())
 
 	match_timer = Timer.new()
 	match_timer.one_shot = true
@@ -28,6 +35,18 @@ func _start() -> void:
 	match_timer.start()
 
 
+func _make_opponent_controller() -> ShipController:
+	if not bot_model_path.is_empty() and FileAccess.file_exists(bot_model_path):
+		var bot := AIShipController.new()
+		bot.model_path = bot_model_path
+		bot.reaction_ticks = bot_reaction_ticks
+		bot.action_noise = bot_action_noise
+		return bot
+	if not bot_model_path.is_empty():
+		push_warning("MatchMode: bot model not found at %s, spawning inert opponent" % bot_model_path)
+	return ShipController.new()  # inert placeholder
+
+
 func _process(_delta):
 	if match_timer and match_timer.time_left > 0:
 		var remaining := ceili(match_timer.time_left)

Game/scripts/policy_network.gd

@@ -0,0 +1,73 @@
+class_name PolicyNetwork
+extends RefCounted
+
+# Minimal MLP forward pass for running trained policies in pure GDScript —
+# no .NET build or ONNX runtime needed. Weights come from a JSON file written
+# by training/export_policy.py (see TRAINING.md). The policy net is tiny
+# (31 → 64 → 64 → 7 by default), and the bot only thinks every few physics
+# ticks, so GDScript is plenty fast.
+#
+# JSON shape:
+# {
+#   "input_size": 31,
+#   "layers": [
+#     {"weights": [[out x in floats]], "biases": [out floats], "activation": "tanh" | "linear"},
+#     ...
+#   ]
+# }
+
+var input_size: int = 0
+var _layers: Array = []
+
+
+static func load_from_file(path: String) -> PolicyNetwork:
+	if not FileAccess.file_exists(path):
+		push_error("PolicyNetwork: model file not found: %s" % path)
+		return null
+	var text := FileAccess.get_file_as_string(path)
+	var data: Variant = JSON.parse_string(text)
+	if data == null or not (data is Dictionary) or not data.has("layers"):
+		push_error("PolicyNetwork: invalid model file: %s" % path)
+		return null
+
+	var net := PolicyNetwork.new()
+	net.input_size = int(data.get("input_size", 0))
+	for layer in data["layers"]:
+		# Flatten each layer's weights into a PackedFloat64Array for speed
+		var out_size: int = layer["biases"].size()
+		var in_size: int = layer["weights"][0].size()
+		var flat := PackedFloat64Array()
+		flat.resize(out_size * in_size)
+		var i := 0
+		for row in layer["weights"]:
+			for value in row:
+				flat[i] = value
+				i += 1
+		var biases := PackedFloat64Array(layer["biases"])
+		net._layers.append({
+			"weights": flat,
+			"biases": biases,
+			"in_size": in_size,
+			"out_size": out_size,
+			"tanh": layer.get("activation", "linear") == "tanh",
+		})
+	return net
+
+
+func forward(observation: Array) -> Array:
+	var x := PackedFloat64Array(observation)
+	for layer in _layers:
+		var in_size: int = layer["in_size"]
+		var out_size: int = layer["out_size"]
+		var weights: PackedFloat64Array = layer["weights"]
+		var biases: PackedFloat64Array = layer["biases"]
+		var y := PackedFloat64Array()
+		y.resize(out_size)
+		for row in out_size:
+			var sum := biases[row]
+			var offset := row * in_size
+			for col in in_size:
+				sum += weights[offset + col] * x[col]
+			y[row] = tanh(sum) if layer["tanh"] else sum
+		x = y
+	return Array(x)

Game/scripts/policy_network.gd.uid

@@ -0,0 +1 @@
+uid://biixjudn05ib2

Game/scripts/rl_ship_controller.gd

@@ -0,0 +1,14 @@
+class_name RLShipController
+extends ShipController
+
+# Controller for externally-driven ships (RL training and trained-policy
+# inference). Something else — a ShipAIController during training, an
+# AIShipController at play time — writes into `action`; the ship pulls it
+# each physics tick like any other controller. The ship never knows it is
+# being trained.
+
+var action: ShipAction = ShipAction.new()
+
+
+func get_action() -> ShipAction:
+	return action

Game/scripts/rl_ship_controller.gd.uid

@@ -0,0 +1 @@
+uid://bctvr5djofbrm

Game/scripts/ship_ai_controller.gd

@@ -0,0 +1,90 @@
+class_name ShipAIController
+extends AIController3D
+
+# Training-side bridge between godot_rl_agents and a ship. This is the only
+# class that touches plugin types (AIController3D / the Sync node protocol) —
+# everything else stays behind the ShipController seam: actions received from
+# the trainer are written into an RLShipController, which the ship pulls like
+# any other controller.
+#
+# Action space is ShipAction verbatim: 6 continuous axes (thrust xyz,
+# rotation xyz, each -1..1) + binary turbo. ShipAction axes are ship-local
+# (body frame), so they need no team mirroring — only observations do
+# (see ShipObservations.canon).
+
+# Reward shaping weights. Dense terms accrue per physics tick (60 sim-ticks
+# per sim-second); event terms fire once. Exported so tuning needs no code
+# edits. Goal rewards are added by TrainingMode, which owns goal events.
+@export var ball_touch_reward := 0.1
+@export var velocity_to_ball_weight := 0.001
+@export var ball_velocity_to_goal_weight := 0.004
+
+var ship: Ship
+var rl_controller: RLShipController
+var ball: RigidBody3D
+var opponent: Ship
+var attack_goal_position: Vector3
+
+
+# Wire up references after the ship is spawned. `attack_goal` is the goal
+# this ship scores into (goal.team == opponent's team).
+func setup(p_ship: Ship, p_rl_controller: RLShipController, p_ball: RigidBody3D, p_opponent: Ship, p_attack_goal_position: Vector3) -> void:
+	ship = p_ship
+	rl_controller = p_rl_controller
+	ball = p_ball
+	opponent = p_opponent
+	attack_goal_position = p_attack_goal_position
+	init(ship)
+
+	# Contact monitoring for the ball-touch reward (training-only cost;
+	# the shipped game leaves contact_monitor off).
+	ship.contact_monitor = true
+	ship.max_contacts_reported = 8
+	ship.body_entered.connect(_on_ship_body_entered)
+
+
+func get_obs() -> Dictionary:
+	return {"obs": ShipObservations.build(ship, opponent, ball, attack_goal_position)}
+
+
+func get_reward() -> float:
+	return reward
+
+
+func get_action_space() -> Dictionary:
+	return {
+		"thrust": {"size": 3, "action_type": "continuous"},
+		"rotation": {"size": 3, "action_type": "continuous"},
+		"turbo": {"size": 2, "action_type": "discrete"},
+	}
+
+
+func set_action(action) -> void:
+	var thrust: Array = action["thrust"]
+	var rot: Array = action["rotation"]
+	rl_controller.action.thrust = Vector3(thrust[0], thrust[1], thrust[2])
+	rl_controller.action.rotation = Vector3(rot[0], rot[1], rot[2])
+	rl_controller.action.turbo = int(action["turbo"]) == 1
+
+
+func _physics_process(delta):
+	super(delta)
+	if not is_instance_valid(ship) or not is_instance_valid(ball):
+		return
+
+	# Dense shaping: own velocity toward the ball
+	var to_ball := ball.global_position - ship.global_position
+	if to_ball.length_squared() > 0.0001:
+		var closing_speed := ship.linear_velocity.dot(to_ball.normalized())
+		reward += velocity_to_ball_weight * closing_speed / ship.max_speed
+
+	# Dense shaping: ball velocity toward the goal we attack
+	var ball_to_goal := attack_goal_position - ball.global_position
+	if ball_to_goal.length_squared() > 0.0001:
+		var ball_progress := ball.linear_velocity.dot(ball_to_goal.normalized())
+		reward += ball_velocity_to_goal_weight * ball_progress / ShipObservations.BALL_SPEED_SCALE
+
+
+func _on_ship_body_entered(body: Node) -> void:
+	if body.is_in_group("ball"):
+		reward += ball_touch_reward

Game/scripts/ship_ai_controller.gd.uid

@@ -0,0 +1 @@
+uid://bql1ixmk23u77

Game/scripts/ship_observations.gd

@@ -0,0 +1,69 @@
+class_name ShipObservations
+extends RefCounted
+
+# Canonical, team-relative observation builder. Shared by training
+# (ShipAIController) and in-game inference (AIShipController) so a trained
+# policy sees byte-identical inputs in both contexts — do not fork this logic.
+#
+# Self-play trick: observations for team 1 are rotated 180° about Y
+# (x → -x, z → -z), so every ship perceives itself attacking toward -Z
+# regardless of which side it spawned on. One policy can then play both teams.
+# The same rotation must be inverted when interpreting actions (see canon —
+# it is its own inverse).
+
+# Normalization scales. Arena bounds: goals at z ≈ ±15.56, ship spawns at
+# z = ±12; positions are soft-normalized to roughly [-1, 1].
+const POSITION_SCALE := Vector3(20.0, 10.0, 20.0)
+const BALL_SPEED_SCALE := 30.0
+const GOAL_DISTANCE_SCALE := 40.0
+
+# Number of floats build() returns; the policy input size.
+const SIZE := 31
+
+
+# 180° rotation about Y for team 1; identity for team 0. A proper rotation
+# (preserves handedness), and its own inverse — used for both observations
+# and mapping canonical-frame actions back to world intent.
+static func canon(v: Vector3, team: int) -> Vector3:
+	return v if team == 0 else Vector3(-v.x, v.y, -v.z)
+
+
+# attack_goal_position: centre of the goal this ship is trying to score in
+# (the goal whose `team` == the opponent's team).
+static func build(ship: Ship, opponent: Ship, ball: RigidBody3D, attack_goal_position: Vector3) -> Array:
+	var team := ship.team
+	var obs := []
+
+	# Own kinematics
+	_append(obs, canon(ship.global_position, team) / POSITION_SCALE)
+	_append(obs, canon(-ship.global_transform.basis.z, team))  # forward
+	_append(obs, canon(ship.global_transform.basis.y, team))   # up
+	_append(obs, canon(ship.linear_velocity, team) / ship.max_speed)
+	_append(obs, canon(ship.angular_velocity, team) / ship.max_angular_speed)
+
+	# Ball, relative to self
+	var ball_rel := ball.global_position - ship.global_position
+	_append(obs, canon(ball_rel, team) / POSITION_SCALE)
+	_append(obs, canon(ball.linear_velocity, team) / BALL_SPEED_SCALE)
+
+	# Opponent, relative to self (zeros if absent, e.g. a 1-ship drill)
+	if is_instance_valid(opponent):
+		var opp_rel := opponent.global_position - ship.global_position
+		_append(obs, canon(opp_rel, team) / POSITION_SCALE)
+		_append(obs, canon(opponent.linear_velocity, team) / ship.max_speed)
+	else:
+		_append(obs, Vector3.ZERO)
+		_append(obs, Vector3.ZERO)
+
+	# Goal we are attacking, relative to self
+	var goal_rel := attack_goal_position - ship.global_position
+	_append(obs, canon(goal_rel, team) / POSITION_SCALE)
+	obs.append(goal_rel.length() / GOAL_DISTANCE_SCALE)
+
+	return obs
+
+
+static func _append(obs: Array, v: Vector3) -> void:
+	obs.append(v.x)
+	obs.append(v.y)
+	obs.append(v.z)

Game/scripts/ship_observations.gd.uid

@@ -0,0 +1 @@
+uid://bqux8vc4eou73

Game/scripts/training_mode.gd

@@ -0,0 +1,244 @@
+class_name TrainingMode
+extends GameMode
+
+# Headless self-play training mode: two RL-driven ships, no HUD, no camera.
+# The scene also contains the godot_rl_agents Sync node, which speaks TCP to
+# the Python trainer; this mode owns the environment rules — episodes, goal
+# rewards, and randomized episode-start states (the RLGym "state setter"
+# lesson: varied starts massively speed up learning versus kickoff-only).
+#
+# Run: godot --headless --path Game res://scenes/training.tscn
+# (started automatically by training/train.py; boots into idle ships with a
+# warning if no trainer is listening).
+#
+# Eval mode (used by training/evaluate.py): pass --eval_model_a=<path> and
+# --eval_model_b=<path> (+ optional --eval_episodes=N) and both ships are
+# instead driven by those exported policies via AIShipController; each episode
+# ends at the first goal (or a draw on timeout), and a final machine-readable
+# "EVAL_RESULT {...}" line is printed before quitting.
+
+@export var episode_length_seconds := 30.0
+@export var goal_reward := 10.0
+
+# Episode-start state mix; remaining probability = fully random state.
+@export_range(0.0, 1.0) var kickoff_state_chance := 0.2
+@export_range(0.0, 1.0) var ball_near_goal_chance := 0.2
+
+# Placement bounds, inset from the arena (goals at z ≈ ±15.56).
+const FIELD_HALF_X := 10.0
+const FIELD_HALF_Z := 13.0
+const FIELD_MIN_Y := 1.5
+const FIELD_MAX_Y := 8.0
+const MAX_RANDOM_BALL_SPEED := 12.0
+const MAX_RANDOM_SHIP_SPEED := 8.0
+
+# Sim runs at 60 physics ticks per sim-second regardless of speedup.
+const TICKS_PER_SIM_SECOND := 60.0
+
+# The arena has no walls/ceiling yet (see TODO.md), so an untrained policy can
+# simply fly away. Leaving this volume ends the episode with a small penalty.
+const BOUNDS_HALF_X := 25.0
+const BOUNDS_HALF_Z := 25.0
+const BOUNDS_MIN_Y := -5.0
+const BOUNDS_MAX_Y := 25.0
+@export var out_of_bounds_penalty := 1.0
+
+var _agents: Array[ShipAIController] = []
+
+# Eval mode state (see header comment)
+var _eval := false
+var _eval_models: Array[String] = ["", ""]
+var _eval_episodes := 20
+var _eval_goals := {0: 0, 1: 0}
+var _eval_draws := 0
+var _eval_episodes_done := 0
+var _episode_ticks := 0
+
+
+func _start() -> void:
+	_parse_eval_args()
+	spawn_ball()
+	if _eval:
+		for team in [0, 1]:
+			var bot := AIShipController.new()
+			bot.model_path = _eval_models[team]
+			spawn_ship(team, 0, bot)
+		return
+	var ship_team0 := spawn_ship(0, 0, RLShipController.new())
+	var ship_team1 := spawn_ship(1, 0, RLShipController.new())
+	_attach_agent(ship_team0, ship_team1)
+	_attach_agent(ship_team1, ship_team0)
+
+
+func _parse_eval_args() -> void:
+	var args := {}
+	for argument in OS.get_cmdline_args():
+		if argument.begins_with("--") and argument.find("=") > -1:
+			var key_value := argument.lstrip("--").split("=", true, 1)
+			args[key_value[0]] = key_value[1]
+	if args.has("eval_model_a") and args.has("eval_model_b"):
+		_eval = true
+		_eval_models[0] = args["eval_model_a"]
+		_eval_models[1] = args["eval_model_b"]
+		_eval_episodes = int(args.get("eval_episodes", str(_eval_episodes)))
+
+
+func _attach_agent(ship: Ship, opponent: Ship) -> void:
+	var agent := ShipAIController.new()
+	agent.name = "ShipAIController"
+	agent.reset_after = int(episode_length_seconds * TICKS_PER_SIM_SECOND)
+	ship.add_child(agent)
+	agent.setup(ship, ship.controller as RLShipController, ball, opponent, _attack_goal_position(ship.team))
+	_agents.append(agent)
+
+
+# The goal this team scores into: the one the opponent defends/concedes.
+func _attack_goal_position(team: int) -> Vector3:
+	for goal in arena.get_goals():
+		if goal.team == 1 - team:
+			return goal.global_position
+	push_error("TrainingMode: no goal found for team %d to attack" % team)
+	return Vector3.ZERO
+
+
+func _physics_process(_delta):
+	if _eval:
+		_episode_ticks += 1
+		for ship in ships:
+			if is_instance_valid(ship) and _out_of_bounds(ship.global_position):
+				_place_body(ship, _ship_spawn_transforms[ship], Vector3.ZERO, Vector3.ZERO)
+		var ball_lost := is_instance_valid(ball) and _out_of_bounds(ball.global_position)
+		if _episode_ticks > int(episode_length_seconds * TICKS_PER_SIM_SECOND) or ball_lost:
+			_eval_draws += 1
+			_end_eval_episode()
+		return
+
+	# Both trainer-requested resets and truncation (reset_after ticks elapsed)
+	# surface as needs_reset. Only truncation is an episode end the trainer
+	# must be told about via done — a trainer-requested reset already knows.
+	var needs_reset := false
+	var truncated := false
+	for agent in _agents:
+		needs_reset = needs_reset or agent.needs_reset
+		truncated = truncated or agent.n_steps > agent.reset_after
+	if needs_reset:
+		if truncated:
+			for agent in _agents:
+				agent.done = true
+		_reset_episode()
+		return
+
+	# A ship leaving the play volume is penalized and respawned at its kickoff
+	# spawn (the episode continues); a lost ball ends the episode for both.
+	for agent in _agents:
+		if _out_of_bounds(agent.ship.global_position):
+			agent.reward -= out_of_bounds_penalty
+			_place_body(agent.ship, _ship_spawn_transforms[agent.ship], Vector3.ZERO, Vector3.ZERO)
+	if is_instance_valid(ball) and _out_of_bounds(ball.global_position) and not _agents.is_empty():
+		for agent in _agents:
+			agent.done = true
+		_reset_episode()
+
+
+func _out_of_bounds(position: Vector3) -> bool:
+	return absf(position.x) > BOUNDS_HALF_X \
+		or absf(position.z) > BOUNDS_HALF_Z \
+		or position.y < BOUNDS_MIN_Y \
+		or position.y > BOUNDS_MAX_Y
+
+
+func _on_goal_scored(conceding_team: int) -> void:
+	if _eval:
+		_eval_goals[1 - conceding_team] += 1
+		_end_eval_episode()
+		return
+	for agent in _agents:
+		agent.reward += goal_reward if agent.ship.team != conceding_team else -goal_reward
+		agent.done = true
+	_reset_episode()
+
+
+func _end_eval_episode() -> void:
+	_eval_episodes_done += 1
+	_episode_ticks = 0
+	if _eval_episodes_done >= _eval_episodes:
+		print("EVAL_RESULT " + JSON.stringify({
+			"model_a": _eval_models[0],
+			"model_b": _eval_models[1],
+			"episodes": _eval_episodes_done,
+			"goals_a": _eval_goals[0],
+			"goals_b": _eval_goals[1],
+			"draws": _eval_draws,
+		}))
+		get_tree().quit()
+		return
+	# Randomized states (not kickoff): deterministic policies would otherwise
+	# replay the identical episode every time.
+	_reset_episode()
+
+
+func _reset_episode() -> void:
+	for agent in _agents:
+		agent.reset()
+
+	var roll := randf()
+	if roll < kickoff_state_chance:
+		reset_ball()
+		reset_ships()
+	elif roll < kickoff_state_chance + ball_near_goal_chance:
+		_place_ships_random()
+		_place_ball_near_goal()
+	else:
+		_place_ships_random()
+		_place_ball_random()
+
+
+func _place_ball_random() -> void:
+	var velocity := _random_direction() * randf_range(0.0, MAX_RANDOM_BALL_SPEED)
+	_place_body(ball, Transform3D(Basis.IDENTITY, _random_position()), velocity, Vector3.ZERO)
+
+
+# Attacking/defending drill states: ball close to a goal, moving toward it.
+func _place_ball_near_goal() -> void:
+	var goals := arena.get_goals()
+	var goal: Goal = goals[randi() % goals.size()]
+	var toward_centre := -signf(goal.global_position.z)
+	var position := Vector3(
+		randf_range(-4.0, 4.0),
+		randf_range(FIELD_MIN_Y, 4.0),
+		goal.global_position.z + toward_centre * randf_range(3.0, 6.0)
+	)
+	var to_goal := (goal.global_position - position).normalized()
+	var velocity := (to_goal + _random_direction() * 0.3).normalized() * randf_range(2.0, MAX_RANDOM_BALL_SPEED)
+	_place_body(ball, Transform3D(Basis.IDENTITY, position), velocity, Vector3.ZERO)
+
+
+func _place_ships_random() -> void:
+	for ship in ships:
+		var orientation := Basis.from_euler(Vector3(
+			randf_range(-0.4, 0.4),
+			randf_range(-PI, PI),
+			randf_range(-0.4, 0.4)
+		))
+		var velocity := _random_direction() * randf_range(0.0, MAX_RANDOM_SHIP_SPEED)
+		_place_body(ship, Transform3D(orientation, _random_position()), velocity, Vector3.ZERO)
+
+
+func _random_position() -> Vector3:
+	return Vector3(
+		randf_range(-FIELD_HALF_X, FIELD_HALF_X),
+		randf_range(FIELD_MIN_Y, FIELD_MAX_Y),
+		randf_range(-FIELD_HALF_Z, FIELD_HALF_Z)
+	)
+
+
+func _random_direction() -> Vector3:
+	var direction := Vector3(randf_range(-1, 1), randf_range(-1, 1), randf_range(-1, 1))
+	return direction.normalized() if direction.length_squared() > 0.001 else Vector3.FORWARD
+
+
+func _place_body(body: RigidBody3D, to: Transform3D, linear_velocity: Vector3, angular_velocity: Vector3) -> void:
+	# Deferred: a RigidBody3D transform can't be set mid-physics-step
+	body.set_deferred("global_transform", to)
+	body.set_deferred("linear_velocity", linear_velocity)
+	body.set_deferred("angular_velocity", angular_velocity)

Game/scripts/training_mode.gd.uid

@@ -0,0 +1 @@
+uid://d2kfcobanp6iv

TODO.md

@@ -4,11 +4,13 @@ Deferred work, in rough priority order. The current architecture (ShipAction/Shi
 
 ## AI opponent (reinforcement learning)
 
-- [ ] `AIShipController extends ShipController` — produces a `ShipAction` per physics tick from observations instead of keyboard input.
-- [ ] Observation builder: self ship state (position, orientation, velocities) + ball state (group `"ball"`) + goal positions/teams (group `"goal"`), normalized for the policy.
-- [ ] Reward shaping: goals scored/conceded, ball touches, ball-toward-opponent-goal velocity, etc.
-- [ ] Headless training scene: a `GameMode` subclass with no HUD/camera, run via `godot --headless`, stepping the sim for training (consider godot-rl-agents or a custom socket bridge).
-- [ ] Swap the inert placeholder opponent in Match mode for the trained `AIShipController`.
+The training pipeline is built — see `TRAINING.md` (self-play PPO via the vendored godot_rl_agents bridge, JSON policy export, in-game GDScript inference, eval ladder). Remaining:
+
+- [ ] Long training runs on the Linux/3090 box to produce actually-good bots; promote checkpoints into `Game/bots/` as `easy`/`medium`/`hard` tiers.
+- [ ] Frozen-opponent league: train the live policy against a pool of past exported checkpoints (via `AIShipController` on the opponent ship in TrainingMode) to prevent self-play strategy collapse on long runs.
+- [ ] Richer state setter / curriculum: aerial states, wall plays, rebound scenarios as skill grows.
+- [ ] Main-menu difficulty picker (Match already takes `bot_model_path`/`bot_reaction_ticks`/`bot_action_noise` exports).
+- [ ] Optional: exported headless Linux build for faster parallel training instances (train.py currently runs the project from source, which is fine but re-parses scripts per instance).
 
 ## Match mode polish
 
@@ -28,3 +30,4 @@ Deferred work, in rough priority order. The current architecture (ShipAction/Shi
 
 - [ ] No autoloads yet by design — add a singleton only when cross-scene state is actually needed (e.g. passing match settings/results between menu, match, and results screens).
 - [ ] More arenas: `arena_01.tscn` is the template — an arena is terrain + lighting + two team-tagged goals + spawn markers, with no rules or state.
+- [ ] Enclose the arena (walls/ceiling): ships and ball can currently fly out of the play volume. TrainingMode papers over this with an out-of-bounds episode-end + penalty; a real enclosed arena fixes it for players too and the training guard can then be removed.

TRAINING.md

@@ -0,0 +1,140 @@
+# Training the AI bot
+
+Cosmic Clash bots are trained with reinforcement learning (self-play PPO): two
+ships in a headless arena share one policy that learns by playing against
+itself. Training runs in Python ([Godot RL Agents](https://github.com/edbeeching/godot_rl_agents)
+bridge + Stable-Baselines3); the trained policy is exported to a small JSON
+file and runs **inside the game** in pure GDScript — shipped bots need no
+Python, no .NET, no network.
+
+## How it fits together
+
+- `Game/scenes/training.tscn` + `scripts/training_mode.gd` — headless self-play
+  environment: two RL ships, randomized episode starts, goal rewards. Contains
+  the vendored godot_rl_agents `Sync` node that talks TCP to the trainer.
+- `scripts/ship_ai_controller.gd` — training-side bridge (observations,
+  rewards, action mapping). `scripts/ship_observations.gd` is the *shared*
+  observation builder — training and in-game inference must stay identical,
+  so never fork it.
+- `training/train.py` — PPO trainer; launches N parallel headless Godot
+  instances (2 agents each).
+- `training/export_policy.py` — SB3 checkpoint → JSON policy for the game.
+- `scripts/ai_ship_controller.gd` + `scripts/policy_network.gd` — in-game
+  inference (GDScript MLP forward pass).
+- `training/evaluate.py` — pits two exported policies against each other and
+  appends to `training/eval_history.json`.
+
+## Hardware
+
+The environment is our own headless Godot sim — fully cross-platform:
+
+- **Any machine (e.g. the M4 Mac mini)**: fine for pipeline development,
+  smoke runs, and short experiments. Env stepping is CPU-bound; the policy is
+  a small MLP, so even CPU-only PPO updates are cheap.
+- **Linux + NVIDIA GPU (e.g. the RTX 3090 box)**: recommended for real
+  multi-hour/overnight runs. PyTorch CUDA works out of the box; more CPU
+  cores also mean more parallel Godot instances (`--n-parallel`).
+
+There is no hard GPU requirement (unlike Rocket League tooling) — a GPU
+mainly speeds up learning updates on long runs.
+
+## Setup
+
+Needs Python 3.10+ and a Godot 4.7 binary.
+
+```bash
+cd training
+python3.12 -m venv .venv          # macOS: brew install python@3.12
+.venv/bin/pip install -r requirements.txt
+```
+
+On Linux, download the Godot 4.7 Linux binary and point at it:
+
+```bash
+export GODOT_BIN=~/godot/Godot_v4.7.1-stable_linux.x86_64
+```
+
+(macOS default is `/Applications/Godot.app/Contents/MacOS/Godot`; override
+with `GODOT_BIN` or `--godot_bin` if yours lives elsewhere.)
+
+## Run a training session
+
+```bash
+cd training
+.venv/bin/python train.py --experiment run01 --timesteps 20000000 --n-parallel 6 --speedup 16
+```
+
+- Checkpoints land in `training/checkpoints/run01/` every `--checkpoint-every`
+  steps (default 100k), plus `final.zip` on exit (also written on Ctrl-C).
+- Resume with `--resume checkpoints/run01/final.zip`.
+- `--n-parallel` = Godot instances (2 agents each). Scale with CPU cores.
+- `--speedup` = in-engine physics speedup. Raise until CPU saturates.
+- `--wandb` mirrors logs to Weights & Biases (`pip install wandb` first).
+
+Expect the smoke-run scale (~100k steps) to only learn crude ball-chasing;
+real behaviour needs tens of millions of steps (hours on the 3090 box).
+
+### Watch progress
+
+```bash
+.venv/bin/tensorboard --logdir training/logs
+```
+
+Key curves: `rollout/ep_rew_mean` (should trend up), `rollout/ep_len_mean`
+(should trend *down* from 225 as goals end episodes early — 225 action steps
+= the 30s episode timeout).
+
+### Reward/observation tuning
+
+Reward weights are exported vars on `ShipAIController` (goal reward on
+`TrainingMode`) — tune in `training.tscn`/scripts without touching the
+trainer. If you change the *observation* layout (`ship_observations.gd`),
+old checkpoints/exports become incompatible: retrain, and bump a note in
+your experiment name.
+
+## Export a checkpoint into the game
+
+```bash
+cd training
+.venv/bin/python export_policy.py checkpoints/run01/final.zip ../Game/bots/hard.json
+```
+
+The exporter runs a parity check (JSON forward pass vs SB3 prediction) before
+writing. Models live in `Game/bots/`.
+
+## Evaluate progress between checkpoints
+
+TensorBoard shows learning, but "is the new checkpoint actually *better*?"
+needs head-to-head play:
+
+```bash
+.venv/bin/python export_policy.py checkpoints/run01/ppo_5000000_steps.zip /tmp/candidate.json
+.venv/bin/python evaluate.py ../Game/bots/hard.json /tmp/candidate.json --episodes 40
+```
+
+Golden-goal episodes (first goal wins, timeout = draw), sides swapped halfway
+for fairness, using the exact inference path that ships in-game. Every run
+appends to `training/eval_history.json` — the long-term progress record.
+Evaluate each new candidate against the previous promoted bot and a fixed
+early reference to see absolute progress over time.
+
+## Difficulty tiers
+
+A bot is `(model, reaction_ticks, action_noise)` — configured on the Match
+mode (`bot_model_path`, `bot_reaction_ticks`, `bot_action_noise` in
+`match.tscn`) or any `AIShipController`:
+
+- **Model**: the main lever. An early checkpoint *is* an easy bot — promote
+  e.g. `easy.json` / `medium.json` / `hard.json` from different stages of one
+  training run (verify the gaps with `evaluate.py`).
+- **reaction_ticks** (default 8 = training cadence): higher = slower
+  reactions, easier.
+- **action_noise**: adds execution error, easier.
+
+## Self-play notes
+
+Both ships share the live policy (mirrored, team-relative observations — see
+`ship_observations.gd`), so training is always against the current self.
+Fixed-opponent training against frozen checkpoints (league play, to avoid
+strategy collapse on long runs) is deferred — see TODO.md; the pieces
+(exported JSON bots + `AIShipController`) already exist.

training/cosmic_env.py

@@ -0,0 +1,76 @@
+"""Godot RL Agents environment wrappers that run Cosmic Clash from source.
+
+Stock GodotEnv expects an *exported* game executable and rewrites its path
+per-platform. These subclasses launch the project straight from the repo with
+a Godot binary instead (no export step), pointing it at the training scene.
+Each Godot instance contributes two agents (one ship per team) that share the
+learning policy: self-play by construction.
+"""
+
+import pathlib
+import subprocess
+
+from godot_rl.core.godot_env import GodotEnv
+from godot_rl.wrappers.stable_baselines_wrapper import StableBaselinesGodotEnv
+
+REPO_ROOT = pathlib.Path(__file__).resolve().parent.parent
+GAME_DIR = REPO_ROOT / "Game"
+TRAINING_SCENE = "res://scenes/training.tscn"
+
+
+class CosmicClashEnv(GodotEnv):
+    """GodotEnv that launches `godot --path Game res://scenes/training.tscn`."""
+
+    # env_path is a Godot binary, not an exported game: skip the suffix and
+    # platform checks stock GodotEnv applies to exported executables.
+    def _set_platform_suffix(self, env_path: str) -> str:
+        return env_path
+
+    def check_platform(self, filename: str):
+        pass
+
+    def _launch_env(self, env_path, port, show_window, framerate, seed, action_repeat, speedup, **kwargs):
+        # sync.gd reads --key=value pairs from the raw command line; they must
+        # NOT go after a `--` separator or OS.get_cmdline_args() drops them.
+        cmd = [
+            env_path,
+            "--path",
+            str(GAME_DIR),
+            TRAINING_SCENE,
+            f"--port={port}",
+            f"--env_seed={seed}",
+        ]
+        if not show_window:
+            cmd += ["--headless", "--disable-render-loop"]
+        if framerate is not None:
+            cmd += ["--fixed-fps", str(framerate)]
+        if action_repeat is not None:
+            cmd.append(f"--action_repeat={action_repeat}")
+        if speedup is not None:
+            cmd.append(f"--speedup={speedup}")
+        for key, value in kwargs.items():
+            cmd.append(f"--{key}={value}")
+        self.proc = subprocess.Popen(cmd, start_new_session=True)
+
+
+class CosmicClashVecEnv(StableBaselinesGodotEnv):
+    """SB3 VecEnv over N parallel CosmicClashEnv instances.
+
+    convert_action_space=True flattens the env's (Box(6), Discrete(2)) action
+    space into a single Box(7): thrust xyz, rotation xyz, turbo (>0 means on).
+    """
+
+    def __init__(self, godot_bin: str, n_parallel: int = 1, seed: int = 0, port: int = GodotEnv.DEFAULT_PORT, **kwargs):
+        self.envs = [
+            CosmicClashEnv(
+                env_path=godot_bin,
+                convert_action_space=True,
+                port=port + p,
+                seed=seed + p,
+                **kwargs,
+            )
+            for p in range(n_parallel)
+        ]
+        self.n_parallel = n_parallel
+        self._check_valid_action_space()
+        self.results = None

training/eval_history.json

@@ -0,0 +1,12 @@
+[
+  {
+    "timestamp": "2026-07-18T18:27:59+00:00",
+    "model_a": "/Users/jcreek/Documents/repos/GitHub/CosmicClash/Game/bots/rookie.json",
+    "model_b": "/Users/jcreek/Documents/repos/GitHub/CosmicClash/Game/bots/rookie.json",
+    "episodes": 6,
+    "wins_a": 1,
+    "wins_b": 0,
+    "draws": 5,
+    "win_rate_a": 0.167
+  }
+]

training/evaluate.py

@@ -0,0 +1,96 @@
+"""Pit two exported policies against each other and record the result.
+
+Uses the same in-Godot inference path that ships in the game
+(AIShipController + PolicyNetwork), so eval strength = in-game strength.
+Episodes are golden-goal: first goal wins, timeout is a draw. Half the
+episodes are played with sides swapped for fairness. Results are appended to
+eval_history.json — the bot-progress-over-time record.
+
+Example:
+    .venv/bin/python evaluate.py ../Game/bots/rookie.json checkpoints/run01/candidate.json --episodes 40
+"""
+
+import argparse
+import datetime
+import json
+import os
+import pathlib
+import subprocess
+
+TRAINING_DIR = pathlib.Path(__file__).resolve().parent
+GAME_DIR = TRAINING_DIR.parent / "Game"
+TRAINING_SCENE = "res://scenes/training.tscn"
+DEFAULT_GODOT_MACOS = "/Applications/Godot.app/Contents/MacOS/Godot"
+
+
+def run_half(godot_bin: str, model_a: str, model_b: str, episodes: int, speedup: int, seed: int) -> dict:
+    cmd = [
+        godot_bin,
+        "--path",
+        str(GAME_DIR),
+        TRAINING_SCENE,
+        "--headless",
+        "--disable-render-loop",
+        f"--eval_model_a={model_a}",
+        f"--eval_model_b={model_b}",
+        f"--eval_episodes={episodes}",
+        f"--speedup={speedup}",
+        f"--env_seed={seed}",
+    ]
+    timeout = episodes * 30 / speedup * 3 + 120  # worst case: all draws, plus margin
+    result = subprocess.run(cmd, capture_output=True, text=True, timeout=timeout)
+    for line in result.stdout.splitlines():
+        if line.startswith("EVAL_RESULT "):
+            return json.loads(line[len("EVAL_RESULT "):])
+    raise RuntimeError(f"No EVAL_RESULT in godot output:\n{result.stdout[-2000:]}\n{result.stderr[-2000:]}")
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("model_a", help="Path to first exported policy .json")
+    parser.add_argument("model_b", help="Path to second exported policy .json")
+    parser.add_argument("--episodes", type=int, default=20, help="Total episodes (split across side swap)")
+    parser.add_argument(
+        "--godot_bin",
+        default=os.environ.get("GODOT_BIN", DEFAULT_GODOT_MACOS),
+        help="Path to the Godot binary (or set GODOT_BIN)",
+    )
+    parser.add_argument("--speedup", type=int, default=16)
+    parser.add_argument("--history", default=str(TRAINING_DIR / "eval_history.json"))
+    args = parser.parse_args()
+
+    model_a = str(pathlib.Path(args.model_a).resolve())
+    model_b = str(pathlib.Path(args.model_b).resolve())
+    half = max(args.episodes // 2, 1)
+
+    # Half the episodes on each side to cancel any residual side asymmetry;
+    # different seeds so the halves see different randomized episode states.
+    first = run_half(args.godot_bin, model_a, model_b, half, args.speedup, seed=1)
+    second = run_half(args.godot_bin, model_b, model_a, half, args.speedup, seed=2)
+
+    record = {
+        "timestamp": datetime.datetime.now(datetime.timezone.utc).isoformat(timespec="seconds"),
+        "model_a": model_a,
+        "model_b": model_b,
+        "episodes": first["episodes"] + second["episodes"],
+        "wins_a": first["goals_a"] + second["goals_b"],
+        "wins_b": first["goals_b"] + second["goals_a"],
+        "draws": first["draws"] + second["draws"],
+    }
+    record["win_rate_a"] = round(record["wins_a"] / record["episodes"], 3)
+
+    history_path = pathlib.Path(args.history)
+    history = json.loads(history_path.read_text()) if history_path.exists() else []
+    history.append(record)
+    history_path.write_text(json.dumps(history, indent=2) + "\n")
+
+    print(
+        f"{pathlib.Path(model_a).name} vs {pathlib.Path(model_b).name} over {record['episodes']} episodes: "
+        f"{record['wins_a']}-{record['wins_b']} ({record['draws']} draws), "
+        f"win rate A = {record['win_rate_a']:.0%}"
+    )
+    print(f"Appended to {history_path}")
+
+
+if __name__ == "__main__":
+    main()

training/export_policy.py

@@ -0,0 +1,86 @@
+"""Export a trained SB3 checkpoint to the JSON format PolicyNetwork.gd loads.
+
+The exported file contains the deterministic policy MLP (obs -> action means);
+the game clamps outputs to [-1, 1] and treats the last value as turbo (> 0).
+A parity self-check compares the JSON forward pass against SB3's own
+deterministic prediction before writing.
+
+Example:
+    .venv/bin/python export_policy.py checkpoints/smoke/final.zip ../Game/bots/rookie.json
+"""
+
+import argparse
+import json
+import pathlib
+
+import numpy as np
+import torch
+from stable_baselines3 import PPO
+
+
+def linear_to_layer(linear: torch.nn.Linear, activation: str) -> dict:
+    return {
+        "weights": linear.weight.detach().cpu().numpy().tolist(),
+        "biases": linear.bias.detach().cpu().numpy().tolist(),
+        "activation": activation,
+    }
+
+
+def extract_layers(policy) -> list[dict]:
+    # Features extractor must be a passthrough (flatten) for this export to
+    # be faithful; it has no parameters for our flat "obs" Box space.
+    n_extractor_params = sum(p.numel() for p in policy.features_extractor.parameters())
+    assert n_extractor_params == 0, "features extractor has weights; export logic needs updating"
+
+    layers = []
+    modules = list(policy.mlp_extractor.policy_net)
+    for i, module in enumerate(modules):
+        if isinstance(module, torch.nn.Linear):
+            next_is_tanh = i + 1 < len(modules) and isinstance(modules[i + 1], torch.nn.Tanh)
+            assert next_is_tanh or i + 1 >= len(modules), (
+                f"unsupported activation after layer {i}: {modules[i + 1] if i + 1 < len(modules) else None}"
+            )
+            layers.append(linear_to_layer(module, "tanh" if next_is_tanh else "linear"))
+        elif not isinstance(module, torch.nn.Tanh):
+            raise AssertionError(f"unsupported module in policy net: {module}")
+    layers.append(linear_to_layer(policy.action_net, "linear"))
+    return layers
+
+
+def json_forward(layers: list[dict], obs: np.ndarray) -> np.ndarray:
+    x = obs
+    for layer in layers:
+        x = np.asarray(layer["weights"]) @ x + np.asarray(layer["biases"])
+        if layer["activation"] == "tanh":
+            x = np.tanh(x)
+    return x
+
+
+def main():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument("checkpoint", help="SB3 checkpoint .zip (e.g. checkpoints/smoke/final.zip)")
+    parser.add_argument("output", help="Output JSON path (e.g. ../Game/bots/rookie.json)")
+    args = parser.parse_args()
+
+    model = PPO.load(args.checkpoint, device="cpu")
+    policy = model.policy
+    layers = extract_layers(policy)
+    input_size = model.observation_space["obs"].shape[0]
+
+    # Parity check: JSON forward pass must match SB3's deterministic action
+    rng = np.random.default_rng(0)
+    for _ in range(16):
+        obs = rng.uniform(-1, 1, input_size).astype(np.float32)
+        expected, _ = model.predict({"obs": obs}, deterministic=True)
+        actual = np.clip(json_forward(layers, obs), -1.0, 1.0)
+        assert np.allclose(actual, expected, atol=1e-5), f"parity check failed: {actual} vs {expected}"
+
+    output = pathlib.Path(args.output)
+    output.parent.mkdir(parents=True, exist_ok=True)
+    with open(output, "w") as f:
+        json.dump({"input_size": int(input_size), "layers": layers}, f)
+    print(f"Exported {args.checkpoint} -> {output} (input size {input_size}, {len(layers)} layers, parity OK)")
+
+
+if __name__ == "__main__":
+    main()

training/requirements.txt

@@ -0,0 +1,5 @@
+godot-rl
+stable-baselines3
+tensorboard
+# Optional, for --wandb logging:
+# wandb

training/train.py

@@ -0,0 +1,105 @@
+"""Train the Cosmic Clash self-play PPO policy.
+
+Example (smoke run):
+    .venv/bin/python train.py --experiment smoke --timesteps 100000
+
+Long run on the Linux/CUDA box:
+    GODOT_BIN=~/godot/Godot_v4.7.1-stable_linux.x86_64 \
+        .venv/bin/python train.py --experiment run01 --timesteps 20000000 \
+        --n-parallel 6 --speedup 16
+
+See TRAINING.md at the repo root for the full workflow.
+"""
+
+import argparse
+import os
+import pathlib
+
+from stable_baselines3 import PPO
+from stable_baselines3.common.callbacks import CheckpointCallback
+from stable_baselines3.common.vec_env.vec_monitor import VecMonitor
+
+from cosmic_env import CosmicClashVecEnv
+
+TRAINING_DIR = pathlib.Path(__file__).resolve().parent
+DEFAULT_GODOT_MACOS = "/Applications/Godot.app/Contents/MacOS/Godot"
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description=__doc__)
+    parser.add_argument(
+        "--godot_bin",
+        default=os.environ.get("GODOT_BIN", DEFAULT_GODOT_MACOS),
+        help="Path to the Godot binary (or set GODOT_BIN)",
+    )
+    parser.add_argument("--experiment", default="default", help="Run name for logs/checkpoints")
+    parser.add_argument("--timesteps", type=int, default=200_000)
+    parser.add_argument("--n-parallel", type=int, default=2, help="Parallel Godot instances (2 agents each)")
+    parser.add_argument("--speedup", type=int, default=8, help="Physics speedup factor inside Godot")
+    parser.add_argument("--port", type=int, default=11008, help="Base TCP port (one per instance)")
+    parser.add_argument("--seed", type=int, default=0)
+    parser.add_argument("--resume", default=None, help="Checkpoint .zip to resume from")
+    parser.add_argument("--checkpoint-every", type=int, default=100_000, help="Timesteps between checkpoints")
+    parser.add_argument("--viz", action="store_true", help="Show game windows (debugging; slow)")
+    parser.add_argument("--wandb", action="store_true", help="Also log to Weights & Biases")
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+    log_dir = TRAINING_DIR / "logs"
+    checkpoint_dir = TRAINING_DIR / "checkpoints" / args.experiment
+    checkpoint_dir.mkdir(parents=True, exist_ok=True)
+
+    if args.wandb:
+        import wandb
+
+        wandb.init(project="cosmic-clash-rl", name=args.experiment, sync_tensorboard=True)
+
+    env = CosmicClashVecEnv(
+        godot_bin=args.godot_bin,
+        n_parallel=args.n_parallel,
+        seed=args.seed,
+        port=args.port,
+        show_window=args.viz,
+        speedup=args.speedup,
+    )
+    env = VecMonitor(env)
+
+    if args.resume:
+        model = PPO.load(args.resume, env=env, tensorboard_log=str(log_dir))
+        print(f"Resumed from {args.resume} at {model.num_timesteps} timesteps")
+    else:
+        model = PPO(
+            "MultiInputPolicy",
+            env,
+            verbose=1,
+            ent_coef=0.0001,
+            n_steps=256,
+            batch_size=256,
+            learning_rate=3e-4,
+            tensorboard_log=str(log_dir),
+        )
+
+    checkpoint_callback = CheckpointCallback(
+        save_freq=max(args.checkpoint_every // env.num_envs, 1),
+        save_path=str(checkpoint_dir),
+        name_prefix="ppo",
+    )
+
+    try:
+        model.learn(
+            args.timesteps,
+            callback=checkpoint_callback,
+            tb_log_name=args.experiment,
+            reset_num_timesteps=not args.resume,
+        )
+    finally:
+        final_path = checkpoint_dir / "final.zip"
+        model.save(str(final_path))
+        print(f"Saved {final_path}")
+        env.close()
+
+
+if __name__ == "__main__":
+    main()