Skip to content
Closed
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
48 changes: 48 additions & 0 deletions src/MintPlayer.AI.ReinforcementLearning.Core/Checkpoints/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,48 @@
# Checkpoints

Versioned, resumable persistence for training runs and deployable models. Everything here sits
between the in-memory runtime objects (networks, optimizer, agents) and raw `Stream`s; the
`FileModelStore` then puts those streams on disk atomically.

## When do I use which?

One rule covers most cases:

- **Resuming training** → save the *composite state* (`DqnTrainingState` for DQN, `TabularCheckpoint`
for tabular). It carries everything needed to continue *bitwise-identically* — nets, optimizer,
replay buffer, RNG streams.
- **Shipping a trained brain for inference** (the web app, an eval-only run) → save the *individual
net* checkpoint (`DuelingQNetCheckpoint` / `MlpCheckpoint` / `ResidualMlpCheckpoint`). Just shape +
weights; no optimizer or buffer.

The campaigns do both: they persist the full resume state under `<env>/<algo>-state` and, keep-best
gated, the deployable net under `<env>/dqn` — the id the web `ModelService` loads.

## The types

| Checkpoint | Serializes | Standalone or embedded | Format |
|---|---|---|---|
| `MlpCheckpoint` / `ResidualMlpCheckpoint` / `DuelingQNetCheckpoint` | one network's architecture + weights | standalone deployable file **or** embedded via `QNetCheckpoint` | binary |
| `QNetCheckpoint` | any `IValueNet` (type-tagged: `Mlp` \| `DuelingQNet`) | embedded (inside a training state) | binary |
| `AdamCheckpoint` | optimizer hyperparameters + moment estimates — needs the net's parameters to reload | embedded only | binary |
| `ReplayBufferCheckpoint` | DQN replay transitions | embedded only | binary |
| `DqnTrainingState` | **the whole DQN run** — online+target nets + Adam + replay buffer + RNGs + env snapshot | top-level resume file | binary |
| `TabularCheckpoint` | a Q-table | standalone resume **and** deploy (small enough to be both) | JSON (human-inspectable, lossless) |

`CheckpointFormat` holds the shared binary primitives (4-byte magic, kind + version header, float/int/
bool arrays, RNG state). Every binary checkpoint above is built from it, so they share one on-disk
convention and one versioning scheme.

## Persistence seam

`IModelStore` / `FileModelStore` (`ModelStore.cs`) is the file layer: one *current* checkpoint per
`(environmentId, algorithmId)` pair, written to `<root>/<env>.<algo>.ckpt` via a temp-file + rename so
a crash mid-save never corrupts the previous checkpoint. It deals only in raw `Stream`s and knows
nothing about net types — the `*Checkpoint` classes above own that knowledge.

## Versioning

Each checkpoint kind carries its own format version and stays backward-compatible: a newer reader
loads older files (e.g. `DqnTrainingState` v2 files load with a default noise RNG; `DuelingQNet` v1
files load as plain non-noisy nets). Bump the version and branch on it in `Read`/`Load` when you add a
field — never reorder existing fields.
222 changes: 75 additions & 147 deletions src/RLDemo.Web/Services/CubeModelService.cs
Original file line number Diff line number Diff line change
Expand Up @@ -9,21 +9,45 @@ namespace RLDemo.Web.Services;
/// <summary>
/// Owns the Rubik's Cube nets: loads the shipped DQN baseline plus the imitation-policy, teacher-free DAVI value,
/// and EfficientCube policy nets from the model store (all trained on a dev machine via the Lab campaigns and
/// committed to <c>models/</c> via Git LFS — the web never trains, PRD §14). Holds the device-resident GPU
/// forwards for the deep solvers when a CUDA device is present.
/// committed to <c>models/</c> via Git LFS — the web never trains, PRD §14). The three deep-solver nets hot-reload
/// during a training campaign via <see cref="RefreshingModel{T}"/>; for the two GPU-scored nets the reload hook
/// rebuilds a device-resident forward (weights on device — the host-span path is transfer-bound and barely beats
/// CPU, so the resident forward is what makes the deep solvers fast).
/// </summary>
public sealed class CubeModelService(IModelStore store, AdaptiveBackend backend, ILogger<CubeModelService> logger) : IModelStartupService, IDisposable
public sealed class CubeModelService : StartupModelService, IDisposable
{
public const string EnvironmentId = "cube";
public const string AlgorithmId = "dqn";
public const string PolicyAlgorithmId = "policy";
public const string ValueDaviAlgorithmId = "value-davi-res";
public const string EfficientPolicyAlgorithmId = "policy-efficient";

public const int MaxMoves = 20;

private readonly object _lock = new();
private readonly IModelStore _store;
private readonly AdaptiveBackend _backend;
private readonly ILogger<CubeModelService> _logger;

private GreedyQAgent? _agent;
private readonly RefreshingModel<CubePolicyNet> _policy;
private readonly RefreshingModel<ResidualMlp> _value;
private readonly RefreshingModel<CubePolicyNet> _efficient;
private DeviceResidualMlp? _residentValue;
private DeviceMlp? _residentEfficient;

public ModelStatus Status { get; private set; } = ModelStatus.Loading;
public string? Error { get; private set; }
public CubeModelService(IModelStore store, AdaptiveBackend backend, ILogger<CubeModelService> logger) : base(logger)
{
_store = store;
_backend = backend;
_logger = logger;
_policy = new(store, EnvironmentId, PolicyAlgorithmId, CubePolicyNet.Load, logger, "cube policy net");
_value = new(store, EnvironmentId, ValueDaviAlgorithmId, ResidualMlpCheckpoint.Load, logger, "cube DAVI value net",
onLoaded: RebuildResidentValue);
_efficient = new(store, EnvironmentId, EfficientPolicyAlgorithmId, CubePolicyNet.Load, logger, "EfficientCube policy net",
onLoaded: RebuildResidentEfficient);
}

protected override string ModelName => "Rubik's Cube";

/// <summary>The greedy inference agent, or null while the model is not ready. Loads lazily from the store.</summary>
public GreedyQAgent? Agent
Expand All @@ -35,177 +59,81 @@ public GreedyQAgent? Agent
}
}

