VaishnaviAgrawal/Radiotherapy_Planning_Environment

RadiotherapyPlanningEnv — RL Environment for Cancer Treatment Planning

A Gymnasium-compatible reinforcement learning environment where an AI agent learns to plan cancer radiotherapy treatment — placing radiation beams to maximize tumor dose while protecting critical organs.

Clinical Motivation

~14 million cancer patients per year require radiotherapy. A radiation oncologist must decide:

• —How many beams to use
• —At what angles
• —With what dose intensity

...while ensuring the tumor receives enough radiation and nearby healthy organs stay below safe limits. This process takes human experts 2–4 hours per patient. This environment simulates that decision-making process for RL agents.

This is NOT a clinical tool. It is a physically-grounded benchmark environment for testing RL algorithms on a meaningful real-world problem.

Live Demo

Try the interactive Gradio demo on HuggingFace Spaces — watch a trained PPO agent plan treatment in real time, or play manually to compete against it:

[https://huggingface.co/spaces/VaishnaviAgrawal/Radiotherapy_Planning_Environment](https://huggingface.co/spaces/VaishnaviAgrawal/Radiotherapy_Planning_Environment)

Quick Start

git clone https://github.com/VaishnaviAgrawal03/Radiotherapy-Planning-Environment.git
cd Radiotherapy-Planning-Environment
pip install -e .

import gymnasium as gym
import radiotherapy_env

env = gym.make("RadiotherapyEnv-prostate-v1", render_mode="rgb_array")
obs, info = env.reset(seed=42)

for _ in range(50):
    action = env.action_space.sample()
    obs, reward, terminated, truncated, info = env.step(action)
    if terminated or truncated:
        break

print(f"Final score: {info['score']:.3f}")
env.close()

Tasks — 3 Difficulty Levels

Baseline Results

PPO Agent (stable-baselines3, MultiInputPolicy)

Score ≥ 0.6 = clinically acceptable treatment plan.

LLM Agent (Llama 3.3 70B via inference.py)

Action Space — Discrete(8)

Observation Space — Dict

All observations normalized to [0, 1] for neural network training stability.

Reward Function

Dense per-step reward in [0.0, 1.0]:

reward = tumor_coverage × 0.55
       − oar_penalty    × 0.40   (priority-weighted: critical=1.5×, moderate=0.5×)
       + plan_efficiency × 0.05

• —Tumor coverage (55%): D95 metric + fraction of tumor receiving ≥ 95% prescription dose
• —OAR penalty (40%): Priority-weighted organ violations (critical organs penalized 1.5×)
• —Plan efficiency (5%): Optimal beam count around 5–7

The reward provides meaningful partial progress signals — distinct from the stricter compute_score() used for final grading (which uses binary pass/fail for critical OARs).

Physics Model

Gaussian pencil-beam dose calculation:

beam_dose = lateral_profile × depth_attenuation × dose_weight × BEAM_SCALE

• —Lateral profile: Gaussian falloff from beam central axis (σ = 4.0)
• —Depth attenuation: Exponential decay through tissue (Beer-Lambert Law, μ = 0.012)
• —Beam superposition: Total dose = sum of all beam contributions
• —Isocenter convergence: All beams aimed at tumor center

Simplified from clinical Monte Carlo (milliseconds vs. hours) but preserves the core trade-off: multiple beams overlap at the tumor for high dose while surrounding organs receive minimal radiation.

Installation

# Core environment
pip install -e .

# With PPO training support
pip install -e ".[training]"

# With LLM inference support
pip install -e ".[inference]"

# With Gradio demo
pip install -e ".[demo]"

Running the LLM Inference Script

Connects an LLM to the environment via an OpenAI-compatible API:

export API_BASE_URL="https://api.groq.com/openai/v1"
export MODEL_NAME="llama-3.3-70b-versatile"
export API_KEY="your_api_key"
python inference.py

Required log output format:

[START] task=prostate env=RadiotherapyEnv-prostate-v1 model=llama-3.3-70b-versatile
[STEP] step=1 action=add_beam reward=0.03 done=false error=null
[STEP] step=2 action=add_beam reward=0.07 done=false error=null
...
[END] success=true steps=50 score=0.635 rewards=0.03,0.07,...

Training a PPO Agent

pip install -e ".[training]"
python baseline/train_ppo.py

Trains on all three tasks using 4 parallel vectorized environments. Saved checkpoints land in baseline/models/.

Auto-Grader

from radiotherapy_env.reward.grader import grade_all

def my_agent(obs, env):
    return env.action_space.sample()

results = grade_all(my_agent, n_episodes=20, seed=42)
print(f"Aggregate: {results['aggregate_score']:.3f}")

Each grader is deterministic with seed; pass threshold is 0.60.

Running Tests

pytest tests/ -v
# 25 tests: Gymnasium compliance, physics, reward, task difficulty

Docker

docker build -t radiotherapy-env:latest .
docker run -p 7860:7860 radiotherapy-env:latest

OpenEnv HTTP Server

The FastAPI server exposes the environment over HTTP (required by openenv validate):

uv run server
# or
uvicorn server.app:app --host 0.0.0.0 --port 8000

Repository Structure

Radiotherapy-Planning-Environment/
├── inference.py               # LLM inference script
├── openenv.yaml               # OpenEnv spec metadata
├── Dockerfile                 # Container build (port 7860)
├── radiotherapy_env/          # Main Python package
│   ├── env.py                 # Core RadiotherapyEnv class
│   ├── physics/               # Dose calculator, DVH, patient models
│   ├── tasks/                 # 3 task definitions (prostate, head_neck, pediatric_brain)
│   ├── reward/                # Reward function, scoring, auto-grader
│   └── rendering/             # Dose heatmap + DVH visualization
├── server/                    # OpenEnv HTTP server (FastAPI)
├── baseline/                  # PPO training, evaluation, saved models, results
├── app/                       # Gradio interactive demo
└── tests/                     # 25-test pytest suite

Key Design Decisions

Author

Vaishnavi Agrawal — vagrawal_be22@thapar.edu