Learning to Predict 3D Rotational Dynamics from Images of a Rigid Body with Unknown Mass Distribution

Abstract

In many real-world settings, image observations of freely rotating 3D rigid bodies may be available when low-dimensional measurements are not. However, the high-dimensionality of image data precludes the use of classical estimation techniques to learn the dynamics. The usefulness of standard deep learning methods is also limited, because an image of a rigid body reveals nothing about the distribution of mass inside the body, which, together with initial angular velocity, is what determines how the body will rotate. We present a physics-based neural network model to estimate and predict 3D rotational dynamics from image sequences. We achieve this using a multi-stage prediction pipeline that maps individual images to a latent representation homeomorphic to $\mathbf{SO}\left(3\right)$, computes angular velocities from latent pairs, and predicts future latent states using the Hamiltonian equations of motion. We demonstrate the efficacy of our approach on new rotating rigid-body datasets of sequences of synthetic images of rotating objects, including cubes, prisms and satellites, with unknown uniform and non-uniform mass distributions. Our model outperforms competing baselines on our datasets, producing better qualitative predictions and reducing the error observed for the state-of-the-art Hamiltonian Generative Network by a factor of 2.

1. Introduction

The study and control of a range of systems can benefit from the means to predict the rotational dynamics of 3D rigid bodies that are only observed through images. A compelling example is the navigation and control of space robotic systems that interact with resident space objects (RSOs). RSOs are natural or designed freely rotating rigid bodies that orbit a planet or moon. Space robotic system missions that involve interaction with RSOs include collecting samples from an asteroid [1], servicing a malfunctioning satellite [2], and removing active space debris [3]. A challenge is that space robotic systems may have limited information on the mass distribution of RSOs. However, they do typically have onboard cameras to observe sequences of RSO movement. Thus, learning to predict the dynamics of the RSOs from onboard images can make a difference for mission success.

Whether a freely rotating 3D rigid body tumbles unstably or spins stably depends on the distribution of mass inside the body and the body’s initial angular velocity (compare Figure 1a and Figure 1b). This means that to predict the body’s rotational dynamics, it is not enough to know the external geometry of the body. That would be insufficient, for instance, to predict the different behavior of two bodies with the same external geometry and different internal mass distribution. Even if the bodies start at the same initial angular velocity, one body could tumble or wobble while the other spins stably (compare Figure 1b and Figure 1d).

Figure 1 shows four simulations of a freely rotating rigid body that illustrate the role of mass distribution and initial velocity. The distribution of mass determines $\mathbf{J}\in {\mathbb{R}}^{3\times 3}$, where $\mathbf{J}$ is the moment-of-inertia matrix for a 3D rigid body expressed with respect to the body-fixed frame, i.e., an orthonormal reference frame $\mathcal{B}=\{\mathbf{i},\mathbf{j},\mathbf{k}\}$ fixed to the body with origin at the body’s center of mass (see Appendix A.1 for details). Figure 1a–c all have the same moment-of-inertia matrix $\mathbf{J}={\mathbf{J}}_1$, which corresponds to that of a rectangular prism with uniform mass distribution (see

Table A1. Table containing the moment-of-inertia tensors, inverse moment-of-inertia tensors, and principal axes used to generate training data for each object.

Object	Moment-of-Inertia Tensor	Inverse Moment-of-Inertia Tensor	Principal Axes
Uniform Cube	${\mathbf{J}}_0=\frac{1}{3}\left(\begin{array}{ccc}1& 0& 0\\ 0& 1& 0\\ 0& 0& 1\end{array}\right)$	${\mathbf{J}}_0^{-1}=\left(\begin{array}{ccc}3& 0& 0\\ 0& 3& 0\\ 0& 0& 3\end{array}\right)$	$\left\{\left(\begin{array}{c}1\\ 0\\ 0\end{array}\right),\left(\begin{array}{c}0\\ 1\\ 0\end{array}\right),\left(\begin{array}{c}0\\ 0\\ 1\end{array}\right)\right\}$
Uniform Prism	${\mathbf{J}}_1=\left(\begin{array}{ccc}0.42& 0& 0.\\ 0& 1.41& 0\\ 0& 0& 1.67\end{array}\right)$	${\mathbf{J}}_1^{-1}=\left(\begin{array}{ccc}2.40& 0& 0\\ 0& 0.71& 0\\ 0& 0& 0.60\end{array}\right)$	$\left\{\left(\begin{array}{c}1\\ 0\\ 0\end{array}\right),\left(\begin{array}{c}0\\ 1\\ 0\end{array}\right),\left(\begin{array}{c}0\\ 0\\ 1\end{array}\right)\right\}$
Non-uniform Cube	${\mathbf{J}}_2=\left(\begin{array}{ccc}0.17& 0& -0.56\\ 0& 0.17& -0.99\\ -0.56& -0.99& 0.17\end{array}\right)$	${\mathbf{J}}_2^{-1}=\left(\begin{array}{ccc}4.53& -2.62& -0.44\\ -2.62& 1.34& -0.78\\ -0.44& -0.78& -0.13\end{array}\right)$	$\left\{\left(\begin{array}{c}-0.35\\ -0.62\\ -0.71\end{array}\right),\left(\begin{array}{c}-0.87\\ 0.49\\ 0\end{array}\right),\left(\begin{array}{c}-0.35\\ -0.62\\ 0.71\end{array}\right)\right\}$
Non-uniform Prism	${\mathbf{J}}_3=\left(\begin{array}{ccc}0.47& 0& -0.28\\ 0& 1.61& -0.49\\ -0.28& -0.49& 1.83\end{array}\right)$	${\mathbf{J}}_3^{-1}=\left(\begin{array}{ccc}2.37& 0.12& 0.39\\ 0.12& 0.68& 0.20\\ 0.39& 0.20& 0.66\end{array}\right)$	$\left\{\left(\begin{array}{c}-0.35\\ -0.62\\ -0.71\end{array}\right),\left(\begin{array}{c}-0.87\\ 0.49\\ 0\end{array}\right),\left(\begin{array}{c}-0.35\\ -0.62\\ 0.71\end{array}\right)\right\}$
CALIPSO	${\mathbf{J}}_4=\left(\begin{array}{ccc}0.33& 0& 0\\ 0& 0.50& 0\\ 0& 0& 1.0\end{array}\right)$	${\mathbf{J}}_4^{-1}=\left(\begin{array}{ccc}3.0& 0& 0\\ 0& 2.0& 0\\ 0& 0& 1.0\end{array}\right)$	$\left\{\left(\begin{array}{c}1\\ 0\\ 0\end{array}\right),\left(\begin{array}{c}0\\ 1\\ 0\end{array}\right),\left(\begin{array}{c}0\\ 0\\ 1\end{array}\right)\right\}$
CloudSat	${\mathbf{J}}_5=\left(\begin{array}{ccc}0.33& 0& 0\\ 0& 0.50& 0\\ 0& 0& 1.0\end{array}\right)$	${\mathbf{J}}_5^{-1}=\left(\begin{array}{ccc}3.0& 0& 0\\ 0& 2.0& 0\\ 0& 0& 1.0\end{array}\right)$	$\left\{\left(\begin{array}{c}1\\ 0\\ 0\end{array}\right),\left(\begin{array}{c}0\\ 1\\ 0\end{array}\right),\left(\begin{array}{c}0\\ 0\\ 1\end{array}\right)\right\}$

in Appendix A.2). Steady spin about the longest and shortest principal axes is stable and about the intermediate principal axis is unstable (see Appendix A.1). So, if the initial angular velocity is near the unstable solution, the body tumbles (Figure 1a), whereas if it is near the stable axis, the body spins (Figure 1b). This is independent of the external geometry, which explains why the satellite in Figure 1c spins identically to the prism in Figure 1b. In Figure 1d, mass is non-uniformly distributed, such that $\mathbf{J}={\mathbf{J}}_3$ (see Table A1 in Appendix A.2) and the same initial velocity as in Figure 1b is no longer close to a stable solution, which explains why the prism wobbles.

