XX XX 8 July/August

# TinyIREE: An ML Execution Environment for Embedded Systems from Compilation
to Deployment

Hsin-I Cindy Liu  Google, Mountain View, CA, USA  Marius Brehler  Fraunhofer
IML, Dortmund, Germany  Mahesh Ravishankar  Google, Seattle, WA, USA  Nicolas
Vasilache  Google, Zürich, Switzerland  Ben Vanik  Google, Seattle, WA, USA
Stella Laurenzo  Google, Seattle, WA, USA

(2022)

###### Abstract

Machine learning model deployment for training and execution has been an
important topic for industry and academic research in the last decade. Much of
the attention has been focused on developing specific toolchains to support
acceleration hardware. In this paper, we present IREE, a unified compiler and
runtime stack with the explicit goal to scale down machine learning programs
to the smallest footprints for mobile and edge devices, while maintaining the
ability to scale up to larger deployment targets. IREE adopts a compiler-based
approach and optimizes for heterogeneous hardware accelerators through the use
of the MLIR compiler infrastructure which provides the means to quickly design
and implement multi-level compiler intermediate representations (IR). More
specifically, this paper is focused on TinyIREE, which is a set of deployment
options in IREE that accommodate the limited memory and computation resources
in embedded systems and bare-metal platforms, while also demonstrating IREE’s
intuitive workflow that generates workloads for different ISA extensions and
ABIs through LLVM.

††journal: IEEE Micro

Machine learning has attracted a lot of attention in the last decade, both in
industry and academic research. As new, deeper machine learning (ML) model
architectures emerge and bigger datasets become available, the demands on the
hardware are also increasing. This has led to the development of new processor
architectures that accelerate ML model training and execution (inference).
However, to fully take advantage of such acceleration hardware, the
development toolchains require new deployment and compilation flows. In
addition, special considerations must be made when deploying models to mobile
and edge devices, due to system resource and power constraints.

A large portion of investment in the industry has targeted use cases that
scale up to the most powerful deployment targets, leaving low-power and
embedded devices poorly catered to. IREE (Intermediate Representation
Execution Environment) was created to address this imbalance by designing a
unified compiler and runtime stack with the explicit goal to scale down to the
smallest footprints while retaining its ability to scale up to larger
deployments (and integrate with higher-level distributed runtimes). This focus
on model portability has driven IREE to effectively target bare-metal/embedded
CPUs and microcontrollers. This paper presents how that portability is
achieved.

Traditionally, targeting microcontrollers by public ML inference frameworks
has largely been done by kernel-based, op-by-op runtimes that are hand-adapted
and optimized for a small set of frequently used operators. One approach to
ease the deployment of models to microcontrollers is TensorFlow Lite for
Microcontrollers (TFLM), a framework specifically designed to support model
execution on microcontrollers [1]. It uses a small runtime library ($\sim
16\text{\,}\mathrm{kB}$ for Arm Cortex-M3) that supports a subset of
TensorFlow operators, with hardware-specific, high-level operator kernel
implementations. The graph optimization still utilizes the TensorFlow Lite
compilation flow and performs optimizations at the operator level.

Whereas TFLM is a runtime library with limited hardware-specific kernels,
compiler-based approaches have also recently been applied to accelerators and
embedded systems. One of those compilers is Glow [2], which uses a two-phase
intermediate representation (IR) to lower a neural network graph. It uses a
high-level intermediate representation to apply domain-specific optimizations.
This is followed by a lower-level IR, allowing the compiler to apply memory-
related optimizations. These optimizations include instruction scheduling,
static memory allocation, and copy elimination.

Another compiler is Apache TVM [3], which performs both high-level graph
reconfiguration and low-level operation optimization in various intermediate
representations. Although TVM targets various device backends such as CPU,
GPU, DSP, etc., for bare-metal microcontrollers, a MicroTVM extension is
required to support the model execution scheduling and memory management, in
order to make up for the lack of an operating system.

Figure 1: A simplified end-to-end flow diagram of IREE.

This brings us to how we approach this challenge with IREE. IREE is an end-to-
end compiler and runtime framework for model execution based on the MLIR
compiler infrastructure [4]. It follows a compiler-based approach that
converts ML models into an intermediate representation that allows for
analysis and optimization of the ML model while generating code to target
heterogeneous hardware accelerators [5]. As shown in Figure 1, IREE can take
various model representations as inputs, and generate executable formats for
different hardware targets. IREE treats the ML model as just another program,
represents it with various MLIR _dialects_ , and eventually generates the
binary that can be run on the target architecture. An MLIR dialect can be
roughly thought of as a logical level of IR. This multi-phase compilation
through a progression of composable IRs [5] is a key differentiator because it
allows us to more effectively retarget and generalize approaches and
components used to target larger scale parts for specialized, deeply embedded
uses. Such generalization is much harder to achieve in classic, emitter based,
two phase compilers and is not possible at all in executors such as TFLM.

IREE is designed to target various edge devices and accelerators, including
CPUs, GPUs, and various ML accelerators. As a result, aspects of mobile
deployment, like small binary size, inter-operability with applications, and
portability are first-class concerns. In this paper, we focus on TinyIREE, the
subset of IREE options that generate a compact workload and runtime library
that’s optimized for embedded systems without an operating system. These
options enable seamless support for different CPU architectures and ISA
extensions, while also providing an interface to support embedded ML
accelerators. In Figure 1, the TinyIREE options are colored, whereas the
general IREE compilation and deployment options are grayed out.

In the following sections, we present IREE’s compilation flow and explain some
of the TinyIREE deployment options. Furthermore, we show results for a
compiled MobileNet Single-Shot Detector (SSD) model before giving a
conclusion.

## IREE COMPILATION FLOW

IREE’s compilation flow uses the MLIR approach of progressively lowering an ML
program into an executable for the target hardware. The primary vehicle in
MLIR to achieve this is _dialects_. A dialect in MLIR is a collection of
operations and associated types that allow you to represent a program.
Different dialects can be used to model the program at different levels of
abstraction.

For example, the IREE importers in Figure 1 show some different dialects used
in IREE to convert an ML program into machine executable code. The TOSA
(Tensor Operator Set Architecture) [6] / MHLO (Meta High Level Optimizer)
dialects model the ML program at the level of tensor operations. The Linalg
dialect represents perfectly nested loop computation in a succinct, easy-to-
manipulate manner, making it easy to implement transformations like fusion,
tiling, loop interchange, etc. The Vector dialect models parts of the program
in terms of operations on virtual vectors. Finally, the LLVM dialect is a
_leaf dialect_ that allows MLIR programs to be translated into LLVM IR for
compilation into an executable binary. Because each of these dialects are
defined using a common infrastructure, at each level you can apply common
compiler optimizations like Common Subexpression Elimination (CSE), Dead Code
Elimination (DCE), etc. It is also possible to apply custom optimizations
within each dialect that use the semantic information captured by operations
in that dialect. For example, in the TOSA dialect, one can easily recognize
$1{\times}1$ 2D convolution operations and convert them into matrix-matrix
multiply operations that can then be executed more efficiently. The process of
converting from a dialect at a higher-level of representation (like TOSA/MHLO)
through intermediate dialects into an executable binary—with each subsequent
dialect using a more general or less specialized representation—is known as
progressive lowering. The rest of the section provides an overview of the
program representation for each of the dialects supported by IREE importers.

Frontend dialects like TOSA and MHLO represent the program with operations
like add, convolution, and dot product. For example, Listing 1 shows a simple
sequence with a dot product followed by bias-add in the TOSA dialect. Notice
that built-in MLIR types like tensor use ? to represent dynamic dimensions
(i.e., dimensions whose extent is not known at compile time). As a result,
MLIR has first-class support for representing dynamically shaped computations.

Listing 1: Dot + Bias add in TOSA dialect.

⬇

%2 = "tosa.matmul"(%0, %1) : ( tensor$<$?x8xf32$>$, tensor$<$8x?xf32$>$ )
$\rightarrow$ tensor$<$?x?xf32$>$

%3 = "tosa.add"(%2, %3) : ( tensor$<$?x?xf32$>$, tensor$<$?x?xf32$>$ )
$\rightarrow$ tensor$<$?x?xf32$>$

The program is then lowered into the Linalg dialect. At this level, operations
are represented as perfectly nested loop computations with the body of the
innermost loop performing a scalar operation. This representation is intended
to make it easier to apply traditional loop transformations like loop fusion,
loop tiling, and loop interchange. This modeling is inspired by the polyhedral
compilation domain where the computation is represented using an orthonormal
space, with one axis for every loop of the computation. This orthonormal space
is referred to as the _iteration space_ of the computation. The values of the
induction variables in a loop are represented as points along the
corresponding axis. Each point in this orthonormal space represents the
sequence of scalar operations performed within the innermost loop. Listing 2
shows the Linalg dialect’s representation of a GEMM operation of the form
$D_{ij}=C_{ij}+A_{ik}\cdot B_{kj}$. The number of loops is captured by the
number of entries in the list _iterator_types_. The data dependence between
the iterations of a loop are specified as either _parallel_ (for no
dependence) or _reduction_ (for reduction-type dependence).

Listing 2: Matrix-matrix multiply in Linalg dialect.

⬇

#map0 = affine_map$<$(i, j, k) $\rightarrow$ (i, k)$>$

#map1 = affine_map$<$(i, j, k) $\rightarrow$ (k, j)$>$

#map2 = affine_map$<$(i, j, k) $\rightarrow$ (i, j)$>$

