genforge 0.1.0__py3-none-any.whl

forge/__init__.py

@@ -0,0 +1,3 @@
+"""genforge — a unified, PyTorch-based framework for generative-modeling techniques."""
+
+__version__ = "0.1.0"

forge/cli.py

@@ -0,0 +1,143 @@
+"""The ``forge`` command-line entrypoint: ``list`` / ``train`` / ``sample``.
+
+``list`` imports the built-ins so registrations fire, then prints the registered components by
+category. ``train`` / ``sample`` compose a Hydra config from an ``experiment=`` selection, build the
+runner, and run. ``sample checkpoint=<path.pt>`` rebuilds everything from the self-contained
+checkpoint alone (Invariant 5). All three fail loudly on misconfiguration.
+"""
+
+from __future__ import annotations
+
+import argparse
+import sys
+from typing import Optional, Sequence
+
+from .core import registry
+from .core.builder import build, import_builtin_components
+
+_EXPERIMENT_HINT = (
+    "Select one with `experiment=<env>/<params>/<method>` "
+    "(e.g. `forge train experiment=distributions/ddpm/base`)."
+)
+
+
+def _split_overrides(overrides: Sequence[str]) -> dict:
+    """Parse ``key=value`` overrides into a flat dict (first '=')."""
+    flat: dict[str, str] = {}
+    for o in overrides:
+        if "=" in o:
+            k, v = o.split("=", 1)
+            flat[k] = v
+    return flat
+
+
+def _cmd_list(_args: argparse.Namespace) -> int:
+    import_builtin_components()
+    # Concrete envs are plugins (no experiment selected here), so import the bundled env packages
+    # too — otherwise `list` would omit environments/datasets/env-preprocessors.
+    from .core.plugins import load_bundled_envs
+
+    load_bundled_envs()
+    reg = registry.registered()
+    print("forge components")
+    print("===================")
+    for category in registry.CATEGORIES:
+        comps = reg.get(category, {})
+        names = ", ".join(comps) if comps else "(none yet)"
+        print(f"  {category:<13} {names}")
+    for category in [c for c in reg if c not in registry.CATEGORIES]:
+        print(f"  {category:<13} {', '.join(reg[category])}")
+    return 0
+
+
+def _run_from_config(overrides: Sequence[str], action: str) -> int:
+    from omegaconf import OmegaConf
+
+    from .core.compose import compose_config
+
+    cfg = compose_config(overrides)
+    runner = build(cfg)
+    runner.resolved_config = OmegaConf.to_container(cfg, resolve=True)
+
+    if action == "train":
+        runner.train()
+        metrics = runner.evaluate()
+        print(f"[train] done. eval: {metrics}")
+        return 0
+
+    # sample from an experiment: load its configured checkpoint if present.
+    ckpt_path = getattr(runner, "ckpt_path", None)
+    if ckpt_path:
+        from pathlib import Path
+
+        from .core.checkpoint import load_checkpoint
+
+        if not Path(ckpt_path).exists():
+            print(
+                f"`forge sample` found no checkpoint at {ckpt_path!r}. Train first "
+                f"(`forge train experiment=...`) or pass `checkpoint=<path.pt>`.",
+                file=sys.stderr,
+            )
+            return 1
+        runner.load_state(load_checkpoint(ckpt_path))
+    metrics = runner.evaluate()
+    print(f"[sample] {metrics}")
+    return 0
+
+
+def _cmd_train(args: argparse.Namespace) -> int:
+    flat = _split_overrides(args.overrides)
+    if "experiment" not in flat:
+        print(f"`forge train` requires an experiment selection. {_EXPERIMENT_HINT}", file=sys.stderr)
+        return 2
+    return _run_from_config(args.overrides, "train")
+
+
+def _cmd_sample(args: argparse.Namespace) -> int:
+    flat = _split_overrides(args.overrides)
+    if "checkpoint" in flat:
+        # Self-contained path: rebuild from the .pt alone (Invariant 5).
+        from .runners.training import TrainingRunner
+
+        runner = TrainingRunner.from_checkpoint(flat["checkpoint"], build_fn=build)
+        metrics = runner.evaluate()
+        print(f"[sample] from checkpoint {flat['checkpoint']}: {metrics}")
+        return 0
+    if "experiment" not in flat:
+        print(
+            f"`forge sample` requires `experiment=...` or `checkpoint=<path.pt>`. {_EXPERIMENT_HINT}",
+            file=sys.stderr,
+        )
+        return 2
+    return _run_from_config(args.overrides, "sample")
+
+
+def _build_parser() -> argparse.ArgumentParser:
+    parser = argparse.ArgumentParser(
+        prog="forge",
+        description="A unified framework for generative modeling with a clean control layer.",
+    )
+    sub = parser.add_subparsers(dest="command", required=True)
+
+    p_list = sub.add_parser("list", help="List registered components by category.")
+    p_list.set_defaults(func=_cmd_list)
+
+    p_train = sub.add_parser("train", help="Train a model from an experiment config.")
+    p_train.add_argument("overrides", nargs="*", help="Hydra-style overrides, e.g. experiment=...")
+    p_train.set_defaults(func=_cmd_train)
+
+    p_sample = sub.add_parser("sample", help="Sample from a trained model or checkpoint.")
+    p_sample.add_argument("overrides", nargs="*", help="experiment=... or checkpoint=<path.pt>")
+    p_sample.set_defaults(func=_cmd_sample)
+
+    return parser
+
+
+def main(argv: Optional[Sequence[str]] = None) -> int:
+    parser = _build_parser()
+    args = parser.parse_args(argv)
+    return args.func(args)
+
+
+if __name__ == "__main__":
+    raise SystemExit(main())

