genforge 0.1.0__tar.gz
This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
- genforge-0.1.0/LICENSE +21 -0
- genforge-0.1.0/PKG-INFO +124 -0
- genforge-0.1.0/README.md +93 -0
- genforge-0.1.0/pyproject.toml +96 -0
- genforge-0.1.0/src/forge/__init__.py +3 -0
- genforge-0.1.0/src/forge/cli.py +143 -0
- genforge-0.1.0/src/forge/configs/config.yaml +22 -0
- genforge-0.1.0/src/forge/configs/control/.gitkeep +1 -0
- genforge-0.1.0/src/forge/configs/control/cbf.yaml +3 -0
- genforge-0.1.0/src/forge/configs/control/fbsde_control.yaml +4 -0
- genforge-0.1.0/src/forge/configs/control/guidance.yaml +4 -0
- genforge-0.1.0/src/forge/configs/control/projection.yaml +2 -0
- genforge-0.1.0/src/forge/configs/control/value_guidance.yaml +5 -0
- genforge-0.1.0/src/forge/configs/cost/.gitkeep +1 -0
- genforge-0.1.0/src/forge/configs/cost/barrier.yaml +4 -0
- genforge-0.1.0/src/forge/configs/cost/halfspace.yaml +4 -0
- genforge-0.1.0/src/forge/configs/cost/reward.yaml +4 -0
- genforge-0.1.0/src/forge/configs/dataset/.gitkeep +1 -0
- genforge-0.1.0/src/forge/configs/environment/.gitkeep +1 -0
- genforge-0.1.0/src/forge/configs/method/.gitkeep +1 -0
- genforge-0.1.0/src/forge/configs/method/conditional.yaml +3 -0
- genforge-0.1.0/src/forge/configs/method/d3pm.yaml +2 -0
- genforge-0.1.0/src/forge/configs/method/ddpm.yaml +2 -0
- genforge-0.1.0/src/forge/configs/method/ddpm_huber.yaml +3 -0
- genforge-0.1.0/src/forge/configs/method/fbsde.yaml +3 -0
- genforge-0.1.0/src/forge/configs/method/flow_matching.yaml +2 -0
- genforge-0.1.0/src/forge/configs/method/mdlm.yaml +2 -0
- genforge-0.1.0/src/forge/configs/method/ot_cfm.yaml +2 -0
- genforge-0.1.0/src/forge/configs/method/sedd.yaml +2 -0
- genforge-0.1.0/src/forge/configs/method/value_training.yaml +4 -0
- genforge-0.1.0/src/forge/configs/model/.gitkeep +1 -0
- genforge-0.1.0/src/forge/configs/model/categorical_mlp.yaml +5 -0
- genforge-0.1.0/src/forge/configs/model/mlp.yaml +6 -0
- genforge-0.1.0/src/forge/configs/model/temporal_unet.yaml +6 -0
- genforge-0.1.0/src/forge/configs/model/temporal_unet_janner.yaml +10 -0
- genforge-0.1.0/src/forge/configs/model/transformer.yaml +8 -0
- genforge-0.1.0/src/forge/configs/model/value_mlp.yaml +5 -0
- genforge-0.1.0/src/forge/configs/preprocessor/.gitkeep +1 -0
- genforge-0.1.0/src/forge/configs/preprocessor/minmax.yaml +2 -0
- genforge-0.1.0/src/forge/configs/preprocessor/standardize.yaml +2 -0
- genforge-0.1.0/src/forge/configs/runner/.gitkeep +1 -0
- genforge-0.1.0/src/forge/configs/runner/planning.yaml +7 -0
- genforge-0.1.0/src/forge/configs/runner/policy_training.yaml +9 -0
- genforge-0.1.0/src/forge/configs/runner/training.yaml +20 -0
- genforge-0.1.0/src/forge/configs/runner/value_training.yaml +7 -0
- genforge-0.1.0/src/forge/configs/sampler/.gitkeep +1 -0
- genforge-0.1.0/src/forge/configs/sampler/ddim.yaml +3 -0
- genforge-0.1.0/src/forge/configs/sampler/ddpm.yaml +2 -0
- genforge-0.1.0/src/forge/configs/sampler/flow.yaml +3 -0
- genforge-0.1.0/src/forge/configs/sampler/interpolant.yaml +3 -0
- genforge-0.1.0/src/forge/configs/sampler/sedd.yaml +2 -0
- genforge-0.1.0/src/forge/configs/sampler/tau_leaping.yaml +2 -0
- genforge-0.1.0/src/forge/configs/schedule/.gitkeep +1 -0
- genforge-0.1.0/src/forge/configs/schedule/absorbing.yaml +3 -0
- genforge-0.1.0/src/forge/configs/schedule/cfm_linear.yaml +3 -0
- genforge-0.1.0/src/forge/configs/schedule/linear_flow.yaml +2 -0
- genforge-0.1.0/src/forge/configs/schedule/si_trig.yaml +2 -0
- genforge-0.1.0/src/forge/configs/schedule/uniform_discrete.yaml +3 -0
- genforge-0.1.0/src/forge/configs/schedule/vp_cosine.yaml +4 -0
- genforge-0.1.0/src/forge/configs/schedule/vp_linear.yaml +4 -0
- genforge-0.1.0/src/forge/configs/space/.gitkeep +1 -0
- genforge-0.1.0/src/forge/configs/space/discrete.yaml +4 -0
- genforge-0.1.0/src/forge/configs/space/euclidean.yaml +3 -0
- genforge-0.1.0/src/forge/configs/visualizer/.gitkeep +1 -0
- genforge-0.1.0/src/forge/configs/visualizer/trajectory.yaml +2 -0
- genforge-0.1.0/src/forge/control/__init__.py +1 -0
- genforge-0.1.0/src/forge/control/cbf.py +35 -0
- genforge-0.1.0/src/forge/control/fbsde_control.py +23 -0
- genforge-0.1.0/src/forge/control/guidance.py +35 -0
- genforge-0.1.0/src/forge/control/projection.py +65 -0
- genforge-0.1.0/src/forge/control/value_guidance.py +79 -0
- genforge-0.1.0/src/forge/core/__init__.py +1 -0
- genforge-0.1.0/src/forge/core/builder.py +236 -0
- genforge-0.1.0/src/forge/core/checkpoint.py +88 -0
- genforge-0.1.0/src/forge/core/compose.py +44 -0
- genforge-0.1.0/src/forge/core/interfaces.py +429 -0
- genforge-0.1.0/src/forge/core/plugins.py +69 -0
- genforge-0.1.0/src/forge/core/protocols.py +143 -0
- genforge-0.1.0/src/forge/core/registry.py +105 -0
- genforge-0.1.0/src/forge/core/resolvers.py +23 -0
- genforge-0.1.0/src/forge/core/types.py +33 -0
- genforge-0.1.0/src/forge/costs/__init__.py +1 -0
- genforge-0.1.0/src/forge/costs/ball.py +44 -0
- genforge-0.1.0/src/forge/costs/barrier.py +40 -0
- genforge-0.1.0/src/forge/costs/box.py +34 -0
- genforge-0.1.0/src/forge/costs/halfspace.py +45 -0
- genforge-0.1.0/src/forge/costs/likelihood.py +32 -0
- genforge-0.1.0/src/forge/costs/reward.py +35 -0
- genforge-0.1.0/src/forge/datasets/__init__.py +7 -0
- genforge-0.1.0/src/forge/environments/__init__.py +7 -0
- genforge-0.1.0/src/forge/methods/__init__.py +1 -0
- genforge-0.1.0/src/forge/methods/conditional.py +54 -0
- genforge-0.1.0/src/forge/methods/d3pm.py +36 -0
- genforge-0.1.0/src/forge/methods/ddpm.py +43 -0
- genforge-0.1.0/src/forge/methods/ddpm_huber.py +48 -0
- genforge-0.1.0/src/forge/methods/fbsde.py +50 -0
- genforge-0.1.0/src/forge/methods/flow_matching.py +40 -0
- genforge-0.1.0/src/forge/methods/mdlm.py +50 -0
- genforge-0.1.0/src/forge/methods/ot_cfm.py +59 -0
- genforge-0.1.0/src/forge/methods/sedd.py +82 -0
- genforge-0.1.0/src/forge/methods/value_training.py +38 -0
- genforge-0.1.0/src/forge/models/__init__.py +1 -0
- genforge-0.1.0/src/forge/models/categorical.py +52 -0
- genforge-0.1.0/src/forge/models/mlp.py +70 -0
- genforge-0.1.0/src/forge/models/temporal_unet.py +104 -0
- genforge-0.1.0/src/forge/models/temporal_unet_janner.py +332 -0
- genforge-0.1.0/src/forge/models/transformer.py +147 -0
- genforge-0.1.0/src/forge/models/value.py +32 -0
- genforge-0.1.0/src/forge/preprocessing/__init__.py +1 -0
- genforge-0.1.0/src/forge/preprocessing/minmax.py +54 -0
- genforge-0.1.0/src/forge/preprocessing/standardize.py +57 -0
- genforge-0.1.0/src/forge/runners/__init__.py +1 -0
- genforge-0.1.0/src/forge/runners/multistep.py +226 -0
- genforge-0.1.0/src/forge/runners/planning.py +66 -0
- genforge-0.1.0/src/forge/runners/policy_training.py +138 -0
- genforge-0.1.0/src/forge/runners/training.py +385 -0
- genforge-0.1.0/src/forge/runners/value_training.py +18 -0
- genforge-0.1.0/src/forge/samplers/__init__.py +1 -0
- genforge-0.1.0/src/forge/samplers/ddim.py +45 -0
- genforge-0.1.0/src/forge/samplers/ddpm.py +56 -0
- genforge-0.1.0/src/forge/samplers/flow.py +55 -0
- genforge-0.1.0/src/forge/samplers/interpolant.py +61 -0
- genforge-0.1.0/src/forge/samplers/sedd.py +45 -0
- genforge-0.1.0/src/forge/samplers/tau_leaping.py +22 -0
- genforge-0.1.0/src/forge/schedules/__init__.py +1 -0
- genforge-0.1.0/src/forge/schedules/discrete.py +146 -0
- genforge-0.1.0/src/forge/schedules/flow.py +125 -0
- genforge-0.1.0/src/forge/schedules/vp.py +110 -0
- genforge-0.1.0/src/forge/spaces/__init__.py +1 -0
- genforge-0.1.0/src/forge/spaces/discrete.py +55 -0
- genforge-0.1.0/src/forge/spaces/euclidean.py +44 -0
- genforge-0.1.0/src/forge/utils/__init__.py +1 -0
- genforge-0.1.0/src/forge/utils/ema.py +82 -0
- genforge-0.1.0/src/forge/utils/logging.py +118 -0
- genforge-0.1.0/src/forge/utils/lora.py +126 -0
- genforge-0.1.0/src/forge/utils/seeding.py +34 -0
- genforge-0.1.0/src/forge/utils/torch_utils.py +23 -0
- genforge-0.1.0/src/forge/visualizations/__init__.py +1 -0
- genforge-0.1.0/src/forge/visualizations/trajectory.py +42 -0
genforge-0.1.0/LICENSE
ADDED
|
@@ -0,0 +1,21 @@
|
|
|
1
|
+
MIT License
|
|
2
|
+
|
|
3
|
+
Copyright (c) 2026 Monish Lokhande
|
|
4
|
+
|
|
5
|
+
Permission is hereby granted, free of charge, to any person obtaining a copy
|
|
6
|
+
of this software and associated documentation files (the "Software"), to deal
|
|
7
|
+
in the Software without restriction, including without limitation the rights
|
|
8
|
+
to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
|
|
9
|
+
copies of the Software, and to permit persons to whom the Software is
|
|
10
|
+
furnished to do so, subject to the following conditions:
|
|
11
|
+
|
|
12
|
+
The above copyright notice and this permission notice shall be included in all
|
|
13
|
+
copies or substantial portions of the Software.
|
|
14
|
+
|
|
15
|
+
THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
|
|
16
|
+
IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
|
|
17
|
+
FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
|
|
18
|
+
AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
|
|
19
|
+
LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
|
|
20
|
+
OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
|
|
21
|
+
SOFTWARE.
|
genforge-0.1.0/PKG-INFO
ADDED
|
@@ -0,0 +1,124 @@
|
|
|
1
|
+
Metadata-Version: 2.4
|
|
2
|
+
Name: genforge
|
|
3
|
+
Version: 0.1.0
|
|
4
|
+
Summary: A unified, PyTorch-based framework for generative-modeling techniques with a clean control layer.
|
|
5
|
+
Keywords: generative-models,diffusion,flow-matching,score-based,discrete-diffusion,guidance,control,pytorch
|
|
6
|
+
Author: Monish Lokhande
|
|
7
|
+
License-Expression: MIT
|
|
8
|
+
License-File: LICENSE
|
|
9
|
+
Classifier: Development Status :: 3 - Alpha
|
|
10
|
+
Classifier: Intended Audience :: Science/Research
|
|
11
|
+
Classifier: Programming Language :: Python :: 3.11
|
|
12
|
+
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
|
|
13
|
+
Requires-Dist: torch>=2.2
|
|
14
|
+
Requires-Dist: numpy>=1.26
|
|
15
|
+
Requires-Dist: hydra-core>=1.3
|
|
16
|
+
Requires-Dist: omegaconf>=2.3
|
|
17
|
+
Requires-Dist: scipy>=1.11 ; extra == 'flow'
|
|
18
|
+
Requires-Dist: wandb>=0.16 ; extra == 'logging'
|
|
19
|
+
Requires-Dist: tqdm>=4.66 ; extra == 'logging'
|
|
20
|
+
Requires-Dist: tiktoken>=0.7 ; extra == 'text'
|
|
21
|
+
Requires-Dist: datasets>=2.18 ; extra == 'text'
|
|
22
|
+
Requires-Python: >=3.11
|
|
23
|
+
Project-URL: Homepage, https://github.com/MonishLokhande/genforge
|
|
24
|
+
Project-URL: Repository, https://github.com/MonishLokhande/genforge
|
|
25
|
+
Project-URL: Documentation, https://github.com/MonishLokhande/genforge#readme
|
|
26
|
+
Project-URL: Issues, https://github.com/MonishLokhande/genforge/issues
|
|
27
|
+
Provides-Extra: flow
|
|
28
|
+
Provides-Extra: logging
|
|
29
|
+
Provides-Extra: text
|
|
30
|
+
Description-Content-Type: text/markdown
|
|
31
|
+
|
|
32
|
+
# genforge
|
|
33
|
+
|
|
34
|
+
A unified, PyTorch framework for generative modeling — score/diffusion SDEs, probability-flow
|
|
35
|
+
ODEs, flow matching / OT-CFM, stochastic interpolants, and discrete (D3PM-style) diffusion — with a
|
|
36
|
+
clean **control layer** for conditioning, guidance, constraints, and amortized control.
|
|
37
|
+
|
|
38
|
+
The idea: a generative model is one process that turns a simple starting distribution (usually noise)
|
|
39
|
+
into data, and every way of steering it —
|
|
40
|
+
conditioning, guidance, constraints, planning — is the **same move**: reweight that process so the
|
|
41
|
+
outcomes you want become more likely. The process, the steering, and how the steering is approximated
|
|
42
|
+
are separate, swappable parts. A full documentation site (architecture, experiments, extending) is coming soon.
|
|
43
|
+
|
|
44
|
+
## Install
|
|
45
|
+
|
|
46
|
+
**Use it as a library:**
|
|
47
|
+
|
|
48
|
+
```bash
|
|
49
|
+
pip install genforge # or: uv add genforge
|
|
50
|
+
```
|
|
51
|
+
|
|
52
|
+
The bundled `experiment/` tree used in the Quickstart below ships with the **source clone**, not
|
|
53
|
+
the PyPI wheel — running `forge train experiment=...` requires `git clone` + `uv sync`, not just
|
|
54
|
+
`pip install genforge`.
|
|
55
|
+
|
|
56
|
+
**Develop / run the bundled experiments:**
|
|
57
|
+
|
|
58
|
+
```bash
|
|
59
|
+
git clone https://github.com/MonishLokhande/genforge
|
|
60
|
+
cd genforge
|
|
61
|
+
uv sync # core (light: 2-D distributions)
|
|
62
|
+
uv sync --extra flow # + OT-CFM (scipy)
|
|
63
|
+
uv sync --extra text # + real BPE / streamed corpora (tiktoken, datasets)
|
|
64
|
+
uv sync --extra logging # + experiment logging (wandb) + progress bars (tqdm)
|
|
65
|
+
```
|
|
66
|
+
|
|
67
|
+
Robotics adapters are a **dependency group**, not an extra (one dependency installs from git):
|
|
68
|
+
`uv sync --group robotics` (mujoco, robomimic, gym-pusht/aloha, minari).
|
|
69
|
+
|
|
70
|
+
## Quickstart
|
|
71
|
+
|
|
72
|
+
```bash
|
|
73
|
+
uv run forge list # registered components
|
|
74
|
+
uv run forge train experiment=distributions/ddpm/base
|
|
75
|
+
uv run forge sample experiment=distributions/ddpm/base # or: sample checkpoint=<path>.pt
|
|
76
|
+
```
|
|
77
|
+
|
|
78
|
+
## Layout
|
|
79
|
+
|
|
80
|
+
Every component registers via `@register(category, name)` and is wired by a config-driven builder in
|
|
81
|
+
dependency order — adding one is a single decorated class plus a config leaf, no other wiring.
|
|
82
|
+
|
|
83
|
+
- **`src/forge/`** — the framework only: `core` (registry · builder · interfaces · protocols ·
|
|
84
|
+
plugins), `spaces`, `schedules`, `models`, `methods`, `samplers`, `costs`, `control`,
|
|
85
|
+
`preprocessing`, `runners`. Protocols, ABCs, and generic utilities — never concrete env code.
|
|
86
|
+
- **`envs/`** — concrete, swappable **data-source plugins** (environment + dataset + processor per
|
|
87
|
+
package). An experiment loads them via its `plugins:` field; they are not baked into the core.
|
|
88
|
+
Contracts: [`core/protocols.py`](src/forge/core/protocols.py).
|
|
89
|
+
- **`experiment/`** — Hydra base+delta bundles, selected with `experiment=<family>/<variant>/<method>`.
|
|
90
|
+
|
|
91
|
+
## Experiments
|
|
92
|
+
|
|
93
|
+
Selected with `experiment=<family>/<variant>/<method>`.
|
|
94
|
+
|
|
95
|
+
| Family | What |
|
|
96
|
+
|---|---|
|
|
97
|
+
| `distributions/*` | Continuous 2-D — DDPM, flow matching, stochastic-interpolant SDE, DDIM; control via projection / guidance / CBF; value guidance. |
|
|
98
|
+
| `discrete/d3pm/base` | Discrete (absorbing) diffusion on a toy categorical target. |
|
|
99
|
+
| `text/char/*` | Discrete diffusion LM, char-level — `d3pm` / `mdlm` / `sedd`. |
|
|
100
|
+
| `text/tinystories/*` | The **same** methods at real GPT-2 BPE (vocab 50258, needs `--extra text`). |
|
|
101
|
+
| `trajectory/plan/base` | Goal-conditioned trajectory planning (flat-tensor windowing, endpoint-pinned). |
|
|
102
|
+
| `robotics/*` | Offline-RL trajectory planning (maze2d, locomotion) and closed-loop diffusion policies (robomimic, pusht, aloha); needs `--group robotics`. |
|
|
103
|
+
|
|
104
|
+
`text/char/*` and `text/tinystories/*` are two variants of one **`text` family** (a single env plugin
|
|
105
|
+
`envs.text` registers both) — the same absorbing + transformer + `{d3pm,mdlm,sedd}` rig; only the
|
|
106
|
+
tokenizer (char vs. real BPE) and scale differ.
|
|
107
|
+
|
|
108
|
+
```bash
|
|
109
|
+
uv run forge train experiment=text/char/d3pm/small # char-level LM
|
|
110
|
+
uv run --extra text forge train experiment=text/tinystories/d3pm/small # same method, real BPE
|
|
111
|
+
```
|
|
112
|
+
|
|
113
|
+
## Documentation
|
|
114
|
+
|
|
115
|
+
A full documentation site — installation, architecture, experiments, and extending — is coming soon.
|
|
116
|
+
Until then, this README plus the inline docstrings across `src/forge/` are the reference.
|
|
117
|
+
|
|
118
|
+
## License
|
|
119
|
+
|
|
120
|
+
MIT — see [LICENSE](LICENSE).
|
|
121
|
+
|
|
122
|
+
## Citation
|
|
123
|
+
|
|
124
|
+
If you use genforge in your work, please cite it — citation metadata is in [CITATION.cff](CITATION.cff).
|
genforge-0.1.0/README.md
ADDED
|
@@ -0,0 +1,93 @@
|
|
|
1
|
+
# genforge
|
|
2
|
+
|
|
3
|
+
A unified, PyTorch framework for generative modeling — score/diffusion SDEs, probability-flow
|
|
4
|
+
ODEs, flow matching / OT-CFM, stochastic interpolants, and discrete (D3PM-style) diffusion — with a
|
|
5
|
+
clean **control layer** for conditioning, guidance, constraints, and amortized control.
|
|
6
|
+
|
|
7
|
+
The idea: a generative model is one process that turns a simple starting distribution (usually noise)
|
|
8
|
+
into data, and every way of steering it —
|
|
9
|
+
conditioning, guidance, constraints, planning — is the **same move**: reweight that process so the
|
|
10
|
+
outcomes you want become more likely. The process, the steering, and how the steering is approximated
|
|
11
|
+
are separate, swappable parts. A full documentation site (architecture, experiments, extending) is coming soon.
|
|
12
|
+
|
|
13
|
+
## Install
|
|
14
|
+
|
|
15
|
+
**Use it as a library:**
|
|
16
|
+
|
|
17
|
+
```bash
|
|
18
|
+
pip install genforge # or: uv add genforge
|
|
19
|
+
```
|
|
20
|
+
|
|
21
|
+
The bundled `experiment/` tree used in the Quickstart below ships with the **source clone**, not
|
|
22
|
+
the PyPI wheel — running `forge train experiment=...` requires `git clone` + `uv sync`, not just
|
|
23
|
+
`pip install genforge`.
|
|
24
|
+
|
|
25
|
+
**Develop / run the bundled experiments:**
|
|
26
|
+
|
|
27
|
+
```bash
|
|
28
|
+
git clone https://github.com/MonishLokhande/genforge
|
|
29
|
+
cd genforge
|
|
30
|
+
uv sync # core (light: 2-D distributions)
|
|
31
|
+
uv sync --extra flow # + OT-CFM (scipy)
|
|
32
|
+
uv sync --extra text # + real BPE / streamed corpora (tiktoken, datasets)
|
|
33
|
+
uv sync --extra logging # + experiment logging (wandb) + progress bars (tqdm)
|
|
34
|
+
```
|
|
35
|
+
|
|
36
|
+
Robotics adapters are a **dependency group**, not an extra (one dependency installs from git):
|
|
37
|
+
`uv sync --group robotics` (mujoco, robomimic, gym-pusht/aloha, minari).
|
|
38
|
+
|
|
39
|
+
## Quickstart
|
|
40
|
+
|
|
41
|
+
```bash
|
|
42
|
+
uv run forge list # registered components
|
|
43
|
+
uv run forge train experiment=distributions/ddpm/base
|
|
44
|
+
uv run forge sample experiment=distributions/ddpm/base # or: sample checkpoint=<path>.pt
|
|
45
|
+
```
|
|
46
|
+
|
|
47
|
+
## Layout
|
|
48
|
+
|
|
49
|
+
Every component registers via `@register(category, name)` and is wired by a config-driven builder in
|
|
50
|
+
dependency order — adding one is a single decorated class plus a config leaf, no other wiring.
|
|
51
|
+
|
|
52
|
+
- **`src/forge/`** — the framework only: `core` (registry · builder · interfaces · protocols ·
|
|
53
|
+
plugins), `spaces`, `schedules`, `models`, `methods`, `samplers`, `costs`, `control`,
|
|
54
|
+
`preprocessing`, `runners`. Protocols, ABCs, and generic utilities — never concrete env code.
|
|
55
|
+
- **`envs/`** — concrete, swappable **data-source plugins** (environment + dataset + processor per
|
|
56
|
+
package). An experiment loads them via its `plugins:` field; they are not baked into the core.
|
|
57
|
+
Contracts: [`core/protocols.py`](src/forge/core/protocols.py).
|
|
58
|
+
- **`experiment/`** — Hydra base+delta bundles, selected with `experiment=<family>/<variant>/<method>`.
|
|
59
|
+
|
|
60
|
+
## Experiments
|
|
61
|
+
|
|
62
|
+
Selected with `experiment=<family>/<variant>/<method>`.
|
|
63
|
+
|
|
64
|
+
| Family | What |
|
|
65
|
+
|---|---|
|
|
66
|
+
| `distributions/*` | Continuous 2-D — DDPM, flow matching, stochastic-interpolant SDE, DDIM; control via projection / guidance / CBF; value guidance. |
|
|
67
|
+
| `discrete/d3pm/base` | Discrete (absorbing) diffusion on a toy categorical target. |
|
|
68
|
+
| `text/char/*` | Discrete diffusion LM, char-level — `d3pm` / `mdlm` / `sedd`. |
|
|
69
|
+
| `text/tinystories/*` | The **same** methods at real GPT-2 BPE (vocab 50258, needs `--extra text`). |
|
|
70
|
+
| `trajectory/plan/base` | Goal-conditioned trajectory planning (flat-tensor windowing, endpoint-pinned). |
|
|
71
|
+
| `robotics/*` | Offline-RL trajectory planning (maze2d, locomotion) and closed-loop diffusion policies (robomimic, pusht, aloha); needs `--group robotics`. |
|
|
72
|
+
|
|
73
|
+
`text/char/*` and `text/tinystories/*` are two variants of one **`text` family** (a single env plugin
|
|
74
|
+
`envs.text` registers both) — the same absorbing + transformer + `{d3pm,mdlm,sedd}` rig; only the
|
|
75
|
+
tokenizer (char vs. real BPE) and scale differ.
|
|
76
|
+
|
|
77
|
+
```bash
|
|
78
|
+
uv run forge train experiment=text/char/d3pm/small # char-level LM
|
|
79
|
+
uv run --extra text forge train experiment=text/tinystories/d3pm/small # same method, real BPE
|
|
80
|
+
```
|
|
81
|
+
|
|
82
|
+
## Documentation
|
|
83
|
+
|
|
84
|
+
A full documentation site — installation, architecture, experiments, and extending — is coming soon.
|
|
85
|
+
Until then, this README plus the inline docstrings across `src/forge/` are the reference.
|
|
86
|
+
|
|
87
|
+
## License
|
|
88
|
+
|
|
89
|
+
MIT — see [LICENSE](LICENSE).
|
|
90
|
+
|
|
91
|
+
## Citation
|
|
92
|
+
|
|
93
|
+
If you use genforge in your work, please cite it — citation metadata is in [CITATION.cff](CITATION.cff).
|
|
@@ -0,0 +1,96 @@
|
|
|
1
|
+
[project]
|
|
2
|
+
name = "genforge"
|
|
3
|
+
version = "0.1.0"
|
|
4
|
+
description = "A unified, PyTorch-based framework for generative-modeling techniques with a clean control layer."
|
|
5
|
+
readme = "README.md"
|
|
6
|
+
license = "MIT"
|
|
7
|
+
license-files = ["LICENSE"]
|
|
8
|
+
authors = [
|
|
9
|
+
{ name = "Monish Lokhande" }
|
|
10
|
+
]
|
|
11
|
+
requires-python = ">=3.11"
|
|
12
|
+
keywords = [
|
|
13
|
+
"generative-models", "diffusion", "flow-matching", "score-based",
|
|
14
|
+
"discrete-diffusion", "guidance", "control", "pytorch",
|
|
15
|
+
]
|
|
16
|
+
classifiers = [
|
|
17
|
+
"Development Status :: 3 - Alpha",
|
|
18
|
+
"Intended Audience :: Science/Research",
|
|
19
|
+
"Programming Language :: Python :: 3.11",
|
|
20
|
+
"Topic :: Scientific/Engineering :: Artificial Intelligence",
|
|
21
|
+
]
|
|
22
|
+
dependencies = [
|
|
23
|
+
"torch>=2.2",
|
|
24
|
+
"numpy>=1.26",
|
|
25
|
+
"hydra-core>=1.3",
|
|
26
|
+
"omegaconf>=2.3",
|
|
27
|
+
]
|
|
28
|
+
|
|
29
|
+
[project.optional-dependencies]
|
|
30
|
+
# Optional extras keep the 2-D distributions core light.
|
|
31
|
+
flow = ["scipy>=1.11"] # minibatch-OT coupling for ot_cfm
|
|
32
|
+
text = ["tiktoken>=0.7", "datasets>=2.18"] # real GPT-2 BPE + streamed corpora (e.g. TinyStories)
|
|
33
|
+
logging = ["wandb>=0.16", "tqdm>=4.66"] # optional experiment logging + progress bars (off by default)
|
|
34
|
+
|
|
35
|
+
[project.urls]
|
|
36
|
+
Homepage = "https://github.com/MonishLokhande/genforge"
|
|
37
|
+
Repository = "https://github.com/MonishLokhande/genforge"
|
|
38
|
+
Documentation = "https://github.com/MonishLokhande/genforge#readme"
|
|
39
|
+
Issues = "https://github.com/MonishLokhande/genforge/issues"
|
|
40
|
+
|
|
41
|
+
[dependency-groups]
|
|
42
|
+
dev = [
|
|
43
|
+
"pytest>=8.0",
|
|
44
|
+
]
|
|
45
|
+
docs = [
|
|
46
|
+
"mkdocs-material>=9.7.6",
|
|
47
|
+
]
|
|
48
|
+
# Offline-RL / robot adapters (maze2d / locomotion / hf_lowdim / robomimic). Pins copied from an offline-RL reference implementation.
|
|
49
|
+
# A dependency-group, not an extra: it contains a git direct-reference (robomimic) that PyPI rejects
|
|
50
|
+
# in published wheel metadata. Sync locally with: uv sync --group robotics
|
|
51
|
+
robotics = [
|
|
52
|
+
"h5py>=3.16", # d4rl / robomimic HDF5 adapters read buffers directly
|
|
53
|
+
"gym-pusht>=0.1.6",
|
|
54
|
+
"gym-aloha>=0.1.3",
|
|
55
|
+
"gymnasium>=1.3.0",
|
|
56
|
+
"gymnasium-robotics>=1.4.2",
|
|
57
|
+
"datasets>=2.19.0",
|
|
58
|
+
"minari[hdf5,hf]>=0.5.3",
|
|
59
|
+
"mujoco>=3.8.0",
|
|
60
|
+
"robomimic @ git+https://github.com/ARISE-Initiative/robomimic.git@v0.5.0",
|
|
61
|
+
"robosuite>=1.5", # unpinned backtracks to 0.3.0 -> deprecated mujoco-py; >=1.5 avoids it
|
|
62
|
+
]
|
|
63
|
+
|
|
64
|
+
[tool.uv]
|
|
65
|
+
constraint-dependencies = ["pymunk<7", "numba>=0.65"] # keep robosuite off the py<3.10 numba 0.53
|
|
66
|
+
# robomimic 0.5 hard-pins an old HF stack; relax to >= so the modern stack is kept.
|
|
67
|
+
override-dependencies = [
|
|
68
|
+
"transformers>=4.41.2",
|
|
69
|
+
"diffusers>=0.11.1",
|
|
70
|
+
"huggingface-hub>=0.23.4",
|
|
71
|
+
]
|
|
72
|
+
|
|
73
|
+
[tool.uv.build-backend]
|
|
74
|
+
module-name = "forge"
|
|
75
|
+
|
|
76
|
+
[[tool.uv.index]]
|
|
77
|
+
name = "testpypi"
|
|
78
|
+
url = "https://test.pypi.org/simple/"
|
|
79
|
+
publish-url = "https://test.pypi.org/legacy/"
|
|
80
|
+
default = false
|
|
81
|
+
# Publish-only: keep this index out of dependency resolution so a stray package on TestPyPI
|
|
82
|
+
# (e.g. an old hydra-core) can't shadow the real one on PyPI. `uv publish --index testpypi`
|
|
83
|
+
# still works via publish-url. Scope any TestPyPI-sourced dep via [tool.uv.sources] instead.
|
|
84
|
+
explicit = true
|
|
85
|
+
|
|
86
|
+
[tool.pytest.ini_options]
|
|
87
|
+
# Put the repo root on sys.path for the test session so the repo-root `envs/*` plugin packages are
|
|
88
|
+
# importable (the package install only adds `src/` via its .pth). conftest.py then registers them.
|
|
89
|
+
pythonpath = ["."]
|
|
90
|
+
|
|
91
|
+
[project.scripts]
|
|
92
|
+
forge = "forge.cli:main"
|
|
93
|
+
|
|
94
|
+
[build-system]
|
|
95
|
+
requires = ["uv_build>=0.11.7,<0.12.0"]
|
|
96
|
+
build-backend = "uv_build"
|
|
@@ -0,0 +1,143 @@
|
|
|
1
|
+
"""The ``forge`` command-line entrypoint: ``list`` / ``train`` / ``sample``.
|
|
2
|
+
|
|
3
|
+
``list`` imports the built-ins so registrations fire, then prints the registered components by
|
|
4
|
+
category. ``train`` / ``sample`` compose a Hydra config from an ``experiment=`` selection, build the
|
|
5
|
+
runner, and run. ``sample checkpoint=<path.pt>`` rebuilds everything from the self-contained
|
|
6
|
+
checkpoint alone (Invariant 5). All three fail loudly on misconfiguration.
|
|
7
|
+
"""
|
|
8
|
+
|
|
9
|
+
from __future__ import annotations
|
|
10
|
+
|
|
11
|
+
import argparse
|
|
12
|
+
import sys
|
|
13
|
+
from typing import Optional, Sequence
|
|
14
|
+
|
|
15
|
+
from .core import registry
|
|
16
|
+
from .core.builder import build, import_builtin_components
|
|
17
|
+
|
|
18
|
+
_EXPERIMENT_HINT = (
|
|
19
|
+
"Select one with `experiment=<env>/<params>/<method>` "
|
|
20
|
+
"(e.g. `forge train experiment=distributions/ddpm/base`)."
|
|
21
|
+
)
|
|
22
|
+
|
|
23
|
+
|
|
24
|
+
def _split_overrides(overrides: Sequence[str]) -> dict:
|
|
25
|
+
"""Parse ``key=value`` overrides into a flat dict (first '=')."""
|
|
26
|
+
flat: dict[str, str] = {}
|
|
27
|
+
for o in overrides:
|
|
28
|
+
if "=" in o:
|
|
29
|
+
k, v = o.split("=", 1)
|
|
30
|
+
flat[k] = v
|
|
31
|
+
return flat
|
|
32
|
+
|
|
33
|
+
|
|
34
|
+
def _cmd_list(_args: argparse.Namespace) -> int:
|
|
35
|
+
import_builtin_components()
|
|
36
|
+
# Concrete envs are plugins (no experiment selected here), so import the bundled env packages
|
|
37
|
+
# too — otherwise `list` would omit environments/datasets/env-preprocessors.
|
|
38
|
+
from .core.plugins import load_bundled_envs
|
|
39
|
+
|
|
40
|
+
load_bundled_envs()
|
|
41
|
+
reg = registry.registered()
|
|
42
|
+
print("forge components")
|
|
43
|
+
print("===================")
|
|
44
|
+
for category in registry.CATEGORIES:
|
|
45
|
+
comps = reg.get(category, {})
|
|
46
|
+
names = ", ".join(comps) if comps else "(none yet)"
|
|
47
|
+
print(f" {category:<13} {names}")
|
|
48
|
+
for category in [c for c in reg if c not in registry.CATEGORIES]:
|
|
49
|
+
print(f" {category:<13} {', '.join(reg[category])}")
|
|
50
|
+
return 0
|
|
51
|
+
|
|
52
|
+
|
|
53
|
+
def _run_from_config(overrides: Sequence[str], action: str) -> int:
|
|
54
|
+
from omegaconf import OmegaConf
|
|
55
|
+
|
|
56
|
+
from .core.compose import compose_config
|
|
57
|
+
|
|
58
|
+
cfg = compose_config(overrides)
|
|
59
|
+
runner = build(cfg)
|
|
60
|
+
runner.resolved_config = OmegaConf.to_container(cfg, resolve=True)
|
|
61
|
+
|
|
62
|
+
if action == "train":
|
|
63
|
+
runner.train()
|
|
64
|
+
metrics = runner.evaluate()
|
|
65
|
+
print(f"[train] done. eval: {metrics}")
|
|
66
|
+
return 0
|
|
67
|
+
|
|
68
|
+
# sample from an experiment: load its configured checkpoint if present.
|
|
69
|
+
ckpt_path = getattr(runner, "ckpt_path", None)
|
|
70
|
+
if ckpt_path:
|
|
71
|
+
from pathlib import Path
|
|
72
|
+
|
|
73
|
+
from .core.checkpoint import load_checkpoint
|
|
74
|
+
|
|
75
|
+
if not Path(ckpt_path).exists():
|
|
76
|
+
print(
|
|
77
|
+
f"`forge sample` found no checkpoint at {ckpt_path!r}. Train first "
|
|
78
|
+
f"(`forge train experiment=...`) or pass `checkpoint=<path.pt>`.",
|
|
79
|
+
file=sys.stderr,
|
|
80
|
+
)
|
|
81
|
+
return 1
|
|
82
|
+
runner.load_state(load_checkpoint(ckpt_path))
|
|
83
|
+
metrics = runner.evaluate()
|
|
84
|
+
print(f"[sample] {metrics}")
|
|
85
|
+
return 0
|
|
86
|
+
|
|
87
|
+
|
|
88
|
+
def _cmd_train(args: argparse.Namespace) -> int:
|
|
89
|
+
flat = _split_overrides(args.overrides)
|
|
90
|
+
if "experiment" not in flat:
|
|
91
|
+
print(f"`forge train` requires an experiment selection. {_EXPERIMENT_HINT}", file=sys.stderr)
|
|
92
|
+
return 2
|
|
93
|
+
return _run_from_config(args.overrides, "train")
|
|
94
|
+
|
|
95
|
+
|
|
96
|
+
def _cmd_sample(args: argparse.Namespace) -> int:
|
|
97
|
+
flat = _split_overrides(args.overrides)
|
|
98
|
+
if "checkpoint" in flat:
|
|
99
|
+
# Self-contained path: rebuild from the .pt alone (Invariant 5).
|
|
100
|
+
from .runners.training import TrainingRunner
|
|
101
|
+
|
|
102
|
+
runner = TrainingRunner.from_checkpoint(flat["checkpoint"], build_fn=build)
|
|
103
|
+
metrics = runner.evaluate()
|
|
104
|
+
print(f"[sample] from checkpoint {flat['checkpoint']}: {metrics}")
|
|
105
|
+
return 0
|
|
106
|
+
if "experiment" not in flat:
|
|
107
|
+
print(
|
|
108
|
+
f"`forge sample` requires `experiment=...` or `checkpoint=<path.pt>`. {_EXPERIMENT_HINT}",
|
|
109
|
+
file=sys.stderr,
|
|
110
|
+
)
|
|
111
|
+
return 2
|
|
112
|
+
return _run_from_config(args.overrides, "sample")
|
|
113
|
+
|
|
114
|
+
|
|
115
|
+
def _build_parser() -> argparse.ArgumentParser:
|
|
116
|
+
parser = argparse.ArgumentParser(
|
|
117
|
+
prog="forge",
|
|
118
|
+
description="A unified framework for generative modeling with a clean control layer.",
|
|
119
|
+
)
|
|
120
|
+
sub = parser.add_subparsers(dest="command", required=True)
|
|
121
|
+
|
|
122
|
+
p_list = sub.add_parser("list", help="List registered components by category.")
|
|
123
|
+
p_list.set_defaults(func=_cmd_list)
|
|
124
|
+
|
|
125
|
+
p_train = sub.add_parser("train", help="Train a model from an experiment config.")
|
|
126
|
+
p_train.add_argument("overrides", nargs="*", help="Hydra-style overrides, e.g. experiment=...")
|
|
127
|
+
p_train.set_defaults(func=_cmd_train)
|
|
128
|
+
|
|
129
|
+
p_sample = sub.add_parser("sample", help="Sample from a trained model or checkpoint.")
|
|
130
|
+
p_sample.add_argument("overrides", nargs="*", help="experiment=... or checkpoint=<path.pt>")
|
|
131
|
+
p_sample.set_defaults(func=_cmd_sample)
|
|
132
|
+
|
|
133
|
+
return parser
|
|
134
|
+
|
|
135
|
+
|
|
136
|
+
def main(argv: Optional[Sequence[str]] = None) -> int:
|
|
137
|
+
parser = _build_parser()
|
|
138
|
+
args = parser.parse_args(argv)
|
|
139
|
+
return args.func(args)
|
|
140
|
+
|
|
141
|
+
|
|
142
|
+
if __name__ == "__main__":
|
|
143
|
+
raise SystemExit(main())
|
|
@@ -0,0 +1,22 @@
|
|
|
1
|
+
# genforge root config (Hydra base+delta).
|
|
2
|
+
#
|
|
3
|
+
# Per-category config groups live in the sibling directories (space/, schedule/, model/, ...).
|
|
4
|
+
# Experiments are base+delta bundles in the repo-root `experiment/` tree, which is added to the
|
|
5
|
+
# Hydra searchpath at runtime by the CLI (env var GENFORGE_EXP_ROOT). Select an experiment with
|
|
6
|
+
# `experiment=<env>/<params>/<method>` — it composes the component groups it needs.
|
|
7
|
+
|
|
8
|
+
defaults:
|
|
9
|
+
- _self_
|
|
10
|
+
- experiment: ??? # mandatory: train/sample require an experiment selection
|
|
11
|
+
|
|
12
|
+
# Global run settings (resolved into the checkpoint, Invariant 5).
|
|
13
|
+
seed: 0
|
|
14
|
+
|
|
15
|
+
hydra:
|
|
16
|
+
searchpath:
|
|
17
|
+
- file://${oc.env:GENFORGE_EXP_ROOT}
|
|
18
|
+
job:
|
|
19
|
+
chdir: false
|
|
20
|
+
output_subdir: null
|
|
21
|
+
run:
|
|
22
|
+
dir: .
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
# Phase 0: empty config group for `control`. Concrete options arrive in later phases.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
# Phase 0: empty config group for `cost`. Concrete options arrive in later phases.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
# Phase 0: empty config group for `dataset`. Concrete options arrive in later phases.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
# Phase 0: empty config group for `environment`. Concrete options arrive in later phases.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
# Phase 0: empty config group for `method`. Concrete options arrive in later phases.
|
|
@@ -0,0 +1 @@
|
|
|
1
|
+
# Phase 0: empty config group for `model`. Concrete options arrive in later phases.
|