Predicting 3D rigid body rotational dynamics is possible if the body’s mass distribution can be learned from observations of the body in motion. This is easier if the observations consist of low-dimensional data, e.g., measurements of the body’s angular velocity and the rotation matrix that defines the body’s orientation. It is much more challenging, however, if the only available measurements consist of images of the body in motion, as in the case of remote observations of a satellite or asteroid or space debris.

We address the challenge of learning and predicting 3D rotational dynamics from image sequences of a rigid body with unknown mass distribution and unknown initial angular velocity. To do so we design a neural network model that leverages Hamiltonian structure associated with 3D rigid body dynamics. We show how our approach outperforms applicable methods from the existing literature.

Deep learning has proven to be an effective tool to learn dynamics from images. Previous work [4,5,6] has made significant progress in using physics-based priors to learn dynamics from images of 2D rigid bodies, such as a pendulum. Learning dynamics of 3D rigid-body motion has also been explored with various types of input data [7,8,9]. We believe our method is the first to use the Hamiltonian formalism to learn 3D rigid-body rotational dynamics from images.

In this work, we introduce a model, with architecture depicted in Figure 2, that (1) learns 3D rigid-body rotational dynamics from images, (2) predicts future image sequences in time, and (3) generates a low-dimensional, interpretable representation of the latent state. During training, our model encodes a sequence of images (input) to a sequence of latent orientations (Figure 2a). The sequence of orientations is processed by two pathways. In one, the sequence is decoded to a sequence of images which are used to compute the auto-encoding reconstruction loss (Figure 2c). In the other, the first element of the sequence is processed by the dynamics pipeline. The resulting sequence is decoded to a sequence of images, which are used to compute the dynamics-based reconstruction loss (Figure 2d). During inference, our model encodes a pair of images (input) to a single latent orientation (Figure 2b). This latent orientation is processed by the dynamics pipeline and decoding pipeline resulting in a predicted image sequence (Figure 2d).

Our model incorporates the Hamiltonian formulation of the dynamics as an inductive bias to facilitate learning the moment-of-inertia matrix, ${\mathbf{J}}_{\phi }$, and an auto-encoding map between images and the special orthogonal group $\mathbf{SO}\left(3\right)=\left\{\mathbf{R}\in {\mathbb{R}}^{3\times 3}|{\mathbf{R}}^T\mathbf{R}={\mathbf{I}}_3,det\left(\mathbf{R}\right)=+1\right\}$. $\mathbf{SO}\left(3\right)$ represents the space of all 3D rotations: the orientation of the rigid body at time t is described by the rotation matrix $\mathbf{R}\left(t\right)\in \mathbf{SO}\left(3\right)$ that maps points on the body from body frame coordinates to inertial frame coordinates at time t.

