# Generating Stable and Collision-Free Policies through Lyapunov Function
Learning

Alexandre Coulombe1 and Hsiu-Chin Lin2 1Alexandre Coulombe is with the
Department of Electrical and Computer Engineering, McGill University, Canada
<EMAIL_ADDRESS>Lin is with the School of Computer
Science and the Department of Electrical and Computer Engineering, McGill
University, Canada<EMAIL_ADDRESS>work is sponsored by Mitacs
Accelerate IT23788.

###### Abstract

The need for rapid and reliable robot deployment is on the rise. Imitation
Learning (IL) has become popular for producing motion planning policies from a
set of demonstrations. However, many methods in IL are not guaranteed to
produce stable policies. The generated policy may not converge to the robot
target, reducing reliability, and may collide with its environment, reducing
the safety of the system. Stable Estimator of Dynamic Systems (SEDS)produces
stable policies by constraining the Lyapunov stability criteria during
learning, but the Lyapunov candidate function had to be manually selected. In
this work, we propose a novel method for learning a Lyapunov function and a
policy using a single neural network model. The method can be equipped with an
obstacle avoidance module for convex object pairs to guarantee no collisions.
We demonstrated our method is capable of finding policies in several
simulation environments and transfer to a real-world scenario.

###### Index Terms:

Imitation Learning, Lyapunov stability, Obstacle Avoidance, Motion Planning,
Neural networks

##  I INTRODUCTION

With the growing number of robots in industry and in human environments, the
need for ease of deployment and safety in robotics becomes increasingly
apparent. Robots are skilled in routine and repetitive tasks. However, it is
hard to handle a new task without tedious modelling, designing, and
programming. Furthermore, safety is crucial to deploying robotic systems since
a mistake may injure the human or damage the robot itself. This is
particularly important when robots need to maneuver in environments with
obstacles, such as the example in Figure 1.

Trajectory optimization is a well-studied method for motion planning. The
method finds a trajectory by minimizing some objective functions while
satisfying its physical limits [1, 2]. In the presence of obstacles, collision
avoidance can be achieved by imposing a minimum distance between the robot and
the obstacles as inequality constraints [3, 4]. While trajectory optimization
is popular for motion planning, the solution is only valid for a specific pair
of initial and target states. Re-planning is a time-consuming process;
therefore, it is not suitable for real-time planning. Additionally, the
planned path is not immune to perturbations.

In recent years, Imitation Learning (IL), the data-driven approach for motion
planning, has shown promising results. IL learns a model to predict the
desired action given the current input, and therefore, produces a global
solution as opposed to trajectory optimization. Many machine learning
approaches have been applied, including supervised learning [5, 6], Inverse
Reinforcement Learning (IRL) [7], and Generative Adversarial Imitation
Learning (GAIL) [8, 9]. However, these methods have no guarantees that the
policy will be stable when deployed.

(a) Initial Position

(b) Policy Rollout

(c) Goal Position

Figure 1: Mecademic Meca500 performing the shelf manipulation task: the robot
(a) starts from an initial position, (b) follows the policy rollout, and (c)
reaches the goal position.

Stability is crucial for robotic systems. Prior work proposed a data-driven
approach with stability constraints to generate stable policies [10, 11].
However, [10, 11] relies on a chosen Lyapunov candidate function to be valid
for the demonstrations. Another prior work constrains the Lyapunov stability
conditions on neural networks and learns the Lyapunov function with the policy
from training data using mixed-integer programming [12]. On the Reinforcement
Learning (RL) front, prior work introduces the stability constraint during
training by widening the region of attraction of a Lyapunov function [13].
However, the challenge is finding a suitable Lyapunov function.

In the presence of obstacles, methods such as virtual potential fields can
modulate the policy to repel the robot from obstacles [14]. However, these
create local minima where the robot is stuck. Recent work avoids extrema by
modulating space around obstacles so that the robot moves around them, but it
only works for point representations of the robot [15, 16, 17].

In this paper, we propose a novel method for learning both the Lyapunov
function and policy with a single neural network. To produce a stable policy,
we train the network using a constrained optimization formulation with
stability conditions as constraints. We deploy the trained policy with an
obstacle avoidance module, augmented by our method to deal with convex
objects. The proposed work is validated in simulation and robotic hardware
with direct sim-to-real transfer. The main contributions include:

  * •

We extend [10] by learning a Lyapunov function from demonstrations instead of
manually choosing a Lyapunov candidate function (Sec. III-B).

  * •

We extend the obstacle avoidance work in [15, 16, 17] to treat robots and
obstacles as convex objects (Sec. III-C).

  * •

We show that policies developed in simulation can be transferred to a real-
world system reliably (Sec. IV).

##  II BACKGROUND

In this section, background knowledge for Lyapunov stability theory and
collision avoidance techniques are provided.

###  II-A Lyapunov Function

From [18], Lyapunov stability dictates that a system with $d$ degree-of-
freedom is locally stable if, for a region around the equilibrium
$\mathbf{x}_{g}\in\mathbb{R}^{d}$, there is a Lyapunov function
$V(\cdot):\mathbb{R}^{d}\rightarrow\mathbb{R}$ that satisfies the conditions:

$\displaystyle V(\mathbf{x}[n])>0,$ $\displaystyle\ \forall\mathbf{x}[n]\in
X,\mathbf{x}[n]\neq\mathbf{x}_{g}$ (1) $\displaystyle
V(\mathbf{x}[n+1])\leq(1-\epsilon)V(\mathbf{x}[n]),$ $\displaystyle\
\forall\mathbf{x}[n]\in X,\mathbf{x}[n]\neq\mathbf{x}_{g}$ (2) $\displaystyle
V(\mathbf{x}_{g})=0$ (3)

where $\epsilon>0$ is a positive scalar and
$X=\\{\mathbf{x}[0]\;|\;V(\mathbf{x}[0])\leq\rho\\}$ is the stable region
around $\mathbf{x}_{g}$, where $\rho>0$ is some positive value. (1) and (3)
specify that $V(\cdot)$ must output a positive value for all positions, except
for $0$ at $\mathbf{x}_{g}$. (2) states that the system must evolve to
positions with lower values, with the lowest value at $\mathbf{x}_{g}$. Thus,
the system will converge to the equilibrium point $\mathbf{x}_{g}$ if a
Lyapunov function can be found.

A common Lyapunov candidate function is the quadratic Lyapunov function

$V(\mathbf{x}[n])=\|\mathbf{x}[n]-\mathbf{x}_{g}\|_{2}$ (4)

The above Lyapunov candidate function imposes constraints on the movement
towards the goal. However, this may not be ideal in some scenarios, such as
moving around obstacles. An example is shown in Figure 2. The quadratic
Lyapunov candidate function (orange) violates the constraint in (2), since
avoiding the obstacles implies moving away from the goal, while a valid
Lyapunov function (blue) is monotonically decreasing along the demonstration
path (red).

Instead of explicitly specifying the Lyapunov candidate function, prior work
[12] learns a function that satisfies the Lyapunov conditions in (1)-(3). In
this work, we will adapt this structure to learn the Lyapunov function.

(a) Demonstrations with an initial (orange) and a goal (green) states while
avoiding obstacles (black).

(b) Outputs of a quadratic Lyapunov function (orange) and a valid Lyapunov
function (blue) for the demonstrations.

Figure 2: An example of a trajectory and its Lyapunov function outputs.

###  II-B Collision Avoidance

Given the distance between the robot and the obstacle, a policy
$\mathbf{\pi}\left(\mathbf{x}[n]\right)$ can be modulated to move around the
obstacles [15, 16, 17] with the following formulation:

$\displaystyle\dot{\mathbf{x}}[n]$
$\displaystyle=\mathbf{M}\left(\mathbf{x}[n]\right)\mathbf{\pi}\left(\mathbf{x}[n]\right)$
(5)

where $\mathbf{M}\left(\mathbf{x}[n]\right)$ is the modulation matrix of the
obstacle. The modulation matrix conserves the existing extrema of the system
dynamics, and does not introduce new extrema while the matrix is full rank
[16]. The modulation matrix $\mathbf{M}\left(\mathbf{x}[n]\right)$ is formed
by the basis matrix $\mathbf{E}\left(\mathbf{x}[n]\right)$ and eigenvalue
diagonal matrix $\mathbf{D}\left(\mathbf{x}[n]\right)$

$\displaystyle\mathbf{M}\left(\mathbf{x}[n]\right)$
$\displaystyle=\mathbf{E}\left(\mathbf{x}[n]\right)\mathbf{D}\left(\mathbf{x}[n]\right)\mathbf{E}\left(\mathbf{x}[n]\right)^{-1}$
(6) $\displaystyle\mathbf{E}\left(\mathbf{x}[n]\right)$
$\displaystyle=\left[r\left(\mathbf{x}[n]\right),e_{1}\left(\mathbf{x}[n]\right),\dots,e_{d-1}\left(\mathbf{x}[n]\right)\right]$
$\displaystyle\mathbf{D}\left(\mathbf{x}[n]\right)$
$\displaystyle=\mathbf{diag}\left(\lambda_{r}\left(\mathbf{x}[n]\right),\lambda_{e}\left(\mathbf{x}[n]\right),\dots,\lambda_{e}\left(\mathbf{x}[n]\right)\right)$

where $r\left(\mathbf{x}[n]\right)$ is a vector from the obstacle reference
point to the robot reference point, and $e_{i}\left(\mathbf{x}[n]\right)$ are
vectors tangent to the obstacle surface.
$\lambda_{r}\left(\mathbf{x}[n]\right),\lambda_{e}\left(\mathbf{x}[n]\right)$
are related to the distance between the robot and the obstacle
$\Gamma\left(\mathbf{x}[n]\right)$ as:

$\lambda_{r}\left(\mathbf{x}[n]\right)=1-\frac{1}{\Gamma\left(\mathbf{x}[n]\right)},\quad\lambda_{e}\left(\mathbf{x}[n]\right)=1+\frac{1}{\Gamma\left(\mathbf{x}[n]\right)}$
(7)

This method prevents motion from penetrating the surface while encouraging the
system to move in the tangent directions of the obstacles. In this work, we
extend this method by treating convex obstacles with a convex hull
representation of the robot.

###  II-C Distance Measurement for Convex Objects

The collision avoidance technique in the previous section requires the signed
distance between the robot and an obstacle. Assuming the robot and the
obstacle can be described as a convex object, the distance between two convex
objects $\mathbf{A}$ and $\mathbf{B}$ can be calculated with the Minkowski
difference,
$\mathbf{A}\ominus\mathbf{B}=\\{\mathbf{a}-\mathbf{b}:\forall\mathbf{a}\in\mathbf{A},\forall\mathbf{b}\in\mathbf{B}\\}$.

To find the shortest distance between $\mathbf{A}$ and $\mathbf{B}$ without
computing the entire Minkowski difference, Recursive Gilbert Johnson Keerthi
(RGJK) [19] and Expanding Polytope Algorithm (EPA) [20] exploit its properties
to obtain the shortest separating and penetrating distance respectively.
Together, they give the signed distance $sd(\mathbf{A},\mathbf{B})$ between a
pair of convex hulls as follows:

$sd(\mathbf{A},\mathbf{B})=RGJK(\mathbf{A},\mathbf{B})-EPA(\mathbf{A},\mathbf{B})$
(8)

Many extensions were introduced to speed up this computation [21, 22].

