Spaces:
Sleeping
Nation Optimizer RL: Teaching an LLM to Budget a Country
What happens if a language model has to run a country for 12.5 years?
Not by writing speeches. Not by predicting the next token in a policy memo. By making budget decisions, surviving crises, keeping a treasury solvent, and explaining its reasoning in public.
That is the idea behind Nation Optimizer RL, an OpenEnv-compatible reinforcement learning environment where a single LLM learns long-horizon resource allocation through a parliamentary planning interface.
The project is inspired by the distributional principle "from each according to ability, to each according to need", but the implementation is a toy economic simulator, not a claim about real-world governance. We use that framing to create a clean research problem: can a model improve collective prosperity when every decision must balance urgent needs, finite public resources, and future uncertainty?
The Core Question
Our main research question is:
Can structured reasoning, in the form of parliamentary debate, proposals, and voting, improve long-horizon resource allocation compared with direct central planning?
The important detail is that this is not a multi-agent RL system. There is one underlying LLM. The "ministers" are role prompts over the same model, not independent learners with separate objectives or gradients.
The parliament is a reasoning scaffold. It forces the model to make public, structured decisions:
- discuss the current situation,
- request discretionary funding,
- vote on other ministries' requests,
- live with the economic consequences.
That gives us something ordinary direct planning often lacks: an interpretable trace of why a budget was passed.
Environment Overview
Each episode simulates a national planning horizon of 50 rounds. Each round represents one quarter, so a full episode covers 12.5 years.
The nation has six departments:
| Department | Baseline demand |
|---|---|
| Social/Municipal | 60 |
| Agriculture | 70 |
| Health | 90 |
| Education/R&D | 80 |
| Defense | 100 |
| Commerce | 75 |
The initial conditions are deliberately simple:
- treasury starts at 1000,
- population starts at 1,000,000,
- productivity starts at 1.0,
- productivity is bounded between 0.5 and 2.0,
- debt is not allowed,
- bankruptcy ends the episode.
Each round has a nine-phase cycle:
1. Event revelation
2. Debate
3. Budget proposals
4. Voting
5. Budget execution
6. Consumption and event impact
7. Revenue calculation
8. Surplus rollover
9. Termination check
The environment is exposed through an OpenEnv-compatible wrapper. The core simulator lives separately from the server wrapper, so the game rules can be tested directly and then served through the Hugging Face Space without changing the mechanics.
The Parliamentary Architecture
The architecture has four main layers:
Hugging Face / local model
|
v
Parliamentary LLM adapter
|
v
Structured JSON actions
|
v
NationGame simulator
|
v
Telemetry, reward, benchmark summaries
The LLM never directly mutates game state. It must emit one of the structured actions allowed in the current phase:
DEBATEFINISH_DEBATEPROPOSE_BUDGETVOTE
Those actions are parsed as JSON, validated against the action schema, and then passed to the simulator. Invalid actions are rejected by the environment, not silently accepted by the adapter.
The same model can be wrapped in two different ways:
- Parliamentary adapter: the model is called as different ministers, producing debate, proposals, and votes.
- Dictator adapter: the same model capacity is used as a direct central planner baseline.
This lets us compare the effect of the reasoning structure, not just the model weights.
What the Model Can See
The environment is intentionally partially observable.
Ministers see public state:
- current treasury,
- population,
- productivity,
- round and phase,
- public sector thresholds,
- current events,
- debate history,
- proposals and votes so far.
They do not receive privileged access to hidden exact event costs. Instead, events are revealed as narratives with severity and affected sectors. For example, a war event may say that defense needs emergency funding and commerce routes are disrupted, but the model must infer how aggressively to allocate from the available public information.
That matters because the main challenge is not a one-step optimization problem. It is planning under uncertainty.
Simulation Logic
The economy is built around a central treasury and six public sectors.
At the start of each round, stochastic events are sampled from an event catalog. Most rounds are quiet, but some contain minor, moderate, critical, or compound events. Events change sector demand through multipliers. Positive events can reduce demand and inject treasury cash; crises can sharply increase demand in affected sectors.
Each sector has four thresholds:
- Critical: the survival floor, equal to 40% of demand.
- Demand: the baseline required operating level.
- Surplus: the peak productive zone, equal to 150% of demand.
- Wastage: the break-even overfunding boundary, equal to 250% of demand.
Demand scales with population and events:
The critical threshold is auto-funded first when the treasury can afford it:
Ministers do not propose the critical floor. They propose discretionary funding above it. This design decision keeps the parliament focused on growth, prioritization, and crisis response, while guaranteeing minimum services when the state is solvent.
If the treasury cannot pay the total critical floor, the country is bankrupt.
The Revenue Curve
The heart of the simulator is a piecewise revenue curve.
An allocation below critical is invalid in the economic model. At exactly critical, the department survives but generates no revenue. From critical to demand, revenue factor rises linearly to 1.0. From demand to surplus, it rises to a peak of 1.8. After surplus, it decays exponentially and eventually falls back below break-even.
Revenue Factor
1.8 | peak
| /\
1.0 |-------- demand/ \------ wastage
| /
0.0 |-- critical--/
|
+-------------------------------- allocation
This creates four strategic zones:
| Zone | Meaning |
|---|---|
| Below critical | not viable |
| Critical to demand | survival, but underproductive |
| Demand to surplus | productive growth zone |
| Beyond surplus | waste and declining returns |
Revenue is calculated in the same round:
This gives the model immediate feedback from budget decisions, while productivity and population carry the long-term consequences forward.
Productivity and Population
Productivity is persistent. It changes based on the average revenue factor across sectors:
Good allocation decisions compound. Sustained allocations in the profit zone push productivity upward, which increases future revenue from the same budget. Bad allocation decisions pull productivity down and make future recovery harder.
Population also changes over time:
Birth rate rises with productivity. Death rate has a base component and receives a penalty during severe crises. As population grows, sector demand grows too, so prosperity is not just "make the treasury bigger." The model must keep per-capita output healthy while the system changes around it.
Reward Design
There are two reward layers in the project.
The first is the environment reward, used to score simulator outcomes. It is collective: every minister receives the same reward. There are no individual department rewards.
The reward is:
Prosperity is revenue per citizen:
The productivity bonus rewards sustained competence. The survival bonus rewards keeping the government alive across long horizons. Under-allocation and over-allocation penalties discourage budgets that leave sectors below demand or push them past the surplus zone into waste.
The second layer is the GRPO training reward. This is a dense reward function for model completions during training. It scores only proposal and voting turns:
- malformed JSON receives a parse penalty,
- illegal actions receive an illegal-action penalty,
- budget proposals are scored by the resulting revenue factor,
- votes are scored according to whether the proposal they support or reject is economically sensible.
Debate is inference-only. We want debate to improve context and interpretability, but we do not directly reward rhetorical style.
GRPO Training Pipeline
The training pipeline uses Hugging Face TRL's GRPOTrainer with LoRA.
The default base model is Qwen/Qwen3.5-9B, configurable through NATION_GRPO_BASE_MODEL. For smoke tests, the code can use a smaller model such as Qwen/Qwen2.5-0.5B-Instruct.
The pipeline is:
- Roll out the
OptimalZoneAdapterheuristic in the real environment. - Collect prompt rows from proposal and voting phases.
- Render those rows with the same minister prompt used at inference time.
- Train the model with GRPO using the dense environment-aligned reward function.
- Evaluate the base model, trained LoRA, and rule-based baselines on shared seeds.
This is not full online RL where every token generation steps the simulator live. Instead, the prompt dataset is generated from environment rollouts, and the reward function is tied to the simulator's actual revenue curve. That made the training loop practical for the hackathon while preserving alignment with the environment mechanics.
The trained artifact is a LoRA adapter, intended to be pushed to:
<hf-user>/nation-parliamentary-grpo-lora
The prompt dataset is intended to be pushed to:
datasets/<hf-user>/nation-parliamentary-prompts
Baselines
We implemented several policy adapters against the same action contract:
- Random: samples noisy behavior and usually fails quickly.
- Greedy: each minister asks for as much discretionary funding as possible.
- Equal split: divides visible resources naively.
- Conservative: requests stable baseline demand.
- Optimal zone: targets roughly 1.3x demand, aiming for the productive band without entering severe waste.
- Parliamentary base LLM: same parliamentary interface, no GRPO LoRA.
- Parliamentary GRPO LLM: base model plus the trained LoRA adapter.
The benchmark runner tracks:
- mean episode return,
- survival rounds,
- bankruptcy and shutdown,
- final prosperity,
- average revenue factor,
- treasury stability,
- productivity growth,
- invalid action and parse error counts for LLM policies.
The project also includes result artifacts for policy comparisons, survival distributions, and reward landscape visualization:
Policy comparison
Survival distribution
Reward landscape
Design Decisions That Matter
One model, many roles. The ministers are not separate agents. This keeps the comparison focused on reasoning structure instead of multi-agent credit assignment.
Critical funding is automatic. The parliament debates discretionary investment, not whether a sector receives a bare survival floor. This avoids brittle episodes where a single missing proposal instantly kills the economy and lets us study prioritization above minimum viability.
No debt. The treasury cannot go negative. This makes failure clear and forces intertemporal tradeoffs.
Same-round revenue. Allocations generate revenue immediately. That creates a usable learning signal while productivity and population preserve long-horizon consequences.
Public reasoning traces. Debate, proposals, and votes are logged. Even when the model is wrong, we can inspect how it reasoned its way there.
Strict JSON actions. The LLM can be creative in debate, but state-changing decisions must be structured and validated.
Shared reward. Every minister succeeds or fails together. This removes individual score-chasing and aligns the game with collective prosperity.
Why This Environment Is Interesting
Nation Optimizer RL sits between a spreadsheet and a political role-play simulation.
It is simpler than a real economy, by design. There is no private market, no debt, no corruption, no private communication, and no individual utility. But those simplifications make the learning problem sharper:
- Can the model infer hidden economic pressure from public event narratives?
- Can it preserve treasury reserves without becoming too conservative?
- Can it keep multiple sectors in the profit zone at the same time?
- Can debate and voting improve planning, or do they just add tokens?
- Can GRPO reduce invalid actions and improve economically meaningful proposals?
The answer is not supposed to come from vibes. The environment logs the decisions, computes the rewards, and benchmarks policies on shared seeds.
What We Learned
The most useful part of the project was separating the story from the contract.
The story is a parliament debating a national budget. That makes episodes readable and gives humans a natural way to inspect the model's behavior.
The contract is stricter:
- every action is typed,
- every budget has thresholds,
- every vote is validated,
- every reward comes from simulator math,
- every benchmark uses the same environment interface.
That combination makes the environment fun to watch, but still measurable.
Reproducing the Project
Run the OpenEnv-compatible server:
uv sync --extra dev --extra viz --extra training
uv run uvicorn server.app:app --host 0.0.0.0 --port 8000
Collect GRPO prompt data:
uv run --extra training python -m scripts.collect_grpo_prompts \
--seeds 50 \
--max-rounds 12 \
--output assets/datasets/grpo_prompts.jsonl
Run a GRPO smoke test:
uv run --extra training python training/train_grpo.py --smoke
Train and push a LoRA adapter with Hugging Face Jobs:
hf jobs uv run --flavor a10g-small --secrets HF_TOKEN \
training/train_grpo.py \
--dataset-id "${NATION_HF_USER}/nation-parliamentary-prompts" \
--hub-model-id "${NATION_HF_USER}/nation-parliamentary-grpo-lora"
Run benchmarks:
uv run python evaluation/benchmark_policies.py \
--seeds 1 2 3 \
--max-rounds 12 \
--output assets/results/benchmark_summary.json
Closing
Nation Optimizer RL is a small experiment in making LLM planning more inspectable.
Instead of asking a model for one opaque allocation, we ask it to deliberate, propose, vote, and then face the consequences in a living economy. The resulting system is still a toy, but it gives us a concrete place to study long-horizon planning, structured reasoning, reward alignment, and the gap between saying a plan makes sense and surviving 50 quarters with a finite treasury.