The efficacy of our approach is demonstrated through long-term image prediction on synthetic datasets. Due to the scarcity of appropriate datasets, we have created publicly available, synthetic datasets of rotating objects (e.g., cubes, prisms, and satellites) applicable for evaluation of our model, as well as other tasks on 3D rigid-body rotation such as pose estimation.

2. Related Work

A growing body of work incorporates Hamiltonian and Lagrangian formalisms to improve the accuracy and interpretability of learned representations in neural network-based dynamical systems forecasting [10,11,12]. Greydanus et al. [10] predict symplectic gradients of a Hamiltonian system using a Hamiltonian parameterized by a neural network. They show that the Hamiltonian neural network (HNN) predicts the evolution of conservative systems better than a baseline black-box model. Chen et al. [11] improve the long-term prediction performance of [10] by minimizing the mean-squared error (MSE) between ground-truth and predicted state trajectories rather than one-step symplectic gradients. Cranmer et al. [12] propose parameterization of the system Lagrangian by a neural network arguing that momentum coordinates may be difficult to compute in some settings. Each of the aforementioned learn from sequences of phase-space measurements; our model learns from images.

The authors of [4,5,6] leverage Hamiltonian and Lagrangian neural networks to learn the dynamics of 2D rigid bodies (e.g., the planar pendulum) from image sequences. Zhong and Leonard [4] introduce a coordinate-aware variational autoencoder (VAE) [13] with a latent Lagrangian neural network (LNN) which learns the underlying dynamics and facilitates control. Allen-Blanchette et al. [6] use a latent LNN in an auto-encoding neural network to learn dynamics without control or prior knowledge of the configuration-space structure. Toth et al. [5] use a latent HNN in a VAE to learn dynamics without control, prior knowledge of the configuration-space structure or dimension. Similarly to Toth et al. [5], we use a latent HNN to learn dynamics. Distinctly, however, we consider 3D rigid body dynamics and incorporate prior knowledge of the configuration-space structure to ensure interpretability of the learned representations.

Others have considered the problem of learning 3D rigid-body dynamics [7,8,9]. Byravan and Fox [7] uses point-cloud data and action vectors (forces) as inputs to a black-box neural network to predict the resulting $\mathbf{SE}\left(3\right)$ transformation matrix, which represents the motion of objects within the input scene. The special Euclidean group $\mathbf{SE}\left(3\right)=\left\{(\mathbf{R},\mathbf{r})|\mathbf{R}\in \mathbf{SO}\left(3\right),\mathbf{r}\in {\mathbb{R}}^3\right\}$ represents the space of all 3D rotations and translations: the orientation and position of the rigid body at time t is described by the rotation matrix and vector pair $\left(\mathbf{R}\right(t),\mathbf{r}(t\left)\right)\in \mathbf{SE}\left(3\right)$ that maps points on the body from body frame coordinates to inertial frame coordinates at time t. Peretroukhin et al. [8] create a novel symmetric matrix representation of $\mathbf{SO}\left(3\right)$ and incorporate it into a neural network to perform orientation prediction on synthetic point-cloud data and images. Duong and Atanasov [9] use low-dimensional measurement data (i.e., the rotation matrix and angular momenta) to learn rigid body dynamics on $\mathbf{SO}\left(3\right)$ and $\mathbf{SE}\left(3\right)$ for control.

The combination of deep learning with physics-based priors allows models to learn dynamics from high-dimensional data such as images [4,5,6]. However, as far as we know, our method is the first to use the Hamiltonian formalism to learn 3D rigid-body rotational dynamics from images.

3. Background

3.1. The ${\mathcal{S}}^2\times {\mathcal{S}}^2$ Parameterization of 3D Rotation Group SO(3)

The ${\mathcal{S}}^2\times {\mathcal{S}}^2$ parameterization of the 3D rotation group $\mathbf{SO}\left(3\right)$ is a surjective and differentiable mapping with a continuous right inverse [14]. Define the n-sphere: ${\mathcal{S}}^n=\left\{\mathbf{v}\in {\mathbb{R}}^{(n+1)}|v_1^2+v_2^2+\cdots +v_{n+1}^2=1\right\}$. The ${\mathcal{S}}^2\times {\mathcal{S}}^2$ parameterization of $\mathbf{SO}\left(3\right)$ is given by $(\mathbf{u},\mathbf{v})\mapsto ({\mathbf{w}}_1,{\mathbf{w}}_2,{\mathbf{w}}_3)$ with ${\mathbf{w}}_1=\mathbf{u},{\mathbf{w}}_2=\mathbf{v}-\mathbf{v}\langle \mathbf{u},\mathbf{v}\rangle ,{\mathbf{w}}_3={\mathbf{w}}_1\times {\mathbf{w}}_2$, where ${\mathbf{w}}_i$ are renormalized to have unit norm.

