# Vid2Robot: End-to-end Video-conditioned Policy Learning with Cross-Attention
Transformers

Vidhi Jain${}^{\text{1,2}}$ Maria Attarian${}^{\text{1,3}}$ Nikhil J
Joshi${}^{\text{1}}$ Ayzaan Wahid${}^{\text{1}}$ Danny Driess${}^{\text{1}}$
Quan Vuong${}^{\text{1}}$ Pannag R Sanketi${}^{\text{1}}$
Pierre Sermanet${}^{\text{1}}$ Stefan Welker${}^{\text{1}}$ Christine
Chan${}^{\text{1}}$ Igor Gilitschenski${}^{\text{3}}$ Yonatan
Bisk${}^{\text{2}}$ Debidatta Dwibedi${}^{\text{1}}$

${}^{\text{1}}$Google DeepMind Robotics ${}^{\text{2}}$Carnegie Mellon
University ${}^{\text{3}}$University of Toronto

###### Abstract

While large-scale robotic systems typically rely on textual instructions for
tasks, this work explores a different approach: can robots infer the task
directly from observing humans? This shift necessitates the robot’s ability to
decode human intent and translate it into executable actions within its own
physical constraints and environment.

We introduce Vid2Robot, a novel end-to-end video-based learning framework for
robots. Given a video demonstration of a manipulation task and current visual
observations, Vid2Robot directly produces robot actions. This is achieved
through a unified representation model trained on a large dataset of human
video and robot trajectory. The model leverages cross-attention mechanisms to
fuse prompt video features to the robot’s current state and generate
appropriate actions that mimic the observed task. To further improve policy
performance, we propose auxiliary contrastive losses that enhance the
alignment between human and robot video representations.

We evaluate Vid2Robot on real-world robots, demonstrating a 20% improvement in
performance compared to other video-conditioned policies when using human
demonstration videos. Additionally, our model exhibits emergent capabilities,
such as successfully transferring observed motions from one object to another,
and long-horizon composition, thus showcasing its potential for real-world
applications. Project website: vid2robot.github.io.

Figure 1: Overview. Vid2Robot is a video-conditioned robot policy. Given a
human demonstration (top), Vid2Robot recognizes the task semantics and
performs the same task based on the robot’s current visual observation (bottom
left). A successful trajectory is presented on the bottom right.

##  I Introduction

The path to creating versatile robots that provide assistance in people’s
daily routines requires them to learn new skills on-the-go. These range from
small preferences like which brand of dishwasher a specific household uses to
entirely different ways to clean the house. For known skills, humans simply
communicate in natural language, but when nuance is required or a skill is
novel, we revert to demonstrations. For example, we might show how a
particular microwave works or how we prefer our cabinets to be organized. To
enable seamless robot deployment, robots need the same ability for
generalization from demonstration for learning new policies that comes so
naturally to humans.

Humans can infer the intentions of other humans based on third-person visual
observations. Oftentimes, we use social reasoning and common sense to
understand others’ goals implicitly. This ability is leveraged both as
children and adults (e.g. via How-To videos [30]) for learning anything where
the mechanical nuance of the task is hard to capture in still images or text
[5] (e.g. how to knead dough or knit). If robots can also be taught to
understand the intentions of other agents, it might allow them to better
interact with humans and perform tasks more efficiently.

This work focuses on visual imitation learning, where robots learn to perform
tasks by watching video demonstrations. This setup offers several advantages.
First, it allows robots to learn from agents with different embodiment,
enabling new skill acquisition without tele-operation. Second, it enables
robots to learn from experts, even if they are not situated with the robot.
Finally, visual imitation learning is ideal for teaching tasks that are
difficult or impossible to describe in words.

Existing multi-task robot manipulation models (e.g. RT-1 [7], RT-2 [8], and
RT-X [33]) use language conditioning to output a robot trajectory. This
reliance of text alone for task specification makes it difficult for robots to
handle polysemy and tasks whose executions vary dramatically based on context.
For example, ‘open drawer’, ‘open cabinet’, ‘open container with lid’ and
‘open jar with screw cap’ might share the same verb but require very different
motor control for each interaction. Here the agent should not generalize its
policy, whereas it should generalize from one drawer to others that vary in
type, color and shape. For this reason, there are a broad range of tasks for
which it is hard to design primitives for high-level planning approaches [25,
2].

Another common approach has been to use a final goal image in goal-conditioned
behavior cloning tasks [31, 24]. While several task specifications can be
defined in terms of the resulting state of the environment, there are others
for which the manner in which the action is performed is part of the
specification. For example, ‘hold the flag’ and ‘wave the flag’ can have the
same final goal image. This ambiguity can be resolved through the use of
several sub-goal frames, that is, video conditioning.

While language conditioned policies achieve somewhat high success rates, video
conditioned policies have lagged behind in performance, as shown in prior work
[21]. Cases of good performance [40] with video conditioning require the
provided video to be from the same workspace with limited variability. Based
on observations, we identify three main challenges for video conditioned
policies: (1) High dimensional data: Raw videos are high dimensional data that
require more compute and memory to process. This makes video conditioned
multi-task policies difficult to train at scale. (2) Variability in task
specification: There can be significant variance in how people perform the
same task. Demonstrations for a task like ‘unstack the cups’ can have both
visually distinctive and physically diverse cups, in addition to changes in
the background distractors and lighting conditions. This leads to high
variability in task specification for a policy that needs to perform the same
task in a new setting. (3) Limited Availability of Training Data: While there
is an abundance of unlabeled video data on the internet, obtaining labeled
video datasets for specific tasks that our robots are capable of doing is
challenging.

Despite these challenges, as noted, video conditioned policy learning is a
core challenge robots need to master. Therefore, to reduce the reliance on
detailed and potentially ambiguous language prompts, we aim to enable physical
visual demonstrations as a another way for task specification. To this end, we
study how end-to-end models with video-conditioning can used to specify tasks
to robot.

We aim to develop an end-to-end system that enables rapid adaptation to tasks
specified in the form of video demonstration. Unlike prior work that either
learned representations from videos for only object and verb recognition [21]
or learned motor control in simulation [44], our work demonstrates the
applicability of end-to-end learned video representations for real-world
multi-task robotic control. We present the key contributions of our work as
follows: (1) We present a transformer-based policy to encode video task
specification, demonstrated by either robot or human agent embodiments (§II).
(2) We encourage alignment between the prompt and robot video representations
using three contrastive auxiliary losses during training (§II-E) (3) Through
real robot experiments, we find our video conditioned policy is better than
baselines on human prompt videos. Furthermore, our policy is better at cross-
object motion transfer (§III).

Figure 2: Dataset creation. (top row) Here we show a Robot-Robot video pair
for placing the rxbar into top drawer. We similarly pair existing robot-robot
videos performing the same task. (middle row) Here we show Hindsight Human-
Robot paired videos for picking a coke can from the bottom drawer and placing
it on the counter task. We use the task instructions from robot trajectories
and ask human participants to perform the task and record a demonstration
video from robot’s perspective/view. (bottom row) Here we show a Co-located
Human-Robot pair of videos for placing the pipe wrench in the toolkit. We
record both a human demonstration and by a robot teleoperation in a same
workspace. Different workspaces can be used to perform the same task
instruction, thus, eventually resulting in pairs with visually diverse prompts
and robot state observations. More details in §II-B.

##  II Approach

###  II-A Preliminaries

Our objective is to design a robotic system that takes in a prompt video of a
manipulation task and outputs actions that accomplish the task demonstrated in
the video. This system needs to infer the underlying task from the prompt
video (which might have a different setup or embodiment than the robot) and
then manipulate the objects in its own environment to achieve the inferred
task. Specifically, we are given a prompt video $V$ and the robot state
$S_{t}=\\{x_{i}\\}_{i=t-k-1}^{t}$ where $x_{i}$ is the frame from the robot’s
camera stream at time $i$, $k$ is the maximum number of historical frames, and
$t$ is the current time-step. We train a policy $\pi(a_{t}|S_{t},V)$ that
infers the underlying task from $V$ and predicts task relevant action $a_{t}$.
To train this model, we need a dataset of paired prompt videos and robot
trajectories. We will discuss in detail how to create paired datasets below.

###  II-B Datasets

To train a video-conditioned robot policy we need a dataset of pairs: prompt
videos and robot trajectories performing the same task. In this work, we
explore prompt videos where the task is performed by both humans and robots.
To create this dataset, we rely on three classes of data:

  1. 1.

Robot-Robot: We pair existing robot-robot videos of the same task. For this
pairing we consider two videos to match if they are performing the same task
in different settings. We define ‘task‘ based on natural language instructions
used when recording robot trajectories. These instructions typically consist
of one or two verbs surrounded by nouns, such as ‘place water bottle upright’,
‘move the coke can to the green chip bag’ or ‘open top drawer’. The objective
of this pairing is two-fold: first, to be able to take advantage of an already
labeled and collected dataset of robot trajectories and second to ensure
robots are able to imitate when the same task is demonstrated in a different
environment.

  2. 2.

Hindsight Human-Robot: Here we use the task instructions from the robot
trajectories dataset and ask one to five human participants to perform the
task and record a demonstration video from the robot’s perspective/view. The
set of instructions are the same as before, but there is a significant
embodiment and speed variability due to different humans performing the task
with left or right hands and at a randomized robot camera angle. This requires
some manual effort but provides us with a lot of paired data for training the
policy for the available set of instructions in the robot dataset without
having to collect new robot trajectories.

  3. 3.

Co-located Human-Robot In this case, a human and a robot perform the same task
in the same workspace. We used this approach to collect human demonstrations
and robot trajectories in diverse spaces such as a living space with sofas, a
meeting room with whiteboards, hardware workstations with toy tools, a kitchen
with a countertop, refrigerator and sink, a storage supplies area, and more.

We show examples of paired prompt and robot videos from each of the three
datasets in Figure 2. As can be seen, there is a considerable difference in
the backgrounds and distractor objects in the Hindsight Human-Robot and Co-
located Human-Robot datasets. A different complexity arises when comparing the
first approach (Robot-Robot) where the actor is a robot with same morphology
to the other two cases where the human is the actor in the prompt videos.

After combining all the datasets, we have $\sim$100k robot videos and
$\sim$10k human videos covering the tasks introduced in RT-1 [7] and RT-2 [8].
We include videos from all three data sources as they represent varying levels
of difficulty and expense to collect. Pairing existing robot datasets requires
less additional effort but lacks diversity in how the task is done. The second
source of data is created by asking humans to mimic existing robot
trajectories. While this adds some diversity in prompt videos, it does not
cover any new tasks on the robot side. Finally, the presumed gold-standard is
to collect data where both humans and robots are co-located in the same
environment and perform diverse tasks. This takes the most amount of time as
labor is required both of the humans and robot trajectories collected through
tele-operation.

###  II-C Model Architecture

Figure 3: Architecture. Our model takes as input frames of the prompt video
and the robot’s current observations, encodes those into prompt video and
robot state token embeddings, which are then processed through into state-
prompt encoder and decoded into a robot action for the current timestep. More
details in §II-C.

