Transformer-based Offline RL Applied to Chess - Part 1: Foundations

This thesis explores how the combination of offline reinforcement learning mechanisms and the Transformer architecture can be applied to chess. The results show that incorporating a reward signal as model input increases the model’s ability to learn from suboptimal data. Furthermore, it is demonstrated that adding such a signal enables the model to adjust playing strength dynamically during inference. This first part covers the theoretical background: the Transformer architecture, online and offline reinforcement learning, the architectures used in the experiments, and the research gap that follows from them.

About this paper

Master Thesis for the Master of Science - Artificial Intelligence at IU International University.

The original LaTeX source, the bibliography, the full code, the data preparation and the compiled PDF are available on GitLab: gitlab.com/iu-msc-ai/transformer-based-offline-reinforcement-learning-applied-to-chess

This is part 1 of 4:

1. Foundations (this part)
2. Experimental setup
3. Results - Chess-Transformer baseline and base UDRL
4. Results - reward targets, manipulation and conclusion

Introduction

Chess is at times considered “a proving ground for artificial intelligence”.¹ While chess algorithms reached a superhuman level already in 1996-1997 with IBM’s Deep Blue defeating the then world champion Garry Kasparov and the currently most sophisticated chess engine “Stockfish” reaching an Elo² of 3.500,³ the goal of this thesis is not to create a superhuman chess engine, but to analyze the effect of specific architectural decisions on a model’s performance and decision-making process. These architectural decisions are mainly based on reinforcement learning (RL) principles and the Transformer architecture.

RL is a type of machine learning where an agent learns to make decisions by interacting with an environment. In online reinforcement learning, the agent learns from its own experiences while interacting with the environment, to then improve its performance, measured by a reward signal. Offline reinforcement learning (ORL) involves training the agent on a dataset of past experiences without further interaction with the environment. Past experiences mean that these experiences are at least not within the same training episode (more on that later). This approach is often used when direct interactions with the environment are not feasible or when the environment is too complex to model accurately. An example of this would be the base training of a self-driving car, which should not interact with the real world during (base) training due to safety concerns for example the passenger or pedestrians. Furthermore, in this example the real world is too complex to model accurately for completely simulated training.⁴ Another field where ORL is useful is when there is already training data available in abundance, which is the case in chess. The goal in this case is to save computationally expensive interactions with the environment, by training on the already available data instead of generating new experiences. The major disadvantage of ORL is that the agent can not explore new possible solutions while training, as it is limited to the data provided. This can lead to suboptimal policies, as the agent may not be able to generalize well to new situations (state-action distributions) that were not present in the training data, called out-of-distribution (OOD) data.

For any RL algorithm a reward signal is needed to guide the learning process. Building the reward signal is one of the major challenges in RL, as it needs to be designed in a way that encourages the agent to learn the desired behavior. In chess this can be done in multiple ways. The simplest would be the game outcome as reward signal, with +1 for a win, 0 for a draw and -1 for a loss. This would be a very sparse reward signal though, as it is only given at the end of a game. Sparse reward signals are still a problem for RL algorithms, as the agent gets very little feedback on its actions during an episode. A more dense reward signal could be the change in material balance after each move, which would provide more frequent feedback to the agent. Another approach would be to rate the position the player is in before and after each move. In chess this can be done via a chess engine, which analyzes and evaluates chess positions and provides a score. The currently most widely used engine in chess is Stockfish, which can play at a superhuman level and therefore is able to provide accurate evaluations of positions. More detail on which reward signal was used in which case is going to be provided later.

In classical RL the reward is optimized for by the policy, while the state of the environment and the executed action is used as input for the policy. In recent years new architectures have been proposed, which turn this around, with the reward being an input token and the action being the target to optimize for. This is a core concept of Upside Down Reinforcement Learning (UDRL), introduced by Schmidhuber et al.⁵ This brings RL and (semi-) supervised learning (SL) closer together, as the model is trained in a supervised way and therefore can even use high throughput architectures like Transformers. The Transformer architecture has been applied very successfully in natural language processing tasks. For RL it has not been applied as widely yet. A possible reason is that most of the artificial and benchmark-focused problems in RL are not complex enough to take full advantage of the attention mechanism.

The Decision Transformer (DT) architecture, introduced by Chen et al. in 2021,⁶ builds upon the UDRL architecture and adds such a Transformer architecture to build an autoregressive mechanism with a return-to-go reward, to model a sequential decision-making process. This thesis takes pieces from both works and applies them to chess.