Intuitively, this mapping constructs an orthonormal frame from the unit vectors $\mathbf{u}$ and $\mathbf{v}$ by Gram–Schmidt orthogonalization. The right inverse of the parameterization is given by $({\mathbf{w}}_1,{\mathbf{w}}_2,{\mathbf{w}}_3)\mapsto ({\mathbf{w}}_1,{\mathbf{w}}_2)$. Other parameterizations of $\mathbf{SO}\left(3\right)$, such as the exponential map ($\mathfrak{so}\left(3\right)\mapsto \mathbf{SO}\left(3\right)$) and the quaternion map (${\mathcal{S}}^3\mapsto \mathbf{SO}\left(3\right)$), do not have continuous inverses and therefore are more difficult to use in deep manifold regression [14,15,16,17].

3.2. 3D Rotating Rigid-Body Kinematics

The orientation of a rotating 3D rigid body $\mathbf{R}\left(t\right)\in \mathbf{SO}\left(3\right)$ changing over time t can be computed from body angular velocity $\mathsf{\Omega }\left(t\right)\in {\mathbb{R}}^3$, i.e., the angular velocity of the body expressed with respect to the body frame $\mathcal{B}$, at time $t\ge 0$ using the kinematic equations given by the time-rate-of-change of $\mathbf{R}\left(t\right)$ shown in Equation (A3). For computational purposes, 3D rigid-body rotational kinematics are commonly expressed in terms of the quaternion representation $\mathbf{q}\left(t\right)\in {\mathcal{S}}^3$ of the rigid-body orientation $\mathbf{R}\left(t\right)$. The kinematics (A3), written in terms of quaternions [18], are

$$\frac{d\mathbf{q}\left(t\right)}{dt}=\mathbf{Q}\left(\mathsf{\Omega }\left(t\right)\right)\mathbf{q}\left(t\right),\mathbf{Q}\left(\mathsf{\Omega }\right)=\left(\begin{array}{cc}-{\mathsf{\Omega }}_{\times }& \mathsf{\Omega }\\ -{\mathsf{\Omega }}^T& 0\end{array}\right),$$

(1)

where ${\mathsf{\Omega }}_{\times }$ is the $3\times 3$ skew-symmetric matrix defined by $\left({\mathsf{\Omega }}_{\times }\right)\mathit{y}=\mathsf{\Omega }\times \mathit{y}$ for $\mathit{y}\in {\mathbb{R}}^3$.

3.3. 3D Rigid-Body Dynamics in Hamiltonian Form

The canonical Hamiltonian formulation derives the equations of motion for a mechanical system using only the symplectic form and a Hamiltonian function, which maps the state of the system to its total (kinetic plus potential) energy [19]. This formulation has been used by several authors to learn unknown dynamics: the Hamiltonian structure (canonical symplectic form) is used as a physics prior and the unknown dynamics are uncovered by learning the Hamiltonian [5,10,20,21,22]. Consider a system with configuration space ${\mathbb{R}}^n$ and a choice of n generalized coordinates that represent configuration. Let $\mathbf{z}\left(t\right)\in {\mathbb{R}}^{2n}$ represent the vector of n generalized coordinates and their n conjugate momenta at time t. Define the Hamiltonian function $\mathcal{H}:{\mathbb{R}}^{2n}\mapsto \mathbb{R}$ such that $\mathcal{H}\left(\mathbf{z}\right)$ is the sum of the kinetic plus potential energy. Then, the equations of motion [19,23] derive as

$$\frac{d\mathbf{z}}{dt}={\mathsf{\Lambda }}_{\mathrm{can}}{\nabla }_{\mathbf{z}}\mathcal{H}\left(\mathbf{z}\right),{\mathsf{\Lambda }}_{\mathrm{can}}=\left(\begin{array}{cc}{\mathbf{0}}_n& {\mathbf{I}}_n\\ -{\mathbf{I}}_n& {\mathbf{0}}_n\end{array}\right)$$

(2)

where ${\mathbf{0}}_n\in {\mathbb{R}}^{n\times n}$ is the matrix of all zeros and ${\mathsf{\Lambda }}_{\mathrm{can}}$ is the matrix representation of the canonical symplectic form.

The Hamiltonian equations of motion for a freely rotating 3D rigid body evolve on the six-dimensional space $T^{*}\mathbf{SO}\left(3\right)$, the co-tangent bundle of $\mathbf{SO}\left(3\right)$. However, because of rotational symmetry in the dynamics, i.e., the invariance of the dynamics of a freely rotating rigid body to the choice of inertial frame, the Hamiltonian formulation of the dynamics can be reduced using the Lie–Poisson Reduction Theorem [24] to the space ${\mathbb{R}}^3\sim {\mathfrak{so}}^{*}\left(3\right)$, the Lie co-algebra of $\mathbf{SO}\left(3\right)$. These reduced Hamiltonian dynamics are equivalent to (A2), where the body angular momentum is $\mathsf{\Pi }\left(t\right)=\mathbf{J}\mathsf{\Omega }\left(t\right)\in {\mathfrak{so}}^{*}\left(3\right)$ for $t\ge 0$. The invariance can be seen by observing that the rotation matrix $\mathbf{R}\left(t\right)$, which describes the orientation of the body at time t, does not appear in (A2). $\mathbf{R}\left(t\right)$ is calculated from the solution of (A2) using (A3).