forge/configs/config.yaml

@@ -0,0 +1,22 @@
+# genforge root config (Hydra base+delta).
+#
+# Per-category config groups live in the sibling directories (space/, schedule/, model/, ...).
+# Experiments are base+delta bundles in the repo-root `experiment/` tree, which is added to the
+# Hydra searchpath at runtime by the CLI (env var GENFORGE_EXP_ROOT). Select an experiment with
+# `experiment=<env>/<params>/<method>` — it composes the component groups it needs.
+
+defaults:
+  - _self_
+  - experiment: ???        # mandatory: train/sample require an experiment selection
+
+# Global run settings (resolved into the checkpoint, Invariant 5).
+seed: 0
+
+hydra:
+  searchpath:
+    - file://${oc.env:GENFORGE_EXP_ROOT}
+  job:
+    chdir: false
+  output_subdir: null
+  run:
+    dir: .

forge/configs/control/.gitkeep

@@ -0,0 +1 @@
+# Phase 0: empty config group for `control`. Concrete options arrive in later phases.

forge/configs/control/cbf.yaml

@@ -0,0 +1,3 @@
+name: cbf
+params:
+  alpha: 1.0

forge/configs/control/fbsde_control.yaml

@@ -0,0 +1,4 @@
+name: fbsde_control
+params:
+  value_checkpoint: checkpoints/lq/fbsde/values.pt
+  scale: 1.0

forge/configs/control/guidance.yaml

@@ -0,0 +1,4 @@
+name: guidance
+params:
+  scale: 2.0
+  sigma_weight: true

forge/configs/control/projection.yaml

@@ -0,0 +1,2 @@
+name: projection
+params: {}

forge/configs/control/value_guidance.yaml

@@ -0,0 +1,5 @@
+name: value_guidance
+params:
+  value_checkpoint: checkpoints/distributions/value/values.pt
+  scale: 3.0
+  sigma_weight: true

forge/configs/cost/.gitkeep

@@ -0,0 +1 @@
+# Phase 0: empty config group for `cost`. Concrete options arrive in later phases.

forge/configs/cost/barrier.yaml

@@ -0,0 +1,4 @@
+name: barrier
+params:
+  normal: [1.0, 0.0]
+  offset: 0.0

forge/configs/cost/halfspace.yaml

@@ -0,0 +1,4 @@
+name: halfspace
+params:
+  normal: [1.0, 0.0]
+  offset: 0.0

forge/configs/cost/reward.yaml

@@ -0,0 +1,4 @@
+name: reward
+params:
+  target: [2.0, 0.0]
+  weight: 1.0

forge/configs/dataset/.gitkeep

@@ -0,0 +1 @@
+# Phase 0: empty config group for `dataset`. Concrete options arrive in later phases.

forge/configs/environment/.gitkeep

@@ -0,0 +1 @@
+# Phase 0: empty config group for `environment`. Concrete options arrive in later phases.

forge/configs/method/.gitkeep

@@ -0,0 +1 @@
+# Phase 0: empty config group for `method`. Concrete options arrive in later phases.

forge/configs/method/conditional.yaml

