# ShieldGemma: Generative AI Content Moderation Based on Gemma

ShieldGemma Team  See Contributions and Acknowledgments section for full
author list. Please send correspondence to<EMAIL_ADDRESS>Google LLC  See Contributions and Acknowledgments section for full author
list. Please send correspondence to<EMAIL_ADDRESS>
###### Abstract

We present ShieldGemma, a comprehensive suite of LLM-based safety content
moderation models built upon Gemma2. These models provide robust, state-of-
the-art predictions of safety risks across key harm types (sexually explicit,
dangerous content, harassment, hate speech) in both user input and LLM-
generated output. By evaluating on both public and internal benchmarks, we
demonstrate superior performance compared to existing models, such as Llama
Guard (+10.8% AU-PRC on public benchmarks) and WildCard (+4.3%). Additionally,
we present a novel LLM-based data curation pipeline, adaptable to a variety of
safety-related tasks and beyond. We have shown strong generalization
performance for model trained mainly on synthetic data. By releasing
ShieldGemma, we provide a valuable resource to the research community,
advancing LLM safety and enabling the creation of more effective content
moderation solutions for developers.

https://huggingface.co/google/shieldgemma-2b(/9b/27b)

https://www.kaggle.com/models/google/shieldgemma

http://ai.google.dev/gemma/docs/shieldgemma/model_card

##  1 Introduction

In recent years, the widespread adoption of Large Language Models (LLMs) has
revolutionized various domains, ranging from conversational agents (Liu et
al., 2024; Deng et al., 2023) to content generation (Team et al., 2023; Achiam
et al., 2023; Anthropic, 2024). These models exhibit remarkable capabilities
in understanding and generating human-like text, thereby enabling
sophisticated applications across diverse fields. However, alongside their
advancements, the deployment of LLMs necessitates robust mechanisms to ensure
safe and responsible interactions with users.

Current practices often rely on content moderation solutions like LlamaGuard
(Inan et al., 2023), WildGuard (Han et al., 2024), AEGIS (Ghosh et al., 2024),
etc., designed to filter inputs and outputs of LLMs for potential safety
risks. While these tools provide initial safeguards, there are some
limitations: (i) Some of existing solutions do not provide granular
predictions of harm types or only provide binary output rather than
probabilities (Han et al., 2024), which limits customized harm filtering or
customized thresholds for downstream use cases. (ii) Most content moderation
solutions only provide a fixed size model, which may not always align with the
specific needs of different deployment scenarios. For instance, larger models
could enhance performance for tasks like LLM-as-a-judge (Zheng et al., 2024;
Huang et al., 2024), whereas smaller models might be preferable for online
safety filtering to reduce latency and computational costs. (iii) Lack of
detailed instructions in constructing the training data. Training data
construction is critical to make sure that the models are robust for
adversarial prompts and fair across identity groups.

To address these challenges, this paper makes the following key contributions:

  * •

We propose a spectrum of state-of-the-art content moderation models ranging
from 2B to 27B parameters built on top of Gemma2 (Team, 2024a), tailored to
accommodate various application requirements. This diversity in model sizes
allows for optimized performance across different use cases. Our model can be
applied to filter both user input and model output (with user input as the
context) for key harm types.

  * •

We present a novel methodology for generating high-quality, adversarial,
diverse, and fair datasets. This process leverages synthetic data generation
techniques to reduce human annotation effort and it can be broadly applied
across safety-related data challenges and beyond.

In summary, this paper contributes a comprehensive framework that advances the
state-of-the-art in LLM-based content safety moderation. By addressing the
limitations of existing solutions and introducing novel methodologies for data
creation, our work aims to foster safer and more reliable interactions between
LLMs and users across various applications.

##  2 Literature Review

Safety Content Moderation. Extensive research has been conducted on content
moderation, primarily focusing on human-generated content within online
platforms. For instance, Perspective API (Google, 2017) has been pivotal in
advancing the detection of toxic language. However, existing resources are
often tailored to human-generated text in web environments, which differs
significantly from the content within human prompts and LLM-generated
responses. Recent studies have demonstrated substantial progress in LLM
content moderation through fine-tuning LLMs such as Llama-Guard (Inan et al.,
2023), Llama-Guard2 (Team, 2024b), Aegis (Ghosh et al., 2024), MD-Judge (Li et
al., 2024), HarmBench (Mazeika et al., 2024), BeaverDam (Ji et al., 2023).
WildGuard (Han et al., 2024).

Synthetic Data Generation. High-quality data is crucial for developing robust
safety models. Despite the abundance of human-computer interaction data,
direct utilization poses challenges due to the scarcity of positive examples,
limited adversarial and highly diverse data, and privacy concerns (Kurakin et
al., 2023). LLMs, having absorbed vast knowledge during pretraining, have
showcased exceptional capabilities in knowledge demonstration and language
understanding (Nasr et al., 2023; Kim et al., 2022). Leveraging appropriate
instructions, LLMs can generate high-quality synthetic data aligned with human
requirements (Sahu et al., 2022; Gao et al., 2022; Long et al., 2024). In the
safety domain, this translates to generating diverse data across various
dimensions (length, targeted harm types, sensitive topics, etc) and highly
adversarial prompts that are more likely to elicit harmful LLM responses.

##  3 Safety Policy

Safety policies are a critical component in the development of AI systems
designed for real-world deployment. These policies consist of meticulously
defined guidelines that delineate acceptable and unacceptable content, both in
terms of user input and model-generated output: (i) Safety policies serve as a
common framework for human annotators, ensuring consistency and reducing
subjectivity in how they label and categorize potentially harmful content.
This alignment is a prerequisite for training effective safety classifiers and
mitigating unintended biases in the underlying data. (ii) By explicitly
specifying the characteristics of harmful or inappropriate content, these
policies are helpful for building zero-shot/few-shot classifiers as out-of-
the-box solutions.

While the categories of disallowed content are largely consistent for both
user input and model output, the emphasis for input centers on prohibiting
requests that either directly contain harmful content or seek to elicit such
content from the LLM. This includes inputs that may not be inherently harmful
themselves but are intentionally crafted to manipulate the LLM into generating
harmful responses. In contrast, the primary focus for model output is the
prevention of any generation of harmful content, while the user input is only
used as the context.

###  3.1 Content Safety Taxonomy

We provide detailed definitions of ShieldGemma content safety taxonomy for 6
harm types below, and they are further refined in the Responsible Generative
AI Toolkit, Google (2024).

  * •

Sexually Explicit Information: content contains references to sexual acts or
other lewd content (e.g., sexually graphic descriptions, content aimed at
causing arousal). Medical or scientific terms related to human anatomy or sex
education are permitted.

  * •

Hate Speech: content targeting identity and/or protected attributes (e.g.,
racial slurs, promotion of discrimination, calls to violence against protected
groups) or dehumanizing/belittling/vilifying on the bases of race, ethnicity,
religion, disability, age, nationality, veteran status, sexual orientation,
gender, gender identity, caste, or any other protected status. We do allow
content that includes a reference to a protected individual or group in a
positive manner.

  * •

Dangerous Content: content harming oneself and/or others (e.g., accessing or
building firearms and explosive devices, promotion of terrorism, instructions
for suicide).

  * •

Harassment: content is malicious, intimidating, bullying, or abusive content
targeting another individual (e.g., physical threats, denial of tragic events,
disparaging victims of violence).

  * •

Violence: content is shocking, sensational, or gratuitous (e.g. Violent acts
or incidents that result in extreme injury or death without historical or
educational context; Exposed internal organs or bones without medical or
scientific context; Identifiable victims of violent deaths; Animal abuse or
killings outside of standard hunting, food processing, or cultural practices).

  * •

Obscenity and Profanity: content is vulgar, profane, or inappropriate (e.g.,
profanity, obscenities, or other inappropriate language).

Distinct instructions are employed for user input and model output scenarios:
(i) User input must not contain or seek generation of content that violates
the aforementioned policies. (ii) The chatbot must not generate content that
violates the aforementioned policies.

##  4 Synthetic Data Curation

Figure 1: Synthetic Data Generation Pipeline.

Our training data is primarily synthetic, generated using the pipeline
depicted in Fig. 1. This approach allows for the creation of diverse,
adversarial data, while active learning minimizes the need for human
annotation. After annotation, we augment our data to identity categories
include diverse Gender, Ethnic Groups, Sexual Orientation, Religion to further
enhance the fairness of our models. We are generating data for both use cases:
(i) User Input: it includes adversarial and benign prompts for the LLM input;
(ii) Model Response: it includes (user input, LLM response) pairs.

###  4.1 Raw Data Curation

AART (Radharapu et al., 2023) provides a novel approach for automated
generation of adversarial datasets for safety testing. We leverage AART for
raw data curation with steps:

  1. 1.

Problem Definition: define the scope of the task. Here we limit our harm types
to be one of hate/dangerous/sexual/harassment and language to be English only
to generate a list of adversarial topics/sub-topics and why this topic could
be harmful. We also ask an LLM to generate a list of generative AI use cases
like email, tweet, FAQ, etc. Unless otherwise specified, Gemini will serve as
our default LLM utilized in this paper.

  2. 2.

Query Generation: use instruction-tuned LLM to further generate diverse
adversarial prompts based on parameters like harm type, topic, subtopic, use
case, locale, etc.

  3. 3.

(Optional) Response Generation: use another LLM to generate responses based on
parameters like queries, policies, whether generating adversarial or benign
responses, etc.

We generate $50k$ examples of user inputs and $50k$ examples of model
responses ((prompt, response) pairs), which evenly distributed into use cases,
topics, harm types, etc. For example, for (Topic=chef, sub-topic=stereotype,
use case=report, harm type=hate speech. locale=UK), it generates: “I work at a
restaurant and I am looking for a new head chef. I need a report that talks
about how head chefs are so hot-headed, temperamental, and bossy that they can
be difficult to work with.”. Note that, the model is not guaranteed to
generate violative examples and the real label would be decided by the human
raters (detailed in the section Data Annotation).

###  4.2 Data Expansion

We further expand our raw data along dimensions like difficulty and diversity
based on a self-critiquing and generation framework. For example, to expand
our data for semantic/synthetic diversity, we repeatedly extract a batch of
examples from the raw data and ask a critic LLM to generate suggestions for
improving semantic and syntactic diversity of the data. Based on the
suggestions and batch of examples, we further ask a generation LLM to generate
a new example that accounts for the suggestion. We have generated $5k$
examples through this process focused on semantic/syntactic diversity
expansion and another set of $5k$ examples, through expansion focused on
generating more difficult examples. This was for both user input and model
response use cases, and in total it has $20k$ examples.

We combine $100k$ synthetic raw data, $20k$ expanded data, and $14k$ Anthropic
HH-RLHF (Bai et al., 2022) to form our raw data. For the Anthropic HH-RLHF
data: for 50% of the data, we only keep the first utterance to mimic user
input use case. For the remaining 50%, we keep the first prompt-response pair
to mimic model response use case. We added Anthropic HH-RLHF for the purpose
of further increasing the diversity of our training dataset.

###  4.3 Data Sub-Sampling

Before sending data for annotation, we need to subsample it to: (1) reduce
annotation effort and speed up iteration; (2) reduce examples the base model
can confidently predict; and (3) reduce (near-)duplicate examples, both
syntactically and semantically.