Chess should work well for these architectures, as it provides a large action space, has large datasets of played games publicly available to execute a training process on, which is needed for Transformer models, and upholds the Markov property under certain conditions, which is an underlying requirement for many RL algorithms. The Markov property states that “the state must include information about all aspects of the past agent-environment interaction that make a difference for the future”.⁷ In chess that is true when elements like castling rights or the move count for the fifty-move rule are included in the state representation.⁸ Further chess is a fully observable environment, which means that the agent has access to the complete state of the environment at each timestep. While RL has been applied to partially observable environments as well, fully observable environments are usually easier to generate good results with.

A further aspect are the attention layers, a key component of Transformers, which allow the model to focus on different parts of the input sequence, in the case of chess e.g. the board state and the reward, when making decisions. These attention layers can be used to analyze how the model makes decisions in chess, by visualizing the attention weights in the game environment, which is the chess board. This can provide insights into the model’s decision-making process and help to identify areas for improvement and experiments.

The starting point for the experiments is the code provided by the Chess-Transformer project on GitHub,⁹ which was initially released in May 2023. The Chess-Transformer project combines the Transformer architecture with the “PyChess” engine¹⁰ and provides pre-trained models, e.g. encoder-only models, and reference results. This code was augmented with new model configurations, mostly based on the UDRL architecture. The data preprocessing was built from scratch and the data loading was heavily improved, to make fast and efficient experiments possible. This code is then used to execute the experiments.¹¹ While the Chess-Transformer projects data is limited to already strong players, it was demonstrated in this thesis that the model can learn from suboptimal data when adding a reward token.

With Transformers being notorious for needing a large amount of training data and iterations, another advantage of utilizing chess as the environment is that despite its complexity it was still possible to train several Transformer-based models within a realistic cost budget and time frame, including testing different architectures and hyperparameters.

Theoretical background

This chapter establishes the theoretical foundation for the experiments conducted in this thesis. It reviews related work in the fields of RL and Transformers, while establishing their relationship to the domain of chess.

Transformer architecture

As the following experiments incorporate the Transformer architecture, a concise introduction is provided here, drawing connections to the chess domain where relevant. The Transformer architecture was introduced by Vaswani et al. in 2017¹² and was defined as a “new simple network architecture” that was based solely on attention mechanisms, without any recurrence or convolution. The Transformer architecture has been very successful in natural language processing tasks via large-scale language models, famously applied among others by OpenAI with their GPT models or Google with their Gemini models.

The original Transformer architecture comprises input tokenization, embeddings, positional encodings, and an encoder-decoder structure composed of multiple layers of self-attention, feed-forward neural networks and layer normalization. The encoder processes the input sequence and generates a set of hidden states. The decoder takes the hidden states from the encoder and the right shifted output to generate the output probabilities. The attention mechanism allows the model to focus on different parts of the input sequence when generating the output sequence. Later visualizations of the attention mechanism are displayed, to show which parts of the input data the model focuses on when making predictions. The below figure shows the key components of the Transformer architecture from the original paper.¹³

Among others Noever et al. applied the Transformer architecture to chess. They used a GPT model to predict the next move in chess. This is very comparable to the mentioned Chess-Transformer project. They used a dataset of 11.291 games with over 800.000 moves, an average Elo by the players of 1.815 and average game length of 73 moves.¹⁴ They demonstrated that with a GPT-2 like model (~774 million model parameters) the cross-entropy log loss goes below 0,1 after 30k training steps.¹⁵

Monroe et al. showed that “Transformers endowed with a sufficiently expressive position representation can match existing chess-playing models at a fraction of the computational cost”.¹ They tested two different encoder-only models. The first one has “15 encoder layers with an embedding depth of 1024, a head count of 32, and a feedforward depth of 4096, for a total of 243 million parameters” and the second one “8 encoder layers, an embedding depth of 256, a head count of 8, and a feedforward depth of 256, for a total of 6 million parameters”. The first experiment did result in an Elo of 2347 (± 10) and the second setup did result in an Elo of 2105 (± 28) respectively. In both cases learned embeddings were used.¹⁶

Ruoss et al. analyzed the planning and reasoning capabilities of Transformers in chess. They did predict action-values based on board states, by using a dataset of 10 million games which they annotated with Stockfish for state-values and the optimal action for each state. They showed that increasing the model size increased playing strength and ruled out that the model can memorize states, “due to the combinatorial explosion of chess board states”.¹⁷