Our policy takes as input the prompt video and the current robot state and
outputs robot actions. It consists of four modules: (1) prompt video encoder
(2) robot state encoder, (3) state-prompt encoder, and (4) robot action
decoder. The full architecture is illustrated in Figure 3 and each of the
modules are detailed below:

(1) Prompt Video Encoder encodes the video demonstration provided as a
reference to convey the desired task semantics. The prompt video encoder
implicitly learns to infer what task should be performed and how it needs to
be done. The prompt encoder consists of a per-frame Image encoder $\phi_{p}$
(ViT [14]) followed by a Perceiver Resampler [1, 19] $\psi_{p}$. The output of
the prompt encoder $\psi_{p}(\phi_{p}(V))=z_{prompt}$ is a set of $N$ tokens
of d-dimension to condition the policy with the task relevant attributes from
the video.

(2) Robot State Encoder encodes the current state of the robot given the
current frame and last $k$ frames as input. Note that this module also encodes
information about the objects and environment of the robot. The architecture
is similar to the prompt encoder, that is, a per-frame Image encoder
$\phi_{s}$ followed by a Perceiver Resampler $\psi_{s}$. Similar to the prompt
encoder’s outputs, the output of the state encoder is
$\psi_{s}(\phi_{s}(S_{t}))=z_{state}$ that encodes the latent environment and
robot state information from the history of recent observations.

We use the same image encoder weights for both (1) and (2), that is,
$\phi_{p}\\!=\\!\phi_{s}\\!=\\!\phi$. The role of the image encoder $\phi$ is
to capture spatial visual information in each frame. The Perceiver Resampler
is used to enable temporal learning across frames as well as reduce the number
of video tokens that must be passed into the action decoder.

(3) State-Prompt Encoder The state-prompt encoder takes the prompt video
encoding $z_{prompt}$ and robot state encoding $z_{state}$ and outputs a task
encoding relevant for action prediction $z_{state|prompt}$. The module is
trained to output robot actions by cross-attending between the state encoding
as queries and the prompt video encoding as keys and values. Intuitively, the
state-prompt encoder enables fusion of the state and prompt information. For
example, if the prompt video demonstrates picking up of an apple in the basket
and the current state contains apple, banana and orange, then the cross
attention between the state and prompt encoding enables learning for which
object to attend to in the state, which is crucial for the next step of action
decoding. We refer to the output of the state-prompt encoder as prompt-aware
state tokens.

(4) Robot Action Decoder The goal of the action decoder is to predict the
action vector $a_{t}$ for the current state $S_{t}$ such that it completes the
task shown in the prompt video $V_{p}$. The action decoder is a transformer
decoder architecture that takes in the fixed action position tokens [49] as
input queries and the prompt-aware state tokens $z_{state|prompt}$ for keys
and values. The size of the action position embedding is $N\times d$ where $N$
is the number of action dimensions and $d$ is the transformer embedding
dimension. More details on the action vector in §II-D.

The action position embeddings cross-attend to the prompt-aware state tokens
to predict the target binned action values as output. Each output token of the
action decoder corresponds to an action dimension for the mode, arm and base.
Specifically, each token embedding is projected to 256 dimensions and a
softmax layer is applied on the top to obtain the bin corresponding to the
target action vector. Unlike prior work [7, 8] that use autoregressive action
decoding that requires multiple forward passes during inference, we use action
position embeddings for one forward pass prediction like in ACT [49]. Instead
of predicting one action for the next timestep, we follow the approach
outlined in [21, 49] and train the policy with a prediction horizon of four
steps. We always use the action bin that has the highest probability, i.e.
argmax over predicted probabilities, to choose the action value for execution.

Cross-Attention Layers. In the Vid2Robot architecture, we use Cross-Attention
Transformer layers extensively. They are used in the following modules: Prompt
Resampler, State Resampler, State-Prompt Encoder and Action Decoder. We found
Cross-Attention layers are helpful in managing the high number of tokens and
the resulting large attention matrices when processing both prompt videos and
robot state videos. This is because the standard self-attention layers would
require orders of magnitude more memory to process the same video. For
example, when using ViT-B/16 the total number of video tokens for a $16$ frame
reference video and a $8$ frame robot state video at $224\times 224$
resolution would be $8\times 196+16\times 196=4704$. A full self-attention
operation on this would lead to an attention matrix with $4704^{2}\sim
22\mathrm{M}$ entries. However, by using two Perceiver Resamplers with 64
latents we were able to train with attention matrices of the size $8\times
196\times 64+16\times 196\times 64\sim.3\mathrm{M}$. Thus, cross attention
layers in Vid2Robot play an important role in reducing attention computation
and enabling training with paired videos.

###  II-D Preprocessing

To handle the varying lengths of videos for efficient training, we randomly
sample $N\\!\\!=\\!\\!16$ frames always including first and last frames and
sort them in increasing order of time. During training, we sample a robot
state $S_{t}$ by sampling a random timestep first. We then select the
preceding $k-1$ frames to create a robot state video comprising of a total of
$k\\!\\!=\\!\\!8$ frames before. In case there are less than $k-1$ frames
before the current time-step, we repeat the first frame to create a fixed size
robot state video. The pixel values in each frame are normalized between 0 to
1. Each frame is resized to $(224,224)$. Photometric distortions like
cropping, brightness, contrast, hue and saturation are applied during
training.

The action vector consists of values indicating the mode, gripper pose and
closedness as well as base displacement and rotation. Each of the values have
different ranges, which we first use to scale the values in between 0 and 1.
We then discretize the values into 256 bins each. In total, we construct
11-dim action vector as target, each of which has value between [0, 255]. In
this study, we train and evaluate in the scenarios where base remains
stationary.

###  II-E Training

Figure 4: Training Setup. We show all the losses Vid2Robot is trained with and
how each loss is connected to its different modules. Along with (1) the main
action prediction loss, we apply three auxiliary losses: (2) temporal video
alignment loss, (3) a contrastive loss between the prompt and robot video
performing the same task, and (4) a contrastive loss between a prompt/robot
video with the language embedding. More details in §II-E.

Action Prediction Loss We train Vid2Robot end-to-end with behavior cloning.
The idea is to learn video representations from raw pixels to recognize task
verb and objects, as well as learn motor control to accomplish it. We use a
classification loss on actions that have been tokenized into
$N\\!\\!=\\!\\!256$ bins. Given the robot trajectory for performing a task
with current visual observations $x_{t}$, we have the corresponding expert
action $a_{t}$. The action prediction loss is Cross Entropy between the
predicted action and the expert action as:

$L_{CE}(a_{t},\hat{a}_{t})=\sum_{\tau}a_{t}\log\hat{a}_{t}$ (1)

This trains all the model parameters, as shown in Fig 3.

Auxiliary Losses. Although our dataset size is substantial, it is insufficient
for training large Transformer based models. In order to prevent over-fitting
by just predicting actions correctly on the training set, we add three
auxiliary losses that encourage learning features that are helpful in
understanding semantics in prompt videos.

Video Alignment Loss: We want to encourage temporal alignment between prompt
videos and robot videos performing that show the same task. By aligning prompt
videos and robot videos, we want the image encoder to learn to be invariant to
different embodiments, lighting, backgrounds, view-angles and distractor
objects while still encoding features relevant to predicting task progress.
Our choice of loss is the temporal-cycle consistency loss introduced in [17].
This loss has been shown to encode task progress when trained on videos of
different agents performing the same task [47]. This loss is applied on per-
frame image embeddings of the prompt $V_{p}$ and robot $V_{r}$ videos during
training. To apply the loss, we average pool the per-frame embeddings output
in spatial dimensions from image encoder $\phi$ and apply a projector head of
2-layer MLP [10]. We call this as alignment pooling layer $\Phi$ on the per-
frame image embeddings, as shown in Fig 4. For each video $V_{i}$, this
results in a sequence of embeddings
$E_{i}=\\{\Phi(v_{i}^{1}),\Phi(v_{i}^{2})...,\Phi(v_{i}^{L_{i}})\\}$, where
$L_{i}$ is the length of the $i^{th}$ video.

We apply TCC loss on encoding $E_{p}$, and $E_{r}$ for prompt and robot video
respectively. The intuitive idea of TCC loss is that we want to ensure the
representation of every frame of $E_{p}$ should have a correspondence in
$E_{r}$ and vice versa. This involves two steps: First, we compute soft
neighbor of $t^{th}$ frame of $E_{p}$ ($E_{p}^{t}$ in short) in $E_{r}$ and
call it $\widetilde{E_{pr}^{t}}$.

$\widetilde{E_{pr}^{t}}=\sum_{k}^{L_{r}}\alpha_{k}E_{r}^{k},\quad\mathrm{where}\quad\alpha_{k}=\frac{e^{-\|E_{i}^{t}-E_{j}^{k}\|^{2}}}{\sum_{k}^{L_{j}}e^{-\|E_{i}^{t}-E_{j}^{k}\|^{2}}}$
(2)

Second, we find the corresponding frame for this newly computed soft-neighbour
in $E_{p}$. This is called cycle-back in [17] and it involves similar soft-
neighbour computation as in Equation 2 to obtain say $\widehat{E_{pr}^{t}}$,
which ideally should be same as $t$, that is, $(\widehat{E_{pr}^{t}}-t)^{2}$
should be minimized. TCC loss minimizes such mean squared error between all
frames for prompt and robot video encodings and vice-versa, that is,

$\displaystyle L_{TCC}(E_{p},E_{r})=\sum_{t\in
V_{p}}(\widehat{E_{pr}^{t}}-t)^{2}$ (3) $\displaystyle
L_{TCC}=\frac{L_{TCC}(E_{p},E_{r})+L_{TCC}(E_{r},E_{p})}{2}$

Prompt-Robot Video Contrastive Loss (VVCL): We apply contrastive loss between
prompt tokens produced by robot or prompt video performing the same task. This
loss encourages the prompt encodings to learn task semantics from video tokens
only in a self-supervised manner. A thing to note here is that the initial
pairing of prompt and robot video has been done using natural language.
However, by using a constrastive loss only on video embeddings with a self-
supervised loss, we hope to encode features not covered by the short natural
language embedding itself. Examples of these features include similar motions
like reaching for objects, and rotating the robot arm. We use a Attention
Pooling layer to merge features from the $N$ prompt tokens to produce a single
embedding for each video. We apply the SigLIP [48] loss between video-video
pairs to encourage videos showing same task, involving similar motions and
interacting objects, to be close to each other while being away from other
videos in the batch. A batch contains the same number of robot videos and
prompt videos, say $B$. We use the prompt encoder $\psi_{p}(\phi(\cdot))$ to
obtain a batch of full robot video embeddings $Z_{robot}$ and prompt video
embeddings $Z_{prompt}$, each of size $B\times d$. We multiply them,
$Z_{robot}\cdot Z_{prompt}^{T}$ to obtain a $B\times B$ matrix. Adding a
learnable temperature $\tau$ and bias $b$, we have our logit matrix as
$\hat{Y}=(Z_{robot}\cdot Z_{prompt}^{T})*\tau+b$. We consider the videos of
robot and prompt performing the same task as positives and assign them a label
of 1 along the diagonal and -1 for off-diagonal pairs, that is, the label
matrix $Y=2\mathrm{I}_{B}-1$. SigLIP loss is the negative loglikelihood,
specifically, $\sigma^{\prime}(Z_{1},Z_{2})=-\sum\log\sigma(Y\cdot(Z_{1}\cdot
Z_{2}^{T})*t+b)\text{, where }Y=2\mathrm{I}_{B}-1$. The video-video
contrastive loss is then defined as:

$L_{VVCL}=\sigma^{\prime}(Z_{prompt},Z_{robot})$ (4)

Video-text Contrastive Loss (VTCL): We add a contrastive loss between prompt
tokens $Z_{prompt}$ and those produced by robot video $Z_{robot}$ and the
embedding of the text instructions of the task $Z_{text}$. This encourages a
part of the embedding space to be aware of object names and verbs present in
the prompt and the robot videos. A version of this loss has been applied
before by BC-Z [21] as auxiliary language regression loss. We use an Attention
Pooling layer [46] with one latent query to merge features from the $N$ prompt
tokens to produce a single embedding for each video. Within a batch, we
retrieve $B$ pairs of video and text embeddings. Similar to Equation 4, we
apply SigLIP [48] loss to get

$L_{VTCL}=\frac{\sigma^{\prime}(Z_{prompt},Z_{text})+\sigma^{\prime}(Z_{robot},Z_{text})}{2}$
(5)

This encourages every video to have similar embeddings to their textual
description embeddings, while being different from the text embeddings
corresponding to other videos in the batch.

Overall, we apply the mean of all four losses for training that is
$L=\frac{1}{4}(L_{CE}+L_{TCC}+L_{VVCL}+L_{VTCL})$.

###  II-F Implementation

We trained the model (implemented in Jax) for 200K iterations. We use AdamW
optimizer with an initial learning rate of 8e-5 using a cosine learning rate
schedule with warmup steps 2,000 and final learning rate of 1e-6. For both the
Prompt and State Resamplers, we use 2 Perceiver Resampler layers with 64
latents. Both state-prompt encoder and action decoder are 4 layer deep cross-
attention transformers.

##  III Experiments

Figure 5: Policy Rollouts. Each row shows a prompt video of a human doing a
task on the left, and on the right we show the corresponding successful robot
rollouts using Vid2Robot. Note how visually different the prompts are, while
the policy rollouts are recorded with different lighting, background, as well
as number and placement of the distractor objects.

We present results with real robot evaluations for our multi-task video-
conditioned policy. One of the key questions that we tackle in this work is
how well robots can imitate humans performing manipulation tasks. Because of
differences in embodiments, humans perform manipulation tasks at a different
speed and style. We study the effect of using robot as well as human videos as
prompts.

Metrics. We refer to a rollout as a sequence of actions inferred from the
policy and executed on the robot from an initial state observation and prompt
video, until the policy terminates or a maximum number of steps are taken,
whichever is lower. We define success for a rollout when the policy executes
the task instruction shown in the prompt video. A successful rollout involves
correct actions to be taken successively in the environment, without any
assistance for resets or recovery. For each task instruction, we record many
rollouts per policy. We take the average of success recorded across all the
rollouts and call it the Success Rate for that task. Aggregated success rate
across tasks is referred as Overall Success Rate.

A mistake made early on in a rollout can result in poor success rate, even if
the model’s offline overall prediction accuracy is high. For example, if a
policy makes an error while grasping a water bottle early on in the task and
it slips to an unreachable location, the rollout will be marked as a failure,
even if the policy had good action predictions for the later steps. To record
partial progress for a rollout, we annotate whether the robot reached the
correct location, grasped the correct object, released the object at the
correct location, and terminated the task correctly. More details on partial
success analysis in §III-A.

Evaluation Setup. We ask human raters to evaluate success for a policy’s
rollout on a robot. We evaluate the policies by varying the robots, lighting
conditions, chest of drawers, and objects. We ensure the policies being
evaluated are shown similar initial object configurations during rollouts. The
initial state is randomized after all policies have been evaluated for a given
initial state. For all rollouts, we sample prompt videos that are not seen by
the models during training. This ensures that the policies are evaluated for
whether they can recognize the task from new prompt videos or not.

Baselines. We compare our model with BC-Z [21], a video conditioned policy
using a ResNet-18 encoder. BC-Z [22] involves demonstration-observation pairs
processed via a FiLM [36] conditioned ResNet encoder and fed into a ResNet
based policy network to predict robot actions. For a fair comparison, we train
the BC-Z model with the same training data used to train the Vid2Robot model.
We run rollouts of BC-Z policy for a fixed maximum number of steps, as BC-Z
doesn’t have a terminate action.

TABLE I: Task Success Rate for Robot and Human prompts. Prompter | Model | pick | pick-place on | place into | open | close | move near | knock over | place upright | Overall
---|---|---|---|---|---|---|---|---|---|---
Robot | BC-Z | 75.0% | 50.0% |  61.5% | 16.7% | 66.7% |  44.0% |  58.3% |  50.0% | 52.6%
Vid2Robot | 75.0% | 58.8% | 50.0% | 91.7% | 100.0% | 33.3% | 41.7% | 16.7% | 54.9%
Human | BC-Z | 50.0% | 12.5% | 12.5% | 0.0% | 50.0% | 43.8% | 12.5% |  50.0% | 30.6%
Vid2Robot | 100.0% | 50.0% | 50.0% | 62.5% | 87.5% | 43.8% | 25.0% | 12.5% | 52.8%

Key Questions and Results We address the following questions in this work:

  1. 1.

How well do video-conditioned policies perform when they are shown a task in
an unseen video? (Fig 5, § III-A)

  2. 2.

What is the gap in success rate due to prompt embodiment difference (robot v/s
human)? (§ III-A)

  3. 3.

Can we leverage the learned motion representations for out-of-distribution
object interactions? (§ III-B)

###  III-A Task-based success

We compare the our Vid2Robot model and baseline BC-Z with robot and human
prompt videos in Table I. Both Vid2Robot and BC-Z were trained on a same data
mixture containing robot-robot and human-robot paired data. Prompt videos
cover a subset of the training tasks but the videos themselves are new for the
models. In this evaluation, we investigate what each model’s ability is to
infer the task specification from prompt video as well as the current observed
state of the robot.

In order to test the capabilities of the model in different settings on real
robot, we evaluate it across eight categories of manipulation tasks as shown
in Table I. Specifically, we evaluate for nine tasks: ‘knock water bottle
over’, ‘move rxbar chocolate near coke can’, ‘move green jalapeno chip bag
near coke can’, ‘pick green rice chip bag’, ‘place coke can upright’, ‘pick
coke can from bottom drawer and place on counter’, ‘open middle drawer’,
‘close middle drawer’, and ‘place apple into top drawer’.

We ask four evaluators to carry out two rollouts per task for a prompt video
dataset and policy setting (a row in Table I), that implies, we have eight
trials per task to evaluate a policy’s task success rate. We report overall
success rate per row over nine tasks with eight trials per task, that is,
$9\\!\times\\!8\\!=\\!72$ trials. In total, our evaluations in Table I
required $72\\!\times\\!4\\!=\\!288$ real robot rollouts.

####  III-A1 What is the gap in success rate due to embodiment difference in
prompt videos?

We compare our model with BC-Z when prompted with robot and human videos. BC-Z
serves as a strong baseline for our comparisons. The overall success rate of
our model Vid2Robot outperforms BC-Z for Human prompt videos by 20%, and is
comparable for Robot prompt videos. Note that there is an order of magnitude
more training samples for robot trajectories than human videos in our training
mixture. Hence, there isn’t a significant gap in performance for robot prompt
videos. For human prompt videos, our model outperforms BC-Z in most tasks,
showing that Vid2Robot captures the task semantics from prompt videos better
than the baseline. Our model outperforms in tasks like picking from drawer and
placing on the counter, and opening/closing drawer tasks by a large margin.
The most challenging task is placing upright and knocking over. We analyze the
failure reasons in §V Fig 9.

####  III-A2 How well do video-conditioned policies perform when they are
shown a task in an unseen video?

In addition to marking a rollout as a success, we recorded partial success
annotations per rollout. In Fig 6, we observe that our model reaches to the
correct object 78% of the time, about 8% more as compared to baseline. The
policies sometimes fail to reach the correct object and go towards a
distractor instead. Next, grasping errors happen, particularly with small and
deformable objects and in collision prone areas like drawer handle or
counter’s edge. Here our model (65%) outperforms BC-Z (45%) by a large margin
of 20%. A successful grasp often the most difficult part in a rollout, and the
most crucial for success. After grasping, most tasks require releasing at a
correct location. There is a slight drop in success rate in both models due to
incorrect release during the rollouts. While BC-Z runs for a fixed number of
steps, our policy Vid2Robot predicts when to terminate. We observe that the
rate of release and terminate is almost identical, about 57% for our model,
that implies, that after releasing at correct location, Vid2Robot mostly
terminates successfully.

Figure 6: Partial Success Rate for BC-Z and Vid2Robot. Our policy Vid2Robot
outperforms BC-Z in terms of reaching the correct object, grasping it,
releasing it at the correct location and then terminating the episode
correctly. Note that BC-Z does not have terminate control.

###  III-B Cross-object motion transfer

TABLE II: Cross-object motion transfer success. |  | pick- | place | place | knock |
---|---|---|---|---|---|---
Model | pick | place on | into | upright | over | Overall
BC-Z | 45.8% |  00.0% | 29.2% | 12.5% |  00.0% | 17.5%
Vid2Robot | 45.8% | 25.0% | 54.2% | 16.7% | 29.2% | 34.2%
Figure 7: Qualitative results for cross-object motion transfer. Given a prompt
video of placing coke can upright, we rollout the policy with a green can,
chips bag, stapler and a soft toy in front of the robot. We observe that our
model can infer the motion of place upright in the prompt video and apply it
on other objects. There is an implicit notion of pragmatics in the policy as
shown by the selection of green can over other objects.

Our policy and baseline were trained with paired videos as discussed in §II-B.
This implies that the training data included only those scenarios where the
interaction object shown in prompt is present in the current robot
observations. But what if we provided a prompt video of one object and tested
on other objects. Does it do the same motion as shown in the prompt video?
Interestingly, we found our model to perform learned manipulation actions on
objects that it has not seen in train set. We call this emergent behavior as
cross-object motion transfer.

We compare Vid2Robot with BC-Z for cross object motion transfer ability with
five prompt videos, namely, ‘knock water bottle over’, ‘pick green rice chip
bag’, ‘place coke can upright’, ‘pick coke can from bottom drawer and place on
counter’, and ‘place apple into top drawer’. Each prompt video is evaluated
with unrelated objects in robot’s initial observation. The objects used for
evaluation are ‘orange’, ‘green can’, ‘chips bag’, ‘banana’, ‘pink piggy soft
toy’, ‘wrist watch’. We selected objects to have diverse shape, size, and
deformability to evaluate situations that require different grasps for
success.

