# MQT Bench: Benchmarking Software and
Design Automation Tools for Quantum Computing

Nils Quetschlich1 Lukas Burgholzer2 Robert Wille13  1Chair for Design
Automation, Technical University of Munich, Germany  2Institute for Integrated
Circuits, Johannes Kepler University Linz, Austria  3Software Competence
Center Hagenberg GmbH (SCCH), Austria<EMAIL_ADDRESS><EMAIL_ADDRESS><EMAIL_ADDRESS>
https://www.cda.cit.tum.de/research/quantum

###### Abstract

Quantum software tools for a wide variety of design tasks on and across
different levels of abstraction are crucial in order to eventually realize
useful quantum applications. This requires practical and relevant benchmarks
for new software tools to be empirically evaluated and compared to the current
state of the art. While benchmarks for specific design tasks are commonly
available, the demand for an overarching cross-level benchmark suite is not
fully met yet and there is no mutual consolidation in how quantum software
tools are evaluated thus far. In this work, we propose the _MQT Bench_
benchmark suite (as part of the _Munich Quantum Toolkit_ (MQT)) based on four
core traits: (1) cross-level support for different abstraction levels, (2)
accessibility via an easy to use web interface
(https://www.cda.cit.tum.de/mqtbench) and a Python package, (3) provision of a
broad selection of benchmarks to facilitate generalizability, as well as (4)
extendability to future algorithms, gate-sets, and hardware architectures. By
comprising more than $50{\small,}000$ benchmark circuits ranging from $2$ up
to $130$ qubits on four abstraction levels, MQT Bench presents a first step
towards benchmarking different abstraction levels with a single benchmark
suite to increase comparability, reproducibility, and transparency.

###### Index Terms:

quantum computing, benchmarking, design automation

##  I Introduction

Quantum computing has gained a lot of attention due to its promising
applications and the advances in quantum computing hardware. But currently,
designing quantum algorithms often means to manually implement quantum
circuits on the gate-level—similar to assembly language in classical
computing. Hence, there is an urgent need to enhance the workflow of designing
and testing quantum algorithms. Without software tools aiding the design of
quantum algorithms, the underlying hardware may not be efficiently utilized.
Thus, researchers and developers have already proposed tools for various
design tasks, such as quantum circuit simulation [1, 2], compilation [3, 4,
5], or verification [6, 7, 8, 9, 10]. This already led to comprehensive
quantum circuit design flows realized through toolkits such as IBM’s Qiskit
[11], Google’s Cirq [12], or Rigetti’s Forest [13]—in addition to numerous
tools and methods developed by other researchers and engineers in the field.

Generally, these software tools operate on and across different levels of
abstraction and all tackle computationally hard problems. As a result, they
have to compromise between resource demands and result quality. Usually, when
a new software tool is proposed, its performance is empirically evaluated. The
question how to evaluate the performance of a software tool is very
challenging and not yet fully answered—especially for tools solving NP-
complete problems. The most common approach is to run certain problem
instances, so-called benchmarks, and compare the performance of the newly
proposed method against state-of-the-art methods regarding a certain
characteristic, e.g., run-time or solution quality.

Currently, this benchmarking is conducted using a wide variety of different
benchmark suites—each with a specific focus. While some benchmarks suites,
e.g., [14], provide application-oriented benchmarks on a higher abstraction
level, other benchmark suites focus on lower abstraction levels, e.g., [15,
16, 17]. Although the intended target levels are properly covered, the demand
for a cross-level benchmark suite is not fully met yet. As another
consequence, there is no mutual consolidation which benchmarks to use for
empirical evaluations yet—leading to lower comparability, reproducibility, and
transparency.

In this paper, we propose _MQT Bench_ —another benchmark suite (as discussed
above not the first one and most likely not the last one) which explicitly
aims to address those drawbacks. It aims at providing a first step towards
benchmarking the whole quantum software stack with a single benchmark suite by
offering the same benchmark algorithms on different levels of abstractions. To
realize such a benchmark suite, several challenges have to be tackled:

  * •

Distinct requirements on different levels: Benchmarks have to fulfill certain
requirements depending on the abstraction level, e.g., only gates of a
device’s native gate-set are allowed.

  * •

Accessibility: To foster adoption of the benchmark suite, it needs to be as
easy to use as possible.

  * •

Generalizability: Providing a comprehensive set of benchmark algorithms to
cover as many use cases as possible.

  * •

Extendability: Future algorithms, gate-sets, and architectures should be easy
to integrate.

MQT Bench has been developed with these challenges in mind and is based on
four core traits:

  1. 1.

Cross-level benchmarking: Provision of benchmarks on four abstraction levels.

  2. 2.

Accessibility: To simplify the usage of MQT Bench, we provide a web interface
(https://www.cda.cit.tum.de/mqtbench) and a Python package in addition to the
open-source repository on GitHub (https://github.com/cda-tum/MQTBench).

  3. 3.

Algorithm selection: Broad selection of different algorithms ranging from
building blocks, i.e., _Quantum Fourier Transform_ (QFT), to applications,
i.e., ground state estimation.

  4. 4.

Extendability: Algorithms, gate-sets, and hardware architectures are easily
extendable and integrable into MQT Bench.

By this, MQT Bench aims to improve the comparability, reproducibility, and
transparency of empirical evaluations for the whole quantum software stack.

The rest of this work is structured as follows: In Section II, we review the
necessary basics of the quantum computing compilation flow and the respective
software stack to keep this work self-contained. Afterwards, Section III
reviews the state of the art on how this quantum software stack is currently
evaluated and benchmarked—motivating the benchmark suite proposed in this
paper. Based on that, Section IV introduces the resulting benchmark suite and
its core traits, before Section V concludes this work.

##  II Background

To keep this paper self contained, this section gives a brief overview of the
quantum circuit compilation flow with its different abstraction levels and the
respective quantum software stack. The levels mentioned in this section are
inspired by the structure proposed by the openQASM 3.0 specification [18] and
can similarly be found in many other compilation flows as well.

###  II-A Quantum Circuit Compilation Flow

Similar to the classical domain, executing a conceptual quantum algorithm on
an actual device requires compiling it to a representation that adheres to all
constraints imposed by the hardware. Usually, quantum algorithms are initially
developed and tested at a hardware-agnostic level. Generally, this is
described as a quantum circuit that consists of high-level building blocks
without any restrictions to a certain gate-set or hardware architecture and is
defined as the _algorithmic level_.

######  Example 1.

_Variational Quantum Algorithms_ (VQAs) are an emerging class of quantum
algorithms with a wide range of applications. A respective circuit is depicted
in 2(a) and shows an example of an ansatz function frequently used for
_Variational Quantum Eigensolvers_ (VQEs), a subclass of VQAs. On this
abstraction level, the circuit is parameterized by the angles $\theta_{i}$ of
the six single-qubit gates.

The first step in realizing a conceptual quantum algorithm on the algorithmic
level for a particular problem is to synthesize the high-level building blocks
and optimize the resulting representation independent from an actual target
architecture. In analogy to a classical compiler, this involves tasks such as
constant propagation and folding, gate modifier evaluation, synthesis, loop
unrolling, and gate simplification. The resulting representation is defined as
the _target-independent level_.

######  Example 2.

VQAs are hybrid quantum-classical algorithms, where the parameters of the
quantum ansatz are iteratively updated by a classical optimizer analogous to
conventional gradient-based optimization. Consider again the circuit from
2(a). Assuming these parameters have been determined, e.g., $\theta_{i}=-\pi$
for $i=0,...,5$, they are now propagated and the resulting quantum circuit is
shown in 2(b).

${q_{2}}$${q_{1}}$${q_{0}}$$R_{y}(\theta_{0})$$R_{y}(\theta_{1})$$R_{y}(\theta_{2})$$R_{y}(\theta_{4})$$R_{y}(\theta_{5})$$R_{y}(\theta_{6})$

(a) Algorithmic level

${q_{2}}$${q_{1}}$${q_{0}}$$R_{y}(\pi)$$R_{y}(\pi)$$R_{y}(\pi)$$R_{y}(\pi)$$R_{y}(\pi)$$R_{y}(\pi)$

(b) Target-independent level

${q_{2}}$${q_{1}}$${q_{0}}$$-\pi$$-\pi$$-\pi$$-\pi$$-\pi$$-\pi$

(c) Target-dependent native gates level

$Q_{0}$$Q_{4}$$Q_{2}$$Q_{1}$$Q_{3}$

(d) Target architecture: IBMQ Manila

${q_{2}}\mapsto Q_{0}$${q_{1}}\mapsto Q_{1}$${q_{0}}\mapsto
Q_{2}$$-\pi$$-\pi$$-\pi$$-\pi$$-\pi$$-\pi$

(e) Target-dependent mapped level

Figure 2: Compilation of a Variational Quantum Algorithm.

Today’s quantum devices impose several constraints on the circuits that may be
executed on them. Thus, a target-dependent compilation phase is necessary. For
one, devices only provide a particular set of native gates—typically
consisting of an entangling gate (like the $CX$ gate) and some family of
single-qubit gates. Hence, the gates of a circuit need to be translated to
this native gate-set, which is typically followed by an optimization pass that
aims to reduce the introduced overhead. This representation level is defined
as the _target-dependent native gates level_.

######  Example 3.

Different quantum computer realizations support different native gate-sets. In
our example, we consider the IBMQ Manila device as the target device which
natively supports $I$, $X$, $\sqrt{X}$, $R_{z}$, and $CX$ gates. Consequently,
the $R_{y}$ gates in 2(b) have to be converted using only these native gates.
In this case, they are substituted by a sequence of $X$ and $R_{z}$ gates
(denoted as $\bullet$ with a phase of $-\pi$) as shown in 2(c).

In addition to the limited gate-set, today’s devices (at least those based on
superconducting qubits) only feature a limited connectivity between their
qubits. Consequently, the circuit’s logical qubits need to be _mapped_ to the
targeted device’s physical qubits, so that multi-qubit gates are only applied
to qubits directly connected on the device. Since such a mapping can rarely be
determined in a static fashion, i.e., globally for the whole circuit, the
mapping has to change dynamically throughout the circuit, which is frequently
referred to as mapping or routing. Again, this is followed by a round of
optimization in order to minimize the overhead caused by the compilation. This
representation level is defined as the _target-dependent mapped level_.

######  Example 4.

Consider again the scenario from Example 3. The architecture of the IBMQ
Manila device is shown in 2(d) and it defines between which qubits a two-qubit
operation may be performed. Since the circuit shown in 2(c) contains $CX$
gates operating between all combination of qubits, there is no mapping
directly matching the target architecture’s layout. As a consequence, a non-
trivial mapping followed by a round of optimization leads to the resulting
circuit shown in 2(e). This is also the reason for the different sequence of
$CX$ gates compared to the previous example.

At this point, the circuit is ready to be sent to the quantum computer’s
backend for execution. From there, the individual gates are scheduled, linked
to a particular calibration, and the resulting circuit is eventually passed to
the target machine code generator. The resulting binaries are then forwarded
to the execution engine to orchestrate the quantum computation on the actual
device.

###  II-B Quantum Software Stack

Quantum software is needed on all levels reviewed above in order to aid
designers in eventually realizing useful quantum applications. In the
following, we review three of the core design tasks for software tools.

As of today, quantum computers are a scarce resource with limited availability
and capability. Until more devices become available which are larger and less
prone to errors, classical means to simulate quantum algorithm circuits (also
known as _quantum circuit simulators_) are essential for fostering the
development of quantum applications. Additionally, and in contrast to actual
quantum computing devices, quantum circuit simulators allow for detailed
insights into the quantum states throughout the circuit execution since the
state’s amplitudes are explicitly tracked during the simulation. On actual
hardware, information can only be extracted in the form of measurements from
the final quantum state, which each yields a classical bit string according to
the distribution described by the state’s amplitudes. Such simulators can be
used on the highest possible level of abstraction, since they can be designed
in a way that does not require restrictions on the circuit’s gate-set or the
qubits’ connectivity, e.g., as proposed in [1, 2]. Also, they might also be
used on the lowest possible level in order to perform noise-aware simulations
which allow to estimate how a circuit is likely to perform on an actual
device, e.g., as proposed in [19, 20, 21]. Either way, on both abstraction
levels, quantum circuit simulation is non-trivial since corresponding
representations of states and operations, in general, require exponential
space or runtime—requiring powerful software tools.

Due to the immense complexity of the tasks involved in quantum circuit
compilation (as it has been illustrated in Section II-A), e.g., mapping being
NP-complete [22], manually conducting compilation often is not an option.
Consequently, efficient software tools are needed across all levels of
abstraction of the aforementioned flow in order to realize a conceptual
algorithm on an actual device. Thus, various compilation tools have been
proposed, e.g., [3, 4, 5].

During compilation, an algorithm’s description is considerably altered and
transformed. Naturally, it is of utmost importance to ensure that the
originally intended functionality is preserved through all levels of
abstraction. This procedure is called verification. While the underlying
principle of comparing the transformations represented by different quantum
circuits is conceptually simple, the exponential size of the underlying
matrices makes this problem challenging—it has been shown to be QMA-complete
[23]. Due to the increasing complexity of today’s compilation flows, tools for
verifying their results become increasingly important. Exemplary verification
tools have been proposed in [6, 7, 8, 9, 10].

These design tasks demonstrate the wide variety of software tools operating on
and across different levels of abstraction—leading to comprehensive quantum
circuit design flows realized through toolkits such as IBM’s Qiskit [11],
Google’s Cirq [12], or Rigetti’s Forest [13]. At the same time, all these
tasks have in common that they have an immensely large complexity. As such, a
multitude of techniques for each task have been proposed—each with its own
trade-off between resource demand and quality of the result. It is key in the
development of scientific methods to empirically evaluate and compare their
performances on practical, relevant benchmarks.

##  III Benchmarking

The benchmark suite proposed in this paper is not the first and certainly not
the last collection of quantum circuit benchmarks. In this section, we review
some of the existing suites and their respective foci. Afterwards, we discuss
how MQT Bench further complements this.

###  III-A Current State of the Art

Providing a comprehensive set of benchmarks to satisfy the needs across all
levels of the quantum circuit compilation flow is a challenging task. Several
approaches to provide benchmark suites for some of the levels within the
quantum circuit compilation flow have already been proposed and a non-
exhaustive overview is given in the following:

_Application-Oriented Performance Benchmarks for Quantum Computing_ [14]:
Starting with an example located on the algorithmic level, this benchmark
suite rather focuses on high-level descriptions. At the moment of writing, it
provides 13 benchmark algorithms which can be generated with a variable number
of qubits via jupyter notebooks using multiple software compilation flows and
are classified into four categories of complexity. The resulting circuits are
available as high-level Python objects without any further compilation steps.

_QASMbench_ [15]: This benchmark suite focusses on the target-independent
level. It offers numerous quantum circuits with a wide but fixed range of both
the number of qubits and depth. It classifies the benchmarks into three
categories of sizes and three categories of algorithm classes. All circuits
are available in an intermediate representation according to the openQASM 2.0
specification [24].

_SupermarQ_ [16]: Similarly to _QASMbench_ , SupermarQ also focusses on the
target-independent level with the contrast of providing benchmarks with
adjustable qubit range for eight algorithms according to the openQASM 2.0
specification [24]. Additionally, six feature vectors are proposed to describe
the characteristics of the benchmarks.

_RevLib_ [17]: This benchmark suite provides a wide variety of reversible
circuits in an intermediate representation format specified by the authors.
Classical functions are frequently embedded into reversible circuits in order
to use them in quantum algorithms, e.g., oracle functions or Boolean building
blocks such as the modular exponentiation in Shor’s algorithm [25]. In
addition, decomposed versions of this benchmark suite have found widespread
use in evaluating the performance of mapping tools, e.g., in [26, 27, 28].
Since reversible circuits merely form a subclass of quantum circuits and do
not employ quantum mechanical effects such as superposition and entanglement,
they do not serve as adequate benchmarks for the whole stack.

###  III-B Motivation

All of the mentioned benchmark suites have in common, that they each target a
specific abstraction level within the quantum circuit compilation flow.
Although the respectively intended target level is well covered, the demand
for a cross-level benchmark suite is not fully met yet. Additionally, so far,
there is no mutual consolidation which benchmarks to use for empirical
evaluations of software tools. This results in a lower comparability,
reproduceability, and transparency of results and may even cause confusion
because benchmarks believed to realize the same particular task are quite
frequently realized in a completely different fashion, e.g., a benchmark
called "grover" could realize Grover’s algorithm with any particular oracle.
Consequently, the demand for a benchmark suite spanning the whole stack of
abstraction levels constitutes the main motivation of this contribution. In
order to realize such a benchmark suite, several challenges have to be
tackled:

  * •

Distinct requirements for benchmark circuits on different levels: The closer a
level is to the actual computing hardware, the more requirements have to be
fulfilled, e.g., only gates from the device’s native gate-set are allowed and
the connectivity of the hardware architecture must be considered.

  * •

Accessibility: In order to facilitate its adoption, the benchmark suite needs
to be as easy to use as possible. While most benchmark suites already provide
open-source access, running and adapting the code to the user’s needs, e.g.,
obtaining a selection of relevant benchmarks, frequently is tedious or
impossible and a solution that keeps this hurdle low is desirable.

  * •

Generalizability of a software tool: This can only be guaranteed if it is
evaluated on a broad set of benchmarks, relating to several characteristics,
e.g., the number of qubits, circuit depth, or the application domain.

  * •

Extendability: Besides that, quantum computing is rapidly advancing and
extendability of the benchmark suite to future algorithms, gate-sets, and
architectures becomes even more important.

MQT Bench has been developed with these challenges in mind and aims to be a
first step towards covering the whole quantum software stack with a benchmark
suite based on the following four core traits (which are described in detail
in Section IV):

  1. 1.

Cross-level benchmarking: All benchmarks are provided on four abstraction
levels reviewed in Section II-A.

  2. 2.

Accessibility: A website to simplify the usage of MQT Bench as much as
possible is provided (https://www.cda.cit.tum.de/mqtbench). Additionally, all
benchmarks can also be generated on-demand using our Python package. While
most needs of users should be covered by this, we additionally give access to
the open-source repository on GitHub (https://github.com/cda-tum/MQTBench).

  3. 3.

Algorithm Selection: Broad selection of different benchmarks with
parameterizable characteristics ranging from building blocks, i.e., QFT, to
applications, i.e., ground state estimation.

  4. 4.

Extendability: MQT Bench is easily extendable regarding the available
benchmarks, native gate-sets, and hardware architectures.

Benchmarking quantum software tools and compilation flows with the same
benchmarks aims to aid comparability, reproducibility and transparency of
empirical evaluations. How this leads to more consistent testing is shown in
the following example.

######  Example 5.

Consider a full quantum compilation flow across all four levels which shall be
developed and tested as depicted in Fig. 3. On the left-hand side of the
levels, the current state of the art of benchmarking is sketched. To
comprehensively test each level, different benchmark suites have to be adapted
and used. In contrast, MQT Bench replaces all of those and enables testing
software tools with the same benchmarks throughout the whole quantum circuit
compilation flow as shown on the right-hand side.

Algorithmic Target-independent Target-dependent native gates Target-dependent
mapped Appl.-orient. Perf. Benchm. [14]QASMbench [15] SupermarQ [16]RevLib
[17]MQT Bench

Directly applicable
Additional effort needed

Figure 3: Benchmarking software tools with and without MQT Bench.

##  IV MQT Bench

In this section, we describe in more detail how each of the four previously
mentioned core traits is implemented.

###  IV-A Cross-Level Benchmarking

As exemplarily illustrated in Section II, software development in quantum
computing takes place at various levels. While there arguably are numerous
possible levels, MQT Bench focusses on four as a balanced measure between
specificity and generizability—inspired by the structure proposed in the
openQASM 3.0 specification [18]. Providing these benchmarks on all of the four
abstraction levels is the main contribution of this work. More precisely, the
following levels are covered:

  1. 1.

_Algorithmic level_ : In this format, all kind of quantum gates may be used
and subsumed into high-level building blocks. Loops are possible and the
circuit structure might be described by constants, variables, and mathematical
expressions. This level provides benchmarks on the most generic level and
specific values are assigned to its variables, e.g., the number of iterations
of a loop. The number of qubits can be specified through the MQT Bench web
interface.

  2. 2.

_Target-independent level_ : Here, loop unrolling, constant folding and
propagation are conducted. Also naive gate simplification is enforced. Again,
the number of qubits can be adapted through the website and the used compiler
can be selected.

  3. 3.

_Target-dependent native gates level_ : On this level, a particular native
gate-set is chosen and the circuit is transpiled and optimized accordingly.
Here, the native gate-set, the used compiler and its settings can be
specified.

  4. 4.

_Target-dependent mapped level_ : Finally, a dedicated architecture layout is
chosen and the circuit is mapped to the targeted device so that it satisfies
all its connectivity constraints and becomes executable on the device. Similar
to the previous level, the used compiler and its settings are selectable—in
addition to the target device.

All relevant information of a generated benchmark are denoted in short in the
file name and in detail within the file itself.

###  IV-B Accessibility

MQT Bench strives to fulfill the needs of many. In order to accomplish this
goal, it has to be as user-friendly as possible. On the one hand, this
involves the benchmark library itself. All of the high-level, algorithmic
benchmark scripts and routines to generate representations for the individual
levels are publicly available on GitHub (https://github.com/cda-tum/MQTBench).
This allows for complete transparency on how the respective circuits have been
generated. Additionally, a Python package is provided such that each benchmark
can be generated on-demand.

On the other hand, and most importantly, this involves the way in which users
interface with the benchmark suite. While most benchmark suites are provided
as an accumulation of benchmark files—making it hard to extract the desired
set of benchmarks, MQT Bench includes an easy-to-use, no-coding-required web
interface (https://www.cda.cit.tum.de/mqtbench) that provides the means to
filter the large number of pre-generated benchmarks according to the specific
needs of the user. This web interface and its underlying server software is
also part of our Python package, such that every user can utilize the same
interface locally without the need to access our publicly available server. A
screenshot of the website’s user interface and its configuration options is
shown in Fig. 4.

Figure 4: User interface of the MQT Bench website.

###  IV-C Algorithm Selection

To provide a broad spectrum of different benchmarks, MQT Bench comprises most
of today’s de-facto standard quantum algorithms. This includes building
blocks, e.g., the Greenberger–Horne–Zeilinger state preparation (GHZ) or
Quantum Fourier Transform (QFT), up to higher-level algorithms, e.g., Shor’s
[25] and Grover’s [29] algorithm. At the time of writing, MQT Bench comprises
the following benchmarks:

  * •

GHZ state preparation

  * •

Graph state preparation

  * •

W-State preparation

  * •

Amplitude Estimation (AE)

  * •

Quantum Fourier Transformation (QFT)

  * •

Quantum Phase Estimation (QPE)

  * •

Deutsch-Jozsa (DJ) algorithm

  * •

Harrow-Hassidim-Lloyd (HHL) algorithm

  * •

Shor’s algorithm

  * •

Grover’s algorithm

  * •

Quantum Walk

  * •

Quantum Approximation Optimization Algorithm (QAOA)

  * •

Variational Quantum Eigensolver (VQE)

  * •

Three VQA ansatz functions with randomly initialized parameters: Two Local,
Real Amplitudes and Efficient SU2

Additionally, MQT Bench provides several application benchmarks specifically
targeting variational quantum algorithms, since those algorithms are
especially promising in the NISQ-era [30]. To this end, we use the
classification provided by IBM Qiskit’s application levels: Optimization,
Machine Learning, Finance, and Nature:

  * •

_Optimization_ : Travelling salesman problem and vehicle routing

  * •

_Machine Learning_ : Quantum Generative Adversarial Network (QGAN)

  * •

_Finance_ : Portfolio optimization and option pricing

  * •

_Nature_ : Ground state estimation

###  IV-D Extendability

Quantum computing is a rapidly evolving area of research and, especially in
the recent years, numerous promising algorithms and fields of application have
emerged. Any benchmark suite has to be extendable and must continuously evolve
in order to remain relevant. MQT Bench is designed to be easily extendable in
a multitude of ways:

  * •

_Algorithms_ : New applications and algorithms can easily be integrated into
MQT Bench by providing the corresponding algorithmic level description. The
generation of all other levels and options (such as the native gate-sets,
architectures, and compilers with their settings) is automatically handled by
the proposed library.

  * •

_Native gate-sets_ : So far, native gate-sets of superconducting quantum
computers from IBM, Rigetti, and Oxford Quantum Circuits are provided—in
addition to the gate-set of IonQ’s ion trap-based quantum computers. In the
future, additional gate-sets (such as those of Google’s and or AQT’s quantum
computers) can be realized by adopting the necessary compilation routines in
the proposed library.

  * •

_Hardware architectures_ : All major players currently rely on a modular
platform for the architecture of their devices in order to further scale their
quantum computers. It is straightforward to extend MQT Bench’s list of
available architectures to accommodate future architectures.

While the four levels MQT Bench is based on cover today’s main use cases, the
number of abstraction levels most likely is going to increase in the future.
While the mapped level has mostly been considered the last one before a
circuit is send to the quantum computer for execution, pulse-level programming
is envisioned to give even more control to software developers as proposed in
[31]. On the other end of the spectrum, higher-level abstractions and
programming languages will be required to foster the adoption of quantum
technology. Currently, a programmatic description for the benchmarks is needed
on the algorithmic level. An even higher level where this programmatic
description will be automatically derived from is also envisioned. _QUARK_
[32], a framework for quantum computing application benchmarking consisting of
four levels where the lowest level considers the whole compilation flow
discussed in Section II-A, constitutes one step towards this direction.

##  V Conclusion

In this work, we proposed _MQT Bench_ —a quantum circuit benchmark suite
comprising different algorithms, compilers, native gate-sets, and target
devices—resulting in more than $50{\small,}000$ benchmark circuits ranging
from $2$ up to $130$ qubits on four abstraction levels. To keep this large
number of benchmarks manageable, we provide an easy-to-use web interface
(https://www.cda.cit.tum.de/mqtbench) allowing users to filter the benchmarks
according to their needs. _MQT Bench_ is also provided as a Python package
(including the server software to start the web interface locally), such that
each of the benchmarks can be easily generated on-demand. Additionally, we
give access to the open-source repository on GitHub (https://github.com/cda-
tum/MQTBench). By this, _MQT Bench_ presents a first step towards serving a
single benchmark suite for the whole quantum software stack—facilitating
empirical evaluations of quantum software tools that are comparable,
reproducible, and transparent.

## Acknowledgments

This work received funding from the European Research Council (ERC) under the
European Union’s Horizon 2020 research and innovation program (grant agreement
No. 101001318), was part of the Munich Quantum Valley, which is supported by
the Bavarian state government with funds from the Hightech Agenda Bayern Plus,
and has been supported by the BMWK on the basis of a decision by the German
Bundestag through project QuaST, as well as by the BMK, BMDW, and the State of
Upper Austria in the frame of the COMET program (managed by the FFG).

## References

## References

  * [1] Alwin Zulehner and Robert Wille  “Advanced Simulation of Quantum Computations”  In _IEEE Trans. on CAD of Integrated Circuits and Systems_ , 2019
  * [2] D Michael Miller, Mitchell A Thornton and David Goodman  “A Decision Diagram Package for Reversible and Quantum Circuit Simulation”  In _Int’l Conf. on Evolutionary Computation_ , 2006
  * [3] Thomas Häner, Damian S. Steiger, Krysta Svore and Matthias Troyer  “A software methodology for compiling quantum programs”  In _Quantum Science and Technology_ 3.2, 2018
  * [4] Matthew Amy and Vlad Gheorghiu  “staq—A full-stack quantum processing toolkit”  In _Quantum Science and Technology_ 5.3 IOP Publishing, 2020  DOI: 10.1088/2058-9565/ab9359
  * [5] Alexander S. Green et al.  “Quipper: A Scalable Quantum Programming Language”  In _SIGPLAN Not._ 48.6 New York, NY, USA: Association for Computing Machinery, 2013  DOI: 10.1145/2499370.2462177
  * [6] Lukas Burgholzer and Robert Wille  “Advanced Equivalence Checking for Quantum Circuits”  In _IEEE Trans. on CAD of Integrated Circuits and Systems_ IEEE, 2020
  * [7] S. Yamashita and I.. Markov  “Fast Equivalence-Checking for Quantum Circuits”  In _Int’l Symp. on Nanoscale Architectures_ , 2010
  * [8] Philipp Niemann et al.  “QMDDs: Efficient Quantum Function Representation and Manipulation”  In _IEEE Trans. on CAD of Integrated Circuits and Systems_ 35.1 IEEE Press, 2016  DOI: 10.1109/TCAD.2015.2459034
  * [9] Ross Duncan, Aleks Kissinger, Simon Perdrix and John Wetering  “Graph-theoretic Simplification of Quantum Circuits with the ZX-calculus”  In _Quantum_ 4 Verein zur Forderung des Open Access Publizierens in den Quantenwissenschaften, 2020  DOI: 10.22331/q-2020-06-04-279
  * [10] Lukas Burgholzer, Rudy Raymond and Robert Wille  “Verifying Results of the IBM Qiskit Quantum Circuit Compilation Flow”  In _Int’l Conf. on Quantum Computing and Engineering_ , 2020
  * [11] MD SAJID ANIS et al.  “Qiskit: An Open-source Framework for Quantum Computing”, 2021  DOI: 10.5281/zenodo.2573505
  * [12] Cirq Developers  “Cirq” See full list of authors on Github: https://github .com/quantumlib/Cirq/graphs/contributors  Zenodo, 2021  DOI: 10.5281/zenodo.5182845
  * [13] Robert S. Smith, Michael J. Curtis and William J. Zeng  “A Practical Quantum Instruction Set Architecture”, 2016  arXiv:1608.03355
  * [14] Thomas Lubinski et al.  “Application-Oriented Performance Benchmarks for Quantum Computing”, 2021  arXiv: http://arxiv.org/abs/2110.03137
  * [15] Ang Li, Samuel Stein, Sriram Krishnamoorthy and James Ang  “QASMBench: A Low-level QASM Benchmark Suite for NISQ Evaluation and Simulation”, 2021  arXiv: http://arxiv.org/abs/2005.13018
  * [16] Teague Tomesh et al.  “SupermarQ: A Scalable Quantum Benchmark Suite”, 2022  arXiv:2202.11045
  * [17] R. Wille et al.  “RevLib: An Online Resource for Reversible Functions and Reversible Circuits” RevLib is available at http://www.revlib.org  In _Int’l Symp. on Multi-Valued Logic_ , 2008
  * [18] Andrew W. Cross et al.  “OpenQASM 3: A Broader and Deeper Quantum Assembly Language”, 2021  arXiv:2104.14722
  * [19] Thomas Grurl, Richard Kueng, Jürgen Fuß and Robert Wille  “Stochastic Quantum Circuit Simulation Using Decision Diagrams”  In _Design, Automation and Test in Europe_ , 2021
  * [20] Benjamin Villalonga et al.  “A flexible high-performance simulator for verifying and benchmarking quantum circuits implemented on real hardware”  In _npj Quantum Information_ 5.1, 2019  DOI: 10.1038/s41534-019-0196-1
  * [21] Tyson Jones, Anna Brown, Ian Bush and Simon C. Benjamin  “QuEST and High Performance Simulation of Quantum Computers”  In _Scientific Reports_ 9.1, 2019  DOI: 10.1038/s41598-019-47174-9
  * [22] Marcos Yukio Siraichi, Vinícius Fernandes Santos, Sylvain Collange and Fernando Magno Quintao Pereira  “Qubit Allocation”  In _Int’l Symp. on Code Generation and Optimization_ , 2018  DOI: 10.1145/3168822
  * [23] Dominik Janzing, Pawel Wocjan and Thomas Beth  ““Non-identity check” is QMA-complete”  In _Int. J. Quantum Inform._ 03.03, 2005, pp. 463–473
  * [24] Andrew W. Cross, Lev S. Bishop, John A. Smolin and Jay M. Gambetta  “Open Quantum Assembly Language”, 2017  arXiv:1707.03429
  * [25] Peter W Shor  “Polynomial-Time Algorithms for Prime Factorization and Discrete Logarithms on a Quantum Computer”  In _SIAM review_ 41.2 SIAM, 1999
  * [26] Kaitlin N. Smith and Mitchell A. Thornton  “A Quantum Computational Compiler and Design Tool for Technology-Specific Targets”  In _Int’l Symp. on Computer Architecture_ , 2019, pp. 579–588
  * [27] Bochen Tan and Jason Cong  “Optimal Layout Synthesis for Quantum Computing”  In _Int’l Conf. on CAD_ , 2020
  * [28] Gushu Li, Yufei Ding and Yuan Xie  “Tackling the Qubit Mapping Problem for NISQ-era Quantum Devices”  In _Int’l Conf. on Architectural Support for Programming Languages and Operating Systems_ , 2019
  * [29] Lov K Grover  “A fast quantum mechanical algorithm for database search”  In _Symp. on Theory of Computing_ , 1996
  * [30] M. Cerezo et al.  “Variational Quantum Algorithms”  In _Nature Reviews Physics_ 3.9, 2021  DOI: 10.1038/s42254-021-00348-9
  * [31] Boxi Li et al.  “Pulse-level noisy quantum circuits with QuTiP”  In _Quantum_ 6 Verein zur Förderung des Open Access Publizierens in den Quantenwissenschaften, 2022, pp. 630  DOI: 10.22331/q-2022-01-24-630
  * [32] Jernej Rudi Finžgar, Philipp Ross, Johannes Klepsch and Andre Luckow  “QUARK: A Framework for Quantum Computing Application Benchmarking”, 2022  arXiv:2202.03028


# Policy-Guided Diffusion

###### Abstract

In many real-world settings, agents must learn from an offline dataset
gathered by some prior behavior policy. Such a setting naturally leads to
distribution shift between the behavior policy and the target policy being
trained—requiring policy conservatism to avoid instability and overestimation
bias. Autoregressive world models offer a different solution to this by
generating synthetic, on-policy experience. However, in practice, model
rollouts must be severely truncated to avoid compounding error. As an
alternative, we propose policy-guided diffusion. Our method uses diffusion
models to generate entire trajectories under the behavior distribution,
applying guidance from the target policy to move synthetic experience further
on-policy. We show that policy-guided diffusion models a regularized form of
the target distribution that balances action likelihood under both the target
and behavior policies, leading to plausible trajectories with high target
policy probability, while retaining a lower dynamics error than an offline
world model baseline. Using synthetic experience from policy-guided diffusion
as a drop-in substitute for real data, we demonstrate significant improvements
in performance across a range of standard offline reinforcement learning
algorithms and environments. Our approach provides an effective alternative to
autoregressive offline world models, opening the door to the controllable
generation of synthetic training data.

Matthew T. Jackson***Correspondence to<EMAIL_ADDRESS>Michael T.
Matthews Cong Lu Benjamin Ellis

Shimon Whiteson†††Equal supervision. Jakob N. Foerster22footnotemark: 2

University of Oxford

##  1 Introduction

A key obstacle to the real-world adoption of reinforcement learning (RL,
Sutton & Barto, 2018) is its notorious sample inefficiency, preventing agents
from being trained on environments with expensive or slow online data
collection. A closely related challenge arises in environments where
exploration, required by standard RL methods, is inherently dangerous,
limiting their applicability. Yet many such settings come with an abundance of
pre-collected or offline experience, gathered under one or more behavior
policies (Yu et al., 2020). These settings enable the application of offline
RL (Levine et al., 2020), where a policy is optimized from an offline dataset
without access to the environment. However, the distribution shift between the
target policy (i.e., the policy being optimized) and the collected data poses
many challenges (Kumar et al., 2020; Kostrikov et al., 2021; Fujimoto et al.,
2019).

In particular, distribution shift between the target and behavior policies
leads to an out-of-sample issue: since the goal of offline RL is to exceed the
performance of the behavior policy, the distribution of state-action pairs
sampled by the target policy necessarily differs from that of the behavior
policy, and its samples are therefore underrepresented (or unavailable) in the
offline dataset. However, the maximizing nature of RL classically leads to
overestimation bias when generalizing to rarely seen state-action pairs,
resulting in an overly optimistic target policy. As a solution, most previous
model-free work has proposed severe regularization of the target policy—such
as penalizing value estimates in uncertain states (Kumar et al., 2020; An et
al., 2021) or regularizing it towards the behavior policy (Fujimoto & Gu,
2021)—sacrificing target policy performance for stability.

Figure 1:  Offline reinforcement learning with policy-guided diffusion.
Offline data from a behavior policy is first used to train a trajectory
diffusion model. Synthetic experience is then generated with diffusion, guided
by the target policy in order to move trajectories further on-policy. An agent
is then trained for multiple steps on the synthetic dataset, before it is
regenerated.

In this paper, we focus on an alternative class of methods: generating
synthetic experience to both augment the offline dataset and lessen the out-
of-sample issue. Prior methods in this area use a model-based approach (Yu et
al., 2020; Kidambi et al., 2020; Ball et al., 2021, see Section 3.1), in which
a single-step world model is learned from the offline dataset, which the
target policy interacts with to generate synthetic on-policy training
experience. While this allows the target policy to sample synthetic
trajectories under its own action distribution, compounding model error
usually forces these methods to severely truncate model rollouts to a handful
of interactions. Thus, there are two options which trade off coverage and
bias. The first is to roll out from the initial state, which is unbiased but
lacks coverage. The second is to roll out from states randomly sampled from
the data set, which increases coverage but introduces bias. Neither option
fully addresses the difference in observed states between the behavior and
target policy when deployed, nor the out-of-sample issue mentioned above.

Instead, we propose policy-guided diffusion (PGD, Figure 1), which avoids
compounding error by modeling entire trajectories (Section 3.2) rather than
single-step transitions. To achieve this, we train a diffusion model on the
offline dataset, from which we can sample synthetic trajectories under the
behavior policy. However, while this addresses data sparsity, these
trajectories are still off-distribution from our target policy. Therefore,
inspired by classifier-guided diffusion (Dhariwal & Nichol, 2021), we apply
guidance from the target policy to shift the sampling distribution towards
that of the target policy. At each diffusion step, our guidance term directly
increases the likelihood of sampled synthetic actions under the target policy,
while the diffusion model updates the entire trajectory towards those in the
dataset. This yields a regularized target distribution that we name the
_behavior-regularized target distribution_ , ensuring actions do not diverge
too far from the behavior policy, limiting generalization error. As a result,
PGD does not suffer from compounding error, while also generating synthetic
trajectories that are more representative of the target policy. We illustrate
this point in Figure 2.

Our approach results in consistent improvements in offline RL performance for
agents trained on policy-guided synthetic data, compared to those trained on
unguided synthetic or real data. We evaluate using the standard TD3+BC
(Fujimoto & Gu, 2021) and IQL (Kostrikov et al., 2021) algorithms across a
variety of D4RL (Fu et al., 2020) datasets. Notably, we see a statistically
significant 11.2% improvement in performance for the TD3+BC algorithm
aggregated across MuJoCo (Todorov et al., 2012) locomotion tasks compared to
training on the real data, with no algorithmic changes. Our results also
extend to even larger improvements for the challenging Maze2d navigation
environments. Furthermore, we analyze synthetic trajectories generated by PGD
and show that PGD achieves lower dynamics error than PETS (Chua et al.,
2018)—a prior offline model-based method—while matching the target policy
likelihood of PETS. Together, our experiments illustrate the potential of PGD
as an effective drop-in replacement for real data—across agents, environments,
and behavior policies.

##  2 Background

###  2.1 Offline Reinforcement Learning

##### Formulation

We adopt the standard reinforcement learning formulation, in which an agent
acts in a Markov Decision Process (MDP, Sutton & Barto, 2018). An MDP is
defined as a tuple $M=\langle\mathcal{S},\mathcal{A},p_{0},T,R,H\rangle$,
where $s\in\mathcal{S}$ and $a\in\mathcal{A}$ are the state and action spaces,
$p_{0}(s_{0})$ is a probability distribution over the initial state,
$T(s_{t+1}|s_{t},a_{t})$ is a conditional probability distribution defining
the transition dynamics,
$R:\mathcal{S}\times\mathcal{A}\xrightarrow{}\mathbb{R}$ is the reward
function, $\gamma$ is the discount factor, and $H$ is the horizon.

In RL, we learn a policy $\pi(a|s)$ that defines a conditional probability
distribution over actions for each state, inducing a distribution over
trajectories $\bm{\tau}\coloneqq(s_{0},a_{1},r_{1},s_{1},\ldots,s_{H})$ given
by

$p_{\pi,M}(\bm{\tau})=p_{0}(s_{0})\prod_{t=0}^{H-1}\pi(a_{t}|s_{t})\cdot
T(s_{t+1}|s_{t},a_{t})\text{,}$ (1)

omitting the reward function throughout our work for conciseness. Our goal is
to learn a policy that maximizes the expected return, defined as
$\mathbb{E}_{p_{\pi,M}}[V(\bm{\tau})]$ where
$V(\bm{\tau})\coloneqq\sum_{t=0}^{H}r_{t}$ is the return of a trajectory. The
offline RL setting (Levine et al., 2020) extends this, preventing the agent
from interacting with the environment and instead presenting it with a dataset
of trajectories $\bm{\tau}\in\mathcal{D}_{\text{off}}$ gathered by some
unknown behavior policy $\pi_{\text{off}}$, with which to optimize a target
policy $\pi_{\text{target}}$.

##### Out-of-Sample Generalization

The core challenge of offline RL emerges from the distribution shift between
the _behavior distribution_ $p_{\pi_{\text{off}},M}(\bm{\tau})$ and the
_target distribution_ $p_{\pi_{\text{target}},M}(\bm{\tau})$, which are
otherwise denoted $p_{\mkern 1.0mu\text{off}}(\bm{\tau})$ and $p_{\mkern
1.0mu\text{target}}(\bm{\tau})$ for conciseness. Optimization of
$\pi_{\text{target}}$ on $\mathcal{D}_{\text{off}}$ can lead to catastrophic
value overestimation at unobserved actions, a problem termed the _out-of-
sample issue_ (Kostrikov et al., 2021). As such, model-free offline algorithms
typically regularize the policy towards the behavior distribution, either
explicitly (Fujimoto & Gu, 2021; Kumar et al., 2020) or implicitly (Kostrikov
et al., 2021).

Alternatively, prior work proposes learning a single-step world model
$\mathcal{M}$ from $\mathcal{D}_{\text{off}}$ (Yu et al., 2020; Kidambi et
al., 2020; Lu et al., 2022). By rolling out the target policy using
$\mathcal{M}$, we generate trajectories $\bm{\tau}\sim p_{\mkern
1.0mu\text{target}}(\bm{\tau})$, with the aim of avoiding distribution shift.
However, in practice, this technique only pushes the generalization issue into
the world model. In particular, RL policies are prone to exploiting errors in
the world model, which can compound over the course of an episode. When
combined with typical maximizing operations used in off-policy RL, this
results in value overestimation bias (Sims et al., 2024).

###  2.2 Diffusion Models

##### Definition

To generate synthetic data, we consider diffusion models (Sohl-Dickstein et
al., 2015; Ho et al., 2020), a class of generative model that allows one to
sample from a distribution $p(x)$ by iteratively reversing a forward noising
process. Karras et al. (2022) present an ODE formulation of diffusion which,
given a noise schedule $\sigma(i)$ indexed by $i$, mutates data according to

$d\bm{x}=-\dot{\sigma}(i)\sigma(i)\nabla_{\bm{x}}\log
p\left(\bm{x};\sigma(i)\right)di\text{,}$ (2)

where $\dot{\sigma}=\frac{\mathrm{d}\sigma}{\mathrm{d}i}$ and
$\nabla_{\bm{x}}\log p\left(\bm{x};\sigma(i)\right)$ is the score function
(Hyvärinen & Dayan, 2005), which points towards areas of high data density.
Intuitively, infinitesimal forward or backward steps of this ODE respectively
nudge a sample away from or towards the data. To generate a sample, we start
with pure noise at the highest noise level $\sigma_{\textrm{max}}$, and
iteratively denoise in discrete timesteps under Equation 2.

##### Classifier Guidance

Our method is designed to augment the data-generating process towards on-
policy trajectories from the target distribution $p_{\mkern
1.0mu\text{target}}(\bm{\tau})$, rather than the behavior distribution
$p_{\mkern 1.0mu\text{off}}(\bm{\tau})$. To achieve this, we take inspiration
from classifier guidance (Dhariwal & Nichol, 2021), which leverages a
differentiable classifier to augment the score function of a pre-trained
diffusion model towards a class-conditional distribution $p(\bm{x}|y)$.
Concretely, this adds a classifier gradient to the score function, giving

$\nabla_{\bm{x}}\log
p_{\lambda}\left(\bm{x}|y;\sigma(i)\right)=\nabla_{\bm{x}}\log
p\left(\bm{x};\sigma(i)\right)+\lambda\nabla_{\bm{x}}\log
p_{\theta}\left(y|\bm{x};\sigma(i)\right)\text{,}$ (3)

where $\nabla_{\bm{x}}\log p_{\theta}\left(y|\bm{x};\sigma(i)\right)$ is the
gradient of the classifier and $\lambda$ is the guidance weight.

Figure 2:  Trajectories from an illustrative 2D environment, in which the
start location is indicated by $\bullet$ and the goals for the behavior and
target policies are indicated by $\mathbf{\times}$ and $\mathbf{\times}$.
Left: Rollouts from the target policy in the real environment. Right: Offline
datasets gathered by the behavior policy suffer from distribution shift and
limited sample size. Truncated world models (Yu et al., 2020; Kidambi et al.,
2020) previously used in offline model-based reinforcement learning offer a
partial solution to this problem but suffer from bias due to short rollouts.
Meanwhile, unguided diffusion (Lu et al., 2023) can increase the sample size,
but maintains the original distribution shift. In contrast, policy-guided
diffusion samples from a regularized target distribution, generating entire
trajectories with low transition error but higher likelihood under the target
distribution.

##  3 On-Policy Sampling from Offline Data

Generating synthetic agent experience is a promising approach to solving out-
of-sample generalization in offline RL. By generating experience that is
unseen in the dataset, the policy may be directly optimized on OOD samples,
thereby moving the generalization problem from the policy to the generative
model. Some prior work has suggested learning a model from the offline dataset
(Lu et al., 2023), thereby sampling synthetic experience from the behavior
distribution. While this improves sample coverage, the approach retains many
of the original challenges of offline RL. As with the behavior policy, the
synthetic trajectories may be suboptimal, meaning we still require
conservative off-policy RL techniques to train the agent.

Instead, we seek to extend this approach by making our generative model sample
from the target distribution. This reduces the need for conservatism and
generates synthetic trajectories with increasing performance as the agent
improves over training. Practically, the effectiveness of this approach
depends on how we parameterize each of the terms of the trajectory
distribution (Equation 1). In this section, we consider two parameterizations:
autoregressive and direct.

###  3.1 Autoregressive Generation — Model $\bm{\mathcal{T}}$, Sample
$\bm{p(s_{0})}$

The autoregressive—or model-based—approach to generating on-policy data is to
use the offline dataset to train a one-step transition model
$\mathcal{T}(s_{t+1}|s_{t},a_{t};\theta)$. To generate unbiased sample
trajectories from the target distribution, we first sample an initial state
(i.e., one that starts an episode) from the offline dataset
$s_{0}\sim\mathcal{D}_{\text{off}}$. Next, we roll out our agent in the
learned model by iteratively sampling actions from the target policy and
approximating environment transitions with the learned dynamics model.
However, compounding error from the transition model usually requires agent
rollouts to be much shorter than the environment horizon—such that the agent
takes $k\ll H$ steps.‡‡‡Typically $k\leq 5$ (Janner et al., 2019; Yu et al.,
2020). Consequently, any states more than $k$ steps away from any initial
state cannot be generated in this manner, limiting the applicability of this
approach.

As an approximation, autoregressive methods typically sample initial states
from any timestep $s_{t}\sim\mathcal{D}_{\text{off}}$ in the offline dataset.
Given a truncated rollout length $k$, this may be seen as approximating the
sub-trajectory distribution—i.e., the trajectory from time $t$ to $t+k$—given
by

$\displaystyle p_{\mkern 1.0mu\text{target}}(\bm{\tau}_{t:t+k};\theta)$
$\displaystyle=p_{\mkern 1.0mu\text{target}}(s_{t})\cdot p_{\mkern
1.0mu\text{target}}(\bm{\tau}_{t:t+k}|s_{t};\theta),$ (4)

by instead modeling

$\displaystyle\mathcal{F}(\bm{\tau}_{t:t+k};\theta)$ $\displaystyle=p_{\mkern
1.0mu\text{off}}(s_{t})\cdot p_{\mkern
1.0mu\text{target}}(\bm{\tau}_{t:t+k}|s_{t};\theta)\text{.}$ (5)

Here, we denote the stationary state distributions of the target and behavior
policies at time $t$ by $p_{\mkern 1.0mu\text{target}}(s_{t})$ and $p_{\mkern
1.0mu\text{off}}(s_{t})$ respectively, and define the conditional sub-
trajectory distribution as

$\displaystyle p_{\mkern
1.0mu\text{target}}(\bm{\tau}_{t:t+k}|s_{t};\theta)\coloneqq\prod_{j=0}^{k-1}$
$\displaystyle\pi_{\text{target}}(a_{t+j}|s_{t+j})\cdot\mathcal{T}(s_{t+j+1}|s_{t+j},a_{t+j};\theta)\text{.}$
(6)

When generating trajectories from this distribution, the difference between
$p_{\mkern 1.0mu\text{target}}(s_{t})$ and $p_{\mkern 1.0mu\text{off}}(s_{t})$
biases the start of rollouts towards states visited by the behavior policy.
Furthermore, we still require $k$ to be small to avoid compounding error. In
combination, sampling from the offline dataset “anchors” synthetic rollouts to
states in the offline dataset, while truncated rollouts prevent synthetic
trajectories from moving far from this anchor. Therefore, the practical
application of autoregressive generation leads to a strong bias towards the
behavior distribution and fails to address the out-of-sample problem.

###  3.2 Direct Generation — Model $\bm{p_{\mkern 1.0mu\text{{off}}}(\tau)}$

As an alternative to autoregressive generation, we can parameterize the target
distribution by directly modeling the behavior distribution, as follows:

$\displaystyle p_{\mkern 1.0mu\text{target}}(\bm{\tau})$
$\displaystyle=p(s_{0})\prod_{t=0}^{H-1}\pi_{\text{target}}(a_{t}|s_{t})\cdot\mathcal{T}(s_{t+1}|s_{t},a_{t})$
$\displaystyle=p(s_{0})\prod_{t=0}^{H-1}\frac{\pi_{\text{target}}(a_{t}|s_{t})}{\pi_{\text{off}}(a_{t}|s_{t})}\cdot\pi_{\text{off}}(a_{t}|s_{t})\cdot\mathcal{T}(s_{t+1}|s_{t},a_{t})$
$\displaystyle=p_{\mkern
1.0mu\text{off}}(\bm{\tau})\prod_{t=0}^{H-1}w_{a_{t},s_{t}}$
$\displaystyle\approx p_{\mkern
1.0mu\text{off}}(\bm{\tau};\theta)\prod_{t=0}^{H-1}w_{a_{t},s_{t}}=p_{\mkern
1.0mu\text{target}}(\bm{\tau};\theta)$ (7)

where $w_{a,s}\coloneqq\frac{\pi_{\text{target}}(a|s)}{\pi_{\text{off}}(a|s)}$
denotes the importance sampling weight for $(a,s)$ (Precup, 2000). This
directly parameterizes the behavior distribution $p_{\mkern
1.0mu\text{off}}(\bm{\tau};\theta)$—which may be learned by modeling entire
trajectories on the offline dataset—and adjusts their likelihoods by the
relative probabilities of actions $w_{a_{t},s_{t}}$ under the target and
behavior policies. By jointly modeling the initial state distribution,
transition function, and behavior policy, such a parameterization is not
required to enforce the Markov property. As a result, it can directly generate
entire trajectories, thereby avoiding the compounding model error suffered by
autoregressive methods when iteratively generating transitions.

However, computing $w_{a_{t},s_{t}}$ requires access to the behavior policy
$\pi_{\text{off}}(a|s)$, which is not assumed in offline RL. Prior work has
explored modeling the behavior policy from the offline dataset and using this
to compute importance sampling corrections. However, products of many
importance weights can lead to problems with high variance (Precup et al.,
2000; Levine et al., 2020).

##  4 Policy-Guided Diffusion

In this work, we propose a method following the direct generation approach
outlined in Section 3.2, named policy-guided diffusion (PGD,
LABEL:alg:generation). Following the success of diffusion models at generating
trajectories (Janner et al., 2022; Lu et al., 2023), we first train a
trajectory-level diffusion model on the offline dataset to model the behavior
distribution. Then, inspired by classifier-guided diffusion (Section 2.2), we
guide the diffusion process using the target policy to move closer to the
target distribution. Specifically, during the denoising process, we compute
the gradient of the action distribution for each action under the target
policy, using it to augment the diffusion process towards high-probability
actions. In doing so, we approximate a regularized target distribution that
equally weights action likelihoods under the behavior and target policies.

In this section, we derive PGD as an approximation of the behavior-regularized
target distribution (Section 4.1), then describe practical details for
controlling and stabilizing policy guidance (Section 4.2). We provide a
summary of PGD against alternative sources of training data in Table 1.

###  4.1 Behavior-Regularized Target Distribution

##### Policy Guidance Derivation

To sample a trajectory via diffusion, we require a noise-conditioned score
function $\nabla_{\bm{\hat{\tau}}}\log p(\bm{\hat{\tau}};\sigma)$ for a noised
trajectory
$\bm{\hat{\tau}}\coloneqq(\hat{s}_{0},\hat{a}_{1},\hat{r}_{1},\hat{s}_{1},\ldots,\hat{s}_{H})$
under a distribution $p(\bm{\tau})$ at a noise level $\sigma$. Given an
offline dataset $\mathcal{D}_{\text{off}}$, it is straightforward to learn
this function under the behavior distribution, $\nabla_{\bm{\hat{\tau}}}\log
p_{\mkern 1.0mu\text{off}}(\bm{\hat{\tau}};\sigma)$, by training a denoiser
model to reconstruct noised trajectories from $\mathcal{D}_{\text{off}}$.
However, there is no apparent method to directly model the noise-conditioned
score function $\nabla_{\bm{\hat{\tau}}}\log p_{\mkern
1.0mu\text{target}}(\bm{\hat{\tau}};\sigma)$ for the target distribution (see
Appendix B for further discussion), meaning we require an approximation.

To achieve this, we consider the score function of a noise-free trajectory
$\bm{\tau}$ under the target distribution, based on the formulation from
Equation 7,

$\nabla_{\bm{\tau}}\log p_{\mkern
1.0mu\text{target}}(\bm{\tau})=\nabla_{\bm{\tau}}\log p_{\mkern
1.0mu\text{off}}(\bm{\tau})+\sum_{t=0}^{H-1}\left(\nabla_{\bm{\tau}}\log\pi_{\text{target}}\left(a_{t}|s_{t}\right)-\nabla_{\bm{\tau}}\log\pi_{\text{off}}\left(a_{t}|s_{t}\right)\right)\text{.}$
(8)

In the limit of noise $\sigma\to 0$, the noise-conditioned score function
$\nabla_{\bm{\hat{\tau}}}\log p_{\mkern
1.0mu\text{target}}(\bm{\hat{\tau}};\sigma)$ clearly approaches
$\nabla_{\bm{\tau}}\log p_{\mkern 1.0mu\text{target}}(\bm{\tau})$. Therefore,
we may approximate this function by

$\nabla_{\bm{\hat{\tau}}}\log p_{\mkern
1.0mu\text{target}}(\bm{\hat{\tau}};\sigma)\approx\nabla_{\bm{\hat{\tau}}}\log
p_{\mkern
1.0mu\text{off}}(\bm{\hat{\tau}};\sigma)+\sum_{t=0}^{H-1}\left(\nabla_{\bm{\hat{\tau}}}\log\pi_{\text{target}}\left(\hat{a}_{t}|\hat{s}_{t}\right)-\nabla_{\bm{\hat{\tau}}}\log\pi_{\text{off}}\left(\hat{a}_{t}|\hat{s}_{t}\right)\right)\text{,}$
(9)

for $\sigma\approx 0$. Whilst iteratively denoising under this function
(Section 2.2) does not model $p_{\mkern 1.0mu\text{target}}(\bm{\tau})$
exactly, the score function approaches $\nabla_{\bm{\hat{\tau}}}\log p_{\mkern
1.0mu\text{target}}(\bm{\hat{\tau}};\sigma)$ towards the end of the denoising
process, which we believe provides an effective approximation.

##### Excluding Behavior Policy Guidance

As discussed, we may directly model the first term of Equation 9 by training a
denoiser model. Furthermore, we may directly compute target policy guidance
$\nabla_{\bm{\hat{\tau}}}\log\pi_{\text{target}}\left(\hat{a}_{t}|\hat{s}_{t}\right)$—the
second term of this approximation—as we assume access to a (differentiable)
target policy in the offline RL setting. However, we generally do not have
access to the behavior policy, preventing us from computing
$\nabla_{\bm{\hat{\tau}}}\log\pi_{\text{off}}\left(\hat{a}_{t}|\hat{s}_{t}\right)$.
Due to this, we exclude behavior policy guidance from our approximation,
resulting in the score function $\nabla_{\bm{\hat{\tau}}}\log p_{\mkern
1.0mu\text{off}}(\bm{\hat{\tau}};\sigma)+\sum_{t=0}^{H-1}\nabla_{\bm{\hat{\tau}}}\log\pi_{\text{target}}\left(\hat{a}_{t}|\hat{s}_{t}\right)$.
As $\sigma\to 0$, this approaches the score function for a proxy distribution
of the form

$\displaystyle\mathcal{F}(\bm{\tau};\pi_{\text{target}})$
$\displaystyle\propto p_{\mkern
1.0mu\text{off}}(\bm{\tau})\prod_{t=0}^{H-1}\pi_{\text{target}}(a_{t}|s_{t})$
$\displaystyle=p_{\mkern 1.0mu\text{off}}(\bm{\tau})\cdot q_{\mkern
1.0mu\text{target}}(\bm{\tau})=p_{\mkern 1.0mu\text{target}}(\bm{\tau})\cdot
q_{\mkern 1.0mu\text{off}}(\bm{\tau})\text{,}$ (10)

where $q_{\mkern
1.0mu\text{target}}(\bm{\tau})\coloneqq\prod_{t=0}^{H-1}\pi_{\text{target}}(a_{t}|s_{t})$
denotes the product of action probabilities under the target policy and
$q_{\mkern 1.0mu\text{off}}(\bm{\tau})$ denotes the same quantity under the
behavior policy. Therefore, we hypothesize that excluding behavior policy
guidance is an effective form of regularization, as it biases trajectories
towards the support of the offline data, thereby limiting model error and the
out-of-sample problem. We refer to
$\mathcal{F}(\bm{\tau};\pi_{\text{target}})$ as the behavior-regularized
target distribution due to it balancing action likelihoods under the behavior
and target policies, and provide further discussion in Appendix C. Finally, as
a promising avenue for future work, we note that the behavior policy may be
modeled by applying behavior cloning to $\mathcal{D}_{\text{off}}$, allowing
for the inclusion of behavior policy guidance in the offline RL setting.

##### Excluding State Guidance

Target policy guidance
$\nabla_{\bm{\hat{\tau}}}\log\pi_{\text{target}}\left(\hat{a}_{t}|\hat{s}_{t}\right)$
has non-zero gradients for the state and action at timestep $t$. In practice,
the action component
$\nabla_{\hat{a}_{t}}\log\pi_{\text{target}}\left(\hat{a}_{t}|\hat{s}_{t}\right)$
typically has an efficient, closed-form solution, with
$\pi_{\text{target}}\left(\hat{a}_{t}|\hat{s}_{t}\right)$ commonly being
Gaussian for continuous action spaces. In contrast, for neural network
policies, the state component
$\nabla_{\hat{s}_{t}}\log\pi_{\text{target}}\left(\hat{a}_{t}|\hat{s}_{t}\right)$
requires backpropagating gradients through the policy network, which is both
expensive to compute and can lead to high variance on noisy, out-of-
distribution states. Due to this, we apply policy guidance to only the noised
action, yielding our policy-guided score function

$\displaystyle s_{\mkern
1.0mu\text{PGD}}(\bm{\hat{\tau}};\sigma)=\underbrace{\nabla_{\bm{\hat{\tau}}}\log
p_{\mkern 1.0mu\text{off}}(\bm{\hat{\tau}};\sigma)}_{\mathclap{\text{Behavior
score function}}}+\underbrace{\nabla_{\bm{\hat{a}}}\log q_{\mkern
1.0mu\text{target}}(\bm{\hat{\tau}})}_{\mathclap{\text{Target policy
guidance}}}\text{,}$ (11)

where (abusing notation) $\nabla_{\bm{\hat{a}}}$ denotes the gradient
$\nabla_{\bm{\hat{\tau}}}$ of $\bm{\hat{\tau}}$, with non-action components
set to 0.

###  4.2 Improving Policy Guidance

##### Controlling Guidance Strength

A standard technique from classifier-guided diffusion is the use of guidance
coefficients (Dhariwal & Nichol, 2021). These augment the guided score
function by introducing a controllable coefficient on the guidance term.
Applied to the PGD score function (Equation 11), this has the form

$s_{\mkern
1.0mu\text{PGD}}(\bm{\hat{\tau}};\sigma,\lambda)=\nabla_{\bm{\hat{\tau}}}\log
p_{\mkern
1.0mu\text{off}}(\bm{\hat{\tau}};\sigma)+\lambda\nabla_{\bm{\hat{a}}}\log
q_{\mkern 1.0mu\text{target}}(\bm{\hat{\tau}})\text{,}$ (12)

where $\lambda$ denotes the guidance coefficient. As $\sigma\to 0$, this
transforms the sampling distribution to

$\mathcal{F}(\bm{\tau}|\pi_{\text{target}};\lambda)\propto p_{\mkern
1.0mu\text{off}}(\bm{\tau})\cdot q_{\mkern
1.0mu\text{target}}(\bm{\tau})^{\lambda}\text{.}$ (13)

Intuitively, $\lambda$ interpolates the actions in the sampling distribution
between the behavior and target distributions. By tuning $\lambda$, we can
therefore control the strength of guidance towards the target policy, avoiding
high dynamics error when the target policy is far from the behavior policy. We
visualize this effect in Figure 3 and analyze its impact on target policy
likelihood in Figure 5.

Figure 3:  Left: Trajectory probability distribution for an example behavior
distribution $p_{\mkern 1.0mu\text{off}}(\bm{\tau})$ and target policy
likelihood $q_{\mkern 1.0mu\text{target}}(\bm{\tau})$. Right: Corresponding
PGD sampling distribution (Equation 13) computed over a range of policy-
guidance coefficients $\lambda$. By increasing $\lambda$, we transform from
the sampling distribution towards the regions of high target policy
likelihood, making PGD an effective mechanism for controlling the level of
regularization towards the behavior distribution.

Following Ma et al. (2023), we also apply a cosine guidance schedule to the
guidance coefficient,

$\lambda_{n}=\lambda\cdot(\sigma_{n}+\beta\sigma_{N}\cdot\sin(\pi\cdot
n/N))\text{,}$ (14)

where $\beta$ is the cosine weight, which is set to 0.3 in all experiments. By
decreasing the strength of guidance in later steps, we find that this schedule
stabilizes guidance and reduces dynamics error.

##### Stabilizing Guided Diffusion

When under distribution shift, RL policies are known to suffer from poor
generalization to unseen states (Kirk et al., 2023). This makes policy
guidance challenging, since the policy must operate on noised states, and
compute action gradients from noised actions. Similar issues have been studied
in classifier-guided diffusion (Ma et al., 2023), where the classifier
gradient can be unstable when exposed to out-of-distribution inputs. Bansal et
al. (2023) alleviate this issue by applying guidance to the denoised sample
estimated by the denoiser model, rather than the original noised sample, in
addition to normalizing the guidance gradient to a unit vector. By applying
these techniques to policy guidance, we lessen the need for the target policy
to generalize to noisy states, which we find decreases dynamics error.

Table 1: Overview of training experience sources in offline RL—for each, we consider the sampling distribution, expected error in transition dynamics, likelihood of actions under the target policy, and state space coverage beyond the behavior distribution. Policy-guided diffusion provides an effective trade-off between each error, likelihood, and coverage. Data source | Distribution | Error ($\bm{\downarrow}$) | Likelihood ($\bm{\uparrow}$) | Coverage ($\bm{\uparrow}$)
---|---|---|---|---
Offline dataset | $p_{\mkern 1.0mu\text{off}}(\bm{\tau})$ | — | Low | Low
Episodic world model | $p_{\mkern 1.0mu\text{target}}(\bm{\tau})$ | High | — | High
Truncated world model | Equation 5 | Low | — | Low
Unguided diffusion | $p_{\mkern 1.0mu\text{off}}(\bm{\tau})$ | Low | Low | High
Policy-guided diffusion | Equation 10 | Low | High | High

##  5 Results

Through our experiments, we first demonstrate that agents trained with
synthetic experience from PGD outperform those trained on unguided synthetic
data or directly on the offline dataset (Section 5.2). We show that this
effect is consistent across agents (TD3+BC and IQL), environments
(HalfCheetah, Walker2d, Hopper, and Maze), behavior policies (random, mixed,
and medium), and modes of data generation (continuous and periodic). Following
this, we demonstrate that tuning the guidance coefficient enables PGD to
sample trajectories with high action likelihood across a range of target
policies. Finally, we verify that PGD retains low dynamics error despite
sampling high-likelihood actions from the policy (Section 5.3).

###  5.1 Experimental Setup

We evaluate PGD on the MuJoCo and Maze2d continuous control datasets from D4RL
(Fu et al., 2020; Todorov et al., 2012). For MuJoCo, we consider the
HalfCheetah, Walker2d, and Hopper environments with random (randomly
initialized behavior policy), medium (suboptimal behavior policy), and medium-
replay (or “mixed”, the replay buffer from medium policy training) datasets.
For Maze2d we consider the original (sparse reward) instances of the umaze,
medium and large layouts. We train 4 trajectory diffusion models on each
dataset, for which we detail hyperparameters in Appendix A. In Section 5.3, we
conduct analysis of PGD against MOPO-style PETS (Chua et al., 2018) models, an
autoregressive world model composed of an ensemble of probabilistic models,
for which we use model weights from OfflineRL-Kit (Sun, 2023).

To demonstrate synthetic experience from PGD as a drop-in substitute for the
real dataset, we transfer the original hyperparameters for IQL (Kostrikov et
al., 2021) and TD3+BC (Fujimoto & Gu, 2021)—as tuned on the real
datasets—without any further tuning. Policy guidance requires a stochastic
target policy, in order to compute the gradient of the action distribution.
Since TD3+BC trains a deterministic policy, we perform guidance by modeling
the action distribution as a unit Gaussian centered on the deterministic
action. We implement all agents and diffusion models from scratch in Jax
(Bradbury et al., 2018), which may be found at
https://github.com/EmptyJackson/policy-guided-diffusion.

###  5.2 Offline Reinforcement Learning

For each D4RL dataset, we train two popular model-free offline algorithms,
TD3+BC (Fujimoto & Gu, 2021) and IQL (Kostrikov et al., 2021) on synthetic
experience generated by trajectory diffusion models with and without policy
guidance, as well as on the real dataset. We first consider periodic
generation of synthetic data, in which the synthetic dataset is regenerated
after extended periods of agent training, such that the agent is near
convergence on the synthetic dataset at the point it is regenerated with the
current policy. Each epoch, we generate a dataset of $2^{14}$ synthetic
trajectories of length 16. Following the notation of LABEL:alg:training, we
set the number of epochs to $N_{\text{epochs}}=4$ with
$N_{\text{policy}}=$250,000$$ train steps per epoch, meaning the agent is
trained to close to convergence before the dataset is regenerated. This can be
viewed as solving a sequence of offline RL tasks with synthetic datasets, in
which the behavior policy is the target policy from the previous generation.

Using periodic generation, performance improves significantly across
benchmarks for both IQL and TD3+BC (Table 2). In MuJoCo, the most consistent
improvement is on mixed datasets, with 4 out of 6 experiments achieving
significant performance improvement. This is to be expected, as these datasets
contain experience from a mixture of behavior policy levels. In this case, the
diffusion model is likely to be able to represent a wide variety of policies,
and on-policy guidance would naturally produce higher return trajectories as
the target policy improves.

Table 2: Final return of IQL and TD3+BC agents trained on real, unguided ($\lambda=0$) synthetic and policy-guided ($\lambda=1$) synthetic data—mean and standard error over 4 seeds (diffusion models and agents) is presented, with significant improvements ($p<0.05$) shaded. | IQL |  | TD3+BC
---|---|---|---
|  | Dataset |  Unguided |  Guided |  | Dataset | Unguided |  Guided
Random | HalfCheetah |  | $\displaystyle 9.1$ | $\pm$ | $2.2$ | $2.6$ | $\pm$ | $0.1$ | $\displaystyle 6.5$ | $\pm$ | $1.7$ |  |  | $11.2$ | $\pm$ | $0.8$ | $11.0$ | $\pm$ | $0.4$ | $\displaystyle 21.1$ | $\pm$ | $0.9$
Walker2d |  | $4.3$ | $\pm$ | $0.5$ | $2.7$ | $\pm$ | $0.7$ | $\displaystyle 5.3$ | $\pm$ | $0.3$ |  |  | $0.5$ | $\pm$ | $0.3$ | $1.1$ | $\pm$ | $1.2$ | $-0.3$ | $\pm$ | $0.1$
Hopper |  | $\displaystyle 7.4$ | $\pm$ | $0.4$ | $5.2$ | $\pm$ | $0.9$ | $4.9$ | $\pm$ | $1.0$ |  |  | $7.4$ | $\pm$ | $0.6$ | $4.2$ | $\pm$ | $1.4$ | $5.5$ | $\pm$ | $2.1$
Mixed | HalfCheetah |  | $\displaystyle 44.2$ | $\pm$ | $0.2$ | $43.6$ | $\pm$ | $0.2$ | $43.6$ | $\pm$ | $0.2$ |  |  | $44.7$ | $\pm$ | $0.1$ | $43.1$ | $\pm$ | $0.2$ | $\displaystyle 46.1$ | $\pm$ | $0.3$
Walker2d |  | $81.3$ | $\pm$ | $2.0$ | $\displaystyle 85.2$ | $\pm$ | $0.3$ | $\displaystyle 84.9$ | $\pm$ | $1.4$ |  |  | $82.7$ | $\pm$ | $1.3$ | $70.7$ | $\pm$ | $10.1$ | $84.0$ | $\pm$ | $1.0$
Hopper |  | $82.9$ | $\pm$ | $3.5$ | $\displaystyle 97.4$ | $\pm$ | $2.7$ | $\displaystyle 100.5$ | $\pm$ | $0.5$ |  |  | $58.6$ | $\pm$ | $11.2$ | $52.1$ | $\pm$ | $1.8$ | $\displaystyle 91.9$ | $\pm$ | $4.3$
Medium | HalfCheetah |  | $\displaystyle 48.4$ | $\pm$ | $0.1$ | $45.4$ | $\pm$ | $0.1$ | $45.1$ | $\pm$ | $0.1$ |  |  | $\displaystyle 48.6$ | $\pm$ | $0.1$ | $45.3$ | $\pm$ | $0.2$ | $47.6$ | $\pm$ | $0.3$
Walker2d |  | $81.7$ | $\pm$ | $1.4$ | $82.1$ | $\pm$ | $0.9$ | $77.8$ | $\pm$ | $3.6$ |  |  | $84.8$ | $\pm$ | $0.1$ | $85.2$ | $\pm$ | $0.2$ | $\displaystyle 86.3$ | $\pm$ | $0.3$
Hopper |  | $63.6$ | $\pm$ | $0.8$ | $59.7$ | $\pm$ | $2.0$ | $62.8$ | $\pm$ | $1.2$ |  |  | $\displaystyle 62.4$ | $\pm$ | $0.9$ | $57.4$ | $\pm$ | $0.4$ | $\displaystyle 63.1$ | $\pm$ | $0.6$
| Total |  | $46.9$ | $\pm$ | $0.4$ | $47.0$ | $\pm$ | $0.4$ | $\displaystyle 47.9$ | $\pm$ | $0.3$ |  |  | $44.5$ | $\pm$ | $1.1$ | $41.1$ | $\pm$ | $1.1$ | $\displaystyle 49.5$ | $\pm$ | $0.9$
Maze2d | UMaze |  | $42.6$ | $\pm$ | $0.4$ | $42.9$ | $\pm$ | $1.8$ | $43.8$ | $\pm$ | $3.5$ |  |  | $50.0$ | $\pm$ | $2.4$ | $33.8$ | $\pm$ | $3.0$ | $\displaystyle 76.2$ | $\pm$ | $17.4$
Medium |  | $38.5$ | $\pm$ | $1.9$ | $33.4$ | $\pm$ | $3.2$ | $\displaystyle 60.0$ | $\pm$ | $13.9$ |  |  | $32.1$ | $\pm$ | $6.8$ | $24.0$ | $\pm$ | $4.0$ | $\displaystyle 89.6$ | $\pm$ | $19.9$
Large |  | $\displaystyle 50.9$ | $\pm$ | $5.8$ | $23.4$ | $\pm$ | $8.0$ | $\displaystyle 45.3$ | $\pm$ | $14.8$ |  |  | $137.2$ | $\pm$ | $20.2$ | $93.3$ | $\pm$ | $31.0$ | $131.1$ | $\pm$ | $37.5$
| Total |  | $\displaystyle 44.0$ | $\pm$ | $2.2$ | $33.2$ | $\pm$ | $1.8$ | $\displaystyle 49.7$ | $\pm$ | $9.5$ |  |  | $73.1$ | $\pm$ | $6.7$ | $50.4$ | $\pm$ | $11.1$ | $\displaystyle 99.0$ | $\pm$ | $14.5$

In order to demonstrate the flexibility of PGD, we also evaluate PGD in a
continuous generation setting, using a data generation rate closer to that of
traditional model-based methods. For this, we set $N_{\text{epochs}}=$100$$
and $N_{\text{policy}}=$10,000$$, then lower the sample size to match the
overall number of synthetic trajectories generated by periodic generation
across training. Due to the decrease in sample size, we maintain each
generated dataset across epochs in a replay buffer, with each dataset being
removed after 10 epochs.

We see similar improvements in performance against real and unguided synthetic
data under this approach, with PGD outperforming real data on 2 out of 3
environments and datasets (Figure 4). Periodic generation outperforms
continuous generation across environments and behavior policies, which we
attribute to training stability, especially when performing guidance early in
training. Regardless, both approaches consistently outperform training on real
and unguided synthetic data, demonstrating the potential of PGD as a drop-in
extension to replay and model-based RL methods.

(a) Environments

(b) Behavior policies

Figure 4: Aggregate MuJoCo performance after training on unguided or policy-
guided synthetic data under continuous and periodic dataset generation, as
well as on the real dataset. For each setting, mean return over TD3+BC and IQL
agents is marked, with standard error over 4 seeds (diffusion models and
agents) highlighted.

###  5.3 Synthetic Trajectory Analysis

We now analyze the quality of trajectories produced by PGD against those from
unguided diffusion and autoregressive world model (PETS) rollouts. In
principle, we seek to evaluate the divergence of these sampling distributions
from the true target distribution. However, this is not tractable to compute
directly, so we instead investigate two proxy objectives:

  1. 1.

Trajectory Likelihood: mean log-likelihood of actions under the target policy;
and

  2. 2.

Dynamics Error: mean squared error between states in the synthetic trajectory
and real environment, when rolled out with the same initial state and action
sequence.

In our experiments, we consider trajectory diffusion and MOPO-style PETS (Chua
et al., 2018) models trained on representative datasets from the D4RL (Fu et
al., 2020) benchmark that were featured in the previous section. Specifically,
we consider the models trained on halfcheetah-medium, before sampling
trajectories with IQL target policies trained on the halfcheetah-random,
-medium, and -expert. This enables us to test the robustness of these models
to target policies far from the behavior policy, both in performance and
policy entropy.

##### Policy Guidance Increases Trajectory Likelihood

In Figure 5, we present the trajectory likelihood of synthetic trajectories
over varying degrees of guidance. Unsurprisingly, unguided diffusion generates
low probability trajectories for all target policies, due to it directly
modeling the behavior distribution. However, as we increase the guidance
coefficient $\lambda$, trajectory likelihood increases monotonically under
each target policy. Furthermore, this effect is robust across target policies,
giving the ability to sample high-probability trajectories with OOD target
policies. The value of $\lambda$ required to achieve the same action
likelihood as direct action sampling (PETS) varies with the target policy.
Since this threshold increases with target policy performance, we hypothesize
that it increases with target policy entropy. Based on this, a promising
avenue for future work is automatically tuning $\lambda$ for hyperparameter-
free guidance.

Figure 5: Action probability of synthetic trajectories generated by diffusion
and PETS models trained on halfcheetah-medium. Target policies are trained on
halfcheetah-random, halfcheetah-medium, and halfcheetah-expert datasets,
demonstrating robustness to OOD actions. Standard error over 4 diffusion model
seeds is shaded (but negligible), with mean computed over 2048 synthetic
trajectories.

##### Policy Guided Diffusion Achieves Lower Error Than Autoregressive Models

In Figure 6, we present the dynamics error of synthetic trajectories over 16
rollout steps. For a fair comparison, we fix the guidance coefficient of PGD
to $\lambda=1.0$, since this was sufficient to match the trajectory likelihood
of PETS (Figure 5). Over all target policies, PGD achieves significantly lower
error than PETS. Furthermore, PGD has similar levels of error across target
policies, while PETS suffers from significantly higher error on OOD (random
and expert) target policies. This highlights the robustness of PGD to target
policy, a critical feature for generating high-likelihood training data
throughout tabula rasa policy training.

Figure 6: Dynamics mean squared error of synthetic trajectories generated by
diffusion and PETS models trained on halfcheetah-medium. Standard error over 4
diffusion model seeds and 3 PETS seeds (via OfflineRL-Kit) is shaded, with
each generating 2048 synthetic trajectories for analysis.

##  6 Related Work

##### Model-based Offline Reinforcement Learning

Model-based methods in offline RL (Yu et al., 2020; Kidambi et al., 2020;
Rigter et al., 2022; Lu et al., 2022) are designed to augment the offline
buffer with additional on-policy samples in order to mitigate distribution
shift. This is typically done by rolling out a policy in a learned world model
(Janner et al., 2019) and applying a suitable pessimism term in order to
account for dynamics model errors. While these methods share the same overall
motivation as our paper, the empirical realization is quite different. In
particular, forward dynamics models are liable to compounding errors over long
horizons, resulting in model exploitation, whereas our trajectories are
generated in a single step.

##### Model-free Offline Reinforcement Learning

Model-free methods in offline RL typically tackle the out-of-sample issue by
applying conservatism to the value function or by constraining the policy to
remain close to the data. For example, CQL (Kumar et al., 2020) and EDAC (An
et al., 2021) both aim to minimize the values of out-of-distribution actions.
Meanwhile, BCQ (Fujimoto et al., 2019) ensures that actions used in value
targets are in-distribution with the behavioral policy using constrained
optimization. We take the opposite approach in this paper: by enabling our
diffusion model to generate on-policy samples without diverging from the
behavior distribution, we reduce the need for conservatism.

##### Diffusion in Reinforcement Learning

Diffusion models are a flexible method for data augmentation in reinforcement
learning. SynthER (Lu et al., 2023) uses unguided diffusion models to upsample
offline or online RL datasets, which are then used by model-free off-policy
algorithms. While this improves performance, SynthER uses unguided diffusion
to model the behavior distribution, resulting in the same issue of
distributional shift. Similarly, MTDiff (He et al., 2023) considers unguided
data generation in multitask settings.

Diffusion models have also been used to train world models. Zhang et al.
(2023) train a world model for sensor observations by first tokenizing using
VQ-VAE and then predicting future observations via discrete diffusion. Alonso
et al. (2023) also train a world model using diffusion and demonstrate it can
more accurately predict future observations. However, neither of these
approaches model the whole trajectory, thereby suffering from compounding
error, nor do they apply policy guidance. Parallel to this work, Rigter et al.
(2023) use guidance from a policy to augment a diffusion world model for
online RL. By contrast, we focus on the offline RL setting, provide a
theoretical derivation and motivation for the trajectory distribution modeled
by policy guidance, and demonstrate improvements in downstream policy
performance.

Diffusion models are also used elsewhere in reinforcement learning. For
example, Diffuser (Janner et al., 2022) and Decision Diffuser (Ajay et al.,
2023) use trajectory diffusion models for planning and to bias planned
trajectories towards high return. By contrast, we use on-policy guidance and
train on the generated data. Diffusion models have also been used as an
expressive policy class (Wang et al., 2023) for $Q$-learning, showing
improvement over MLPs.

##  7 Conclusion

We presented policy-guided diffusion, a method for controllable generation of
synthetic trajectories in offline RL. We provided a theoretical analysis of
existing approaches to synthetic experience generation, identifying the
advantages of direct trajectory generation compared to autoregressive methods.
Motivated by this, we proposed PGD under the direct approach, deriving the
regularized target distribution modeled by policy guidance.

Evaluating against PETS deep ensembles, a state-of-the-art autoregressive
approach, we found that PGD can generate synthetic experience at the same
target policy likelihood with significantly lower dynamics error. Furthermore,
we found consistent improvements in downstream agent performance over a range
of environments and behavior policies when trained on policy-guided synthetic
data, against real and unguided synthetic experience.

By addressing the out-of-sample issue through synthetic data, we hope that
this work enables the development of less conservative algorithms for offline
RL. There are a range of promising avenues for future work, including
automatically tuning the guidance coefficient for hyperparameter-free
guidance, leveraging on-policy RL techniques with policy-guided data, and
extending this approach to large-scale video generation models.

#### Acknowledgments

We thank Mattie Fellows and Sebastian Towers for their insights regarding our
method’s theoretical underpinning, as well as Alex Goldie and the NeurIPS 2023
Robot Learning Workshop reviewers for their helpful feedback. Matthew Jackson
is funded by the EPSRC Centre for Doctoral Training in Autonomous Intelligent
Machines and Systems, and Amazon Web Services.

## References

  * Ajay et al. (2023) Anurag Ajay, Yilun Du, Abhi Gupta, Joshua B. Tenenbaum, Tommi S. Jaakkola, and Pulkit Agrawal.  Is conditional generative modeling all you need for decision making?  In _The Eleventh International Conference on Learning Representations_ , 2023.  URL https://openreview.net/forum?id=sP1fo2K9DFG.
  * Alonso et al. (2023) Eloi Alonso, Adam Jelley, Anssi Kanervisto, and Tim Pearce.  Diffusion world models.  2023\.
  * An et al. (2021) Gaon An, Seungyong Moon, Jang-Hyun Kim, and Hyun Oh Song.  Uncertainty-based offline reinforcement learning with diversified q-ensemble.  _Advances in neural information processing systems_ , 34:7436–7447, 2021.
  * Ball et al. (2021) Philip J Ball, Cong Lu, Jack Parker-Holder, and Stephen Roberts.  Augmented world models facilitate zero-shot dynamics generalization from a single offline environment.  In Marina Meila and Tong Zhang (eds.), _Proceedings of the 38th International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , pp. 619–629. PMLR, 18–24 Jul 2021.
  * Bansal et al. (2023) Arpit Bansal, Hong-Min Chu, Avi Schwarzschild, Soumyadip Sengupta, Micah Goldblum, Jonas Geiping, and Tom Goldstein.  Universal guidance for diffusion models.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pp. 843–852, 2023.
  * Bradbury et al. (2018) James Bradbury, Roy Frostig, Peter Hawkins, Matthew James Johnson, Chris Leary, Dougal Maclaurin, George Necula, Adam Paszke, Jake VanderPlas, Skye Wanderman-Milne, and Qiao Zhang.  JAX: composable transformations of Python+NumPy programs, 2018.  URL http://github.com/google/jax.
  * Chua et al. (2018) Kurtland Chua, Roberto Calandra, Rowan McAllister, and Sergey Levine.  Deep reinforcement learning in a handful of trials using probabilistic dynamics models, 2018.
  * Dhariwal & Nichol (2021) Prafulla Dhariwal and Alexander Nichol.  Diffusion models beat gans on image synthesis.  _Advances in neural information processing systems_ , 34:8780–8794, 2021.
  * Fu et al. (2020) 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.
  * Fujimoto & Gu (2021) Scott Fujimoto and Shixiang Shane Gu.  A minimalist approach to offline reinforcement learning.  In _Thirty-Fifth Conference on Neural Information Processing Systems_ , 2021.
  * Fujimoto et al. (2019) Scott Fujimoto, David Meger, and Doina Precup.  Off-policy deep reinforcement learning without exploration.  In Kamalika Chaudhuri and Ruslan Salakhutdinov (eds.), _Proceedings of the 36th International Conference on Machine Learning_ , volume 97 of _Proceedings of Machine Learning Research_ , pp. 2052–2062. PMLR, 09–15 Jun 2019.  URL https://proceedings.mlr.press/v97/fujimoto19a.html.
  * He et al. (2023) Haoran He, Chenjia Bai, Kang Xu, Zhuoran Yang, Weinan Zhang, Dong Wang, Bin Zhao, and Xuelong Li.  Diffusion model is an effective planner and data synthesizer for multi-task reinforcement learning.  _arXiv preprint arXiv:2305.18459_ , 2023.
  * Ho et al. (2020) Jonathan Ho, Ajay Jain, and Pieter Abbeel.  Denoising diffusion probabilistic models.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 6840–6851. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper/2020/file/4c5bcfec8584af0d967f1ab10179ca4b-Paper.pdf.
  * Hyvärinen & Dayan (2005) Aapo Hyvärinen and Peter Dayan.  Estimation of non-normalized statistical models by score matching.  _Journal of Machine Learning Research_ , 6(4), 2005.
  * Janner et al. (2019) Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine.  When to trust your model: Model-based policy optimization.  In _Advances in Neural Information Processing Systems_ , 2019.
  * Janner et al. (2022) Michael Janner, Yilun Du, Joshua Tenenbaum, and Sergey Levine.  Planning with diffusion for flexible behavior synthesis.  In _International Conference on Machine Learning_ , 2022.
  * Karras et al. (2022) Tero Karras, Miika Aittala, Timo Aila, and Samuli Laine.  Elucidating the design space of diffusion-based generative models.  _Advances in Neural Information Processing Systems_ , 35:26565–26577, 2022.
  * Kidambi et al. (2020) Rahul Kidambi, Aravind Rajeswaran, Praneeth Netrapalli, and Thorsten Joachims.  Morel: Model-based offline reinforcement learning.  _Advances in neural information processing systems_ , 33:21810–21823, 2020.
  * Kirk et al. (2023) Robert Kirk, Amy Zhang, Edward Grefenstette, and Tim Rocktäschel.  A survey of zero-shot generalisation in deep reinforcement learning.  _Journal of Artificial Intelligence Research_ , 76:201–264, January 2023.  ISSN 1076-9757.  doi: 10.1613/jair.1.14174.  URL http://dx.doi.org/10.1613/jair.1.14174.
  * Kostrikov et al. (2021) Ilya Kostrikov, Ashvin Nair, and Sergey Levine.  Offline reinforcement learning with implicit q-learning.  _arXiv preprint arXiv:2110.06169_ , 2021.
  * Kumar et al. (2020) Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine.  Conservative q-learning for offline reinforcement learning.  _Advances in Neural Information Processing Systems_ , 33:1179–1191, 2020.
  * Levine et al. (2020) 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.
  * Lu et al. (2022) Cong Lu, Philip Ball, Jack Parker-Holder, Michael Osborne, and Stephen J. Roberts.  Revisiting design choices in offline model based reinforcement learning.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=zz9hXVhf40.
  * Lu et al. (2023) Cong Lu, Philip J. Ball, Yee Whye Teh, and Jack Parker-Holder.  Synthetic experience replay.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.  URL https://openreview.net/forum?id=6jNQ1AY1Uf.
  * Ma et al. (2023) Jiajun Ma, Tianyang Hu, Wenjia Wang, and Jiacheng Sun.  Elucidating the design space of classifier-guided diffusion generation.  _arXiv preprint arXiv:2310.11311_ , 2023.
  * Precup (2000) Doina Precup.  Eligibility traces for off-policy policy evaluation.  _Computer Science Department Faculty Publication Series_ , pp. 80, 2000.
  * Precup et al. (2000) Doina Precup, Richard S. Sutton, and Satinder P. Singh.  Eligibility traces for off-policy policy evaluation.  In _Proceedings of the Seventeenth International Conference on Machine Learning_ , ICML ’00, pp. 759–766, San Francisco, CA, USA, 2000. Morgan Kaufmann Publishers Inc.  ISBN 1558607072.
  * Rigter et al. (2022) Marc Rigter, Bruno Lacerda, and Nick Hawes.  RAMBO-RL: Robust adversarial model-based offline reinforcement learning.  In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho (eds.), _Advances in Neural Information Processing Systems_ , 2022.  URL https://openreview.net/forum?id=nrksGSRT7kX.
  * Rigter et al. (2023) Marc Rigter, Jun Yamada, and Ingmar Posner.  World models via policy-guided trajectory diffusion, 2023.
  * Ronneberger et al. (2015) Olaf Ronneberger, Philipp Fischer, and Thomas Brox.  U-net: Convolutional networks for biomedical image segmentation.  In _Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18_ , pp. 234–241. Springer, 2015.
  * Sims et al. (2024) Anya Sims, Cong Lu, and Yee Whye Teh.  The edge-of-reach problem in offline model-based reinforcement learning, 2024.
  * Sohl-Dickstein et al. (2015) Jascha Sohl-Dickstein, Eric Weiss, Niru Maheswaranathan, and Surya Ganguli.  Deep unsupervised learning using nonequilibrium thermodynamics.  In Francis Bach and David Blei (eds.), _Proceedings of the 32nd International Conference on Machine Learning_ , volume 37 of _Proceedings of Machine Learning Research_ , pp. 2256–2265, Lille, France, 07–09 Jul 2015. PMLR.  URL https://proceedings.mlr.press/v37/sohl-dickstein15.html.
  * Sun (2023) Yihao Sun.  Offlinerl-kit: An elegant pytorch offline reinforcement learning library.  https://github.com/yihaosun1124/OfflineRL-Kit, 2023.
  * Sutton & Barto (2018) Richard S. Sutton and Andrew G. Barto.  _Reinforcement Learning: An Introduction_.  The MIT Press, second edition, 2018.  URL http://incompleteideas.net/book/the-book-2nd.html.
  * Todorov et al. (2012) Emanuel Todorov, Tom Erez, and Yuval Tassa.  Mujoco: A physics engine for model-based control.  _IEEE_ , pp. 5026–5033, 2012.  URL http://dblp.uni-trier.de/db/conf/iros/iros2012.html#TodorovET12.
  * Wang et al. (2023) Zhendong Wang, Jonathan J Hunt, and Mingyuan Zhou.  Diffusion policies as an expressive policy class for offline reinforcement learning, 2023.
  * Yu et al. (2020) Tianhe Yu, Garrett Thomas, Lantao Yu, Stefano Ermon, James Y Zou, Sergey Levine, Chelsea Finn, and Tengyu Ma.  Mopo: Model-based offline policy optimization.  In H. Larochelle, M. Ranzato, R. Hadsell, M.F. Balcan, and H. Lin (eds.), _Advances in Neural Information Processing Systems_ , volume 33, pp. 14129–14142. Curran Associates, Inc., 2020.  URL https://proceedings.neurips.cc/paper/2020/file/a322852ce0df73e204b7e67cbbef0d0a-Paper.pdf.
  * Zhang et al. (2023) Lunjun Zhang, Yuwen Xiong, Ze Yang, Sergio Casas, Rui Hu, and Raquel Urtasun.  Learning unsupervised world models for autonomous driving via discrete diffusion.  _arXiv preprint arXiv:2311.01017_ , 2023.

## Appendix

##  Appendix A Hyperparameters

We open-source our implementation at https://github.com/EmptyJackson/policy-
guided-diffusion.

###  A.1 Diffusion Model

For the diffusion model, we used a U-Net architecture (Ronneberger et al.,
2015) with hyperparameters outlined in Table 3. We transformed the trajectory
by stacking the observation, action, reward, and done flags for each
transition, before performing 1D convolution across the sequence of
transitions.

Table 3: U-Net hyperparameters Hyperparameter | Value
---|---
Trajectory length | 16
Kernel size | 3
Features | 1024
U-Net blocks | 3
Batch size | 16
Dataset epochs | 250
Optimizer | Adam
Learning rate | $2\times 10^{-3}$
LR schedule | Cosine decay

###  A.2 Diffusion Sampling

We use EDM (Karras et al., 2022) for diffusion sampling, retaining many of the
default hyperparameters from Lu et al. (2023) (Table 4). We tuned the number
of diffusion timesteps, finding diminishing improvement in dynamics error
beyond 256 timesteps.

Table 4: EDM hyperparameters Hyperparameter | Value
---|---
Diffusion timesteps | 256
$\text{S}_{\text{churn}}$ | 80
$\text{S}_{\text{noise}}$ | $1.003$
$\text{S}_{\text{tmax}}$ | 50
$\text{S}_{\text{tmin}}$ | $0.05$
$\sigma_{\text{max}}$ | 80
$\sigma_{\text{min}}$ | $0.002$

##  Appendix B Noised Target Distribution

To model the target distribution with diffusion, we require the noise-
conditioned score function $\nabla_{\bm{\hat{\tau}}}\log p_{\mkern
1.0mu\text{target}}(\bm{\hat{\tau}};\sigma)$ for the target distribution.
However, since we do not have access to samples from
$p_{\text{target}}(\bm{\hat{\tau}};\sigma)$, one might wish to apply a
factorization of the target distribution, such as

$\displaystyle p_{\mkern 1.0mu\text{target}}(\bm{\tau})$
$\displaystyle=p_{\mkern
1.0mu\text{off}}(\bm{\tau})\prod_{t=0}^{H-1}\frac{\pi_{\text{target}}(a|s)}{\pi_{\text{off}}(a|s)}\text{,}$
(15)

before modeling its terms separately. However, by applying independent
Gaussian noise to each of the elements in $\bm{\hat{\tau}}$, we lose
conditional independence between contiguous states and actions—i.e.,
$p_{\mkern
1.0mu\text{target}}(\hat{a}_{t}|\bm{\hat{\tau}}\\!\setminus\\!\hat{a}_{t};\sigma)\neq
p_{\mkern 1.0mu\text{target}}(\hat{a}_{t}|\hat{s}_{t};\sigma)$—preventing us
from applying an equivalent factorization. Due to this, we must approximate
$\nabla_{\bm{\hat{\tau}}}\log p_{\mkern
1.0mu\text{target}}(\bm{\hat{\tau}};\sigma)$ directly, as we propose in
Section 4.1.

##  Appendix C Behavior-Regularized Target Distribution

Intuitively, the behavior-regularized target distribution transforms the
target distribution by increasing the likelihood of actions under the behavior
policy. As is typical in offline RL (Kumar et al., 2020; Fujimoto & Gu, 2021;
Fujimoto et al., 2019), regularizing the policy towards the behavior
distribution is required in order to avoid out-of-sample states and
consequently minimize value overestimation. Rather than regularizing the
policy, PGD shifts this regularization to the data generation process, which
helps our guided samples remain in-distribution with respect to the diffusion
model, and thus less susceptible to model error.

Moreover, we note that this type of regularization is not immediately
available for prior autoregressive world models, and thus they typically
penalize reward by dynamics error (Yu et al., 2020; Kidambi et al., 2020; Lu
et al., 2022) in an ad-hoc fashion in order to avoid model exploitation. In
contrast, PGD presents a _natural mechanism for behavioral regularization
during data generation_ , making offline policy optimization without
regularization a promising path for future work.

##  Appendix D Agent Training with Policy-Guided Diffusion

In LABEL:alg:training, we present pseudocode for training an agent with
synthetic experience generated by PGD. PGD is agnostic to the underlying
offline RL algorithm used to train the target policy, making it a drop-in
extension to any model-free method.