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.
Files changed (139) hide show
  1. genforge-0.1.0/LICENSE +21 -0
  2. genforge-0.1.0/PKG-INFO +124 -0
  3. genforge-0.1.0/README.md +93 -0
  4. genforge-0.1.0/pyproject.toml +96 -0
  5. genforge-0.1.0/src/forge/__init__.py +3 -0
  6. genforge-0.1.0/src/forge/cli.py +143 -0
  7. genforge-0.1.0/src/forge/configs/config.yaml +22 -0
  8. genforge-0.1.0/src/forge/configs/control/.gitkeep +1 -0
  9. genforge-0.1.0/src/forge/configs/control/cbf.yaml +3 -0
  10. genforge-0.1.0/src/forge/configs/control/fbsde_control.yaml +4 -0
  11. genforge-0.1.0/src/forge/configs/control/guidance.yaml +4 -0
  12. genforge-0.1.0/src/forge/configs/control/projection.yaml +2 -0
  13. genforge-0.1.0/src/forge/configs/control/value_guidance.yaml +5 -0
  14. genforge-0.1.0/src/forge/configs/cost/.gitkeep +1 -0
  15. genforge-0.1.0/src/forge/configs/cost/barrier.yaml +4 -0
  16. genforge-0.1.0/src/forge/configs/cost/halfspace.yaml +4 -0
  17. genforge-0.1.0/src/forge/configs/cost/reward.yaml +4 -0
  18. genforge-0.1.0/src/forge/configs/dataset/.gitkeep +1 -0
  19. genforge-0.1.0/src/forge/configs/environment/.gitkeep +1 -0
  20. genforge-0.1.0/src/forge/configs/method/.gitkeep +1 -0
  21. genforge-0.1.0/src/forge/configs/method/conditional.yaml +3 -0
  22. genforge-0.1.0/src/forge/configs/method/d3pm.yaml +2 -0
  23. genforge-0.1.0/src/forge/configs/method/ddpm.yaml +2 -0
  24. genforge-0.1.0/src/forge/configs/method/ddpm_huber.yaml +3 -0
  25. genforge-0.1.0/src/forge/configs/method/fbsde.yaml +3 -0
  26. genforge-0.1.0/src/forge/configs/method/flow_matching.yaml +2 -0
  27. genforge-0.1.0/src/forge/configs/method/mdlm.yaml +2 -0
  28. genforge-0.1.0/src/forge/configs/method/ot_cfm.yaml +2 -0
  29. genforge-0.1.0/src/forge/configs/method/sedd.yaml +2 -0
  30. genforge-0.1.0/src/forge/configs/method/value_training.yaml +4 -0
  31. genforge-0.1.0/src/forge/configs/model/.gitkeep +1 -0
  32. genforge-0.1.0/src/forge/configs/model/categorical_mlp.yaml +5 -0
  33. genforge-0.1.0/src/forge/configs/model/mlp.yaml +6 -0
  34. genforge-0.1.0/src/forge/configs/model/temporal_unet.yaml +6 -0
  35. genforge-0.1.0/src/forge/configs/model/temporal_unet_janner.yaml +10 -0
  36. genforge-0.1.0/src/forge/configs/model/transformer.yaml +8 -0
  37. genforge-0.1.0/src/forge/configs/model/value_mlp.yaml +5 -0
  38. genforge-0.1.0/src/forge/configs/preprocessor/.gitkeep +1 -0
  39. genforge-0.1.0/src/forge/configs/preprocessor/minmax.yaml +2 -0
  40. genforge-0.1.0/src/forge/configs/preprocessor/standardize.yaml +2 -0
  41. genforge-0.1.0/src/forge/configs/runner/.gitkeep +1 -0
  42. genforge-0.1.0/src/forge/configs/runner/planning.yaml +7 -0
  43. genforge-0.1.0/src/forge/configs/runner/policy_training.yaml +9 -0
  44. genforge-0.1.0/src/forge/configs/runner/training.yaml +20 -0
  45. genforge-0.1.0/src/forge/configs/runner/value_training.yaml +7 -0
  46. genforge-0.1.0/src/forge/configs/sampler/.gitkeep +1 -0
  47. genforge-0.1.0/src/forge/configs/sampler/ddim.yaml +3 -0
  48. genforge-0.1.0/src/forge/configs/sampler/ddpm.yaml +2 -0
  49. genforge-0.1.0/src/forge/configs/sampler/flow.yaml +3 -0
  50. genforge-0.1.0/src/forge/configs/sampler/interpolant.yaml +3 -0
  51. genforge-0.1.0/src/forge/configs/sampler/sedd.yaml +2 -0
  52. genforge-0.1.0/src/forge/configs/sampler/tau_leaping.yaml +2 -0
  53. genforge-0.1.0/src/forge/configs/schedule/.gitkeep +1 -0
  54. genforge-0.1.0/src/forge/configs/schedule/absorbing.yaml +3 -0
  55. genforge-0.1.0/src/forge/configs/schedule/cfm_linear.yaml +3 -0
  56. genforge-0.1.0/src/forge/configs/schedule/linear_flow.yaml +2 -0
  57. genforge-0.1.0/src/forge/configs/schedule/si_trig.yaml +2 -0
  58. genforge-0.1.0/src/forge/configs/schedule/uniform_discrete.yaml +3 -0
  59. genforge-0.1.0/src/forge/configs/schedule/vp_cosine.yaml +4 -0
  60. genforge-0.1.0/src/forge/configs/schedule/vp_linear.yaml +4 -0
  61. genforge-0.1.0/src/forge/configs/space/.gitkeep +1 -0
  62. genforge-0.1.0/src/forge/configs/space/discrete.yaml +4 -0
  63. genforge-0.1.0/src/forge/configs/space/euclidean.yaml +3 -0
  64. genforge-0.1.0/src/forge/configs/visualizer/.gitkeep +1 -0
  65. genforge-0.1.0/src/forge/configs/visualizer/trajectory.yaml +2 -0
  66. genforge-0.1.0/src/forge/control/__init__.py +1 -0
  67. genforge-0.1.0/src/forge/control/cbf.py +35 -0
  68. genforge-0.1.0/src/forge/control/fbsde_control.py +23 -0
  69. genforge-0.1.0/src/forge/control/guidance.py +35 -0
  70. genforge-0.1.0/src/forge/control/projection.py +65 -0
  71. genforge-0.1.0/src/forge/control/value_guidance.py +79 -0
  72. genforge-0.1.0/src/forge/core/__init__.py +1 -0
  73. genforge-0.1.0/src/forge/core/builder.py +236 -0
  74. genforge-0.1.0/src/forge/core/checkpoint.py +88 -0
  75. genforge-0.1.0/src/forge/core/compose.py +44 -0
  76. genforge-0.1.0/src/forge/core/interfaces.py +429 -0
  77. genforge-0.1.0/src/forge/core/plugins.py +69 -0
  78. genforge-0.1.0/src/forge/core/protocols.py +143 -0
  79. genforge-0.1.0/src/forge/core/registry.py +105 -0
  80. genforge-0.1.0/src/forge/core/resolvers.py +23 -0
  81. genforge-0.1.0/src/forge/core/types.py +33 -0
  82. genforge-0.1.0/src/forge/costs/__init__.py +1 -0
  83. genforge-0.1.0/src/forge/costs/ball.py +44 -0
  84. genforge-0.1.0/src/forge/costs/barrier.py +40 -0
  85. genforge-0.1.0/src/forge/costs/box.py +34 -0
  86. genforge-0.1.0/src/forge/costs/halfspace.py +45 -0
  87. genforge-0.1.0/src/forge/costs/likelihood.py +32 -0
  88. genforge-0.1.0/src/forge/costs/reward.py +35 -0
  89. genforge-0.1.0/src/forge/datasets/__init__.py +7 -0
  90. genforge-0.1.0/src/forge/environments/__init__.py +7 -0
  91. genforge-0.1.0/src/forge/methods/__init__.py +1 -0
  92. genforge-0.1.0/src/forge/methods/conditional.py +54 -0
  93. genforge-0.1.0/src/forge/methods/d3pm.py +36 -0
  94. genforge-0.1.0/src/forge/methods/ddpm.py +43 -0
  95. genforge-0.1.0/src/forge/methods/ddpm_huber.py +48 -0
  96. genforge-0.1.0/src/forge/methods/fbsde.py +50 -0
  97. genforge-0.1.0/src/forge/methods/flow_matching.py +40 -0
  98. genforge-0.1.0/src/forge/methods/mdlm.py +50 -0
  99. genforge-0.1.0/src/forge/methods/ot_cfm.py +59 -0
  100. genforge-0.1.0/src/forge/methods/sedd.py +82 -0
  101. genforge-0.1.0/src/forge/methods/value_training.py +38 -0
  102. genforge-0.1.0/src/forge/models/__init__.py +1 -0
  103. genforge-0.1.0/src/forge/models/categorical.py +52 -0
  104. genforge-0.1.0/src/forge/models/mlp.py +70 -0
  105. genforge-0.1.0/src/forge/models/temporal_unet.py +104 -0
  106. genforge-0.1.0/src/forge/models/temporal_unet_janner.py +332 -0
  107. genforge-0.1.0/src/forge/models/transformer.py +147 -0
  108. genforge-0.1.0/src/forge/models/value.py +32 -0
  109. genforge-0.1.0/src/forge/preprocessing/__init__.py +1 -0
  110. genforge-0.1.0/src/forge/preprocessing/minmax.py +54 -0
  111. genforge-0.1.0/src/forge/preprocessing/standardize.py +57 -0
  112. genforge-0.1.0/src/forge/runners/__init__.py +1 -0
  113. genforge-0.1.0/src/forge/runners/multistep.py +226 -0
  114. genforge-0.1.0/src/forge/runners/planning.py +66 -0
  115. genforge-0.1.0/src/forge/runners/policy_training.py +138 -0
  116. genforge-0.1.0/src/forge/runners/training.py +385 -0
  117. genforge-0.1.0/src/forge/runners/value_training.py +18 -0
  118. genforge-0.1.0/src/forge/samplers/__init__.py +1 -0
  119. genforge-0.1.0/src/forge/samplers/ddim.py +45 -0
  120. genforge-0.1.0/src/forge/samplers/ddpm.py +56 -0
  121. genforge-0.1.0/src/forge/samplers/flow.py +55 -0
  122. genforge-0.1.0/src/forge/samplers/interpolant.py +61 -0
  123. genforge-0.1.0/src/forge/samplers/sedd.py +45 -0
  124. genforge-0.1.0/src/forge/samplers/tau_leaping.py +22 -0
  125. genforge-0.1.0/src/forge/schedules/__init__.py +1 -0
  126. genforge-0.1.0/src/forge/schedules/discrete.py +146 -0
  127. genforge-0.1.0/src/forge/schedules/flow.py +125 -0
  128. genforge-0.1.0/src/forge/schedules/vp.py +110 -0
  129. genforge-0.1.0/src/forge/spaces/__init__.py +1 -0
  130. genforge-0.1.0/src/forge/spaces/discrete.py +55 -0
  131. genforge-0.1.0/src/forge/spaces/euclidean.py +44 -0
  132. genforge-0.1.0/src/forge/utils/__init__.py +1 -0
  133. genforge-0.1.0/src/forge/utils/ema.py +82 -0
  134. genforge-0.1.0/src/forge/utils/logging.py +118 -0
  135. genforge-0.1.0/src/forge/utils/lora.py +126 -0
  136. genforge-0.1.0/src/forge/utils/seeding.py +34 -0
  137. genforge-0.1.0/src/forge/utils/torch_utils.py +23 -0
  138. genforge-0.1.0/src/forge/visualizations/__init__.py +1 -0
  139. 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.