The evaluation setup is similar to §III-A. Here the evaluator sets up one of
the object for a task and records rollouts for each model. We compare 2 models
on 5 tasks with 6 objects, so every evaluator runs
$2\\!\times\\!5\\!\times\\!6\\!=\\!60$ rollouts. We repeat the evaluation with
four raters, thus reporting results in Table II on a total of
$4\\!\times\\!60=240$ rollouts.

####  III-B1 Can we provide a prompt video of one object and test it on other
objects? Does the policy do the same motion as shown in the prompt video?

In Fig 7, we show the above experimental setup qualitatively. We use a prompt
video to ‘place coke can upright’. We observe that the policy is able transfer
the action of ‘placing upright’ to several objects, like a green can, a chips
bag, a stapler, and a soft toy. Note that the policy adheres to the prompt
video and chooses green can over chips bag or banana for placing upright.

Quantitatively, we observe that BC-Z is often unable to successfully complete
the tasks when testing cross=object motion transfer, as shown in each task in
Table II. In contrast, our model (34%) performs better than BC-Z (17%) in this
setting and performs the motion indicated in the prompt video. Our model is
comparable to BC-Z with 45% success rate on picking out-of-distribution
objects. More importantly, tasks involving placing into drawers demonstrates
significant improvement ($29\%\rightarrow 54\%$). For certain tasks like
picking from drawers and placing on counters and knocking over, BC-Z is unable
to perform at all whereas Vid2Robot is able to complete the task $25\%-29\%$
of the time.

###  III-C Ablations

In §II-E, we presented action prediction loss and three auxiliary losses. Here
we analyze the role of these additional loss functions to the overall success
rate. We investigate the impact of (1) not using any auxiliary loss, and (2)
adding auxiliary language loss.

We consider the tasks similar to that described in §III-A, that is, 9 tasks
for evaluating each policy.

We have 3 model variants, namely, the original Vid2Robot, the one without
video-text contrastive loss (CL) and the one with only action prediction loss.
We ask 3 human evaluators to run the each model variant with 2 rollouts each.
In total, we report results with
$3\\!\times\\!3\\!\times\\!9\\!\times\\!2\\!=\\!162$ rollouts in Fig 8. The
error bars indicate the standard deviation for success reported on rollouts
with each model variant.

####  III-C1 What is the impact of not using any auxiliary loss?

We observe that the performance of our model (61%) is significantly improved
by enforcing representation constraints through auxiliary losses, in
comparison to using only action prediction loss (45%). It highlights the
importance of the proposed auxiliary losses in §II-E.

####  III-C2 What is the impact of the auxiliary language loss?

BC-Z proposed to use language representations to improve video representations
for conditioning the policy. We compare our policy with another variant
trained with all losses but the Video-Text CL. We observe only marginal
improvement of 1-2% in success rate when using the language loss. This implies
that video alignment and video contrastive loss contribute significantly
towards performance improvement. Our results hope to serve as a promising
evidence that effective video representations can be learned without auxiliary
losses that use pre-trained language embeddings.

Figure 8: Ablation for auxilliary losses used in Vid2Robot. We compare our
proposed approach that has all auxiliary losses (green, left) with a variant
without language contrastive loss that was originally proposed in BC-Z
(orange, middle) and a version with no auxilliary losses (blue, right). More
details in (§III-C) Figure 9: Failure analysis with policy rollouts. (Top)
Policy predicts gripper pose and depends on the IK solver to move the arm.
Sometimes, the IK solution can block the robot’s camera view. (Middle)
Grasping failures happen, especially with transparent and deformable objects.
(Bottom) Distractor objects as well as difference in lighting and background
may cause recognition errors, where policy might perform the correct motion
but with incorrect object(s).

##  IV Related Work

Task Specifications for Robots The development of general-purpose robots
hinges on effectively grounding task specifications. Videos are a dense source
of information that not only provide what to do but also how to do it in
physical world. Recent works have used videos for task specification [4, 23,
40]. Another line of work uses videos to learn world models to predict future
visual observations [29, 26, 9, 31, 15]. While language [45, 7, 33, 34], final
goal images [24, 6], and others like hand-drawn inputs [43] have been proposed
as means for task specification, learning from prompt videos is complementary
to these approaches and inevitable for rapid adaptation of trained polices to
perform new manipulation skills at deployment.

Learning from Human Demonstrations As videos of humans performing various
tasks proliferate the internet, several works aim to address how to best
leverage this information for robot learning. The difference in robot vs human
embodiment poses a significant challenge, for which existing approaches range
from translating image of a human into the robot [42] to inpainting for agent-
agnostic representations [3]. Many prior works propose to leverage off-the-
shelf models for hand pose estimation and contact tracking [4, 12, 37],
object-centric representations [38, 20], as well as reward functions for
reinforcement learning [3, 27, 42]. Other methods [32, 44, 4] cast this
problem into visual representation learning to accelerate learning of
downstream motor control tasks. While these modular learning solutions work
well in limited datasets, these are prone to compounding error of each of its
component, and thus, not efficiently scalable. End-to-end training approaches
for goal-conditioned imitation learning [11, 41, 18, 13] and reinforcement
learning [39, 35] are promising alternatives to these techniques, but these
results have been largely limited in simulation and hindered by sim-to-real
gap. In contrast, we choose to tackle this as an end-to-end large multi-task
learning from human videos with real robot evaluations.

Imitation via Paired Demonstrations Our setup of paired prompt videos and
robot trajectory is most similar to One-Shot Visual Imitation literature. Many
prior works assume access to pairs, where the first is used as the
demonstration of the task to be performed, and the second as the observation
of the agent. Some of the early works [16] proposed training a demonstration
network via temporal convolution and neighorhood attention to condition a
manipulation policy network. In more recent approaches like [11, 28, 20],
paired demonstrations and observations are used to train a transformer policy,
often with additional constraints like inverse dynamics prediction[11] or
contrastive representation learning [28]. However, these approaches are
largely evaluated in specific set of simulated tasks, and not compared on real
robots. Most similar to our work is BC-Z [22] which reports evaluations with
real robot tasks. While our setup is similar to some of this prior art, our
model Vid2Robot couples large image encoders, cross-attention layers, and
contrastive auxiliary losses to learn a manipulation policy that imitates a
human showing a task.

##  V Limitations and Future Directions

In §III, we show that our approach has improved over previous work but there
is a gap in performance for video-conditioned policies. Language conditioned
policies like [8] shows a higher success for known set of tasks with several
hundreds of teleoperation trajectories for training. We, on the other hand,
accomplish the first milestone of evaluating the video-conditioned policies in
the similar setup. We discuss three limitations of our work and provide
insights for future directions here.

First, we qualitatively investigate some reasons for failure of a policy
rollout. In Fig 9, we illustrate and explain 3 examples showing how self
occlusion, grasping errors and presence of distractors can lead to failure
during any rollout. Second, we observe a significant drop in the grasping
success in Fig 6. While we use robot camera observation to estimate the state
and implicitly learn depth estimation, it is often incomplete when there is
occlusion or when the robot gripper is out of camera view. By enhancing the
state information with multimodal sensor fusion, we may improve the grasp
success rate. Third, we consider carefully collected short task instruction
demonstrations from three different sources as shown in §II-B, all of which
are 5 to 20 seconds videos. To test our models on long horizon demonstrations
or ‘in-the-wild’ videos online, we need effective pairing strategies for
videos and a few corresponding robot trajectories to train the policy.

##  VI Conclusion

We demonstrate novel methods for both data collection and modeling for video
conditioned skill learning. These skills generalize to novel object
configurations and more abstracted verb meanings when no immediately obvious
object is visible. The skills and generality provided by our model complement
other approaches to widen the set of skills that robots have access to, and to
include skills not otherwise easily acquired. Future work can leverage these
learned primitives to execute novel task plans. We hope our cross-object
motion transfer experiments will encourage further research in transferring
motion to new objects and settings for bootstrapping data collection, and
enabling human-robot interaction with rapid adaptation to new skills.

TABLE III: Author Contributions Name | Data Collection | Model Training | Evaluations | Infrastructure | Leadership | Paper Writing
---|---|---|---|---|---|---
Ayzaan Wahid |  | ✓ |  |  |  |
Christine W.Y. Chan | ✓ |  |  |  |  |
Danny Driess |  |  |  | ✓ |  | ✓
Debidatta Dwibedi | ✓ | ✓ | ✓ | ✓ | ✓ | ✓
Igor Gilitschenski |  |  |  |  | ✓ | ✓
Maria Attarian | ✓ | ✓ |  | ✓ |  | ✓
Nikhil J Joshi | ✓ |  | ✓ | ✓ | ✓ | ✓
Pannag R Sanketi | ✓ |  |  | ✓ | ✓ |
Pierre Sermanet | ✓ |  |  |  |  |
Quan Vuong |  | ✓ |  | ✓ |  |
Stefan Welker | ✓ |  |  | ✓ |  |
Vidhi Jain | ✓ | ✓ | ✓ | ✓ | ✓ | ✓
Yonatan Bisk |  |  |  |  | ✓ | ✓

## Acknowledgments

We would like to thank Yansong Pang, Grecia Salazar, Utsav Malla, Deeksha
Manjunath, Jornell Quiambao, Sarah Nguyen, Sangeetha Ramesh, Tran Pham, Samuel
Wan, Tomas Jackson, Jodilyn Peralta, Celeste Barajas, Elio Prado, Rochelle
Dela Cruz, Alex Luong and Krista Reymann for supporting data collection via
teleoperation. Special thanks to Jornell Quiambao, Grecia Salazar, Utsav
Malla, Deeksha Manjunath, Sarah Nguyen, Sangeetha Ramesh, and Jaspiar Singh
for evaluations on robot; Michael Ahn, Anthony Brohan and Keerthana
Gopalakrishnan for policy evaluation infrastructure; Suneel Belkhale, Dorsa
Sadigh, Chelsea Finn, and Sergey Levine for helpful discussions; Jonathan
Tompson, and Vincent Vanhouke for thorough feedback on the writing. This work
was also supported by Google Robotics Funding at Carnegie Mellon University.

