Graph Reinforcement Learning as an MDP over Tensor States
Reinforcement learning is usually taught over vector states — a cart's position, a game screen. But many graph problems are naturally sequential decision problems: colour the nodes one at a time, grow a graph edge by edge, walk a knowledge graph to answer a query. These are Markov decision processes whose state is a graph. TGraphX provides an Experimental graph-RL subsystem built on exactly this framing, and seeing the MDP structure makes the API legible. This note lays out the formalism and maps it onto tgraphx/rl.
A maturity note first: graph RL is labelled Experimental — correct foundations, but an API and semantics that may evolve between minor releases. Treat it as a research scaffold, pin your version, and read the introductory graph-RL article for context.
The MDP, with graphs as states
An MDP is the tuple (S, A, P, R, γ): states, actions, transition dynamics, reward, and discount factor. The agent picks actions to maximise expected discounted return E[ Σ_t γ^t r_t ]. In graph RL the pieces specialise:
S : a graph (nodes, edges, possibly tensor-valued features)
A : a graph operation — pick a node, add an edge, choose a colour, take a step
P : how the graph changes after an action (often deterministic edits)
R : a task reward — cut size, validity, path success
Because the state is a graph, the policy and value functions cannot be plain MLPs — they must read graph structure. That single requirement drives the whole design: every learnable component is a graph network.
Environments
tgraphx/rl ships a set of environments, each a small MDP over graphs with the familiar reset() / step(action) loop and an action_space:
| Environment | Task | Reward signal |
|---|---|---|
GraphColoringEnv |
proper graph colouring | fewer colours / valid colouring |
MaxCutEnv |
maximum cut | cut weight |
ShortestPathEnv |
reach a target node | path cost |
GraphNavigationEnv |
navigate the graph | reaching goals |
KGPathReasoningEnv |
multi-hop KG reasoning | answer correctness |
GraphGenerationEnv |
build a graph | generation reward |
There are continuous variants too (ContinuousNavigationEnv, ContinuousGraphEditEnv) for continuous action spaces. Several of these are classic combinatorial-optimisation problems (MaxCut, colouring) recast as learning problems — a well-studied use of RL.
Graph-aware networks and action masking
The agent's brain is built from graph networks: GraphPolicyNetwork, GraphValueNetwork, GraphQNetwork, and GraphActorCriticNetwork. Each consumes the graph state and outputs action logits or values. The detail that matters most in practice is MaskedCategoricalPolicy: in graph problems many actions are invalid in a given state (you cannot recolour a fixed node, cannot add an existing edge), and a masked categorical policy zeroes the probability of illegal actions before sampling. Without masking, an agent wastes enormous effort learning not to make impossible moves; with it, the action distribution is restricted to the legal set by construction.
The algorithm registry
Rather than hard-coding one method, TGraphX exposes a registry. list_graph_rl_algorithms() (in tgraphx/rl/high_level_api.py) returns the available algorithms, split into discrete and continuous action types:
import tgraphx as tgx
algos = tgx.list_graph_rl_algorithms() # name -> metadata
# discrete: reinforce, a2c, dqn, double_dqn, ppo, actor_critic
# continuous: ddpg, td3, sac
result = tgx.train_graph_rl(env, algorithm="ppo") # one-call entry point
The discrete agents include policy-gradient methods (REINFORCE, A2C, PPO), value-based methods (DQN, Double DQN), and actor-critic; the continuous side adds DDPG, TD3, and SAC; and there are RandomPolicy / GreedyPolicy no-learn baselines. The README summarises this as "13 algorithms," and the baselines are usefully labelled Beta even though the learning agents are Experimental — a sanity baseline you can trust is exactly what you want when debugging an RL run.
Honest framing
This is the section that matters most for an Experimental subsystem. Reinforcement learning is notoriously sample-inefficient and seed-sensitive; results vary across runs and require careful baselines (which is why the Random/Greedy policies are there). TGraphX's graph-RL is a research scaffold for graph-structured MDPs, not a tuned solver — and for non-graph RL, mature libraries like Stable-Baselines3 or RLlib remain the right tools, as the README's positioning makes explicit. Report seeds, compare against the built-in baselines, and treat any single run as a sample, not a verdict.
Why baselines are non-negotiable
In reinforcement learning a single training curve proves almost nothing — variance across seeds is large, and a method can look good or bad by luck. This is why TGraphX ships RandomPolicy and GreedyPolicy as Beta, environment-agnostic baselines: they give you a floor to beat. A learned agent that does not clearly beat a greedy heuristic on a graph problem is not yet demonstrating learning, no matter how its loss curve looks.
The disciplined protocol is to run several seeds for both the agent and the baselines, report the spread (not just the best run), and compare against the greedy and random floors on the same environment. For combinatorial problems like MaxCut, a strong greedy baseline can be hard to beat, which is exactly the kind of honest check that keeps an Experimental subsystem from overstating itself. Treat any single run as one sample, profile the wall-clock cost from the source, and let the baseline comparison — not the raw reward — decide whether the agent is actually learning.
Related guides
Conclusion
Graph RL is cleanest when you see it as an MDP whose states are graphs: that framing dictates graph-aware policy/value networks and action masking for legal moves. TGraphX implements this in tgraphx/rl with a registry of discrete and continuous algorithms and trustworthy no-learn baselines. It is Experimental and RL is hard, so lean on the baselines and report your seeds — but the MDP-over-graphs structure is a sound and well-grounded way to attack sequential graph problems.