GTN-Base-V6.0 Reward-Calibrated Surrogate Model

This repository contains the trained gtn-base-v6.0 reward-calibrated GTN surrogate checkpoint for shared-energy power-converter topology evaluation.

Unlike gtn-base-v5, which predicts only physical simulator targets, gtn-base-v6.0 keeps the v5 GTN surrogate backbone and adds a target-gamma-conditioned reward calibration head:

f_theta(T, d) -> (efficiency, Vout)
g_phi(embedding(T, d), target_gamma, f_theta(T, d)) -> calibrated_reward

Where:

• T is the power-converter topology graph.
• d is the duty cycle.
• target_gamma is the target normalized output-voltage ratio.
• calibrated_reward is the reward-calibrated fan reward used for topology/duty scoring.

The canonical model class is:

shared_energy.gtn_surrogate.GTNRewardCalibratedSurrogateModel

Loading

This is a project-native PyTorch checkpoint, not a Transformers model. The local shared_energy codebase must be importable.

from huggingface_hub import hf_hub_download
import torch
from shared_energy.gtn_surrogate import GTNRewardCalibratedSurrogateModel

repo_id = "DanielJeongsooLee/gtn-base-v6.0"
checkpoint_path = hf_hub_download(
    repo_id=repo_id,
    filename="gtn_reward_calibrated_surrogate_model.pt",
)

payload = torch.load(checkpoint_path, map_location="cpu")
model = GTNRewardCalibratedSurrogateModel(**payload["model_config"])
model.load_state_dict(payload["model_state_dict"])
model.eval()

For reward-aware inference, call model.predict_with_reward(...) with graph tensors and a target_gamma tensor. The returned dictionary contains:

prediction          # [efficiency, Vout]
analytic_reward    # eta * fan_delta(Vout, target_gamma)
reward             # calibrated reward, clamped to [0, 1]
reward_residual    # learned residual added to analytic_reward
reward_features    # reward-head input features derived from prediction and target_gamma

Training Summary

• Stage: gtn_base_v6_0_reward_calibrated
• Dataset: dataset/gtn_dataset_5comp_v2_corrected_5000c_cleaned.jsonl
• Base rows: 57024
• Expanded reward samples: 741312
• Topology split counts: train 564, validation 70, test 70
• Row split counts: train 593892, validation 73710, test 73710
• Topology overlap counts: train/val 0, train/test 0, val/test 0
• Completed epochs: 80
• Best epoch: 78
• Frozen v5 surrogate base: True
• Initialization checkpoint: artifacts/gtn_base_v5_toposplit_cleaned/gtn_surrogate_model.pt
• Target gammas: [-3.0, -2.5, -2.0, -1.5, -1.0, -0.5, 0.25, 0.5, 0.75, 1.5, 2.0, 2.5, 3.0]

Reward Calibration Objective

The v6.0 model was initialized from gtn-base-v5 and trained with the surrogate base frozen. The newly added reward calibration head predicts a residual on top of the analytic fan reward:

calibrated_reward = clamp(analytic_reward + residual, 0, 1)

The training objective used reward supervision with a small residual regularization term:

loss = lambda_reward * reward_loss + lambda_residual * residual_regularization + lambda_phys * physical_prediction_loss

Training coefficients from the manifest:

{
  "lambda_phys": 0.0,
  "lambda_residual": 0.01,
  "lambda_reward": 1.0,
  "warmup_epochs": 0
}

Metrics

Final Validation Metrics

{
  "val_analytic_reward_mae": 0.006970888003706932,
  "val_analytic_reward_mae_ge_0p1": 0.09695614129304886,
  "val_analytic_reward_mae_ge_0p5": 0.11839184165000916,
  "val_analytic_reward_rmse": 0.03858204558491707,
  "val_mse": 70.92378234863281,
  "val_reward_mae": 0.004141189623624086,
  "val_reward_mae_ge_0p1": 0.04318135604262352,
  "val_reward_mae_ge_0p5": 0.05086888372898102,
  "val_reward_rmse": 0.026247916743159294,
  "val_rse_efficiency": 0.05786222591996193,
  "val_rse_mean": 0.034220974426716566,
  "val_rse_vout": 0.010579722933471203
}

Final Test Metrics

{
  "test_analytic_reward_mae": 0.008163928054273129,
  "test_analytic_reward_mae_ge_0p1": 0.11714919656515121,
  "test_analytic_reward_mae_ge_0p5": 0.12885914742946625,
  "test_analytic_reward_rmse": 0.04117625579237938,
  "test_mse": 36.0149040222168,
  "test_reward_mae": 0.004424168728291988,
  "test_reward_mae_ge_0p1": 0.05401118844747543,
  "test_reward_mae_ge_0p5": 0.061262767761945724,
  "test_reward_rmse": 0.021638687700033188,
  "test_rse_efficiency": 0.05090036243200302,
  "test_rse_mean": 0.027539070695638657,
  "test_rse_vout": 0.004177778959274292
}

The calibrated reward head reduces reward error relative to the analytic reward baseline in the held-out topology split. For example, on the test split:

• analytic reward RMSE: 0.04117625579237938
• calibrated reward RMSE: 0.021638687700033188
• analytic reward MAE: 0.008163928054273129
• calibrated reward MAE: 0.004424168728291988

Files

• gtn_reward_calibrated_surrogate_model.pt: PyTorch checkpoint with model config, state dict, target names, target gammas, and training metadata.
• manifest.json: Training summary and final validation/test metrics.
• training_history.jsonl: Per-epoch training and validation history.
• README.md: This model card.

Generated from local artifact mapping:

{
  "manifest": "artifacts/gtn_base_v6_0_reward_calibrated/manifest.json",
  "model": "artifacts/gtn_base_v6_0_reward_calibrated/gtn_reward_calibrated_surrogate_model.pt",
  "test_predictions": "artifacts/gtn_base_v6_0_reward_calibrated/test_predictions.jsonl",
  "training_history": "artifacts/gtn_base_v6_0_reward_calibrated/training_history.jsonl",
  "validation_predictions": "artifacts/gtn_base_v6_0_reward_calibrated/validation_predictions.jsonl"
}