Dual-Head Physics-Informed Graph DT

• The paper introduces DH-PGDT, which combines a causal transformer, subgoal guidance, and graph-feasible constraint enforcement to overcome limitations of DRL and standard Decision Transformers.
• The architecture features a dual-head design where the Guidance Head predicts intermediate subgoals and the Action Head generates control actions, enabling scalability in few-shot and zero-shot scenarios.
• Empirical results on IEEE test systems demonstrate that DH-PGDT achieves optimal restoration with zero variance, outperforming benchmarks like PPO and A2C in dynamic, constrained environments.

Dual-Head Physics-informed Graph Decision Transformer (DH-PGDT) is a sequence-modeling architecture for distribution system restoration (DSR) under uncertainty that combines a GPT-style causal transformer with explicit physical and topological constraint handling. Introduced in "Dual-Head Physics-Informed Graph Decision Transformer for Distribution System Restoration" (Zhao et al., 8 Aug 2025), it is designed to address limitations attributed in the paper to conventional deep reinforcement learning (DRL) and to standard Decision Transformers (DTs), particularly the data-intensive nature of DRL, reliance on the Markov Decision Process (MDP) assumption, and the dependence of DTs on return-to-go (RTG) cloning. The model integrates a dual-head transformer, subgoal-based guidance, and an operational constraint-aware graph reasoning module so that action generation is guided by intermediate restoration targets and filtered through graph-feasible operational constraints.

1. Problem setting and conceptual motivation

The problem domain is DSR, where the objective is to recover service after disruptions in a distribution network while respecting operational constraints. The paper situates DH-PGDT against two families of methods. First, DRL technologies are described as showing "great potential" for DSR under uncertainty, but their "data-intensive nature" and "reliance on the Markov Decision Process (MDP) assumption" are identified as limiting factors when the task requires "long-term temporal dependencies" or "few-shot and zero-shot decision making" (Zhao et al., 8 Aug 2025). Second, emerging DTs are presented as a promising sequence-modeling alternative for DRL tasks, yet their "reliance on return-to-go (RTG) cloning and limited generalization capacity" is described as restricting effectiveness in dynamic power system environments.

DH-PGDT is proposed to address those issues by combining "physical modeling, structural reasoning, and subgoal-based guidance." In the formulation summarized in the paper, the architecture is intended to support scalable and robust DSR, including "zero-shot or few-shot scenarios," by separating intermediate-goal prediction from control-action prediction and by enforcing graph-feasible actions online (Zhao et al., 8 Aug 2025).

A useful way to interpret the method is as a DT variant that weakens strict per-step RTG conditioning. The paper explicitly attributes generalization to two design choices: subgoal guidance, which "decouples the model from strict per-step RTG conditioning," and the graph-mask module, which enforces "physics/operational constraints in-line." This suggests that the model is not merely a sequence imitator over state-action-return tuples, but a constrained sequence planner over restoration trajectories.

2. Architectural organization

DH-PGDT consists of three major blocks: an encoder, a GPT-style causal transformer body with two parallel output heads, and a decoder (Zhao et al., 8 Aug 2025). The encoder embeds raw state, action, and physics-informed RTG tokens into a joint token space. The transformer body is causal and uses two heads in parallel: a Guidance Head for subgoal vector regression and an Action Head for control-action logits. The decoder maps the Action Head logits to switch-activation confidence probabilities.

Block	Role	Details
Encoder	Tokenization	Embeds state, action, and RTG tokens
Transformer body	Sequence modeling	GPT-style causal transformer with Guidance Head and Action Head
Decoder	Action realization	Maps action logits to switch-activation confidence probabilities

The token embeddings are defined by separate linear maps for states, actions, and RTG values:

$$\mathbf{e}_t^{s} = W_s\,\mathbf{s}_t + b_s,\quad \mathbf{e}_t^{a} = W_a\,\mathbf{a}_t + b_a,\quad \mathbf{e}_t^{R} = W_R\,\hat R_t + b_R.$$

Causal positional embeddings are then added "in the usual way," and the transformer operates with standard GPT-style causal self-attention. After $L$ stacked causal attention-plus-MLP blocks, the final hidden sequence $H^{(L)}$ is dispatched to two separate linear heads:

$$[\hat{\mathbf g}_1,\dots,\hat{\mathbf g}_q] = H^{(L)}W_g + b_g,$$

$\tilde{\mathbf a}_t = H_t^{(L)}W_a' + b_a'.$