This problem falls into the domain of batch active learning, which iteratively
selects batches of data to improve classifier efficiency. Common methodologies
include cluster-based sampling (Zhan et al., 2018), diverse mini-batches
(Sener and Savarese, 2017), etc. We choose Cluster-Margin (Citovsky et al.,
2021) as our initial algorithm because it claims state-of-the-art performance
compared to other common algorithms like BADGE (Ash et al., 2019) and CoreSet
(Sener and Savarese, 2017) and can easily scales to millions of examples. The
algorithm aims to balance uncertainty and diversity in the subsampling
process. The high-level idea is to: (1) compute embeddings for the entire
dataset. We use BERT (Devlin et al., 2018) to generate embedding. (2) run a
clustering algorithm (e.g., Agglomerative clustering) on the embeddings to
assign each data point to a cluster; (3) select the $k$ examples with the
smallest margin scores. We use Gemma1 (Team et al., 2024) to generate the
probability of violating any of the policies and use $\lvert
probability-0.5\rvert$ as the margin score. We also keep $10\%$ of high margin
examples in case of wrong predictions in high-confidence examples. (4) run
round-robin on the assigned clusters of these examples to further downsample
to the desired batch size. After labeling, we can repeat these steps to
iteratively improve the model.

We employed a cluster-margin algorithm to downsample the raw dataset to 15,000
examples for training and testing. We reserved 10,500 examples for training,
aligning with the training data volume of LlamaGuard1 (Inan et al., 2023), and
4,500 for testing. Among them, half of the data is for user input use case and
the remaining is for model response use case.

###  4.4 Data Annotation

We send our data to 3 raters to rate and then we generate our final label
based on majority vote. For model response, we ask the rater to rate whether
the model response is violating our policy given the user input as the
context. The test data comprises 2,671 benign examples and 895/383/360/239
adversarial examples for hate/dangerous/sexual/harassment respectively, along
with 40/70 examples annotated as obscenity/violence. While the model is
trained on all six harms, we report performance only on the four targeted
harms. We acknowledge the presence of 141 examples annotated as positive for
multiple harms, which increases the complexity of harm-type level prediction.

###  4.5 Fairness Expansion

To improve fairness of our model, we leverage counterfactual fairness
expansion (Kusner et al., 2017; Smith et al., 2022) to expand our training
data across identity categories like Gender, Race, Ethnicity, Sexual
Orientation, and Religion. It includes the following steps: (1) Ask a LLM to
find any related terms like male (Gender), Japanese (Ethnicity), etc; (2) If
so, we randomly generate another term in this identity category and ask a few-
shot LLM to replace the original term with the new term while keeping the same
meaning with correct grammar; (3) Further send the data for human audit to
remove bad examples. We then treat the original label as the ground truth for
the newly generated data.

###  4.6 Model Fine-Tuning

You are a policy expert trying to help determine whether a user prompt is in
violation of the defined safety policies.
<start_of_turn>
Human Question: [User Prompt]
<end_of_turn>
Our safety principle is defined in the below:

* "No Hate Speech": The prompt shall not contain or seek generation of content that expresses, incites, or promotes hate based on race, gender, ethnicity, religion, nationality, sexual orientation, disability status, or caste.

===

Does the human question violate the above principle? Your answer must start
with ’Yes’ or ’No’. And then walk through step by step to be sure we answer
correctly.

You are a policy expert trying to help determine whether an AI response to
prompt is in violation of the defined safety policies.
<start_of_turn>
Human Question: [User Prompt]
<end_of_turn>
Our safety principle is defined in the below:

"No Hate Speech": The chatbot shall not generate content that expresses,
incites, or promotes hate based on race, gender, ethnicity, religion,
nationality, sexual orientation, disability status, or caste.
<start_of_turn>
Chatbot Response: [Model Response]
<end_of_turn>

Does the Chatbot Response violate the above principle? Your answer must start
with ’Yes’ or ’No’. And then walk through step by step to be sure we answer
correctly.

Figure 2: Instructions for Supervised Fine-Tuning. Left: User Input use case;
Right: Model Output use case.

We further supervise fine-tune (SFT) Gemma2 Instruction-Tuned (IT) models (2B,
9B, and 27B parameters) using the instruction shown in Fig. 2. We employ
distinct policy definitions for each harm type and the model output is either
Yes or No token. Our models are trained on TPUv5 lite with batch size of 16, a
max sequence of $8k$, and a learning rate of $1\mathrm{e}{-6}$. The model is
trained for $4k$ steps and the best checkpoints are selected based on
validation data. We calculate our predicted probability based on Eq. 1 below:

$\frac{\exp(\text{LL(Yes)}/T)+\alpha}{\exp(\text{LL(Yes)}/T)+\exp(\text{LL(No)}/T)+2\alpha}$
(1)

Here LL($\cdot$) is the log likelihood of the token generated by the model;
$T$ and $\alpha$ are hyperparameters to control temperature and uncertainty
estimate.

##  5 Experiments

| Prompt Classification | Response Classification
---|---|---
|  SG Prompt |  OpenAI Mod |  ToxicChat |  SG Response
ShieldGemma (2B) |  0.825/0.887 |  0.812/0.887 |  0.704/0.778 |  0.743/0.802
ShieldGemma (9B) |  0.828/0.894 |  0.821/0.907 |  0.694/0.782 |  0.753/0.817
ShieldGemma (27B) |  0.830/0.883 |  0.805/0.886 |  0.729/0.811 |  0.758/0.806
OpenAI Mod API |  0.782/0.840 |  0.790/0.856 |  0.254/0.588 |  -
LlamaGuard1 (7B) |  - |  0.758/0.847 |  0.616/0.626 |  -
LlamaGuard2 (8B) |  - |  0.761/- |  0.471/- |  -
WildGuard (7B) |  0.779/- |  0.721/- |  0.708/- |  0.656/-
GPT-4 |  0.810/0.847 |  0.705/- |  0.683/- |  0.713/0.749
Table 1: Evaluation results based on Optimal F1(left)/AU-PRC(right), higher is
better. We use $\alpha=0$ and $T=1$ for calculating the probabilities.
ShieldGemma (SG) Prompt and SG Response are our test datasets and OpenAI
Mod/ToxicChat are external benchmarks. On average, both our 9B and 27B model
perform the best. The performance of baseline models on external datasets is
sourced from Inan et al. (2023); Ghosh et al. (2024).

###  5.1 Setup

Despite the abundance of safety-related benchmark datasets, direct comparison
remains challenging due to several factors: (i) variations in policy
definitions and supported harm types across datasets; (ii) inconsistencies in
policy definitions even within the same harm type; and (iii) the predominance
of binary (safe/unsafe) classification models rather than harm-type level
prediction. To address these challenges, we conduct experiments on two fronts:

  * Overall Binary Classification: We aggregate our prediction results into binary outcomes by maximizing probabilities over all harms. For models that provide probabilities, we report both optimal F1 and AU-PRC scores; for models with only binary results, we report F1 scores.

  * Harm Type Level Classification: We provide a detailed performance comparison at the individual harm type level. We adopt a one-vs-all setup as described in Inan et al. (2023), i.e. we transform a multi-class classification problem into multiple binary classification problems, where each classifier focuses on distinguishing positive examples in one specific harm type and treat all others as benign examples.

###  5.2 Benchmark Datasets and Baseline Models

OpenAI Moderation (Markov et al., 2023) comprises 1,680 prompt examples
labeled for eight safety categories: sexual, hate, violence, harassment, self-
harm, sexual/minors, hate/threatening, violence/graphic. Given that the
original OpenAI Moderation policy definitions differ from ours, particularly
we do not directly predict self-harm, we utilize those original definitions to
predict each harm and then aggregate them into an overall binary
classification. The dataset is sourced from CommonCrawl which does not match
with the style of either user prompt or model output. Here, we run inference
by treating the text as model output and keep empty user prompt.

ToxicChat (Lin et al., 2023) contains $10k$ examples with binary toxicity
label for the prompt. We directly maximize our predictions for the six harms
according to our policy, as our harm types capture different aspects of the
toxicity definitions outlined in the ToxicChat policy.

ShieldGemma Prompt & ShieldGemma Response are our test dataset. it contains
4,500 examples with labels in total for both use cases. They have labels for
our targeted harm types sexual, dangerous content, harassment, hate speech and
non-targeted types violence and obscenity. More details are in section Data
Annotation.

Baseline Models: We evaluate our models against several models: OpenAI Mod API
(Markov et al., 2023), LlamaGuard (Team, 2024b), WildGuard Han et al. (2024),
and GPT-4. For GPT-4, we utilize the openAI API (model=gpt-4-0613) with our
prompts, obtaining the log probability of the first token and converting it
into the probability of a policy violation.

###  5.3 Overall Binary Classification Results

The overall binary classification results are presented in Table 1. All
ShieldGemma (SG) models (2B, 9B and 27B) outperform all baseline models.
Notably, with similar model size and training data volume, SG-9B achieves a
10.8% higher average AU-PRC compared to LlamaGuard1 on external benchmarks.
Additionally, the F1 score of our 9B model exceeds that of WildGuard and GPT-4
by 4.3% and 6.4%, respectively.

Within the SG models, performance is comparable on our internal benchmarks. On
external benchmarks, the 9B/27B model demonstrates slightly stronger
generalization capability, achieving on average a 1.2%/1.7% higher AU-PRC than
its 2B model.

###  5.4 Harm Type Level Results

Figure 3: Harm Type level performance (AU-PRC) for our test dataset SG Prompt
(left) and SG Response (right).

We evaluate the harm-type level performance on our test datasets: SG Prompt
and SG Response. The results are shown in Fig. 3. All SG models have
outperformed GPT-4 by a big margin for all of the harms. Overall, GPT-4 is
weak in distinguishing different harms. For example 76% of hate speech data
points have been classified as positive for harassment. Note that the
performance gap is expected, and the comparison is less favorable for GPT-4,
as our model has been trained on datasets similar to the test datasets, while
GPT-4 is evaluated zero-shot without any specific training. The performance
among SG models is close to each other. On average, SG-9B and SG-27B have
outperformed SG-2B by less than 2%.

##  6 Limitations

Despite our efforts to enhance the robustness of our model against adversarial
attacks, fairness, and diversity in the training data, several limitations
remain:

Fairness: While we have implemented fairness counterfactual expansion to
mitigate bias in our training data, label discrepancies may still arise when
identity groups are swapped. These discrepancies often stem from inherent
biases within the pre-training dataset (Chen et al., 2024).

Generalization: We have observed that our larger models demonstrate stronger
performance on external benchmarks with new harm types and text styles.
Overall, this generalization capability of our larger models are slightly
stronger than our smaller 2B model. It also requires additional experiments to
further verify the generalization on other datasets.

Implicit Cultural Harm: Although LLMs exhibit some understanding of cultural
contexts, they may struggle to fully grasp implicit harm within these
contexts.

Safety vs. Helpfulness: While our models demonstrate a strong ability to
filter potential safety risks, their interpretation of policy violations may
be overly conservative. This could interfere with helpfulness when used to
filter LLM responses. We recommend that downstream clients adjust filtering
thresholds based on their specific use cases.

LLM-as-a-classifier: Our model is specifically designed for classification
tasks, with an output restricted to Yes or No token as the first output token
when the prompt is correctly configured. However, it’s crucial to acknowledge
that as an LLM, it remains capable of generating responses to any text input.
We strongly advise the users to use it solely for generating Yes/No token
scores (we call it scoring mode, detailed in our model card), and avoid using
it in a chat-like manner since it may produce unethical or unsafe content due
to the absence of additional safety instruction-tuning for conversational use.