## References

  * Alayrac et al. [2022] Jean-Baptiste Alayrac, Jeff Donahue, Pauline Luc, Antoine Miech, Iain Barr, Yana Hasson, Karel Lenc, Arthur Mensch, Katherine Millican, Malcolm Reynolds, Roman Ring, Eliza Rutherford, Serkan Cabi, Tengda Han, Zhitao Gong, Sina Samangooei, Marianne Monteiro, Jacob Menick, Sebastian Borgeaud, Andrew Brock, Aida Nematzadeh, Sahand Sharifzadeh, Mikolaj Binkowski, Ricardo Barreira, Oriol Vinyals, Andrew Zisserman, and Karen Simonyan.  Flamingo: a visual language model for few-shot learning.  In Alice H. Oh, Alekh Agarwal, Danielle Belgrave, and Kyunghyun Cho, editors, _Advances in Neural Information Processing Systems_ , 2022.  URL https://openreview.net/forum?id=EbMuimAbPbs.
  * Arenas et al. [2023] Montserrat Gonzalez Arenas, Ted Xiao, Sumeet Singh, Vidhi Jain, Allen Z. Ren, Quan Vuong, Jake Varley, Alexander Herzog, Isabel Leal, Sean Kirmani, Dorsa Sadigh, Vikas Sindhwani, Kanishka Rao, Jacky Liang, and Andy Zeng.  How to prompt your robot: A promptbook for manipulation skills with code as policies.  In _2nd Workshop on Language and Robot Learning: Language as Grounding_ , 2023.  URL https://openreview.net/forum?id=T8AiZj1QdN.
  * Bahl et al. [2022] Shikhar Bahl, Abhinav Gupta, and Deepak Pathak.  Human-to-robot imitation in the wild.  In _Robotics Science and Systems (RSS)_ , 2022.
  * Bahl et al. [2023] Shikhar Bahl, Russell Mendonca, Lili Chen, Unnat Jain, and Deepak Pathak.  Affordances from human videos as a versatile representation for robotics.  In _Computer Vision and Pattern Recognition_ , April 2023.
  * Bisk et al. [2020] Yonatan Bisk, Rowan Zellers, Ronan Le Bras, Jianfeng Gao, and Yejin Choi.  PIQA: Reasoning about Physical Commonsense in Natural Language.  In _Thirty-Fourth AAAI Conference on Artificial Intelligence_ , 2020\.  URL https://yonatanbisk.com/piqa.
  * Bousmalis et al. [2023] Konstantinos Bousmalis, Giulia Vezzani, Dushyant Rao, Coline Devin, Alex X. Lee, Maria Bauzá, Todor Davchev, Yuxiang Zhou, Agrim Gupta, Akhil S. Raju, Antoine Laurens, Claudio Fantacci, Valentin Dalibard, Martina Zambelli, Murilo Fernandes Martins, Rugile Pevceviciute, Michiel Blokzijl, Misha Denil, Nathan Batchelor, Thomas Lampe, Emilio Parisotto, Konrad Zolna, Scott E. Reed, Sergio Gomez Colmenarejo, Jonathan Scholz, Abbas Abdolmaleki, Oliver Groth, Jean-Baptiste Regli, Oleg O. Sushkov, Tom Rothorl, José Enrique Chen, Yusuf Aytar, David Barker, Joy Ortiz, Martin A. Riedmiller, Jost Tobias Springenberg, Raia Hadsell, Francesco Nori, and Nicolas Manfred Otto Heess.  Robocat: A self-improving generalist agent for robotic manipulation.  In _Transactions on Machine Learning Research_ , September 2023.
  * Brohan et al. [2022] Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, Karol Hausman, Alex Herzog, Jasmine Hsu, Julian Ibarz, Brian Ichter, Alex Irpan, Tomas Jackson, Sally Jesmonth, Nikhil J Joshi, Ryan Julian, Dmitry Kalashnikov, Yuheng Kuang, Isabel Leal, Kuang-Huei Lee, Sergey Levine, Yao Lu, Utsav Malla, Deeksha Manjunath, Igor Mordatch, Ofir Nachum, Carolina Parada, Jodilyn Peralta, Emily Perez, Karl Pertsch, Jornell Quiambao, Kanishka Rao, Michael Ryoo, Grecia Salazar, Pannag Sanketi, Kevin Sayed, Jaspiar Singh, Sumedh Sontakke, Austin Stone, Clayton Tan, Huong Tran, Vincent Vanhoucke, Steve Vega, Quan Vuong, Fei Xia, Ted Xiao, Peng Xu, Sichun Xu, Tianhe Yu, and Brianna Zitkovich.  RT-1: Robotics transformer for Real-World control at scale.  In _arXiv preprint arXiv:2212.06817_ , December 2022.
  * Brohan et al. [2023] Anthony Brohan, Noah Brown, Justice Carbajal, Yevgen Chebotar, Xi Chen, Krzysztof Choromanski, Tianli Ding, Danny Driess, Avinava Dubey, Chelsea Finn, Pete Florence, Chuyuan Fu, Montse Gonzalez Arenas, Keerthana Gopalakrishnan, Kehang Han, Karol Hausman, Alex Herzog, Jasmine Hsu, Brian Ichter, Alex Irpan, Nikhil Joshi, Ryan Julian, Dmitry Kalashnikov, Yuheng Kuang, Isabel Leal, Lisa Lee, Tsang-Wei Edward Lee, Sergey Levine, Yao Lu, Henryk Michalewski, Igor Mordatch, Karl Pertsch, Kanishka Rao, Krista Reymann, Michael Ryoo, Grecia Salazar, Pannag Sanketi, Pierre Sermanet, Jaspiar Singh, Anikait Singh, Radu Soricut, Huong Tran, Vincent Vanhoucke, Quan Vuong, Ayzaan Wahid, Stefan Welker, Paul Wohlhart, Jialin Wu, Fei Xia, Ted Xiao, Peng Xu, Sichun Xu, Tianhe Yu, and Brianna Zitkovich.  RT-2: Vision-Language-Action Models Transfer Web Knowledge to Robotic Control.  In _arXiv preprint arXiv:2307.15818_ , 2023.
  * Chane-Sane et al. [2021] Elliot Chane-Sane, Cordelia Schmid, and Ivan Laptev.  Goal-conditioned reinforcement learning with imagined subgoals.  In Marina Meila and Tong Zhang, editors, _Proceedings of the 38th International Conference on Machine Learning_ , volume 139 of _Proceedings of Machine Learning Research_ , pages 1430–1440. PMLR, 18–24 Jul 2021.  URL https://proceedings.mlr.press/v139/chane-sane21a.html.
  * Chen et al. [2020] Ting Chen, Simon Kornblith, Kevin Swersky, Mohammad Norouzi, and Geoffrey E Hinton.  Big self-supervised models are strong semi-supervised learners.  _Advances in neural information processing systems_ , 33:22243–22255, 2020.
  * Dasari and Gupta [2021] Sudeep Dasari and Abhinav Gupta.  Transformers for one-shot visual imitation.  In _Conference on Robot Learning_ , pages 2071–2084. PMLR, 2021.
  * Dessalene et al. [2023] Eadom Dessalene, Chinmaya Devaraj, Michael Maynord, Cornelia Fermuller, and Yiannis Aloimonos.  Forecasting action through contact representations from first person video.  _IEEE Trans. Pattern Anal. Mach. Intell._ , 45(6):6703–6714, June 2023.
  * Ding et al. [2019] Yiming Ding, Carlos Florensa, Mariano Phielipp, and Pieter Abbeel.  Goal-conditioned imitation learning.  In _Proceedings of the 33rd International Conference on Neural Information Processing Systems_ , pages 15324–15335. Curran Associates Inc., Red Hook, NY, USA, December 2019.
  * Dosovitskiy et al. [2021] Alexey Dosovitskiy, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, Matthias Minderer, Georg Heigold, Sylvain Gelly, Jakob Uszkoreit, and Neil Houlsby.  An image is worth 16x16 words: Transformers for image recognition at scale.  In _International Conference on Learning Representations_ , 2021.  URL https://openreview.net/forum?id=YicbFdNTTy.
  * Du et al. [2023] Yilun Du, Sherry Yang, Bo Dai, Hanjun Dai, Ofir Nachum, Joshua B. Tenenbaum, Dale Schuurmans, and Pieter Abbeel.  Learning universal policies via text-guided video generation.  In _Thirty-seventh Conference on Neural Information Processing Systems_ , 2023.  URL https://openreview.net/forum?id=bo8q5MRcwy.
  * Duan et al. [2017] Yan Duan, Marcin Andrychowicz, Bradly Stadie, OpenAI Jonathan Ho, Jonas Schneider, Ilya Sutskever, Pieter Abbeel, and Wojciech Zaremba.  One-shot imitation learning.  _Advances in neural information processing systems_ , 30, 2017.
  * Dwibedi et al. [2019] Debidatta Dwibedi, Yusuf Aytar, Jonathan Tompson, Pierre Sermanet, and Andrew Zisserman.  Temporal Cycle-Consistency learning.  In _Computer Vision and Pattern Recognition_ , 2019.
  * Groth et al. [2021] Oliver Groth, Chia-Man Hung, Andrea Vedaldi, and Ingmar Posner.  Goal-Conditioned End-to-End visuomotor control for versatile skill primitives.  In _2021 IEEE International Conference on Robotics and Automation (ICRA)_ , pages 1319–1325. IEEE, May 2021.
  * Jaegle et al. [2022] Andrew Jaegle, Sebastian Borgeaud, Jean-Baptiste Alayrac, Carl Doersch, Catalin Ionescu, David Ding, Skanda Koppula, Daniel Zoran, Andrew Brock, Evan Shelhamer, Olivier J Henaff, Matthew Botvinick, Andrew Zisserman, Oriol Vinyals, and Joao Carreira.  Perceiver IO: A general architecture for structured inputs & outputs.  In _International Conference on Learning Representations_ , 2022.  URL https://openreview.net/forum?id=fILj7WpI-g.
  * Jain et al. [2023] Vidhi Jain, Yixin Lin, Eric Undersander, Yonatan Bisk, and Akshara Rai.  Transformers are adaptable task planners.  In Karen Liu, Dana Kulic, and Jeff Ichnowski, editors, _Proceedings of The 6th Conference on Robot Learning_ , volume 205 of _Proceedings of Machine Learning Research_ , pages 1011–1037. PMLR, 14–18 Dec 2023.  URL https://proceedings.mlr.press/v205/jain23a.html.
  * Jang et al. [2022a] Eric Jang, Alex Irpan, Mohi Khansari, Daniel Kappler, Frederik Ebert, Corey Lynch, Sergey Levine, and Chelsea Finn.  BC-Z: Zero-Shot Task Generalization with Robotic Imitation Learning.  In _Proceedings of the 5th Conference on Robot Learning_ , pages 991–1002, 2022a.
  * Jang et al. [2022b] Eric Jang, Alex Irpan, Mohi Khansari, Daniel Kappler, Frederik Ebert, Corey Lynch, Sergey Levine, and Chelsea Finn.  Bc-z: Zero-shot task generalization with robotic imitation learning.  In _Conference on Robot Learning_ , pages 991–1002. PMLR, 2022b.
  * Jiang et al. [2023] Yunfan Jiang, Agrim Gupta, Zichen Zhang, Guanzhi Wang, Yongqiang Dou, Yanjun Chen, Li Fei-Fei, Anima Anandkumar, Yuke Zhu, and Linxi Fan.  VIMA: General robot manipulation with multimodal prompts.  In _Fortieth International Conference on Machine Learning_ , 2023\.
  * Krantz et al. [2023] Jacob Krantz, Theophile Gervet, Karmesh Yadav, Austin Wang, Chris Paxton, Roozbeh Mottaghi, Dhruv Batra, Jitendra Malik, Stefan Lee, and Devendra Singh Chaplot.  Navigating to objects specified by images.  April 2023.
  * Liang et al. [2023] Jacky Liang, Wenlong Huang, Fei Xia, Peng Xu, Karol Hausman, Brian Ichter, Pete Florence, and Andy Zeng.  Code as policies: Language model programs for embodied control.  In _2023 IEEE International Conference on Robotics and Automation (ICRA)_ , 2023.
  * 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 _2018 IEEE International Conference on Robotics and Automation (ICRA)_ , pages 1118–1125. IEEE, May 2018.
  * Ma et al. [2023] Yecheng Jason Ma, Shagun Sodhani, Dinesh Jayaraman, Osbert Bastani, Vikash Kumar, and Amy Zhang.  VIP: Towards universal visual reward and representation via Value-Implicit Pre-Training.  In _International Conference on Learning Representations_ , 2023.
  * Mandi et al. [2022] Zhao Mandi, Fangchen Liu, Kimin Lee, and Pieter Abbeel.  Towards more generalizable one-shot visual imitation learning.  In _2022 International Conference on Robotics and Automation (ICRA)_ , pages 2434–2444. IEEE, 2022.
  * Mendonca et al. [2023] Russell Mendonca, Shikhar Bahl, and Deepak Pathak.  Structured world models from human videos.  2023\.
  * Miech et al. [2019] Antoine Miech, Dimitri Zhukov, Jean-Baptiste Alayrac, Makarand Tapaswi, Ivan Laptev, and Josef Sivic.  HowTo100M: Learning a Text-Video Embedding by Watching Hundred Million Narrated Video Clips.  In _ICCV_ , 2019.
  * Nair et al. [2018] Ashvin Nair, Vitchyr Pong, Murtaza Dalal, Shikhar Bahl, Steven Lin, and Sergey Levine.  Visual reinforcement learning with imagined goals.  In _NeurIPS_ , July 2018.
  * Nair et al. [2022] Suraj Nair, Aravind Rajeswaran, Vikash Kumar, Chelsea Finn, and Abhinav Gupta.  R3m: A universal visual representation for robot manipulation.  In _6th Annual Conference on Robot Learning_ , 2022.  URL https://openreview.net/forum?id=tGbpgz6yOrI.
  * Open X-Embodiment Collaboration et al. [2023] Open X-Embodiment Collaboration, Abhishek Padalkar, Acorn Pooley, Ajay Mandlekar, Ajinkya Jain, Albert Tung, Alex Bewley, Alex Herzog, Alex Irpan, Alexander Khazatsky, Anant Rai, Anikait Singh, Animesh Garg, Anthony Brohan, Antonin Raffin, Ayzaan Wahid, Ben Burgess-Limerick, Beomjoon Kim, Bernhard Schölkopf, Brian Ichter, Cewu Lu, Charles Xu, Chelsea Finn, Chenfeng Xu, Cheng Chi, Chenguang Huang, Christine Chan, Chuer Pan, Chuyuan Fu, Coline Devin, Danny Driess, Deepak Pathak, Dhruv Shah, Dieter Büchler, Dmitry Kalashnikov, Dorsa Sadigh, Edward Johns, Federico Ceola, Fei Xia, Freek Stulp, Gaoyue Zhou, Gaurav S Sukhatme, Gautam Salhotra, Ge Yan, Giulio Schiavi, Gregory Kahn, Hao Su, Hao-Shu Fang, Haochen Shi, Heni Ben Amor, Henrik I Christensen, Hiroki Furuta, Homer Walke, Hongjie Fang, Igor Mordatch, Ilija Radosavovic, Isabel Leal, Jacky Liang, Jad Abou-Chakra, Jaehyung Kim, Jan Peters, Jan Schneider, Jasmine Hsu, Jeannette Bohg, Jeffrey Bingham, Jiajun Wu, Jialin Wu, Jianlan Luo, Jiayuan Gu, Jie Tan, Jihoon Oh, Jitendra Malik, Jonathan Booher, Jonathan Tompson, Jonathan Yang, Joseph J Lim, João Silvério, Junhyek Han, Kanishka Rao, Karl Pertsch, Karol Hausman, Keegan Go, Keerthana Gopalakrishnan, Ken Goldberg, Kendra Byrne, Kenneth Oslund, Kento Kawaharazuka, Kevin Zhang, Krishan Rana, Krishnan Srinivasan, Lawrence Yunliang Chen, Lerrel Pinto, Li Fei-Fei, Liam Tan, Lionel Ott, Lisa Lee, Masayoshi Tomizuka, Max Spero, Maximilian Du, Michael Ahn, Mingtong Zhang, Mingyu Ding, Mohan Kumar Srirama, Mohit Sharma, Moo Jin Kim, Naoaki Kanazawa, Nicklas Hansen, Nicolas Heess, Nikhil J Joshi, Niko Suenderhauf, Norman Di Palo, Nur Muhammad Mahi Shafiullah, Oier Mees, Oliver Kroemer, Pannag R Sanketi, Paul Wohlhart, Peng Xu, Pierre Sermanet, Priya Sundaresan, Quan Vuong, Rafael Rafailov, Ran Tian, Ria Doshi, Roberto Martín-Martín, Russell Mendonca, Rutav Shah, Ryan Hoque, Ryan Julian, Samuel Bustamante, Sean Kirmani, Sergey Levine, Sherry Moore, Shikhar Bahl, Shivin Dass, Shubham Sonawani, Shuran Song, Sichun Xu, Siddhant Haldar, Simeon Adebola, Simon Guist, Soroush Nasiriany, Stefan Schaal, Stefan Welker, Stephen Tian, Sudeep Dasari, Suneel Belkhale, Takayuki Osa, Tatsuya Harada, Tatsuya Matsushima, Ted Xiao, Tianhe Yu, Tianli Ding, Todor Davchev, Tony Z Zhao, Travis Armstrong, Trevor Darrell, Vidhi Jain, Vincent Vanhoucke, Wei Zhan, Wenxuan Zhou, Wolfram Burgard, Xi Chen, Xiaolong Wang, Xinghao Zhu, Xuanlin Li, Yao Lu, Yevgen Chebotar, Yifan Zhou, Yifeng Zhu, Ying Xu, Yixuan Wang, Yonatan Bisk, Yoonyoung Cho, Youngwoon Lee, Yuchen Cui, Yueh-Hua Wu, Yujin Tang, Yuke Zhu, Yunzhu Li, Yusuke Iwasawa, Yutaka Matsuo, Zhuo Xu, and Zichen Jeff Cui.  Open X-Embodiment: Robotic learning datasets and RT-X models.  _arXiv 2310.08864_ , 2023.
  * Parashar et al. [2023] Priyam Parashar, Vidhi Jain, Xiaohan Zhang, Jay Vakil, Sam Powers, Yonatan Bisk, and Chris Paxton.  Spatial-Language Attention Policies for Efficient Robot Learning.  In _Conference on Robot Learning_ , 2023.  URL https://arxiv.org/abs/2304.11235.
  * Peng et al. [2018] Xue Bin Peng, Angjoo Kanazawa, Jitendra Malik, Pieter Abbeel, and Sergey Levine.  SFV: reinforcement learning of physical skills from videos.  _ACM Trans. Graph._ , 37(6):1–14, December 2018\.
  * Perez et al. [2018] Ethan Perez, Florian Strub, Harm De Vries, Vincent Dumoulin, and Aaron Courville.  Film: Visual reasoning with a general conditioning layer.  In _Proceedings of the AAAI conference on artificial intelligence_ , 2018.
  * Petrík et al. [2021] Vladimír Petrík, Makarand Tapaswi, Ivan Laptev, and Josef Sivic.  Learning object manipulation skills via approximate state estimation from real videos.  In Jens Kober, Fabio Ramos, and Claire Tomlin, editors, _Proceedings of the 2020 Conference on Robot Learning_ , volume 155 of _Proceedings of Machine Learning Research_ , pages 296–312. PMLR, 2021.
  * Pirk et al. [2019] Sören Pirk, Mohi Khansari, Yunfei Bai, Corey Lynch, and Pierre Sermanet.  Online object representations with contrastive learning.  2019\.
  * Schmeckpeper et al. [2020] Karl Schmeckpeper, Oleh Rybkin, Kostas Daniilidis, Sergey Levine, and Chelsea Finn.  Reinforcement learning with videos: Combining offline observations with interaction.  November 2020.
  * Shah et al. [2023] Rutav Shah, Roberto Martín-Martín, and Yuke Zhu.  MUTEX: Learning unified policies from multimodal task specifications.  In _7th Annual Conference on Robot Learning_ , 2023.
  * Sharma et al. [2019] Pratyusha Sharma, Deepak Pathak, and Abhinav Gupta.  Third-Person visual imitation learning via decoupled hierarchical controller.  _Adv. Neural Inf. Process. Syst._ , 32, 2019.
  * Smith et al. [2020] Laura Smith, Nikita Dhawan, Marvin Zhang, Pieter Abbeel, and Sergey Levine.  AVID: Learning Multi-Stage tasks via Pixel-Level translation of human videos.  In _Robotics: Science and Systems (RSS)_ , December 2020.
  * Stone et al. [2023] Austin Stone, Ted Xiao, Yao Lu, Keerthana Gopalakrishnan, Kuang-Huei Lee, Quan Vuong, Paul Wohlhart, Sean Kirmani, Brianna Zitkovich, Fei Xia, Chelsea Finn, and Karol Hausman.  Open-World object manipulation using pre-trained Vision-Language models.  March 2023.
  * Xiao et al. [2023] Tete Xiao, Ilija Radosavovic, Trevor Darrell, and Jitendra Malik.  Masked visual pre-training for motor control.  _arXiv preprint arXiv:2203.06173_ , 2023.
  * Yenamandra et al. [2023] Sriram Yenamandra, Arun Ramachandran, Karmesh Yadav, Austin Wang, Mukul Khanna, Theophile Gervet, Tsung-Yen Yang, Vidhi Jain, Alexander William Clegg, John Turner, Zsolt Kira, Manolis Savva, Angel Chang, Devendra Singh Chaplot, Dhruv Batra, Roozbeh Mottaghi, Yonatan Bisk, and Chris Paxton.  HomeRobot: Open-Vocabulary Mobile Manipulation.  In _Conference on Robot Learning_ , 2023.  URL https://arxiv.org/abs/2306.11565.
  * Yu et al. [2022] Jiahui Yu, Zirui Wang, Vijay Vasudevan, Legg Yeung, Mojtaba Seyedhosseini, and Yonghui Wu.  Coca: Contrastive captioners are image-text foundation models.  _arXiv preprint arXiv:2205. 01917_ , 2022.
  * Zakka et al. [2022] Kevin Zakka, Andy Zeng, Pete Florence, Jonathan Tompson, Jeannette Bohg, and Debidatta Dwibedi.  Xirl: Cross-embodiment inverse reinforcement learning.  In _Conference on Robot Learning_ , pages 537–546. PMLR, 2022.
  * Zhai et al. [2023] Xiaohua Zhai, Basil Mustafa, Alexander Kolesnikov, and Lucas Beyer.  Sigmoid loss for language image pre-training.  In _2023 IEEE/CVF International Conference on Computer Vision (ICCV)_ , pages 11941–11952, Los Alamitos, CA, USA, oct 2023. IEEE Computer Society.  doi: 10.1109/ICCV51070.2023.01100.
  * Zhao et al. [2023] Tony Z Zhao, Vikash Kumar, Sergey Levine, and Chelsea Finn.  Learning Fine-Grained bimanual manipulation with Low-Cost hardware.  In _Robotics: Science and Systems_ , 2023.


