Skill: reedcgx/deep-hjb-solver-skill (version 1.0.2)

--- name: deep-hjb-solver description: "Create or refactor code for solving HJB equations with this repository's TensorFlow DGM framework. Use when users ask to generate new HJB training code, add a new problem (config/problem/loss/train script), adapt sampling/training hyperparameters, or create plotting/analysis code from trainer CSV outputs." --- # Deep HJB Solver **Before writing any loss or config code**, read `references/repo-conventions.md` — it contains the exact constructor signatures and return-type contracts that must be followed. **Before writing any CSV or plotting code**, read `references/training-output-contract.md`. Never modify `DGMTrainer`, `DGMNet`, or any file under `src/trainers/` or `src/models/` unless explicitly asked. --- ## Output layout Every new HJB problem is created as a **self-contained folder** named after the problem slug. Nothing is written outside it. ``` <slug>/ ├── src/ │ ├── __init__.py │ ├── configs/ │ │ ├── __init__.py │ │ ├── common_config.py ← copied from assets │ │ └── <slug>_config.py ← generated │ ├── models/ │ │ ├── __init__.py │ │ └── dgm_net.py ← copied from assets │ ├── problems/ │ │ ├── __init__.py │ │ ├── base_problem.py ← copied from assets │ │ └── <slug>_problem.py ← generated │ ├── losses/ │ │ ├── __init__.py │ │ └── <slug>_loss.py ← generated │ ├── samplers/ │ │ ├── __init__.py │ │ ├── base_sampler.py ← copied from assets │ │ ├── uniform_sampler.py ← copied from assets │ │ └── uniform_sampler_2d.py ← copied from assets │ ├── trainers/ │ │ ├── __init__.py │ │ └── dgm_trainer.py ← copied from assets │ └── utils/ │ ├── __init__.py │ └── visualization.py ← copied from assets ├── examples/ │ └── <slug>_train.py ← generated ├── plot_training_csv.py ← copied from assets └── requirements.txt ← copied from assets ``` --- ## Workflow Use the following steps in order. In every template below, replace: - `<slug>` → snake_case problem name, e.g. `linear_control` - `<Prefix>` → CamelCase version, e.g. `LinearControl` - `<dim>` → spatial dimension (`1` or `2`) - `<N>` → number of control variables - `<controls>` → Python list literal of control names, e.g. `['Z', 'W']` ### Step 1 — Copy the DGM framework into `<slug>/src/` **This step is mandatory and must be executed immediately without asking the user for permission or confirmation.** Do not say "should I copy the assets?" — just do it. Run the following shell commands to copy the bundled framework. Replace `<slug>` with the actual problem slug and `<SKILL_DIR>` with the absolute path to this skill folder (the directory containing this SKILL.md): ```bash mkdir -p <slug>/src cp -r <SKILL_DIR>/assets/src/. <slug>/src/ cp <SKILL_DIR>/assets/plot_training_csv.py <slug>/plot_training_csv.py cp <SKILL_DIR>/assets/requirements.txt <slug>/requirements.txt ``` **Do not proceed to Step 2 until the copy commands have completed successfully.** ### Step 2 — Create `<slug>/src/configs/<slug>_config.py` **1D domain** (`<dim>` = 1): ```python """Configuration for <Prefix>.""" from dataclasses import dataclass, field from .common_config import CommonConfig @dataclass class <Prefix>Config(CommonConfig): dimension: int = 1 T: float = 1.0 t_low: float = 0.0 X_low: float = 0.0 X_high: float = 1.0 num_controls: int = <N> control_names: list = field(default_factory=lambda: <controls>) metrics_config: list = field(default_factory=lambda: ['maxdiff_V', 'maxdiff_terminal']) extra_info_mapping: dict = field(default_factory=dict) early_stop_metric: str = 'maxdiff_V' early_stop_threshold: float = 1e-4 problem_params_keys: list = field(default_factory=list) saveName: str = '<slug>' ``` **2D domain** (`<dim>` = 2) — replace the bounds lines with: ```python dimension: int = 2 X_low: list = field(default_factory=lambda: [0.0, 0.0]) X_high: list = field(default_factory=lambda: [1.0, 1.0]) ``` ### Step 3 — Create `<slug>/src/problems/<slug>_problem.py` ```python """Problem definition for <Prefix>.""" from .base_problem import BaseProblem def terminal_utility_<slug>(x): """TODO: implement terminal payoff g(x). x shape: (batch, dim).""" return -x[:, :1] class <Prefix>Problem(BaseProblem): def get_terminal_condition(self, x): return terminal_utility_<slug>(x) ``` ### Step 4 — Create `<slug>/src/losses/<slug>_loss.py` ```python """Loss functions for <Prefix>.""" import tensorflow as tf class <Prefix>Loss: def __init__(self, problem): self.problem = problem def compute_value_loss(self, model, control, t_interior, X_interior, t_terminal, X_terminal): # TODO: replace with real HJB PDE residual. # IMPORTANT: if you need more than one gradient from the same tape (e.g. # both V_t and V_x), you MUST use persistent=True and delete the tape # afterward. A non-persistent tape raises RuntimeError on the second call. with tf.GradientTape(persistent=True, watch_accessed_variables=False) as gt: gt.watch(t_interior) gt.watch(X_interior) V = model(t_interior, X_interior) V_t = gt.gradient(V, t_interior) # ∂V/∂t V_x = gt.gradient(V, X_interior) # ∂V/∂x (use if needed by the HJB) del gt # release persistent tape immediately after use ctrl = control(t_interior, X_interior) # u from control network — substitute into HJB residual = V_t # TODO: replace with actual HJB residual, e.g. V_t + ctrl * V_x + ... L1 = tf.reduce_mean(tf.square(residual)) target_terminal = self.problem.get_terminal_condition(X_terminal) fitted_terminal = model(t_terminal, X_terminal) diff_terminal = fitted_terminal - target_terminal L3 = tf.reduce_mean(tf.square(diff_terminal)) # diff_V can be either: # (a) a plain tensor — the HJB residual (e.g. diff_V = residual), OR # (b) a dict of debug tensors that MUST contain a 'residual' key # (e.g. {'residual': residual, 'V': V, 'V_t': V_t, 'V_x': V_x}) # base_problem.extract_metrics handles both forms automatically. # MUST return exactly this 4-tuple — DGMTrainer unpacks it directly. diff_V = residual # option (a): simplest form # diff_V = {'residual': residual, 'V': V, 'V_t': V_t, 'V_x': V_x} # option (b) return L1, L3, diff_V, diff_terminal def compute_control_loss(self, model, control, t_interior, X_interior, t_terminal, X_terminal): # TODO: implement FOC / control objective. # Always use persistent=True here: you need at least V_x, and some problems # also need V_xx (second-order), which requires a nested tape inside this one. # persistent=True lets you reuse the outer tape multiple times safely. with tf.GradientTape(persistent=True, watch_accessed_variables=False) as gt: gt.watch(X_interior) V = model(t_interior, X_interior) V_x = gt.gradient(V, X_interior) # ∂V/∂x del gt ctrl = control(t_interior, X_interior) # u from control network # TODO: compute L2 from the inf{...} terms of the HJB using ctrl and V_x. # See "How to translate an HJB equation into loss functions" below for the rule. # Example (LQ, inf_u { u*V_x + ½u² }): L2 = tf.reduce_mean(ctrl * V_x + 0.5 * tf.square(ctrl)) L2 = tf.reduce_mean(ctrl) # placeholder — replace with real Hamiltonian terms # MUST return exactly this 2-tuple — DGMTrainer unpacks it directly return L2, {'V_x': V_x, 'ctrl': ctrl} ``` ### Step 5 — Register new classes in `<slug>/src/configs/__init__.py`, `<slug>/src/problems/__init__.py`, `<slug>/src/losses/__init__.py` Append one import line to each file: ```python # src/configs/__init__.py from .<slug>_config import <Prefix>Config # src/problems/__init__.py from .<slug>_problem import <Prefix>Problem # src/losses/__init__.py from .<slug>_loss import <Prefix>Loss ``` ### Step 6 — Create `<slug>/examples/<slug>_train.py` All imports are relative to `<slug>/` (the script is run from inside that folder). ```python """Training script for <slug>.""" import os import sys import tensorflow as tf # Ensure the problem folder is on the path when run directly sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__)))) from src.configs.<slug>_config import <Prefix>Config from src.problems.<slug>_problem import <Prefix>Problem from src.losses.<slug>_loss import <Prefix>Loss from src.models import DGMNet from src.samplers import UniformSampler, UniformSampler2D from src.trainers import DGMTrainer from plot_training_csv import plot_training_history def main(): config = <Prefix>Config() config.save_dir = './results' os.makedirs(config.save_dir, exist_ok=True) problem = <Prefix>Problem(config) # UniformSampler: scalar 1D bounds only. # UniformSampler2D: supports any dimension >= 2 (array bounds). sampler = UniformSampler(config, problem) if config.dimension == 1 else UniformSampler2D(config, problem) model = DGMNet( layer_width=config.nodes_per_layer, n_layers=config.num_layers, input_dim=config.dimension, output_dim=1, control_output=False, activation=config.activation, ) control = DGMNet( layer_width=config.nodes_per_layer, n_layers=config.num_layers, input_dim=config.dimension, output_dim=config.num_controls, control_output=True, activation=config.activation, ) loss_fn = <Prefix>Loss(problem) lr_schedule = config.get_optimizer_config() trainer = DGMTrainer( model=model, control=control, loss_fn=loss_fn, sampler=sampler, optimizer_value=tf.keras.optimizers.Adam(lr_schedule), optimizer_control=tf.keras.optimizers.Adam(lr_schedule), config=config, problem=problem, ) history = trainer.train() if config.saveOutput: trainer.save_models(config.save_dir) csv_path = f"{config.save_dir}/{config.saveName}_training_history.csv" if os.path.exists(csv_path): plot_training_history(csv_path, config.save_dir) print('final maxdiff_V:', history['maxdiff_V'][-1]) if __name__ == '__main__': main() ``` ### Step 7 — Fill the real PDE In `<slug>/src/losses/<slug>_loss.py`, replace the placeholder body of `compute_value_loss` with the actual HJB residual. The return signature **must not change**: `return L1, L3, diff_V, diff_terminal`. In `<slug>/src/problems/<slug>_problem.py`, implement `terminal_utility_<slug>` with the correct payoff. See **"How to translate an HJB equation into loss functions"** below for the exact decomposition rule. --- ## How to translate an HJB equation into loss functions Given an HJB equation with an inf (or sup) operator, split it into two losses following this fixed rule. ### The rule | Loss | What to put in it | |------|-------------------| | `compute_value_loss` → `L1` | Drop the inf/sup symbol; substitute the **control network** output for u; sum **all** remaining terms into a residual; squared mean. | | `compute_control_loss` → `L2` | Keep **only the terms inside** inf{…}; evaluate them with the **control network**; take the mean directly (no square). For sup{…}, negate first. | The intuition: `L1` trains the value network so the PDE residual → 0 (equation is satisfied). `L2` trains the control network to actually minimise (or maximise) the Hamiltonian — it IS a gradient-descent step on the inf objective, so no squaring. --- ## Running the training ```bash cd <slug> python examples/<slug>_train.py ``` Results are saved to `<slug>/results/`. --- ## Common pitfall: GradientTape reuse ### Rule **Always use `persistent=True` in both `compute_value_loss` and `compute_control_loss`.** Even if you think you only need one gradient today, problems that include `V_xx` (second-order / diffusion terms) require a nested tape inside `compute_control_loss`, which only works when the outer tape is persistent. Using `persistent=True` consistently prevents hard-to-diagnose errors. | Method | Minimum gradients | Must use `persistent=True`? | |---|---|---| | `compute_value_loss` | `V_t` + `V_x` (2 gradients) | Yes | | `compute_control_loss` | `V_x` + possibly `V_xx` via nested tape | Yes | ### `compute_value_loss` — always `persistent=True` ```python with tf.GradientTape(persistent=True, watch_accessed_variables=False) as gt: gt.watch(t_interior) gt.watch(X_interior) V = model(t_interior, X_interior) V_t = gt.gradient(V, t_interior) V_x = gt.gradient(V, X_interior) del gt # required — frees memory held by persistent tape ``` ### `compute_control_loss` — always `persistent=True` ```python with tf.GradientTape(persistent=True, watch_accessed_variables=False) as gt: gt.watch(X_interior) V = model(t_interior, X_interior) V_x = gt.gradient(V, X_interior) del gt # always delete, even if only one gradient was taken ``` Omitting `persistent=True` when computing two or more gradients raises: ``` RuntimeError: A non-persistent GradientTape can only be used to compute one set of gradients ``` --- ## Guardrails - Never modify `DGMTrainer`, `DGMNet`, or sampler internals. - Keep all tensors as `tf.float32`. - Never hardcode CSV column names in plotting code — detect control columns dynamically. --- ## Environment setup (before running) If the user wants to run the training script, make sure the Python environment is correctly configured first. ### 1. Install dependencies ```bash cd <slug> pip install -r requirements.txt ``` `requirements.txt` includes: `tensorflow`, `numpy`, `matplotlib`, `tqdm`, `pandas`. ### 2. Verify TensorFlow can see the GPU (optional but recommended) ```python import tensorflow as tf print(tf.__version__) # should be >= 2.10 print(tf.config.list_physical_devices('GPU')) # empty list = CPU-only mode ``` If GPU is available but not listed, install the matching `tensorflow-gpu` or `cuda`/`cudnn` version. ### 3. Run from the correct directory The training script uses relative imports rooted at `<slug>/`. **Always `cd` into the problem folder first:** ```bash cd <slug> python examples/<slug>_train.py ``` Running from the workspace root will cause `ModuleNotFoundError` for `src.*`. ### 4. Results location Output is written to `<slug>/results/` (configurable via `CommonConfig.save_dir`). CSV training history: `<slug>/results/<saveName>_training_history.csv` Saved models: `<slug>/results/<saveName>_value_model/` and `<slug>/results/<saveName>_control_model/`