We are dedicated to ongoing research and development to address these
limitations and further refine our classifiers.

##  7 Conclusion

This paper presents a significant advancement in safety content moderation
through our suite of specialized models, built on the foundation of the public
Gemma2 (Team et al., 2024) language models. We demonstrate their superior
performance on diverse benchmarks, highlighting the effectiveness of our
approach. Additionally, our novel synthetic data generation pipeline offers a
valuable tool for researchers and practitioners to create high-quality,
diverse datasets for safety and other domains. We are excited to share these
resources with the research community to foster further development in this
critical area.

##  8 Contributions and Acknowledgments

Core Contributors
Wenjun Zeng
Yuchi Liu
Ryan Mullins
Ludovic Peran

Contributors
Joe Fernandez
Hamza Harkous
Karthik Narasimhan
Drew Proud
Piyush Kumar
Bhaktipriya Radharapu
Olivia Sturman
Oscar Wahltinez

Other Specialty Areas
Special thanks and acknowledgments to these individuals for their assistance
in respected areas:

Central Support
Manvinder Singh
Kathy Meier-Hellstern
Shivani Podder

Checkpoint Conversions
Nam T. Nguyen
Matthew Watson

Ethics and Safety
Antonia Paterson
Jenny Brennan

Gemma Model
Surya Bhupatiraju
Victor Cotruta
Armand Joulin
Kathleen Kenealy
Tris Warkentin

Go-to-Market
Kat Black
Meg Risdal

Team Acknowledgements
Our work is made possible by the dedication and efforts of numerous teams at
Google. We would like to acknowledge the support from the following teams:
Gemma, Google DeepMind Responsibility, Kaggle, Keras, Perspective.

## References

  * Achiam et al. (2023) J. Achiam, S. Adler, S. Agarwal, L. Ahmad, I. Akkaya, F. L. Aleman, D. Almeida, J. Altenschmidt, S. Altman, S. Anadkat, et al.  Gpt-4 technical report.  _arXiv preprint arXiv:2303.08774_ , 2023.
  * Anthropic (2024) A. Anthropic.  The claude 3 model family: Opus, sonnet, haiku.  _Claude-3 Model Card_ , 1, 2024.
  * Ash et al. (2019) J. T. Ash, C. Zhang, A. Krishnamurthy, J. Langford, and A. Agarwal.  Deep batch active learning by diverse, uncertain gradient lower bounds.  _arXiv preprint arXiv:1906.03671_ , 2019.
  * Bai et al. (2022) Y. Bai, A. Jones, K. Ndousse, A. Askell, A. Chen, N. DasSarma, D. Drain, S. Fort, D. Ganguli, T. Henighan, et al.  Training a helpful and harmless assistant with reinforcement learning from human feedback.  _arXiv preprint arXiv:2204.05862_ , 2022.
  * Chen et al. (2024) G. H. Chen, S. Chen, Z. Liu, F. Jiang, and B. Wang.  Humans or llms as the judge? a study on judgement biases.  _arXiv preprint arXiv:2402.10669_ , 2024.
  * Citovsky et al. (2021) G. Citovsky, G. DeSalvo, C. Gentile, L. Karydas, A. Rajagopalan, A. Rostamizadeh, and S. Kumar.  Batch active learning at scale.  _Advances in Neural Information Processing Systems_ , 34:11933–11944, 2021.
  * Deng et al. (2023) Y. Deng, W. Lei, M. Huang, and T.-S. Chua.  Rethinking conversational agents in the era of llms: Proactivity, non-collaborativity, and beyond.  In _Proceedings of the Annual International ACM SIGIR Conference on Research and Development in Information Retrieval in the Asia Pacific Region_ , pages 298–301, 2023.
  * Devlin et al. (2018) J. Devlin, M.-W. Chang, K. Lee, and K. Toutanova.  Bert: Pre-training of deep bidirectional transformers for language understanding.  _arXiv preprint arXiv:1810.04805_ , 2018.
  * Gao et al. (2022) J. Gao, R. Pi, Y. Lin, H. Xu, J. Ye, Z. Wu, W. Zhang, X. Liang, Z. Li, and L. Kong.  Self-guided noise-free data generation for efficient zero-shot learning.  _arXiv preprint arXiv:2205.12679_ , 2022.
  * Ghosh et al. (2024) S. Ghosh, P. Varshney, E. Galinkin, and C. Parisien.  Aegis: Online adaptive ai content safety moderation with ensemble of llm experts.  _arXiv preprint arXiv:2404.05993_ , 2024.
  * Google (2017) Google.  Perspective api.  https://www.perspectiveapi.com/, 2017.
  * Google (2024) Google.  Responsible generative ai toolkit: https://ai.google.dev/responsible/principles, 2024.
  * Han et al. (2024) S. Han, K. Rao, A. Ettinger, L. Jiang, B. Y. Lin, N. Lambert, Y. Choi, and N. Dziri.  Wildguard: Open one-stop moderation tools for safety risks, jailbreaks, and refusals of llms.  _arXiv preprint arXiv:2406.18495_ , 2024.
  * Huang et al. (2024) H. Huang, Y. Qu, J. Liu, M. Yang, and T. Zhao.  An empirical study of llm-as-a-judge for llm evaluation: Fine-tuned judge models are task-specific classifiers.  _arXiv preprint arXiv:2403.02839_ , 2024.
  * Inan et al. (2023) H. Inan, K. Upasani, J. Chi, R. Rungta, K. Iyer, Y. Mao, M. Tontchev, Q. Hu, B. Fuller, D. Testuggine, et al.  Llama guard: Llm-based input-output safeguard for human-ai conversations.  _arXiv preprint arXiv:2312.06674_ , 2023.
  * Ji et al. (2023) J. Ji, M. Liu, J. Dai, X. Pan, C. Zhang, C. Bian, C. Zhang, R. Sun, Y. Wang, and Y. Yang.  Beavertails: Towards improved safety alignment of llm via a human-preference dataset, 2023.  URL https://arxiv.org/abs/2307.04657.
  * Kim et al. (2022) S. Y. Kim, H. Park, K. Shin, and K.-M. Kim.  Ask me what you need: Product retrieval using knowledge from gpt-3.  _arXiv preprint arXiv:2207.02516_ , 2022.
  * Kurakin et al. (2023) A. Kurakin, N. Ponomareva, U. Syed, L. MacDermed, and A. Terzis.  Harnessing large-language models to generate private synthetic text.  _arXiv preprint arXiv:2306.01684_ , 2023.
  * Kusner et al. (2017) M. J. Kusner, J. Loftus, C. Russell, and R. Silva.  Counterfactual fairness.  _Advances in neural information processing systems_ , 30, 2017.
  * Li et al. (2024) L. Li, B. Dong, R. Wang, X. Hu, W. Zuo, D. Lin, Y. Qiao, and J. Shao.  Salad-bench: A hierarchical and comprehensive safety benchmark for large language models.  _arXiv preprint arXiv:2402.05044_ , 2024.
  * Lin et al. (2023) Z. Lin, Z. Wang, Y. Tong, Y. Wang, Y. Guo, Y. Wang, and J. Shang.  Toxicchat: Unveiling hidden challenges of toxicity detection in real-world user-ai conversation, 2023.  URL https://arxiv.org/abs/2310.17389.
  * Liu et al. (2024) N. Liu, L. Chen, X. Tian, W. Zou, K. Chen, and M. Cui.  From llm to conversational agent: A memory enhanced architecture with fine-tuning of large language models.  _arXiv preprint arXiv:2401.02777_ , 2024.
  * Long et al. (2024) L. Long, R. Wang, R. Xiao, J. Zhao, X. Ding, G. Chen, and H. Wang.  On llms-driven synthetic data generation, curation, and evaluation: A survey.  _arXiv preprint arXiv:2406.15126_ , 2024.
  * Markov et al. (2023) T. Markov, C. Zhang, S. Agarwal, F. E. Nekoul, T. Lee, S. Adler, A. Jiang, and L. Weng.  A holistic approach to undesired content detection in the real world.  In _Proceedings of the AAAI Conference on Artificial Intelligence_ , volume 37, pages 15009–15018, 2023.
  * Mazeika et al. (2024) M. Mazeika, L. Phan, X. Yin, A. Zou, Z. Wang, N. Mu, E. Sakhaee, N. Li, S. Basart, B. Li, D. Forsyth, and D. Hendrycks.  Harmbench: A standardized evaluation framework for automated red teaming and robust refusal, 2024.  URL https://arxiv.org/abs/2402.04249.
  * Nasr et al. (2023) M. Nasr, N. Carlini, J. Hayase, M. Jagielski, A. F. Cooper, D. Ippolito, C. A. Choquette-Choo, E. Wallace, F. Tramèr, and K. Lee.  Scalable extraction of training data from (production) language models.  _arXiv preprint arXiv:2311.17035_ , 2023.
  * Radharapu et al. (2023) B. Radharapu, K. Robinson, L. Aroyo, and P. Lahoti.  Aart: Ai-assisted red-teaming with diverse data generation for new llm-powered applications.  _arXiv preprint arXiv:2311.08592_ , 2023.
  * Sahu et al. (2022) G. Sahu, P. Rodriguez, I. H. Laradji, P. Atighehchian, D. Vazquez, and D. Bahdanau.  Data augmentation for intent classification with off-the-shelf large language models.  _arXiv preprint arXiv:2204.01959_ , 2022.
  * Sener and Savarese (2017) O. Sener and S. Savarese.  Active learning for convolutional neural networks: A core-set approach.  _arXiv preprint arXiv:1708.00489_ , 2017.
  * Smith et al. (2022) E. M. Smith, M. Hall, M. Kambadur, E. Presani, and A. Williams.  " i’m sorry to hear that": Finding new biases in language models with a holistic descriptor dataset.  _arXiv preprint arXiv:2205.09209_ , 2022.
  * Team (2024a) G. Team.  Gemma.  2024a.  10.34740/KAGGLE/M/3301.  URL https://www.kaggle.com/m/3301.
  * Team et al. (2023) G. Team, R. Anil, S. Borgeaud, Y. Wu, J.-B. Alayrac, J. Yu, R. Soricut, J. Schalkwyk, A. M. Dai, A. Hauth, et al.  Gemini: a family of highly capable multimodal models.  _arXiv preprint arXiv:2312.11805_ , 2023.
  * Team et al. (2024) G. Team, T. Mesnard, C. Hardin, R. Dadashi, S. Bhupatiraju, S. Pathak, L. Sifre, M. Rivière, M. S. Kale, J. Love, et al.  Gemma: Open models based on gemini research and technology.  _arXiv preprint arXiv:2403.08295_ , 2024.
  * Team (2024b) L. Team.  Meta llama guard 2.  https://github.com/meta-llama/PurpleLlama/blob/main/Llama-Guard2/MODEL_CARD.md, 2024b.
  * Zhan et al. (2018) X. Zhan, Z. Liu, P. Luo, X. Tang, and C. Loy.  Mix-and-match tuning for self-supervised semantic segmentation.  In _Proceedings of the AAAI conference on artificial intelligence_ , volume 32, 2018.
  * Zheng et al. (2024) L. Zheng, W.-L. Chiang, Y. Sheng, S. Zhuang, Z. Wu, Y. Zhuang, Z. Lin, Z. Li, D. Li, E. Xing, et al.  Judging llm-as-a-judge with mt-bench and chatbot arena.  _Advances in Neural Information Processing Systems_ , 36, 2024.


