Skip to content

Latest commit

 

History

5 Commits

Folders and files

NameName
Last commit message
Last commit date
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

The RL Pipeline — A Gentle Walkthrough

This document explains the RL (Reinforcement Learning) pipeline step by step, assuming almost no background in machine learning or RL. Every concept is introduced when first encountered, with simple examples.


1. What Are We Trying to Do?

We have a physics simulation that prepares quantum states (Dicke states). The simulation is controlled by two knobs:

  • gamma_in ($\gamma_\text{in}$) — a number between 0 and $\pi$ (roughly 0 to 3.14)
  • gamma_sh ($\gamma_\text{sh}$) — a number between 0 and $\pi$

Think of them like two volume knobs on a machine. Depending on where you set them, the machine produces a different quantum state. Some settings produce a state that's very close to what we want (high "fidelity"), and most settings produce garbage.

The goal: Given a quantum state we want to prepare (described by two numbers $N$ and $K$), find the best knob settings (gamma values) that make the machine produce the closest match.

Why is this hard? There's no formula to calculate the best gammas. You have to actually run the simulation and check. The simulation is slow (especially for large $N$), and the landscape is bumpy — there are many local "hills" and "valleys." Turning the knobs a tiny bit can change the result a lot.

What we have:

  • For small states ($N = 5$ to $10$), we've already tried thousands of gamma pairs via brute force. We know the good gammas.
  • For larger states ($N = 11$ to $14$), we haven't searched yet. We need a way to predict good gammas without trying them all.

This is where the RL pipeline comes in: learn patterns from the small states, and use those patterns to predict gammas for the larger ones.


2. The Full Pipeline at a Glance

Step 1: Load data         →  "Here are gammas that worked well for small states"
Step 2: Train a network   →  "Learn to predict good gammas from (N,K)"
Step 3: Make predictions  →  "Given a new (N,K), suggest some gamma candidates"
Step 4: Test them          →  "Run each candidate in the simulation, keep the best"

3. Step 1: Loading Training Data

Where it lives: data/optimal_with_states_D5_10_005sp/

We have files full of entries like:

For D(5,2):   gamma_in=0.000, gamma_sh=1.515  →  fidelity = 0.9895  (great!)
For D(5,2):   gamma_in=1.570, gamma_sh=1.570  →  fidelity = 0.0200  (terrible!)
For D(10,3):  gamma_in=1.619, gamma_sh=0.010  →  fidelity = 0.8872  (decent)

Each entry tells us: "for this $(N,K)$, this gamma pair gave this fidelity score." Fidelity is a number from 0 (no match at all) to 1 (perfect match).

We load about 4,032 gamma pairs for each of 15 different (N,K) states — that's about 60,000 entries total.

Training states (15): D(5,2), D(6,2), D(6,3), D(7,2), D(7,3), D(8,2), D(8,3), D(8,4), D(9,2), D(9,3), D(9,4), D(10,2), D(10,3), D(10,4), D(10,5)

Test states (24): All valid (N,K) for N = 11 through 14. These are the states we want to predict good gammas for — we have NO pre-computed data for these.


4. Step 2: Training the Neural Network

4.1 What is a Neural Network?

A neural network is a math function that takes some numbers in and produces some numbers out. What makes it special is that it has thousands of adjustable parameters (called "weights") inside. By adjusting these weights during "training," we can make it produce the outputs we want for given inputs.

Analogy: Imagine a very flexible curve-fitting tool. You show it examples: "when the input is X, the output should be Y." It adjusts itself to match these examples, and then it can predict outputs for inputs it hasn't seen before.

4.2 What Goes In: Input Features

We don't feed raw $N$ and $K$ directly. Instead, we compute 6 numbers (called "features") that describe the state in a way the network can use:

# Feature Formula What it means Example for D(8,3)
1 N normalized $2 \times \frac{N - 5}{20 - 5} - 1$ How big the system is, scaled to [-1, 1]. N=5 gives -1, N=20 gives +1. $2 \times \frac{3}{15} - 1 = -0.6$
2 K normalized $2 \times \frac{K}{N} - 1$ How "excited" the system is relative to its size. $2 \times \frac{3}{8} - 1 = -0.25$
3 K/N $\frac{K}{N}$ Fraction of excitations. 0.5 means half-excited. $0.375$
4 (N-K)/N $\frac{N-K}{N}$ Fraction of "register" qubits. $0.625$
5 log C(N,K) $\frac{\log \binom{N}{K}}{\log \binom{20}{10}}$ How big the relevant part of the quantum state space is (normalized). $\frac{\log 56}{\log 184756} \approx 0.332$
6 sqrt C(N,K) $\frac{\sqrt{\binom{N}{K}}}{\sqrt{\binom{20}{10}}}$ Related to the target amplitude of each term. $\frac{\sqrt{56}}{\sqrt{184756}} \approx 0.0174$

Why not just use N and K? Because the network needs help understanding what matters. For example, D(8,3) and D(12,3) have the same K but very different physics. The ratio K/N, the Hilbert space size, etc., capture the relevant structure better than raw numbers.

4.3 What Comes Out: The 4-Head Design

This is the most important design decision. Let me explain with an analogy.

The problem with one prediction: Suppose you ask someone: "What's a good restaurant in town?" If there are two great restaurants — one Italian and one Japanese — and you force them to give ONE answer, they might average the two locations and point you to an empty parking lot in between. That's useless! The same thing happens with gamma prediction: for D(8,4), good gammas exist near (0, 0.4) AND near (π, 2.75). A single prediction would average them to roughly (π/2, π/2), which is the WORST possible region.

The solution — 4 heads: Instead of making one prediction, the network makes 4 independent predictions through 4 separate "heads." Each head can focus on a different good region.

Common misconception: You might picture the trunk (256 neurons) connecting to 4 single nodes at the end — like a funnel narrowing to 4 endpoints. That's NOT what happens. Each head is itself a small neural network (called a "branch"), not a single node. Here's the actual structure:

Input: 6 features describing (N, K)
            │
            ▼
  ┌────────────────────┐
  │  Trunk Layer 1     │   6 → 256 neurons (+ ReLU + LayerNorm)
  └────────┬───────────┘
           │  256 numbers flow down
           ▼
  ┌────────────────────┐
  │  Trunk Layer 2     │   256 → 256 neurons (+ ReLU + LayerNorm)
  └────────┬───────────┘
           │  256 numbers flow down
           │  (same 256 numbers are COPIED to all 4 heads)
           │
     ┌─────┼─────────┬───────────┐
     │     │         │           │
     ▼     ▼         ▼           ▼
  Head 0  Head 1    Head 2     Head 3
     │     │         │           │
     ▼     ▼         ▼           ▼
  4 nums  4 nums    4 nums     4 nums

But what's INSIDE each head? Each head is NOT a single node. It's two mini-networks ("branches") running in parallel:

                       Head 0 (zoomed in)
                 ┌──────────────────────────────────────────┐
                 │                                          │
  256 numbers ───┤                                          │
  from trunk     │   mu (center) branch:                    │
                 │   256 → [128 neurons → ReLU] → [2 neurons → Tanh]  →  (μ_in, μ_sh)
                 │                                          │
                 │   sigma (spread) branch:                 │
                 │   256 → [128 neurons → ReLU] → [2 neurons]         →  (σ_in, σ_sh)
                 │                                          │
                 └──────────────────────────────────────────┘
                                    │
                          Outputs 4 numbers:
                          μ_in, μ_sh  (center guesses for gamma_in and gamma_sh)
                          σ_in, σ_sh  (uncertainty for each)

So the network is much wider at the end than you might think:

  • The trunk part: 6 → 256 → 256 (a funnel that WIDENS)
  • Then it splits into 4 heads, each containing:
    • A mu branch: 256 → 128 → 2 (narrows toward the center guess)
    • A sigma branch: 256 → 128 → 2 (narrows toward the uncertainty)
  • Each head has 130 neurons in the mu path + 130 in the sigma path = 260 neurons
  • 4 heads × 260 = 1040 neurons after the trunk
  • Total output: 4 heads × 4 numbers = 16 numbers

What is a "neuron"? A neuron is the basic building block of a neural network. It takes some numbers, multiplies each by a weight, adds them up, and produces one output number. A "layer of 256 neurons" means 256 of these units working in parallel. More neurons = more capacity to learn complex patterns.

