Skip to main content

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(), and close() 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 name
  • seed: Random seed for reproducibility
  • log_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 run
  • godot_path: Path to Godot executable
  • auto_launch: Auto-launch Godot if not running
  • headless: Run Godot in headless mode

Game configuration (game.*):

  • base_health: Starting base health
  • starting_gold: Starting gold amount
  • enemy_spawn_config.*: Enemy spawning parameters
  • reward_config.*: Reward shaping settings
  • difficulty_config.*: Difficulty multipliers

Environment (env.*):

  • observation_space.*: Observation space configuration
  • action_space.*: Action space settings
  • episode.*: Episode settings

Training (training.* and ppo.*):

  • See configs/training/ppo.yaml for 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:

  1. Add RLBridge Node: Open base_level.tscn and add the RLBridge node as a child of the root node.

  2. Configure References: Select the RLBridge node and set the following exported variables:

    • bank: Reference to the Bank node
    • base: Reference to the Base node
    • enemy_manager: Reference to the EnemyManager node
    • tower_manager: Reference to the TowerManager node
    • difficulty_manager: Reference to the DifficultyManager node
  3. Configure Port (optional): Set the port variable to match your Python environment.

    • Hydra config default: 9090 (configs/default.yaml)
    • Environment default: 9876 (utdg_gym/env.py)
    • Override via CLI: websocket.port=9876 or websocket.port=9090
    • Ensure Godot and Python use the same port!
  4. Auto-start: Ensure auto_start is 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_kills in 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_damage in 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_wave in 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_built in 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-Baselines3
  • evaluate.py: Evaluate trained policies
  • record_video.py: Record video demonstrations of trained agents
  • rollout.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 policies
  • run_gail.py: Run inference with trained GAIL policies
  • run_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

  1. Create a new sweep:
wandb sweep --project utdg --entity rl4aa configs/sweeps/sweep.yaml
  1. 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

Project details


Download files

Download the file for your platform. If you're not sure which to choose, learn more about installing packages.

Source Distribution

utdg_env-0.1.6.tar.gz (19.4 kB view details)

Uploaded Source

Built Distribution

If you're not sure about the file name format, learn more about wheel file names.

utdg_env-0.1.6-py3-none-any.whl (15.9 kB view details)

Uploaded Python 3

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

Hashes for utdg_env-0.1.6.tar.gz
Algorithm Hash digest
SHA256 1f1aa0c699b2cff7e315794ffcc2b0fc5066ed6f6da4a3ff32ed42376c8309a9
MD5 90ce2f88abca998ab95c6b9f3e42772f
BLAKE2b-256 c4a8b6407b590218521ab46dd54033f8ad72e1382ef897263ff48ab32c9eb4df

See more details on using hashes here.

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

Hashes for utdg_env-0.1.6-py3-none-any.whl
Algorithm Hash digest
SHA256 e5425230573696a07c05620e53b7c45c9240fb66fb59596c59c41964db378487
MD5 226ef89ea19b5a8ae78b01ea7d048883
BLAKE2b-256 94cb917b1bb9f575c0e11ce2b1fd4ecc0490d08dbef3fd1988ddb47e185e5cee

See more details on using hashes here.

Supported by

AWS Cloud computing and Security Sponsor Datadog Monitoring Depot Continuous Integration Fastly CDN Google Download Analytics Pingdom Monitoring Sentry Error logging StatusPage Status page