public const string PolicyAlgorithmId = "policy";
private CubePolicyNet? _policyNet;
private DateTime _policyLoadedUtc = DateTime.MinValue;
private static readonly TimeSpan PolicyRefresh = TimeSpan.FromMinutes(5);

/// <summary>
/// The Kociemba-imitation policy/value net (preferred over the DQN when present, PLAN
/// M16), or null if the store has none. Re-read every few minutes so a long-running
/// Lab campaign's improving checkpoints are picked up without restarting the host.
/// The Kociemba-imitation policy/value net (preferred over the DQN when present, PLAN M16), or null if the store
/// has none. Re-read every few minutes so a long-running Lab campaign's checkpoints are picked up live.
/// </summary>
public CubePolicyNet? PolicyNet
{
get
{
if (DateTime.UtcNow - _policyLoadedUtc < PolicyRefresh) return _policyNet;
lock (_lock)
{
if (DateTime.UtcNow - _policyLoadedUtc < PolicyRefresh) return _policyNet;
try
{
using var stream = store.TryOpenRead(EnvironmentId, PolicyAlgorithmId);
if (stream is not null)
{
_policyNet = CubePolicyNet.Load(stream);
logger.LogInformation("Loaded cube policy net from the store.");
}
}
catch (Exception ex)
{
// A mid-write or corrupt checkpoint must not break solving; keep the previous net.
logger.LogWarning(ex, "Failed to (re)load the cube policy net; keeping the previous one.");
}
_policyLoadedUtc = DateTime.UtcNow;
return _policyNet;
}
}
}
public CubePolicyNet? PolicyNet => _policy.Value;

public const string ValueDaviAlgorithmId = "value-davi-res";
private ResidualMlp? _valueNet;
private DeviceResidualMlp? _residentValue;
private DateTime _valueLoadedUtc = DateTime.MinValue;
/// <summary>
/// The teacher-free DAVI value net (the "self-taught AI" shortest-move solver, PLAN M21), or null if the store
/// has none. Reading it also keeps <see cref="ResidentValueForward"/> in sync (rebuilt on the GPU on reload).
/// </summary>
public ResidualMlp? ValueNet => _value.Value;

/// <summary>
/// The teacher-free DAVI value net (the "self-taught AI" shortest-move solver, PLAN M21), or null
/// if the store has none. Re-read on the same cadence as <see cref="PolicyNet"/> so an improving
/// Lab campaign's checkpoints are picked up without restarting the host. Accessing this also keeps
/// <see cref="ResidentValueForward"/> in sync (rebuilt on the GPU whenever the net reloads).
/// A device-resident GPU forward over the current value net, or null when no CUDA device is present (the solver
/// then scores on the CPU). Touching it applies any pending reload first. Thread-safe: the forward serializes
/// GPU access internally.
/// </summary>
public ResidualMlp? ValueNet
public CubeValueSearch.BatchForward? ResidentValueForward
{
get
{
if (DateTime.UtcNow - _valueLoadedUtc < PolicyRefresh) return _valueNet;
lock (_lock)
{
if (DateTime.UtcNow - _valueLoadedUtc < PolicyRefresh) return _valueNet;
try
{
using var stream = store.TryOpenRead(EnvironmentId, ValueDaviAlgorithmId);
if (stream is not null)
{
_valueNet = ResidualMlpCheckpoint.Load(stream);
// Mirror the freshly-loaded weights onto the GPU as a resident forward when a CUDA
// device is present — the host-span path is transfer-bound and barely beats CPU, so
// the resident forward (weights on device) is what makes the deep solver fast.
_residentValue?.Dispose();
_residentValue = backend.Gpu is { } gpu ? gpu.CreateResidentForward(_valueNet) : null;
logger.LogInformation("Loaded cube DAVI value net from the store ({Mode}).",
_residentValue is not null ? "resident GPU forward" : "CPU forward");
}
}
catch (Exception ex)
{
// A mid-write or corrupt checkpoint must not break solving; keep the previous net.
logger.LogWarning(ex, "Failed to (re)load the cube DAVI value net; keeping the previous one.");
}
_valueLoadedUtc = DateTime.UtcNow;
return _valueNet;
}
}
get { _ = ValueNet; return _residentValue is { } dev ? dev.Forward : null; }
}