@@ -0,0 +1,3 @@
+name: conditional
+params:
+  pin_positions: [0, -1]

forge/configs/method/d3pm.yaml

@@ -0,0 +1,2 @@
+name: d3pm
+params: {}

forge/configs/method/ddpm.yaml

@@ -0,0 +1,2 @@
+name: ddpm
+params: {}

forge/configs/method/ddpm_huber.yaml

@@ -0,0 +1,3 @@
+name: ddpm_huber
+params:
+  delta: 1.0   # residual threshold; |r| > delta is penalized linearly instead of quadratically

forge/configs/method/fbsde.yaml

@@ -0,0 +1,3 @@
+name: fbsde
+params:
+  q: 4.0

forge/configs/method/flow_matching.yaml

@@ -0,0 +1,2 @@
+name: flow_matching
+params: {}

forge/configs/method/mdlm.yaml

@@ -0,0 +1,2 @@
+name: mdlm
+params: {}

forge/configs/method/ot_cfm.yaml

@@ -0,0 +1,2 @@
+name: ot_cfm
+params: {}

forge/configs/method/sedd.yaml

@@ -0,0 +1,2 @@
+name: sedd
+params: {}

forge/configs/method/value_training.yaml

@@ -0,0 +1,4 @@
+name: value_training
+params:
+  target: [2.0, 0.0]
+  weight: 1.0

forge/configs/model/.gitkeep

@@ -0,0 +1 @@
+# Phase 0: empty config group for `model`. Concrete options arrive in later phases.

forge/configs/model/categorical_mlp.yaml

@@ -0,0 +1,5 @@
+name: categorical_mlp
+params:
+  num_classes: 5
+  hidden: 128
+  depth: 3

forge/configs/model/mlp.yaml

@@ -0,0 +1,6 @@
+name: mlp
+params:
+  dim: 2
+  hidden: 128
+  depth: 3
+  output_type: eps

forge/configs/model/temporal_unet.yaml

@@ -0,0 +1,6 @@
+name: temporal_unet
+params:
+  dim: 2
+  horizon: 32
+  base: 32
+  output_type: eps

forge/configs/model/temporal_unet_janner.yaml

@@ -0,0 +1,10 @@
+name: temporal_unet_janner
+params:
+  dim: 32
+  dim_mults: [1, 2, 4, 8]
+  horizon: 32
+  transition_dim: 14
+  cond_dim: 0
+  cond_predict_scale: false
+  output_type: x0
+  obs_dim: 0

forge/configs/model/transformer.yaml

@@ -0,0 +1,8 @@
+name: transformer
+params:
+  vocab_size: 32        # = char_text unique chars (31) + 1 mask token
+  length: 32
+  d_model: 128
+  depth: 4
+  n_heads: 4
+  output_type: logits

forge/configs/model/value_mlp.yaml

@@ -0,0 +1,5 @@
+name: value_mlp
+params:
+  dim: 2
+  hidden: 128
+  depth: 3

forge/configs/preprocessor/.gitkeep

@@ -0,0 +1 @@
+# Phase 0: empty config group for `preprocessor`. Concrete options arrive in later phases.

forge/configs/preprocessor/minmax.yaml

@@ -0,0 +1,2 @@
+name: minmax
+params: {}

forge/configs/preprocessor/standardize.yaml

@@ -0,0 +1,2 @@
+name: standardize
+params: {}

forge/configs/runner/.gitkeep

@@ -0,0 +1 @@
+# Phase 0: empty config group for `runner`. Concrete options arrive in later phases.

forge/configs/runner/planning.yaml

@@ -0,0 +1,7 @@
+name: planning
+params:
+  steps: 2000
+  batch_size: 128
+  lr: 1.0e-3
+  device: cpu
+  seed: 0

forge/configs/runner/policy_training.yaml

@@ -0,0 +1,9 @@
+name: policy_training
+params:
+  steps: 100000
+  batch_size: 256
+  lr: 1.0e-4
+  ema_decay: 0.9999
+  n_sample_steps: 100
+  device: cuda
+  seed: 0

forge/configs/runner/training.yaml

