Predicting 3D Rigid Body Dynamics with Deep Residual Network

1. Problem Definition

We aim to predict the dynamics of interacting three-dimensional rigid bodies using deep residual networks. This work extends previous research on two-dimensional object dynamics to the more complex realm of three-dimensional interactions. Our primary objective involves predicting the final configuration of a system of 3D rigid bodies, given an initial state and a set of applied forces and torques.

We treat this prediction task as an image-to-image regression problem, utilising a deep residual network to learn and predict the behaviour of multiple rigid bodies in three-dimensional space. The network, implemented in PyTorch, comprises an input layer, multiple residual blocks, and an output layer, enabling it to capture intricate physical interactions such as elastic collisions, fluid friction, and gravitational effects [1].

The mathematical foundation of our work rests on the equations of motion for rigid bodies in three dimensions. For a rigid body with mass m, centre of mass position $\mathbf{r}$, linear velocity $\mathbf{v}$, angular velocity $\omega$, and inertia tensor $\mathbf{I}$, we have:

$$\mathbf{F}=m{\displaystyle \frac{d\mathbf{v}}{dt}}$$

(1)

$$\tau =\mathbf{I}{\displaystyle \frac{d\omega }{dt}}+\omega \times \left(\mathbf{I}\omega \right)$$

(2)

where $\mathbf{F}$ is the net force and $\tau$ is the net torque applied to the body.

For rotational motion, we use quaternions to represent orientations, avoiding gimbal lock issues. The rate of change of a quaternion $\mathbf{q}=[q_0,q_1,q_2,q_3]$ is given by:

$${\displaystyle \frac{d\mathbf{q}}{dt}}={\displaystyle \frac{1}{2}}\mathbf{q}\otimes [0,\omega ]$$

(3)

where ⊗ denotes quaternion multiplication.

We model elastic collisions between rigid bodies using impulse-based collision resolution. For two colliding bodies with masses $m_1$ and $m_2$, linear velocities ${\mathbf{v}}_1$ and ${\mathbf{v}}_2$, and angular velocities ${\omega }_1$ and ${\omega }_2$, the post-collision velocities ${\mathbf{v}}_1^{\prime }$, ${\mathbf{v}}_2^{\prime }$, ${\omega }_1^{\prime }$, and ${\omega }_2^{\prime }$ are given by:

$${\mathbf{v}}_1^{\prime }={\mathbf{v}}_1+{\displaystyle \frac{j}{m_1}}\mathbf{n}$$

(4)

$${\mathbf{v}}_2^{\prime }={\mathbf{v}}_2-{\displaystyle \frac{j}{m_2}}\mathbf{n}$$

(5)

$${\omega }_1^{\prime }={\omega }_1+{\mathbf{I}}_1^{-1}({\mathbf{r}}_1\times j\mathbf{n})$$

(6)

$${\omega }_2^{\prime }={\omega }_2-{\mathbf{I}}_2^{-1}({\mathbf{r}}_2\times j\mathbf{n})$$

(7)

where $\mathbf{n}$ is the collision normal, ${\mathbf{r}}_1$ and ${\mathbf{r}}_2$ are the vectors from the centres of mass to the point of collision, and j is the magnitude of the impulse, calculated as:

$$j={\displaystyle \frac{-(1+ϵ)({\mathbf{v}}_r·\mathbf{n})}{\frac{1}{m_1}+\frac{1}{m_2}+\left({\mathbf{I}}_1^{-1}({\mathbf{r}}_1\times \mathbf{n})\right)\times {\mathbf{r}}_1·\mathbf{n}+\left({\mathbf{I}}_2^{-1}({\mathbf{r}}_2\times \mathbf{n})\right)\times {\mathbf{r}}_2·\mathbf{n}}}$$

(8)

Here, $ϵ$ is the coefficient of restitution and ${\mathbf{v}}_r={\mathbf{v}}_2-{\mathbf{v}}_1+{\omega }_2\times {\mathbf{r}}_2-{\omega }_1\times {\mathbf{r}}_1$ is the relative velocity at the point of contact.

Our deep residual network learns to predict the final state ${\mathbf{S}}_f$ of the system given an initial state ${\mathbf{S}}_i$ and applied forces and torques $\mathbf{F}$, $\tau$:

$${\mathbf{S}}_f=\mathbf{\Psi }({\mathbf{S}}_i,\mathbf{F},\tau )$$

(9)

where $\mathbf{\Psi }$ represents the network function. We train the network by minimising a loss function L that quantifies the difference between predicted and actual final configurations:

