# 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: 1. discuss the current situation, 2. request discretionary funding, 3. vote on other ministries' requests, 4. 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: ```text 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: ```text 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: - `DEBATE` - `FINISH_DEBATE` - `PROPOSE_BUDGET` - `VOTE` 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: $$ Demand_d = Baseline_d \times \frac{Population_t}{Population_0} \times EventMultiplier_d $$ The critical threshold is auto-funded first when the treasury can afford it: $$ Critical_d = 0.4 \times Demand_d $$ 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. ```text 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: $$ Revenue_d = Allocation_d \times RevenueFactor_d \times Productivity_t $$ 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: $$ Productivity_t = clamp(Productivity_{t-1} + 0.05 \times (AvgRF - 1.0), 0.5, 2.0) $$ 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: $$ Population_t = Population_{t-1} \times (1 + BirthRate - DeathRate) $$ 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: $$ R_t = Prosperity_t + ProductivityBonus_t + SurvivalBonus_t + AllocationPenalties_t + BankruptcyAdjustment_t $$ Prosperity is revenue per citizen: $$ Prosperity_t = \frac{TotalRevenue_t}{Population_t} $$ 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: 1. Roll out the `OptimalZoneAdapter` heuristic in the real environment. 2. Collect prompt rows from proposal and voting phases. 3. Render those rows with the same minister prompt used at inference time. 4. Train the model with GRPO using the dense environment-aligned reward function. 5. 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: ```text /nation-parliamentary-grpo-lora ``` The prompt dataset is intended to be pushed to: ```text datasets//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: ```bash 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: ```bash 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: ```bash uv run --extra training python training/train_grpo.py --smoke ``` Train and push a LoRA adapter with Hugging Face Jobs: ```bash 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: ```bash 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.