##  III PROPOSED APPROACH

We consider a task with $d$ degree-of-freedom where
$\mathbf{x}[\cdot]\in\mathbb{R}^{d}$ represents the state, and
$\dot{\mathbf{x}}[\cdot]\in\mathbb{R}^{d}$ represents the velocities. Assuming
we have full knowledge of the environment, our goal is to learn a policy
$\dot{\mathbf{x}}[\cdot]=\pi(\mathbf{x}[\cdot]):\mathbb{R}^{d}\rightarrow\mathbb{R}^{d}$
that outputs the velocities which guides the robot towards a target state
$\mathbf{x}_{g}\in\mathbb{R}^{d}$ while enforcing stability on the motion.

###  III-A Automated Demonstration Collection

The data is automatically generated via trajectory optimization. Given an
initial position $\mathbf{x}_{0}$ and a target position $\mathbf{x}_{g}$, the
objective of the optimization is to find a trajectory of positions
$\mathbf{x}[\cdot]$ and a trajectory of velocities $\dot{\mathbf{x}}[\cdot]$,
that travel from $\mathbf{x}_{0}$ and $\mathbf{x}_{g}$ with the shortest
distance. In addition, the solution should be constrained by the position
limits $\mathbf{x}^{-},\mathbf{x}^{+}$ and the velocity limits
$\dot{\mathbf{x}}^{-},\dot{\mathbf{x}}^{+}$.

Trajectory optimization allows us to introduce additional constraints such as
avoiding collisions with the environment. Assuming that the robot has
$N_{link}$ links and performs some tasks in an environment of $N_{obstacles}$
obstacles, we need to avoid collisions between each link of the robot and the
obstacles. For collision avoidance, we model each object as a convex hull,
since they are less prone to overestimating the volume of the object. For non-
convex objects, the object can be represented as a collection of its convex
components. The trajectory optimization is formulated as follows:

$\displaystyle\min_{\mathbf{x}[\cdot],\dot{\mathbf{x}}[\cdot]}$
$\displaystyle~{}~{}~{}\sum_{n=0}^{N-1}\|\mathbf{x}[n+1]-\mathbf{x}[n]\|$ (9)
subject to
$\displaystyle\mathbf{x}[n+1]=\mathbf{f}(\mathbf{x}[n],\dot{\mathbf{x}}[n]),\forall
n\in[0,N-1]$
$\displaystyle\mathbf{x}[0]=\mathbf{x}_{0},\quad\mathbf{x}[N]=\mathbf{x}_{g}$
$\displaystyle\mathbf{x}^{-}\leq\mathbf{x}[n]\leq\mathbf{x}^{+}$
$\displaystyle\dot{\mathbf{x}}^{-}\leq\dot{\mathbf{x}}[n]\leq\dot{\mathbf{x}}^{+}$
$\displaystyle sd(A_{i},O_{j})\geq d_{safe},\forall i\in[1,N_{link}]$
$\displaystyle\hskip 60.0pt\forall j\in[1,N_{obstacles]}$

where $N$ is the number of via points, $\mathbf{f}$ is the forward dynamics of
the robot, $sd$ is the signed distance for convex hulls in (8), $A_{i}$ is the
convex hull of the $i^{th}$ link of the robot, $O_{i}$ is the convex hull of
the $i^{th}$ obstacle, and $d_{safe}>0$ is the minimum distance that must not
be violated for obstacles.

With this information, the formulation in (9) can be passed to an interior
point optimization solver, such as IPOPT [23]. This can be run multiple times
from different initial positions to produce a set of demonstrations.

###  III-B Value Function and Policy Learning

We assume that data are generated using Sec. III-A as a set of positions
$\mathbf{x}[\cdot]$ and velocities $\dot{\mathbf{x}}[\cdot]$. Our goal is to
learn a policy $\pi$ that predicts the most suitable velocity given the
current state $\dot{\mathbf{x}}[\cdot]=\pi(\mathbf{x}[\cdot])$ and satisfies
the Lyapunov stability conditions. The proposed supervised learning method
aims to shape a Lyapunov function
$\tilde{V}\left({\cdot}\right):\mathbb{R}^{d}\rightarrow\mathbb{R}$.

We follow the formulation in [12] to structure a Lyapunov candidate function
as follows:

$\tilde{V}\left({\mathbf{x}[n]}\right)=\phi\left(\mathbf{x}[n]-\mathbf{x}_{g},\theta_{V}\right)-\phi\left(\mathbf{0},\theta_{V}\right)+\|\mathbf{x}[n]-\mathbf{x}_{g}\|$
(10)

where $\phi(\cdot,\theta_{V}):\mathbb{R}^{d}\rightarrow\mathbb{R}$ is an
Artificial Neural Network (ANN) parameterized by $\theta_{V}$. This
formulation imposes that the value at the target state is $0$, which satisfies
(3).

(a) ANN for the Lyapunov function

(b) Policy

Figure 3: Learning Lyapunov policy: (a) an ANN is learned (b) the action is
the negative gradient of the Lyapunov function

The policy can be taken to be related to the negative gradient of the learned
Lyapunov function

$\pi\left(\mathbf{x}[n]\right)\propto-\nabla\tilde{V}\left({\mathbf{x}[n]}\right)$
(11)

The idea is illustrated in Fig. 3. The Lyapunov function
$\tilde{V}\left({\mathbf{x}[n]}\right)$ is nonlinear with respect to positions
$\mathbf{x}[n]$, so we deploy an ANN to learn the relationship (Fig. 3(a)).
Once a Lyapunov function is learned, the negative gradient
$-\nabla\tilde{V}\left({\mathbf{x}[n]}\right)$ will guide the robot toward the
target (Fig. 3(b)).

To train the neural network, the objective is to minimize the discrepancy
between the policy output $-\nabla\tilde{V}\left({\mathbf{x}[n]}\right)$ and
the demonstrated action $\dot{\mathbf{x}}[n]$. This discrepancy $e(n)$ can be
taken from the dot product between two vectors.

$-\nabla\tilde{V}\left({\mathbf{x}[n]}\right)\cdot\dot{\mathbf{x}}[n]=\|\nabla\tilde{V}\left({\mathbf{x}[n]}\right)\|\|\dot{\mathbf{x}}[n]\|\cos(e(n))$

Dividing both sides by the magnitude of the two vectors and find the inverse
cosine, the objective function becomes

$\min_{\theta_{V}}\sum_{n=0}^{N-1}\Bigg{\|}\arccos{\left(\frac{-\nabla\tilde{V}\left({\mathbf{x}[n]}\right)\cdot\dot{\mathbf{x}}[n]}{\|\nabla\tilde{V}\left({\mathbf{x}[n]}\right)\|\|\dot{\mathbf{x}}[n]\|}\right)}\Bigg{\|}$
(12)

In order to generate a valid Lyapunov function, the stability conditions in
(1) and (2) are added to the objective in (12). The problem is converted into:

$\begin{split}\min_{\theta_{V}}&\sum_{n=0}^{N-1}\Bigg{\|}\arccos{\left(\frac{-\nabla\tilde{V}\left({\mathbf{x}[n]}\right)\cdot\dot{\mathbf{x}}[n]}{\|\nabla\tilde{V}\left({\mathbf{x}[n]}\right)\|\|\dot{\mathbf{x}}[n]\|}\right)}\Bigg{\|}\\\
\text{subject to}\
&\tilde{V}\left({\mathbf{x}[n]}\right)>0,\forall\mathbf{x}[n]\neq\mathbf{x}_{g},\\\
&\tilde{V}\left({\mathbf{x}[n+1]}\right)-(1-\epsilon)\tilde{V}\left({\mathbf{x}[n]}\right)\leq
0\end{split}$ (13)

To make the optimization easier for implementation in tools such as PyTorch
[24], the formulation is modified to be an augmented Lagrangian. Also, the
trigonometry is simplified to be a difference by rearranging the terms of the
dot product and inserting the desired angle. These changes produce the
following optimization:

$\displaystyle\min_{\theta_{V}}$
$\displaystyle\sum_{n=0}^{N-1}\left(1+\frac{\nabla\tilde{V}\left({\mathbf{x}[n]}\right)\cdot\dot{\mathbf{x}}[n]}{\|\nabla\tilde{V}\left({\mathbf{x}[n]}\right)\|\|\dot{\mathbf{x}}[n]\|}\right)$
(14)
$\displaystyle+\lambda_{1}\max\Big{(}-\tilde{V}\left({\mathbf{x}[n]}\right),0\Big{)}$
$\displaystyle+\lambda_{2}\max\Big{(}\tilde{V}\left({\mathbf{x}[n+1]}\right)-(1-\epsilon)\tilde{V}\left({\mathbf{x}[n]}\right),0\Big{)}$

where $\lambda_{1}$ and $\lambda_{2}$ are the Lagrange multipliers. The model
will be a valid Lyapunov function for the region around the provided
demonstrations. Outside the demonstration regions, there is no guarantee on
the stability of the model. Termination before the constraints are satisfied
will result in an invalid Lyapunov function. Finally, to ensure that the
policy output is within the physical limits, the output is scaled by

$\pi\left(\mathbf{x}[n]\right)=\frac{-\nabla\tilde{V}\left({\mathbf{x}[n]}\right)}{\|\nabla\tilde{V}\left({\mathbf{x}[n]}\right)\|}\dot{\mathbf{x}}_{max}$
(15)

###  III-C Obstacle Avoidance for Convex object pairs

The learning approach introduced in the previous section learns a policy that
reproduces the demonstrated data with stability guarantee. However, it has no
information about obstacles in the environment. To further guarantee a
collision-free policy, the obstacle avoidance method from [17] was improved
upon. While [17] treats the robot as a point, we augmented the method for
handling convex representations. We represent each link of the robot with the
convex hull of the link, and the obstacles are convex objects.

Figure 4: The reference points on the convex hulls used to get the distance
function output.

For convex to convex comparison, the signed distance between the convex
objects is required. An example is shown in Fig. 4. Assuming $\mathbf{A}$ and
$\mathbf{B}$ are two convex objects, $\mathbf{A}_{ref}$ and $\mathbf{B}_{ref}$
are the reference position of object $\mathbf{A}$ and $\mathbf{B}$
respectively, $p_{a}$ is the closest point on $\mathbf{A}$ to $\mathbf{B}$,
and $p_{b}$ is the closest point on $\mathbf{B}$ to $\mathbf{A}$. The signed
distance $sd(\mathbf{A},\mathbf{B})$ between $\mathbf{A}$ and $\mathbf{B}$ can
be computed from (8).

For handling convex objects, we define the distance function for the
modulation matrix between convex objects as

$\Gamma\left(\mathbf{x}[n]\right)=\frac{\|\mathbf{A}_{ref}-\mathbf{B}_{ref}\|}{\|\mathbf{A}_{ref}-\mathbf{B}_{ref}\|-sd(\mathbf{A},\mathbf{B})}$
(16)

In (16), the signed distance $sd(\mathbf{A},\mathbf{B})$ is used to form a
ratio. When $sd(\mathbf{A},\mathbf{B})=0$, the two objects collide. The result
is $\Gamma\left(\mathbf{x}[n]\right)=1$, so that
$\lambda_{r}\left(\mathbf{x}[n]\right)=0$ in (7) and makes (5) block motion
towards the obstacle. When $sd(\mathbf{A},\mathbf{B})\geq 0$,
$\Gamma\left(\mathbf{x}[n]\right)\geq 1$, which gives the desired eigenvalues
for the obstacle modulation when they are not intersecting.