Total network size:

  • 6 input features
  • Trunk: two layers of 256 neurons (shared by all heads)
  • 4 heads, each with two branches of 128 hidden neurons → 2 output neurons
  • Grand total: roughly 200,000 adjustable weights

Why does each head need 128 hidden neurons? Why not just 2 neurons that directly output μ and σ?

You might wonder: the trunk already processed the input through 256 neurons twice — isn't that enough "thinking"? Can't each head just be a simple 256 → 2 connection?

The answer is no, because of what "simple connection" means mathematically. A direct 256 → 2 connection (without a hidden layer) computes:

$$\mu = w_1 \times \text{trunk}_1 + w_2 \times \text{trunk}_2 + \cdots + w_{256} \times \text{trunk}_{256} + b$$

That's a weighted sum — a linear function. It can only draw a flat hyperplane through the feature space. It's like trying to separate regions on a contour map using only straight lines.

But the relationship between the trunk's features and the best gammas is curved and bumpy. Different heads need to carve out different irregularly-shaped regions of gamma space. A linear function simply can't do this.

By adding a 128-neuron hidden layer with ReLU, each head gets its own non-linear transformation:

Without hidden layer (linear):       trunk → [weighted sum] → 2 outputs
                                      Can only learn: "if feature A is big, gamma is big"

With 128 hidden neurons + ReLU:       trunk → [128 neurons + ReLU] → [2 outputs]
                                      Can learn: "if feature A is big AND feature B
                                      is small, gamma is big — BUT only when feature C
                                      is in a certain range"

Analogy: Imagine you have 4 delivery drivers (heads) and a city map. Without the hidden layer, each driver can only describe their delivery zone as "everything north of Main Street" (a straight-line boundary). With the hidden layer, each driver can say "the irregularly-shaped neighborhood between the park, the river, and 5th Avenue" — a much more precise zone.

Why 128 specifically? There's no deep reason — it's a practical choice. 128 is half the trunk width (256), which is a common pattern in neural network design. Using 64 would also work but with slightly less capacity. Using 256 would work too but risks overfitting (learning noise instead of patterns, since we only have 15 training states). 128 is a reasonable middle ground.

