# Deep Hjb Solver

> 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.

- Skill: `dvcrn/deep-hjb-solver` (Agent Skill, multi-file: 22 files)
- Install (CLI): `npx skillmds@latest add dvcrn/deep-hjb-solver`
- Raw SKILL.md: https://api.skillmd.com/api/skills/dvcrn/deep-hjb-solver/raw
- Safety review: pending (external: skill-scanner PASS, skillspector PASS)
- Works with: Claude Code, Claude.ai, OpenAI Codex
- Category: Data & Analytics
- Author: dvcrn (https://skillmd.com/u/dvcrn)
- Updated: 2026-09-08
- Page: https://skillmd.com/skills/dvcrn/deep-hjb-solver

---


# 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/`