The reduced Hamiltonian $h:{\mathfrak{so}}^{*}\left(3\right)\mapsto \mathbb{R}$ for the freely rotating 3D rigid body (kinetic energy) is

$$h\left(\mathsf{\Pi }\right)=\frac{1}{2}\mathsf{\Pi }·{\mathbf{J}}^{-1}\mathsf{\Pi }.$$

(3)

The reduced Hamiltonian formulation [24] is

$$\frac{d\mathsf{\Pi }}{dt}={\mathsf{\Lambda }}_{{\mathfrak{so}}^{*}\left(3\right)}\left(\mathsf{\Pi }\right){\nabla }_{\mathsf{\Pi }}h\left(\mathsf{\Pi }\right),{\mathsf{\Lambda }}_{{\mathfrak{so}}^{*}\left(3\right)}\left(\mathsf{\Pi }\right)={\mathsf{\Pi }}_{\times },$$

(4)

which can be seen to be equivalent to (A2). Equation (4), called the Lie–Poisson equation, generalizes the canonical Hamiltonian formulation. The generalization allows for different symplectic forms, i.e., ${\mathsf{\Lambda }}_{{\mathfrak{so}}^{*}\left(3\right)}$ instead of ${\mathsf{\Lambda }}_{\mathrm{can}}$ in this case, each of which is only related to the latent space and symmetry. Our physics prior is the generalized symplectic form and learning the unknown dynamics means learning the reduced Hamiltonian. This is a generalization of the existing literature, where dynamics of canonical Hamiltonian systems are learned with the canonical symplectic form as the physics prior [5,10,11,12]. Using the generalized Hamiltonian formulation allows extension of the approach to a much larger class of systems than those described by Hamilton’s canonical equations, including rotating and translating 3D rigid bodies, rigid bodies in a gravitational field, multi-body systems, and more.

4. Materials and Methods

In this section, we outline our approach for learning and predicting rigid-body dynamics from image sequences. The multi-stage prediction pipeline maps individual images to an $\mathbf{SO}\left(3\right)$ latent space where angular velocities are computed from latent pairs. Future latent states are computed using the generalized Hamiltonian equations of motion (4) and a learned representation of the reduced Hamiltonian (3). Finally, the predicted latent representations are mapped to images giving a predicted image sequence.

4.1. Notation

N denotes the number of image sequences in the dataset, and $T+1$ is the length of each image sequence. Image sequences are written ${\mathbf{x}}_k=\{x_0^k,\cdots ,x_T^k\}$, sequences of latent rotation matrices are written ${\mathbf{R}}_k=\{{\mathit{R}}_0^k,\cdots ,{\mathit{R}}_T^k\}$ with ${\mathit{R}}_i^k\in \mathbf{SO}\left(3\right)$, and quaternion latent sequences are written ${\mathbf{q}}_k=\{q_0^k,\cdots ,q_T^k\}$ with $q_i^k\in \mathcal{S}\left(3\right)$. Each element $y_i^k$ represents the quantity y at time step $t=i$ for sequence k from the dataset, where $k\in \{1,\cdots ,N\}$. Quantities generated with the learned dynamics are denoted with a hat (e.g., $\widehat{q}$).

4.2. Embedding Images to an SO(3) Latent Space

In the first stage of our prediction pipeline, we embed image observations of a freely rotating rigid body to a low-dimensional latent representation to facilitate computation of the dynamics. The latent representation is constrained to have the same $\mathbf{SO}\left(3\right)$ structure as the configuration space of the rigid body, making learned representations interpretable and compatible with the equations of motion. Our embedding network $\mathsf{\Phi }$ is given by the composition of functions $\mathsf{\Phi }:=f\circ \pi \circ E_{\varphi }:\mathcal{I}\mapsto \mathbf{SO}\left(3\right)$. The convolutional encoding neural network $E_{\varphi }:\mathcal{I}\mapsto {\mathbb{R}}^6$ parameterized by $\varphi$ maps image observations from image space $\mathcal{I}$ to a vector $\mathbf{z}\in {\mathbb{R}}^6$. The projection operator $\pi :{\mathbb{R}}^6\mapsto {\mathcal{S}}^2\times {\mathcal{S}}^2$ decomposes the vector z into the vectors $\mathbf{u},\mathbf{v}\in {\mathbb{R}}^3$ and normalizes them, i.e., $\pi \left(\mathbf{z}\right)=(\mathbf{u}/\parallel \mathbf{u}\parallel ,\mathbf{v}/\parallel \mathbf{v}\parallel )$. Finally, the function $f:{\mathcal{S}}^2\times {\mathcal{S}}^2\mapsto \mathbf{SO}\left(3\right)$ maps the normalized vectors $\mathbf{u}$ and $\mathbf{v}$ to the configuration space using the surjective and differentiable ${\mathcal{S}}^2\times {\mathcal{S}}^2$ parameterization of $\mathbf{SO}\left(3\right)$ (see Section 3.1).

4.3. Computing Dynamics in the Latent Space

In the second stage of our prediction pipeline, we compute the dynamics of the freely rotating rigid body using a Hamiltonian with a learned moment-of-inertia tensor, ${\mathbf{J}}_{\phi }$. The moment-of-inertia tensor, ${\mathbf{J}}_{\phi }$, is parameterized by the vectors ${\mathit{\phi }}_1,{\mathit{\phi }}_2\in {\mathbb{R}}^3$, representing the diagonal and off-diagonal components of the matrix, and computed using the Cholesky decomposition [25].