# SAFIN: Arbitrary Style Transfer with Self-Attentive Factorized Instance
Normalization

###### Abstract

Artistic style transfer aims to transfer the style characteristics of one
image onto another image while retaining its content. Existing approaches
commonly leverage various normalization techniques, although these face
limitations in adequately transferring diverse textures to different spatial
locations. Self-Attention based approaches have tackled this issue with
partial success but suffer from unwanted artifacts. Motivated by these
observations, this paper aims to combine the best of both worlds: self-
attention and normalization. That yields a new plug-and-play module that we
name Self-Attentive Factorized Instance Normalization (SAFIN). SAFIN is
essentially a spatially-adaptive normalization module whose parameters are
inferred through attention on the content and style image. We demonstrate that
plugging SAFIN into the base network of another state-of-the-art method
results in enhanced stylization. We also develop a novel base network composed
of Wavelet Transform for multi-scale style transfer, which when combined with
SAFIN, produces visually appealing results with lesser unwanted textures.

##  1 Introduction

Artistic style transfer is the task of extracting style and texture patterns
from one image and transferring them onto the content of another. Initial work
by [3] employed an optimization based approach. This approach was replaced by
feed-forward networks that could generate the images in a single forward pass
[4] [5], with the major limitation being that these had to be retrained for
every new style image. Further work on doing multiple-styles-per-model
followed [6], though these could not effectively generalize to new unseen
styles.