As Transformer architectures can have a large number of trainable parameters, e.g. for large language models, they are able to learn complex patterns in the data. They should still not be able to memorize all possible board states and best actions in chess, due to the combinatorial explosion of chess board states. For the explosion of chess board states additional research exists. The most famous number is the Shannon number, which estimates the lower bound of the game-tree complexity of chess to be at least $10^{120}$ possible game variations.¹⁸ Later research found a 95% confidence level at $(4,822 \pm 0,028) \times 10^{44}$ for this value.¹⁹ Especially for the later experiments the training dataset would be too small to contain all relevant state-action transitions (compare data preprocessing). Further the model’s trainable parameters would be too few to memorize these transitions (compare hyperparameters). Therefore, the model needs to generalize from the training data to unseen board states during inference.

The main building blocks of the Transformer architecture are shown in the table below. These building blocks will get discussed in detail and again applied to the chess domain in the following sections.

Transformer architecture components overview - source: own depiction

Input representation

First, tokenization of the input can be needed, depending on the input type. For e.g. text input the tokenization is needed to convert the input values into a computable format. In this case the text gets split into smaller units, called tokens, which can be words, parts of words, or characters and then mapped to a numerical representation. For chess the input consists of the board state, which in itself consists of the board squares and the pieces occupying them, the current move, and possible additional metadata, like castling rights or the mentioned fifty-move rule. In chess the board state is usually represented via Forsyth-Edwards Notation (FEN) strings. A FEN string consists of six parts, which are separated by spaces. The first part describes the board state row by row, starting from the 8th rank to the 1st rank. Each piece is represented by a letter, with uppercase letters for white pieces and lowercase letters for black pieces. Empty squares are represented by numbers, which indicate the number of consecutive empty squares. The second part indicates whose turn it is to move, with “w” for white and “b” for black. The third part describes the castling rights for both players, with “K” and “Q” for white’s kingside and queenside castling rights, and “k” and “q” for black’s kingside and queenside castling rights. A “-” indicates that no castling rights are available. The fourth part indicates the en-passant target square, if applicable, or a “-” if there is no en-passant target square. The fifth part indicates the halfmove clock, which counts the number of halfmoves since the last capture or pawn move. The sixth part indicates the full move number, which counts the number of full moves in the game.

The FEN of the starting position is rnbqkbnr/pppppppp/8/8/8/8/PPPPPPPP/RNBQKBNR w KQkq - 0 1. The encoding for the board state results in a 8x8 or 1x64 array and setting the values in the array to a predefined number for each individual piece. The metadata can be encoded mostly as boolean values, with the en-passant moves as integers.²⁰ This is shown in the table below, with the FEN string, which is usually shown as one line, broken up into parts which represent a board row and then the additional metadata. The encoded part is shown in the same line as the FEN part. More details on this and how further things like a move would be encoded is shown later when the data preprocessing is discussed (compare data preprocessing).

Example of a tokenized Forsyth-Edwards Notation - source: own depiction

These numerical representations are not used directly by the Transformer. Instead, an embedding layer maps each token to a learned dense vector of fixed size. Each token is converted to its corresponding embedding vector during training and inference. During training, the embedding vectors are optimized together with the rest of the model such that they become useful for predicting the next action.

Positional information is then added to the token embeddings to provide information about the position of each token in the sequence. This is crucial for the Transformer architecture, as no other information about the order of the tokens exists. In chess the position of the board state is two-dimensional. This positional information can be either learned during training via e.g. a linear layer or can be fixed, e.g. as sinusoidal positional encoding. The latter are typically defined using sine and cosine functions of different frequencies. This has the advantage that the positional information can be computed for any sequence length, depending on the chosen frequency of the function, and the values are between -1 and 1, which makes them numerically suitable as input for the Transformer architecture. The initial introduction of the Transformer architecture stated that learned positional embeddings and fixed positional encodings had nearly the same results.²¹ For chess the sequence length would be similar for most inputs. Exceptions would be if e.g. the legal moves in a board state would be used as input, as the amount of possible legal moves can vary. This could be solved via e.g. padding the input to a fixed length.

Attention mechanism

Attention was already discussed briefly in the beginning of the chapter. Transformer models usually apply multi-head-attention, which is computed over multiple layers and heads. The layers are calculated sequentially, with the first layers usually focusing on low-level features, such as local patterns in the input sequence, while the later layers may focus on higher-level features, such as global patterns. The heads are calculated in parallel, which allows the model to focus on different parts of the input sequence at the same time. This sequential and parallel processing of the input sequence was already shown in the figure above.

