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.

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 -> 576) (embed_state): Linear(18 -> 576) (embed_action): Linear(9 -> 576) (embed_timestep): Embedding(100000 -> 576) (embed_ln): RMSNorm(576) (blocks): 8x ModernBlock( (norm1): RMSNorm(576) (attn): CausalSelfAttention( q_proj: Linear(576 -> 576) # 12 heads × 48 head_dim k_proj: Linear(576 -> 288) # 6 KV heads × 48 head_dim v_proj: Linear(576 -> 288) # 6 KV heads × 48 head_dim out_proj: Linear(576 -> 576) qk_norm: RMSNorm(48) per head ✅ n_rep: 2 (each KV head serves 2 Q heads) ) (norm2): RMSNorm(576) (ffn): SwiGLU(576 -> 2304 -> 576, dropout=0.15) ) (ln_f): RMSNorm(576) (pred_act): Linear(576 -> 9) -> Tanh [tied with embed_act weights] (pred_state): Linear(576 -> 18) [tied with embed_state weights] (pred_return): Linear(576 -> 1) [tied with embed_return weights] )

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.

Usage

Loading the Model

import torch
from huggingface_hub import hf_hub_download

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

# Load checkpoint
checkpoint = torch.load(model_path, map_location="cpu")
state_dict = checkpoint["model_state_dict"]
config = checkpoint.get("config", {})