@@ -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).
@@ -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,3 @@
1
+ """genforge — a unified, PyTorch-based framework for generative-modeling techniques."""
2
+
3
+ __version__ = "0.1.0"
@@ -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,3 @@
1
+ name: cbf
2
+ params:
3
+ alpha: 1.0
@@ -0,0 +1,4 @@
1
+ name: fbsde_control
2
+ params:
3
+ value_checkpoint: checkpoints/lq/fbsde/values.pt
4
+ scale: 1.0
@@ -0,0 +1,4 @@
1
+ name: guidance
2
+ params:
3
+ scale: 2.0
4
+ sigma_weight: true
@@ -0,0 +1,2 @@
1
+ name: projection
2
+ params: {}
@@ -0,0 +1,5 @@
1
+ name: value_guidance
2
+ params:
3
+ value_checkpoint: checkpoints/distributions/value/values.pt
4
+ scale: 3.0
5
+ sigma_weight: true
@@ -0,0 +1 @@
1
+ # Phase 0: empty config group for `cost`. Concrete options arrive in later phases.
@@ -0,0 +1,4 @@
1
+ name: barrier
2
+ params:
3
+ normal: [1.0, 0.0]
4
+ offset: 0.0
@@ -0,0 +1,4 @@
1
+ name: halfspace
2
+ params:
3
+ normal: [1.0, 0.0]
4
+ offset: 0.0
@@ -0,0 +1,4 @@
1
+ name: reward
2
+ params:
3
+ target: [2.0, 0.0]
4
+ weight: 1.0
@@ -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,3 @@
1
+ name: conditional
2
+ params:
3
+ pin_positions: [0, -1]
@@ -0,0 +1,2 @@
1
+ name: d3pm
2
+ params: {}
@@ -0,0 +1,2 @@
1
+ name: ddpm
2
+ params: {}
@@ -0,0 +1,3 @@
1
+ name: ddpm_huber
2
+ params:
3
+ delta: 1.0 # residual threshold; |r| > delta is penalized linearly instead of quadratically
@@ -0,0 +1,3 @@
1
+ name: fbsde
2
+ params:
3
+ q: 4.0
@@ -0,0 +1,2 @@
1
+ name: flow_matching
2
+ params: {}
@@ -0,0 +1,2 @@
1
+ name: mdlm
2
+ params: {}
@@ -0,0 +1,2 @@
1
+ name: ot_cfm
2
+ params: {}
@@ -0,0 +1,2 @@
1
+ name: sedd
2
+ params: {}
@@ -0,0 +1,4 @@
1
+ name: value_training
2
+ params:
3
+ target: [2.0, 0.0]
4
+ weight: 1.0
@@ -0,0 +1 @@
1
+ # Phase 0: empty config group for `model`. Concrete options arrive in later phases.
@@ -0,0 +1,5 @@
1
+ name: categorical_mlp
2
+ params:
3
+ num_classes: 5
4
+ hidden: 128
5
+ depth: 3
@@ -0,0 +1,6 @@
1
+ name: mlp
2
+ params:
3
+ dim: 2
4
+ hidden: 128
5
+ depth: 3
6
+ output_type: eps
@@ -0,0 +1,6 @@
1
+ name: temporal_unet
2
+ params:
3
+ dim: 2
4
+ horizon: 32
5
+ base: 32
6
+ output_type: eps
@@ -0,0 +1,10 @@
1
+ name: temporal_unet_janner
2
+ params:
3
+ dim: 32
4
+ dim_mults: [1, 2, 4, 8]
5
+ horizon: 32
6
+ transition_dim: 14
7
+ cond_dim: 0
8
+ cond_predict_scale: false
9
+ output_type: x0
10
+ obs_dim: 0