As the attention mechanism introduces $O(n^2)$ complexity due to the pairwise interactions between all tokens in the input sequence, it can be computationally expensive for long sequences. To mitigate this issue, various modified attention mechanisms have been proposed. Compare e.g. Lin et al.,²² which shows different types of attention mechanisms, such as sparse attention, low-rank attention, and linearized attention. Another mechanism to reduce computational load is to save the attention weights between inference runs. This can be applied in e.g. in large language models, so the calculated attention weights can get reused for each new token generation. This moves the load to the memory instead. As the input sequences in chess are not extremely long and the models are just generating the single next action for each state, such modified attention mechanisms were not applied in the experiments.

Attention is commonly visualized as a matrix, where the row represents the query and the columns represent the keys. The attention weights are computed by taking the dot product of the query and key vectors, followed by a softmax operation to normalize the weights. Different visual shapes in these matrices can provide a better understanding of what the attention mechanism focuses on. For example, a diagonal shape, excluding the main diagonal, in the attention matrix indicates that the model is focusing on the tokens that are close to each other in the sequence, while a more uniform distribution indicates that the model is considering all tokens equally. Some examples are shown below.

In chess the attention can be mapped back to the game environment, namely the chess board, to understand which squares the model focuses on. Monroe and Chalmers did exactly that and named one advantage of the Transformer architecture “the interpretability of their attention maps”.²³ One example of such a visualization is shown below.²⁴ The later visualizations of the experiments use a similar approach as starting point (compare attention and logit visualization).

Jenner et al. discussed look-ahead capabilities for the attention mechanism in chess. They found that Leela Chess Zero (LCZero), which applies attention,²⁵ “internally represents future optimal moves and that these representations are crucial for its final output in certain board states”.²⁶

Ballout et al. used the Tic-Tac-Toe game to analyze attention. They found that their model-based on BERT struggles to attend to the correct tokens, possibly “due to the way the 2D structure is flattened into a 1D sequence, making it more challenging for the model to recognize vertical and diagonal patterns, even in short sequences”.²⁷

Additional components

Layer normalization is a technique used to normalize the inputs of each layer in the Transformer architecture. It helps to stabilize the training process and improve the convergence of the model. Dropout is a regularization technique used to prevent overfitting in neural networks. It works by randomly dropping out a fraction of the neurons during training, which forces the model to learn more robust features. During inference, dropout is turned off. Residual connections are used to connect the input of a layer to its output, which helps to mitigate the vanishing gradient problem and improve the flow of information through the network. All of these additional components were used in the experiments.

Online and offline reinforcement learning

This section discusses RL and ORL, as these form the basis for the later experiments. RL provides “mathematical formalism for learning-based control”, while the “reward function defines what an agent should do”, and the “reinforcement learning algorithm determines how to do it”.²⁸ One of the primary dimensions along which RL is categorized is whether the interaction between the agent and the environment, which creates the rewards, is possible or not in the training process. The further case is termed online learning, the other case is termed offline learning. This section is not going to discuss RL and ORL in all its facets, but focuses on the parts which are relevant for the later experiments.

Overall a RL system has four main elements:²⁹

• A policy, which defines the learning agent’s way of behaving at a given time.
• A reward signal, which defines the goal of a RL problem.
• A value function, which specifies what is good in the long run.
• A model of the environment (optional), which allows inferences to be made about how the environment will behave.

The following image shows the classical online RL process, where interaction with the environment is possible:³⁰

This dynamical system is described by a Markov Decision Process (MDP). A MDP is a “classical formalization of sequential decision making, where actions influence not just immediate rewards, but also subsequent situations, or states, and through those future rewards”.³¹ The MDP is defined by the tuple (S, A, P, R), where S is the state space, A is the action space, P is the transition probability function, and R is the reward function. Sometimes a discount factor $\gamma$ is used to weight future rewards, to ensure convergence of the learning process.

Sutton and Barto provide an example of such a MDP with a recycling robot. The figure below graphically describes the state, action, transition probabilities, new state and reward.³² It is important to mention that “the transition distribution completely describes the environments’ dynamics”.³³ As mentioned earlier, “the property that a state $s_t$ must have all the information required to infer $s_{t+1}$ after taking action $a_t$ is known as the Markov property.”³³

