EnergyDecision-DT-V2: Decision Transformer for AEMO FCAS Battery Trading

Model Description

EnergyDecision-DT-V2 is a Decision Transformer model trained on simulated battery dispatch data from the Australian Energy Market Operator (AEMO) Frequency Control Ancillary Services (FCAS) market. It models optimal battery dispatch as a sequence prediction problem, conditioning on returns-to-go, observed states, and past actions to predict the next action.

The model learns to dispatch battery energy storage (charge/discharge) and bid into 8 FCAS contingency markets simultaneously, using a modern transformer architecture with Grouped-Query Attention, QK-Norm, SwiGLU activations, and weight-tied embeddings.

This is the current SOTA model in the benchmark: it beats PPO, dispatch replay, and all GRPO-tuned variants on the fair same-asset dispatch-matched benchmark without any online RL fine-tuning. See Benchmark results.

Key Features

Modern architecture: Grouped-Query Attention (6 KV heads, 12 Q heads), QK-Norm for training stability, SwiGLU FFN, RMSNorm pre-norm
Weight tying: Embedding and prediction layers share weights for parameter efficiency
Action Space (9-dim):
- Dim 0: Energy dispatch in $$[-1, 1]$$ (negative = charge, positive = discharge)
- Dims 1-8: FCAS contingency bids in $$[0, 1]$$
State Space (18-dim): Normalized market observations including prices, demand, renewables penetration, and battery state-of-charge
Context Length: 210 timesteps (looks back ~17.5 hours of 5-minute dispatch intervals)

Intended Use

This model is intended for:

Research into offline RL for energy markets
Simulation of battery trading strategies in the AEMO FCAS market
Baseline for comparing decision transformer approaches against traditional RL (Stablebaselines3 based model, Decision Transformer GRPO fine-tuned)

It is not intended for live trading without further validation, risk management, and regulatory compliance.

Training Data

Source: AEMO simulated trade dataset
Size: 86,412,124 rows after filtering (2,449,631 rows from old_rule policy excluded due to mismatched 3D action space)
Episodes: 2,401 episodes (after filtering for minimum context length)
Source policies: A2C (76.9M rows) + GRPO-DT (11.9M rows)

Dataset Schema

Column	Type	Description
`episode_id`	`i32`	Unique episode identifier
`step`	`i64`	Timestep within the episode (0-indexed, 5-minute intervals)
`norm_observation`	`list[f32]` (18-dim)	Normalized market observations
`action`	`list[f32]` (9-dim)	Battery dispatch action
`reward`	`f32`	Scalar reward from the simulated market interaction
`source_policy`	`str`	Policy that generated the episode (`a2c` or `grpo_dt`)

Preprocessing

Observations already normalized to zero mean, unit variance per feature
Rows with incorrect action dimensionality (3D instead of 9D) filtered out
Returns-to-go computed with discount factor $$\gamma = 0.95$$
Overlapping trajectory chunks with stride = context_len / 2 (105 timesteps)

Model Architecture

DecisionTransformer(
  (embed_return):  Linear(1 -> 768)
  (embed_state):   Linear(18 -> 768)
  (embed_action):  Linear(9 -> 768)
  (embed_timestep): Embedding(100000 -> 768)
  (embed_ln):      RMSNorm(768)
  (blocks): 8x ModernBlock(
    (norm1): RMSNorm(768)
    (attn): CausalSelfAttention(
      q_proj: Linear(768 -> 768)   # 12 Q heads × 64 head_dim
      k_proj: Linear(768 -> 384)   # 6 KV heads × 64 head_dim
      v_proj: Linear(768 -> 384)   # 6 KV heads × 64 head_dim
      out_proj: Linear(768 -> 768)
      qk_norm: RMSNorm(64) per head  ✅
      n_rep: 2 (each KV head serves 2 Q heads — GQA)
    )
    (norm2): RMSNorm(768)
    (ffn): SwiGLU(768 -> 3072 -> 768, dropout=0.15)
  )
  (ln_f): RMSNorm(768)
  (pred_act):   Linear(768 -> 9)  -> Tanh   [tied with embed_act weights]
  (pred_state): Linear(768 -> 18)         [tied with embed_state weights]
  (pred_return): Linear(768 -> 1)         [tied with embed_return weights]
)