Further research led to models that perform style transfer for arbitrary
styles [2] [7]. Adaptive Instance Normalization (AdaIN) [2] was one of the
enabling techniques, which adjusted the mean and variance of the content image
feature-maps to match those of the style image. The transferred results,
though often decent, also appear to be over-simplistic and lack fine details,
since only simple statistic moments were drawn for content-style matching.
Following this idea, many other normalization based approaches such as Batch-
IN [9] and Dynamic-IN [8] were proposed.

Fig. 1: Stylization results. A few stylization results generated using SAFIN
with our base network.

Despite much progress being made on normalization-based techniques, one
remaining challenge lies in the lack of ability to transfer different textures
to different spatial locations of the image with semantic awareness. Although
spatially adaptive normalization has been recently used to good effect in
other tasks [17, 1], similar ideas have not been explored in style transfer
yet.

Attention-based models [10] [11] have shown another successful stream in style
transfer. They attempt to resolve this limitation by employing self-attention
over the spatial dimensions when transforming the content features. However,
these sometimes cause unwanted patterns in output (Fig. 3).

Motivated by these observations, this paper aims to combine the best of both
worlds: self-attention and normalization. Our underlying motivation is to make
normalization more spatially flexible and semantically aware. We achieve this
by using a self-attention based module to learn to generate the parameters of
our spatially-adaptive normalization module. We further introduce a new
factorized parameter structure to the normalization layer. The resulting new
module, called Self-Attentive Factorized Instance Normalization (SAFIN), has a
plug-and-play nature and can boost other state-of-the-art normalization-based
style transfer methods. In addition, we develop a novel base network composed
of Wavelet Transform for multi-scale style transfer. When combined with SAFIN,
visually appealing stylization results can be obtained without additional
artifacts and distortions.

##  2 Related Work

###  2.1 Normalization in NST

The essential task in Neural Style Transfer (NST) is the transformation of
content-image features in a way that their style characteristics match those
of the style image.

In optimization based methods (i.e. single-style-per-model) [5] [4], the use
of batch-normalization (BN) [13] was purely for it’s standard purpose of
easing neural network training. The work by [14] introduced instance-
normalization (IN) and its advantage over batch normalization due to it’s
ability to discard instance-specific contrast information. With new models
performing multiple-styles-per-model or arbitrary-styles-per-model,
normalization (i.e affine shifting and scaling of mean-variance-normalized
content features) became a way of imbuing the content features with the style
characteristics of the input style. For instance, distinct scaling/shifting
parameters are learned for each new style during training in Conditional-IN
[6], thus allowing the model to achieve stylization for multiple styles.

The main task in accomplishing arbitrary style transfer using the
normalization based approach is to compute the normalization parameters at
test time. AdaIN [2] showed that even parameters as simple as the channel-wise
mean and variance of the style-image features could be effective. Batch-IN [9]
used a combination of BN and IN whereas Dynamic-IN [8] used a learned
convolutional network to extract the scaling and shifting parameters given the
style and content images.

Fig. 2: Semantic awareness. Attention-based models show some capability of
matching patterns to decide what texture to transfer to different spatial
locations. SAFIN has no obvious spillover, whereas WCT and AdaIN suffer from
spillovers.

A major limitation of the current normalization based approaches is their
inability to transfer diverse and fine stroke patterns from different parts of
the style image to the semantically similar regions in the content image (Fig
2). This is because the normalization is either not spatially adaptive [2] [9]
or lacks an appropriate mechanism to infer and transfer the diversity in style
patterns onto the normalization parameters.

###  2.2 Other Techniques

WCT [7] performs the transformation on the content features by using the
whitening and coloring transform. Not only does WCT suffer from transfer of
global style characteristics, leading to less varied stroke patterns, but it
is also computationally expensive and especially slow and memory consuming for
high-resolution images.

Initial attempts to solve the problem of transferring different style patterns
used a patch-based approach by matching local statistics between semantically
matching patches [18] [19]. Though these methods achieve visually pleasing
results, their application is limited to content-style combinations with dense
semantic correspondence.

###  2.3 Attention in NST

Self-Attention [20] has been applied to the problem of style transfer in SANet
[10] since it suggests a potential way to designate different style patterns
to different regions in an image. However, in SANet, the fine-grained nature
of the attention target and direct incorporation within the content feature
maps without further processing sometimes leads to extraneous artifacts and
distortions (Fig. 3).

Fig. 3: Extraneous artifacts. SANet produces eye-like artifacts (zoom in for
greater clarity) due to the fine-grained nature of its attention target while
SAFIN does not.

Our method builds upon the ability of SANet and attempts to avert its
drawbacks by using it in a novel manner. We use the self-attention module as a
part of a convolutional network trained to infer the normalization parameters.
This is in contrast to SANet, which directly uses the attention output to
infer the feature map for the stylized output.

##  3 Approach

Fig. 4: Network Architecture. We train a multi-scale encoder-decoder network
with two SAFIN modules and Wavelet Transform to stylize the low frequency
component (LL, shown through bold lines) as well as the high frequency
components (LH, HL, and HH, shown through dotted lines). AdaIN [2] is used at
different scales to further enrich the stylization.

###  3.1 Factorized Instance Normalization (FIN)

We present a novel normalization module called Factorized Instance
Normalization (FIN). Given the content feature maps
$F_{c}\in\mathbb{R}^{H*W*C}$ we transform the style characteristics of these
features to match the style characteristics of the style image while
maintaining it’s semantic information. We achieve this via the following
transformation:

$\displaystyle\text{FIN}(F_{c})$
$\displaystyle=\gamma_{s}\odot\big{(}\gamma_{ind}\odot\overline{F_{c}}+\beta_{ind}\big{)}+\beta_{s}$
(1)

where $\overline{F_{c}}\in\mathbb{R}^{H*W*C}$ is the normalized content
feature obtained by applying channel-wise instance normalization on $F_{c}$.
$\odot$ denotes element-wise multiplication.

For the shifting and scaling operations (Eq. 1), we use two types of learned
normalization parameters:

  1. 1.

$\gamma_{ind}$ and $\beta_{ind}$ represent the style independent parameters
and are single dimension vectors along the channel dimension (i.e.
$\gamma_{ind},\beta_{ind}\in\mathbb{R}^{C}$). These parameters are meant to
capture the shared transformation between the different style images, by
transforming $\overline{F_{c}}$ before the style dependent transformation.

  2. 2.

On the other hand, $\gamma_{s}$ and $\beta_{s}$ are inferred from the input
style image. $\gamma_{s}$, $\beta_{s}$ are 3 dimensional; along the channel
dimension as well as the spatial dimensions (i.e.
$\gamma_{s},\beta_{s}\in\mathbb{R}^{H*W*C}$). Thus, element-wise shifting and
scaling of the content features using $\gamma_{s}$ and $\beta_{s}$ allows for
greater flexibility in transferring diverse textures to different spatial
regions on the content image. However, to generate rich stylization, it is not
enough to simply use spatially adaptive normalization, since a mechanism is
required that can adequately learn to derive these normalization-parameters
from the input style image. Thus, $\gamma_{s}$ and $\beta_{s}$ are inferred
from the style image using a network employing self-attention followed by
convolution. This mechanism is explained in subsection 3.2.

###  3.2 Self-Attention Mechanism (SA)

We employ Self-Attention [20] to capture the semantic correspondence between
the content feature maps $F_{c}$ and the style feature maps $F_{s}$ (extracted
from the encoder) for generating the style-based normalization parameters
$\gamma_{s}$ and $\beta_{s}$.

$\displaystyle\gamma_{s}=\textsc{ReLU}(W_{\gamma}\otimes\text{SA}(\overline{F_{c}},\overline{F_{s}}))$
(2)
$\displaystyle\beta_{s}=\textsc{ReLU}(W_{\beta}\otimes\text{SA}(\overline{F_{c}},\overline{F_{s}}))$
(3)

Here $\overline{F_{c}},\overline{F_{s}}\in\mathbb{R}^{H*W*C}$ are channel-wise
instance normalized content and style feature maps respectively. $W_{\gamma}$,
$W_{\beta}$ represent $1*1$ convolutional layers and $\otimes$ denotes the
convolution operation. Our Self-Attention module (SA) is similar to that used
by [10], though it is employed for a different purpose. [10] used self-
attention to directly modify the content feature maps, whereas we instead use
it to generate the style-dependent normalization parameters of our FIN module.

###  3.3 Wavelet Transform

For better preservation of content features in the stylized outputs, we use
Wavelet Pooling (to replace Max Pooling) and Wavelet Unpooling (to replace
Upsampling) [24], a pooling technique which allows exact reconstruction due to
it’s invertible nature. This technique has been previously employed in photo-
realistic style transfer [12], but not, to the best of our knowledge, in
artistic style transfer.