path: references/training-output-contract.md
content: # Training Output Contract

## `trainer.train()` return value

Returns `self.history` — a `dict[str, list[float | int]]` with the same keys and ordering as the training history CSV. Values are Python lists, one entry per completed stage.

## CSV files

Written by `trainer.save_history_csv()` when `config.saveOutput = True`.

### `{config.save_dir}/{config.saveName}_training_history.csv`

One row per training stage. Stage count = `config.sampling_stages`, unless early stopping triggers.

| Column | Type | Description |
|--------|------|-------------|
| `stage` | int | Stage index 0, 1, …, N−1 |
| `loss_total` | float | `L1 + L3` (control loss L2 is excluded) |
| `loss_L1` | float | PDE / HJB interior residual loss |
| `loss_L2` | float | Control / FOC residual loss |
| `loss_L3` | float | Terminal condition loss |
| `maxdiff_V` | float | `max|diff_V|` — max abs HJB residual over the validation batch |
| `maxdiff_terminal` | float | `max|diff_terminal|` — max abs terminal residual |
| `V_at_origin` | float | Value network output at `t=0, x=0` (or `x=[1,0]` for 2D) |
| `<control_name>` | float | Control network output at the eval point — one column per name in `config.control_names` |

**Optional columns** (present when listed in `config.metrics_config` and populated via `config.extra_info_mapping`):