In the online setting the agent attempts to visit the entire state space by exploring different actions and observing the resulting states and rewards. In the offline setting it is not guaranteed that all states have been observed in the available data. Technically, this lack of coverage can be viewed analogous to missing information about the MDP’s dynamics in parts of the state-action space, making the problem similarly to a Partially Observable MDP (POMDP), as parts of the environments’ dynamics remain hidden from the learner. This is mostly disregarded in the ORL literature though, but is one of the main downsides of just learning from offline data.³³ As in chess the opponent’s actions do influence the next state as well, it could be argued that chess is a POMDP. However, as the opponent’s actions are part of the environments’ dynamics and therefore included in the transition probabilities, chess can still be modeled as a MDP. Therefore chess is treated as a MDP in this thesis.

Overall the goal of RL is to learn a policy that maximizes the expected cumulative reward over time and therefore “learn to control a dynamical system”.³⁴ Formally this is defined by the RL objective function:

$$ J(\pi) = \mathbb{E}{\tau \sim p{\pi}(\tau)} \left[ \sum_{t=0}^{H} \gamma^t r(s_t, a_t) \right] $$

The value function, which is “the total amount of reward an agent can expect to accumulate over the future”, is commonly built via dynamic programming.³⁵ The value function is defined as following:³⁶

$$ v_\pi(s) = \mathbb{E}_{\pi} \left[ G_t | S_t = s \right] $$

As mentioned, in certain scenarios’ interaction with the environment may not be possible or desired for various reasons, such as interactive data collection being too expensive or the interaction being not safe. On the other hand it could be just more efficient to train on previously collected data if e.g. a large dataset is available. This is the case in chess, where a lot of data on games played by humans with different skill levels is made available online by e.g. chess platforms like “Lichess” (compare data preprocessing). In such a case ORL can be used to train the agent on a fixed dataset of past experiences without further interaction with the environment.³⁷

In certain dimensions offline RL is comparable to off-policy RL, which is an online learning method that uses a different policy to generate the training data than the one being trained. The policy to generate data is called the behavior policy, while the policy being trained is called the target policy. The behavior policy is then updated during training based on the learned policy. In the case of ORL the policy to generate data never updates during training though, as it was fixed during data generation. An example would be using Q-learning without updating the behavior policy until the end of training. Off-policy RL is not designed to never update the behavior policy and therefore could create suboptimal results.³⁷ Therefore ORL is sometimes called “fully off-policy” to make this distinction clear.³⁸ Another term ORL is sometimes referred to is batch RL, which could lead to confusion, as online RL uses batches generated by the agent itself during training as well. In summary, that means that “in online reinforcement learning, the agent learns from its own experiences in real-time, while in ORL, the agent learns from a fixed dataset of past experiences without further interaction with the environment.”³⁹

The following image shows the difference between online and offline RL as a visual overview:⁴⁰

ORL needs to “derive a sufficient understanding of the dynamical system underlying the MDP […] entirely from a fixed dataset, and then construct a policy $π(a|s)$ that attains the largest possible cumulative reward when it is actually used to interact with the MDP”. The dataset in this case is a static dataset of transitions, $D = \{(s^i_t , a^i_t , s^i_{t+1} , r^i_t )\}$.⁴¹

This leads to the trajectory distribution $p_\pi(s, a)$, which is the distribution of state-action pairs (s, a) under the policy $\pi$. The trajectory distribution is defined as:

$$ p_\pi(s, a) = d_0(s_0) \prod_{t=0}^{H} \pi(a_t | s_t) T(s_{t+1} | s_t, a_t) $$

where $d_0(s_0)$ is the distribution of initial states, $\pi(a_t | s_t)$ is the policy, and $T(s_{t+1} | s_t, a_t)$ is the transition probability function.

As ORL relies on previously collected data to learn, OOD data can be a problem, as during inference the agent may poorly adapt to situations not available in the training data. If the data “does not contain transitions that illustrate high-reward regions of the state space, it may be impossible to discover those high-reward regions”.⁴² This can lead to suboptimal policies, as the agent may not be able to generalize well to new situations that were not present in the training data. As exploration is not possible one way to address this issue is a larger and more complete dataset. In the case of chess large amounts of training data are readily available, which can help to minimize this issue.

Prudencio et al. propose different methods to mitigate this issue. Either to restrict the learned policy to the distribution of the training data or to using conservative estimates of the value function to avoid overestimation of the expected rewards.³⁷ Limiting the action space to the states and actions present in the training data is not possible in chess, as the model needs to continue playing even during OOD situations. Levine et al. discuss this issue and argue for an upper limit of the possible error term of $O(H^2)$, where $H$ is the horizon of the task, which is the number of steps the agent can take in the environment.³⁴

Levine et al. and Prudencio et al. describe three main ways for ORL to learn, namely importance sampling, dynamic programming or model-based methods. Prudencio et al. shows this visually in the image below.⁴³ The three mentioned methods will be discussed briefly in the following sections.