What are ReLU, Tanh, and LayerNorm?

  • ReLU: A simple rule: if the number is positive, keep it; if negative, make it 0. Purpose: lets the network learn non-linear patterns (straight lines alone can't capture the bumpy gamma landscape).
  • Tanh: Squashes any number into the range [-1, +1]. Purpose: forces the output to be bounded, which we then convert to [0, π] for gammas.
  • LayerNorm: Adjusts the numbers in a layer so they have mean 0 and standard deviation 1. Purpose: makes training more stable and faster. Without it, numbers can grow very large or very small during training, causing problems.

4.4 What Each Head Actually Outputs

Each head produces 4 numbers — a pair for gamma_in and a pair for gamma_sh:

  • μ (mu) = the center — "My best guess for the gamma value." This is a number in [-1, 1] (thanks to Tanh), which we convert to [0, π] for the actual gamma.

    Conversion: $\gamma = \frac{\mu + 1}{2} \times \pi$

    So μ = -1 → γ = 0, μ = 0 → γ = π/2, μ = +1 → γ = π.

  • σ (sigma) = the spread — "How uncertain I am." A bigger sigma means "I'm not very sure — you should explore around my guess." A smaller sigma means "I'm confident — my guess is precise."

    Technically, the network outputs log(σ) (called log_sigma), which is clamped to the range [-3.0, 0.5]. This means:

    • Minimum σ = $e^{-3.0}$ ≈ 0.05 (very confident)
    • Maximum σ = $e^{0.5}$ ≈ 1.65 (very uncertain)

    In practice, after training, σ settles around 0.22 for most heads.

Together, μ and σ define a bell curve (Gaussian/normal distribution). We can either:

  • Take μ directly (the "best guess") — called a deterministic prediction
  • Roll dice using the bell curve centered at μ with width σ — called a stochastic sample (useful for exploring alternatives)

Simple example: If Head 0 outputs μ = (0.5, -0.8) and σ = (0.2, 0.2), it's saying: "I think gamma_in ≈ 2.36 and gamma_sh ≈ 0.31, but you should explore within about ±0.2 of my action-space values."

4.5 Training Phase 1: Supervised Pre-Training (3000 epochs)

"Supervised" means we have answers (the known-good gammas from training data) and we tell the network: "here are examples of what good outputs look like — learn to produce them."

What is an epoch? One pass through all the training data. So 3000 epochs means we show the network all 15 training states 3000 times, adjusting the weights a little each time.

Parameters:

  • TRAIN_EPOCHS = 3000 — number of passes through the data
  • LEARNING_RATE = 3e-3 (= 0.003) — how big each adjustment step is. Too large → the network overshoots and never settles. Too small → training is painfully slow. 0.003 is a common starting value.
  • weight_decay = 1e-5 (= 0.00001) — a tiny penalty on large weight values. Prevents the network from memorizing noise.

How does it learn? The loss function.

A "loss function" is a number that measures "how bad is the network right now." Training tries to make this number as small as possible. Our loss has three parts:

Part 1: Winner-Take-All Loss (WTA)

For each known-good gamma pair in the training data:

  1. Measure the distance from EACH head's prediction to that target
  2. Only blame the CLOSEST head (the "winner") — update it to get closer
  3. Leave the other heads alone

Why? This prevents all 4 heads from chasing the same target. Over time, each head naturally specializes on a different region of the gamma space.

Analogy: Imagine 4 pizza delivery drivers. When an order comes in, only the closest driver picks it up. Over time, each driver naturally covers a different neighborhood.

The distance is weighted by the fidelity of the target — high-fidelity targets matter more.

Part 2: Sigma Regularization Loss

Formula: $\text{sigma_loss} = \left(\left(\log \sigma + 1.5\right)^2\right)_\text{mean} \times 0.01$

What this does: Pushes σ toward $e^{-1.5}$ ≈ 0.22. Without this, σ could collapse to near-zero (network becomes overconfident and stops exploring) or blow up (network guesses randomly).

The coefficient 0.01 makes this a gentle nudge, not a hard constraint. Sigma can deviate from 0.22 if there's a good reason.

Part 3: Diversity Loss

Formula: For every pair of heads, compute: $-\log(\text{distance}^2 + 0.01)$, averaged and scaled by 0.05.

What this does: If two heads predict almost the same thing, this loss increases (because $-\log(\text{small number})$ is a large positive number). So heads are repelled from each other, like magnets with the same pole.

The coefficient 0.05 controls how strongly heads push apart.

Total loss = WTA + Sigma Reg + Diversity

The optimizer (Adam) adjusts all ~200,000 weights to reduce this total loss.

What is Adam? An optimizer algorithm that adjusts the network's weights. It's smarter than simple "go downhill" because it adapts: parameters that haven't changed much get bigger updates, and parameters that have been oscillating get smaller updates. Think of it as a ball rolling downhill that remembers its momentum.

Learning rate schedule: CosineAnnealingWarmRestarts

  • Starts at lr = 0.003
  • Gradually decreases (so the network takes smaller and smaller steps)
  • Every T_0 = 500 epochs, resets back to the original learning rate
  • After each restart, the period doubles (T_mult = 2): 500, 1000, 2000, ...

Why restart? Sometimes the network gets stuck in a bad spot. Resetting the learning rate gives it a "kick" to escape and explore other solutions. The longer warmup periods mean later restarts are gentler.

Gradient clipping: max_norm = 1.0 — If any weight update is too large, it gets scaled down so its total size is at most 1.0. This prevents the network from making wild jumps that could undo previous learning.

Best model selection: After each epoch, if the loss is the lowest we've ever seen, we save the network's weights. At the end of training, we restore the best weights — not the final ones (which might be slightly worse due to noise).

4.6 Training Phase 2: REINFORCE Fine-Tuning (300 episodes)

After supervised pre-training, the network is decent but it was only trained on the recorded good gammas from the dataset. It might be able to do better by actually trying gammas in the simulation and learning from the results. This is the Reinforcement Learning part.

What is Reinforcement Learning? A learning paradigm where an agent:

  1. Observes the situation (here: the (N,K) state)
  2. Takes an action (here: picks a gamma pair)
  3. Gets a reward (here: the fidelity score from the simulation)
  4. Adjusts its behavior to get more reward next time

No labeled answers needed — just trial and error.

What is REINFORCE / Policy Gradient?

REINFORCE is one of the simplest RL algorithms. The idea:

  • If you tried an action and got a HIGH reward → increase the probability of that action
  • If you tried an action and got a LOW reward → decrease the probability of that action

Mathematically: adjust the network weights in the direction of $\text{reward} \times \nabla \log \pi(a|s)$, where $\pi(a|s)$ is the probability the network assigned to action $a$ in state $s$.

Simple analogy: You're blindfolded throwing darts. Each time, someone tells you the score. If you scored high, you try to throw more like that. If you scored low, you adjust. Over time, you get better even without seeing the target.

Parameters:

  • REINFORCE_EPISODES = 300 — how many "throws" to make
  • lr = 1e-4 (= 0.0001) — much smaller than supervised training! We don't want to forget what we learned, just fine-tune it
  • gradient clip = 0.5 — even tighter than supervised training, for stability

What happens in each episode:

  1. Pick a random (N,K) from the 15 training states
  2. Pick a random head (0, 1, 2, or 3)
  3. Sample an action from that head's distribution (mu + sigma × random_noise, then tanh)
  4. Convert the action to gammas and run the actual simulation → get fidelity (this is the reward)
  5. Compute the "advantage" = reward minus baseline

What is a baseline? If you just use raw reward, the network will reinforce EVERY action (because fidelity is always ≥ 0). The baseline is the average of recent rewards (last 50). Actions above average get reinforced; actions below average get discouraged. This dramatically reduces noise.

Example: If average fidelity is 0.5, then:

  • Action with fidelity 0.8 → advantage = +0.3 → reinforce it!
  • Action with fidelity 0.3 → advantage = -0.2 → discourage it
  1. Update the network weights: $ \Delta w \propto \text{advantage} \times \nabla \log \pi(a|s) $

Why only 300 episodes? Each episode requires running a full physics simulation (which is expensive). 300 episodes takes about 2-3 minutes. More episodes give diminishing returns because there are only 15 training states — the network quickly learns the main patterns.

Training time: ~5.6 minutes total (3000 supervised epochs + 300 REINFORCE episodes)


5. Step 3: Generating Predictions (Candidate Generation)

After training, we have a network that can predict gamma distributions for any (N,K). To predict for a new state (say D(12,4)):

5.1 Generate up to 28 gamma candidates

We generate candidates from several sources to balance exploitation (using what the network learned) and exploration (trying things the network might not have considered):

Source How many How it works
Deterministic 4 Take μ from each head directly. These are the network's "best guesses."
Stochastic 8 Sample from each head's bell curve (2 samples per head). Like rolling dice near each best guess.
High-entropy (2×) 4 Double each head's σ, then sample. This explores a wider area.
High-entropy (3×) 4 Triple each head's σ, then sample. Even wider — catches distant optima.
Random 4-8 Uniform random in [0, π] × [0, π]. Pure exploration — no network guidance. These catch cases where the network is completely wrong.
Total ≤ 28

Deduplication: If two candidates fall within 0.01 of each other (in rounded gamma values), only one is kept. This avoids wasting simulation time on near-identical candidates.

5.2 Evaluate each candidate

For each of the 28 candidates, we run the actual physics pipeline:

gamma pair → dicke_evolution(max_rounds=200) → find best round → rz_optimize → rz_fidelity

What is max_rounds=200? The collision simulation runs for up to 200 collision rounds. At each round, we check how close the quantum state is to the target. Different gammas peak at different rounds — some might peak at round 5, others at round 150.

What is "best round"? The round where the fidelity is highest BEFORE applying RZ optimization. Note: the array has more entries than 200 because each round produces K intermediate states (one per ancilla), so the total entries = 1 + 200 × K.

What is RZ optimization? After finding the best round, we apply rotation gates (RZ gates) to each qubit to correct the phases of the quantum state. This is a separate numerical optimization (L-BFGS-B algorithm) that adjusts N angles to maximize fidelity. It typically improves the score significantly.

5.3 Local refinement

After evaluating all 28 candidates, we take the best one and try to improve it by making small random perturbations:

For each noise_scale in [0.05, 0.1, 0.2]:
    Try 8 random perturbations:
        new_gamma_in = best_gamma_in + random_normal × noise_scale
        new_gamma_sh = best_gamma_sh + random_normal × noise_scale
        Run simulation
        If better → update best

Parameters:

  • LOCAL_REFINE_STEPS = 8 — number of random perturbations per noise scale
  • 3 noise scales:
    • 0.05 — fine-grained search nearby (precision tuning)
    • 0.1 — medium-range search
    • 0.2 — wider search (in case we're near the wrong local peak)
  • Total: 3 × 8 = 24 additional simulations on top of the initial 28

Why random perturbation instead of a smarter optimizer? We tried Nelder-Mead (a gradient-free optimizer) in v5, but it gave similar results while using ~45 simulations (vs 24). The random approach is simpler and sufficient because the network already gets close.

5.4 Total simulation budget per state

Phase Simulations
Candidate evaluation 28
Local refinement 24
Total ~52

Each simulation takes ~1-4 seconds depending on N, so evaluation of all 24 test states takes about 60-90 minutes.


6. Understanding the Training Data Flow

Let me trace a concrete example through the entire pipeline:

Example: Predicting gammas for D(12,4) (N=12, K=4) — a test state we have NO data for.

Step 1 — Compute features:

N_norm     = 2 × (12-5)/(20-5) - 1 = -0.067
K_norm     = 2 × 4/12 - 1 = -0.333
K/N        = 0.333
(N-K)/N    = 0.667
log C(12,4) normalized = log(496)/log(184756) = 0.512
sqrt C(12,4) normalized = √496/√184756 = 0.052
→ input = [-0.067, -0.333, 0.333, 0.667, 0.512, 0.052]

Step 2 — Network forward pass: These 6 numbers flow through the trunk (two 256-neuron layers), then branch into 4 heads. Each head outputs μ (2 values in [-1,1]) and σ (2 values ≈ 0.22).

Step 3 — Generate candidates: The 4 deterministic predictions, 8 stochastic samples, 8 high-entropy samples, and some random gammas → 28 candidates.

Step 4 — Simulate each: Each candidate runs through gen_dicke_sim (200 rounds) + optimize_state_rz → get rz_fidelity.

Step 5 — Refine best: 24 random perturbations around the best candidate.

Step 6 — Result: The gamma pair with the highest fidelity wins.

In the latest run, D(12,4) achieved rz_fidelity = 0.8749 with gammas (2.918, 3.107), found via local_refine_0.2 (meaning it was a perturbation of an initially good candidate).


7. What Went Wrong in Earlier Versions (and Why v4 Works)

V1: All predictions collapse to (π/2, π/2) — fidelity 0.39

Used a standard single-output RL algorithm (SAC). With one output, the network averaged all good gammas into one prediction — the center of the [0,π] range. But gammas near (π/2, π/2) are consistently the WORST region.

Lesson: You need multiple outputs to handle multiple good answers.

V2: Better but still mediocre — fidelity 0.63

Added supervised pre-training before SAC. This helped, but the single-output network still couldn't represent multiple modes. Good results came mostly from lucky random perturbations, not from the network.

Lesson: Pre-training helps, but the architecture needs to support multimodality.

V4: 4-head stochastic policy — fidelity 0.86

The current version fixes both problems:

  • 4 heads → can represent 4 different good regions
  • Stochastic outputs (μ + σ) → natural exploration around predictions
  • REINFORCE → improves predictions beyond the training data
  • No data lookups → generalizes purely from learned patterns

V5: Tried to improve v4, but couldn't — fidelity 0.85

Tried: softer sigma regularization, systematic REINFORCE, entropy bonus, Nelder-Mead optimizer, complement gammas. None helped. Reverted back to v4.

Why? The bottleneck is the training data (only 15 states), not the architecture. Both v4 and v5 plateau at ~0.85.


8. Complete Parameter Reference

Network Architecture

Parameter Value Meaning
input_dim 6 Number of input features describing (N,K)
n_heads 4 Number of prediction heads
hidden_dims [256, 256] Two trunk layers, 256 neurons each
Head mu branch 128 → 2 128 hidden neurons, 2 output (gamma_in, gamma_sh)
Head sigma branch 128 → 2 128 hidden neurons, 2 output (log_sigma_in, log_sigma_sh)
log_sigma clamp [-3.0, 0.5] σ range: [0.05, 1.65]. Prevents extreme values
Activation ReLU + Tanh (mu output) ReLU in hidden layers, Tanh to bound output
Normalization LayerNorm Applied after each trunk layer

Supervised Pre-Training

Parameter Value Meaning
TRAIN_EPOCHS 3000 Number of full passes through the data
LEARNING_RATE 0.003 Initial learning rate for Adam optimizer
weight_decay 1e-5 L2 regularization strength (prevents overfitting)
Scheduler CosineAnnealingWarmRestarts Learning rate decay with periodic resets
T_0 500 First restart after 500 epochs
T_mult 2 Restart intervals double: 500, 1000, 2000
Gradient clip 1.0 Maximum gradient norm
WTA loss weight 1.0 Weight of winner-take-all loss (implicit, it's the main loss)
Sigma reg coefficient 0.01 How strongly sigma is pushed toward ≈0.22
Sigma reg target $e^{-1.5}$ ≈ 0.22 Target standard deviation
Diversity coefficient 0.05 How strongly heads are pushed apart
min_dist for clustering 0.3 Minimum distance between diverse training targets

REINFORCE Fine-Tuning

Parameter Value Meaning
REINFORCE_EPISODES 300 Number of trial-and-error episodes
Learning rate 0.0001 30x smaller than supervised — gentle fine-tuning
Gradient clip 0.5 Tighter than supervised — extra stability
Baseline window 50 Average reward over last 50 episodes
Head selection Random Each episode picks a random head
State selection Random Each episode picks a random (N,K)

Candidate Generation

Parameter Value Meaning
N_CANDIDATES 28 Maximum candidates to generate
Deterministic 4 One per head (μ directly)
Stochastic per head 2 Samples from learned distribution
High-entropy scales 2×, 3× Sigma inflation for wider exploration
Random minimum 4 Minimum uniform random candidates
Dedup threshold 0.01 Candidates within 0.01 are merged

Simulation

Parameter Value Meaning
MAX_ROUNDS_TRAIN 200 Collision rounds for both training and test evaluation
LOCAL_REFINE_STEPS 8 Random perturbations per noise scale
Noise scales [0.05, 0.1, 0.2] Three levels of perturbation for refinement
Total sims/state ~52 28 candidates + 24 refinement

Training Data

Parameter Value Meaning
Training states 15 D(5,2) through D(10,5)
Test states 24 D(11,1) through D(14,7)
Samples per state ~4032 Grid search results from brute-force
Total training entries ~60,480 15 states × 4032 entries
rz_fidelity_threshold 0.0 Keep all entries (no filtering)

9. Results

Version Test Mean Fidelity Eval Time Approach
V1 ~0.39 (training only) Single-output SAC
V2 0.63 Supervised + SAC
V4 0.86 64 min 4-Head Policy + REINFORCE
V5 0.85 84 min V4 + Nelder-Mead (reverted)

The ±0.02 variation between runs is normal — it comes from the randomness in candidate generation and REINFORCE training.


10. File Guide

File What it does
agent.py The neural network + agent. All the architecture, training loops, and candidate generation.
train.py Orchestrates training: load data → train network → evaluate on training states → save model.
evaluate.py Evaluates the saved model on test states (N=11-14). Generates candidates, simulates, refines.
data_loading.py Reads the brute-force gamma data from disk and organizes it by (N,K).
models/agent.pt Saved network weights (~200K parameters, less than 1 MB file).
models/training_results.json Full-precision training results (gammas, fidelities, best rounds).
models/evaluation_results.json Full-precision test results for all 24 states.
EXPLANATION.md The concise technical explanation.
progress_log.txt Detailed history of all versions and experiments.

About

No description, website, or topics provided.

Resources

Stars

0 stars

Watchers

1 watching

Forks

Releases

Packages

Contributors

Languages