@@ -0,0 +1,20 @@
+name: training
+params:
+  steps: 2000
+  batch_size: 256
+  lr: 1.0e-3
+  ema_decay: 0.999
+  n_sample_steps: 100
+  n_eval_samples: 2000
+  eval_radius: 0.6
+  device: cpu
+  seed: 0
+  # Optional experiment logging — needs the `logging` extra (pip install genforge[logging]).
+  # Absent/false => no-op logger + plain loop; a default run is unchanged. Any runner accepts
+  # these via CLI too, e.g. runner.params.log.wandb=true (or env FORGE_WANDB=1).
+  log:
+    wandb: false      # off by default; no wandb import, no login prompt
+    project: forge
+    mode: null        # online | offline | disabled
+    name: null        # run name; defaults to <method>-<steps>steps
+    progress: false   # tqdm bar; plain loop if tqdm absent or stderr is non-TTY

forge/configs/runner/value_training.yaml

@@ -0,0 +1,7 @@
+name: value_training
+params:
+  steps: 2000
+  batch_size: 512
+  lr: 1.0e-3
+  device: cpu
+  seed: 0

forge/configs/sampler/.gitkeep

@@ -0,0 +1 @@
+# Phase 0: empty config group for `sampler`. Concrete options arrive in later phases.

forge/configs/sampler/ddim.yaml

@@ -0,0 +1,3 @@
+name: ddim
+params:
+  eta: 0.0

forge/configs/sampler/ddpm.yaml

@@ -0,0 +1,2 @@
+name: ddpm
+params: {}

forge/configs/sampler/flow.yaml

@@ -0,0 +1,3 @@
+name: flow
+params:
+  integrator: heun

forge/configs/sampler/interpolant.yaml

@@ -0,0 +1,3 @@
+name: interpolant
+params:
+  epsilon: 1.0   # free diffusion-coefficient scale; ε(t)=epsilon·σ(t)². 0 ⇒ probability-flow ODE.

forge/configs/sampler/sedd.yaml

@@ -0,0 +1,2 @@
+name: sedd
+params: {}

forge/configs/sampler/tau_leaping.yaml

@@ -0,0 +1,2 @@
+name: tau_leaping
+params: {}

forge/configs/schedule/.gitkeep

@@ -0,0 +1 @@
+# Phase 0: empty config group for `schedule`. Concrete options arrive in later phases.

forge/configs/schedule/absorbing.yaml

@@ -0,0 +1,3 @@
+name: absorbing
+params:
+  num_classes: 5

forge/configs/schedule/cfm_linear.yaml

@@ -0,0 +1,3 @@
+name: cfm_linear
+params:
+  sigma_min: 0.01

forge/configs/schedule/linear_flow.yaml

@@ -0,0 +1,2 @@
+name: linear_flow
+params: {}

forge/configs/schedule/si_trig.yaml

@@ -0,0 +1,2 @@
+name: si_trig
+params: {}

forge/configs/schedule/uniform_discrete.yaml

@@ -0,0 +1,3 @@
+name: uniform_discrete
+params:
+  num_classes: 4

forge/configs/schedule/vp_cosine.yaml

@@ -0,0 +1,4 @@
+name: vp_cosine
+params:
+  s: 0.008
+  parameterization: sqrt_alpha_bar

forge/configs/schedule/vp_linear.yaml

@@ -0,0 +1,4 @@
+name: vp_linear
+params:
+  beta_min: 0.1
+  beta_max: 20.0

forge/configs/space/.gitkeep

@@ -0,0 +1 @@
+# Phase 0: empty config group for `space`. Concrete options arrive in later phases.

forge/configs/space/discrete.yaml

@@ -0,0 +1,4 @@
+name: discrete
+params:
+  num_classes: 5
+  length: 1

forge/configs/space/euclidean.yaml

@@ -0,0 +1,3 @@
+name: euclidean
+params:
+  dim: 2

forge/configs/visualizer/.gitkeep

@@ -0,0 +1 @@
+# Phase 0: empty config group for `visualizer`. Concrete options arrive in later phases.

forge/configs/visualizer/trajectory.yaml

@@ -0,0 +1,2 @@
+name: trajectory
+params: {}

forge/control/__init__.py

@@ -0,0 +1 @@
+"""The control layer: HOW you approximate the tilt Q ∝ exp(log_h)·P (the thesis surface)."""

forge/control/cbf.py