/// <summary>
/// A device-resident GPU forward over the current value net, or null when no CUDA device is present
/// (the solver then scores on the CPU). Kept in sync by the <see cref="ValueNet"/> getter — read that
/// first. Thread-safe: the underlying forward serializes GPU access internally.
/// The teacher-free EfficientCube policy net (the website's "self-taught AI", trained self-supervised on
/// scramble reversals — no Kociemba, no value-iteration bootstrap), or null if the store has none. Solved by
/// beam search; reading it keeps <see cref="ResidentEfficientForward"/> in sync.
/// </summary>
public CubeValueSearch.BatchForward? ResidentValueForward => _residentValue is { } dev ? dev.Forward : null;

public const string EfficientPolicyAlgorithmId = "policy-efficient";
private CubePolicyNet? _efficientNet;
private DeviceMlp? _residentEfficient;
private DateTime _efficientLoadedUtc = DateTime.MinValue;
public CubePolicyNet? EfficientPolicyNet => _efficient.Value;

/// <summary>
/// The teacher-free EfficientCube policy net (the website's "self-taught AI", trained self-supervised on
/// scramble reversals — no Kociemba, no value-iteration bootstrap), or null if the store has none. Solved
/// by beam search (<see cref="CubePolicySearch.BeamSearch(System.Func{float[],int,float[]}, FaceletCube, int, int)"/>).
/// Re-read on the same cadence as the other nets so an improving Lab campaign's checkpoints are picked up
/// without restarting the host; accessing this keeps <see cref="ResidentEfficientForward"/> in sync (the
/// policy head's weights mirrored onto the GPU whenever the net reloads).
/// A device-resident GPU forward over the current EfficientCube policy head (rows × 12 logits), or null when no
/// CUDA device is present (beam search then scores on the CPU). Touching it applies any pending reload first.
/// </summary>
public CubePolicyNet? EfficientPolicyNet
public Func<float[], int, float[]>? ResidentEfficientForward
{
get
get { _ = EfficientPolicyNet; return _residentEfficient is { } dev ? dev.Forward : null; }
}

private void RebuildResidentValue(ResidualMlp net)
{
// Serialized against Dispose via Lock (both are the only writer/disposer of _residentValue).
lock (Lock)
{
if (DateTime.UtcNow - _efficientLoadedUtc < PolicyRefresh) return _efficientNet;
lock (_lock)
{
if (DateTime.UtcNow - _efficientLoadedUtc < PolicyRefresh) return _efficientNet;
try
{
using var stream = store.TryOpenRead(EnvironmentId, EfficientPolicyAlgorithmId);
if (stream is not null)
{
_efficientNet = CubePolicyNet.Load(stream);
// Beam search runs the bulk of the forwards; mirror the policy head onto the GPU as a
// resident forward when a CUDA device is present (host-span is transfer-bound).
_residentEfficient?.Dispose();
_residentEfficient = backend.Gpu is { } gpu ? gpu.CreateResidentForward(_efficientNet.PolicyAsMlp()) : null;
logger.LogInformation("Loaded EfficientCube policy net from the store ({Mode}).",
_residentEfficient is not null ? "resident GPU forward" : "CPU forward");
}
}
catch (Exception ex)
{
// A mid-write or corrupt checkpoint must not break solving; keep the previous net.
logger.LogWarning(ex, "Failed to (re)load the EfficientCube policy net; keeping the previous one.");
}
_efficientLoadedUtc = DateTime.UtcNow;
return _efficientNet;
}
_residentValue?.Dispose();
_residentValue = _backend.Gpu is { } gpu ? gpu.CreateResidentForward(net) : null;
}
_logger.LogInformation("Cube DAVI value net scored on {Mode}.", _residentValue is not null ? "resident GPU forward" : "CPU");
}