To compute the dynamics, we first construct an initial condition $({\mathit{R}}_0^k,{\mathsf{\Pi }}_0^k)\in T^{*}\mathbf{SO}\left(3\right)$. Given the sequential pair $({\mathit{R}}_0^k,{\mathit{R}}_1^k)=(\mathsf{\Phi }\left(x_0^k\right),\mathsf{\Phi }\left(x_1^k\right))$, we perform this in two steps. First, we compute the angular velocity ${\mathsf{\Omega }}_0^k$ by Algorithm 1. Then, we compute the angular momentum by the matrix product of the learned moment-of-inertia and angular velocity, i.e., ${\mathsf{\Pi }}_0^k={\mathbf{J}}_{\phi }{\mathsf{\Omega }}_0^k$. With the initial condition $({\mathit{R}}_0^k,{\mathsf{\Pi }}_0^k)$, subsequent angular momenta ${\left\{{\widehat{\mathsf{\Pi }}}_i^k\right\}}_{i=1}^T$ are computed using the Lie–Poisson Equation (4) and the reduced Hamiltonian formed using the learned momentum-of-inertia ${\mathbf{J}}_{\phi }$. We integrate the Lie–Poisson equations forward in time using a Runge–Kutta fourth-order (RK45) numerical solver [26].

Subsequent rotations ${\left\{{\widehat{\mathit{R}}}_i^k\right\}}_{i=1}^T$ are computed in two steps. First, we compute the sequence of quaternions ${\left\{{\widehat{q}}_i^k\right\}}_{i=1}^T$ by Equation (1), using the quaternion representation $q_0^k$ of the initial rotation ${\mathit{R}}_0^k$ and the initial angular velocity ${\mathsf{\Omega }}_0^k$. We integrate Equation (1) forward in time using an RK45 solver with a normalization step [18] that ensures elements of the resulting sequence are valid quaternions. Then, we transform the sequence of quaternions ${\left\{{\widehat{q}}_i^k\right\}}_{i=1}^T$ to a sequence of rotations rotations ${\left\{{\widehat{\mathit{R}}}_i^k\right\}}_{i=1}^T$ using a modified Shepperd’s algorithm [27].

Algorithm 1: An algorithm to calculate the body angular velocity given two sequential orientation matrices and the time step in between them.

4.4. Decoding SO(3) Latent States to Images

In the final stage of our prediction pipeline, we decode the sequence of $\mathbf{SO}\left(3\right)$ latent states produced by the dynamics pipeline to a sequence of images (see Figure 2d). Our decoding network $\mathsf{\Psi }$ is given as the composition of functions $\mathsf{\Psi }: =D_{\psi }\circ {\pi }^{-1}\circ f^{-1}:\mathbf{SO}\left(3\right)\mapsto \mathcal{I}$, where the convolutional decoding network $D_{\psi }:{\mathbb{R}}^6\mapsto \mathcal{I}$ parameterized by $\psi$ maps a vector $\mathbf{z}=(\mathbf{u},\mathbf{v})$, where $(\mathbf{u},\mathbf{v})\in {\mathcal{S}}^2\times {\mathcal{S}}^2$, to the image space $\mathcal{I}$.

4.5. Training Methodology

In this section, we describe the loss functions used to optimize our model: the auto-encoding reconstruction loss (${\mathcal{L}}_{\mathrm{ae}}$), the dynamics-based reconstruction loss (${\mathcal{L}}_{\mathrm{dyn}}$), the latent orientation loss (${\mathcal{L}}_{\mathrm{latent},\mathrm{R}}$), and latent momentum loss (${\mathcal{L}}_{\mathrm{latent},\mathsf{\Pi }}$). ${\mathcal{L}}_{\mathrm{ae}}$ ensures the embedding to $\mathbf{SO}\left(3\right)$ is sufficiently expressive to represent the entire image dataset, and ${\mathcal{L}}_{\mathrm{dyn}}$ ensures correspondence between the input image sequences and the images sequences produced by the learned dynamics. The latent loss functions, ${\mathcal{L}}_{\mathrm{latent},\mathrm{R}}$ and ${\mathcal{L}}_{\mathrm{latent},\mathsf{\Pi }}$, ensure consistency between the latent states produced by the encoding pipeline and those produced by the dynamics pipeline.

For notational convenience, we denote the encoding pipeline $\mathcal{E}:\mathcal{I}\mapsto {\mathcal{S}}^3$ and the decoding pipeline $\mathcal{D}:{\mathcal{S}}^3\mapsto \mathcal{I}$. Quantities computed in the encoding pipeline use subscript ae (e.g., ${{\mathit{R}}_{\mathrm{ae}}^k}_i$), while those computed in the dynamics pipeline use subscript dyn (e.g., ${{\mathit{R}}_{\mathrm{dyn}}^k}_i$).

4.5.1. Reconstruction Losses

The auto-encoding reconstruction loss is the mean squared error (MSE) between the ground-truth image sequence and the reconstructed image sequence without dynamics:

$${\mathcal{L}}_{\mathrm{ae}}=\frac{1}{NT}\sum _{k=1}^N\sum _{i=0}^{T-1}{∥x_i^k-(\mathcal{D}\circ \mathcal{E})\left(x_i^k\right)∥}_2^2.$$