@@ -0,0 +1,35 @@
+"""CBF — a control-barrier-function safety filter on the reverse DRIFT (the rate ẋ).
+
+The first real `drift`-surface controller (Invariant 6), and the proof the surface works: a CBF
+*reads the rate*, so it has no faithful x̂₀ reduction. Given a barrier ``h(x)≥0``, it minimally
+edits the drift so the barrier's forward-invariance condition ``ḣ = ∇h·ẋ ≥ −α·h(x)`` holds. For one
+linear barrier this is the closed-form CBF-QP solution (project the drift onto the safe halfspace):
+
+    slack = ∇h·drift + α·h ;  if slack < 0:  drift += (−slack/‖∇h‖²)·∇h   (else unchanged).
+
+With the data-ward heading drift and α=1 this keeps the clean estimate feasible every step
+(``h(x̂₀') ≥ (1−α)·h(xₜ) = 0``), so the formed sample lands in the feasible set.
+"""
+
+from __future__ import annotations
+
+import torch
+
+from ..core.interfaces import Controller
+from ..core.registry import register
+
+
+@register("control", "cbf")
+class CBF(Controller):
+    surface = "drift"
+
+    def __init__(self, cost, alpha: float = 1.0):
+        super().__init__(cost)
+        self.alpha = float(alpha)
+
+    def modify_drift(self, drift: torch.Tensor, x: torch.Tensor, t, schedule) -> torch.Tensor:
+        grad = self.cost.grad_h(x)                              # ∇h(x), (..., dim)
+        h = self.cost.value(x)                                 # h(x),   (...,)
+        slack = (grad * drift).sum(dim=-1) + self.alpha * h    # ḣ + α·h ≥ 0 required
+        coef = torch.clamp(-slack, min=0.0) / (grad * grad).sum(dim=-1).clamp_min(1e-12)
+        return drift + coef.unsqueeze(-1) * grad               # minimal safe correction

forge/control/fbsde_control.py

@@ -0,0 +1,23 @@
+"""FBSDEControl — consume an FBSDE-learned value (cost-to-go) to steer sampling.
+
+Same amortized machinery as `ValueGuidance` (loads the value from a checkpoint, never imports the
+method), but descends the cost-to-go: the optimal control follows ``−∇V`` (``sign = −1``).
+
+FBSDE control belongs on the ``drift`` surface (``u* = Gᵀ∇V`` added to b_θ). This implementation
+keeps it on the inherited ``x0`` surface (a −∇V shift of x̂₀) with identical behavior; drift-level
+FBSDE is future work (it depends on the dual-surface refactor landing first).
+"""
+
+from __future__ import annotations
+
+from ..core.registry import register
+from .value_guidance import AmortizedValueController
+
+
+@register("control", "fbsde_control")
+class FBSDEControl(AmortizedValueController):
+    surface = "x0"  # TODO: reclassify to the "drift" surface.
+
+    def __init__(self, value_checkpoint: str, cost=None, scale: float = 1.0, sigma_weight: bool = True):
+        # Descend the cost-to-go: optimal control is −∇V.
+        super().__init__(value_checkpoint, cost=cost, scale=scale, sigma_weight=sigma_weight, sign=-1.0)

forge/control/guidance.py

@@ -0,0 +1,35 @@
+"""Guidance controller — first-order ∇log h on the clean-sample estimate (DPS/RePaint style).
+
+Each reverse step, nudge x̂₀ along ∇_{x̂₀} log h (computed by autograd), never the noisy iterate.
+The step is scaled by ``scale`` and, optionally, by σ(t)² so the correction is gentler near the data
+(t→0) — a soft tilt rather than a hard projection. The base model is never touched (Invariant 6).
+"""
+
+from __future__ import annotations
+
+import torch
+
+from ..core.interfaces import Controller
+from ..core.registry import register
+
+
+@register("control", "guidance")
+class Guidance(Controller):
+    surface = "x0"
+
+    def __init__(self, cost, scale: float = 1.0, sigma_weight: bool = True):
+        super().__init__(cost)
+        self.scale = float(scale)
+        self.sigma_weight = bool(sigma_weight)
+
+    def modify_x0(self, x0_hat: torch.Tensor, x: torch.Tensor, t, schedule) -> torch.Tensor:
+        with torch.enable_grad():
+            z = x0_hat.detach().requires_grad_(True)
+            lh = self.cost.log_h(z, t).sum()
+            (grad,) = torch.autograd.grad(lh, z)
+        step = self.scale
+        if self.sigma_weight:
+            # σ(t)² fades the correction as the estimate sharpens toward the data manifold.
+            sigma = torch.as_tensor(schedule.sigma(t), device=x0_hat.device)
+            step = step * (sigma**2)
+        return (x0_hat + step * grad).detach()