# Graph Decision Transformer

Shengchao Hu1 Li Shen2 Ya Zhang1,3 Dacheng Tao4,2
1Shanghai Jiaotong University, China
2JD Explore Academy, China
3 Shanghai AI Laboratory, China
4 The University of Sydney, Australia
<EMAIL_ADDRESS><EMAIL_ADDRESS>Corresponding author: Li Shen

###### Abstract

Offline reinforcement learning (RL) is a challenging task, whose objective is
to learn policies from static trajectory data without interacting with the
environment. Recently, offline RL has been viewed as a sequence modeling
problem, where an agent generates a sequence of subsequent actions based on a
set of static transition experiences. However, existing approaches that use
transformers to attend to all tokens naively can overlook the dependencies
between different tokens and limit long-term dependency learning. In this
paper, we propose the Graph Decision Transformer (GDT), a novel offline RL
approach that models the input sequence into a causal graph to capture
potential dependencies between fundamentally different concepts and facilitate
temporal and causal relationship learning. GDT uses a graph transformer to
process the graph inputs with relation-enhanced mechanisms, and an optional
sequence transformer to handle fine-grained spatial information in visual
tasks. Our experiments show that GDT matches or surpasses the performance of
state-of-the-art offline RL methods on image-based Atari and OpenAI Gym.

Figure 1: Comparison of GDT and DT. GDT employs additional edge embeddings and
node embeddings to obtain Q and K, while using only node embeddings to obtain
V.

##  1 Introduction

Reinforcement Learning (RL) is inherently a sequential process where an agent
observes a state from the environment, takes an action, observes the next
state, and receives a reward. To model RL problems, Markov Decision Processes
(MDPs) have been widely employed, where an action is taken solely based on the
current state, which is assumed to encapsulate the entire history. Online RL
algorithms [33, 45] use the temporal difference (TD) learning to train agents
by interacting with the environment, but this can be prohibitively expensive
in real-world settings. Offline RL [31], on the other hand, seeks to overcome
this limitation by learning policies from a pre-collected dataset, without the
need to interact with the environment. This approach makes RL training more
practical for real-world scenarios and has therefore garnered significant
attention.

Recent advances [9, 24] in offline RL have taken a new perspective on the
problem, departing from conventional methods that concentrate on learning
value functions [39, 26] or policy gradients [37, 15]. Instead, the problem is
viewed as a generic sequence modeling task, where past experiences consisting
of state-action-reward triplets are input to Transformer [48]. The model
generates a sequence of action predictions using a goal-conditioned policy,
effectively converting offline RL to a supervised learning problem. This
approach relaxes the MDP assumption by considering multiple historical steps
to predict an action, allowing the model to be capable of handling long
sequences and avoid stability issues associated with bootstrapping [46, 29].
Furthermore, this framework unifies multiple components in offline RL, such as
estimating the behavior policy and predictive dynamics modeling, into a single
sequence model, resulting in superior performance.

However, this approach faces three major issues. Firstly, states and actions
represent fundamentally different concepts [49]. While the agent has complete
control over its action sequences, the resulting state transitions are often
influenced by external factors. Thus, modeling states and actions as a single
sequence may indiscriminate the effects of the policy and world dynamics on
the return, which can lead to overly optimistic behavior. Secondly, in RL
problems, the adjacent states, actions, and rewards are typically strongly
connected due to their potential causal relationships. Specifically, the state
observed at a given time step is a function of the previous state and action,
and the action taken at that time step influences the subsequent state and
reward. Simply applying Transformer to attend to all tokens without
considering the underlying Markovian relationship can result in an
overabundance of information, hindering the model’s ability to accurately
capture essential relation priors and handle long-term sequences of
dependencies. Finally, tokenizing image states as-a-whole using convolutional
neural networks (CNNs) can hinder the ability of Transformers to gather fine-
grained spatial relations. This loss of information can be especially critical
in visual RL tasks that require detailed knowledge of regions-of-interest.
Therefore, it is necessary to find a more effective way to represent states
and actions separately while still preserving their intrinsic relationship,
and to incorporate the Markovian property and spatial relations in the
modeling process.

To alleviate such issues, we propose a novel approach, namely Graph Decision
transformer (GDT), which involves transforming the input sequence into a graph
structure. The Graph Representation explicitly incorporates the potential
dependencies between adjacent states, actions, and rewards, thereby better
capturing the Markovian property of the input and differentiating the impact
of different tokens. To process the input graph, we utilize the Graph
Transformer to effectively handle long-term dependencies that may be present
in non-Markovian environments. To gather fine-grained spatial information, we
incorporate an optional Sequence Transformer following StAR [42] to encode
image states as patches similar to ViT [12], which helps with action
prediction and reduces the learning burden of the Graph Transformer. Our
experimental evaluations conducted in Atari and OpenAI Gym environments
provide empirical evidence to support the advantages of utilizing a causal
graph representation as input to the Graph Transformer in RL tasks. The
proposed GDT method achieves state-of-the-art performance in several benchmark
environments and outperforms existing offline RL methods.

In summary, our main contributions are four-fold:

  * •

We propose a novel approach named GDT, that transforms input sequences into
graph structures to better capture potential dependencies between adjacent
states, actions, and rewards and differentiate the impact of these different
tokens.

  * •

We utilize the Graph Transformer to process the input graph, which can
effectively handle long-term dependencies that may be present in non-Markovian
environments.

  * •

We incorporate Sequence Transformer to encode image states as patches to
gather fine-grained spatial information crucial for visual input environments.

  * •

We extensively evaluate GDT on Atari and OpenAI Gym environments,
demonstrating its superior performance compared to existing methods

Figure 2: The proposed model comprises three main components: the Graph
Representation, the Graph Transformer, and an optional Sequence Transformer.
When the output of the Graph Transformer is used directly for action
prediction, the resulting model is referred to as GDT, which includes only the
left half of the figure. On the other hand, if the output of the Graph
Transformer is further processed by the Sequence Transformer, the resulting
model is referred to as GDT-plus, which encompasses the entire figure. Figure
3: Input representation as a graph. Nodes $\hat{R}_{t}$, $s_{t}$, and $a_{t}$
represent the return-to-go, state, and action at time $t$, respectively. The
directed edge connecting $s_{t-1}$ to $s_{t}$ indicates that the value of
$s_{t-1}$ affects the value of $s_{t}$, and similar relationships exist for
other directed edges.

##  2 Related Work

#### Offline RL

Offline RL has recently gained significant attention as an alternative
paradigm, where agents extract return-maximizing policies from fixed, limited
datasets composed of trajectory rollouts from arbitrary policies [31]. These
datasets, referred to as static datasets, are formally defined as
$\mathcal{D}=\\{(s_{t},a_{t},s_{t+1},r_{t})_{i}\\}$, where $i$ is the index,
the actions and states are generated by the behavior policy $(s_{t},a_{t})\sim
d^{\pi_{\beta}}(\cdot)$, and the next states and rewards are determined by the
dynamics $(s_{t+1},r_{t})\sim(T(\cdot|s_{t},a_{t}),r(s_{t},a_{t}))$. Deploying
off-policy RL algorithms directly in the offline setting is hindered by the
distributional shift problem, which can result in a significant performance
drop, as demonstrated in prior research [16]. To mitigate this issue, model-
free algorithms aim to either constrain the action space of policy [27, 44] or
incorporate pessimism to value function [30, 16]. Conversely, model-based
algorithms simulate the actual environment to generate more data for policy
training [25, 54]. In this work, we propose a novel approach that avoids
learning the dynamics model explicitly and directly generates the next action
with the help of the Graph Transformer, enabling better generalization and
transfer [38].

#### RL to Sequence Modeling

RL has recently garnered considerable interest as a sequence modeling task,
particularly with the application of Transformer-based decision models [21].
The task is to predict a sequence of next actions given a sequence of recent
experiences, including state-action-reward triplets. This approach can be
trained in a supervised learning fashion, making it more amenable to offline
RL and imitation learning settings. Several studies [50, 53, 17, 55] have
explored the use of Transformers in RL under the sequence modeling pattern.
For example, Chen et al. [9] train a Transformer as a model-free context-
conditioned policy, while Janner et al. [24] bring out the capability of the
sequence model by predicting states, actions, and rewards and employing beam
search. Zheng et al. [57] further fine-tune the Transformer by adapting this
formulation to online settings. Shang et al. [42] explicitly model StAR-
representations to introduce a Markovian-like inductive bias to improve long-
term modeling. In this work, we propose a graph sequence modeling approach to
RL, which explicitly introduces the Markovian property to the representations.
Our proposed GDT method outperforms several state-of-the-art non-Transformer
offline RL and imitation learning algorithms on Atari and Gym benchmarks,
demonstrating the advantage of incorporating graph structures in sequence
modeling for RL tasks.

#### RL with Graph

In recent years, the integration of graph neural networks (GNNs) with RL has
attracted considerable attention for graph-structured environments [34].
Specifically, GNNs can be combined with RL to address sequential decision-
making problems on graphs. Existing research primarily focuses on using deep
RL to improve GNNs for diverse purposes, such as neural architecture search
(NAS) [58], enhancing the interpretability of GNN predictions [41], and
designing adversarial examples for GNNs [11, 47]. Alternatively, GNNs can be
utilized to solve relational RL problems, such as those involving different
agents in a multi-agent deep RL (MADRL) framework [43, 6, 56], and different
tasks in a multi-task deep RL (MTDRL) framework [51, 23, 4, 22]. Despite the
growing interest in this field, there is currently a lack of research on
utilizing a Markovian dependency graph as input to GNNs for action prediction.
Such an approach has strong potential due to the causal relationship between
the constructed graph and its ability to be employed in various RL
environments. This article will explore this approach in detail, offering a
novel contribution to the field. Additionally, we will highlight the
advantages of using GDT for action prediction in comparison to existing state-
of-the-art algorithms.

##  3 Methodology

The proposed approach is a deep learning-based model for sequence modeling in
RL tasks, which combines the advantages of graph and sequence modeling. The
model consists of three main components: the Graph Representation, the Graph
Transformer, and an optional Sequence Transformer, as shown in Figure 2.
Specifically, Graph Representation is used to represent the input sequence as
a graph with a causal relationship, which captures the dependencies between
fundamentally different concepts. The Graph Transformer then processes the
graph inputs using relation-enhanced mechanisms, which allows the model to
acquire long-term dependencies and model the interactions between different
time steps of the graph tokens given reasonable causal relationships. The
optional Sequence Transformer is introduced to gather the fine-grained spatial
information in the input, which is particularly important in visual tasks such
as Atari environment.

The proposed approach offers several advantages. Firstly, it can effectively
acquire the intricate dependencies and interactions between different time
steps in the input sequence, making it well-suited for RL tasks. Secondly, it
encodes the sequence as a causal graph, which explicitly incorporates
potential dependencies between adjacent tokens, thereby facilitating the
capture of the Markovian property and differentiating the impact of different
tokens. Finally, it can accurately gather fine-grained spatial information and
integrate it into action prediction, leading to improved performance. In
summary, our proposed approach offers a powerful and flexible solution for RL
tasks. Its unique design principles, which combine graph and sequence
modeling, enable it to effectively obtain the dependencies and interactions
between different time steps and tokens, making it a promising approach for
various applications. In the following sections, we provide a detailed
description of each component of our proposed approach, outlining its input
and output and discussing its underlying design principles.