Compared with the legacy energydecision-dt (8×384, standard MHA, untied weights), the v2 model uses 768-dim hidden, Grouped-Query Attention (12 Q / 6 KV heads), QK-Norm, SwiGLU, and weight tying across all three prediction heads.

Hyperparameters

Parameter	Value
Blocks	8
Hidden dim	768
Attention heads (Q)	12
KV heads (GQA)	6
Context length	210
Dropout	0.15
QK-Norm	✅ Enabled
Weight tying	✅ Enabled
State dim	18
Action dim	9
Discount factor	0.95
Return scale	2.0
Loss weights	action=0.999, state=0.002, return=0.0001

Training Procedure

Hardware: CUDA GPU (AMP mixed precision)
Optimizer: AdamW (lr=3e-5, weight_decay=1e-4)
Batch size: 128
Epochs: 3
Total training time: 2h 31m (9032.7 seconds)
Throughput: ~251 samples/sec, ~1.97 batches/sec
Gradient clipping: 1.0
Strategy: Overlapping context windows with stride = context_len / 2 = 105

Training Metrics

Epoch	Train Loss	Val Loss	Action Loss (end)	Duration
1	0.057152	0.019429	0.056399	2910.0s
2	0.015032	0.009449	0.014647	2911.3s
3	0.008868	0.007034	0.008644	2913.4s

Val loss dropped 63.8% from epoch 1 to epoch 3 (0.019429 → 0.007034), with action loss improving 84.7% (0.056399 → 0.008644). The model was still learning at epoch 3 end — further training would likely yield additional gains.

Benchmark Results (SOTA)

Evaluated on the fair same-asset dispatch-matched benchmark (every policy runs on the identical battery derived from Dalrymple North: 8 MWh / 30 MW, 3.75 C, full_fcas action space, 5-minute resolution, Q4 2024 SA1, 144 h / 1728 steps per episode) and on the broader standard cross-region surface (5 regions, medium batteries).

Model	Standard	Dispatch-matched (rtg=0.5)	Dispatch-matched (rtg=0.0)
Modern v2 pretrained (this model)	$4,630	$6,793	$10,138
Phase C GRPO (2 bat, 3 region, 144h)	$4,102	$6,445	$6,183
Legacy Phase 1 GRPO (8×384, overfit)	$1,533	$8,242	$5,451
PPO reference	$2,353	$7,757	—
Dispatch Dalrymple North	$4,660	$3,663	—

Key findings:

Modern v2 pretrained is SOTA. It achieves the highest profit on the broad standard surface ($4,630/ep) and the highest dispatch-matched profit ($10,138/ep at RTG=0.0). The architecture improvements (GQA, RMSNorm, weight tying) captured the benefits that online RL fine-tuning (GRPO) once provided the legacy model.
GRPO does not improve this model. The best GRPO variant reaches $4,102 standard / $6,445 dispatch-matched — within 5–11% of pretrained, never exceeding it.
FCAS capability comes from the offline data, not GRPO. Every DT variant earns 3–5× more FCAS revenue than the real dispatch strategy.
PPO retains a degradation edge ($310/ep vs every DT variant) due to more conservative cycling. Closing this gap while keeping FCAS revenue is the primary open problem.

RTG Calibration — The Transformer as a Tunable Controller

The DT's behaviour is adjusted at inference time via the return-to-go (RTG) prompt — no retraining. The optimal RTG differs by architecture, so it must be calibrated per model.

Modern v2 model (8×768 GQA) RTG calibration (dispatch-matched surface):

RTG	Profit/ep	FCAS/ep
0.0	$10,138	$10,068
0.5	$6,793	$6,703
1.0	$6,877	$6,101
1.5	$6,999	$6,074
2.0	$6,329	$6,092