For the basis vectors in (6),
$r\left(\mathbf{x}[n]\right)=p_{a}-p_{b}/\|p_{a}-p_{b}\|$ is taken to be the
normal direction in which to impede robot motion and
$e_{i}\left(\mathbf{x}[n]\right)$ are chosen to be orthogonal to the normal.
In 2D, this can simply be the normal rotated by $90^{\circ}$. In 3D, this can
be done using spherical coordinates and taking
$\hat{r}=r\left(\mathbf{x}[n]\right)$, and $\hat{\theta}$ and $\hat{\phi}$ to
be the tangents.

A summary of our proposed method is illustrated in Fig. 5.

Figure 5: A summary of our proposed method

##  IV EXPERIMENTS

To verify the performance of the proposed method, we evaluate the method on
several different environment scenarios. The scenarios explored are 2D
examples, a manipulation task in simulation and on a real robot. For each
case, a set of demonstrations are collected through trajectory optimization
method from Section III-A. We use LASA Stable Estimator of Dynamic Systems
(SEDS) as our baseline, where the chosen Lyapunov function is (4). Details
about the baseline and our implementations can be found in the Appendix. We
used the obstacle avoidance module in all experiments to show the difference
in learned system dynamics.

###  IV-A Evaluation Criteria

The performance of the method is evaluated based on its convergence to the
goal, its ability to avoid collisions, its ability to satisfy the Lyapunov
stability conditions and the prediction error. The prediction error is taken
as the mean squared error (MSE) of the unit vectors of the policy outputs and
demonstration actions.

###  IV-B Experiment 1: 2D examples

(a) Demonstrations (red arrows) with the goal is marked with a green dot.

(b) Values and vector field from SEDS with obstacle avoidance.

(c) Values and vector field from our method with obstacle avoidance.

Figure 6: 2D hallway environment. The red arrows denote the velocities. The
background colour gradient and contour lines show the Lyapunov values for the
given position, and the purple streamlines show the motion field of the
policy.

(a) Demonstrations (red arrows) with the goal is marked with a green dot.

(b) Values and vector field from SEDS with obstacle avoidance

(c) Values and Vector field from our Method with obstacle avoidance

Figure 7: 2D cross environment. The red arrows denote the velocities. The
background colour gradient and contour lines show the Lyapunov values for the
given position, and the purple streamlines show the motion field of the
policy.

As a proof-of-concept, the method was first validated on 2D examples. We
designed two scenarios: (a) a hallway case and (b) a cross-shaped case, which
can be viewed in Fig. 6 and in Fig. 7. The target position is marked as the
green dot.

We used the automatic data collection method in Section III-A to collect 75
and 2818 demonstrations for the hallway and cross environments respectively,
where the initial position is a valid position from a grid on the environment
around the target position. The collected demonstrations can be viewed as the
red arrows in Fig. 6(a) and in Fig. 7(a).

The baseline uses a quadratic Lyapunov function. Fig.6(b) and Fig.7(b) are the
results for the baseline, where the background color is the Lyapunov function
values and the purple streamlines are the vector fields. In Fig. 6(b), the
Lyapunov function cannot explain the demonstrations at the top left corner.
Similarly, in Fig.7(b), the top right quadrant fails due to the trajectory
moving away from the target (which violates the quadratic Lyapunov
conditions).

For both cases, our proposed method finds a solution that brings the system to
the target position no matter the starting position, as seen in Fig. 6(c) and
Fig. 7(c). In terms of the Lyapunov conditions, we see clearly that the target
position has a value of $0$, the values are all above $0$ for the environment
and the policy motion field is in the direction of decreasing values, thus
satisfying the criteria for stability.

As for the prediction error, we deploy the policy on our validation
demonstration dataset and compare the prediction velocities with the
demonstrations. The errors for our method and the baseline on the 2D
environments are summarized in Fig 8. We can see that our method has a much
lower MSE than the baseline. This is due to the policy not being restricted by
the quadratic Lyapunov function.

Figure 8: Prediction error of the learned policy using the baseline and our
method on the validation set. The y-axis is the MSE in log scale.

###  IV-C Experiment 2: Manipulation in Simulation

Next, we explore the performance of our method in a manipulation task where
the demonstrations are defined in task space. We use Pybullet [25] to simulate
the environment and use the Mecademic Meca500 as the robot to perform the
task. The task is to take an object from the top of a shelf to a target
position under the shelf (see Fig. 9(a)).

(a) Pybullet simulation

(b) Demonstrations

Figure 9: Simulation and data for the shelf experiment. (a) Mecademic Meca500
in Pybullet simulation (b) Example data collected from trajectory optimization
with multiple initial positions (green) and a target position (orange).

We collected $1209$ demonstrations by sampling a grid of start positions and
select valid positions for trajectory optimization. Some example trajectories
are shown in Fig. 9(b). The network architecture used to learn the model is
found in Appendix. The obstacle avoidance module was added so that the payload
does not collide with the shelf.

The results are summarized in Fig. 10. The outcome of the baseline does not
match the demonstrations and collide with the shelf (Fig. 10(a)). With
collision avoidance, the motion looks closer to the demonstrations, but solely
due to the modulation forcing the movement in the same direction of the
obstacle (Fig. 10(c)). From Fig. 10(b) and Fig. 10(d), our method reproduces
the motion with or without the obstacle avoidance module. In all cases, the
policy rollouts reach the goal. From the prediction error in Fig. 8, we see
that, again, our method has much lower MSE than the baseline.

(a) Baseline without collision avoidance

(b) Our method without collision avoidance

(c) Baseline with collision avoidance

(d) Our method with collision avoidance

Figure 10: Planned paths by the learned Policy with SEDS and our method where
the green points show the initial positions and the orange point shows the
goal position.

Fig. 11 is a visualization of our learnt Lyapunov function. The x-axis is the
time-step $n$, and the y-axis is the Lyapunov function output
$\tilde{V}\left({\mathbf{x}[n]}\right)$ normalized by the output at initial
positions $\tilde{V}\left({\mathbf{x}[0]}\right)$. Fig. 11(a) and Fig. 11(b)
are the outputs of $\tilde{V}\left({\mathbf{x}[n]}\right)$ with the
demonstration data and with the policy rollout, respectively. We can see that
the Lyapunov function outputs are positive and monotonically decreasing until
they reach $0$ at the goal. This shows that our method found a suitable
Lyapunov function that explains the data.

(a) Demonstrations

(b) Policy Rollout

Figure 11: Trajectories of the learned Lyapunov function outputs
$\tilde{V}\left({\mathbf{x}[.]}\right)$ normalized by the output at the
initial positions $\tilde{V}\left({\mathbf{x}[0]}\right)$, where x-axis is the
time-step $n$ and y-axis is the Lyapunov function output
$\tilde{V}\left({\mathbf{x}[n]}\right)$.