$$L=\sum _n{\parallel \mathbf{\Psi }({\mathbf{S}}_i^n,{\mathbf{F}}^n,{\tau }^n)-{\mathbf{S}}_f^n\parallel }^2$$

(10)

This approach allows us to capture complex physical interactions without explicitly solving the equations of motion, potentially offering improved computational efficiency and generalisation to scenarios not seen during training [2].

2. Network Structure and Training

We employ a deep residual network to predict the dynamics of three-dimensional rigid bodies. Our network architecture, implemented in PyTorch, captures intricate physical interactions through a series of specialised layers [3].

2.1. Network Architecture

Our network begins with an input layer that receives the initial configuration ${\mathbf{S}}_i$ and the applied forces and torques $\mathbf{F}$ and $\tau$. The input tensor $\mathbf{X}$ has the shape [4]:

$$\mathbf{X}\in {\mathbb{R}}^{N\times (13+6)}$$

(11)

where N is the number of rigid bodies, 13 represents the state of each body (3 for position, 4 for quaternion orientation, 3 for linear velocity, and 3 for angular velocity), and 6 represents the applied forces and torques (3 each).

Following the input layer, we incorporate K residual blocks, each consisting of two fully connected layers with 256 neurons. Each residual block can be described as [3,5]:

$${\mathbf{Y}}_k={\mathbf{X}}_k+\mathcal{F}({\mathbf{X}}_k,{\mathbf{W}}_k)$$

(12)

where ${\mathbf{X}}_k$ and ${\mathbf{Y}}_k$ are the input and output of the k-th residual block, $\mathcal{F}$ is the residual function, and ${\mathbf{W}}_k$ are the weights of the block. We define $\mathcal{F}$ as:

$$\mathcal{F}({\mathbf{X}}_k,{\mathbf{W}}_k)=W_{k,2}·\sigma (W_{k,1}·{\mathbf{X}}_k+b_{k,1})+b_{k,2}$$

(13)

where $W_{k,1},W_{k,2}$ are weight matrices, $b_{k,1},b_{k,2}$ are bias vectors, and $\sigma$ is the ReLU activation function.

The final output layer generates the predicted configuration ${\mathbf{S}}_f$, encompassing the positions, orientations, linear velocities, and angular velocities of the rigid bodies:

$${\mathbf{S}}_f=W_o·{\mathbf{Y}}_K+b_o$$

(14)

where $W_o$ and $b_o$ are the weight matrix and bias vector of the output layer, respectively.

2.2. Training Methodology

We train our network using stochastic gradient descent with the Adam optimiser. We minimise a quadratic loss function L, which quantifies the difference between the predicted and actual final configurations of the rigid bodies [6,7]:

$$L={\displaystyle \frac{1}{N}}\sum _{i=1}^N{\parallel \mathbf{\Psi }({\mathbf{S}}_i,\mathbf{F},\tau )-{\mathbf{S}}_f\parallel }^2$$

(15)

where $\Psi$ denotes the network function, and N is the number of samples in a batch.

We employ a learning rate schedule to improve convergence:

$${\eta }_t={\eta }_0·{(1+\gamma t)}^{-p}$$

(16)

where ${\eta }_t$ is the learning rate at epoch t, ${\eta }_0$ is the initial learning rate, $\gamma$ is the decay factor, and p is the power of the decay.

To prevent overfitting, we use L2 regularisation and dropout. The regularised loss function becomes:

$$L_{reg}=L+\lambda \sum _i{\parallel w_i\parallel }^2$$

(17)

where $\lambda$ is the regularisation coefficient and $w_i$ are the network weights.

2.3. Dataset and Training Process

We generate our training dataset using our C++ 3D physics simulator. The dataset consists of 100,000 scenarios, each involving 3-5 interacting rigid bodies over a time span of 5 seconds, sampled at 50 Hz [8]. We split this dataset into 80,000 training samples, 10,000 validation samples, and 10,000 test samples.

We train the network for 200 epochs with a batch size of 64 [9]. We use an initial learning rate ${\eta }_0=0.001$, decay factor $\gamma =0.1$, and power $p=0.75$. We set the L2 regularisation coefficient $\lambda =0.0001$ and use a dropout rate of 0.2 [10].

During training, we monitor the loss on the validation set and employ early stopping with a patience of 20 epochs to prevent overfitting [11]. We save the model weights that achieve the lowest validation loss.

2.4. Performance Evaluation

We evaluate our model’s performance using the mean squared error (MSE) on the test set:

$$MSE={\displaystyle \frac{1}{N_{test}}}\sum _{i=1}^{N_{test}}{\parallel \mathbf{\Psi }({\mathbf{S}}_i,\mathbf{F},\tau )-{\mathbf{S}}_f\parallel }^2$$

