File size: 9,634 Bytes
6759502
 
8270f31
6759502
 
 
 
 
 
 
 
 
 
8270f31
6759502
8270f31
6759502
8270f31
6759502
ede0718
 
03fcaa7
ede0718
6759502
8270f31
6759502
8270f31
6759502
8270f31
6759502
8270f31
 
 
 
 
6759502
 
8270f31
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6759502
 
8270f31
6759502
8270f31
6759502
8270f31
 
 
 
 
 
 
6759502
ede0718
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
8270f31
6759502
8270f31
6759502
a34d4a2
6759502
8270f31
 
 
 
 
6759502
8c686ce
 
ede0718
8c686ce
a34d4a2
8c686ce
 
 
 
 
 
8270f31
6759502
8270f31
6759502
be9f416
8270f31
 
 
6759502
 
8270f31
6759502
be9f416
8270f31
be9f416
6759502
8270f31
6759502
be9f416
8270f31
be9f416
6759502
8270f31
6759502
be9f416
8270f31
 
 
 
be9f416
6759502
8270f31
6759502
be9f416
8270f31
 
 
 
ede0718
be9f416
6759502
8270f31
 
 
6759502
be9f416
ede0718
 
8270f31
 
 
 
ede0718
 
8270f31
ede0718
 
 
 
8270f31
ede0718
 
8270f31
ede0718
 
 
 
be9f416
6759502
8270f31
6759502
be9f416
 
 
6759502
 
be9f416
6759502
8270f31
 
 
ede0718
8270f31
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
6759502
8270f31
6759502
8270f31
6759502
be9f416
 
6759502
be9f416
 
 
 
6759502
be9f416
8270f31
 
6759502
be9f416
 
6759502
be9f416
 
6759502
be9f416
8270f31
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
---
title: Ontology Architect Environment Server
emoji: πŸ”¬
colorFrom: purple
colorTo: green
sdk: docker
pinned: false
app_port: 8000
base_path: /web
tags:
  - openenv
---

# πŸ”¬ Ontology Architect

> **Train an LLM to discover hidden physical laws from raw sensor data.**

Ontology Architect is an OpenEnv-compatible RL environment that challenges LLM agents to perform *unsupervised scientific discovery*. An agent observes noisy sensor logs from a simulated alien universe governed by hidden ODEs and must iteratively write compact theories that predict future observations.