Modern optimal: rtg_value=0.0. This is the inverse of the legacy 8×384 model (which peaks at rtg_value=0.5). The modern model internalizes the reward structure more directly and needs less prompt-based guidance. Use rtg_value=0.0 unless you have recalibrated on your own data.

Usage

RTG prompt: When evaluating, use rtg_value=0.0 (see RTG Calibration). This is the optimal prompt for the modern v2 architecture.

Installation

The model is loaded with the DecisionTransformer class from the energydecision repository:

git clone https://github.com/mrvictoru/energydecision.git
cd energydecision
pip install -r requirements.txt

Loading the Model

import torch
from huggingface_hub import hf_hub_download
from decision_transformer import DecisionTransformer  # from the energydecision repo

# Canonical model config (also shipped at configs/aemo_decision_transformer_model_kwargs_modern_v2_full_fcas.json)
model_kwargs = {
    "state_dim": 18,
    "act_dim": 9,
    "n_block": 8,
    "h_dim": 768,
    "n_heads": 12,
    "n_kv_heads": 6,      # Grouped-Query Attention
    "context_len": 210,
    "drop_p": 0.15,
    "max_timestep": 100000,
    "qk_norm": True,
    "rope_enabled": False,
    "tie_weights": True,
}

# Download weights
model_path = hf_hub_download(
    repo_id="mrvictoru/energydecision-dt-v2",
    filename="aemo_dt_fcas_model.pt",
)

model = DecisionTransformer(**model_kwargs)
model.load_from_checkpoint(model_path)   # handles the HF checkpoint format automatically
model.eval()

Inference (single step)

The DT conditions on a return-to-go (RTG) prompt, the recent state/action history, and the timestep. Below is a minimal single-step call — in practice you maintain rolling buffers of the past context_len steps as the episode unfolds (see src/decision.py::AEMOAgent).

import torch

# Prepare inputs: (batch, context_len, dim)
T = model_kwargs["context_len"]
states = torch.zeros(1, T, 18)                  # normalized 18-dim market observations
actions = torch.zeros(1, T, 9)                  # past actions (zero-padded at start)
returns_to_go = torch.full((1, T), 0.0)         # RTG prompt — use 0.0 for modern v2
timesteps = torch.arange(T).unsqueeze(0)        # 0 .. T-1

with torch.no_grad():
    action = model.get_action(states, actions, returns_to_go, timesteps)
    # action shape: (1, 9)
    #   dim 0 : energy dispatch in [-1, 1]  (negative = charge, positive = discharge)
    #   dims 1-8: FCAS bids in [0, 1]

Inference via the bundled agent (recommended)

For a full episode with automatic rolling context and RTG updates, use the repository's AEMOAgent:

from decision import AEMOAgent
from AEMOBatteryEnv import AEMOBatteryTradingEnv

env = AEMOBatteryTradingEnv(...)               # configure for full_fcas, your region/battery
agent = AEMOAgent(env, algorithm="dt", model=model, rtg_value=0.0)

obs, _ = env.reset()
done = False
while not done:
    action = agent.act(obs)
    obs, reward, terminated, truncated, info = env.step(action)
    done = terminated or truncated

Environment / observation details

The 18-dim observation and 9-dim action are defined by AEMOBatteryTradingEnv with action_mode="full_fcas". Observations are expected normalized (zero mean, unit variance per feature) as produced by AEMODataPreprocessor. Do not feed raw market prices directly.

Standard Tier Leaderboard (Jul 2026)

Oct 2024 · 5 NEM regions · 144h · medium_1c · full_fcas · best RTG per model

Model	Profit/ep	FCAS/ep	Deg/ep	Best RTG
Modern v2	$4,991	$4,836	$229	10.0
Dispatch Dalrymple North	$4,660	$2,287	$1,020	—
Forecast DT	$4,564	$3,663	$270	50.0
Phase C GRPO (mod v2)	$4,322	$2,508	$1,058	10.0
Phase 1 GRPO (legacy)	$2,678	$2,914	$384	50.0
PPO reference	$2,353	$2,192	$236	—