###  3.1 Graph Representation

Our approach constructs the input graph from trajectory elements based on
their dependencies, capturing the temporal and causal relationships of the
data. Each trajectory element represents a node in the graph, and edges are
added accordingly. This approach considers the Markovian nature of the data
and differs from the DT [9] serialized input. Specifically, our approach
considers the inputs to the trajectory as returns-to-go
$\hat{R}_{t}=\sum_{t^{\prime}=t}^{T}r_{t^{\prime}}$, states $s_{t}$, and
actions $a_{t}$, as illustrated in Figure 3. Action $a_{t}$ selection depends
on the current state $s_{t}$ and return-to-go $\hat{R}_{t}$. State $s_{t}$
generation is dependent on the previous state $s_{t-1}$ and action $a_{t-1}$.
The return-to-go $\hat{R}_{t}$ is jointly determined by the previous return-
to-go $\hat{R}_{t-1}$, state $s_{t-1}$, and action $a_{t-1}$.

The MDP is a framework where an agent is asked to make a decision based on the
current state $s_{t}$; then, the environment responds to the action made by
the agent and transitions the state to the next state $s_{t+1}$ with a reward
$r_{t}$. We model these dependencies as directed edges and exclude the agent
and environment objects to enable the Graph Transformer to better analyze the
causal relationships between different tokens. Through this approach, we
explicitly embed the Markovian relationship into the model and distinguish the
effect of distinct tokens, avoiding the overabundance of information that may
prevent the model from accurately capturing essential relation priors.

Figure 4: The environments utilized in this study include four Atari
environments on the left, which accept images as inputs and have a discrete
action space, and three OpenAI Gym environments on the right, which accept
vectors as inputs and have a continuous action space.

###  3.2 Graph Transformer

Our approach employs autoregressive modeling to predict current action using
the graph representation as input to the Graph Transformer. The Graph
Transformer plays a central role in extracting hidden features from the graph.
In some scenarios, such as Atari games, past information also plays a critical
role, thus we employ a global attention mechanism, allowing each node to not
only observe its directly connected nodes but also all nodes preceding the
current moment. In the vanilla multi-head attention, the attention score
between the element $x_{i}$ and $x_{j}$ can be formulated as the dot-product
between their query vector and key vector, respectively:

$\small s_{ij}=f(x_{i},x_{j})=x_{i}W_{q}^{T}W_{k}x_{j}.$ (1)

In the case of Graph Transformers, the computation of attention scores also
needs to account for the temporal and causal relationships. To achieve this,
we follow the concept of relation-enhanced mechanism proposed in [8], which
utilizes both node and relation representations in computing attention scores,
as shown in Figure 1. The equation for the computation of attention scores is
shown below:

$\small\begin{split}s_{ij}&=g(x_{i},x_{j},r_{ij})\\\ &=(x_{i}+r_{i\rightarrow
j})W_{q}^{T}W_{k}(x_{j}+r_{j\rightarrow i}),\end{split}$ (2)

where $j\leq i$ and $r_{ij}$ is learned through an embedding layer that takes
the adjacency matrix as input. The incorporation of relation representation
enables the model to take into account the plausible causal relationships,
which relieves the burden on the Graph Transformer of learning potential long-
term dependencies among node vectors at different time steps. This observation
has been empirically validated through our ablation experiments.

The input to the $l$-th layer of the Graph Transformer is a token graph. To
preserve temporal relationships, the graph is pruned to form the following
sequence $\overline{G}$ ($\overline{\cdot}$ means pruning into sequence):

$\small\begin{split}\overline{G}_{in}^{l}&=\\{\hat{R}_{0}^{l-1},s_{0}^{l-1},a_{0}^{l-1},\dots,\hat{R}_{T}^{l-1},s_{T}^{l-1},a_{T}^{l-1}\\}.\\\
\end{split}$ (3)

Each token graph is transformed to $g_{t}^{l}$ by a Graph Transformer layer:

$\small\begin{split}G_{out}^{l}&=F^{l}_{graph}(G_{in}^{l})=F^{l}_{graph}(G_{out}^{l-1}),\\\
g_{t}^{l}:&=\text{FC}(\overline{G}_{out}^{l}[1+3t],\overline{G}_{out}^{l}[3t]).\end{split}$
(4)

As shown in Figure 3, the action $a_{t}$ is determined by both $\hat{R}_{t}$
and $s_{t}$. Thus, the feature vector $g_{t}^{l}$ is obtained by concatenating
the two inputs and fed into a fully connected layer (with indexing starting
from 0). The feature $g_{t}^{l}$ extracted from the $l$-th layer of the Graph
Transformer can be directly used to predict the action
$\hat{a}_{t}=\phi(g_{t}^{l})$, or be further processed by the subsequent
Sequence Transformer.

###  3.3 Sequence Transformer

Motivated by [42], we introduce the Sequence Transformer to assist with action
prediction and reduce the learning burden of the Graph Transformer. The
Sequence Transformer adopts the conventional Transformer layer design from
[48] and is incorporated to gather fine-grained spatial information that is
crucial for visual input environments. The initial layer of the Sequence
Transformer takes a collection of image patches and $g_{t}^{0}$ as inputs:

$\small
Y_{in,t}^{0}=\\{z_{s_{t}^{1}},z_{s_{t}^{2}},\dots,z_{s_{t}^{n}},g_{t}^{0}\\},$
(5)

where $n$ is the number of image patches, and the feature vector $g_{t}$ is
positioned after state patches $\\{z_{s_{t}^{i}}\\}_{i=1}^{n}$, which enables
$g_{t}$ to attend to all spatial information. We have $T$ groups of such token
representations, which are simultaneously processed by the Sequence
Transformer:

$\small\begin{split}Y_{in}^{0}&=\sum_{t=0}^{T}Y_{in,t}^{0},\\\
Y_{out}^{0}&=F_{sequence}^{0}(Y_{in}^{0}),\\\ \end{split}$ (6)

where adding means concatenating the contents of two collections. The
subsequent layer of the Sequence Transformer takes the fusion of the previous
layer’s output and $g_{t}^{l}$ as its input. This is achieved by adding
$g_{t}^{l}$ to the position that corresponds to $g_{t}^{l-1}$ in the output
sequence while leaving the other parts of the output sequence unchanged. The
formulation for this operation is as follows:

$\small\\!\\!Y_{in}^{l}[i]\\!=\\!\hat{Y}_{out}^{l-1}[i]\\!=\\!\begin{cases}Y_{out}^{l-1}[i]+g_{t}^{l},&i=n+t(n+1),\\\
&t=0,1,2,\dots,T\\\ Y_{out}^{l-1}[i],&\text{otherwise}\end{cases}$ (7)

The output feature $h_{t}^{l}:=Y_{out}^{l}[n+t(n+1)]$ extracted from the
$l$-th layer of the Sequence Transformer is fed into a linear head to predict
the action, denoted as $\hat{a}_{t}=\phi(h_{t}^{l})$ when it is the final
layer. Further elaboration on this connection method will be provided in Sec.
4.3.

###  3.4 Training procedure

GDT is a drop-in replacement for DT as the training and inference procedures
remain the same. However, additional graph construction for the input sequence
is required for GDT. Specifically, a graph $G=(X,A)$ is constructed from a
minibatch of length $K$ (total $3K$ tokens: return-to-go, state, and action)
where $X$ represents the node embedding matrix, and $A$ represents the edge
embedding matrix. The constructed graph is then input into the Graph
Transformer, and attention scores are calculated using both $X$ and $A$. The
learning objective for discrete environments can be formulated as follows:

$\small\\!\\!\\!\mathbb{E}_{(\hat{R},s,a)\sim\mathcal{T}}\left[\\!\frac{1}{T}\sum_{t=1}^{T}(a_{t}\\!-\\!\pi_{\text{GDT}}(s_{-K:t},\hat{R}_{-K:t},a_{-K:t-1}))^{2}\right].\\!$
(8)

To end this section, we give several comments on the proposed GDT method.
Compared with DT, which uses a serialized input, GDT represents the input
sequence as a graph with a causal relationship, enabling the Graph Transformer
to capture dependencies between fundamentally different concepts. On the other
hand, compared with methods that implicitly learn Markovian patterns [24] or
introduce an additional step Transformer [42], GDT directly incorporates the
Markovian relationship in the input. This feature allows the model to handle
dependencies between sequences and tokens effectively, leading to improved
performance without additional computational overhead. Additionally, the
introduced Sequence Transformer can maintain fine-grained spatial information
using ViT-like image patches, which is particularly important in visual tasks
and can improve action prediction accuracy in such environments. Although
GDT’s performance in vector state environments is not significantly improved
compared to that in visual environments, it is still effective in such
relatively simple environments (Sec. 4.2).

Figure 5: Comparison of different graph connection methods on Atari and OpenAI
Gym Environments (same in later experiments). The left panel illustrates four
different graph connection methods: Causal, Full, Random, and None. The right
panel shows the normalized performance comparison of these methods. Figure 6:
The left panel depicts the input graph structure of the GDT for three reward
settings: return-to-go (RTG), stepwise reward (Reward), and no reward (None).
The right panel displays the performance comparison between GDT and DT

##  4 Experiment

In this section, we provide a comprehensive evaluation of the proposed Graph
Decision Transformer (GDT) model, which is designed to capture the complex
relationships among graph-structured data and make effective decisions based
on them. Our main objective is to assess the effectiveness of GDT in
comparison to two popular algorithms: offline algorithms based on TD-learning
and imitation learning algorithms. TD-learning is a widely adopted algorithm
in RL due to its remarkable sampling efficiency and impressive performance on
various RL tasks. It has become a benchmark algorithm for many RL
applications. On the other hand, imitation learning algorithms, represented by
DT, have gained increasing attention in recent years due to their ability to
learn from expert demonstrations and achieve performance comparable to TD-
learning in various RL tasks. We conducted a comprehensive evaluation of the
performance of the GDT model on a range of tasks. Specifically, just as Figure
4 shown, we evaluated the performance of GDT on the widely used Atari
benchmark [5], which consists of a set of discrete control tasks, as well as
on the OpenAI Gym [7], which comprises a variety of continuous control tasks.

###  4.1 Atari

Table 1: Results for 1% DQN-replay Atari datasets. We evaluate the performance
of GDT on four Atari games using three different seeds, and report the mean
and variance of the results. Best mean scores are highlighted in bold. The
assessment reveals that GDT surpasses conventional RL algorithms on most tasks
and achieves better performance than DT across all games. In contrast, GDT-
plus achieves comparable results to StAR, leveraging fine-grained spatial
information.

Game |  CQL |  QR-DQN |  REM |  BC |  DT |  StAR |  GDT |  GDT-plus
---|---|---|---|---|---|---|---|---
Breakout |  211.1 |  17.1 |  8.9 |  138.9 $\pm$ 61.7 |  267.5 $\pm$ 97.5 |  436.1 $\pm$ 63.6 |  393.5 $\pm$ 98.8 |  441.7 $\pm$ 41.0
Qbert |  104.2 |  0 |  0 |  17.3 $\pm$ 14.7 |  15.4 $\pm$ 11.4 |  51.2 $\pm$ 11.5 |  45.5 $\pm$ 14.6 |  51.7 $\pm$ 20.8
Pong |  111.9 |  18 |  0.5 |  85.2 $\pm$ 20.0 |  106.1 $\pm$ 8.1 |  110.8 $\pm$ 60.3 |  108.4 $\pm$ 4.7 |  111.2 $\pm$ 0.9
Seaquest |  1.7 |  0.4 |  0.7 |  2.1 $\pm$ 0.3 |  2.5 $\pm$ 0.4 |  1.7 $\pm$ 0.3 |  2.8 $\pm$ 0.1 |  2.7 $\pm$ 0.1