The dynamics-based reconstruction loss is the MSE between the ground-truth image sequence and the image sequence produced by the dynamics pipeline:

$${\mathcal{L}}_{\mathrm{dyn}}=\frac{1}{NT}\sum _{k=1}^N\sum _{i=1}^T{∥x_i^k-\mathcal{D}\left({q_{\mathrm{dyn}}^k}_i\right)∥}_2^2.$$

4.5.2. Latent Losses

We define ${{\mathcal{L}}_{\mathrm{latent}}}_{\mathit{R}}$ as the $\mathbf{SO}\left(3\right)$ distance [19] between the $3\times 3$ identity matrix and right-difference of orientations produced in the encoding pipeline and the orientations produced in the dynamics pipeline:

$${{\mathcal{L}}_{\mathrm{latent}}}_{\mathit{R}}=\frac{1}{NT}\sum _{k=1}^N\sum _{i=1}^T{∥{\mathbf{I}}_3-{\left({{\mathit{R}}_{\mathrm{ae}}^k}_i\right)}^T{{\mathit{R}}_{\mathrm{dyn}}^k}_i∥}_F^2.$$

We define ${{\mathcal{L}}_{\mathrm{latent}}}_{\mathsf{\Pi }}$ as the MSE between the angular momenta estimated in the encoding pipeline and the angular momenta computed in the dynamics pipeline (see Figure 2):

$${{\mathcal{L}}_{\mathrm{latent}}}_{\mathsf{\Pi }}=\frac{1}{NT}\sum _{k=1}^N\sum _{i=1}^T{∥{{\mathsf{\Pi }}_{\mathrm{ae}}^k}_i-{{\mathsf{\Pi }}_{\mathrm{dyn}}^k}_i∥}_2^2.$$

The hyperparameters we use to train our model are given in

Table A2. Hyperparameters used to train model for the non-uniform mass density prism experiment. Only values differing from default values are given in the table.

Experiment Hyperparameters
Parameter Name	Value
Random seed	0
Test dataset split	0.2
Validation dataset split	0.1
Number of epochs	1000
Batch size	256
Autoencoder learning rate	$1\times {10}^{-3}$
Dynamics learning rate	$1\times {10}^{-3}$
Sequence length	10
Time step	$1\times {10}^{-3}$

in Appendix A.3. We train our model for 500 epochs, on a single NVIDIA A100 SXM4 GPU. Our training time is approximately 12 h, and our inference time is approximately 300 milliseconds.

4.6. 3D Rotating Rigid-Body Datasets

To evaluate our model, we introduce six synthetic datasets of freely rotating objects. Previous efforts in learning dynamics from images [4,5,6,10] consider only 2D planar systems (e.g., the simple pendulum, Acrobot, and cart-pole); existing datasets of freely rotating rigid bodies in 3D such as SPEED [28,29], SPEED+ [30], and URSO [31], contain random image-pose pairs rather than sequential pairs needed for video prediction and dynamics extraction. Our datasets showcase the rich dynamical behaviors of 3D rotational dynamics through images, and can be used for 3D dynamics learning tasks. Specifically, we introduce the following five datasets (see Table A1 in Appendix A.2 for moment-of-inertia matrices):

• Uniform mass density cube: a multi-colored cube of uniform mass density;
• Uniform mass density prism: a multi-colored rectangular prism with uniform mass density;
• Non-uniform mass density cube: a multi-colored cube with non-uniform mass density;
• Non-uniform mass density prism: a multi-colored prism with non-uniform mass density;
• Uniform density synthetic-satellites: renderings of CALIPSO and CloudSat satellites with uniform mass density.

For each dataset, $N=1000$ trajectories are created. Each trajectory consists of an initial condition ${\mathbf{x}}_0=({\mathbf{R}}_0,{\mathsf{\Pi }}_0)$ that is integrated forward in time using a Python-based Runge–Kutta solver for $T=100$ time steps with spacing $\Delta t={10}^{-3}$. Initial conditions are chosen such that $({\mathbf{R}}_0,{\mathsf{\Pi }}_0)\sim \mathrm{Uniform}\left(\mathbf{SO}\left(3\right)\times S^2\right)$ with ${\mathsf{\Pi }}_0$ scaled to have ${\parallel {\mathsf{\Pi }}_0\parallel }_2=50$. The orientations $\widehat{\mathbf{q}}$ from the integrated trajectories are passed to Blender [32] to render images of $28\times 28$ pixels (as shown in Figure 1).

The synthetic image datasets are generated using Blender [32] with ideal and fixed lighting conditions. Models trained on this dataset may exhibit sensitivity to variations in lighting conditions, and may not generalize to real data.

5. Results

Figure 3 and Figure 4 show the model’s performance on the datasets for both short- and long-term predictions. Figure 3 results show that the model is capable of predicting into the future at least five fold longer than the length of the time horizon used at training time. Figure 4 results show that the model is capable of predicting the future with images of more complex geometries and surface properties, i.e., those of the CALIPSO and CloudSat satellites, at least ten fold longer than the length of the time horizon used at training time. The model’s performance on the datasets is indicative of its capabilities to predict dynamics and map them to image space.

