Gymnasium environment for UTDG (Untitled Tower Defense Game)
Project description
UTDG Gymnasium Environment
A Gymnasium-compatible environment for training reinforcement learning agents on the Untitled Tower Defense Game (UTDG).
Overview
This package provides a WebSocket-based interface between the Godot game engine and Python, allowing you to train RL agents using popular frameworks like Stable-Baselines3 or Ray RLlib.
Features
- Gymnasium API: Standard
reset(),step(), andclose()methods - WebSocket Communication: Real-time bidirectional communication with Godot
- Flexible Action Space: Discrete actions for tower placement
- Rich Observations: Game state including gold, health, enemy/tower positions
- Easy Integration: Compatible with Stable-Baselines3, RLlib, and other RL libraries
Installation
From source:
cd python_env
pip install -e .
With training dependencies:
pip install -e ".[train]"
From requirements file:
pip install -r requirements.txt
Configuration System
UTDG uses Hydra for configuration management, providing a flexible and composable way to manage experiment parameters.
Configuration Structure
python_env/configs/
├── default.yaml # Main composition file (for testing)
├── default_training.yaml # Training composition file (includes PPO config)
├── env/
│ └── default.yaml # Environment settings
├── game/
│ ├── default.yaml # Game configuration (consolidated)
│ ├── enemy_spawn.yaml # Alternative enemy spawn config
│ ├── rewards.yaml # Alternative reward config
│ └── difficulty.yaml # Alternative difficulty config
└── training/
├── common.yaml # Common training settings
└── ppo.yaml # PPO hyperparameters
Note: simple_test.py uses default.yaml while train_ppo.py uses default_training.yaml which includes PPO hyperparameters by default.
Using Hydra CLI
Run scripts with default configuration:
# From python_env/ directory
python examples/simple_test.py
# From root UTDG directory
python python_env/examples/simple_test.py
Override specific parameters:
# Change game parameters
python examples/simple_test.py game.base_health=200 experiment.seed=123
# Configure runtime settings (Godot launch)
python examples/simple_test.py \
runtime.auto_launch=true \
runtime.godot_path="builds/UTDG-macOS.app/Contents/MacOS/UntitledTowerDefenseGame" \
runtime.episodes=5
# Training with custom hyperparameters
python examples/train_ppo.py \
ppo.learning_rate=0.0001 \
training.training.total_timesteps=500000
Multi-run for parameter sweeps:
python examples/simple_test.py -m \
game.difficulty_config.enemy_health_multiplier=1.0,1.5,2.0
Configuration Parameters
Experiment metadata (experiment.*):
name: Experiment nameseed: Random seed for reproducibilitylog_dir: Directory for logs
WebSocket connection (websocket.*):
host: WebSocket host (default: localhost)port: WebSocket port (default: 9090)timeout: Connection timeout in seconds
Runtime settings (runtime.*):
episodes: Number of episodes to rungodot_path: Path to Godot executableauto_launch: Auto-launch Godot if not runningheadless: Run Godot in headless mode
Game configuration (game.*):
base_health: Starting base healthstarting_gold: Starting gold amountenemy_spawn_config.*: Enemy spawning parametersreward_config.*: Reward shaping settingsdifficulty_config.*: Difficulty multipliers
Environment (env.*):
observation_space.*: Observation space configurationaction_space.*: Action space settingsepisode.*: Episode settings
Training (training.* and ppo.*):
- See
configs/training/ppo.yamlfor PPO hyperparameters
Quick Start
1. Start Godot Game
First, ensure the Godot game is running with the RL Bridge enabled:
# Headless mode (no GUI)
godot --headless --path /path/to/UTDG
# Or with GUI for visualization
godot --path /path/to/UTDG
Note: The game must have the RLBridge node added to the main scene (see Setup section below).
2. Train an Agent
from utdg_gym import UntitledTowerDefenseEnv
from stable_baselines3 import PPO
# Create environment
env = UntitledTowerDefenseEnv(
host="127.0.0.1",
port=9876,
max_episode_steps=1000
)
# Create PPO agent
model = PPO("MultiInputPolicy", env, verbose=1)
# Train the agent
model.learn(total_timesteps=100000)
# Save the model
model.save("ppo_utdg")
# Close environment
env.close()
3. Use the Example Scripts
# Simple test (displays configuration)
python examples/simple_test.py
# Simple test with Godot auto-launch
python examples/simple_test.py \
runtime.auto_launch=true \
runtime.godot_path="path/to/godot/executable"
# Train PPO agent with custom parameters
python examples/train_ppo.py \
ppo.learning_rate=0.0003 \
training.training.total_timesteps=100000
Godot Setup
To enable RL integration in your Godot project:
-
Add RLBridge Node: Open
base_level.tscnand add theRLBridgenode as a child of the root node. -
Configure References: Select the RLBridge node and set the following exported variables:
bank: Reference to the Bank nodebase: Reference to the Base nodeenemy_manager: Reference to the EnemyManager nodetower_manager: Reference to the TowerManager nodedifficulty_manager: Reference to the DifficultyManager node
-
Configure Port (optional): Set the
portvariable to match your Python environment.- Hydra config default: 9090 (
configs/default.yaml) - Environment default: 9876 (
utdg_gym/env.py) - Override via CLI:
websocket.port=9876orwebsocket.port=9090 - Ensure Godot and Python use the same port!
- Hydra config default: 9090 (
-
Auto-start: Ensure
auto_startis enabled to start the WebSocket server automatically.
Environment Details
Observation Space
The observation is a dictionary containing:
gold: Current gold amount (Box: [0, 10000])base_health: Current base health (Box: [0, 100])base_max_health: Maximum base health (Box: [0, 100])num_enemies: Number of active enemies (Box: [0, max_enemies])num_towers: Number of placed towers (Box: [0, max_towers])game_time: Elapsed game time in seconds (Box: [0, 1000])enemy_positions: Flattened array of enemy positions (Box: [-100, 100] × max_enemies × 3)tower_positions: Flattened array of tower positions (Box: [-100, 100] × max_towers × 3)
Action Space
Discrete action space with the following mapping:
- Action 0: Wait/Skip (do nothing)
- Actions 1-99: Place tower at grid position (mapped to world coordinates)
Reward Structure
The reward function is designed to align agent behavior with the core tower defense objective: protect the castle by building towers that eliminate enemies.
Reward Components
The total reward is computed from four components:
total_reward = (+1.0 × kills) + (-10 × base_damage) + (50 × waves_cleared) + (5 × towers_built)
1. Kill Reward: +1.0 per enemy killed
- Rewards tower effectiveness and enemy elimination
- Scaled from game economy (15 gold per kill × 0.067 scaling factor)
- Reduced from +1.5 to emphasize base protection over kill-chasing
- Provides continuous feedback on defensive performance
- Tracked as
custom/reward_killsin W&B
2. Base Damage Penalty: -10 per damage point
- Heavily penalizes enemies reaching the castle
- Primary failure signal that drives defensive behavior
- Calibrated to be strong (6.7× kill reward) without suppressing exploration
- Tracked as
custom/reward_damagein W&B
3. Wave Clear Bonus: +50 per wave cleared
- Rewards overall progress and wave completion
- Sparse but significant milestone reward
- Encourages long-term survival and advancement
- Tracked as
custom/reward_wavein W&B
4. Tower Building Reward: +5 per tower placed
- Encourages proactive tower placement and defensive expansion
- Addresses risk-aversion by providing immediate positive feedback for building
- Small enough not to dominate but helps overcome exploration barriers
- Tracked as
custom/reward_towers_builtin W&B
Economic Balance
The reward structure is calibrated to encourage tower building while maintaining base protection as the top priority:
| Metric | Value | Implication |
|---|---|---|
| Tower cost | 100 gold | Agent must justify investment |
| Tower building reward | +5 | Immediate feedback (computed from observation delta) |
| Reward per kill | +1.0 | Experiment 1: Reduced from +1.5 |
| Net tower value | +5 initial, then +1.0 per kill | Encourages strategic over opportunistic kills |
| Base damage penalty | -10 | 10× kill reward (stronger than 6.7× baseline) |
| Wave clear bonus | +50 | ~50 kills equivalent |
Design rationale (Experiment 1):
- Reduced kill reward (+1.0 vs +1.5) shifts focus toward base protection
- Achieves 10× damage dominance (same ratio as Option B) with proven -10 penalty
- Tests hypothesis: lower kill reward → more strategic placement, less kill-greedy behavior
- Maintains proven penalty level (-10) to avoid exploration suppression seen with -15
- Tower reward (+5) now active via observation-based tracking (tower_count delta)
- Wave bonuses reward long-term survival and proper defense scaling
Key difference from baseline:
- Baseline: 6.7× damage dominance (kills more rewarding)
- Experiment 1: 10× damage dominance (protection more important)
- Option B failure: 10× but with -15 penalty (too harsh, suppressed exploration)
Expected Behavioral Changes
Agents trained with this reward structure should exhibit:
- 🛡️ Stronger focus on base protection (10× damage dominance vs 6.7× baseline)
- 🎯 More strategic tower placement (less emphasis on high-kill locations)
- 📊 Reduced kill-chasing behavior (lower kill reward discourages opportunistic placement)
- ⚖️ Better long-term planning (wave bonus relatively more valuable: 50 kills equivalent)
- 🔍 Careful defensive positioning (proven -10 penalty maintains exploration)
Experiment 1 hypothesis:
- Lower kill reward (+1.0 vs +1.5) shifts agent toward defensive strategies
- Maintains proven -10 penalty to avoid Option B's exploration suppression
- Tests if 10× damage dominance improves performance without risk-aversion
The reward function emphasizes base protection as paramount (-10 penalty, 10× kill reward), while providing feedback on defensive effectiveness: kills (+1.0), waves cleared (+50), and towers built (+5 when event available).
Episode Termination
An episode ends when:
- Base health reaches zero (Loss)
- All waves complete and all enemies defeated (Win)
- Maximum episode steps reached (Truncated)
Advanced Usage
Custom Reward Function
You can modify the reward calculation in RLBridge/rl_bridge.gd:
func _calculate_reward(action_success: bool) -> float:
var reward: float = 0.0
# Your custom reward logic here
return reward
Multiple Parallel Environments
from stable_baselines3.common.env_util import make_vec_env
def make_env(rank):
def _init():
return UntitledTowerDefenseEnv(
port=9876 + rank # Different port per environment
)
return _init
# Create 4 parallel environments
env = make_vec_env(make_env, n_envs=4)
Note: You'll need to run multiple Godot instances, each on a different port.
Using with Ray RLlib
from ray.rllib.algorithms.ppo import PPOConfig
from utdg_gym import UntitledTowerDefenseEnv
config = (
PPOConfig()
.environment(env=UntitledTowerDefenseEnv)
.framework("torch")
.training(train_batch_size=4000)
)
algo = config.build()
for i in range(100):
result = algo.train()
print(f"Iteration {i}: reward={result['episode_reward_mean']}")
Message Protocol
Communication uses JSON messages with the following format:
{
"type": "message_type",
"data": { ... }
}
Message Types
- reset: Request environment reset
- reset_response: Initial observation after reset
- step: Execute action
- step_response: Observation, reward, done, info
- close: Close connection
- config: Configuration parameters
- error: Error message
See utdg_gym/protocol.py for detailed message schemas.
Examples
The examples/ directory contains:
Reinforcement Learning
trainer.py: Complete PPO training pipeline with Stable-Baselines3evaluate.py: Evaluate trained policiesrecord_video.py: Record video demonstrations of trained agentsrollout.py: Generate policy rollouts
Imitation Learning
Train agents from human demonstrations using various imitation learning algorithms:
Data Collection
record_demos.py: Collect human demonstration trajectories for imitation learning
Training Scripts
train_bc.py: Behavioral Cloning (supervised learning from demonstrations)train_gail.py: GAIL (Generative Adversarial Imitation Learning)train_airl.py: AIRL (Adversarial Inverse Reinforcement Learning)
Inference Scripts
run_bc.py: Run inference with trained BC policiesrun_gail.py: Run inference with trained GAIL policiesrun_airl.py: Run inference with trained AIRL policies
Requirements: pip install imitation
Basic workflow:
# 1. Collect human demonstrations
python examples/record_demos.py
# 2. Train using one of the imitation learning algorithms
python examples/train_bc.py # Behavioral Cloning (fastest, simplest)
python examples/train_gail.py # GAIL (adversarial training)
python examples/train_airl.py # AIRL (learns reward function)
# 3. Run inference with trained policy
python examples/run_bc.py # Run BC policy
python examples/run_gail.py # Run GAIL policy
python examples/run_airl.py # Run AIRL policy
Algorithm Comparison:
- BC (Behavioral Cloning): Direct supervised learning from demonstrations. Fastest to train, but may not generalize well to novel situations.
- GAIL: Uses adversarial training to match expert behavior. More robust than BC, uses BasicRewardNet.
- AIRL: Also adversarial but learns a shaped reward function (BasicShapedRewardNet). Can transfer learned rewards to different tasks.
Troubleshooting
Connection Issues
- Ensure Godot game is running before starting Python script
- Check that ports match between Godot (RLBridge) and Python environment
- Verify firewall settings allow local WebSocket connections
Performance
- Use headless mode for faster training:
godot --headless - Reduce observation space size if needed (adjust
max_enemies,max_towers) - Consider running multiple parallel environments for faster data collection
Godot Errors
- Ensure all node references in RLBridge are properly set
- Check that
try_place_tower()method exists in TowerManager - Verify
is_spawning()method exists in DifficultyManager
Hyperparameter Sweeps
This project uses Weights & Biases Sweeps for automated hyperparameter optimization. Sweep configurations are stored in configs/sweeps/ and use Bayesian optimization to search over model, training, and environment parameters.
Running a Sweep
- Create a new sweep:
wandb sweep --project utdg --entity rl4aa configs/sweeps/sweep.yaml
- Start the sweep agent using the ID returned from the previous command:
wandb agent rl4aa/utdg/<SWEEP_ID>
You can run multiple agents in parallel (on different machines or terminals) to speed up the search.
Sweep Configuration
The sweep uses Hydra override syntax to pass hyperparameters to the training script. Key parameters include:
| Parameter | Range | Description |
|---|---|---|
model.learning_rate |
1e-5 to 1e-3 | Learning rate (log-uniform) |
model.gamma |
0.9 to 0.999 | Discount factor |
model.batch_size |
32, 64, 128, 256 | Batch size |
model.n_steps |
1024 to 4096 | Steps per rollout |
training.total_timesteps |
100k to 300k | Total training steps |
See configs/sweeps/sweep.yaml for the full parameter specification.
Contributing
Contributions are welcome! Please ensure:
- Code follows Google-style docstrings
- Type hints are used throughout
- Code is modular and well-documented
License
See the main UTDG repository for license information.
Acknowledgments
- Built with Godot Engine
- Uses Gymnasium
- Training examples use Stable-Baselines3
Project details
Release history Release notifications | RSS feed
Download files
Download the file for your platform. If you're not sure which to choose, learn more about installing packages.
Source Distribution
Built Distribution
Filter files by name, interpreter, ABI, and platform.
If you're not sure about the file name format, learn more about wheel file names.
Copy a direct link to the current filters
File details
Details for the file utdg_env-0.1.6.tar.gz.
File metadata
- Download URL: utdg_env-0.1.6.tar.gz
- Upload date:
- Size: 19.4 kB
- Tags: Source
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.11.13
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
1f1aa0c699b2cff7e315794ffcc2b0fc5066ed6f6da4a3ff32ed42376c8309a9
|
|
| MD5 |
90ce2f88abca998ab95c6b9f3e42772f
|
|
| BLAKE2b-256 |
c4a8b6407b590218521ab46dd54033f8ad72e1382ef897263ff48ab32c9eb4df
|
File details
Details for the file utdg_env-0.1.6-py3-none-any.whl.
File metadata
- Download URL: utdg_env-0.1.6-py3-none-any.whl
- Upload date:
- Size: 15.9 kB
- Tags: Python 3
- Uploaded using Trusted Publishing? No
- Uploaded via: twine/6.2.0 CPython/3.11.13
File hashes
| Algorithm | Hash digest | |
|---|---|---|
| SHA256 |
e5425230573696a07c05620e53b7c45c9240fb66fb59596c59c41964db378487
|
|
| MD5 |
226ef89ea19b5a8ae78b01ea7d048883
|
|
| BLAKE2b-256 |
94cb917b1bb9f575c0e11ce2b1fd4ecc0490d08dbef3fd1988ddb47e185e5cee
|