The Guidance Head predicts a block of $q$ subgoal vectors, each of dimension $|\mathcal S|$, whereas the Action Head emits logits over all switch actions. In the abstract, the Guidance Head is described as generating subgoal representations, and the Action Head is described as using these subgoals to generate actions "independently of RTG" (Zhao et al., 8 Aug 2025). That separation is central to the architecture’s claimed departure from conventional RTG-conditioned DT formulations.

3. Subgoal representation and dual-head sequence design

The subgoal mechanism is the defining feature of the Guidance Head. The paper specifies a fixed number $q$ of intermediate subgoals. If there are $N$ total node cells, then the $n$-th subgoal corresponds to the state when $L$0 cells are energized. The associated time index is denoted $L$1. If a trajectory ends prematurely, the remaining subgoals are set equal to the terminal state $L$2 (Zhao et al., 8 Aug 2025).

The Guidance Head consumes a sparsely sampled, subgoal-focused trajectory:

$L$3

where $L$4 is the final fully-energized goal state. Training of the Guidance Head uses mean squared error against the ground-truth subgoals:

$L$5

The Action Head uses a different input construction. At training time, a contiguous length-$L$6 window of the full trajectory is sampled and augmented with the final goal $L$7 and the $L$8 subgoal offsets $L$9 at the front:

$H^{(L)}$0

More generally, for time $H^{(L)}$1:

$H^{(L)}$2

This dual sequence design is significant because the two heads are not merely parallel predictors over the same token stream. Instead, the Guidance Head learns intermediate restoration structure from sparsified subgoal-centric trajectories, while the Action Head conditions on the goal and subgoal offsets when generating the next control decision. A plausible implication is that the model can reuse high-level restoration structure even when the exact trajectory statistics differ from the training set.

4. Physics-informed graph reasoning and constraint enforcement

DH-PGDT integrates an online graph reasoning module to enforce operational constraints, specifically "no loops, no revisits" (Zhao et al., 8 Aug 2025). At step $H^{(L)}$3, the system forms an undirected graph $H^{(L)}$4 whose nodes are the node cells. The adjacency matrix $H^{(L)}$5 encodes connectivity among energized cells and candidate neighbors, while the node-feature matrix $H^{(L)}$6 is the state vector $H^{(L)}$7 stacked.

The constraint mechanism is represented by a mask function

$H^{(L)}$8

with $H^{(L)}$9 if activating switch $$[\hat{\mathbf g}_1,\dots,\hat{\mathbf g}_q] = H^{(L)}W_g + b_g,$$0 would create a loop or revisit. The paper states that, in practice, $$[\hat{\mathbf g}_1,\dots,\hat{\mathbf g}_q] = H^{(L)}W_g + b_g,$$1 is implemented by checking connectivity and prior-visit flags. The Action Head logits are then filtered by the operational-constraint mask and normalized:

$$[\hat{\mathbf g}_1,\dots,\hat{\mathbf g}_q] = H^{(L)}W_g + b_g,$$2

The paper also describes the graph reasoning module as "operational constraint-aware" and as generating a "confidence-weighted action vector for refining DT trajectories" (Zhao et al., 8 Aug 2025). In functional terms, the mask ensures that the learned policy produces only graph-feasible moves. This is important for interpreting the reported generalization behavior: the model does not rely solely on learned action preferences, because the action distribution is constrained online by the current topology and operational state.

A common misconception in reading transformer-based control models is to treat feasibility as a post-processing concern. In DH-PGDT, feasibility is internal to the action-generation pathway through $$[\hat{\mathbf g}_1,\dots,\hat{\mathbf g}_q] = H^{(L)}W_g + b_g,$$3, rather than being an external repair step. The paper’s formulation therefore ties policy validity directly to topology-aware masking.

5. Training pipeline and optimization objective

The training process begins with offline trajectories $$[\hat{\mathbf g}_1,\dots,\hat{\mathbf g}_q] = H^{(L)}W_g + b_g,$$4 collected "via random walks or expert logs" (Zhao et al., 8 Aug 2025). For each trajectory, the method extracts subgoal-focused subsequences $$[\hat{\mathbf g}_1,\dots,\hat{\mathbf g}_q] = H^{(L)}W_g + b_g,$$5 and constructs a per-timestep modified sequence

$$[\hat{\mathbf g}_1,\dots,\hat{\mathbf g}_q] = H^{(L)}W_g + b_g,$$6

The transformed dataset is then

$$[\hat{\mathbf g}_1,\dots,\hat{\mathbf g}_q] = H^{(L)}W_g + b_g,$$7

Optimization is joint over the two heads. Minibatches of $$[\hat{\mathbf g}_1,\dots,\hat{\mathbf g}_q] = H^{(L)}W_g + b_g,$$8 sequences of length $$[\hat{\mathbf g}_1,\dots,\hat{\mathbf g}_q] = H^{(L)}W_g + b_g,$$9 are sampled, run through the transformer, and updated using

