nation_optimizer / BLOG_POST.md
abanwild's picture
Fix blog post reward formula
465032a verified
|
Raw
History Blame Contribute Delete
15.7 kB

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:

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:

  • 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:

Demandd=BaselinedΓ—PopulationtPopulation0Γ—EventMultiplierd 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:

Criticald=0.4Γ—Demandd 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.

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:

Revenued=AllocationdΓ—RevenueFactordΓ—Productivityt 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:

Productivityt=clamp(Productivitytβˆ’1+0.05Γ—(AvgRFβˆ’1.0),0.5,2.0) 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:

Populationt=Populationtβˆ’1Γ—(1+BirthRateβˆ’DeathRate) 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:

Rt=Prosperityt+ProductivityBonust+SurvivalBonust+AllocationPenaltiest+BankruptcyAdjustmentt R_t = Prosperity_t + ProductivityBonus_t + SurvivalBonus_t + AllocationPenalties_t + BankruptcyAdjustment_t

Prosperity is revenue per citizen:

Prosperityt=TotalRevenuetPopulationt 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:

<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.