Importance sampling. “Almost all off-policy methods utilize importance sampling” according to Sutton and Barto, by “weighting returns according to the relative probability of their trajectories occurring under the target and behavior policies”.⁴⁴ Levine et al. sees major challenges to apply this method in the offline setting, as importance sampling “already suffer from high variance, and this variance increases dramatically in the sequential setting, since the importance weights at successive time steps are multiplied together”.⁴⁵ Even as there exist mitigating methods for this problem, via e.g. marginalized importance sampling and therefore not using per-action importance weighting, we will not discuss this method further here.

Dynamic programming. Levine et al. describes here mainly Q-learning and actor-critic algorithms, which technically can be adjusted to the offline setting, by just “setting the number of collection steps S to zero, and initializing the buffer to be non-empty”.⁴⁶ These methods are still susceptible to distributional shifts though, which can be mitigated just via policy constraints or the estimation of the possible shift.

While state distribution shift just appears during inference, as by definition the training data does not include these possible states, action distribution shift can already appear during training, as the policy being trained can select actions which are not present in the training data. This can lead to “extrapolation error”, as the value function is estimated for state-action pairs which are not present in the training data and therefore can be wrong. While this would be mitigated by the return in an online setting, this is not possible in the offline approach. This is called the “unlearning” effect and looks similar to overtraining, but can not be mitigated by a larger dataset. This effect can just be mitigated by policy constraints.⁴⁷

Model-based methods for ORL. Levine et al. describe here mainly trajectory sampling, with model-based methods having the ability “to estimate $T(s_{t+1} |s_t , a_t )$ via a parameterized model $T_\psi (s_{t+1} |s_t , a_t )$”.⁴⁸ In trajectory sampling “sample state transitions and rewards are given by the model, and sample actions are given by the current policy”, according to Sutton and Barto. Therefore, “one simulates explicit individual trajectories and performs updates at the state or state-action pairs encountered along the way”.⁴⁹ Here they even provide chess as an example for such a model and emphasize that rather real situations should be sampled. This already aligns well with the offline setting, as the model can be trained on real situations and then used to generate new trajectories. Further this is trainable via a supervised approach according to Levine et al., which “can benefit from the same kind of ‘blessing of scale’ that has proven so effective across a range of supervised learning application areas”.⁵⁰

Offline reinforcement learning architectures for experiments

Having established the fundamentals of ORL, this section discusses two specific architectures which were used in the experiments shown later, in combination with ORL. All architectures discussed in this section optimize a supervised sequence prediction objective, while leveraging reinforcement learning concepts through reward-conditioning rather than explicit policy evaluation or improvement. The experiments primarily implement the UDRL architecture by Schmidhuber et al., while incorporating sequence modeling via a Transformer from the DT architecture by Chen et al. One of the experiments implements the DT architecture directly. The Trajectory Transformer (TT) architecture by Janner et al. was not used in the experiments, but is discussed here to provide another example of a Transformer-based ORL architecture.

Upside Down Reinforcement Learning

In June 2020, Schmidhuber et al. published a paper describing a RL architecture they called UDRL. While standard RL aims to create a policy that optimizes the expected reward, the UDRL approach tries to learn the next action given the current state and a target reward. While it still uses the tuple (S, A, P, R) of a MDP, it does not use the value function⁵¹ or action-value function to build a policy. Instead, it uses a supervised learning approach to learn a mapping from states and rewards to actions.⁵ As with all RL methods, a reward function is needed to evaluate the behavior and to be able to distinguish desired from undesired behavior by assigning a reward. How the reward function for the chess problem gets defined is discussed in detail later (compare experiments and evaluation).

Sutton and Barto state that supervised learning is not capable of fully replacing RL, among other things because “in supervised learning, the goal is to reconstruct the unknown function $(f)$ that assigns output values $(y)$ to data points $(x)$, in RL, the goal is to find the input $(x^)$ that gives the maximum reward $R(x^)$”.⁵² In other words “feedback from the environment provides error signals in SL but evaluation signals in RL”.⁵³ UDRL trains “agents to directly map commands to actions that can fulfill those commands by retrospectively interpreting past experiences as successful examples of following commands”, which leads to the cumulative reward being “an input to the agent rather than a prediction as in value-based RL” and learning being “based on optimizing a true SL objective, unlike many RL algorithms where targets are non-stationary or the data becomes ‘stale’ after a single learning step.”⁵⁴