| Typical column | Source |
|---------------|--------|
| `max_foc_residual` | returned in `extra_info` from `compute_control_loss` |
| `max_control_bound` | returned in `extra_info` |
| `max_penalty` | returned in `extra_info` |

Any key in `extra_info_mapping` that maps from `compute_control_loss` → metric name will appear as a column if also listed in `config.metrics_config`.

### `{config.save_dir}/{config.saveName}_training_summary.csv`

One row per tracked metric. Columns: `metric, final_value, min_value, max_value, mean_value`.

Covers `loss_total`, `loss_L1`, `loss_L2`, `loss_L3`, plus all entries in `config.metrics_config`.

## `problem.extract_metrics()` interface

Called by `DGMTrainer` after every validation eval step.

```python
def extract_metrics(self, validation_results, t_eval, X_eval, model, control):
    ...
    return dict[str, float]
```

Input `validation_results` tuple (in order):
```
(L1_val, L3_val, diff_V_val, diff_terminal_val, L2_val, extra_info)
```
- `extra_info`: `dict[str, tf.Tensor]` — returned by `compute_control_loss`

**Default implementation** (no override needed for standard metrics):
1. Computes `maxdiff_V = max|diff_V_val|` and `maxdiff_terminal = max|diff_terminal_val|`
2. Evaluates `V_at_origin = model(t_eval, X_eval)[0, 0]`
3. Evaluates each control at `(t_eval, X_eval)` and records under `control_names`
4. Applies `config.extra_info_mapping`: for each `{extra_key: metric_name}`, pulls `extra_info[extra_key]` as `max|·|`
5. Computes `config.custom_metrics` from arithmetic combinations of already-computed metrics