Wavelet pooling uses four distinct pooling kernels
$\\{LL^{\intercal},LH^{\intercal},HL^{\intercal},HH^{\intercal}\\}$ at every
pooling layer, made by combining the high-pass filter
$H^{\intercal}=\frac{1}{\sqrt{2}}[\hskip 2.84526pt-1\hskip 8.53581pt1\hskip
2.84526pt]$ and the low-pass filter $L^{\intercal}=\frac{1}{\sqrt{2}}[\hskip
2.84526pt1\hskip 8.53581pt1\hskip 2.84526pt]$ . We use the terms LL, LH, HL,
and HH to refer to the output of the respective kernel.

###  3.4 Network Architecture

Figure 4 shows the network architecture, which follows the commonly employed
encoder-decoder structure [2], [7], [10]. The encoder has the same structure
as the VGG-19 network, and the decoder is symmetric to the encoder.

The loss function for training the SAFIN module and the decoder is given as:
$L=L_{c}+\lambda_{s}*L_{s}$, where $L_{c}$ and $L_{s}$ are content and style
losses and $\lambda_{s}$ is the weight assigned to the style loss relative to
the content loss.

The content loss is the Euclidean distance between LL components of stylised
output image encoding $\mathcal{E}(I_{cs})$ and content image encoding
$\mathcal{E}(I_{c})$ extracted from the layer ReLU_4_1:
$L_{c}=||\mathcal{E}_{LL}^{4.1}(I_{cs})-\mathcal{E}_{LL}^{4.1}(I_{c})||_{2}$

The style loss is the Euclidean distance between mean and variance of LL
components of stylised output image encoding $\mathcal{E}(I_{cs})$ and style
image encoding $\mathcal{E}(I_{s})$ from the layers ReLU_X_1 where
$\text{X}\in\\{1,2,3,4\\}$:

$\displaystyle L_{s}=$
$\displaystyle\sum\limits_{k=1}^{4}||\mu(\mathcal{E}_{LL}^{k.1}(I_{cs}))-\mu(\mathcal{E}_{LL}^{k.1}(I_{s}))||_{2}+$
$\displaystyle\hskip
85.35826pt||\sigma(\mathcal{E}_{LL}^{k.1}(I_{cs}))-\sigma(\mathcal{E}_{LL}^{k.1}(I_{s}))||_{2}$

Further details of the network architecture are described in the appendix.

##  4 Experiments

We trained our network using MS-COCO dataset [15] for content images and
WikiArt dataset [16] for style images. The detailed experimental settings are
included in the appendix. Code and pretrained model are available here
https://github.com/Aaditya-Singh/SAFIN.

###  4.1 Generalizability

(a) SAFIN module when substituted in place of the SANet module in SANet’s base
network, leads to reduced unwanted distortions.

(b) Training loss comparison: SANet’s base network with SANet’s style transfer
module vs SAFIN module.

Fig. 5: SAFIN (with SANet’s base network) vs SANet.

To demonstrate the superiority as well as flexibility of our spatially
adaptive normalization technique, we show how SAFIN module when substituted in
to replace the SANet [10] module within the SANet base network, also mitigates
the extraneous artifacts to some extent (Fig. 5(a)). The training-loss
comparison plot (Fig. 5(b)) shows a lower content loss with the SAFIN module
compared to the SANet module thus demonstrating that the usage of SAFIN
results in better content preservation. Style losses are similar. Indeed, our
module when inserted into SANet’s base network provides a comparable quality
of stylization along with reduced unwanted distortions.

###  4.2 Ablation Study

####  4.2.1 Effect of removal of Self-Attention from SAFIN

In order to understand the effectiveness of Self-Attention (SA) mechanism
within the overall SAFIN module, we compare the results when SAFIN module is
replaced with FIN module (normalization without self-attention) in our
network. From figure 6, it can be observed that while FIN module enables the
transfer of general color and texture from the style image, self-attention
allows the network to transfer the style characteristics from semantically
similar regions in the style image leading to an enhanced stylization.

Further ablation studies and more stylization outputs for comparison are
included in the appendix.

Fig. 6: Effect of removal of Self-Attention (SA) module. FIN transfers the
general color and texture patterns but causes spillovers, whereas SAFIN
produces a more semantically coherent stylization.

###  4.3 Quantitative Analysis

Fig. 7: User Survey Results. SAFIN with our base network causes the least
amount of artifacts and distortions without compromising the stylization
quality.

We also conducted an extensive survey for a thorough quantitative analysis.
There were 25 participants and 40 sets of content-style image pairs along with
their stylization outputs using 4 different methods - SAFIN (with the base
network), SANet, WCT and AdaIN. For every set, each participant was asked two
questions:

  * •

To choose the best stylization, i.e. which method leads to the best color and
texture transfer.

  * •

To choose which stylization caused the least amount of unwanted artifacts and
distortions.

The results (Fig. 7) show that our method is preferred by users against other
methods, providing better or comparable stylization without leading to
additional distortions.

##  5 Conclusion

In this paper, we propose a novel module called SAFIN which combines
normalization and self-attention in a unique way to mitigate the issues with
previous style transfer techniques. Extensive experiments demonstrate its
effectiveness. We show the generalization capability of the proposed module by
showing how it improves the stylization and losses of another state-of-the-art
method when incorporated within its base network. We also develop a novel base
network composed of Wavelet Transform for multi-scale style transfer that
helps in preserving more content information. Experimental results illustrate
the effectiveness of the proposed method.

## References

  * [1] Taesung Park, Ming-Yu Liu, Ting-Chun Wang, and Jun-Yan Zhu,  “Semantic image synthesis with spatially-adaptive normalization,”  in Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2019, pp. 2337–2346.
  * [2] Xun Huang and Serge Belongie,  “Arbitrary style transfer in real-time with adaptive instance normalization,”  in Proceedings of the IEEE International Conference on Computer Vision, 2017, pp. 1501–1510.
  * [3] Leon A Gatys, Alexander S Ecker, and Matthias Bethge,  “Image style transfer using convolutional neural networks,”  in Proceedings of the IEEE conference on computer vision and pattern recognition, 2016, pp. 2414–2423.
  * [4] Justin Johnson, Alexandre Alahi, and Li Fei-Fei,  “Perceptual losses for real-time style transfer and super-resolution,”  in European conference on computer vision. Springer, 2016, pp. 694–711.
  * [5] Dmitry Ulyanov, Vadim Lebedev, Andrea Vedaldi, and Victor S Lempitsky,  “Texture networks: Feed-forward synthesis of textures and stylized images.,”  in ICML, 2016, vol. 1, p. 4.
  * [6] Vincent Dumoulin, Jonathon Shlens, and Manjunath Kudlur,  “A learned representation for artistic style,”  arXiv preprint arXiv:1610.07629, 2016.
  * [7] Yijun Li, Chen Fang, Jimei Yang, Zhaowen Wang, Xin Lu, and Ming-Hsuan Yang,  “Universal style transfer via feature transforms,”  in Advances in neural information processing systems, 2017, pp. 386–396.
  * [8] Yongcheng Jing, Xiao Liu, Yukang Ding, Xinchao Wang, Errui Ding, Mingli Song, and Shilei Wen,  “Dynamic instance normalization for arbitrary style transfer,”  arXiv preprint arXiv:1911.06953, 2019.
  * [9] Hyeonseob Nam and Hyo-Eun Kim,  “Batch-instance normalization for adaptively style-invariant neural networks,”  in Advances in Neural Information Processing Systems, 2018, pp. 2558–2567.
  * [10] Dae Young Park and Kwang Hee Lee,  “Arbitrary style transfer with style-attentional networks,”  in Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2019, pp. 5880–5888.
  * [11] Yuan Yao, Jianqiang Ren, Xuansong Xie, Weidong Liu, Yong-Jin Liu, and Jun Wang,  “Attention-aware multi-stroke style transfer,”  in Proceedings of the IEEE Conference on Computer Vision and Pattern Recognition, 2019, pp. 1467–1475.
  * [12] Jaejun Yoo, Youngjung Uh, Sanghyuk Chun, Byeongkyu Kang, and Jung-Woo Ha,  “Photorealistic style transfer via wavelet transforms,”  in Proceedings of the IEEE International Conference on Computer Vision, 2019, pp. 9036–9045.
  * [13] Sergey Ioffe and Christian Szegedy,  “Batch normalization: Accelerating deep network training by reducing internal covariate shift,”  arXiv preprint arXiv:1502.03167, 2015.
  * [14] Dmitry Ulyanov, Andrea Vedaldi, and Victor Lempitsky,  “Instance normalization: The missing ingredient for fast stylization,”  arXiv preprint arXiv:1607.08022, 2016.
  * [15] Tsung-Yi Lin, Michael Maire, Serge Belongie, James Hays, Pietro Perona, Deva Ramanan, Piotr Dollár, and C Lawrence Zitnick,  “Microsoft coco: Common objects in context,”  in European conference on computer vision. Springer, 2014, pp. 740–755.
  * [16] Fred Phillips and Brandy Mackintosh,  “Wiki art gallery, inc.: A case for critical thinking,”  Issues in Accounting Education, vol. 26, no. 3, pp. 593–608, 2011\.
  * [17] Jiashun Wang, Chao Wen, Yanwei Fu, Haitao Lin, Tianyun Zou, Xiangyang Xue, and Yinda Zhang,  “Neural pose transfer by spatially adaptive instance normalization,”  in Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2020, pp. 5831–5839.
  * [18] YiChang Shih, Sylvain Paris, Connelly Barnes, William T Freeman, and Frédo Durand,  “Style transfer for headshot portraits,”  2014\.
  * [19] Jing Liao, Yuan Yao, Lu Yuan, Gang Hua, and Sing Bing Kang,  “Visual attribute transfer through deep image analogy,”  arXiv preprint arXiv:1705.01088, 2017.
  * [20] Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Lukasz Kaiser, and Illia Polosukhin,  “Attention is all you need,”  in Advances in neural information processing systems, 2017, pp. 5998–6008.
  * [21] Dzmitry Bahdanau, Kyunghyun Cho, and Yoshua Bengio,  “Neural machine translation by jointly learning to align and translate,”  arXiv preprint arXiv:1409.0473, 2014.
  * [22] Kelvin Xu, Jimmy Ba, Ryan Kiros, Kyunghyun Cho, Aaron Courville, Ruslan Salakhudinov, Rich Zemel, and Yoshua Bengio,  “Show, attend and tell: Neural image caption generation with visual attention,”  in International conference on machine learning, 2015, pp. 2048–2057.
  * [23] Zichao Yang, Xiaodong He, Jianfeng Gao, Li Deng, and Alex Smola,  “Stacked attention networks for image question answering,”  in Proceedings of the IEEE conference on computer vision and pattern recognition, 2016, pp. 21–29.
  * [24] Jin-Sheng Guf and Wei-Sun Jiang,  “The haar wavelets operational matrix of integration,”  International Journal of Systems Science, vol. 27, no. 7, pp. 623–628, 1996.