$\tilde{\mathbf a}_t = H_t^{(L)}W_a' + b_a'.$0

where the Action Head loss is teacher-forced cross-entropy against the recorded expert switch action:

$\tilde{\mathbf a}_t = H_t^{(L)}W_a' + b_a'.$1

The algorithmic loop is given in the paper as follows: initialize model parameters $\tilde{\mathbf a}_t = H_t^{(L)}W_a' + b_a'.$2, goal $\tilde{\mathbf a}_t = H_t^{(L)}W_a' + b_a'.$3, and dataset $\tilde{\mathbf a}_t = H_t^{(L)}W_a' + b_a'.$4; for $\tilde{\mathbf a}_t = H_t^{(L)}W_a' + b_a'.$5 to $\tilde{\mathbf a}_t = H_t^{(L)}W_a' + b_a'.$6 epochs, sample minibatch $\tilde{\mathbf a}_t = H_t^{(L)}W_a' + b_a'.$7 of $\tilde{\mathbf a}_t = H_t^{(L)}W_a' + b_a'.$8 sequences $\tilde{\mathbf a}_t = H_t^{(L)}W_a' + b_a'.$9; compute guidance outputs $q$0 and action logits $q$1; compute $q$2 and $q$3; and update

$q$4

The optimizer is "Adam or a similar optimizer," and "no extra regularization terms were introduced beyond standard weight decay" (Zhao et al., 8 Aug 2025).

This training design makes the model simultaneously a subgoal regressor and an action imitator. The paper’s stated rationale is that the joint loss drives both heads together, while the subgoal pathway reduces dependence on strict RTG conditioning. This suggests that DH-PGDT treats restoration not just as token prediction over past returns, but as coordinated prediction over goal structure and feasible control.

6. Empirical results, zero-shot behavior, and reported scope

The reported experiments use OpenDSS for "AC power-flow and constraint checking" on modified IEEE 13-node, 123-node, and Iowa 240-node systems; the Iowa 240-node case includes DGs, substations, and time-varying loads, and the dynamic test uses time horizon $q$5 (Zhao et al., 8 Aug 2025). Baselines are the Decision Transformer "PIDT" from prior work, PPO, and A2C. The reported metrics are average return, standard deviation of return, average power restoration (APR), standard deviation (SDPR), and the number of optimal solutions over 50 trials.

The quantitative summaries reported in the paper are specific. On the IEEE 13-node system, DH-PGDT achieves mean return $q$6 with $q$7 standard deviation and $q$8 optimal solutions, whereas PPO attains $q$9 optimal solutions and A2C attains $|\mathcal S|$0. On the IEEE 123-node system, DH-PGDT with $|\mathcal S|$1 or $|\mathcal S|$2 yields $|\mathcal S|$3 kW, $|\mathcal S|$4 standard deviation, and $|\mathcal S|$5 optimal solutions, whereas PIDT attains only $|\mathcal S|$6. On the Iowa 240-node zero-shot test, DH-PGDT with $|\mathcal S|$7 yields $|\mathcal S|$8 kW, $|\mathcal S|$9 standard deviation, and $q$0 optimal solutions, while PPO and A2C "fail to reliably satisfy ramp-rate constraints" (Zhao et al., 8 Aug 2025).

The paper also reports that Figures 5, 8, and 11 show DH-PGDT converging faster and to higher return than PPO and A2C, and exhibiting "no variance once converged." An ablation on the hyperparameter $q$1 shows a trade-off: more subgoals "slows convergence slightly but does not harm final performance" (Zhao et al., 8 Aug 2025).

The zero-shot analysis is central to the paper’s claims. In a dynamic-load Iowa 240 test feeder scenario, the model is trained on static loads and tested directly on unseen time-varying loads; DH-PGDT reportedly obtains $q$2 optimal solutions, while PPO and A2C fail to generalize. In switch-contingency scenarios on the 240-node feeder, arbitrary switch outages produce new island shapes; the paper states that the Guidance Head recomputes subgoals and the mask enforces feasibility, and that both scenarios recover full restoration with no fine-tuning (Zhao et al., 8 Aug 2025).

The reported scope is broader than DSR alone. Although the work is explicitly focused on DSR, the paper states that the underlying computing model of the proposed PGDT is "broadly applicable to sequential decision making across various power system operations and other complex engineering domains" (Zhao et al., 8 Aug 2025). This should be read as a scope claim about the computing model rather than as an experimentally validated result beyond DSR.