**Override only** when you need metrics that cannot be expressed through `extra_info_mapping` or `custom_metrics`.

## Plotting requirements

- Never assume fixed control column names. Detect them dynamically: any column that is not in the core set (`stage, loss_total, loss_L1, loss_L2, loss_L3, maxdiff_V, maxdiff_terminal, V_at_origin`) is a dynamic column.
- Handle absent optional columns gracefully (check `col in df.columns` before plotting).
- Minimum required plots: loss curves (`loss_total`, `L1`, `L2`, `L3`), `maxdiff_V`, `V_at_origin`.
- Recommended output files: `training_history.png`, `residuals.png`.


path: assets/requirements.txt
content: tensorflow>=2.15.0
numpy>=1.21.0
matplotlib>=3.5.0


path: assets/src/losses/__init__.py
content: """Loss functions package."""


path: assets/src/__init__.py
content: """HJB solver core package with standardized Example1/2/3 modules."""

__version__ = '1.0.0'


path: assets/src/problems/__init__.py
content: """Problems module."""

from .base_problem import BaseProblem

__all__ = ['BaseProblem']


path: assets/src/utils/__init__.py
content: """Utilities module."""

from .visualization import plot_control_surface, plot_control_sum, plot_loss_history

__all__ = ['plot_control_surface', 'plot_control_sum', 'plot_loss_history']


