case study
RL Environment Auto-Designer
From a plain-English problem to a trained reinforcement learning policy
Designing a reinforcement learning environment by hand takes a long time, and asking an LLM to write one in a single shot rarely produces something that trains. I split the job across five agents with a typed contract between them and a repair loop around the generated code.
- Type
- Personal project
- Stack
- FastAPI, Gymnasium, Stable Baselines3, Next.js
- Models
- Claude or GPT, picked from whichever API key is set
The problem
A valid environment needs more than a Python class. State bounds, the action space, reward components, termination rules and the training algorithm all have to agree with each other. When one LLM call writes everything at once, a small mismatch usually means the environment crashes or trains on nonsense.
How it works
- Designer: reads the problem (for example "optimise hospital shift scheduling") and returns a structured design with the state space, action space, reward function, episode logic and its reasoning.
- Coder: turns that design into a complete, runnable Gymnasium environment.
- Executor: runs the generated environment in a separate process as a sanity check. If it fails, the error goes back to the Coder for up to two automatic fixes.
- Trainer: trains SAC for continuous actions or PPO for discrete ones with Stable Baselines3, and saves the model, a reward curve and a results CSV. Training crashes are fed back to the Coder in the same way.
- Reporter: writes a markdown report explaining the design decisions and how training went.
The FastAPI backend streams progress as NDJSON, so the Next.js frontend shows each stage live and then the report, code, reward curve and design in tabs.
Key decisions
- Design before code. The Designer's output is a typed contract (Pydantic models shared by every agent), so the Coder works from a spec, and when something breaks I can tell whether the design or the code was wrong.
- Run it, don't trust it. Every environment is executed before training, and errors go back with the traceback rather than a vague retry.
- Pick the algorithm from the action space. Continuous Box actions go to SAC; Discrete, MultiDiscrete and MultiBinary go to PPO.
- Log every reward component per step, so the report can explain why the agent learned what it did.
Limitations and next steps
- The sanity check catches code that crashes, not designs that are valid but unhelpful. A reward that trains smoothly can still encode the wrong goal.
- Next: publish a few complete example runs (design, reward curve and report) and add automated tests around the Designer and Executor contracts.