The following visualization shows the conceptual difference between traditional RL and UDRL:⁵⁵

Because of this change supervised learning with gradient descent can be used to train the model, which is usually more stable than RL methods and fits very well with the ORL approach. The authors call it “trivial to combine UDRL and SL, since both share the same basic framework”, which will result in “imitate[ing] teacher-given trajectories”.⁵⁶

For this to work the reward is “assumed to be independent of the history of previous actions and observations”, which is a consequence of the Markov assumption. Further randomness is considered a “separate, unobservable oracle injecting extra bits of information into the observations” and “instead of learning to map expected rewards to actions […] it can learn the concrete reward’s dependence on time, and is not misled by a few lucky past experiences”.⁵⁷

In the accompanying paper to UDRL by Srivastava et al. experiments with this approach are conducted, to show “training can produce a range of rewarding behaviors for multiple episodic environments.”⁵³ They assume that their architecture “can naturally learn from delayed rewards” and “can adjust their behavior after training in response to commands”. Especially the response to commands during inference is interesting, as this means the agent can adapt its behavior based on commands it receives as reward targets. This provides an external adjustment mechanism within an episode and was tested in the experiments.

For this to work they assume that the state and actions are elements of finite sets and “episodic Markovian environments with scalar rewards”, to build a behavior function $B$. The paper shows a simple toy example with two different trajectories. Based on the desired return, in combination with the horizon, the agent prefers one over the other.

The following algorithm they derived can “approximate $B_T$ that compresses the agent’s experience, makes computation of the conditional probabilities efficient and enables generalization to unseen states or commands”.⁵⁸ For this it uses a loss function, based on the problem. In their case the loss function is cross entropy, which works for discrete actions and is used in the later experiments therefore as well.

This loss function is supposed to be minimized by a function approximator like a neural network and takes in the return of $d^r$ desired in the next $d^h$ steps. Further “B is not used to act like a typical policy, instead it induces a policy at each step”.⁵⁹

$$ \begin{aligned} & B_\tau = \underset{B}{\mathrm{argmin}} \sum_{(\tau, t_1, t_2)} L(B(a_{t_1}, s_{t_1}, d^r, d^h), a_{t_1}), \\ & \text{where } 0 \leq t_1 \leq t_2 \leq \text{len}(\tau) \text{ for } \tau \in T , d^r = \sum_{t=t_1}^{t_2} r_t \text{ and } d^h = t_2 - t_1. \end{aligned} $$

This is supposed to work in an online setting with a replay buffer, but as already stated, could work well in an offline setting too. In the replay buffer they filter on good episodes with wanted returns. A replay buffer to just take good examples with e.g. a high reward would limit the agent to just good actions though, but there are cases where that might not be wanted or, in case the reward is not being a classification in good or bad, learning negative rewards could be used for external steering. The authors even state that this is a goal of this method. In a newer paper from 2025 Ventura et al. build a policy generator for UDRL, which “learn[s] to decode input commands representing a desired expected return into command-specific weight matrices”.⁶⁰ In chess that could be an input signal for preferring specific actions, general emphasis on specific pieces like the queen, or defined strategies. Such a sophisticated additionally encoded steering signal was not implemented in the experiments, but multiple experiments conducted steering via the reward signal. This can be done by changing the reward signal for individual games, within a game (compare 2_UDRL-eval) or via manipulating the reward in the training data (compare 8_UDRL-eval-manipulated_reward).

A further relevant aspect of this approach is that the reward can be delayed. According to the paper, models trained via UDRL “retained much of its performance in this challenging setting without modification because by design, it directly assigns credit across long time horizons”.⁶¹ In chess this could help to learn strategies, where a specific move only pays off several moves later.

UDRL is the basis for the experiments. The models were trained to predict the next action, given the current state and a desired reward. Further, the experiments investigated the impact of different reward structures on the learning process. The mentioned ability to adapt behavior based on commands during inference was also tested.

Decision Transformer

With the UDRL architecture established, now an architecture which adds trajectory conditioning in combination with a Transformer architecture to this setup will be discussed. The Decision architecture is such a sequence modeling approach based on UDRL and ORL. It has been applied to various tasks, such as Atari games, OpenAI Gym environments, and robotic control tasks. Chen et al. showed that the DT can achieve state-of-the-art performance on these tasks.⁶² The DT architecture takes states, actions, and a so called return-to-go (RTG) as input into a causal Transformer, which then predicts the action for the next time step.⁶³ The method is described as being autoregressive and uses the past trajectories of states and RTGs to predict the next action. The image below shows the described process of the DT from a high level perspective.