path: assets/src/models/__init__.py
content: """Models module for different neural network architectures."""

from .dgm_net import DenseLayer, DGMNet

__all__ = ['DenseLayer', 'DGMNet']


path: assets/src/trainers/__init__.py
content: """Trainers module."""

from .dgm_trainer import DGMTrainer

__all__ = ['DGMTrainer']


path: assets/src/samplers/base_sampler.py
content: """Base class for sampling strategies."""

from abc import ABC, abstractmethod


class BaseSampler(ABC):
    """Base class for sampling strategies."""
    
    def __init__(self, config, problem):
        """
        Initialize sampler with configuration.
        
        Args:
            config: Configuration object
            problem: Problem object with domain information
        """
        self.config = config
        self.problem = problem
        self.domain_bounds = problem.get_domain_bounds()
    
    @abstractmethod
    def sample(self, nSim_interior, nSim_terminal):
        """
        Sample points from the domain.
        
        Args:
            nSim_interior: Number of interior points to sample
            nSim_terminal: Number of terminal points to sample
            
        Returns:
            Tuple of (t_interior, X_interior, t_terminal, X_terminal)
        """
        pass


path: assets/src/samplers/__init__.py
content: """Samplers module for different sampling strategies."""

from .base_sampler import BaseSampler
from .uniform_sampler import UniformSampler
from .uniform_sampler_2d import UniformSampler2D

__all__ = ['BaseSampler', 'UniformSampler', 'UniformSampler2D']


path: assets/src/configs/__init__.py
content: """Configurations module."""

from .common_config import CommonConfig

__all__ = ['CommonConfig']


Security scan verdicts — ClawScan: clean | static analysis: clean | VirusTotal: clean | SkillSpector: suspicious (MEDIUM)