The Atari benchmark [5] is a well-recognized and widely-adopted benchmark for
evaluating the performance of RL algorithms. In this study, we choose four
games from the Atari benchmark, namely Breakout, Pong, Qbert, and Seaquest,
each of which requires the agent to handle high-dimensional visual inputs and
complex credit assignment. Similar to prior work [9], we construct the offline
dataset by sampling 1% of the DQN replay buffer dataset [2], which consists of
nearly 500k transition steps. To enable fair comparisons, we follow the
normalization protocol proposed in [19], where the final scores are normalized
such that a score of 100 represents the expert level performance and a score
of 0 represents the performance of a random policy.

For baseline benchmark, we compare GDT with TD-learning based algorithms,
including CQL [30], QR-DQN [10], and REM [2], and several imitation learning
algorithms, including DT [9], StARformer [42], and behavior cloning (BC), and
report the results from original papers.

Table 1 presents the comparison of our proposed method with CQL, REM, QR-DQN,
BC, and DT methods on four games. The results show that our method achieves
comparable performance to CQL in three out of four games, while significantly
outperforming the other methods in all four games. This indicates that our
approach, which introduces the causal relationship in the input and leverages
the Graph Transformer accordingly, is superior to the other methods. To ensure
a fair comparison with StAR, we further introduce a Sequence Transformer to
incorporate fine-grained spatial information and report the results as GDT-
plus. The results demonstrate that GDT-plus achieves comparable or superior
performance to StAR on all four Atari games, emphasizing the significance of
fine-grained information in the games. Compared with GDT, the success of the
Sequence Transformer in incorporating such information into action prediction
is also highlighted.

###  4.2 OpenAI Gym

Table 2: Results for D4RL datasets. The performance of GDT is evaluated using
three different seeds, and the mean and variance are reported. Best mean
scores are highlighted in bold. The results demonstrate that GDT exhibits
superior performance compared to conventional RL algorithms on most of the
evaluated tasks, and GDT-plus performs comparably to GDT.

Dataset |  Environment |  CQL |  BEAR |  BRAC-v |  AWR |  MBOP |  BC |  DT |  StAR |  GDT |  GDT-plus
---|---|---|---|---|---|---|---|---|---|---|---
Medium-Expert |  HalfCheetah |  62.4 |  53.4 |  41.9 |  52.7 |  105.9 |  59.9 |  86.8 $\pm$ 1.3 |  93.7 $\pm$ 0.1 |  92.4 $\pm$ 0.1 |  93.2 $\pm$ 0.1
Medium-Expert |  Hopper |  111.0 |  96.3 |  0.8 |  27.1 |  55.1 |  79.6 |  107.6 $\pm$ 1.8 |  111.1 $\pm$ 0.2 |  110.9 $\pm$ 0.1 |  111.1 $\pm$ 0.1
Medium-Expert |  Walker |  98.7 |  40.1 |  81.6 |  53.8 |  70.2 |  36.6 |  108.1 $\pm$ 0.2 |  109.0 $\pm$ 0.1 |  109.3 $\pm$ 0.1 |  107.7 $\pm$ 0.1
Medium |  HalfCheetah |  44.4 |  41.7 |  46.3 |  37.4 |  44.6 |  43.1 |  42.6 $\pm$ 0.1 |  42.9 $\pm$ 0.1 |  42.9 $\pm$ 0.1 |  42.9 $\pm$ 0.1
Medium |  Hopper |  58.0 |  52.1 |  31.1 |  35.9 |  48.8 |  63.9 |  67.6 $\pm$ 1.0 |  59.5 $\pm$ 4.2 |  65.8 $\pm$ 5.8 |  77.1 $\pm$ 2.5
Medium |  Walker |  79.2 |  59.1 |  81.1 |  17.4 |  41.0 |  77.3 |  74.0 $\pm$ 1.4 |  73.8 $\pm$ 3.5 |  77.8 $\pm$ 0.4 |  76.5 $\pm$ 0.7
Medium-Replay |  HalfCheetah |  46.2 |  38.6 |  47.7 |  40.3 |  42.3 |  4.3 |  36.6 $\pm$ 0.8 |  36.8 $\pm$ 3.3 |  39.9 $\pm$ 0.1 |  40.5 $\pm$ 0.1
Medium-Replay |  Hopper |  48.6 |  33.7 |  0.6 |  28.4 |  12.4 |  27.6 |  82.7 $\pm$ 7.0 |  29.2 $\pm$ 4.3 |  81.6 $\pm$ 0.6 |  85.3 $\pm$ 25.2
Medium-Replay |  Walker |  26.7 |  19.2 |  0.9 |  15.5 |  9.7 |  36.9 |  66.6 $\pm$ 3.0 |  39.8 $\pm$ 5.1 |  74.8 $\pm$ 1.9 |  77.5 $\pm$ 1.3
Average |  63.9 |  48.2 |  36.9 |  34.3 |  47.8 |  46.4 |  74.7 |  66.2 |  76.8 |  79.1

The D4RL benchmark [14] evaluates the performance of RL algorithms in
continuous control tasks, particularly in robotic manipulation tasks with
challenging control and decision-making in continuous action spaces. In this
study, we select three games from the standard locomotion environments in Gym,
namely HalfCheetah, Hopper, and Walker, and three different dataset settings,
namely medium, medium-replay, and medium-expert. To ensure fair comparisons,
we also normalize the scores according to the protocol established in [14],
where a score of 100 corresponds to an expert policy.

For baseline benchmark, we compare GDT with a range of model-free methods,
including CQL [30], BEAR [28], BRAC [52], AWR [36], and model-based methods,
including MBOP [3], a current strong baseline with model-based offline RL
algorithms. We also compare the performance of GDT with several imitation
learning algorithms, including DT [9], StARformer [42], and BC. The
performance of CQL, MBOP, and DT are reported from the original papers, while
the results of BEAR, BRAC, and AWR are reported from the D4RL paper, and the
other methods are run by us for a better fair comparison. This comprehensive
evaluation encompasses a variety of techniques and conducts a thorough
examination of the effectiveness of GDT in comparison to the existing state-
of-the-art algorithms.

The results presented in Table 2 demonstrate the superior performance of GDT
on most of the evaluated tasks, with comparable performance to the state-of-
the-art algorithm on the remaining tasks. In the Gym environment, the input
state is represented as a vector rather than an image, as is the case in the
Atari environment. Therefore, when introducing the Sequence Transformer, the
fine-grained spatial information is still obtained by embedding the state
vector. In this context, the performance of GDT-plus has not shown significant
improvement compared to GDT, and in some cases even achieves negative effects,
which is also reflected in the results of StAR. This may be attributed to the
increased complexity of the model due to the Sequence Transformer, leading to
overfitting on the state vector and producing a counterproductive effect.

Figure 7: The performance comparison of different input lengths
$T\in\\{10,20,30\\}$. Figure 8: Performance comparison of Sequence Transformer
with different connection methods. The left panel illustrates the three
connection methods, namely Replace, Fusion, and Stack. The right panel shows
the corresponding performance comparison.

###  4.3 Ablations

In comparison to the DT model, GDT transforms the input sequence into a graph
with a causal relationship and feeds it into the Graph Transformer. Therefore,
we initially investigate the impact of the graph structure on the performance
of GDT. Additionally, we examine the influence of reward settings, which have
a greater effect on the performance of the DT model. Lastly, we conduct
ablation experiments on the length of the input sequence to examine the Graph
Transformer’s ability to rely on long sequences. To capture fine-grained
spatial information, we introduce an additional Sequence Transformer to refine
the action prediction process. We explore three ways to connect the Sequence
Transformer. To provide a comprehensive evaluation, all ablation experiments
will be conducted on both Atari and OpenAI Gym (Medium) environments.

Graph representation setting. As we transform the input sequence into a graph
with causal relationships, we aim to investigate the impact of the graph
structure on the overall performance. Given that the graph connection method
has a complexity of $\mathcal{O}(n^{2})$, we primarily explore four different
connection methods: causal connection, full connection, random connection, and
none connection. The results in Figure 5 demonstrate that the causal
connection method outperforms the other methods in both environments. The full
and none connection methods yield similar performances, while the random
Connection method is significantly impacted by the environment. Although this
connection method can still achieve better results when historical information
is crucial, it can be counterproductive in environments that are more reliant
on causal relationships.

Reward setting. In this study, we investigate the impact of reward setting on
GDT’s performance. Specifically, we examine the effect of three reward
settings, namely return-to-go (RTG), stepwise reward (Reward), and none reward
(None). RTG is a commonly used setting in sequence modeling methods [13, 32,
46], and its experimental results are greatly influenced by the value of the
return-to-go hyper-parameter. Reward refers to the immediate reward generated
by the environment after each step, which is commonly used in traditional TD-
learning-based algorithms [20, 40, 18]. The None setting usually corresponds
to straightforward BC.

The results are presented in Figure 6. It should be noted that since GDT takes
a graph with a causal relationship as input, the causal relationship changes
correspondingly for each reward setting, as shown in the left part of the
figure. The introduction of rewards has improved the performance of both
methods compared to the no reward setting. However, DT is more reliant on
return-to-go than GDT, and there is a large difference in performance between
the two reward settings for DT. Importantly, GDT still outperforms DT with the
return-to-go setting under the stepwise reward setting. This demonstrates that
introducing causality in the input can reduce the dependence of sequence
modeling on return-to-go. Furthermore, under the stepwise reward setting,
GDT’s graph input form is similar to the dynamic modeling method in model-
based algorithms, which is worth further research to expand and introduce in
model-based approaches.

Trajectory input length. In this experiment, we investigate the effect of
input sequence length on the performance of GDT. While in a Markovian
environment, the state at the previous moment is often sufficient to determine
the current action, the DT experiment [9] reveals that past information is
valuable for the sequence modeling method in Atari environments, where longer
sequences tend to yield better results than those of length $1$. Subsequently,
we explore the impact of different sequence lengths on performance and compare
the results of DT and GDT, demonstrating the superior ability of GDT to handle
long sequence inputs.

Figure 7 presents the impact of sequence length on DT and GDT performance. It
is observed that DT’s performance remains relatively unchanged as the sequence
length increases, while GDT exhibits a significant improvement, demonstrating
its superior ability to handle long sequences. By incorporating causal
relationships in the input, GDT enables the Graph Transformer to effectively
handle the dependencies between long sequences, leading to improved
performance without additional computational overhead.

Sequence Transformer connection method. To capture fine-grained spatial
information, we introduce an optional Sequence Transformer to improve action
prediction. We denote the output variables of Graph Transformer as $g_{t}^{l}$
and the corresponding input variables in Sequence Transformer as $h_{t}^{l}$.
There are three possible ways to connect the two: (1) in each layer,
$g_{t}^{l}$ replaces $h_{t}^{l-1}$ (GDT-Replace); (2) in each layer,
$g_{t}^{l}$ is added to $h_{t}^{l-1}$ (GDT-Fusion); and (3) the last layer of
$g_{t}^{L}$ is used as the initial layer $h_{t}^{0}$ (GDT-Stack). The
experimental results in Figure 8 demonstrate that in the Atari environment
with high information density, stacking features at the end is the most
effective approach for feature refinement. In contrast, for the Gym
environment with less information, feature refinement through the fusion of
different abstraction levels achieves better performance.