We also look at the effect of perturbations on the policy. We find that the
policy can robustly recover from external perturbations on the robot, since
the policy learns a state dependant motion field from the demonstrations and
is guaranteed to reach the goal (see the supplementary video
https://youtu.be/_NzE6LxoSx4).

###  IV-D Experiment 3: Manipulation with a real robot

We deploy the policy learnt from Experiment 2 on a robotic hardware, the
Meca500 robot from Mecademic. Similar to the simulation, the task is to bring
an object from anywhere in the shelf environment to a goal position below the
shelf.

We start the robot from several initial positions and deploy the policy to
bring the system to the goal. The policy is used as the motion planner and
provides the physical robot the sequence of task-space via points toward the
goal. Then, this is controlled through the standard inverse kinematics
control. Fig. 1 is an example of the outcome; where Fig. 1(a) shows an example
of initial position, Fig. 1(b) is the policy rollout (red), and Fig. 1(c)
shows that the robot reaches the target position.

The performance of the policy on the real-world robot can be viewed in the
supplementary video. We see that the robot is able to perform the task without
making contact with the shelf. In all cases, the policy performs the same in
task-space on the real-world system as it did in simulation.

Our method can achieve sim-to-real transfer reliably. This is due to our
policy being kinematically planned, (i.e., mapping the positions to the
velocities), which is less prone to modelling error. Also, when the Lyapunov
stability conditions are enforced, the movement are restricted to the
directions of demonstrations and less likely to diverge.

##  V CONCLUSION & DISCUSSION

We propose a novel method for learning a Lyapunov function and a policy using
a single neural network through imitation learning. Our method is able to
learn a policy that satisfy the Lyapunov stability conditions and reproduce
the demonstrations. With our extension of the previous collision avoidance
module, our method is capable of avoiding collisions between convex
representations of the robot and environment obstacles. The policies were
successfully applied to simulated environments and a real-world scenario.

In future work, the obstacle information will be incorporated into the learned
neural network. This will remove the need for the obstacle avoidance module to
prevent collisions and have the obstacles be learned by the network. With this
improvement, the gradient of the Lyapunov function can have information for
how to move away from obstacles and prevent colliding with them. This would
make our method self-contained and capable of obstacle avoidance on its own.

## APPENDIX

For the baseline, the data were automatically clustered into $6$, $11$ and $3$
Gaussian Mixture Model (GMM) for the hallway, cross, and shelf environment
respectively. Thus, is required to learn $6$, $11$ and $3$ linear dynamical
systems for each segment.

For our experiments, the activation function used for all the neurons is the
tanh function. The architectures used can be found in Table A1.

TABLE A1: The neural network structures used for each experiment. The information is the number of neurons in each layer: input layer, hidden layers, output layer. Experiment | Neural Network Architecture
---|---
Hallway | $\left(2,128,128,128,1\right)$
Cross | $\left(2,128,128,128,1\right)$
Shelf | $\left(3,256,256,256,256,1\right)$

## References

  * [1] J. T. Betts, “Survey of numerical methods for trajectory optimization,” _Journal of guidance, control, and dynamics_ , vol. 21, no. 2, pp. 193–207, 1998.
  * [2] M. Kalakrishnan, S. Chitta, E. Theodorou, P. Pastor, and S. Schaal, “Stomp: Stochastic trajectory optimization for motion planning,” in _IEEE international conference on robotics and automation_ , 2011, pp. 4569–4574.
  * [3] J. Schulman, J. Ho, A. X. Lee, I. Awwal, H. Bradlow, and P. Abbeel, “Finding locally optimal, collision-free trajectories with sequential convex optimization.” in _Robotics: science and systems_ , vol. 9, no. 1. Citeseer, 2013, pp. 1–10.
  * [4] J. Schulman, Y. Duan, J. Ho, A. Lee, I. Awwal, H. Bradlow, J. Pan, S. Patil, K. Goldberg, and P. Abbeel, “Motion planning with sequential convex optimization and convex collision checking,” _The International Journal of Robotics Research_ , vol. 33, no. 9, pp. 1251–1270, 2014.
  * [5] M. Bojarski, D. Del Testa, D. Dworakowski, B. Firner, B. Flepp, P. Goyal, L. D. Jackel, M. Monfort, U. Muller, J. Zhang, _et al._ , “End to end learning for self-driving cars,” _arXiv preprint arXiv:1604.07316_ , 2016.
  * [6] H.-C. Lin, M. Howard, and S. Vijayakumar, “A novel approach for representing and generalising periodic gaits,” _Robotica_ , vol. 32, no. 8, pp. 1225–1244, 2014.
  * [7] P. Abbeel and A. Y. Ng, “Apprenticeship learning via inverse reinforcement learning,” in _Proceedings of the twenty-first international conference on Machine learning_ , 2004, pp. 1–8.
  * [8] J. Ho and S. Ermon, “Generative adversarial imitation learning,” _Advances in neural information processing systems_ , vol. 29, 2016.
  * [9] I. Goodfellow, J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio, “Generative adversarial nets,” in _Advances in neural information processing systems_ , vol. 27, 2014.
  * [10] S. M. Khansari-Zadeh and A. Billard, “Learning stable nonlinear dynamical systems with gaussian mixture models,” _IEEE Transactions on Robotics_ , vol. 27, no. 5, pp. 943–957, 2011.
  * [11] N. Figueroa and A. Billard, “Locally active globally stable dynamical systems: Theory, learning, and experiments,” _The International Journal of Robotics Research_ , 2022.
  * [12] H. Dai, B. Landry, L. Yang, M. Pavone, and R. Tedrake, “Lyapunov-stable neural-network control,” _Robotics: Science and Systems_ , 2021.
  * [13] F. Berkenkamp, M. Turchetta, A. Schoellig, and A. Krause, “Safe model-based reinforcement learning with stability guarantees,” in _Advances in neural information processing systems_ , vol. 30, 2017.
  * [14] C. W. Warren, “Global path planning using artificial potential fields,” in _IEEE International Conference on Robotics and Automation_ , 1989, pp. 316–317.
  * [15] S. M. Khansari-Zadeh and A. Billard, “A dynamical system approach to realtime obstacle avoidance,” _Autonomous Robots_ , vol. 32, no. 4, pp. 433–454, 2012\.
  * [16] L. Huber, A. Billard, and J.-J. Slotine, “Avoidance of convex and concave obstacles with convergence ensured through contraction,” _IEEE Robotics and Automation Letters_ , vol. 4, no. 2, pp. 1462–1469, 2019.
  * [17] L. Huber, J.-J. Slotine, and A. Billard, “Avoiding dense and dynamic obstacles in enclosed spaces: Application to moving in crowds,” _IEEE Transactions on Robotics_ , 2022.
  * [18] N. P. Bhatia and G. P. Szegö, _Stability theory of dynamical systems_. Springer Science and Business Media, 2002.
  * [19] E. G. Gilbert, D. W. Johnson, and S. S. Keerthi, “A fast procedure for computing the distance between complex objects in three-dimensional space,” _IEEE Journal on Robotics and Automation_ , vol. 4, no. 2, pp. 193–203, 1988\.
  * [20] G. Van Den Bergen, _Collision detection in interactive 3D environments_. CRC Press, 2003.
  * [21] C. J. Ong and E. G. Gilbert, “The Gilbert-Johnson-Keerthi distance algorithm: a fast version for incremental motions,” vol. 2, pp. 1183–1189, 1997\.
  * [22] A. Coulombe and H.-C. Lin, “High precision real time collision detection,” _arXiv preprint arXiv:2007.12045_ , 2020.
  * [23] A. Wächter and L. T. Biegler, “On the implementation of an interior-point filter line-search algorithm for large-scale nonlinear programming,” _Mathematical programming_ , vol. 106, no. 1, pp. 25–57, 2006.
  * [24] A. Paszke, S. Gross, F. Massa, A. Lerer, J. Bradbury, G. Chanan, T. Killeen, Z. Lin, N. Gimelshein, L. Antiga, A. Desmaison, A. Kopf, E. Yang, Z. DeVito, M. Raison, A. Tejani, S. Chilamkurthy, B. Steiner, L. Fang, J. Bai, and S. Chintala, “Pytorch: An imperative style, high-performance deep learning library,” vol. 32, pp. 8024–8035, 2019.
  * [25] E. Coumans and Y. Bai, “Pybullet, a python module for physics simulation for games, robotics and machine learning,” http://pybullet.org, 2016–2021.


# Mamba4Rec: Towards Efficient Sequential Recommendation
with Selective State Space Models

Chengkai Liu  Texas A&M UniversityCollege Station, TexasUSA
<EMAIL_ADDRESS>,  Jianghao Lin  Shanghai Jiao Tong
UniversityShanghaiChina<EMAIL_ADDRESS>,  Jianling Wang  Texas A&M
UniversityCollege Station, TexasUSA<EMAIL_ADDRESS>,  Hanzhou Liu  Texas
A&M UniversityCollege Station, TexasUSA<EMAIL_ADDRESS>and  James
Caverlee  Texas A&M UniversityCollege Station, TexasUSA<EMAIL_ADDRESS>
###### Abstract.

Sequential recommendation aims to estimate the dynamic user preferences and
sequential dependencies among historical user behaviors. Although Transformer-
based models have proven to be effective for sequential recommendation, they
suffer from the inference inefficiency problem stemming from the quadratic
computational complexity of attention operators, especially for long-range
behavior sequences. Inspired by the recent success of state space models
(SSMs), we propose Mamba4Rec, which is the first work to explore the potential
of selective SSMs for efficient sequential recommendation. Built upon the
basic Mamba block which is a selective SSM with an efficient hardware-aware
parallel algorithm, we incorporate a series of sequential modeling techniques
to further promote the model performance and meanwhile maintain the inference
efficiency. Experiments on two public datasets demonstrate that Mamba4Rec is
able to well address the effectiveness-efficiency dilemma, and defeat both
RNN- and attention-based baselines in terms of both effectiveness and
efficiency. The code is available at https://github.com/chengkai-
liu/Mamba4Rec.

Sequential Recommendation, State Space Models

##  1\. Introduction

Personalized online services rely heavily on sequential recommender systems to
capture dynamic user preferences (wang2021survey; lu2015recommender;
lin2023map). These systems aim to predict users’ future interactions by
effectively modeling sequential dependencies among historical user behaviors.
Earlier neural approaches adopt Convolutional Neural Networks (CNNs)
(tang2018personalized) and Recurrent Neural Networks (RNNs)
(hidasi2015session; li2017neural) which have pioneered the application of
neural networks in sequential recommendation yet suffer from the catastrophic
forgetting issue (schak2019study; kirkpatrick2017overcoming). More recently,
researchers start to introduce Transformer (vaswani2017attention) as the
backbone module and leverage the self-attention mechanism to model the dynamic
user behavior sequences (kang2018self; yuan2022multi). In spite of their
remarkable performance, these attention-based methods generally suffer from
the inference inefficiency problem due to the quadratic computational
complexity inherent in attention operators, especially for long user behavior
sequences or large-scale datasets.

To address the dilemma of recommendation performance and inference efficiency,
in this paper, we consider the state space model (SSM) (gu2021efficiently) as
the core operator for sequential recommendation. SSMs (_e.g._ , S4
(gu2021efficiently), S5 (smith2022simplified), Mega (ma2022mega)) have been
widely adopted as alternatives to RNNs, CNNs, and Transformers in various
natural language tasks. They inherently possess inference efficiency due to
their recurrent nature, and meanwhile enjoy strong performance to handle long-
range dependencies through structured state matrices. Moreover, a recently
proposed SSM variant (_i.e._ , Mamba (gu2023mamba)) further introduces the
data-dependent selective mechanism with efficient hardware-aware design to
tackle the data- and time-invariant issue of prior SSMs. This allows the model
to selectively extract essential knowledge and filter out irrelevant noise
according to the input data, thus leading to superior sequential modeling
performance.

In this work, we propose Mamba4Rec, which is the first work to explore the
potential of selective SSM for efficient sequential recommendation. Built upon
the basic Mamba block (gu2023mamba), we incorporate and discuss a series of
techniques (_e.g._ , position embedding (vaswani2017attention), residual
connection (he2016deep), layer normalization (ba2016layer), and position-wise
feed-forward network (vaswani2017attention)), which supplements the system
non-linearity, stabilizes the training dynamic, and thereby further promotes
the sequential modeling capability.

The main contributions of this paper are as follows:

  * •

We are the first to explore the potential of selective SSMs for sequential
recommendation, which addresses the dilemma of recommendation performance and
inference efficiency.

  * •

We propose Mamba4Rec by incorporating the basic Mamba block with various
techniques like layer normalization and feed-forward networks, which further
improves the sequential modeling capability without sacrificing the inference
efficiency.

  * •

Experiments on two public datasets demonstrate the superiority of our proposed
Mamba4Rec compared with RNN- and attention-based baselines in terms of both
effectiveness and efficiency.

##  2\. Preliminaries

###  2.1. Problem Statement

In sequential recommendation, let
$\mathcal{U}=\left\\{u_{1},u_{2},\ldots,u_{|\mathcal{U}|}\right\\}$ denote the
user set, $\mathcal{V}=\left\\{v_{1},v_{2},\ldots,v_{|\mathcal{V}|}\right\\}$
denote the item set, and
$\mathcal{S}_{u}=\left\\{v_{1},v_{2},\ldots,v_{n_{u}}\right\\}$ denote the
chronologically ordered interaction sequence for user $u\in\mathcal{U}$, where
$n_{u}$ is the the length of the sequence. Given the interaction history
$\mathcal{S}_{u}$, the task is to predict the next interacted item, denoted as
$v_{n_{u+1}}$ for user $u$.

###  2.2. State Space Models

The State Space Model (SSM) is a sequence modeling framework based on linear
ordinary differential equations. It maps an input sequence
$x(t)\in\mathbb{R}^{D}$ to the output sequence $y(t)\in\mathbb{R}^{N}$ through
the latent state $h(t)\in\mathbb{R}^{N}$:

(1) $\displaystyle h^{\prime}(t)$ $\displaystyle=\bm{A}h(t)+\bm{B}x(t),$
$\displaystyle y(t)$ $\displaystyle=\bm{C}h(t),$

where $\bm{A}\in\mathbb{R}^{N\times N}$ and
$\bm{B},\bm{C}\in\mathbb{R}^{N\times D}$ are learnable matrices. To model
discrete sequences instead of continuous functions, it must be discretized by
a step size $\Delta$. Alternatively, the discretized SSM can be expressed as
follows:

(2) $\displaystyle h_{t}$
$\displaystyle=\overline{\bm{A}}h_{t-1}+\overline{\bm{B}}x_{t},$
$\displaystyle y_{t}$ $\displaystyle=\bm{C}h_{t},$

where $\overline{\bm{A}}=\exp(\Delta\bm{A})$ and
$\overline{\bm{B}}=(\Delta\bm{A})^{-1}(\exp(\Delta\bm{A})-\bm{I})\cdot\Delta\bm{B}$.
After converting the parameters from a continuous form
$(\Delta,\bm{A},\bm{B},\bm{C})$ to a discrete form
$(\overline{\bm{A}},\overline{\bm{B}},\bm{C})$, the model can be computed in a
linear recurrence way (gu2021combining), which enhances computational
efficiency. Deriving from the vanilla SSM, the structured state space model
(S4) (gu2021efficiently) imposes structure on the state matrix $\bm{A}$ with
HiPPO (gu2020hippo) initialization to further improve the long-range
dependency modeling.

Building upon S4, Mamba (gu2023mamba) introduces a data-dependent selection
mechanism and leverages a hardware-aware parallel algorithm in recurrent mode.
This combined approach empowers Mamba to capture contextual information
effectively, particularly in long sequences, while maintaining computational
efficiency. As a linear-time sequence model, Mamba achieves Transformer-
quality performance with better efficiency, especially on long sequences.

##  3\. Mamba4Rec

In this section, we introduce our proposed framework, Mamba4Rec. We begin with
a high-level overview of the framework, followed by an exploration of its
technical components. We describe how Mamba4Rec constructs a sequential
recommendation model through an embedding layer, selective state space models,
and a prediction layer. Additionally, we discuss the key components widely
used in sequential recommendation, such as positional embeddings
(vaswani2017attention), feed-forward networks, dropout
(srivastava2014dropout), and layer normalization (ba2016layer).

###  3.1. Framework Overview

As illustrated in Figure 1, the proposed Mamba4Rec is a sequential
recommendation model that leverages selective state space models through Mamba
blocks. The core element of Mamba4Rec is the Mamba layer, which combines a
Mamba block with a position-wise feed-forward network. Mamba4Rec offers
flexibility in its architecture: while it can be stacked with multiple Mamba
layers, a single Mamba layer often proves sufficient. Each layer comprises a
Mamba block followed by a feed-forward network, enabling the model to capture
both item-specific information and sequential context effectively from the
user’s interaction history.

Figure 1. The overview of Mamba4Rec.

###  3.2. Embedding Layer

Similar to existing models, our approach utilizes an embedding layer to map
item IDs to a high-dimensional space. This embedding layer uses a learnable
embedding matrix $\bm{E}\in\mathbb{R}^{|\mathcal{V}|\times D}$, where $D$ is
the embedding dimension. By applying the embedding layer to the input item
sequence $\mathcal{S}_{u}$, we obtain the item embeddings $H$. To enhance
robustness and prevent overfitting, we incorporate both embedding dropout and
layer normalization after retrieving the embeddings:

(3) $\displaystyle H$
$\displaystyle=\operatorname{LayerNorm}\left(\operatorname{Dropout}\left(H_{o}\right)\right)\in\mathbb{R}^{n_{u}\times
D}.$

During the training and inference stage, we merge multiple samples into one
mini-batch, resulting the input data $x\in\mathbb{R}^{B\times L\times D}$ for
the subsequent Mamba block, where $B$ is the batch size and $L$ is padded
sequence length.

Positional Embedding. We investigate the impact of positional embeddings on
our model’s performance. This includes evaluating learnable positional
embeddings as used in (kang2018self; sun2019bert4rec). These embeddings are
typically injected into the input item embedding before layer normalization.
However, unlike Transformer-based models that rely on positional embeddings to
capture the sequence order, state space models inherently handle sequential
information through their recurrent nature. While theoretically unnecessary,
we analyze the effect of positional embeddings on our model’s behavior in our
experiments.

###  3.3. Mamba Layer

Table 1. Recommendation performance. The best results are bold, and the second-best are underlined. Method | MovieLens-1M | Amazon-Beauty | Amazon-Video-Games
---|---|---|---
HR@10 | NDCG@10 | MRR@10 | HR@10 | NDCG@10 | MRR@10 | HR@10 | NDCG@10 | MRR@10
| NARM | 0.2735 | 0.1506 | 0.1132 | 0.0627 | 0.0347 | 0.0262 | 0.1032 | 0.0530 | 0.0379
| GRU4Rec | 0.2934 | 0.1642 | 0.1249 | 0.0606 | 0.0332 | 0.0249 | 0.1030 | 0.0536 | 0.0380
| SASRec | 0.2977 | 0.1687 | 0.1294 | 0.0847 | 0.0425 | 0.0296 | 0.1168 | 0.0571 | 0.0390
| BERT4Rec | 0.3098 | 0.1764 | 0.1357 | 0.0760 | 0.0393 | 0.0282 | 0.1053 | 0.0538 | 0.0381
| Mamba4Rec | 0.3121 | 0.1822 | 0.1425 | 0.0812 | 0.0451 | 0.0362 | 0.1152 | 0.0603 | 0.0438

While the Mamba block (gu2023mamba) has potential as a sequence-to-sequence
model, we leverage its strengths for sequential recommendation tasks. We
introduce the Mamba layer, which incorporates both the Mamba block and a
position-wise feed-forward network. Instead of simply stacking Mamba blocks,
we use the Mamba layer as the fundamental building block for our architecture.

Mamba Block. As detailed in Algorithm 1, the Mamba block operates on an input
$x_{t}\in\mathbb{R}^{B\times L\times D}$ with dimensions of batch size $B$,
sequence length $L$, and hidden dimension $D$ at each time step $t$. It begins
by performing linear projections on the input $x_{t}$ with an expanded hidden
dimension of $D$ to obtain $x$ and $z$. These projections are then processed
through a 1D convolution and a SiLU (elfwing2018sigmoid) activation. The core
of the block involves a selective state space model with parameters
discretized based on the input. This discretized SSM, along with $x^{\prime}$,
generates the state representation $y$. Finally, $y$ is combined with a
residual connection from $z$ after applying SiLU, and a final linear
projection delivers the final output $y_{t}$ at time step $t$. Overall, the
Mamba block leverages input-dependent adaptations and selective SSM to process
sequential information effectively. The parameters of the Mamba block consist
of an SSM state expansion factor $N$, a kernel size $K$ for the convolution,
and a block expansion factor $E$ for input and output linear projections.
Selecting the expansion factors $N$ and $E$ involves a tradeoff between
(capturing complex relationships) and training efficiency (higher $N,E$
increase computational cost).

The Mamba block addresses the prevalent effectiveness-efficiency dilemma
existing in RNNs and Transformers. The recurrent SSM in the Mamba block
operates in linear time with respect to the sequence length, making it as
efficient as RNNs for inference, achieved by discretized parameters.
Additionally, the selection mechanism makes the parameters of SSM data-
dependent. This enables Mamba to selectively remember or forget information,
crucial for capturing complex dependencies within sequential data. With the
structured state matrix addressing long-range dependencies with HiPPO, it
outperforms self-attention blocks in modeling long sequences.

Algorithm 1 Mamba Block with Selective SSM

1:Input: $x_{t}:(B,L,D)$

2:Output: $y_{t}:(B,L,D)$

3:$x,z:(B,L,ED)\leftarrow\operatorname{Linear}(x_{t})$ $\triangleright$ Linear
projections of input

4:$x^{\prime}:(B,L,ED)\leftarrow\operatorname{SiLU}(\operatorname{Conv1d}(x)$

5:$\bm{A}:(D,N)\leftarrow\text{Parameter}$ $\triangleright$ Structured state
matrix

6:$\bm{B,C}:(B,L,N)\leftarrow\operatorname{Linear}(x^{\prime}),\operatorname{Linear}(x^{\prime})$

7:$\Delta:(B,L,D)\leftarrow\operatorname{Softplus}(\text{Parameter}+\operatorname{Broadcast}(\operatorname{Linear}(x^{\prime})))$

8:$\overline{\bm{A}},\overline{\bm{B}}:(B,L,D,N)\leftarrow\text{discretize}(\Delta,\bm{A},\bm{B})$
$\triangleright$ Input-dependent parameters and discretization

9:$y:(B,L,ED)\leftarrow\text{SelectiveSSM}(\overline{\bm{A}},\overline{\bm{B}},\bm{C})(x^{\prime})$

10:$y^{\prime}:(B,L,ED)\leftarrow y\otimes\operatorname{SiLU}(z)$

11:$y_{t}:(B,L,D)\leftarrow\operatorname{Linear}(y^{\prime})$ $\triangleright$
Final linear projection

12:return $y_{t}$

Position-wise Feed-Forward Network. We leverage a position-wise feed-forward
network (PFFN) to improve the modeling of user actions in the hidden
dimension:

(4)
$\operatorname{PFFN}(H)=\operatorname{GELU}\left(H\bm{W}^{(1)}+\bm{b}^{(1)}\right)\bm{W}^{(2)}+\bm{b}^{(2)},$

where $\bm{W}^{(1)}\in\mathbb{R}^{D\times
4D},\bm{W}^{(2)}\in\mathbb{R}^{4D\times D},\bm{b}^{(1)}\in\mathbb{R}^{4D}$ and
$\bm{b}^{(2)}\in\mathbb{R}^{D}$ are parameters of two dense layers and we use
GELU (hendrycks2016gaussian) activation. The PFFN aims to capture complex
patterns and interactions within the sequential data by applying two non-
linear transformations using dense layers with the activation function.
Additionally, to enhance model robustness and prevent overfitting, we utilize
dropout and layer normalization, similar to Equation 3, after each Mamba block
and feed-forward network within the Mamba layers. This helps to regularize the
model and accelerate training convergence.

Stacking Mamba Layers We also explore using stacked Mamba layers in sequential
recommendation. While deeper networks with multiple layers are not inherently
beneficial for performance, residual connections (he2016deep) play a crucial
role in facilitating the propagation of low-level features to higher layers
throughout the multi-layer architecture. Recognizing the potential benefits of
stacked Mamba layers, especially on long sequences, we explore two
configurations for Mamba layers that balance effectiveness and efficiency:

  * •

Single Layer: When using only one Mamba layer, we apply layer normalization
without residual connections. A single Mamba layer demonstrates comparable
effectiveness to stacking Mamba layers while achieving higher efficiency in
our tasks.

  * •

Stacked Layers: For stacked Mamba layers, residual connections become
essential and enable the propagation of the last visited item’s embedding to
the final layer.

To explore the tradeoff between effectiveness and efficiency, we further
investigate the potential of stacked layers in experiments.

###  3.4. Prediction Layer

In the prediction layer, we adopt the approach of SASRec, using the last item
embedding to generate the final output prediction scores:

(5)
$\hat{y}=\operatorname{Softmax}\left(\operatorname{Linear}\left(h\right)\right)\in\mathbb{R}^{|\mathcal{V}|},$

where $h\in\mathbb{R}^{D}$ is the last item embedding from the Mamba layer.
$\hat{y}\in\mathbb{R}^{|\mathcal{V}|}$ represents the probability distribution
over the next item in the item set $\mathcal{V}$.

##  4\. Experiments

###  4.1. Experimental Setup

Datasets. We conduct experiments and evaluate our models using three real-
world datasets:

  * •

MovieLens-1M (harper2015movielens): A benchmark dataset contains about 1
million movie ratings from users.

  * •

Amazon-Beauty and Amazon-Video-Games (mcauley2015image): Two datasets contain
product reviews and ratings for the “Beauty” and “Video Games” categories on
Amazon.

For each user, we construct an interaction sequence by sorting their
interaction records based on timestamps. We filter users and items with less
than 5 recorded interactions, following the setting in previous works
(kang2018self). The detailed statistics for each dataset after preprocessing
are summarized in Table 2.

Table 2. Statistics of the experimented datasets.

Dataset | # Users | # Items | # Interactions | Avg. Length
---|---|---|---|---
ML-1M | $6,040$ | $3,416$ | $999,611$ | $165.4$
Beauty | $22,363$ | $12,101$ | $198,502$ | $8.9$
Video-Games | $14,494$ | $6,950$ | $132,209$ | $9.1$

Baselines. We compare Mamba4Rec against several baseline methods, which
include RNN-based models like GRU4Rec (hidasi2015session) and NARM
(li2017neural), as well as Transformer-based models such as SASRec
(kang2018self) and BERT4Rec (sun2019bert4rec).

Evaluation Metrics. We adopt Hit Ratio (HR), Normalized Discounted Cumulative
Gain (NDCG), and Mean Reciprocal Rank (MRR) with truncation at 10 as the
evaluation metrics, _i.e._ , HR@10, NDCG@10, and MRR@10.

Implementation Details. In the default architecture of our model, a single
Mamba layer is employed without position embeddings. The Adam optimizer
(kingma2014adam) is used with a learning rate 0.001. The training batch size
is 2048, and the evaluation batch size is 4096. All models use an embedding
dimension of 64. To address the sparsity of the Amazon datasets, a dropout
rate of 0.4 is used, compared to 0.2 for MovieLens-1M. The maximum sequence
length is set proportionally to the mean number of actions per user: 200 for
MovieLens-1M and 50 for Amazon datasets. We follow RecBole (zhao2021recbole)
for further implementation details. For the parameters of the Mamba block, the
SSM state expansion factor is 32, the kernel size for 1D convolution is 4, and
the block expansion factor for linear projections is 2.

###  4.2. Overall Performance

Table 1 demonstrates the overall recommendation performance. Among sequential
recommendation methods, Transformer-based models SASRec and BERT4Rec generally
outperform RNN-based baselines. Notably, the state-space-based Mamba4Rec
achieves superior performance compared to both RNN and Transformer-based
baselines across sparse and dense datasets, as well as varying maximum
sequence lengths. Compared to SASRec, Mamba4Rec achieves significantly better
performance on the MovieLens-1M dataset, with longer sequence lengths. This
highlights the potential of state space models as a compelling alternative to
Transformers and RNNs architectures for sequential recommendation.

###  4.3. Model Efficiency

Table 3 details the efficiency performance of Transformer-based models and
Mamba4Rec, all measured on a single Nvidia A5000 GPU. Considering the tradeoff
between efficiency and effectiveness, we use the model with a single Mamba
layer instead of models with only a Mamba block or stacked Mamba layers on the
long-sequence MovieLens-1M dataset. Mamba4Rec exhibits superior performance in
terms of GPU memory usage, training time per epoch, and inference time per
batch. Additionally, it demonstrates faster convergence, reaching its best
validation performance around epoch 50, compared to SASRec (epoch 70) and
BERT4Rec (epoch 180). These results suggest that Mamba4Rec achieves superior
efficiency and reduces GPU memory cost, particularly on long-sequence datasets
when compared to Transformer-based models.

Table 3. Efficiency performance on MovieLens-1M Method | GPU memory | Training time | Inference time
---|---|---|---
SASRec | 14.98GB | 131.24s | 0.27s
BERT4Rec | 15.48GB | 207.60s | 0.90s
Mamba4Rec | 4.82GB | 75.52s | 0.14s

###  4.4. Ablation Study

To understand the contributions of each component in our architecture, we
conduct an ablation study. Table 4 presents the results. We introduce each
variant and analyze its respective impact:

  * •

Positional embedding: The results demonstrate that adding learnable positional
embeddings does not improve the performance, due to the recurrent nature of
the Mamba block.

  * •

Layer normalization: Following each Mamba block and feed-forward network, we
employ dropout and layer normalization. Our results demonstrate their
effectiveness in mitigating overfitting and enhancing performance.

  * •

Feed-forward network: The results indicate that the PFFN layer is particularly
beneficial for the longer sequence dataset, MovieLens-1M. This suggests the
PFFN’s ability to model nonlinearity leads to improvement in performance.

  * •

Mamba block only: Replacing the Mamba layer with a Mamba block (and removing
PFFN, dropout, and layer normalization) results in faster computation but
sacrifices overall performance.

  * •

Number of Mamba layers: Stacking two Mamba layers with residual connections
yields performance improvements, particularly on the larger, long-sequence
MovieLens-1M dataset. This shows the model’s ability to handle complex user
behavior at the cost of reduced computational efficiency.

Table 4. Ablation analysis on ML-1M and Beauty.

Architecture | Movielens-1M | Amazon-Beauty
---|---|---
NDCG@10 | MRR@10 | NDCG@10 | MRR@10
Default | 0.1822 | 0.1425 | 0.0451 | 0.0362
Block Only | 0.1786 | 0.1404 | 0.0405 | 0.0331
2 Layers | 0.1845 | 0.1437 | 0.0440 | 0.0314
w/ PE | 0.1738 | 0.1342 | 0.0408 | 0.0319
w/o PFFN | 0.1740 | 0.1357 | 0.0437 | 0.0340
w/o LayerNorm | 0.1820 | 0.1396 | 0.0419 | 0.0335

##  5\. Conclusion and Future Work

In this work, we tackle the limitations of Transformer-based sequential
recommendation models. We introduce Mamba4Rec built upon the selective
structured state space model. Our experiments demonstrate that Mamba4Rec
achieves strong performance on datasets with varying sparsity and sequence
lengths. Additionally, it has significant improvements in both computational
efficiency and memory cost, making it effective and efficient for sequential
recommendation tasks. In the future, we aim to design state space models
tailored for recommendation systems and promote the SSM-based models in this
domain.

## References

  * (1)
  * Ba et al. (2016) Jimmy Lei Ba, Jamie Ryan Kiros, and Geoffrey E Hinton. 2016.  Layer normalization.  _arXiv preprint arXiv:1607.06450_ (2016).
  * Elfwing et al. (2018) Stefan Elfwing, Eiji Uchibe, and Kenji Doya. 2018.  Sigmoid-weighted linear units for neural network function approximation in reinforcement learning.  _Neural networks_ 107 (2018), 3–11.
  * Gu and Dao (2023) Albert Gu and Tri Dao. 2023.  Mamba: Linear-time sequence modeling with selective state spaces.  _arXiv preprint arXiv:2312.00752_ (2023).
  * Gu et al. (2020) Albert Gu, Tri Dao, Stefano Ermon, Atri Rudra, and Christopher Ré. 2020.  Hippo: Recurrent memory with optimal polynomial projections.  _Advances in neural information processing systems_ 33 (2020), 1474–1487.
  * Gu et al. (2021a) Albert Gu, Karan Goel, and Christopher Ré. 2021a.  Efficiently modeling long sequences with structured state spaces.  _arXiv preprint arXiv:2111.00396_ (2021).
  * Gu et al. (2021b) Albert Gu, Isys Johnson, Karan Goel, Khaled Saab, Tri Dao, Atri Rudra, and Christopher Ré. 2021b.  Combining recurrent, convolutional, and continuous-time models with linear state space layers.  _Advances in neural information processing systems_ 34 (2021), 572–585.
  * Harper and Konstan (2015) F Maxwell Harper and Joseph A Konstan. 2015.  The movielens datasets: History and context.  _Acm transactions on interactive intelligent systems (tiis)_ 5, 4 (2015), 1–19.
  * He et al. (2016) Kaiming He, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2016.  Deep residual learning for image recognition. In _Proceedings of the IEEE conference on computer vision and pattern recognition_. 770–778.
  * Hendrycks and Gimpel (2016) Dan Hendrycks and Kevin Gimpel. 2016.  Gaussian error linear units (gelus).  _arXiv preprint arXiv:1606.08415_ (2016).
  * Hidasi et al. (2015) Balázs Hidasi, Alexandros Karatzoglou, Linas Baltrunas, and Domonkos Tikk. 2015.  Session-based recommendations with recurrent neural networks.  _arXiv preprint arXiv:1511.06939_ (2015).
  * Kang and McAuley (2018) Wang-Cheng Kang and Julian McAuley. 2018.  Self-attentive sequential recommendation. In _2018 IEEE international conference on data mining (ICDM)_. IEEE, 197–206.
  * Kingma and Ba (2014) Diederik P Kingma and Jimmy Ba. 2014.  Adam: A method for stochastic optimization.  _arXiv preprint arXiv:1412.6980_ (2014).
  * Kirkpatrick et al. (2017) James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Grabska-Barwinska, et al. 2017\.  Overcoming catastrophic forgetting in neural networks.  _Proceedings of the national academy of sciences_ 114, 13 (2017), 3521–3526.
  * Li et al. (2017) Jing Li, Pengjie Ren, Zhumin Chen, Zhaochun Ren, Tao Lian, and Jun Ma. 2017.  Neural attentive session-based recommendation. In _Proceedings of the 2017 ACM on Conference on Information and Knowledge Management_. 1419–1428.
  * Lin et al. (2023) Jianghao Lin, Yanru Qu, Wei Guo, Xinyi Dai, Ruiming Tang, Yong Yu, and Weinan Zhang. 2023.  MAP: A Model-agnostic Pretraining Framework for Click-through Rate Prediction. In _Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining_. 1384–1395.
  * Lu et al. (2015) Jie Lu, Dianshuang Wu, Mingsong Mao, Wei Wang, and Guangquan Zhang. 2015.  Recommender system application developments: a survey.  _Decision support systems_ 74 (2015), 12–32.
  * Ma et al. (2022) Xuezhe Ma, Chunting Zhou, Xiang Kong, Junxian He, Liangke Gui, Graham Neubig, Jonathan May, and Luke Zettlemoyer. 2022.  Mega: moving average equipped gated attention.  _arXiv preprint arXiv:2209.10655_ (2022).
  * McAuley et al. (2015) Julian McAuley, Christopher Targett, Qinfeng Shi, and Anton Van Den Hengel. 2015.  Image-based recommendations on styles and substitutes. In _Proceedings of the 38th international ACM SIGIR conference on research and development in information retrieval_. 43–52.
  * Schak and Gepperth (2019) Monika Schak and Alexander Gepperth. 2019.  A study on catastrophic forgetting in deep LSTM networks. In _Artificial Neural Networks and Machine Learning–ICANN 2019: Deep Learning: 28th International Conference on Artificial Neural Networks, Munich, Germany, September 17–19, 2019, Proceedings, Part II 28_. Springer, 714–728.
  * Smith et al. (2022) Jimmy TH Smith, Andrew Warrington, and Scott W Linderman. 2022.  Simplified state space layers for sequence modeling.  _arXiv preprint arXiv:2208.04933_ (2022).
  * Srivastava et al. (2014) Nitish Srivastava, Geoffrey Hinton, Alex Krizhevsky, Ilya Sutskever, and Ruslan Salakhutdinov. 2014.  Dropout: a simple way to prevent neural networks from overfitting.  _The journal of machine learning research_ 15, 1 (2014), 1929–1958.
  * Sun et al. (2019) Fei Sun, Jun Liu, Jian Wu, Changhua Pei, Xiao Lin, Wenwu Ou, and Peng Jiang. 2019.  BERT4Rec: Sequential recommendation with bidirectional encoder representations from transformer. In _Proceedings of the 28th ACM international conference on information and knowledge management_. 1441–1450.
  * Tang and Wang (2018) Jiaxi Tang and Ke Wang. 2018.  Personalized top-n sequential recommendation via convolutional sequence embedding. In _Proceedings of the eleventh ACM international conference on web search and data mining_. 565–573.
  * Vaswani et al. (2017) Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Łukasz Kaiser, and Illia Polosukhin. 2017.  Attention is all you need.  _Advances in neural information processing systems_ 30 (2017).
  * Wang et al. (2021) Shoujin Wang, Longbing Cao, Yan Wang, Quan Z Sheng, Mehmet A Orgun, and Defu Lian. 2021.  A survey on session-based recommender systems.  _ACM Computing Surveys (CSUR)_ 54, 7 (2021), 1–38.
  * Yuan et al. (2022) Enming Yuan, Wei Guo, Zhicheng He, Huifeng Guo, Chengkai Liu, and Ruiming Tang. 2022.  Multi-behavior sequential transformer recommender. In _Proceedings of the 45th international ACM SIGIR conference on research and development in information retrieval_. 1642–1652.
  * Zhao et al. (2021) Wayne Xin Zhao, Shanlei Mu, Yupeng Hou, Zihan Lin, Yushuo Chen, Xingyu Pan, Kaiyuan Li, Yujie Lu, Hui Wang, Changxin Tian, et al. 2021\.  Recbole: Towards a unified, comprehensive and efficient framework for recommendation algorithms. In _proceedings of the 30th acm international conference on information & knowledge management_. 4653–4664.


Copyright for this paper by its authors. Use permitted under Creative Commons
License Attribution 4.0 International (CC BY 4.0).

AutomationXP23: Intervening, Teaming, Delegating Creating Engaging Automation
Experiences, April 23rd, Hamburg, Germany

[1]

[1]

[1]

[1]

[1]

[1]

[2]

[2]

[1]These authors contributed equally. [2]These authors advised equally.
[1]Corresponding author<EMAIL_ADDRESS>
# ReBound: An Open-Source 3D Bounding Box Annotation Tool for Active Learning

Wesley Chen  Andrew Edgley  Raunak Hota  Joshua Liu  Ezra Schwartz  Aminah
Yizar  Neehar Peri  James Purtilo

(2023)

###### Abstract

In recent years, supervised learning has become the dominant paradigm for
training deep-learning based methods for 3D object detection. Lately, the
academic community has studied 3D object detection in the context of
autonomous vehicles (AVs) using publicly available datasets such as nuScenes
and Argoverse 2.0. However, these datasets may have incomplete annotations,
often only labeling a small subset of objects in a scene. Although commercial
services exists for 3D bounding box annotation, these are often prohibitively
expensive. To address these limitations, we propose ReBound, an open-source 3D
visualization and dataset re-annotation tool that works across different
datasets. In this paper, we detail the design of our tool and present survey
results that highlight the usability of our software. Further, we show that
ReBound is effective for exploratory data analysis and can facilitate active-
learning. Our code and documentation is available on GitHub.

###### keywords:

Autonomous Driving, Active Learning, 3D Annotation Tools, Human Computer
Interaction, Data Visualization

##  1 Introduction

3D object detection is a critical component of the autonomous vehicle (AV)
perception stack [1, 2]. To facilitate research in 3D object detection, the AV
industry has released large-scale 3D annotated multimodal datasets [2, 3, 4].
These datasets include LiDAR sweeps and multi-camera RGB images from diverse
driving logs, which captures detailed information about the surrounding
environment. Crucially, objects of interest are annotated by drawing 3D
bounding boxes and labeling them as part of a particular category.
Contemporary 3D detectors [5, 6, 7, 8, 9] are trained using supervised
learning, and are limited by the annotations provided with the dataset. For
example, objects in the nuScenes dataset are inconsistently labeled beyond
50m, making it challenging to evaluate long-range detection. Further, the
nuScenes dataset does not label street signs and traffic lights, which are
critical for safe navigation.

Manually re-annotating datasets with new 3D bounding box annotations is
particularly challenging. Annotating 3D bounding boxes using 2D RGB images is
difficult because it is not possible to accurately estimate bounding-box
depth. Similarly, annotating 3D bounding boxes using LiDAR point clouds is
difficult because LiDAR returns are sparse, making it tough to identify
individual objects at long-range and in cluttered scenes, as shown in Fig. 2.
Although commercial services can annotate 3D data at scale, it is often
prohibitively expensive. As a result, several tools have been created for
quick, efficient, and easy point cloud annotation [10, 11, 12, 13], with the
intention of making it simpler for researchers to create their own datasets.
However, these tools do not support a wide variety of data formats, multi-
modal exploratory data analysis and active learning.

Figure 1: We present ReBound, an open-source 3D annotation tool that allows
users to add, delete, and modify annotations from existing datasets or model
predictions to support active learning.

In this paper, we introduce ReBound, an open-source annotation tool which aims
to simplify the process of multi-modal 3D visualization and bounding box
annotation. Our framework is designed to provide users with the ability to
import and modify annotations from various datasets. To achieve this, we
propose a generic data type that can be extended to accommodate different data
formats. Additionally, ReBound enables active learning by allowing users to
correct bounding box predictions and labels, create new custom labels, analyze
model predictions, and export new annotations back to dataset-specific formats
for model re-training. We measure the effectiveness of our tool through user
studies, focusing on the ease of using the data conversion and annotation
editing features. Our results show that the tool is both intuitive to use and
useful for rapidly adding new annotations.

Figure 2: Long-range 3D detection is crucial for safe navigation in autonomous
vehicles. Modern benchmarks have greatly improved 3D object detection.
However, standard benchmarks like nuScenes only evaluate detections up to 50m,
masking the poor far-field object detection performance of 3D detectors (b &
c). Existing benchmarks do not sufficiently annotate far-field objects (d &
e), partly due to fewer LiDAR returns. Our proposed annotation framework
facilitates far-field annotation, and will enable researchers to better study
this problem. We adopt this figure from [14].

##  2 Related Work

In this section we present a brief overview of existing 3D annotation tools
and active learning methods for object detection.

3D Annotation Tools. Modern 3D detectors require diverse, large-scale 3D
annotations for supervised learning. A variety of tools have been created to
address this challenge [15], including labelCloud[10], SAnE[11], LATTE[12],
and 3D BAT[13], which all allow for efficient manual or semi-automatic dataset
annotation. However, unlike these existing tools, ReBound also allows users to
manually update annotations from existing datasets for evolving use cases like
detecting a new category or updating annotations for far-field objects. Our
tool currently supports re-annotation for the nuScenes[16], Waymo [17], and
Argoverse 2.0[18] datasets.

Active Learning for 3D Object Detection. Although manually labeling tens of
thousands of images or LiDAR sweeps can be prohibitively expensive, recent
work in active learning suggests that we can bootstrap deep learning models by
iteratively annotating and training on a targeted subset of informative
examples to significantly improve model performance. Selecting the most
relevant data to be annotated by human experts is a primary challenge for
active learning. Recent works define a scoring function, using either an
uncertainty-based [19, 20, 21, 22, 23, 24, 25, 26] or a diversity-based
approach [27, 26], to determine the most informative training samples. When
selecting samples based on uncertainty, informativeness is measured by the
predictive uncertainty, and the samples with the highest uncertainty are
provided to the human annotators. When selecting samples based on diversity,
scores are assigned based on spatial and temporal diversity [27]. Both the
uncertainty-based and diversity-based approaches have been used for
bootstrapping 3D object detection, but the uncertainty-based approach has been
shown to be more effective in practice. More recently, [26] attempts to
combine both strategies. Despite substantial work in active learning for image
recognition and 2D object detection, exploration into active learning for 3D
object detection is limited. ReBound helps support active learning by allowing
users to filtering predicted annotations based on detection confidence score,
which is useful for measuring uncertainty.

##  3 ReBound: 3D Bounding-Box Re-Annotation

In this section, we describe the functionality and software architecture of
ReBound. Our code, documentation, and videos demonstrating our tool are
available on GitHub.

Data Format. ReBound currently supports three different datasets: nuScenes,
Waymo Open Dataset, and Argoverse 2.0. Annotations for these three datasets
are converted into our generic ReBound format using a separate command line
tool built in Python. Specifically, the generic data format captures the
minimum information required to annotate a 3D bounding box (e.g. object
center, box size, box rotation, and sensor extrinsics), which allows us to
scale this format across multiple datasets. Users can support new datasets by
extending our command line tool for their use case.

Data Visualization. The visualization tool has three windows: a control
window, a point cloud viewer, and an RGB image viewer, as shown in Figure 1.
The control window allows the user to navigate through different frames in a
driving log, switch between different camera views in the RGB image window,
and filter through both ground truth annotations and model predictions. The
point cloud viewer displays the point cloud corresponding to the current
frame, as well as all ground truth annotations and predictions (if available)
for that frame as wireframe cuboids. The user can rotate, translate, and zoom
in/out to interact with the scene. We use Open3D as our rendering back-end as
this natively supports 3D rendering on top of LiDAR sweeps and RGB images.

Editing Tool. Users can directly click on a location in the point cloud window
to add, edit, or delete an annotation. To edit a box’s properties, users must
first click on a desired cuboid in the point cloud window. This will highlight
the bounding box and display the position, rotation, size and annotation class
of the selected box in the control window. These fields can all be directly
updated, allowing the user to make precise changes. Users can also make coarse
changes to the location of the selected box by clicking and dragging the box
within the LiDAR viewer. To make it easier to interact with 3D objects, the
mouse tool is restricted to two different control modes:

  * •

Horizontal Translation Mode: The user can articulate objects along the $X$ and
$Y$ axes with the $Z$ axis locked

  * •

Vertical Translation Mode: The user can click and drag objects along the $Z$
axis and rotate objects about the $Z$ axis.

Finer-grained modifications can be manually entered in the control window.
Users can also delete the selected bounding box, and create new bounding boxes
with a single mouse-click. All bounding box transformations are updated in
both the LiDAR and RGB image viewers in real-time. In practice, the editing
tool can be used to update bounding boxes that are too large, are misaligned
with the true object location (c.f. Figure 4), update mis-classified objects,
and add new categories like stop signs and traffic lights.

Figure 3: Our tool supports data conversion, visualization, and modification
of the nuScenes, Waymo and Argoverse 2.0 datasets. First, dataset-specific RGB
images, LiDAR sweeps, sensor extrinsics, and annotations are converted to the
ReBound data format using our conversion scripts and visualized in the GUI.
Using this GUI, we can add, edit, or delete existing annotations. Lastly, we
can export data from the ReBound data format back to the respective dataset-
specific formats using the provided export scripts.

Exporting. Users can export updated annotations back to the original dataset-
specific formats using a command line tool. Importantly, ReBound only exports
the modified bounding box annotations back to the original format. Since
ReBound does not modify the LiDAR sweeps or RGB images, we can directly use
the data from the original dataset. We find that this dramatically increases
data export speed. Interestingly, since we convert all datasets into a unified
format, we can easily export annotations between datasets. Concretely, this
means we can convert nuScenes annotations into the Argoverse 2.0 format using
the ReBound format as an intermediary, making it easier to evaluate models
across different datasets. Further, this can facilitate future research in
model robustness and domain adaptation between different datasets. Similar to
the data conversion tool that converts from dataset-specific formats to the
ReBound format, users that wish to support a new dataset can extend our
command line tool for their use case.

##  4 Experiments

In this section we present the results of our user survey to evaluate the
effectiveness of our tool.

Ease of Use. We conducted surveys to evaluate the tool’s ease of use in
facilitating data conversion between data formats and bounding box adjustments
among both people familiar and unfamiliar with autonomous vehicle datasets and
3D annotation tools.

We asked ten participants to perform 13 tasks (shown in Table 1) over a video
conference. Participants were shown a demonstration of the tool before being
asked to complete a new, but related task using the tool. This method of
demonstration allows us to avoid re-downloading the datasets and re-installing
the software on a new computer for each trial. After completing all 13 tasks,
users were asked to complete a questionnaire designed to rate different parts
of the user experience on a scale of one to five (five is highest).

Based on our survey, we identified that the greatest challenge for users was
translating and rotating bounding boxes. Both of these operations require
fine-grained controls and navigating in a 3D space, which may be not be
intuitive for some users. However these features also had high standard
deviation compared to the other tasks, indicating that the utility of our tool
may depend on prior use of 3D visualization and annotation tools. In general,
participants stated that the application can be useful for autonomous vehicle
research.

Task List Avg. Std. Dev Convert nuScenes data to the ReBound generic type 4.1
0.74 Enter annotation mode 4.2 0.79 Add a new annotation bounding box 3.8 1.4
Add a custom annotation type 4.7 0.67 Select an existing bounding box 4.5 0.97
Translate a bounding box 3.8 1.03 Rotate a bounding box 3.6 1.17 Change the
label of a bounding box 4.8 0.42 Delete a bounding box 4.8 0.63 Use the
control viewer to edit an annotation 4 1.15 Save modified annotations 4.5 0.71
Exit the application 4.5 0.53 Export data back to the nuScenes data format 3.9
0.88

Table 1: Our survey results show that users found it easier to create, delete,
and modify annotations but found it more difficult to rotate, translate, and
export bounding boxes from the ReBound data format to dataset-specific
formats.

Bounding Box Adjustments. We visualize ground truth annotations in both the
LiDAR and RGB image viewers. We are able to identify examples where the ground
truth annotation is misaligned with the real object, as shown in Figure 4,
highlighting a potential application of our tool. This is practically
meaningful as training 3D detectors with noisy labels can result in degraded
model performance.

Figure 4: We visualize a ground truth annotation from the Argoverse 2.0
dataset (left). We find that the annotation is misaligned with the true car
location. We adjust this ground truth annotation using ReBound (right).
Visually inspecting the updated bounding box in both the LiDAR and RGB image
viewer confirms that the updated annotation correctly localizes the car.

##  5 Conclusion

The academic community studies 3D object detection using publicly available AV
datasets, which often provide LiDAR sweeps, RGB images, and 3D bounding box
annotations. However, these datasets are limited by the annotations provided
by their curators. In some cases, we find that these annotations may be
incorrect or incomplete. Existing annotation tools cannot be easily used to
update annotations for existing datasets, which makes it difficult to fix
incorrect annotations, annotate new categories of interest, or iteratively
improve 3D object detection models through active learning. In this paper, we
propose ReBound, an open-source alternative that simplifies the process of
bounding box re-annotation for existing datasets. We validate the utility of
our tool through user surveys and find that users are able to rapidly add,
modify, and delete 3D bounding box annotations.

###### Acknowledgements.

Author James Purtilo was supported by N000142112821 while working on this
project. This work is supported by the SEAM Lab at the University of Maryland
and the Argo AI Center for Autonomous Vehicle Research at Carnegie Mellon
University.

## References

  * Geiger et al. [2012] A. Geiger, P. Lenz, R. Urtasun,  Are we ready for autonomous driving? the kitti vision benchmark suite,  in: IEEE Conference on Computer Vision and Pattern Recognition, 2012.
  * Caesar et al. [2020] H. Caesar, V. Bankiti, A. H. Lang, S. Vora, V. E. Liong, Q. Xu, A. Krishnan, Y. Pan, G. Baldan, O. Beijbom,  nuscenes: A multimodal dataset for autonomous driving,  in: Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2020.
  * Wilson et al. [2021] B. Wilson, W. Qi, T. Agarwal, J. Lambert, J. Singh, S. Khandelwal, B. Pan, R. Kumar, A. Hartnett, J. K. Pontes, D. Ramanan, P. Carr, J. Hays,  Argoverse 2: Next generation datasets for self-driving perception and forecasting,  in: Neural Information Processing Systems Datasets and Benchmarks Track, 2021.
  * Sun et al. [2020] P. Sun, H. Kretzschmar, X. Dotiwalla, A. Chouard, V. Patnaik, P. Tsui, J. Guo, Y. Zhou, Y. Chai, B. Caine, V. Vasudevan, W. Han, J. Ngiam, H. Zhao, A. Timofeev, S. Ettinger, M. Krivokon, A. Gao, A. Joshi, Y. Zhang, J. Shlens, Z. Chen, D. Anguelov,  Scalability in perception for autonomous driving: Waymo open dataset,  in: IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
  * Yin et al. [2020] T. Yin, X. Zhou, P. Krähenbühl,  Center-based 3d object detection and tracking,  arXiv preprint arXiv:2006.11275 (2020).
  * Peri et al. [2022] N. Peri, A. Dave, D. Ramanan, S. Kong,  Towards long-tailed 3d detection,  in: Conference on Robot Learning (CoRL), 2022\.
  * Lang et al. [2019] A. H. Lang, S. Vora, H. Caesar, L. Zhou, J. Yang, O. Beijbom,  Pointpillars: Fast encoders for object detection from point clouds,  in: IEEE Conference on Computer Vision and Pattern Recognition (CVPR), 2019.
  * Zhu et al. [2019] B. Zhu, Z. Jiang, X. Zhou, Z. Li, G. Yu,  Class-balanced grouping and sampling for point cloud 3d object detection,  arXiv preprint arXiv:1908.09492 (2019).
  * Peri et al. [2022] N. Peri, J. Luiten, M. Li, A. Osep, L. Leal-Taixe, D. Ramanan,  Forecasting from lidar via future object detection,  arXiv:2203.16297 (2022).
  * Sager et al. [2021] C. Sager, P. Zschech, N. Kühl,  labelcloud: A lightweight domain-independent labeling tool for 3d object detection in point clouds,  CoRR abs/2103.04970 (2021).
  * Arief et al. [2020] H. A. Arief, M. Arief, G. Zhang, Z. Liu, M. Bhat, U. G. Indahl, H. Tveite, D. Zhao,  Sane: Smart annotation and evaluation tools for point cloud data,  IEEE Access 8 (2020) 131848–131858.
  * Wang et al. [2019] B. Wang, V. Wu, B. Wu, K. Keutzer,  LATTE: accelerating lidar point cloud annotation via sensor fusion, one-click annotation, and tracking,  in: 2019 IEEE Intelligent Transportation Systems Conference, ITSC 2019, Auckland, New Zealand, October 27-30, 2019, IEEE, 2019, pp. 265–272.
  * Zimmer et al. [2019] W. Zimmer, A. Rangesh, M. M. Trivedi,  3d BAT: A semi-automatic, web-based 3d annotation toolbox for full-surround, multi-modal data streams,  in: 2019 IEEE Intelligent Vehicles Symposium, IV 2019, Paris, France, June 9-12, 2019, IEEE, 2019, pp. 1816–1821.
  * Gupta et al. [2022] S. Gupta, J. Kanjani, M. Li, F. Ferroni, J. Hays, D. Ramanan, S. Kong,  Far3det: Towards far-field 3d detection,  in: NeurIPS, 2022.
  * Li and Ibañez-Guzmán [2020] Y. Li, J. Ibañez-Guzmán,  Lidar for autonomous driving: The principles, challenges, and trends for automotive lidar and perception systems,  IEEE Signal Process. Mag. 37 (2020) 50–61.
  * Caesar et al. [2020] H. Caesar, V. Bankiti, A. H. Lang, S. Vora, V. E. Liong, Q. Xu, A. Krishnan, Y. Pan, G. Baldan, O. Beijbom,  nuscenes: A multimodal dataset for autonomous driving,  in: 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020, Computer Vision Foundation / IEEE, 2020, pp. 11618–11628.
  * Sun et al. [2020] P. Sun, H. Kretzschmar, X. Dotiwalla, A. Chouard, V. Patnaik, P. Tsui, J. Guo, Y. Zhou, Y. Chai, B. Caine, V. Vasudevan, W. Han, J. Ngiam, H. Zhao, A. Timofeev, S. Ettinger, M. Krivokon, A. Gao, A. Joshi, Y. Zhang, J. Shlens, Z. Chen, D. Anguelov,  Scalability in perception for autonomous driving: Waymo open dataset,  in: 2020 IEEE/CVF Conference on Computer Vision and Pattern Recognition, CVPR 2020, Seattle, WA, USA, June 13-19, 2020, Computer Vision Foundation / IEEE, 2020, pp. 2443–2451.
  * Wilson et al. [2021] B. Wilson, W. Qi, T. Agarwal, J. Lambert, J. Singh, S. Khandelwal, B. Pan, R. Kumar, A. Hartnett, J. K. Pontes, D. Ramanan, P. Carr, J. Hays,  Argoverse 2: Next generation datasets for self-driving perception and forecasting,  in: J. Vanschoren, S. Yeung (Eds.), Proceedings of the Neural Information Processing Systems Track on Datasets and Benchmarks 1, NeurIPS Datasets and Benchmarks 2021, December 2021, virtual, 2021.
  * Haussmann et al. [2020] E. Haussmann, M. Fenzi, K. Chitta, J. Ivanecky, H. Xu, D. Roy, A. Mittel, N. Koumchatzky, C. Farabet, J. M. Alvarez,  Scalable active learning for object detection,  in: IEEE Intelligent Vehicles Symposium, IV 2020, Las Vegas, NV, USA, October 19 - November 13, 2020, IEEE, 2020, pp. 1430–1435.
  * Yuan et al. [2021] T. Yuan, F. Wan, M. Fu, J. Liu, S. Xu, X. Ji, Q. Ye,  Multiple instance active learning for object detection,  in: CVPR, 2021.
  * Feng et al. [2019] D. Feng, X. Wei, L. Rosenbaum, A. Maki, K. Dietmayer,  Deep active learning for efficient training of a lidar 3d object detector,  in: 2019 IEEE Intelligent Vehicles Symposium, IV 2019, Paris, France, June 9-12, 2019, IEEE, 2019, pp. 667–674.
  * Roy et al. [2018] S. Roy, A. Unmesh, V. P. Namboodiri,  Deep active learning for object detection,  in: British Machine Vision Conference 2018, BMVC 2018, Newcastle, UK, September 3-6, 2018, BMVA Press, 2018, p. 91.
  * Hekimoglu et al. [2022] A. Hekimoglu, M. Schmidt, A. Marcos-Ramiro, G. Rigoll,  Efficient active learning strategies for monocular 3d object detection,  in: 2022 IEEE Intelligent Vehicles Symposium, IV 2022, Aachen, Germany, June 4-9, 2022, IEEE, 2022, pp. 295–302.
  * Schmidt et al. [2020] S. Schmidt, Q. Rao, J. Tatsch, A. C. Knoll,  Advanced active learning strategies for object detection,  in: IEEE Intelligent Vehicles Symposium, IV 2020, Las Vegas, NV, USA, October 19 - November 13, 2020, IEEE, 2020, pp. 871–876.
  * Meyer and Kuschk [2019] M. Meyer, G. Kuschk,  Automotive radar dataset for deep learning based 3d object detection,  in: 2019 16th European Radar Conference (EuRAD), 2019\.
  * Luo et al. [2023] Y. Luo, Z. Chen, Z. Wang, X. Yu, Z. Huang, M. Baktashmotlagh,  Exploring active 3d object detection from a generalization perspective,  CoRR abs/2301.09249 (2023).
  * Liang et al. [2022] Z. Liang, X. Xu, S. Deng, L. Cai, T. Jiang, K. Jia,  Exploring diversity-based active learning for 3d object detection in autonomous driving,  CoRR abs/2205.07708 (2022).