The uniform mass density cube and prism datasets are used to demonstrate baseline capabilities of our approach for image prediction. The non-uniform mass density datasets validate the model’s capability to predict a mass distribution that may not be visible from the exterior, e.g., for an asteroid or space debris or as part of failure diagnostics on a satellite where there may be broken or shifted internal components. The satellite datasets are used to validate the model’s capability to handle bodies with less regular and more realistic external geometries.

We compare the performance of our model to three baseline models: (1) the Long Short Term Memory (LSTM) network, (2) the Neural ODE [33] network, and (3) the Hamiltonian Generative Network (HGN) [5]. Recurrent neural networks like the LSTM-baseline provide a discrete dynamics model. Neural ODE can be combined with a multi-layer perceptron to predict continuous dynamics. HGN is a generative model with a Hamiltonian inductive bias. Architecture and training details for each baseline are given in Appendix A.4. The prediction performances of our model and the baselines are shown in

Table 1. Average pixel mean square error over a 30-step prediction on the train and test data on six datasets. All values are multiplied by 1 $\times {10}^3$. We evaluate our model and compare to three baseline models: (1) recurrent model (LSTM [34]), (2) Neural ODE ([33]), and (3) HGN ([5]).

Dataset	Ours		LSTM-Baseline		Neural ODE-Baseline		HGN
	TRAIN	TEST	TRAIN	TEST	TRAIN	TEST	TRAIN	TEST
Uniform Prism	2.66 ± 0.10	2.71 ± 0.08	3.46 ± 0.59	3.47 ± 0.61	3.96 ± 0.68	4.00 ± 0.68	4.18 ± 0.0	7.80 ± 0.30
Uniform Cube	3.54 ± 0.17	3.97 ± 0.16	21.55 ± 1.98	21.64 ± 2.12	9.48 ± 1.19	9.43 ± 1.20	17.43 ± 0.00	18.69 ± 0.12
Non-uniform Prism	4.27 ± 0.18	6.61 ± 0.88	4.50 ± 1.31	4.52 ± 1.34	4.67 ± 0.58	4.75 ± 0.59	6.16 ± 0.08	8.33 ± 0.26
Non-uniform Cube	6.24 ± 0.29	4.85 ± 0.35	7.47 ± 0.51	7.51 ± 0.50	7.89 ± 1.50	7.94 ± 1.59	14.11 ± 0.13	18.14 ± 0.36
CALIPSO	0.79 ± 0.53	0.87 ± 0.50	0.62 ± 0.21	0.65 ± 0.22	0.69 ± 0.26	0.71 ± 0.27	1.18 ± 0.02	1.34 ± 0.05
CloudSat	0.64 ± 0.45	0.65 ± 0.29	0.89 ± 0.36	0.93 ± 0.43	0.65 ± 0.22	0.66 ± 0.25	1.48 ± 0.04	1.66 ± 0.11
Number of Parameters	6		52,400		11,400		-

. Our model has the lowest MSE on the majority of our datasets with good prediction performance on all of our datasets. Our model outperforms the state-of-the-art HGN model, reducing the expected MSE by nearly half on all datasets. Overall, our model outperforms the baseline models on the majority of the datasets with a more interpretable latent space, continuous dynamics, and fewer model parameters.

In Appendix A.5, we present the results of ablations studies and provide discussion. We find that the latent losses improve performance. However, the model may be over constrained with both the dynamics-based and auto-encoding based reconstruction losses.

6. Summary and Conclusions

6.1. Summary

In this work, we have presented the first physics-informed deep learning framework for predicting image sequences of 3D rotating rigid-bodies by embedding the images as measurements in the configuration space $\mathbf{SO}\left(3\right)$ and propagating the Hamiltonian dynamics forward in time. We have evaluated our approach on new datasets of free-rotating 3D bodies with different inertial properties, and have demonstrated the ability to perform long-term image predictions. We outperform the LSTM, Neural ODE and Hamiltonian Generative Network (HGN) baselines on our datasets, producing better qualitative predictions and reducing the error observed for the state-of-the-art HGN by a factor of 2.

6.2. Conclusions

By enforcing the representation of the latent space to be $\mathbf{SO}\left(3\right)$, this work provides the advantage of interpretability over black-box physics-informed approaches. The extra interpretability of our approach is a step towards placing additional trust into sophisticated deep learning models. This work provides a natural path to investigating how to incorporate and evaluate the effect of classical model-based control directly to trajectories in the latent space.

7. Future Work

Although our approach so far has been limited to embedding RGB-images of rotating rigid-bodies with configuration spaces in $\mathbf{SO}\left(3\right)$, there are natural extensions to a wider variety of problems. For instance, this framework can be extended to embed different high-dimensional sensor measurements, such as point clouds, by modifying the feature extraction layers of the autoencoder. The latent space can be chosen to reflect generic rigid bodies in $\mathbf{SE}\left(3\right)$ or systems in more complicated spaces, such as the n-jointed robotic arm on a restricted subspace of ${\mathsf{\Pi }}_{i=1}^n\left(\mathbf{SO}\left(3\right)\right)$. Another possible extension includes multibody systems, i.e., systems with rigid and flexible body dynamics, which would have applications to systems such as spacecraft with flexible solar panels and aircraft with flexible wings.