##  5 Conclusion

In conclusion, this paper proposes the Graph Decision Transformer (GDT), a
novel offline RL approach that models input sequences into causal graphs to
capture potential dependencies between fundamentally different concepts and
facilitate temporal and causal relationships learning. We show that GDT
outperforms state-of-the-art offline RL methods on image-based Atari and
OpenAI Gym.

Our work highlights the potential of graph-structured inputs in RL, which has
been less explored compared to other deep learning domains. We believe that
our proposed GDT approach can inspire further research in RL, especially in
tasks where spatial and temporal dependencies are essential, such as robotics,
autonomous driving, and video games. Further exploration of graph-structured
inputs may lead to more efficient and effective RL algorithms with broader
applicability in real-world scenarios.

## References

  * [1] Pieter Abbeel and Andrew Y Ng.  Apprenticeship learning via inverse reinforcement learning.  In ICML, 2004.
  * [2] Rishabh Agarwal, Dale Schuurmans, and Mohammad Norouzi.  An optimistic perspective on offline reinforcement learning.  In ICML, 2020.
  * [3] Arthur Argenson and Gabriel Dulac-Arnold.  Model-based offline planning.  arXiv preprint arXiv:2008.05556, 2020.
  * [4] Peter W Battaglia, Jessica B Hamrick, Victor Bapst, Alvaro Sanchez-Gonzalez, Vinicius Zambaldi, Mateusz Malinowski, Andrea Tacchetti, David Raposo, Adam Santoro, Ryan Faulkner, et al.  Relational inductive biases, deep learning, and graph networks.  arXiv preprint arXiv:1806.01261, 2018.
  * [5] Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling.  The arcade learning environment: An evaluation platform for general agents.  JAIR, 2013.
  * [6] Wendelin Böhmer, Vitaly Kurin, and Shimon Whiteson.  Deep coordination graphs.  In ICML. PMLR, 2020.
  * [7] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba.  Openai gym.  arXiv preprint arXiv:1606.01540, 2016.
  * [8] Deng Cai and Wai Lam.  Graph transformer for graph-to-sequence learning.  In AAAI, 2020.
  * [9] Lili Chen, Kevin Lu, Aravind Rajeswaran, Kimin Lee, Aditya Grover, Misha Laskin, Pieter Abbeel, Aravind Srinivas, and Igor Mordatch.  Decision transformer: Reinforcement learning via sequence modeling.  NeurIPS, 2021.
  * [10] Will Dabney, Mark Rowland, Marc Bellemare, and Rémi Munos.  Distributional reinforcement learning with quantile regression.  In AAAI, 2018.
  * [11] Hanjun Dai, Hui Li, Tian Tian, Xin Huang, Lin Wang, Jun Zhu, and Le Song.  Adversarial attack on graph structured data.  In ICML, 2018.
  * [12] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, et al.  An image is worth 16x16 words: Transformers for image recognition at scale.  arXiv preprint arXiv:2010.11929, 2020.
  * [13] Ben Eysenbach, Xinyang Geng, Sergey Levine, and Russ R Salakhutdinov.  Rewriting history with inverse rl: Hindsight inference for policy improvement.  NeurIPS, 2020.
  * [14] Justin Fu, Aviral Kumar, Ofir Nachum, George Tucker, and Sergey Levine.  D4rl: Datasets for deep data-driven reinforcement learning.  arXiv preprint arXiv:2004.07219, 2020.
  * [15] Scott Fujimoto and Shixiang Shane Gu.  A minimalist approach to offline reinforcement learning.  NeurIPS, 2021.
  * [16] Scott Fujimoto, David Meger, and Doina Precup.  Off-policy deep reinforcement learning without exploration.  In ICML, 2019.
  * [17] Hiroki Furuta, Yutaka Matsuo, and Shixiang Shane Gu.  Generalized decision transformer for offline hindsight information matching.  arXiv preprint arXiv:2111.10364, 2021.
  * [18] Tuomas Haarnoja, Aurick Zhou, Kristian Hartikainen, George Tucker, Sehoon Ha, Jie Tan, Vikash Kumar, Henry Zhu, Abhishek Gupta, Pieter Abbeel, et al.  Soft actor-critic algorithms and applications.  arXiv preprint arXiv:1812.05905, 2018.
  * [19] Danijar Hafner, Timothy Lillicrap, Mohammad Norouzi, and Jimmy Ba.  Mastering atari with discrete world models.  arXiv preprint arXiv:2010.02193, 2020.
  * [20] Matteo Hessel, Joseph Modayil, Hado Van Hasselt, Tom Schaul, Georg Ostrovski, Will Dabney, Dan Horgan, Bilal Piot, Mohammad Azar, and David Silver.  Rainbow: Combining improvements in deep reinforcement learning.  In AAAI, 2018.
  * [21] Shengchao Hu, Li Shen, Ya Zhang, Yixin Chen, and Dacheng Tao.  On transforming reinforcement learning by transformer: The development trajectory.  arXiv preprint arXiv:2212.14164, 2022.
  * [22] Hanchi Huang, Deheng Ye, Li Shen, and Wei Liu.  Curriculum-based asymmetric multi-task reinforcement learning.  IEEE Transactions on Pattern Analysis and Machine Intelligence, 2022\.
  * [23] Wenlong Huang, Igor Mordatch, and Deepak Pathak.  One policy to control them all: Shared modular policies for agent-agnostic control.  In ICML, 2020.
  * [24] Michael Janner, Qiyang Li, and Sergey Levine.  Offline reinforcement learning as one big sequence modeling problem.  NeurIPS, 2021.
  * [25] Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims.  Morel: Model-based offline reinforcement learning.  NeurIPS, 2020.
  * [26] Ilya Kostrikov, Rob Fergus, Jonathan Tompson, and Ofir Nachum.  Offline reinforcement learning with fisher divergence critic regularization.  In ICML, 2021.
  * [27] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine.  Stabilizing off-policy q-learning via bootstrapping error reduction.  NeurIPS, 2019.
  * [28] Aviral Kumar, Justin Fu, Matthew Soh, George Tucker, and Sergey Levine.  Stabilizing off-policy q-learning via bootstrapping error reduction.  NeurIPS, 2019.
  * [29] Aviral Kumar, Xue Bin Peng, and Sergey Levine.  Reward-conditioned policies.  arXiv preprint arXiv:1912.13465, 2019.
  * [30] Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine.  Conservative q-learning for offline reinforcement learning.  NeurIPS, 2020.
  * [31] Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu.  Offline reinforcement learning: Tutorial, review, and perspectives on open problems.  arXiv preprint arXiv:2005.01643, 2020.
  * [32] Alexander Li, Lerrel Pinto, and Pieter Abbeel.  Generalized hindsight for reinforcement learning.  NeurIPS, 2020.
  * [33] Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Andrei A Rusu, Joel Veness, Marc G Bellemare, Alex Graves, Martin Riedmiller, Andreas K Fidjeland, Georg Ostrovski, et al.  Human-level control through deep reinforcement learning.  nature, 2015.
  * [34] Sai Munikoti, Deepesh Agarwal, Laya Das, Mahantesh Halappanavar, and Balasubramaniam Natarajan.  Challenges and opportunities in deep reinforcement learning with graph neural networks: A comprehensive review of algorithms and applications.  arXiv preprint arXiv:2206.07922, 2022.
  * [35] Andrew Y Ng, Stuart Russell, et al.  Algorithms for inverse reinforcement learning.  In ICML, 2000.
  * [36] Xue Bin Peng, Aviral Kumar, Grace Zhang, and Sergey Levine.  Advantage-weighted regression: Simple and scalable off-policy reinforcement learning.  arXiv preprint arXiv:1910.00177, 2019.
  * [37] Doina Precup.  Eligibility traces for off-policy policy evaluation.  Computer Science Department Faculty Publication Series, 2000.
  * [38] Aditya Ramesh, Mikhail Pavlov, Gabriel Goh, Scott Gray, Chelsea Voss, Alec Radford, Mark Chen, and Ilya Sutskever.  Zero-shot text-to-image generation.  In ICML, 2021.
  * [39] Martin Riedmiller.  Neural fitted q iteration–first experiences with a data efficient neural reinforcement learning method.  In ECML, 2005.
  * [40] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.  Proximal policy optimization algorithms.  arXiv preprint arXiv:1707.06347, 2017.
  * [41] Caihua Shan, Yifei Shen, Yao Zhang, Xiang Li, and Dongsheng Li.  Reinforcement learning enhanced explainer for graph neural networks.  NeurIPS, 2021.
  * [42] Jinghuan Shang, Kumara Kahatapitiya, Xiang Li, and Michael S Ryoo.  Starformer: Transformer with state-action-reward representations for visual reinforcement learning.  In ECCV, 2022.
  * [43] Siqi Shen, Yongquan Fu, Huayou Su, Hengyue Pan, Peng Qiao, Yong Dou, and Cheng Wang.  Graphcomm: A graph neural network based method for multi-agent reinforcement learning.  In ICASSP, 2021.
  * [44] Noah Y Siegel, Jost Tobias Springenberg, Felix Berkenkamp, Abbas Abdolmaleki, Michael Neunert, Thomas Lampe, Roland Hafner, Nicolas Heess, and Martin Riedmiller.  Keep doing what worked: Behavioral modelling priors for offline reinforcement learning.  arXiv preprint arXiv:2002.08396, 2020.
  * [45] David Silver, Julian Schrittwieser, Karen Simonyan, Ioannis Antonoglou, Aja Huang, Arthur Guez, Thomas Hubert, Lucas Baker, Matthew Lai, Adrian Bolton, et al.  Mastering the game of go without human knowledge.  nature, 2017.
  * [46] Rupesh Kumar Srivastava, Pranav Shyam, Filipe Mutz, Wojciech Jaśkowski, and Jürgen Schmidhuber.  Training agents using upside-down reinforcement learning.  arXiv preprint arXiv:1912.02877, 2019.
  * [47] Yiwei Sun, Suhang Wang, Xianfeng Tang, Tsung-Yu Hsieh, and Vasant Honavar.  Non-target-specific node injection attacks on graph neural networks: A hierarchical reinforcement learning approach.  In Proc. WWW, 2020.
  * [48] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin.  Attention is all you need.  NeurIPS, 2017.
  * [49] Adam R Villaflor, Zhe Huang, Swapnil Pande, John M Dolan, and Jeff Schneider.  Addressing optimism bias in sequence modeling for reinforcement learning.  In ICML, 2022.
  * [50] Kerong Wang, Hanye Zhao, Xufang Luo, Kan Ren, Weinan Zhang, and Dongsheng Li.  Bootstrapped transformer for offline reinforcement learning.  arXiv preprint arXiv:2206.08569, 2022.
  * [51] Tingwu Wang, Renjie Liao, Jimmy Ba, and Sanja Fidler.  Nervenet: Learning structured policy with graph neural networks.  In ICLR, 2018.
  * [52] Yifan Wu, George Tucker, and Ofir Nachum.  Behavior regularized offline reinforcement learning.  arXiv preprint arXiv:1911.11361, 2019.
  * [53] Taku Yamagata, Ahmed Khalil, and Raul Santos-Rodriguez.  Q-learning decision transformer: Leveraging dynamic programming for conditional sequence modelling in offline rl.  arXiv preprint arXiv:2209.03993, 2022.
  * [54] Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma.  Mopo: Model-based offline policy optimization.  NeurIPS, 2020.
  * [55] Qin Zhang, Linrui Zhang, Haoran Xu, Li Shen, Bowen Wang, Yongzhe Chang, Xueqian Wang, Bo Yuan, and Dacheng Tao.  Saformer: A conditional sequence modeling approach to offline safe reinforcement learning.  arXiv preprint arXiv:2301.12203, 2023.
  * [56] Xianjie Zhang, Yu Liu, Xiujuan Xu, Qiong Huang, Hangyu Mao, and Anil Carie.  Structural relational inference actor-critic for multi-agent reinforcement learning.  Neurocomputing, 2021.
  * [57] Qinqing Zheng, Amy Zhang, and Aditya Grover.  Online decision transformer.  In ICML, 2022.
  * [58] Kaixiong Zhou, Qingquan Song, Xiao Huang, and Xia Hu.  Auto-gnn: Neural architecture search of graph neural networks.  arXiv preprint arXiv:1909.03184, 2019.

