# ponseko/econojax architecture

> Architecture diagram of the GitHub repository ponseko/econojax, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-09-29). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/ponseko/econojax
- Repository: https://github.com/ponseko/econojax

## Overview

EconoJax is a JAX-based economic simulation environment, trained in the supplied entry point with PPO. The workflow initializes the environment, trains agents through its environment interface, and simulates gathering/crafting, market trading, government taxation, and utility outcomes. The sampled source confirms these mechanics and their step-level calls; unsampled implementation details are not expanded.

Main entry points: [`econojax/econojax.py`](https://github.com/ponseko/econojax/blob/main/econojax/econojax.py) (EconoJax), [`main.py`](https://github.com/ponseko/econojax/blob/main/main.py) (Training entry), [`econojax/util.py`](https://github.com/ponseko/econojax/blob/main/econojax/util.py) (Gini metric)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_training["Training"]
  node_main["Training entry<br/>[main.py]"]
  node_ppo["PPO trainer"]
  node_logwrapper["Log wrapper"]
end

subgraph group_environment["Environment"]
  node_env["EconoJax<br/>[econojax.py]"]
  node_state["Environment state<br/>[econojax.py]"]
  node_interface["Observations and actions<br/>[econojax.py]"]
end

subgraph group_economy["Economic mechanics"]
  node_gathercraft["Gather and craft<br/>[econojax.py]"]
  node_trading["Market trading<br/>[econojax.py]"]
  node_taxation["Government taxation<br/>[econojax.py]"]
end

subgraph group_metrics["Outcomes"]
  node_utility["Utility calculation<br/>[econojax.py]"]
  node_gini["Gini metric<br/>[util.py]"]
end

node_user(("Researcher"))

node_user -->|"starts training"| node_main
node_main -->|"creates environment"| node_env
node_main -->|"wraps environment"| node_logwrapper
node_main -->|"starts training"| node_ppo
node_main -->|"passes to trainer"| node_logwrapper
node_ppo -->|"trains through"| node_logwrapper
node_env -->|"initializes and updates"| node_state
node_env -->|"provides"| node_interface
node_env -->|"dispatches"| node_gathercraft
node_env -->|"dispatches"| node_trading
node_env -->|"dispatches"| node_taxation
node_env -->|"calculates"| node_utility
node_gathercraft -->|"updates"| node_state
node_trading -->|"updates"| node_state
node_taxation -->|"updates"| node_state
node_utility -->|"updates"| node_state
node_utility -->|"calculates equality with"| node_gini

click node_main "https://github.com/ponseko/econojax/blob/main/main.py"
click node_env "https://github.com/ponseko/econojax/blob/main/econojax/econojax.py"
click node_state "https://github.com/ponseko/econojax/blob/main/econojax/econojax.py"
click node_interface "https://github.com/ponseko/econojax/blob/main/econojax/econojax.py"
click node_gathercraft "https://github.com/ponseko/econojax/blob/main/econojax/econojax.py"
click node_trading "https://github.com/ponseko/econojax/blob/main/econojax/econojax.py"
click node_taxation "https://github.com/ponseko/econojax/blob/main/econojax/econojax.py"
click node_utility "https://github.com/ponseko/econojax/blob/main/econojax/econojax.py"
click node_gini "https://github.com/ponseko/econojax/blob/main/econojax/util.py"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_main,node_ppo,node_logwrapper toneBlue
class node_env,node_state,node_interface toneAmber
class node_gathercraft,node_trading,node_taxation toneMint
class node_utility,node_gini toneRose
class node_user toneIndigo
```

## Components

### Training

- [Training entry](https://github.com/ponseko/econojax/blob/main/main.py): `main.py`
- **PPO trainer**
- **Log wrapper**

### Environment

- [EconoJax](https://github.com/ponseko/econojax/blob/main/econojax/econojax.py): `econojax/econojax.py`
- [Environment state](https://github.com/ponseko/econojax/blob/main/econojax/econojax.py): `econojax/econojax.py`
- [Observations and actions](https://github.com/ponseko/econojax/blob/main/econojax/econojax.py): `econojax/econojax.py`

### Economic mechanics

- [Gather and craft](https://github.com/ponseko/econojax/blob/main/econojax/econojax.py): `econojax/econojax.py`
- [Market trading](https://github.com/ponseko/econojax/blob/main/econojax/econojax.py): `econojax/econojax.py`
- [Government taxation](https://github.com/ponseko/econojax/blob/main/econojax/econojax.py): `econojax/econojax.py`

### Outcomes

- [Utility calculation](https://github.com/ponseko/econojax/blob/main/econojax/econojax.py): `econojax/econojax.py`
- [Gini metric](https://github.com/ponseko/econojax/blob/main/econojax/util.py): `econojax/util.py`

### Other components

- **Researcher**

## Connections

- Researcher → Training entry: starts training
- Training entry → EconoJax: creates environment
- Training entry → Log wrapper: wraps environment
- Training entry → PPO trainer: starts training
- Training entry → Log wrapper: passes to trainer
- PPO trainer → Log wrapper: trains through
- EconoJax → Environment state: initializes and updates
- EconoJax → Observations and actions: provides
- EconoJax → Gather and craft: dispatches
- EconoJax → Market trading: dispatches
- EconoJax → Government taxation: dispatches
- EconoJax → Utility calculation: calculates
- Gather and craft → Environment state: updates
- Market trading → Environment state: updates
- Government taxation → Environment state: updates
- Utility calculation → Environment state: updates
- Utility calculation → Gini metric: calculates equality with