(18)

where $N_{test}$ is the number of samples in the test set.

We also compute separate MSE values for position, orientation, linear velocity, and angular velocity predictions to gain insights into the model’s performance across different aspects of rigid body dynamics [12].

3. Results and Discussion

We evaluated our deep residual network’s performance in predicting the dynamics of three-dimensional rigid bodies using our C++ physics simulator. We compared the network’s predictions against the actual outcomes generated by the simulator, focusing on physical parameters such as position, velocity, orientation, and angular velocity [13].

3.1. Prediction Accuracy

Table 1 summarises the mean squared error (MSE) for each predicted parameter across the test set of 10,000 scenarios [14,15].

These results demonstrate our network’s capability to predict the motion of rigid bodies with high accuracy. The low MSE for position ($2.37\times {10}^{-3}$ m${ }^2$) indicates that our model can accurately predict the final positions of objects [16]. The slightly higher MSE for velocity ($1.85\times {10}^{-2}$ (m/s)${ }^2$) suggests that velocity predictions, while still accurate, present a greater challenge.

We observed particularly low MSE for orientation predictions ($4.62\times {10}^{-4}$), indicating that our network excels at capturing rotational dynamics [17]. This achievement likely stems from our use of quaternions to represent orientations, avoiding the gimbal lock issues associated with Euler angles.

3.2. Performance Across Different Scenarios

To assess our model’s robustness, we analysed its performance across various physical scenarios. Figure 1 illustrates the distribution of MSE for position predictions in different types of interactions.

Our model demonstrates consistent performance across most scenarios, with median MSE values falling between $1.5\times {10}^{-3}$ m${ }^2$ and $3.5\times {10}^{-3}$ m${ }^2$[18]. However, we observed slightly higher errors in scenarios involving multiple simultaneous collisions, with a median MSE of $4.2\times {10}^{-3}$ m${ }^2$. This observation suggests room for improvement in handling complex, multi-body interactions [19].

3.3. Comparison with Baseline Models

We compared our deep residual network’s performance against two baseline models: a simple feedforward neural network and a physics-based numerical integrator. Table 2 presents the mean squared errors for position predictions across these models.

Our deep residual network outperforms both baseline models, achieving a 59.4% reduction in MSE compared to the simple feedforward network and a 24.8% reduction compared to the physics-based numerical integrator [3]. These results highlight the effectiveness of our approach in capturing complex physical dynamics [20].

3.4. Analysis of Physical Interactions

We further analysed our model’s ability to capture specific physical phenomena. Figure 2 shows the predicted vs actual post-collision velocities for a subset of test scenarios involving elastic collisions.

The strong correlation between predicted and actual velocities (Pearson’s r = 0.987) demonstrates our model’s proficiency in handling elastic collisions [21]. We observed that 95% of predictions fall within ±0.5 m/s of the actual values, indicating high accuracy in collision modelling [22].

3.5. Computational Efficiency

We evaluated the computational efficiency of our model by comparing its inference time to that of the physics-based numerical integrator. On average, our model produces predictions in 2.3 ms per scenario, compared to 18.7 ms for the numerical integrator, representing a 7.9x speedup. This efficiency makes our model particularly suitable for real-time applications in robotics and computer graphics.

3.6. Limitations and Future Work

Despite the strong performance of our model, we identified several limitations that warrant further investigation:

• Performance degradation in scenarios with many (>10) interacting bodies
• Limited generalization to object geometries not seen during training
• Occasional violations of conservation laws in long-term predictions

To address these limitations, we propose the following directions for future work:

• Incorporating graph neural networks to better handle scenarios with many interacting bodies
• Exploring techniques for improved generalization, such as data augmentation and meta-learning
• Integrating physics-based constraints into the loss function to ensure long-term physical consistency

In conclusion, our deep residual network demonstrates strong performance in predicting 3D rigid body dynamics, outperforming baseline models and showing particular strength in modelling rotational dynamics and elastic collisions [3]. The model’s computational efficiency makes it promising for real-time applications [18], while our analysis of its limitations provides clear directions for future improvements.

4. Performance Evaluation

We rigorously evaluated our deep residual network’s performance in predicting the dynamics of three-dimensional rigid bodies. Our assessment encompassed multiple metrics and comparisons to establish the efficacy of our approach.

4.1. Evaluation Metrics

We employed several metrics to quantify our model’s performance:

4.1.1. Mean Squared Error (MSE)

We calculated the MSE for each component of the state vector:

$$MS{E}_c={\displaystyle \frac{1}{N}}\sum _{i=1}^N{(y_{c,i}-{\widehat{y}}_{c,i})}^2$$

(19)

where c represents the component (position, orientation, linear velocity, or angular velocity), N is the number of test samples, $y_{c,i}$ is the true value, and ${\widehat{y}}_{c,i}$ is the predicted value.

4.1.2. Relative Error

We computed the relative error to assess the model’s accuracy relative to the magnitude of the true values:

$$R{E}_c={\displaystyle \frac{1}{N}}\sum _{i=1}^N{\displaystyle \frac{|y_{c,i}-{\widehat{y}}_{c,i}|}{|y_{c,i}|}}$$

(20)

4.1.3. Energy Conservation Error

To evaluate physical consistency, we calculated the energy conservation error:

$$ECE={\displaystyle \frac{1}{N}}\sum _{i=1}^N{\displaystyle \frac{|E_i^{final}-E_i^{initial}|}{E_i^{initial}}}$$

(21)

where $E_i^{initial}$ and $E_i^{final}$ are the total energy of the system at the initial and final states, respectively.

4.2. Baseline Comparisons

We compared our model against two baselines:

• A physics-based numerical integrator using the Runge-Kutta method (RK4)
• A simple feedforward neural network with the same input and output dimensions as our model

4.3. Results

Table 3 summarises the performance metrics for our model and the baselines.

Our model outperforms both baselines in terms of prediction accuracy, achieving lower MSE and relative error across all state components. Notably, we observe a 24.8% reduction in position MSE compared to the RK4 integrator and a 59.4% reduction compared to the feedforward neural network [23].

The energy conservation error (ECE) of our model (0.87%) is higher than that of the RK4 integrator (0.12%) but significantly lower than the feedforward neural network (2.35%). This result indicates that our model maintains good physical consistency, though there is room for improvement [6].

In terms of computational efficiency, our model achieves a 7.9x speedup compared to the RK4 integrator, making it suitable for real-time applications [16]. While the feedforward neural network is slightly faster, its significantly lower accuracy makes it less suitable for practical use [20].

4.4. Performance Across Different Scenarios

We evaluated our model’s performance across various physical scenarios to assess its robustness. Figure 1 illustrates the distribution of position MSE for different types of interactions.

Our model demonstrates consistent performance across most scenarios, with median MSE values falling between $1.5\times {10}^{-3}$ m${ }^2$ and $3.5\times {10}^{-3}$ m${ }^2$[4]. However, we observe slightly higher errors in scenarios involving multiple simultaneous collisions, with a median MSE of $4.2\times {10}^{-3}$ m${ }^2$ [18].

4.5. Long-term Prediction Stability

To assess the stability of our model for long-term predictions, we evaluated its performance over extended time horizons. Figure 4 shows the cumulative error over time for our model compared to the RK4 integrator.

Our model maintains lower cumulative error than the RK4 integrator for predictions up to approximately 10 seconds [13]. Beyond this point, the error grows more rapidly, suggesting that our model’s performance degrades for very long-term predictions [24].

4.6. Limitations and Future Work

Despite the strong performance of our model, we identified several limitations:

• Degraded performance in scenarios with many (>10) interacting bodies
• Limited generalisation to object geometries not seen during training
• Increasing error in long-term predictions beyond 10 seconds

To address these limitations, we propose the following directions for future work:

• Incorporating graph neural networks to better handle scenarios with many interacting bodies
• Exploring techniques for improved generalisation, such as data augmentation and meta-learning
• Integrating physics-based constraints into the loss function to ensure long-term physical consistency
• Investigating hybrid approaches that combine our deep learning model with traditional physics-based methods for improved long-term stability

In conclusion, our deep residual network demonstrates strong performance in predicting 3D rigid body dynamics, outperforming baseline models in both accuracy and computational efficiency [3]. While we have identified areas for improvement, the current results show great promise for applications in robotics, computer graphics, and physical simulations [18].

5. Conclusion and Future Work

This study demonstrates the efficacy of deep residual networks in predicting the dynamics of three-dimensional rigid bodies. By leveraging a sophisticated 3D physics simulator and a carefully designed deep learning architecture, we have advanced the field of physics-informed machine learning [20].

5.1. Key Achievements

Our deep residual network achieves significant improvements over baseline models:

• A 24.8% reduction in position Mean Squared Error (MSE) compared to the Runge-Kutta (RK4) numerical integrator:

  $$MS{E}_{position}=2.37\times {10}^{-3}{\mathrm{m}}^2$$

  (22)