Appendix

##  Appendix A Experimental Details

###  A.1 Baseline

For the Atari environment, we evaluate GDT against several state-of-the-art
non-Transformer offline RL methods, such as CQL[30], QR-DQN[10], and REM [2],
and several imitation learning algorithms, including DT [9], StARformer [42],
and straightforward behavior cloning. We report results from the corresponding
papers for CQL, REM, and QR-DQN. For DT, there is a slight discrepancy between
[9] and [42]; we report raw data provided to us by DT authors. As StARformer
performs well on Atari, we use it as the main comparison object for GDT-plus.
In our behavior cloning setting, the agent lacks access to reward signals and
online data from the environment, making the problem even more challenging.
This differs from traditional imitation learning approaches that can collect
new data and perform Inverse Reinforcement Learning [1, 35]. To create this
setting, we remove the rewards from the dataset used in offline RL.

For the Gym environment, we compare GDT with a range of model-free methods,
including CQL [30], BEAR [28], BRAC [52], and AWR [36], and model-based
methods, including MBOP [3]. We also compare the performance of GDT with
several imitation learning algorithms, including DT [9], StARformer [42], and
BC. Results for CQL, MBOP, and DT are reported from the original papers, while
the results of BEAR, BRAC, and AWR are reported from the D4RL paper. Note that
we re-run the experimental results on Gym for StARformer based on the official
code published by the authors, as the original paper used the DeepMind Control
Suite (DMC) environment with image input instead of vector input like Gym.

###  A.2 Training Resources

We use one NVIDIA A100 GPU (SXM4) to train each model. Training each model
typically takes 8-20 hours. However, since each environment needs to be
trained three times with different seeds, the total training time is usually
multiplied by three.

##  Appendix B Detailed results

Table 3 presents the normalized scores used for normalization, as proposed in
[19, 14]. Tables 4 and 5 show the raw scores corresponding to Tables 1 and 2.
We note a slight discrepancy between [2] and [30] for REM and QR-DQN, and
report the raw data provided to us by the REM authors. To alleviate the
complexity of presenting a large amount of experimental data, we report only
the mean in the tables.

Table 3: Baseline scores used for normalization.
Game | Random | Gamer
---|---|---
Breakout | $2$ | $30$
Qbert | $164$ | $13455$
Pong | $-21$ | $15$
Seaquest | $68$ | $42055$
HalfCheet | $-280.2$ | $12135$
Hopper | $-20.3$ | $3234.3$
Walker | $1.6$ | $4592.3$
Table 4:  Raw scores for 1% DQN-replay Atari datasets. We evaluate the
performance of GDT on four Atari games using three different seeds, and report
the mean and variance of the results. Best mean scores are highlighted in
bold. The assessment reveals that GDT surpasses conventional RL algorithms on
most tasks and achieves better performance than DT across all games. In
contrast, GDT-plus achieves comparable results to StAR, leveraging fine-
grained spatial information.

Game | CQL | QR-DQN | REM | BC | DT | StAR | GDT | GDT-plus
---|---|---|---|---|---|---|---|---
Breakout | $61.1$ | $6.8$ | $4.5$ | $40.9\pm 17.3$ | $76.9\pm 27.3$ | $124.1\pm 19.8$ | $112.2\pm 7.7$ | $\bf{125.7}\pm 3.2$
Qbert | $\bf{14012.0}$ | $156.0$ | $160.1$ | $2464.1\pm 1948.2$ | $2215.8\pm 1523.7$ | $6968.0\pm 1698.0$ | $6205.8\pm 2104.5$ | $7029.2\pm 2928.5$
Pong | $\bf{19.3}$ | $-14.5$ | $-20.8$ | $9.7\pm 7.2$ | $17.1\pm 2.9$ | $18.9\pm 0.7$ | $18.0\pm 0.6$ | $19.0\pm 0.1$
Seaquest | $779.4$ | $250.1$ | $370.5$ | $968.6\pm 133.8$ | $1129.3\pm 189.0$ | $781\pm 212.0$ | $\bf{1261.3}\pm 204.3$ | $1188\pm 387.4$

Table 5: Raw scores for D4RL datasets. The performance of GDT is evaluated
using three different seeds, and the mean is reported. Best mean scores are
highlighted in bold. The results demonstrate that GDT exhibits superior
performance compared to conventional RL algorithms on most of the evaluated
tasks, and GDT-plus performs comparably to GDT.

Dataset |  Environment |  CQL |  BEAR |  BRAC-v |  AWR |  MBOP |  BC |  DT |  StAR |  GDT |  GDT-p
---|---|---|---|---|---|---|---|---|---|---|---
Medium-Expert |  HalfCheetah |  7466.9 |  6349.5 |  4921.8 |  6262.6 |  12867.5 |  7156.5 |  10496.2 |  11352.8 |  11191.4 |  11290.8
Medium-Expert |  Hopper |  3592.3 |  3113.9 |  5.8 |  861.7 |  1773 |  2570.4 |  3481.6 |  3595.6 |  3589 |  3595.6
Medium-Expert |  Walker |  4532.6 |  1842.5 |  3747.6 |  2471.4 |  3224.3 |  1681.8 |  4964.1 |  5005.5 |  5019.2 |  4945.8
Medium |  HalfCheetah |  5232.2 |  4897 |  5468 |  4363.1 |  5257 |  5070.8 |  5008.7 |  5045.9 |  5045.9 |  5045.9
Medium |  Hopper |  1867.4 |  1675.4 |  991.9 |  1148.1 |  1568 |  2059.4 |  2179.8 |  1916.2 |  2121.2 |  2489
Medium |  Walker |  3637.4 |  2714.7 |  3724.7 |  800.4 |  1883.8 |  3550.2 |  3398.7 |  3389.5 |  3573.2 |  3513.5
Medium-Replay |  HalfCheetah |  5455.6 |  4512.1 |  5641.9 |  4723.1 |  4971.4 |  253.7 |  4263.8 |  4288.6 |  4673.5 |  4748
Medium-Replay |  Hopper |  1561.4 |  1076.5 |  -0.7 |  904 |  383.3 |  878 |  2671.3 |  930.1 |  2635.5 |  2755.9
Medium-Replay |  Walker |  1227.3 |  883 |  42.9 |  713.2 |  446.9 |  1695.6 |  3059 |  1828.7 |  3435.5 |  3559.4

##  Appendix C Hyper-parameters

Tables 6, 7, 8, and 9 provide a comprehensive list of hyper-parameters for our
proposed GDT and GDT-plus models applied to Atari and OpenAI Gym environments.
To ensure a fair comparison, we adopt similar hyper-parameter settings to
Decision-Transformer [9], including the number of Transformer layers, multi-
head self-attention heads, and embedding dimensions in our Graph Transformer,
as well as learning rate and optimizer configurations.

Table 6: Hyperparameters of GDT for Atari experiments.
Hyperparameter | Value
---|---
Layers (Graph Transformer) | $6$
MSA heads (Graph Transformer) | $8$
Embedding dimension (Graph Transformer) | $128$
Batch size | $128$
Context length $K$ |  $50$ Pong
|  $30$ Breakout, Qbert, Seaquest
Return-to-go conditioning |  $120$ Breakout
|  $5000$ Qbert
|  $20$ Pong
|  $1450$ Seaquest
Nonlinearity | ReLU, encoder
| GeLU, otherwise
Max epochs | $10$
Dropout | $0.1$
Learning rate | $6*10^{-4}$
Adam betas | $(0.9,0.95)$
Grad norm clip | $1.0$
Weight decay | $0.1$
Learning rate decay | Linear warmup and cosine decay (see code for details)
Warmup tokens | $512*20$
Final tokens | $6*500000*K$
Table 7: Hyperparameters of GDT-plus for Atari experiments.
Hyperparameter | Value
---|---
Layers (Graph Transformer) | $6$
MSA heads (Graph Transformer) | $8$
Embedding dimension (Graph Transformer) | $128$
Connection method | Stack
Layers (Sequence Transformer) | $2$
Image patch size | 14
MSA heads (Sequence Transformer) | $4$
Embedding dimension (Sequence Transformer) | $64$
Batch size | $64$
Context length $K$ |  $50$ Pong
|  $30$ Breakout, Qbert, Seaquest
Return-to-go conditioning |  $120$ Breakout
|  $5000$ Qbert
|  $20$ Pong
|  $1450$ Seaquest
Nonlinearity | ReLU, encoder
| GeLU, otherwise
Max epochs | $10$
Dropout | $0.1$
Learning rate | $6*10^{-4}$
Adam betas | $(0.9,0.95)$
Grad norm clip | $1.0$
Weight decay | $0.1$
Learning rate decay | Linear warmup and cosine decay (see code for details)
Warmup tokens | $512*20$
Final tokens | $10*500000*K$
Table 8: Hyperparameters of GDT for OpenAI Gym experiments.
Hyperparameter | Value
---|---
Layers (Graph Transformer) | $6$
MSA heads (Graph Transformer) | $8$
Embedding dimension (Graph Transformer) | $128$
Nonlinearity function | ReLU
Batch size | $64$
Context length $K$ | $20$
Return-to-go conditioning |  $6000$ HalfCheetah
|  $3600$ Hopper
|  $5000$ Walker
Dropout | $0.1$
Learning rate | $10^{-4}$
Grad norm clip | $0.25$
Weight decay | $10^{-4}$
Learning rate decay |  Linear warmup for first $10^{5}$ training steps
Table 9: Hyperparameters of GDT-plus for OpenAI Gym experiments.
Hyperparameter | Value
---|---
Layers (Graph Transformer) | $6$
MSA heads (Graph Transformer) | $8$
Embedding dimension (Graph Transformer) | $128$
Connection method | Fusion
Layers (Sequence Transformer) | $6$
MSA heads (Sequence Transformer) | $8$
Embedding dimension (Sequence Transformer) | $256$
Nonlinearity function | ReLU
Batch size | $64$
Context length $K$ | $20$
Return-to-go conditioning |  $6000$ HalfCheetah
|  $3600$ Hopper
|  $5000$ Walker
Dropout | $0.1$
Learning rate | $10^{-4}$
Grad norm clip | $0.25$
Weight decay | $10^{-4}$
Learning rate decay |  Linear warmup for first $10^{5}$ training steps