/// <summary>
/// A device-resident GPU forward over the current EfficientCube policy head (rows × 12 logits), or null when
/// no CUDA device is present (beam search then scores on the CPU). Kept in sync by the
/// <see cref="EfficientPolicyNet"/> getter — read that first. Thread-safe: the forward serializes GPU access.
/// </summary>
public Func<float[], int, float[]>? ResidentEfficientForward => _residentEfficient is { } dev ? dev.Forward : null;
private void RebuildResidentEfficient(CubePolicyNet net)
{
lock (Lock)
{
_residentEfficient?.Dispose();
// Beam search runs the bulk of the forwards; mirror the policy head onto the GPU when a device is present.
_residentEfficient = _backend.Gpu is { } gpu ? gpu.CreateResidentForward(net.PolicyAsMlp()) : null;
}
_logger.LogInformation("EfficientCube policy net scored on {Mode}.", _residentEfficient is not null ? "resident GPU forward" : "CPU");
}

public bool TryLoadFromStore()
public override bool TryLoadFromStore()
{
lock (_lock)
lock (Lock)
{
if (_agent is not null) return true;
using var stream = store.TryOpenRead(EnvironmentId, AlgorithmId);
using var stream = _store.TryOpenRead(EnvironmentId, AlgorithmId);
if (stream is null) return false;
var network = MlpCheckpoint.Load(stream);
_agent = new GreedyQAgent(network, RubiksCubeEnv.ActionCount);
Status = ModelStatus.Ready;
logger.LogInformation("Loaded Rubik's Cube model from the store.");
_logger.LogInformation("Loaded Rubik's Cube model from the store.");
return true;
}
}

/// <summary>Loads the shipped DQN baseline at startup. The web does not train (PRD §14); the deep solver
/// nets (policy / DAVI value / EfficientCube) load lazily through their getters.</summary>
public void Initialize(CancellationToken cancellationToken)
{
if (TryLoadFromStore()) return;
lock (_lock)
{
Status = ModelStatus.Failed;
Error = "No trained Rubik's Cube model in the store.";
}
logger.LogWarning("No Rubik's Cube DQN model in the store — train it on a dev machine and commit the checkpoint to models/ (Git LFS).");
}

/// <summary>
/// Greedy rollout on a drawn cube: up to <see cref="MaxMoves"/> quarter-turns under
/// the same no-undo mask the agent was trained with, reported honestly whether they
Expand All @@ -232,7 +160,7 @@ public static (bool Solved, List<string> Moves) Rollout(GreedyQAgent agent, Face

public void Dispose()
{
lock (_lock)
lock (Lock)
{
_residentValue?.Dispose();
_residentValue = null;
Expand Down
24 changes: 5 additions & 19 deletions src/RLDemo.Web/Services/Game2048ModelService.cs
Original file line number Diff line number Diff line change
Expand Up @@ -6,18 +6,16 @@ namespace RLDemo.Web.Services;
/// <summary>
/// Owns the 2048 n-tuple agent: loads the shipped checkpoint from the model store at startup (trained on a dev
/// machine via self-play, the M5 recipe, and committed to <c>models/</c> via Git LFS — the web never trains,
/// PRD §14).
/// PRD §14). Startup/status plumbing lives in <see cref="StartupModelService"/>.
/// </summary>
public sealed class Game2048ModelService(IModelStore store, ILogger<Game2048ModelService> logger) : IModelStartupService
public sealed class Game2048ModelService(IModelStore store, ILogger<Game2048ModelService> logger) : StartupModelService(logger)
{
public const string EnvironmentId = "2048";
public const string AlgorithmId = "ntuple";

private readonly object _lock = new();
private NTuple2048Agent? _agent;

public ModelStatus Status { get; private set; } = ModelStatus.Loading;
public string? Error { get; private set; }
protected override string ModelName => "2048";

/// <summary>The trained agent, or null while not ready. Loads lazily from the store.</summary>
public NTuple2048Agent? Agent
Expand All @@ -29,9 +27,9 @@ public NTuple2048Agent? Agent
}
}

public bool TryLoadFromStore()
public override bool TryLoadFromStore()
{
lock (_lock)
lock (Lock)
{
if (_agent is not null) return true;
using var stream = store.TryOpenRead(EnvironmentId, AlgorithmId);
Expand All @@ -42,16 +40,4 @@ public bool TryLoadFromStore()
return true;
}
}

/// <summary>Loads the shipped checkpoint at startup. The web does not train (PRD §14).</summary>
public void Initialize(CancellationToken cancellationToken)
{
if (TryLoadFromStore()) return;
lock (_lock)
{
Status = ModelStatus.Failed;
Error = "No trained 2048 model in the store.";
}
logger.LogWarning("No 2048 model in the store — game unavailable. Train it on a dev machine and commit the checkpoint to models/ (Git LFS).");
}
}
Loading
Loading