%3 = linalg.generic {

iterator_types = ["parallel", "parallel", "reduction"],

indexing_maps = [#map0, #map1, #map2]}

ins(%0, %1 : tensor$<$?x?xf32$>$, tensor$<$?x?xf32$>$)

outs(%2 : tensor$<$?x?xf32$>$) {

^bb0(%arg0 : f32, %arg1 : f32, %arg2 : f32):

%0 = mulf %arg0, %arg1 : f32

%1 = addf %arg1, %arg2 : f32

linalg.yield %1 : f32

} $\rightarrow$ tensor$<$?x?xf32$>$<

Similar to iteration space, each operand is also represented by an orthonormal
space of the same dimensionality as the operand. This space is referred to as
the _data space_ of the operand. For example, in a GEMM operation, all the
operands are represented by a 2D data space. The maps in _indexing_maps_
capture the data access pattern for each operand. In each map, the domain
represents a point in the iteration space, and the range represents the point
in the operand’s data space that is accessed by the map.

The _region_ of the operation (the sequence of operations between the { and })
represents the computation performed at each point in the iteration space.
Each Linalg operation has a region with a single basic block whose _arguments_
represent the scalar values, which are obtained from the operands using the
indexing map. The yielded value is the result of the computation at that point
in the iteration space. For example, in Listing 2, each point in the iteration
space performs a multiply-add operation.

All tensor operations of TOSA/MHLO can be lowered to the linalg.generic
operation of the Linalg dialect. Two such operations that have a producer-
consumer relationship can be _fused_ into another linalg.generic operation by
using just the information in the iterator types and the indexing maps. This
fusion encompasses all _element-wise fusion_ optimizations (also known as loop
fusion in XLA) without having to reason about the actual computation performed
by each operation. Doing the same fusion at, say, the TOSA/MHLO level, would
result in a combinatorial explosion in the number of fusion possibilities.
Notice that the class of operations that can be fused in this manner is larger
than pure elementwise operations, including permutations and broadcasts.

Operations in Linalg dialect can also be tiled using just the information
provided by the iterator types and indexing maps. This transformation allows
splitting the computation at each layer into smaller tiles of a similar
computation. For example, a GEMM operation can be tiled into several smaller
GEMM operations. These smaller tiles can be computed in parallel, and can
therefore be distributed to different threads. Then, the computation
represented by each tile is encapsulated within a dispatch region. Each
dispatch region contains code that has to be executed on the device in an
atomic fashion. At this point, each computation is split into two parts for
execution: The dispatch region (the tiled computation) and the code to execute
host-side computation (VM commands), which specifies the order in which
dispatch regions are executed. The code within dispatch regions is then
transformed further using fusion, loop interchange, etc., to achieve better
cache locality and data access patterns.

Finally, the code is lowered into the _Vector_ dialect. This dialect
represents high-level, retargetable, vector instructions available on the
target hardware. Transformations within this dialect allow the generated code
to use efficient vector reads/writes and special fixed function units
available on the target architecture, like the Arm sdot instructions.

The last step in the device-side compilation is lowering the program into the
LLVM dialect, which allows the program to be mechanically converted into LLVM
IR. The LLVM compilation stack can then generate the binary code for the
target architecture. The LLVM target flags can be used to select which CPU
architecture, ABI, and ISA extensions to use. For example, to build the module
for x86_64 CPU, apply the following flag:

⬇

- -iree-llvm-target-triple=x86_64-pc-linux-elf

Whereas for RISC-V 32-bit CPU with multiplication and floating point ISA
extension support, use these flags:

⬇

- -iree-llvm-target-triple=riscv32-pc-linux-elf

- -iree-llvm-target-cpu=generic-rv32

- -iree-llvm-target-cpu-features=+m,+f

- -iree-llvm-target-abi=ilp32

In addition, for Armv7E-M CPU, use the following flags:

⬇

- -iree-llvm-target-triple=armv7em-pc-linux-elf

- -iree-llvm-target-float-abi=hard

The aforementioned optimization process can easily be applied to other non-ML-
specific linear algebra operations, as long as these operations are in the
appropriate format. In other words, IREE can also optimize operations that
perform input pre-processing and output post-processing as part of the
program. For example, it is typical to perform color space conversion or image
resizing of the image data stream for vision ML programs, and these operations
can be fused within the program by following the same compilation flow.

## TINYIREE DEPLOYMENT OPTIONS

The output of the compilation flow described above is an IREE module, which
contains a segment of VM commands to control the buffer setup and a collection
of dispatch regions as workloads. These workloads enfold the ML program that
can be dispatched to the target devices. As shown in the bottom part of Figure
1, when targeting embedded systems, the user needs to build an IREE runtime
application, which includes the IREE runtime library to configure the VM
setup, the workload loader, and a workload scheduler. The latter two are
defined as the device hardware abstraction layer (HAL) driver.

Before compiling the final executable, we can further break down the IREE
module structure with different deployment strategies, as described in the
following subsections. As shown in Figure 2, IREE provides a flexible set of
tools for various deployment scenarios on CPUs, including static and dynamic
embedded libraries, whereas the VM control support can be implemented with a
bytecode renderer or C source code. Here, the VM bytecode is stored in a file
in FlatBuffer format (vmfb), with the option of embedding the device workload
in the format of the shared dynamic library.

Figure 2: IREE executable deployment options for TinyIREE.

In addition to these deployment paths for embedded CPUs, developers can also
utilize the existing HAL interface to interact with customized ML
accelerators. Moreover, while IREE runtime can support dynamic deployments,
embedded systems can even bypass the IREE runtime entirely and deploy a
specific workload to achieve the most efficient execution.

### Virtual machine control with EmitC

VM commands control the resource ownership and execution flow, and they are
encoded as another MLIR dialect: the VM dialect. These operations can be
serialized to a bytecode module that is interpreted at runtime. Another option
is to serialize these operations into C source code. This is achieved by
converting the VM operations into the EmitC dialect, an MLIR dialect that
allows the generation of C/C++ from MLIR. After converting the VM operations
into EmitC operations, these are translated into C via the Cpp emitter. EmitC
contains a call operation that represents a C function call. Amongst other
things, this type of operation is used to directly call the VM API instead of
serializing the calls into an equivalent bytecode representation. Therefore,
the bytecode interpreter is not linked into the executable, thus reducing its
size.
An example of using EmitC is provided in the static_library_demo_c sample in
the IREE repository.

### ML workload as a static or dynamic library

The ML workload can be included either in a static or a dynamic library. In
both cases, IREE uses LLVM to compile the program into highly optimized
instruction streams for the particular target. The static ML workload library
can be used in combination with the serialized VM bytecode module to form an
IREE module. However, to create a binary optimized for size, the static
library should instead be used in combination with the VM operations
translated to C calls.

The dynamic ML workload library, on the other hand, provides more flexibility.
For example, it is possible to create multiple dynamic libraries optimized for
different architectures with the same entry and exit points. It is then
possible to decide at runtime which of the dynamic libraries to use. Because
embedded systems often do not provide dynamic library support, IREE uses a
fixed dynamic library ABI, and integrates its own low-overhead dynamic library
loader in the runtime library.

### Scheduler of the HAL driver

Once the command buffer is set properly by the host VM, a workload scheduler
is responsible to dispatch the device workload set in the command buffer.
Inspired by the GPU scheduler and compute APIs, IREE adopts a 3D grid topology
for the workload scheduling with regard to a specific worker structure that
contains the data buffer and schedule information, as shown in Listing 3.

Listing 3: Workload dispatch loop

⬇

for (int z = 0; z < worker.cnt.z; ++z) {

for (int y = 0; y < worker.cnt.y; ++y) {

for (int x = 0; x < worker.cnt.x; ++x) {

vec3_t work_id = {{x, y, z}};

int ret = dispatch_ptr(&st, &work_id);

}

}

}

For systems with multi-threading support, IREE provides an asynchronous task
scheduler to split up the workload with more efficient execution of directed
acyclic graphs (DAG). This allows for out-of-order and pipelined execution,
which can achieve better parallelism and better utilization of the target
device. However, IREE also supports synchronous scheduler execution by issuing
the workload sequentially. For embedded systems without the threading handlers
or a CPU without the operating system support, the synchronous scheduler
provides an intuitive path to dispatch the workload.

### Stream execution

Work scheduled by the IREE runtime is provided to the underlying system
scheduler—whether a hardware command processor or CPU thread pool—with
dependency information ordering the work submissions. This provides for a
final level of just-in-time scheduling, work balancing, and safe preemption
points for when the system is sharing limited device resources. The memory
needed during the asynchronous work is reserved from a pool that is aware of
the streaming behavior, and only as much memory as is required for the
concurrently executable work is ever allocated. This enables deep pipelines of
work to be submitted to devices without over-committing limited memory
resources. Because a majority of the memory used within a typical inference is
either constant or transient (local to only a particular invocation), this
approach reduces the at-rest memory consumption of a loaded program to mere
kilobytes, in addition to the memory required by constants, which can often be
mapped into discardable memory. This allows for several programs to be
loaded—even when retaining state across invocations—which can then be
scheduled for interleaved execution while only consuming the peak memory used
by any individual invocation.

### Buffer allocation and permission control

The IREE runtime uses device-provided memory allocators to handle the device
buffer configuration and preparation, accompanied by a system allocator to
handle the host-side VM buffers. The allocators can set the visibility of the
buffer between the host (VM) and the device explicitly, while using the
standard memory allocation libraries for the proper alignment. A memory block
can be allocated on one side (host or device), while being selectively
accessible from the other end. For example, the input buffer for the ML
program may be allocated at the host so it can be initiated with the correct
input value, while set to be visible to the device with read access so its
content can be consumed by the workload. This visibility and permission
control—in conjunction with systems supporting enclave computing—provides a
secured ML execution environment.

## RESULTS

By applying IREE’s compilation flow and choosing bare-metal-friendly
configurations to deploy ML programs, we can reduce the artifact size of both
the IREE module and the IREE runtime library.

Table 1 shows the result of passing MobileNet SSD V2 [7] through the IREE
compilation flow for different compilation targets and modes, using IREE
snapshot 20211203.686. To compile for an Armv7E-M (e.g., Arm Cortex-M4)
target, the parameters are quantized. It shows both the size of the IREE
module, packed into the FlatBuffer file type, and the size of the executable
workload (embedded within the module FlatBuffer) across different LLVM
compilation targets, each with four IREE compilation modes: Debug-dylib
(dynamic library with debug symbols included; the default), Dylib (no debug
symbols), Embedded (embedded system-friendly dynamic library), and Static
(static library). The latter two modes are the TinyIREE modes. The result
shows the workload sizes vary with the LLVM targets. In particular, the
workload size is much smaller for the targets for the embedded systems, such
as Armv7E-M, as a result of the LLVM optimization for such targets.

Table 1: IREE module FlatBuffer and the embedded device executable workload size for Mobilenet V2 SSD model compiled for different CPU architecture targets. Mode | FlatBuffer size | Workload size
---|---|---
|  ($\mathrm{kB}$) |  ($\mathrm{kB}$)
x86_64 |  |
Debug-dylib | 5,588 | 256.55
Dylib | 5,525 | 208.35
Embedded | 5,521 | 203.80
Static | 5,317 | –
Armv7E-M |  |
Debug-dylib | 5,468 | 136.55
Dylib | 5,420 | 102.98
Embedded | 5,406 | 88.84
Static | 5,317 | –
RISC-V32imf |  |
Debug-dylib | 5,500 | 168.42
Dylib | 5,437 | 120.18
Embedded | 5,437 | 120.18
Static | 5,317 | –

Figure 3 plots the workload size comparison across different targets and
modes, and the benefit of TinyIREE modes for embedded systems. Notice for the
static library mode, the ML workload becomes a library that can be optimized
during the link time, so its final size depends on how the user application
integrates the workload library. For example, when switching from the static
library to the virtual machine control with EmitC, the bytecode interpreter
can be dropped, which saves about $15\text{\,}\mathrm{kB}$ when compiling for
Armv7E-M.

Figure 3: Workload size comparison of MobileNet SSD V2 model generated by IREE flow for different CPU architectures in different deployment modes. Table 2: Runtime peak memory usage and the host library size for running MobileNet V2 SSD model with TFLite interpreter vs. IREE runtime application. |  Peak memory ($\mathrm{MB}$) |  Host library ($\mathrm{kB}$)
---|---|---
TFLite |  |
x86_64 | 20.05 | 2971.06
IREE |  |
(Embedded) |  |
x86_64 | 5.93 | 96.31
Armv7E-M | 5.93 | 79.96
RISC-V32imf | 5.93 | 152.86

In terms of execution, IREE runtime also provides more efficient memory usage,
while requiring a small runtime library. Table 2 shows the peak memory usage
comparison between x86_64 TensorFlow Lite (TFLite) benchmark tool [8] and IREE
application for Mobilenet V2 SSD model built with the dynamic embedded library
mode and synchronous HAL driver. Compared with the TFLite interpreter library,
the same workload can be deployed with a much smaller runtime library. Notice
that TFLM could provide a host library with a similar size, but it only
supports a subset of operators. Thus, not all the TFLite artifacts are
compatible with TFLM. For example, the pre-2.5.0-rc0 TFLM build can deploy the
same MobileNet V2 SSD model shown in Table 2 with $4.91\text{\,}\mathrm{MB}$
of peak memory usage and a $80.02\text{\,}\mathrm{kB}$ x86_64 host library,
but TFLM has since dropped its support for unsigned operators in favor of the
recommended signed ones, hence the support for this particular model.

## CONCLUSION

We presented IREE, a compiler-based ML execution framework that takes the full
advantage of MLIR and that can be easily configured for different target
devices. Moreover, TinyIREE options can use the same program compilation flow,
but easily adjust its runtime to target embedded systems with reduced
overhead. We also demonstrated the capability of supporting different CPU
architectures and ISA extensions by adjusting the LLVM compilation flags. As a
result, the presented framework can be optimized for multiple target backends
at the same time, while keeping enough flexibility for customization.

## APPENDIX

IREE is an open-source project, and we hope to encourage the ML community to
contribute to the work so the framework can be applicable for more targets and
applications.

Please visit https://github.com/google/iree and
https://github.com/iml130/iree-bare-metal-arm to access the source code. A set
of samples associated with IREE, including the quantized model shown in this
paper, along with other supported ML models, are available at
https://github.com/google/iree-samples.

## ACKNOWLEDGMENT

We would like to thank the members of the IREE team as well as members of the
team at Fraunhofer IML. Especially, we would like to thank Scott Todd for the
work on runtime library development and Simon Camphausen for the work on the
integration of EmitC. We would also like to thank Scott Main for the review
and feedback of the manuscript.
The work of M. Brehler was supported by the German Federal Ministry of
Education and Research (BMBF) as part of the AIA project under grant number
01IS19060A and the Competence Center Machine Learning Rhine-Ruhr (ML2R) under
grant number 01IS18038B.

## References

  * [1] P. Warden and D. Situnayake, “TinyML: Machine Learning with TensorFlow on Arduino, and Ultra-Low Power Micro-Controllers,” Sebastopol, CA: O’Reilly Media, 2019.
  * [2] N. Rotem _et al._ , “Glow: Graph Lowering Compiler Techniques for Neural Networks,” arXiv:1805.00907 [cs.PL], 2018.
  * [3] T. Chen _et al._ , “TVM: An Automated End-to-End Optimizing Compiler for Deep Learning,” 13th USENIX Symp. on Operating Syst. Design and Implementation (OSDI), pp. 579–594, 2018.
  * [4] C. Lattner _et al._ , “MLIR: Scaling Compiler Infrastructure for Domain Specific Computation,” IEEE/ACM Int. Symp. on Code Gener. and Optimization (CGO), pp. 2-14, 2021.
  * [5] N. Vasilache _et al._ , “Composable and Modular Code Generation in MLIR: A Structured and Retargetable Approach to Tensor Compiler Construction,” arXiv:2202:03293 [cs.PL], 2022.
  * [6] https://developer.mlplatform.org/w/tosa/
  * [7] M. Sandler, _et al._ , “MobileNetV2: Inverted Residuals and Linear Bottlenecks,” arxiv:1801.04381 [cs.CV], 2019.
  * [8] https://www.tensorflow.org/lite/performance/measurement

Hsin-I Cindy Liu is currently with Google, Mountain View, CA, USA. She
received the Ph. D. degree in Electrical Engineering and Computer Sciences
from University of California, Berkeley in 2010. Her research interests
include ambient intelligence sensing, DNN model architecture
research/deployment, etc. Contact her at<EMAIL_ADDRESS>
Marius Brehler is currently with Fraunhofer IML, Dortmund, Germany. He
received the B.Sc., M.Sc., and Dr.-Ing. degrees in Electrical Engineering from
Technische Universität Dortmund, Dortmund, Germany, in 2011, 2013, and 2019,
respectively. Amongst other things, his current work includes compilers for
machine learning. Contact him at<EMAIL_ADDRESS>
Mahesh Ravishankar is currently with Google, Seattle, WA, USA. He received his
Ph. D. degree in Computer Science and Engineering from The Ohio State
University, Columbus, OH in 2014, after which he worked at NVIDIA Corporation
in various teams, including the CUDA compiler team. His work primarily focuses
on compilation of ML models, with a focus on code generation for CPU and GPUs.
Contact him at<EMAIL_ADDRESS>
Nicolas Vasilache is currently with Google, Zürich, Switzerland. After a Ph.
D. at INRIA, France and a stint at Reservoir Labs, NYC on polyhedral
compilation, he joined the Facebook AI Research Lab, NYC where he worked on
fast distributed GPU libraries and led the Tensor Comprehensions systems
research project. He joined Google as an early core member of MLIR in 2018 and
now leads Google’s Structured and Retargetable Codegen efforts in MLIR.
Contact him at<EMAIL_ADDRESS>
Ben Vanik is currently with Google, Seattle, WA, USA. His experience is in
real-time graphics and hardware-accelerated compute with recent work focusing
on helping ML take advantage of the wealth of research that has been done in
those domains. Contact him at<EMAIL_ADDRESS>
Stella Laurenzo is currently with Google, Seattle, WA, USA. Her experience is
with scaling machine learning applications and approaches to small form
factors. Contact her at<EMAIL_ADDRESS>


L[1]Q[l,#1] C[1]Q[c,#1] R[1]Q[r,#1] mytabular[1]

# Video Prediction Models as Rewards
for Reinforcement Learning

Alejandro Escontrela† Ademi Adeniji† Wilson Yan†
Ajay Jain Xue Bin Peng Ken Goldberg
Youngwoon Lee Danijar Hafner Pieter Abbeel
University of California Corresponding author<EMAIL_ADDRESS>Berkeley
†Equal contribution

###### Abstract

Specifying reward signals that allow agents to learn complex behaviors is a
long-standing challenge in reinforcement learning. A promising approach is to
extract preferences for behaviors from unlabeled videos, which are widely
available on the internet. We present Video Prediction Rewards (VIPER), an
algorithm that leverages pretrained video prediction models as action-free
reward signals for reinforcement learning. Specifically, we first train an
autoregressive transformer on expert videos and then use the video prediction
likelihoods as reward signals for a reinforcement learning agent. VIPER
enables expert-level control without programmatic task rewards across a wide
range of DMC, Atari, and RLBench tasks. Moreover, generalization of the video
prediction model allows us to derive rewards for an out-of-distribution
environment where no expert data is available, enabling cross-embodiment
generalization for tabletop manipulation. We see our work as starting point
for scalable reward specification from unlabeled videos that will benefit from
the rapid advances in generative modeling. Source code and datasets are
available on the project website: https://escontrela.me/viper

##  1 Introduction

Figure 1: VIPER uses the next-token likelihoods of a frozen video prediction
model as a general reward function for various tasks.

Manually designing a reward function is laborious and often leads to
undesirable outcomes [32]. This is a major bottleneck for developing general
decision making agents with reinforcement learning (RL). A more scalable
approach is to learn complex behaviors from videos, which can be easily
acquired at scale for many applications (e.g., Youtube videos).

Previous approaches that learn behaviors from videos reward the similarity
between the agent’s current observation and the expert data distribution [37,
38, 30, 43]. Since their rewards only condition on the current observation,
they cannot capture temporally meaningful behaviors. Moreover, the approaches
with adversarial training schemes [30, 43] often result in mode collapse,
which hinders generalization.

Other works fill in actions for the (action-free) videos using an inverse
dynamics model [42, 8]. Dai et al. [8] leverage recent advances in generative
modeling to capture multi-modal and temporally-coherent behaviors from large-
scale video data. However, this multi-stage approach requires performing
expensive video model rollouts to then label actions with a learned inverse
dynamics model.

In this paper, we propose using Video Prediction Rewards (VIPER) for
reinforcement learning. VIPER first learns a video prediction model from
expert videos. We then train an agent using reinforcement learning to maximize
the log-likelihood of agent trajectories estimated by the video prediction
model, as illustrated in Figure 1. Directly leveraging the video model’s
likelihoods as a reward signal encourages the agent to match the video model’s
trajectory distribution. Additionally, rewards specified by video models
inherently measure the temporal consistency of behavior, unlike observation-
level rewards. Further, evaluating likelihoods is significantly faster than
performing video model rollouts, enabling faster training times and more
interactions with the environment.

We summarize the three key contributions of this paper as follows:

  * •

We present VIPER: a novel, scalable reward specification algorithm which
leverages rapid improvements in generative modeling to provide RL agents with
rewards from unlabeled videos.

  * •

We perform an extensive evaluation, and show that VIPER can achieve expert-
level control without task rewards on 15 DMC tasks [44], 6 RLBench tasks [20],
and 7 Atari tasks [3] (see examples in Figure 2 and Section A.5).

  * •

We demonstrate that VIPER generalizes to different environments for which no
training data was provided, enabling cross-embodiment generalization for
tabletop manipulation.

Along the way, we discuss important implementation details that improve the
robustness of VIPER.

##  2 Related Work

((a)) DeepMind Control Suite
(15 tasks)

((b)) Robot Learning Benchmark
(6 tasks)

((c)) Atari
(7 tasks)

Figure 2:  VIPER achieves expert-level control directly from pixels without
access to ground truth rewards or expert actions on 28 reinforcement learning
benchmark tasks.

Learning from observations is an active research area which has led to many
advances in imitation learning and inverse reinforcement learning [27, 1, 49].
This line of research is motivated by the fact that learning policies from
expert videos is a scalable way of learning a wide variety of complex tasks,
which does not require access to ground truth rewards, expert actions, or
interaction with the expert. Scalable solutions to this problem would enable
learning policies using the vast quantity of videos on the internet. Enabling
policy learning from expert videos can be largely categorized into two
approaches: (1) Behavioral cloning [31] on expert videos labelled with
predicted actions and (2) reinforcement learning with a reward function
learned from expert videos.

#### Labelling videos with predicted actions

An intuitive way of leveraging action-free expert videos is to guess which
action leads to each transition in expert videos, then to mimic the predicted
actions. Labelling videos with actions can be done using an inverse dynamics
model, $p(a|s,s^{\prime})$, which models an action distribution given the
current and next observations. An inverse dynamics model can be learned from
environment interactions [28, 42, 29] or a curated action-labelled dataset [2,
34]. Offline reinforcement learning [24] can be also used instead of
behavioral cloning for more efficient use of video data with predicted actions
[34]. However, the performance of this approach heavily depends on the quality
of action labels predicted by an inverse dynamics model and the quantity and
diversity of training data.

#### Reinforcement learning with videos

Leveraging data from online interaction can further improve policies trained
using unlabelled video data. To guide policy learning, many approaches have
learned a reward function from videos by estimating the progress of tasks [37,
38, 25, 23] or the divergence between expert and agent trajectories [43, 30].
Adversarial imitation learning approaches [43, 30] learn a discriminator that
discriminates transitions from the expert data and the rollouts of the current
policy. Training a policy to maximize the discriminator error leads to similar
expert and policy behaviors. However, the discriminator is prone to mode
collapse, as it often finds spurious associations between task-irrelevant
features and expert/agent labels [50, 19], which requires a variety of
techniques to stabilize the adversarial training process [7]. In contrast,
VIPER directly models the expert video distribution using recent generative
modeling techniques [10, 47], which offers stable training and strong
generalization.

#### Using video models as policies

Recently, UniPi [8] uses advances in text-conditioned video generation models
[33] to plan a trajectory. Once the future video trajectory is generated,
UniPi executes the plan by inferring low-level controls using a learned
inverse dynamics model. Instead of using slow video generations for planning,
VIPER uses video prediction likelihoods to guide online learning of a policy.

##  3 Video Prediction Rewards

Train video prediction model $p_{\theta}$ on expert videos.

while  _not converged_ do

Choose action: $a_{t}\sim\pi(x_{t})$

Step environment: $x_{t+1}\leftarrow\operatorname{env}(a_{t})$

Fill in reward: $r_{t}\leftarrow\ln p_{\theta}(x_{t+1}\;|\;x_{t-k:t})+\beta
r^{\text{expl}}_{t}$

Add transition $(x_{t},a_{t},r_{t},x_{t+1})$ to replay buffer.

Train $\pi$ from replay buffer using any RL algorithm.

Algorithm 1 VIPER

In this section, we propose Video Prediction Rewards (VIPER), which learns
complex behaviors by leveraging the log-likelihoods of pre-trained video
prediction models for reward specification. Our method does not require any
ground truth rewards or action annotations, and only requires videos of
desired agent behaviors. VIPER implements rewards as part of the environment,
and can be paired with any RL algorithm. We overview the key components of our
method below.

###  3.1 Video Modeling

Likelihood-based video models are a popular paradigm of generative models
trained to model the data distribution by maximizing an exact or lower bound
on the log-likelihood of the data. These models have demonstrated their
ability to fit highly multi-modal distributions and produce samples depicting
complex dynamics, motion, and behaviors [40, 18, 17, 45].

Our method can integrate any video model that supports computing likelihoods
over the joint distribution factorized in the following form:

$\displaystyle\begin{aligned} \log p(x_{1:T})=\sum_{t=1}^{T}\log p(x_{t}\mid
x_{1:t-1}),\end{aligned}$ (1)

where $x_{1:T}$ is the full video consisting of $T$ frames,
$x_{1},\dots,x_{T}$. When using video models with limited context length $k$,
it is common to approximate likelihoods of an entire video sequence with its
subsequences of length $k$ as follows:

$\displaystyle\begin{aligned} \log p(x_{1:T})\approx\sum_{t=1}^{T}\log
p(x_{t}\mid x_{\max(1,t-k):t-1}).\end{aligned}$ (2)

In this paper, we use an autoregressive transformer model based on VideoGPT
[47, 36] as our video generation model. We first train a VQ-GAN [10] to encode
individual frames $x_{t}$ into discrete codes $z_{t}$. Next, we learn an
autoregressive transformer to model the distribution of codes $z$ through the
following maximum likelihood objective:

$\displaystyle\begin{aligned} \max_{\theta}\sum_{t=1}^{T}\sum_{i=1}^{Z}\log
p_{\theta}(z^{i}_{t}\mid z^{1:i-1}_{t},z_{1:t-1}),\end{aligned}$ (3)

where $z_{t}^{i}$ is the $i$-th code of the $t$-th frame, and $Z$ is the total
number of codes per frame. Computing the exact conditional likelihoods
$p_{\theta}(x_{t}\mid x_{1:t-1})$ is intractable, as it requires marginalizing
over all possible combinations of codes. Instead, we use the conditional
likelihoods over the latent codes $p_{\theta}(z_{t}\mid z_{1:t-1})$ as an
approximation. In our experiments, we show that these likelihoods are
sufficient to capture the underlying dynamics of the videos.

Note that our choice of video model does not preclude the use of other video
generation models, such as MaskGIT-based models [48, 45, 13] or diffusion
models [18, 40]. However, we opted for an autoregressive model due to the
favorable properties of being able to model complex distributions while
retaining fast likelihood computation. Video model comparisons are performed
in Section 4.4.

Figure 3: Aggregate results across 15 DMC tasks, 7 Atari games, and 6 RLBench
tasks. DMC results are provided for DrQ and DreamerV3 (Dv3) RL agents. Atari
and RLBench results are reported for DreamerV3. Atari scores are computed
using Human-Normalized Mean.

###  3.2 Reward Formulation

Given a pretrained video model, VIPER proposes an intuitive reward that
maximizes the conditional log-likelihoods for each transition
$(x_{t},a_{t},x_{t+1})$ observed by the agent:

$\displaystyle\begin{aligned} r^{\mathrm{VIPER}}_{t}\doteq\ln
p_{\theta}(x_{t+1}\;|\;x_{1:t}).\end{aligned}$ (4)

This reward incentivizes the agent to find the most likely trajectory under
the expert video distribution as modeled by the video model. However, the most
probable sequence does not necessarily capture the distribution of behaviors
we want the agent to learn.

For example, when flipping a weighted coin with $p(\mathrm{heads}=0.6)$ 1000
times, _typical sequences_ will count roughly 600 heads and 400 tails, in
contrast to the most probable sequence of 1000 heads that will basically never
be seen in practice [39]. Similarly, the most likely image under a density
model trained on MNIST images is often the image of only background without a
digit, despite this never occurring in the dataset [26]. In the reinforcement
learning setting, an additional issue is that solely optimizing a dense reward
such as $r^{\mathrm{VIPER}}_{t}$ can lead to early convergence to local
optima.

To overcome these challenges, we take the more principled approach of matching
the agent’s trajectory distribution $q(x_{1:T})$ to the sequence distribution
$p_{\theta}(x_{1:T})$ of the video model by minimizing the KL-divergence
between the two distributions [41, 15]:

$\displaystyle\begin{aligned}
\operatorname{KL}\\!\big{[}q(x_{1:T})\;\big{\|}\;p(x_{1:T})\big{]}=&\hskip
1.19995pt\underbracket{\operatorname{E}_{q}\big{[}-\ln
p_{\theta}(x_{1:T})\big{]}}_{\text{cross-entropy}}-\hskip
1.19995pt\underbracket{\operatorname{H}\big{[}q(x_{1:T})\big{]}}_{\text{entropy}}\\\
=&-\sum_{t=1}^{T}\bigg{[}\operatorname{E}_{q}\big{[}\hskip
1.19995pt\underbracket{\ln
p_{\theta}(x_{t+1}\;\big{|}\;x_{1:t})}_{\text{$r_{t}^{\mathrm{VIPER}}$}}\big{]}+\hskip
1.19995pt\underbracket{\operatorname{H}\big{[}q(x_{t+1}\;\big{|}\;x_{1:t})\big{]}}_{\text{exploration
term}}\bigg{]},\end{aligned}$ (5)

The KL objective shows that for the agent to match the video distribution
under the video prediction model, it has to not only maximize the VIPER reward
but also balance this reward while maintaining high entropy over its input
sequences [21]. In the reinforcement learning literature, the entropy bonus
over input trajectories corresponds to an exploration term that encourages the
agent to explore and inhabit a diverse distribution of sequences within the
regions of high probability under the video prediction model. This results in
the final reward function that the agent maximizes:

$\displaystyle\begin{aligned}
r^{\mathrm{KL}}_{t}\,\doteq\,r^{\mathrm{VIPER}}_{t}\,+\,\beta\,r^{\mathrm{expl}}_{t},\end{aligned}$
(6)

where $\beta$ determines the amount of exploration. To efficiently compute the
likelihood reward, we approximate the context frames with a sliding window as
discussed in Equation 2:

$\displaystyle\begin{aligned} r^{\mathrm{VIPER}}_{t}\approx\ln
p_{\theta}(x_{t+1}\;|\;x_{\max(1,t-k):t}).\end{aligned}$ (7)

Figure 1 shows the process of computing rewards using log probabilities under
the video prediction model. VIPER is agnostic to the choice of exploration
reward and in this paper we opt for Plan2Explore [35] and RND [5].

###  3.3 Data Curation

In this work, we explore whether VIPER provides adequate reward signal for
learning low-level control. We utilize the video model likelihoods provided by
an autoregressive video prediction model pre-trained on data from a wide
variety of environments. We curate data by collecting expert video
trajectories from task oracles and motion planning algorithms with access to
state information. Fine-tuning large text-conditioned video models [33, 45,
40] on expert videos would likely lead to improved generalization performance
beyond the curated dataset, and would make for an interesting future research
direction. We explore the favorable generalization capabilities of video
models trained on small datasets, and explore how this leads to more general
reward functions in Section 4.3.

##  4 Experiments

We evaluate VIPER on 28 different tasks across the three domains shown in
Figure 2. We utilize 15 tasks from the DeepMind Control (DMC) suite [44], 7
tasks from the Atari Gym suite [4], and 6 tasks from the Robot Learning
Benchmark (RLBench) [20]. We compare VIPER agents to variants of Adversarial
Motion Priors (AMP) [30], which uses adversarial training to learn behaviors
from reward-free and action-free expert data. All agents are trained using raw
pixel observations from the environment, with no access to state information
or task rewards. In our experiments, we aim to answer the following questions:

  1. 1.

Does VIPER provide an adequate learning signal for solving a variety of tasks?
(Section 4.2)

  2. 2.

Do video models trained on many different tasks still provide useful rewards?
(Section 4.2)

  3. 3.

Can the rewards generalize to novel scenarios where no expert data is
available? (Section 4.3)

  4. 4.

What implementation details matter when using video model likelihoods as
rewards? (Section 4.4)

###  4.1 Video Model Training Details

#### Training data

To collect expert videos in DMC and Atari tasks, Task Oracle RL agents are
trained until convergence with access to ground truth state information and
task rewards. After training, videos of the top $k$ episodes out of 1000
episode rollouts for each task are sampled as expert videos, where $k=50$ and
$100$ for DMC and Atari, respectively. For RLBench, a sampling-based motion
planning algorithm with access to full state information is used to gather 100
demos for each task, where selected tasks range from “easy” to “medium”.

Figure 4: Video model rollouts for 3 different evaluation environments.

#### Video model training

We train a single autoregressive video model for each suite of tasks. Example
images for each suite are shown in Section A.5. For DMC, we train a single VQ-
GAN across all tasks that encodes $64\times 64$ images to $8\times 8$ discrete
codes. Similarly for RLBench and Atari, we train VQ-GANs across all tasks
within each domain, but encode $64\times 64$ images to $16\times 16$ discrete
codes. In practice, the level of VQ compression depends on the visual
complexity of the environment – more texturally simple environments (e.g.,
DMC) allow for higher levels of spatial compression. We follow the original
VQ-GAN architecture [10], consisting of a CNN encoder and decoder. We train
VQ-GAN with a batch size of 128 and learning rate $10^{-4}$ for 200k
iterations.

To train our video models, we encode each frame of the trajectory and stack
the codes temporally to form an encoded 3D video tensor. Similar to VideoGPT
[47], the encoded VQ codes for all video frames are jointly modeled using an
autoregressive transformer in raster scan order. For DMC and Atari, we train
on 16 frames at a time with a one-hot label for task conditioning. For
RLBench, we train both single-task and multi-task video models on 4 frames
with frame skip 4. We add the one-hot task label for the multi-task model used
in our cross-embodiment generalization experiments in Section 4.3. We perform
task conditioning identical to the class conditioning mechanism in VideoGPT,
with learned gains and biases specific to each task ID for each normalization
layer.

Figure 4 shows example video model rollouts for each domain. In general, our
video models are able to accurately capture the dynamics of each environment
to produce probable futures. Further details on model architecture and
hyperparameters can be found in Section A.2. All models are trained on TPUv3-8
instances which are approximately similar to 4 Nvidia V100 GPUs.

###  4.2 Video Model Likelihoods as Rewards for Reinforcement Learning

Figure 5: VIPER incentivizes the agent to maximize trajectory likelihood under
the video model. As such, it provides high rewards for reference (expert)
sequences, and low rewards for unlikely behaviors. Mean rewards
$r^{\mathrm{VIPER}}$ and returns are computed across 40 trajectories.

To learn behaviors from expert videos, we provide reward signals using
Equation 6 in VIPER and the discriminator error in AMP. Both VIPER and AMP are
agnostic to the choice of RL algorithm; but, in this paper, we evaluate our
approach and baselines with two popular RL algorithms: DrQ [22] and DreamerV3
[16]. We use Random Network Distillation [5] (model-free) as the exploration
objective for DrQ, and Plan2Explore [35] (model-based) for DreamerV3.
Hyperparameters for each choice of RL algorithm are shown in Section A.3. We
compare VIPER to two variants of the AMP algorithm: the single-task variant
where only expert videos for the specific task are provided to the agent, and
the multi-task case where expert videos for all tasks are provided to the
agent.

As outlined in Algorithm 1, we compute Equation 6 for every environment step
to label the transition with a reward before adding it to the replay buffer
for policy optimization. In practice, batching Equation 6 across multiple
environments and leveraging parallel likelihood evaluation leads to only a
small decrease in training speed. DMC agents were trained using 1 Nvidia V100
GPU, while Atari and RLBench agents were trained using 1 Nvidia A100 GPU.

We first verify whether VIPER provides meaningful rewards aligned with the
ground truth task rewards. Figure 5 visualizes the ground truth task rewards
and our log-likelihood rewards (Equation 7) for a reference (expert)
trajectory and a random out-of-distribution trajectory. In the reward curves
on the left, VIPER starts to predict high rewards for the expert transitions
once the transitions becomes distinguishable from a random trajectory, while
it consistently outputs low rewards for out-of-distribution transitions. The
return plot on the right clearly shows positive correlation between returns of
the ground truth reward and our proposed reward.

Then, we measure the task rewards when training RL agents with predicted
rewards from VIPER and baselines and report the aggregated results for each
suite and algorithm in Figure 3. The full results can be found in Section A.4.

In DMC, VIPER achieves near expert-level performance from pixels with our
video prediction rewards alone. Although VIPER slightly underperforms Task
Oracle, this is surprising as the Task Oracle uses full state information
along with dense task rewards. VIPER outperforms both variants of AMP. Worth
noting is the drastic performance difference between the single-task and
multi-task AMP algorithms. This performance gap can possibly be attributed to
mode collapse, whereby the discriminator classifies all frames for the current
task as fake samples, leading to an uninformative reward signal for the agent.
Likelihood-based video models, such as VIPER, are less susceptible to mode
collapse than adversarial methods.

In Atari, VIPER approaches the performance of the Task Oracle trained with the
original sparse task reward, and outperforms the AMP baseline. Shown in Figure
6, we found that masking the scoreboard in each Atari environment when
training the video model improved downstream RL performance.

Figure 6:  RL training curves on Atari Pong when using VIPER trained with or
without masking the scoreboard.

Since the video model learns to predict all aspects of the data distribution,
including the scoreboard, it tends to provide noisier reward signal during
policy training when the agent encounters out-of-distributions scores not seen
by the video model. For example, expert demos for Pong only contain videos
where the player scores and the opponent’s score is always zero. During policy
training, we observe that the Pong agent tends to exhibit more erratic
behaviors as soon as the opponent scores at least one point, whereas when
masking the scoreboard, learned policies are generally more stable. These
results suggest the potential benefits of finetuning large video models on
expert demos to learn more generalizable priors, as opposed to training from
scratch.

For RLBench, VIPER outperforms the Task Oracle because RLBench tasks provide
very sparse rewards after long sequences of actions, which pose a challenging
objective for RL agents. VIPER instead provides a dense reward extracted from
the expert videos, which helps learn these challenging tasks. When training
the video model, we found it beneficial to train at a reduced frame rate,
accomplished by subsampling video sequences by a factor of 4. Otherwise, we
observed the video model would assign high likelihoods to stationary
trajectories, resulting in learned policies rarely moving and interacting with
the scene. We hypothesize that this may be partially due to the high control
frequency of the environment, along with the initial slow acceleration of the
robot arm in demonstrations, resulting in very little movement between
adjacent frames. When calculating likelihoods for reward computation, we
similarly input observations strided by time, e.g., $p(x_{t}\mid
x_{t-4},x_{t-8},\dots)$.

###  4.3 Generalization of Video Prediction Rewards

Figure 7: Sampled video predictions for in distribution reference videos
(Train) and an OOD arm/task combination (OOD). The video model displays cross-
embodiment generalization to arm/task combination not observed in the training
data. Video model generalization can enable specifying new tasks where no
reference data is available.  Figure 8: (Left) Training curve for RL agent
trained with VIPER on OOD task. (Right) Task-conditional likelihood for
reference and random trajectory for an OOD task.

Prior works in the space of large, pretrained generative models have shown a
powerful ability to generalize beyond the original training data, demonstrated
by their ability to produce novel generations (e.g., unseen text captions for
image or video generation [33, 45, 40]) as well as learn more generalizable
models from limited finetuning data [8, 46]. Both capabilities provide
promising directions for extending large video generation models to VIPER,
where we can leverage text-video models to capture priors specified by text
commands, or finetune on a small set of demonstrations to learn a task-
specific prior that better generalizes to novel states.

In this section, we seek to understand how this generalization can be used to
learn more general reward functions. We train a model on two datasets of
different robot arms, and evaluate the cross-embodiment generalization
capabilities of the model. Specifically, we gather demonstrations for 23 tasks
on the Rethink Robotics Sawyer Arm, and demonstrations for 30 tasks on the
Franka Panda robotic arm, where only 20 tasks are overlapping between arms. We
then train a task-conditioned autoregressive video model on these
demonstration videos and evaluate the video model by querying unseen arm/task
combinations, where a single initial frame is used for open loop predictions.

Sample video model rollouts for in distribution training tasks and an OOD
arm/task combination are shown in Figure 7. Even though the video model was
not directly trained on demonstrations of the Franka Panda arm to solve the
saucepan task in RLBench, it is able to generate reasonable trajectories for
the arm and task combination. Figure 8 further validates this observation by
assigning higher likelihood to expert videos (Ref) compared to random robot
trajectories (Rand). We observe that these generalization capabilities also
extend to downstream RL, where we use our trained video model with VIPER to
learn a policy for the Franka Robot arm to solve an OOD task without requiring
demos for that specific task and arm combination. These results demonstrate a
promising direction for future work in applying VIPER to larger scale video
models that will be able to better generalize and learn desired behaviors only
through a few demonstrations.

###  4.4 Ablation Studies

Figure 9: Effect of exploration reward term, video model choice, and context
length on downstream RL performance. An equally weighted exploration reward
term and longer video model context leads to improved performance. MaskGIT
substantially underperforms VideoGPT as a choice of video model. The BYOL
model performs moderately due to the deterministic architecture not properly
handling multi-modality.

In this section, we study the contributions of various design decisions when
using video prediction models as rewards: (1) how to weight the exploration
objective, (2) which video model to use, and (3) what context length to use.
Ablation studies are performed across DMC tasks.

#### Exploration objective

As discussed in subsection 3.2, an exploration objective may help the RL agent
learn a distribution of behaviors that closely matches the distribution of the
original data and generalizes, while preventing locally optimal solutions. In
Figure 9, we ablate the $\beta$ parameter introduced in Equation 6 using a
VideoGPT-like model as the VIPER backbone. $\beta=0$ corresponds to no
exploration objective, whereas $\beta=1$ signifies equal weighting between
$r^{\mathrm{VIPER}}$ and $r^{\mathrm{expl}}$. We ablate the exploration
objective using the Plan2Explore [35] reward, which provides the agent with a
reward proportional to the disagreement between an ensemble of one-step
dynamics models. Using no exploration objective causes the policy’s behavior
to collapse, while increasing the weight of the exploration objective leads to
improved performance.

#### Video model

Although our experiments focus on using an autoregressive VideoGPT-like model
to compute likelihoods, VIPER generally allows for any video model that
supports computing conditional likelihoods or implicit densities.

Figure 9 shows additional ablations replacing our VideoGPT model with a
similarly sized MaskGIT [6] model, where frames are modeled with MaskGIT over
space, and autoregressive over time. MaskGIT performs substantially worse than
the VideoGPT model, which is possibly due to noisier likelihoods from parallel
likelihood computation. In addition, while VideoGPT only requires 1 forward
pass to compute likelihood, MaskGIT requires as many as 8 forward passes on
the sequence of tokens for a frame, resulting in an approximately $8\times$
slowdown in reward computation, and $2.5\times$ in overall RL training.

Finally, we evaluate the performance of a video model that computes implicit
densities using a negative distance metric between online predictions and
target encodings with a recurrent Bootstrap Your Own Latent (BYOL)
architecture [12, 11]. We refer the reader to Section A.1 for more details
about the implementation. While BYOL outperforms MaskGIT in Figure 9, its
deterministic recurrent architecture is unable to predict accurate embeddings
more than a few steps into the future. This limits the ability of the learned
reward function to capture temporal dynamics. We observe that agents trained
with BYOL often end up in a stationary pose, which achieves high reward from
the BYOL model.

#### Context length

The context length $k$ of the video model is an important choice in
determining how much history to incorporate into the conditional likelihood.
At the limit where $k=0$, the reward is the unconditional likelihood over each
frame. Such a choice would lead to stationary behavior on many tasks. Figure 9
shows that increasing the context length can help improve performance when
leveraging a VideoGPT-based model for downstream RL tasks, subject to
diminishing returns. We hypothesize that a longer context length may help for
long-horizon tasks.

##  5 Conclusion

This paper presents Video Prediction Rewards (VIPER), a general algorithm that
enables agents to learn behaviors from videos of experts. To achieve this, we
leverage a simple pre-trained autoregressive video model to provide rewards
for a reinforcement learning agent. These rewards are parameterized as the
conditional probabilities of a particular frame given a context past of
frames. In addition, we include an entropy maximization objective to ensure
that the agent learns diverse behaviors which match the video model’s
trajectory distribution. VIPER succeeds across 3 benchmarks and 28 tasks,
including the DeepMind Control Suite, Atari, and the Reinforcement Learning
Benchmark. We find that simple data augmentation techniques can dramatically
improve the effectiveness of VIPER. We also show that VIPER generalizes to
out-of-distribution tasks for which no demonstrations were provided. Moreover,
VIPER can learn reward functions for a wide variety of tasks using a single
video model, outperforming AMP [30], which suffers from mode collapse as the
diversity of the expert videos increases.

Limitations of our work include that VIPER is trained using in-domain data of
expert agents, which is not a readily-available data source in the real world.
Additionally, video prediction rewards trained on stochastic data may lead the
agent to prefer states that minimize the video model’s uncertainty, which may
lead to sub-optimal behaviors. This challenge may be present for cases where
either the environment is stochastic or the demonstrator provides noisy
demonstrations. Moreover, selecting the right trade-off between VQCode size
and context length has a significant impact on the quality of the learned
rewards, with small VQCodes failing to model important components of the
environment (e.g., the ball in Atari Breakout) and short context lengths
leading to myopic behaviors which results in poor performance. Exploring
alternative video model architectures for VIPER would likely be a fruitful
research direction.

To improve the generalization capabilities of VIPER, larger pre-trained video
models are necessary. Future work will explore how fine-tuning or human
preferences can be used to improve video prediction rewards. Using text-to-
video models remain another interesting direction for future work in
extracting text-conditioned task-specific priors from a pretrained video
model. In addition, extending this line of work to leverage Video Diffusion
Models [18] as video prediction rewards may lead to interesting outcomes.

## Acknowledgments and Disclosure of Funding

This work was supported in part by an NSF Fellowship, NSF NRI #2024675, ONR
MURI N00014-22-1-2773, Komatsu, and the Vanier Canada Graduate Scholarship. We
also thank Google TPU Research Cloud and Cirrascale (https://cirrascale.com)
for providing compute resources.

## References

  * Abbeel and Ng [2004] Pieter Abbeel and Andrew Y Ng.  Apprenticeship learning via inverse reinforcement learning.  In _International Conference on Machine Learning_ , 2004.
  * Baker et al. [2022] Bowen Baker, Ilge Akkaya, Peter Zhokhov, Joost Huizinga, Jie Tang, Adrien Ecoffet, Brandon Houghton, Raul Sampedro, and Jeff Clune.  Video pretraining (vpt): Learning to act by watching unlabeled online videos.  In _Advances in Neural Information Processing Systems_ , 2022.
  * Bellemare et al. [2013] Marc G Bellemare, Yavar Naddaf, Joel Veness, and Michael Bowling.  The arcade learning environment: An evaluation platform for general agents.  _Journal of Artificial Intelligence Research_ , 47:253–279, 2013.
  * Brockman et al. [2016] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba.  Openai gym.  _arXiv preprint arXiv:1606.01540_ , 2016.
  * Burda et al. [2019] Yuri Burda, Harrison Edwards, Amos Storkey, and Oleg Klimov.  Exploration by random network distillation.  In _International Conference on Learning Representations_ , 2019.
  * Chang et al. [2022] Huiwen Chang, Han Zhang, Lu Jiang, Ce Liu, and William T Freeman.  Maskgit: Masked generative image transformer.  In _Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition_ , pages 11315–11325, 2022.
  * Chintala et al. [2020] Soumith Chintala, Emily Denton, Martin Arjovsky, and Michael Mathieu.  How to train a gan? tips and tricks to make gans work.  https://github.com/soumith/ganhacks, 2020.
  * Dai et al. [2023] Yilun Dai, Mengjiao Yang, Bo Dai, Hanjun Dai, Ofir Nachum, Josh Tenenbaum, Dale Schuurmans, and Pieter Abbeel.  Learning universal policies via text-guided video generation.  _arXiv preprint arXiv:2302.00111_ , 2023.
  * Doerr et al. [2018] Andreas Doerr, Christian Daniel, Martin Schiegg, Nguyen-Tuong Duy, Stefan Schaal, Marc Toussaint, and Trimpe Sebastian.  Probabilistic recurrent state-space models.  In _International conference on machine learning_ , pages 1280–1289. PMLR, 2018.
  * Esser et al. [2021] Patrick Esser, Robin Rombach, and Bjorn Ommer.  Taming transformers for high-resolution image synthesis.  In _Proceedings of the IEEE/CVF conference on computer vision and pattern recognition_ , pages 12873–12883, 2021.
  * Grill et al. [2020] Jean-Bastien Grill, Florian Strub, Florent Altché, Corentin Tallec, Pierre Richemond, Elena Buchatskaya, Carl Doersch, Bernardo Avila Pires, Zhaohan Guo, Mohammad Gheshlaghi Azar, et al.  Bootstrap your own latent-a new approach to self-supervised learning.  In _Advances in neural information processing systems_ , volume 33, pages 21271–21284, 2020.
  * Guo et al. [2022] Zhaohan Guo, Shantanu Thakoor, Miruna Pîslar, Bernardo Avila Pires, Florent Altché, Corentin Tallec, Alaa Saade, Daniele Calandriello, Jean-Bastien Grill, Yunhao Tang, et al.  Byol-explore: Exploration by bootstrapped prediction.  In _Advances in neural information processing systems_ , volume 35, pages 31855–31870, 2022.
  * Gupta et al. [2022] Agrim Gupta, Stephen Tian, Yunzhi Zhang, Jiajun Wu, Roberto Martín-Martín, and Li Fei-Fei.  Maskvit: Masked visual pre-training for video prediction.  _arXiv preprint arXiv:2206.11894_ , 2022.
  * Hafner et al. [2019] Danijar Hafner, Timothy Lillicrap, Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson.  Learning latent dynamics for planning from pixels.  In _International conference on machine learning_ , pages 2555–2565. PMLR, 2019.
  * Hafner et al. [2020] Danijar Hafner, Pedro A Ortega, Jimmy Ba, Thomas Parr, Karl Friston, and Nicolas Heess.  Action and perception as divergence minimization.  _arXiv preprint arXiv:2009.01791_ , 2020.
  * Hafner et al. [2023] Danijar Hafner, Jurgis Pasukonis, Jimmy Ba, and Timothy Lillicrap.  Mastering diverse domains through world models.  _arXiv preprint arXiv:2301.04104_ , 2023.
  * Ho et al. [2022a] Jonathan Ho, William Chan, Chitwan Saharia, Jay Whang, Ruiqi Gao, Alexey Gritsenko, Diederik P Kingma, Ben Poole, Mohammad Norouzi, David J Fleet, et al.  Imagen video: High definition video generation with diffusion models.  _arXiv preprint arXiv:2210.02303_ , 2022a.
  * Ho et al. [2022b] Jonathan Ho, Tim Salimans, Alexey Gritsenko, William Chan, Mohammad Norouzi, and David J Fleet.  Video diffusion models.  In _Advances in Neural Information Processing Systems_ , 2022b.
  * Jaegle et al. [2021] Andrew Jaegle, Yury Sulsky, Arun Ahuja, Jake Bruce, Rob Fergus, and Greg Wayne.  Imitation by predicting observations.  In _International Conference on Machine Learning_ , pages 4665–4676. PMLR, 2021.
  * James et al. [2020] Stephen James, Zicong Ma, David Rovick Arrojo, and Andrew J. Davison.  Rlbench: The robot learning benchmark & learning environment.  _IEEE Robotics and Automation Letters_ , 2020.
  * Jaynes [1957] Edwin T Jaynes.  Information theory and statistical mechanics.  _Physical review_ , 106(4):620, 1957.
  * Kostrikov et al. [2021] Ilya Kostrikov, Denis Yarats, and Rob Fergus.  Image augmentation is all you need: Regularizing deep reinforcement learning from pixels.  In _International Conference on Learning Representations_ , 2021.
  * Lee et al. [2021] Youngwoon Lee, Andrew Szot, Shao-Hua Sun, and Joseph J. Lim.  Generalizable imitation learning from observation via inferring goal proximity.  In _Neural Information Processing Systems_ , 2021.
  * 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.
  * Liu et al. [2018] YuXuan Liu, Abhishek Gupta, Pieter Abbeel, and Sergey Levine.  Imitation from observation: Learning to imitate behaviors from raw video via context translation.  In _IEEE International Conference on Robotics and Automation_ , 2018\.
  * Nalisnick et al. [2019] Eric Nalisnick, Akihiro Matsukawa, Yee Whye Teh, Dilan Gorur, and Balaji Lakshminarayanan.  Do deep generative models know what they don’t know?  In _International Conference on Learning Representations_ , 2019.
  * Ng and Russell [2000] Andrew Y. Ng and Stuart J. Russell.  Algorithms for inverse reinforcement learning.  In _International Conference on Machine Learning_ , 2000.
  * Niekum et al. [2015] Scott Niekum, Sarah Osentoski, George Konidaris, Sachin Chitta, Bhaskara Marthi, and Andrew G Barto.  Learning grounded finite-state representations from unstructured demonstrations.  _The International Journal of Robotics Research_ , 34(2):131–157, 2015.
  * Pathak et al. [2018] Deepak Pathak, Parsa Mahmoudieh, Michael Luo, Pulkit Agrawal, Dian Chen, Fred Shentu, Evan Shelhamer, Jitendra Malik, Alexei A. Efros, and Trevor Darrell.  Zero-shot visual imitation.  In _International Conference on Learning Representations_ , 2018.
  * Peng et al. [2021] Xue Bin Peng, Ze Ma, Pieter Abbeel, Sergey Levine, and Angjoo Kanazawa.  AMP: adversarial motion priors for stylized physics-based character control.  _ACM Transactions on Graphics (TOG)_ , 40(4):1–20, 2021.
  * Pomerleau [1989] Dean A Pomerleau.  Alvinn: An autonomous land vehicle in a neural network.  In _Advances in Neural Information Processing Systems_ , pages 305–313, 1989.
  * Riedmiller et al. [2018] Martin Riedmiller, Roland Hafner, Thomas Lampe, Michael Neunert, Jonas Degrave, Tom Wiele, Vlad Mnih, Nicolas Heess, and Jost Tobias Springenberg.  Learning by playing solving sparse reward tasks from scratch.  In _International conference on machine learning_ , pages 4344–4353. PMLR, 2018.
  * Saharia et al. [2022] Chitwan Saharia, William Chan, Saurabh Saxena, Lala Li, Jay Whang, Emily Denton, Seyed Kamyar Seyed Ghasemipour, Burcu Karagol Ayan, S. Sara Mahdavi, Rapha Gontijo Lopes, Tim Salimans, Jonathan Ho, David J Fleet, and Mohammad Norouzi.  Photorealistic text-to-image diffusion models with deep language understanding.  In _Advances in Neural Information Processing Systems_ , 2022.
  * Schmeckpeper et al. [2021] Karl Schmeckpeper, Oleh Rybkin, Kostas Daniilidis, Sergey Levine, and Chelsea Finn.  Reinforcement learning with videos: Combining offline observations with interaction.  In _Conference on Robot Learning_ , 2021.
  * Sekar et al. [2020] Ramanan Sekar, Oleh Rybkin, Kostas Daniilidis, Pieter Abbeel, Danijar Hafner, and Deepak Pathak.  Planning to explore via self-supervised world models.  In _International Conference on Machine Learning_ , 2020.
  * Seo et al. [2022] Younggyo Seo, Kimin Lee, Fangchen Liu, Stephen James, and Pieter Abbeel.  Harp: Autoregressive latent video prediction with high-fidelity image generator.  In _2022 IEEE International Conference on Image Processing (ICIP)_ , pages 3943–3947. IEEE, 2022.
  * Sermanet et al. [2017] Pierre Sermanet, Kelvin Xu, and Sergey Levine.  Unsupervised perceptual rewards for imitation learning.  _Robotics: Science and Systems_ , 2017.
  * Sermanet et al. [2018] Pierre Sermanet, Corey Lynch, Yevgen Chebotar, Jasmine Hsu, Eric Jang, Stefan Schaal, and Sergey Levine.  Time-contrastive networks: Self-supervised learning from video.  In _IEEE International Conference on Robotics and Automation_ , pages 1134–1141, 2018.
  * Shannon [2001] Claude Elwood Shannon.  A mathematical theory of communication.  _ACM SIGMOBILE mobile computing and communications review_ , 5(1):3–55, 2001.
  * Singer et al. [2023] Uriel Singer, Adam Polyak, Thomas Hayes, Xi Yin, Jie An, Songyang Zhang, Qiyuan Hu, Harry Yang, Oron Ashual, Oran Gafni, Devi Parikh, Sonal Gupta, and Yaniv Taigman.  Make-a-video: Text-to-video generation without text-video data.  In _International Conference on Learning Representations_ , 2023.
  * Todorov [2009] Emanuel Todorov.  Efficient computation of optimal actions.  _Proceedings of the national academy of sciences_ , 106(28):11478–11483, 2009.
  * Torabi et al. [2018a] Faraz Torabi, Garrett Warnell, and Peter Stone.  Behavioral cloning from observation.  In _International Joint Conferences on Artificial Intelligence_ , 2018a.
  * Torabi et al. [2018b] Faraz Torabi, Garrett Warnell, and Peter Stone.  Generative adversarial imitation from observation.  _arXiv preprint arXiv:1807.06158_ , 2018b.
  * Tunyasuvunakool et al. [2020] Saran Tunyasuvunakool, Alistair Muldal, Yotam Doron, Siqi Liu, Steven Bohez, Josh Merel, Tom Erez, Timothy Lillicrap, Nicolas Heess, and Yuval Tassa.  dm_control: Software and tasks for continuous control.  _Software Impacts_ , 6:100022, 2020.  ISSN 2665-9638.
  * Villegas et al. [2022] Ruben Villegas, Mohammad Babaeizadeh, Pieter-Jan Kindermans, Hernan Moraldo, Han Zhang, Mohammad Taghi Saffar, Santiago Castro, Julius Kunze, and Dumitru Erhan.  Phenaki: Variable length video generation from open domain textual description.  _arXiv preprint arXiv:2210.02399_ , 2022.
  * Wei et al. [2022] Jason Wei, Maarten Bosma, Vincent Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M. Dai, and Quoc V Le.  Finetuned language models are zero-shot learners.  In _International Conference on Learning Representations_ , 2022.
  * Yan et al. [2021] Wilson Yan, Yunzhi Zhang, Pieter Abbeel, and Aravind Srinivas.  Videogpt: Video generation using vq-vae and transformers.  _arXiv preprint arXiv:2104.10157_ , 2021.
  * Yu et al. [2022] Lijun Yu, Yong Cheng, Kihyuk Sohn, José Lezama, Han Zhang, Huiwen Chang, Alexander G Hauptmann, Ming-Hsuan Yang, Yuan Hao, Irfan Essa, et al.  Magvit: Masked generative video transformer.  _arXiv preprint arXiv:2212.05199_ , 2022.
  * Ziebart et al. [2008] Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, and Anind K Dey.  Maximum entropy inverse reinforcement learning.  In _Association for the Advancement of Artificial Intelligence_ , 2008\.
  * Zolna et al. [2020] Konrad Zolna, Scott Reed, Alexander Novikov, Sergio Gomez Colmenarej, David Budden, Serkan Cabi, Misha Denil, Nando de Freitas, and Ziyu Wang.  Task-relevant adversarial imitation learning.  In _Conference on Robot Learning_ , 2020.

##  Appendix A Appendix

###  A.1 Bootstrap Your Own Latent Details

Model Architecture. The BYOL model used in this work derives from the
recurrent BYOL-Explore architecture proposed in [12], with a few
modifications. Namely Guo et al. propose BYOL-Explore, a multi-step predictive
latent world model. The BYOL-Explore architecture trains an online network
using targets generated by an exponential moving average target network on
future observations. Observations $o_{t}$ are first encoded into an
observation representation $f_{\theta}(o_{t})$. Online predictions are
computed by then encoding the observation representation using a closed-loop
RNN $h_{\theta}^{c}(f_{\theta}(o_{t}))$. Guo et al. also pass actions into the
closed loop RNN, but we wish to learn an action-free reward mechanism that can
be trained solely from videos. At each timestep, the carry $b_{t}$ from the
closed loop RNN is then fed into an open-loop RNN $h_{\theta}^{o}$ which
predicts open-loop representations $K$ steps into the future:
$(b_{t,k}\in\mathcal{R}^{\mathcal{M}})_{k=1}^{K-1}$, where
$b_{t,k}=h_{\theta}^{o}(b_{t,k-1})$. The original BYOL-Explore architecture
feeds actions $a_{t+k+1}$ as inputs to the open loop cell $h_{\theta}^{o}$ as
well, but we opt to omit these inputs for the same reason outlined above.

Given the deterministic RNN architecture used to predict online targets,
omitting the action conditioning may lead to poor performance in multi-modal
or stochastic environments. A more principled approach would utilize a
probabilistic recurrent state space model [9, 14] to account for the multi-
modality of futue states. We leave this approach to future work.

Finally, a predictor head $g_{\theta}(b_{t,k})$ outputs online predictions. We
refer the reader to [12] for a figure of the outlined architecture.

The target network is an observation encoder $f_{\phi}$ whose parameters are
an exponential moving average of $f_{\theta}$. The loss function for the
online network is then defined as:

$\displaystyle\mathcal{L}_{\text{BYOL-Explore}}(\theta,t,k)$
$\displaystyle=\left\|\frac{g_{\theta}(b_{t,k})}{\|g_{\theta}(b_{t,k})\|_{2}}-\texttt{sg}\left(\frac{f_{\phi}(o_{t+k})}{\|f_{\phi}(o_{t+k})\|_{2}}\right)\right\|_{2}^{2},$
$\displaystyle\mathcal{L}_{\text{BYOL-Explore}}(\theta)$
$\displaystyle=\frac{1}{B(T-1)}\sum_{t=0}^{T-2}\frac{1}{K(t)}\sum_{k=1}^{K(t)}\mathcal{L}_{\text{BYOL-
Explore}}(\theta,t,k),$

where $K(t)=\min(K,T-1-t)$ is the valid open-loop horizon for a trajectory of
length $T$ and sg is the stop-gradient operator.

Computing Rewards. The uncertainty associated with the transition
$(o_{t},a_{t},o_{t+1})$ is the sum of the corresponding prediction losses:

$\ell_{t}=\sum_{p+q=t+1}\mathcal{L}_{\text{BYOL-Explore}}(\theta,j,p,q),$

where $0\leq p\leq T-2$, $1\leq q\leq K$ and $0\leq t\leq T-2$. This
accumulates all the losses corresponding to the latent dynamics model
uncertainties relative to the observation $o_{t+1}$. For our purposes, we
calculate rewards as $r_{t}=-\ell_{t}$, which is equivalent to negating the
reward used in the original BYOL-Explore formulation, and can be interpreted
as negating the exploration objective to ensure that the agent stays close to
the trajectory distribution found in the original dataset.

###  A.2 Video Model Hyperparameters

Table 1: Hyperparameters and training details for all VQ-GAN models | DMC | Atari |  | RLBench
---
(single + multi task)
Input size | $64\times 64\times 3$ | $64\times 64\times 3$ | $64\times 64\times 3$
Latent size | $8\times 8$ | $16\times 16$ | $16\times 16$
$\beta$ (commitment loss coefficient) | 0.25 | 0.25 | 0.25
Batch size | 128 | 128 | 128
Learning rate | $10^{-4}$ | $10^{-4}$ | $10^{-4}$
Learning rate schedule | constant | constant | constant
Training steps | 200k | 200k | 200k
Base channels | 128 | 128 | 128
Ch mult | [1, 1, 2, 2] | [1, 2, 2] | [1, 2, 2]
Num res blocks | 1 | 1 | 2
Codebook size | 1024 | 1024 | 1024
Codebook dimension | 64 | 64 | 64
Perceptual loss weight | 0.1 | 0.1 | 0.1
Disc base features | 32 | 32 | 32
Disc gradient penalty weight | $10^{8}$ | $10^{8}$ | $10^{8}$
Disc max hidden feature size | 512 | 512 | 512
Disc mbstd group size | 4 | 4 | 4
Disc mbstd num features | 1 | 1 | 1
Disc loss weight | 0.1 | 0.1 | 0.1
Table 2: Hyperparameters and training details for all VideoGPT and MaskGit models | DMC | Atari |  | RLBench
---
(single task)
| RLBench
---
(multi-task)
Sequence length | 16 | 16 | 4 | 4
Frame skip | 1 | 1 | 4 | 4
Latent size | $8\times 8$ | $16\times 16$ | $16\times 16$ | $16\times 16$
Number of classes | 16 | 8 | n/a | 34
Batch size | 32 | 32 | 32 | 32
Learning rate | $10^{-4}$ | $10^{-4}$ | $10^{-4}$ | $10^{-4}$
Learning rate schedule | constant | constant | constant | constant
Training steps | 500k | 500k | 500k | 500k
Adam ($\beta_{1},\beta_{2}$) | (0.9, 0.999) | (0.9, 0.999) | (0.9, 0.999) | (0.9, 0.999)
Hidden dim | 256 | 512 | 256 | 512
Feedforward dim | 1024 | 2048 | 1024 | 2048
Number of heads | 8 | 8 | 8 | 8
Number of layers | 8 | 8 | 8 | 8
Dropout | 0 | 0 | 0 | 0

###  A.3 Reinforcement Learning Hyperparameters

Table 3: Hyperparameters and training details for DreamerV3 | DMC | Atari |  | RLBench
---
(single task)
| RLBench
---
(multi-task)
General |  |  |  |
Replay Capacity (FIFO) | $10^{6}$ | $10^{6}$ | $10^{6}$ | $10^{6}$
Start learning (prefill) | 0 | 0 | 0 | 0
Batch Size | 16 | 16 | 16 | 16
Batch length | 64 | 64 | 64 | 64
MLP Size | $2\times 512$ | $4\times 1024$ | $4\times 1024$ | $4\times 1024$
Activation | LayerNorm + swish |  |  |
World Model |  |  |  |
RSSM Size | 512 | 4096 | 4096 | 4096
Number of Latents | 32 | 32 | 32 | 32
Classes per Latent | 32 | 32 | 32 | 32
KL Balancing | 0.667 | 0.667 | 0.667 | 0.667
Actor Critic |  |  |  |
Imagination Horizon | 15 | 15 | 15 | 15
Discount | 0.995 | 0.995 | 0.995 | 0.995
Return Lambda | 0.95 | 0.95 | 0.95 | 0.95
Target Update Interval | 50 | 50 | 50 | 50
All Optimizers |  |  |  |
Gradient Clipping | 100 | 100 | 100 | 100
Learning Rate | $10^{-4}$ | $10^{-4}$ | $10^{-4}$ | $10^{-4}$
Adam epsilon | $10^{-6}$ | $10^{-6}$ | $10^{-6}$ | $10^{-6}$
Plan2Explore |  |  |  |
Ensemble size | 10 | 10 | 10 | 10
Table 4: Hyperparameters and training details for DrQ | DMC
---|---
Actor Critic |
MLP Size | $256\times 256$
CNN Features | $32\times 64\times 128\times 256$
CNN Filters | $3\times 3\times 3\times 3$
CNN Strides | $2\times 2\times 2\times 2$
Learning Rate | $3e-4$
Discount | $0.99$
Number of Critics | 2
$\tau$ | 0.005
Init Temperature | 0.1
Augmentations | Random crop
Random Network Distillation |
RND CNN Features | $32\times 64\times 64$
RND MLP Size | $512\times 512$
RND learning rate | $3e-4$
RND Exploration Weight | $1$

###  A.4 Environment Training Curves

Figure 10: Results across 6 Reinforcement Learning Benchmark tasks, with mean
and standard deviation computed across 3 seeds. Figure 11: Results across 7
Atari tasks. Scores computed using Human-Normalized Mean across 3 seeds.
Figure 12: Results across 15 DeepMind Control Suite tasks, with mean and
standard deviation computed across 3 seeds. AMP runs were stopped early due to
poor performance.

###  A.5 Training Environments

Figure 13: A single autoregressive video model is trained on 30 tasks for the
Franka Panda and 23 tasks for the ReThink Robotics Sawyer, using $16\times 16$
VQCodes and a context length of 4, with frame skip 4. Figure 14: A single
task-conditioned autoregressive video model is trained on 15 DeepMind Control
Tasks, using $8\times 8$ VQCodes and a context length of 16. Figure 15: A
single autoregressive video model is trained on 7 Atari tasks, using $16\times
16$ VQCodes and a context length of 16.

###  A.6 Visualizing Video Prediction Uncertainty

Video model uncertainty can be visualized by upsampling the conditional
likelihoods over VQCodes. Since VQCodes tend to model local features, this
provides a useful tool for analyzing which regions of the image the video
model is uncertain about. We visualize video prediction uncertainty for three
environments below:

((a)) Reference Trajectory

((b)) Random Trajectory

Figure 16: Uncertainty visualized for a reference trajectory (top) and a
random trajectory (bottom). Brighter values correspond to higher likelihoods.
For the random trajectory, the video model assigns lower likelihoods to
regions of the image containing the ball. This is especially true for random
trajectories, where the video model, trained solely on expert trajectories,
cannot accurately predict what happens when the agent plays sub-optimally.
Figure 17: Uncertainty visualized for a random trajectory from the dmc
cartpole balance task. Brighter values correspond to higher likelihoods.

((a)) Reference Trajectory

((b)) Random Trajectory

Figure 18: Uncertainty visualized for a reference trajectory (top) and a
random trajectory (bottom). Brighter values correspond to higher likelihoods.
Notice the high uncertainty over the position of objects on the table for the
first frame. This corresponds to the unconditional likelihood (with no
context). Since the position of the saucepan is randomized for every episode,
the video model assigns some uncertainty to this initial configuration. For
the random trajectory, there is high uncertainty assigned to the position of
the arm since the video model has not seen trajectories that display poor
behavior.