**β†’ [πŸ€— HuggingFace Space (Environment)](https://huggingface.co/spaces/LostDecimal/ontology-architect)**  
**β†’ [πŸ““ Colab Training Notebook](https://colab.research.google.com/github/rohan-27p/ontology_architect/blob/main/Alien_Physics_Discovery_OpenEnv.ipynb)**  
**β†’ [πŸ’» GitHub Repository](https://github.com/rohan-27p/Ontology-architect)**  
**β†’ [πŸ€– Trained Model (HuggingFace)](https://huggingface.co/LostDecimal/alien-physics-qwen2.5-1.5b)**

---

## The Challenge: Kuhnian Scientific Discovery

The environment implements a 3-phase discovery challenge inspired by Thomas Kuhn's *Structure of Scientific Revolutions*:

```mermaid
graph LR
    A["Phase 1: Normal Science"] -->|"drift event"| B["Phase 2: Crisis"]
    B -->|"new latent variable"| C["Phase 3: Paradigm Shift"]
    C -->|"new equilibrium"| A
```

1. **Normal Science:** The agent sees 4 sensor channels (`pressure`, `turbulence`, `thermal_radiation`, `magnetic_flux`) and must discover that two hidden latent variables (A, B) drive all observations through coupled ODEs.
2. **Crisis:** The forcing frequency shifts β€” predictions break. The agent must detect the anomaly and adapt.
3. **Paradigm Shift:** A *third* latent variable (C) activates. The old 2-variable theory is fundamentally incomplete. The agent must discover this new ontological structure.

## How It Works

The agent submits a **structured Theory DSL** β€” a JSON spec defining latent state variables, dynamics equations, and observation projections. The environment:

1. Executes the theory in a secure sandbox
2. Compares predictions against hidden future observations
3. Scores using a composite reward: `log_likelihood - MDL_penalty + anomaly_bonus + drift_bonus + stability_bonus`
4. Returns peer review feedback showing where predictions diverged

### Theory DSL Example

```json
{
  "dsl_version": 1,
  "name": "dual fluid sketch",
  "state": ["A", "B"],
  "dynamics": {
    "A": {"linear": {"terms": {"A": -0.3, "B": 0.25}}},
    "B": {"add": [{"sin": "A"}, {"linear": {"terms": {"A": -0.08}}}]}
  },
  "observations": {
    "pressure": {"var": "A"},
    "turbulence": {"pow": [{"add": [{"var": "A"}, {"neg": {"var": "B"}}]}, 2]},
    "thermal_radiation": {"mul": [{"var": "B"}, {"exp": {"neg": {"abs": {"var": "A"}}}}]},
    "magnetic_flux": {"mul": [0.1, {"var": "A"}, {"var": "B"}]}
  },
  "integrator": {"dt": 0.15, "substeps": 2},
  "noise": 0.1
}
```

The DSL supports: `var`, `const`, `linear`, `add`, `mul`, `sin`, `cos`, `tanh`, `exp`, `abs`, `pow`, `sqrt`, `neg`.

## Reward Structure

| Component | Type | Signal |
|---|---|---|
| `r_prediction` | Dense | Gaussian log-likelihood on hidden future window |
| `r_complexity` | Dense | AST/DSL Minimum Description Length penalty |
| `r_anomaly` | Sparse | Bonus for detecting rare anomaly events |
| `r_drift` | Sparse | Bonus for adapting when hidden laws shift |
| `r_stability` | Dense | Rewards incremental refinement, penalizes wild oscillation |

## Training Results

We trained **Qwen2.5-Coder-1.5B-Instruct** on an NVIDIA A100-SXM4-80GB GPU using a two-phase pipeline:

### Phase 1: Supervised Fine-Tuning (SFT)

The model first learns the Theory DSL format by imitating oracle traces from the `dual_fluid` universe. Loss drops from **1.25 β†’ 0.017** in 200 steps (~8.6 min).

![SFT Training Loss](artifacts/results/sft_training_loss.png)

### Phase 2: Group Reward Optimization (GRPO)

Using the SFT-warm-started model, we run GRPO (group_size=4, temperature=1.2) for 50 steps. The model explores diverse theory structures and is rewarded based on real sandbox evaluation.

![GRPO Training Curve](artifacts/results/grpo_training_curve.png)

Key results:
- **Best single theory reward: -0.94** (72% improvement over SFT baseline of -3.41)
- Valid theory generation rate: ~75% (up from 0% without SFT)
- Model discovers improved coefficient values and observation mappings

### Combined Training Overview

![Training Overview](artifacts/results/training_overview.png)

## Results: Latent Variable Discovery

The following plot shows the environment running with the DSL oracle baseline (purple) vs a static persistence baseline (gray). The top rows show the *true hidden latent variables* (A, B, C) β€” which the agent never sees. The bottom rows show sensor predictions.

![Discovery Plot](artifacts/results/discovery_plot.png)

Key observations:
- **Latent A** exhibits autocatalytic damping β€” the oracle captures its trajectory
- **Latent B** is a forced oscillator driven by `sin(ω·t)` β€” after fixing the ODE bug, this now genuinely oscillates
- **Latent C** activates after the second drift event β€” the agent must discover this new dimension
- Red dashed lines mark drift events (hidden from the agent)

## Results: Agent Learning & Reward Curve

The environment evaluates agents by tracking their `total_reward` across episodes, directly capturing their ability to adapt to falsification. The generated reward curve below plots the performance of various baseline agents (Random, Heuristic, LLM, and Oracle) over time.

![Reward Curve Plot](artifacts/results/reward_curve.png)

Key observations:
- **Phase Transitions**: Highlighted in shaded vertical regions (e.g., green for Normal Science, red for Crisis).
- **Falsification & Recovery**: When the alien physics undergo a paradigm shift, agents experience a sharp drop in reward (falsification) and must structurally refactor their theories to recover.
- **Theory Diversity**: The grader actively tracks structural distance between consecutive theories to ensure the agent is not stuck in a local minimum repeating the same flawed equations.

## How To Run

### Install Dependencies

```bash
uv sync                    # core deps
uv sync --extra dev        # + pytest
uv sync --extra train      # + transformers, torch, trl
```

### Run Tests

```bash
uv run pytest              # 27 tests β€” universe, sandbox, reward, DSL, dual_fluid
```

### Start the OpenEnv Server

```bash
uv run uvicorn server.app:app --reload
```

### Smoke Test (Dual-Fluid Universe)

```bash
uv run python -m ontology_architect.scripts.smoke \
  --config configs/dual_fluid_demo.json \
  --baseline dual_fluid_dsl \
  --steps 5
```

### Generate Discovery Visualization

```bash
uv run python -m ontology_architect.scripts.visualize_discovery \
  --config configs/dual_fluid_demo.json \
  --baseline dual_fluid_dsl \
  --compare-baseline static \
  --output artifacts/results/discovery_plot.png
```

### Training (Colab / GPU)

See the [Colab notebook](https://colab.research.google.com/github/rohan-27p/ontology_architect/blob/main/Alien_Physics_Discovery_OpenEnv.ipynb) for the full GRPO training pipeline:

```bash
# Step 1: Oracle curriculum generation
python -m scripts.generate_curriculum \
  --config configs/dual_fluid_demo.json \
  --output artifacts/curriculum/oracle.jsonl \
  --episodes 60

# Step 2: SFT warm-start (teaches DSL format)
python -m scripts.train_sft \
  --config configs/dual_fluid_demo.json \
  --model-id Qwen/Qwen2.5-Coder-1.5B-Instruct \
  --data artifacts/curriculum/oracle.jsonl \
  --output-dir artifacts/sft_checkpoint \
  --max-steps 200 --batch-size 2

# Step 3: Group Reward Optimization (learns to write better theories)
python -m scripts.train_gro \
  --config configs/dual_fluid_demo.json \
  --model-id artifacts/sft_checkpoint \
  --group-size 4 \
  --max-steps 50 \
  --output-dir artifacts/checkpoints
```

### Docker

```bash
docker build -t ontology_architect-env:latest -f server/Dockerfile .
docker run --rm -p 8000:8000 ontology_architect-env:latest
```

## Configuration

| Config | Purpose |
|---|---|
| `configs/tiny_smoke.json` | Fast local testing (thermal_split, 3 steps) |
| `configs/dual_fluid_demo.json` | **Demo showcase** (dual_fluid, 100 steps, paradigm shifts) |
| `configs/full_research.json` | Extended research benchmarks |

## Architecture

```mermaid
graph TD
    Agent["LLM Agent"] -->|"Theory DSL JSON"| Env["OntologyArchitectEnvironment"]
    Env -->|"validate + render"| DSL["TheoryDSL Parser"]
    DSL -->|"Python Theory class"| Sandbox["TheorySandbox"]
    Sandbox -->|"predictions"| Reward["RewardScorer"]
    Env -->|"hidden future"| Reward
    Reward -->|"decomposed reward"| Env
    Env -->|"observation + peer review"| Agent
    
    Universe["ProceduralAlienUniverse"] -->|"sensor records"| Env
    Universe -.->|"hidden latent state"| Viz["visualize_discovery.py"]
```

## Theory Module API (Python Mode)

For advanced agents, raw Python `Theory` classes are also supported:

```python
import math

class Theory:
    def fit(self, history):
        self.last = dict(history[-1]["sensors"]) if history else {}
        return {"records": len(history)}

    def predict(self, window):
        return [{"sensors": dict(self.last), "anomaly_prob": 0.05}
                for _ in range(window["steps"])]

    def log_prob(self, observations):
        return -float(len(observations))

    def describe(self):
        return "Static persistence theory."
```

## License

BSD-style license. See LICENSE file.