The RTG is defined as a sum of future rewards, with a possible discount factor. This RTG formulation is conceptually similar to the summed future reward used in UDRL; however, the optional inclusion of discounting allows the RTG to provide a more sophisticated signal. In chess that could lead to better play, as delayed counter actions to a specific move by the opponent could be incorporated into the reward signal. A RTG seems like a valuable addition to the experiments and was tested compared to a standard reward signal.

The DT uses a sequence modeling objective, to “bypass the need for bootstrapping for long term credit assignment”.⁶⁴ This should avoid the so-called deadly triad problem of approximation, bootstrapping, and off-policy learning, which can lead to the value estimates becoming unbounded.⁶⁵

Hu et al. takes a distinction between feature representation and environment representation, which can be enhanced by using a Transformer. In the first case the Transformer is used to represent the state and action space, while the latter is used to represent the environment dynamics and reward modeling.⁶⁶

Similar to regular RL the DT relies on the MDP assumption described by the tuple (S, A, P, R) as well. The paper suggests a cross entropy loss function to train the model for discrete problems, which is similar to the loss function used in supervised learning. Therefore, the DT is sometimes considered a supervised learning problem during training.⁶⁷ The loss function would be applicable to the chess task as well.

Siebenborn et al. tested the influence of the Transformer part in the DT, by comparing it to a LSTM architecture and behavioral cloning on a continuous control task. They found that the Transformer had better results than the behavioral cloning policy, but not better than the LSTM model and did conclude that “advantages of DT observed in prior works may be rather due to the sequence modeling approach than to the particular choice of the prediction module”.⁶⁸ The experiment was conducted with the same amount of training steps for all three approaches. Transformer models are known to require more training steps to achieve good results though. Therefore, a longer training time for the Transformer model would have been interesting. Further they acknowledge that Transformer architecture would be better suited for tasks with discrete action spaces.⁶⁸

Trajectory Transformer

A similar approach to the DT is the TT by Janner et al., which calls ORL “One Big Sequence Modeling Problem” in the title of the paper.⁶⁹ It uses a GPT-like architecture to model the trajectory distribution of the MDP. The model is trained to predict the next state and reward, given the previous states, actions, and rewards. The authors show that the TT can achieve state-of-the-art performance on various tasks, such as Atari games and OpenAI Gym environments.

Comparable to the mentioned DT, the TT uses a trajectory approach, by which the policy is not trained by model-free or model-based methods, but rather sequence modeling. Instead of the more common way of using a trajectory optimizer they use beam search as the planning algorithm.

In this setting they use similar to the DT the RTG which “estimates are functions of the behavior policy that collected the training data and do not, in general, correspond to the values achieved by the Trajectory Transformer-derived policy”. Further they state that “a value function for an improved policy would provide a better search heuristic, though requires invoking the tools of dynamic programming”, but a RTG is sufficient in cases with a non-sparse reward.⁷⁰

Research gap

With the theoretical background established regarding UDRL and the Transformer architecture, here the research gap will be discussed and the research questions defined. It was already shown that the Transformer has been applied to chess in various ways and the combination of ORL, UDRL-like methods and Transformers has been explored in the DT architecture. However, a direct application of UDRL, in combination with ORL and the Transformer architecture has not been investigated in a complex and high-branching environment like chess yet.

Therefore, the main focus of the later experiments was on applying UDRL together with a Transformer-based architecture and offline data. The experiments were evaluated with respect to a potential increase of playing strength compared to a supervised learning baseline and the ability to externally steer the model via the reward signal during inference. Additionally, a DT architecture was implemented and compared to a standard UDRL architecture to analyze the impact of trajectory modeling on the mentioned dimensions.

Different reward setups were tested, to see how they impact the learning process. This includes a RTG as reward signal, similar to how it is used in the DT architecture. The steering capability was analyzed for all implemented architectures and different reward setups. Furthermore, the attention and resulting logits were visualized to understand how the model makes decisions and whether certain models yield clearer attention patterns. Which exact experiments were conducted and how the experiments were designed in detail gets discussed in the next chapter.

In summary the research gap can be defined as two main research questions:

1. Does the UDRL architecture increase the playing strength for specific reward targets, compared to a supervised learning baseline?
2. Does the UDRL architecture allow to externally steer the models behavior during inference, by changing the desired reward target?

Continue with Part 2: Experimental setup, which introduces the eleven trained models, the evaluation against Stockfish, the attention and logit visualizations, the codebase, the data preprocessing and the reward distributions.

References