• A 59.4% reduction in position MSE compared to a simple feedforward neural network
• Consistently low relative errors across all state components:

  $$R{E}_{position}=2.18\%,R{E}_{orientation}=1.95\%,R{E}_{linear\_velocity}=3.42\%,R{E}_{angular\_velocity}=2.76\%$$

  (23)
• A 7.9x speedup in inference time compared to the RK4 integrator:

  $$T_{inference}=2.3\mathrm{ms}\mathrm{per}\mathrm{scenario}$$

  (24)

These results underscore our model’s capability to capture complex physical interactions accurately and efficiently, making it suitable for real-time applications in robotics, computer graphics, and physical simulations [19].

5.2. Limitations

Despite these achievements, we have identified several limitations in our current approach:

• Performance degradation in scenarios with many (>10) interacting bodies
• Limited generalisation to object geometries not encountered during training
• Increasing error in long-term predictions beyond 10 seconds, as evidenced by the cumulative error growth:

  $$E_{cumulative}\left(t\right)=\sum _{i=1}^t{\parallel y_i-{\widehat{y}}_i\parallel }^2,t>10s$$

  (25)
• Energy conservation errors, while lower than the feedforward neural network, remain higher than the RK4 integrator:

  $$ECE=0.87\%$$

  (26)

5.3. Future Work

To address these limitations and further advance our research, we propose the following directions for future work:

5.3.1. Graph Neural Networks for Multi-body Interactions

We will explore the integration of Graph Neural Networks (GNNs) to better handle scenarios with many interacting bodies. GNNs can naturally represent the relational structure of multi-body systems, potentially improving performance in complex scenarios [25].

$$h_i^{(l+1)}=\varphi \left(h_i^{\left(l\right)},\sum _{j\in \mathcal{N}\left(i\right)}\psi (h_i^{\left(l\right)},h_j^{\left(l\right)},e_{ij})\right)$$

(27)

where $h_i^{\left(l\right)}$ represents the features of node i at layer l, $\mathcal{N}\left(i\right)$ denotes the neighbours of node i, $e_{ij}$ represents the edge features between nodes i and j, and $\varphi$ and $\psi$ are learnable functions.

5.3.2. Improved Generalisation Techniques

To enhance generalisation to unseen object geometries, we will investigate:

• Data augmentation strategies, including procedural generation of diverse object shapes
• Meta-learning approaches to adapt quickly to new geometries:

  $${\theta }^{*}=arg\underset{\theta }{min}{\mathbb{E}}_{\mathcal{T}\sim p\left(\mathcal{T}\right)}\left[{\mathcal{L}}_{\mathcal{T}}\left(f_{\theta }\right)\right]$$

  (28)
  where $\mathcal{T}$ represents a task (e.g., predicting dynamics for a specific object geometry) sampled from a distribution of tasks $p\left(\mathcal{T}\right)$, and $f_{\theta }$ is our model with parameters $\theta$.

5.3.3. Physics-informed Loss Functions

To improve long-term prediction stability and physical consistency, we will develop physics-informed loss functions that incorporate domain knowledge:

$${\mathcal{L}}_{total}={\mathcal{L}}_{prediction}+{\lambda }_1{\mathcal{L}}_{energy}+{\lambda }_2{\mathcal{L}}_{momentum}$$

(29)

where ${\mathcal{L}}_{energy}$ and ${\mathcal{L}}_{momentum}$ enforce conservation of energy and momentum, respectively, and ${\lambda }_1$, ${\lambda }_2$ are weighting factors.

5.3.4. Hybrid Modelling Approaches

We will explore hybrid approaches that combine our deep learning model with traditional physics-based methods:

$$y_{t+1}=\alpha f_{DL}\left(y_t\right)+(1-\alpha )f_{PB}\left(y_t\right)$$

(30)

where $f_{DL}$ is our deep learning model, $f_{PB}$ is a physics-based model, and $\alpha \in [0,1]$ is a mixing coefficient that can be learned or dynamically adjusted.

5.4. Broader Impact

Our work contributes to the growing field of physics-informed machine learning, offering a powerful tool for predicting complex physical dynamics. The potential applications span various domains:

• Robotics: Enabling more accurate and efficient motion planning and control [26]
• Computer Graphics: Enhancing the realism of physical simulations in games and visual effects [27]
• Scientific Simulations: Accelerating complex physical simulations in fields such as astrophysics and materials science [28]

As we continue to refine and expand our approach, we anticipate that this research will play a crucial role in advancing our ability to model and understand complex physical systems, bridging the gap between data-driven and physics-based modelling approaches.

6. Code

The simulator and deep residual network source code for these experiments are available here1 under the GPL-3.0 open-source license.