Sequence-based imitation learning for surgical robot operations

Generation of 3D environments from patient CT scans

The creation of realistic 3D environments from patients’ CT scans is essential for accurate surgical simulations. In this section, we detail the process of transforming CT scans into anatomically accurate 3D models of the kidney and tumor, employing advanced AI methods and simulation software. The first step in the process involves segmenting the kidney and tumor regions from the CT scans. To accomplish this, we employed a state-of-the-art deep learning architecture known as U-Net^[25]. The U-Net model is trained to perform semantic segmentation, accurately delineating organ boundaries within medical images. By training the U-Net on annotated CT scans, we obtained precise segmentation masks outlining the kidney and tumor regions^[26]. With the segmented regions extracted, we proceeded to reconstruct 3D models of the kidney and tumor. Each segmented region was converted into a three-dimensional mesh representation using standard techniques. The mesh generation process involved interpolating the segmented contours to create a continuous surface, resulting in anatomically faithful 3D representations. Following mesh generation, we conducted refinement and optimization procedures to enhance the quality of the 3D models. This included smoothing the mesh surfaces to remove artifacts and irregularities introduced during the reconstruction process. Additionally, we optimized the mesh topology to ensure computational efficiency and stability during simulation. The 3D models are then refined and colorized, imported into the SOFA software environment, as shown in Figure 1. SOFA facilitated the integration of the 3D models into dynamic simulation scenarios, enabling real-time deformation and interaction with surgical instruments. The anatomically accurate 3D environments created from patient CT scans served as the foundation for realistic surgical simulations and were used to perform the task, i.e., tumor cutting and removal from the kidney. In our work, we generated a total of 53 distinct kidney tumor environments, each characterized by unique variations in tumor morphology and relative spatial arrangement since each of them corresponds to a different patient. To introduce diversity in surgical scenarios, we varied the shape and size of tumors across different environments. This variation in tumor morphology influenced the trajectories of surgical instruments, as surgeons navigated through different tissue structures and encountered varying degrees of resistance. Additionally, the size and shape variations ensured that each simulation presented unique challenges, reflecting the heterogeneity observed in clinical scenarios. Another factor contributing to the diversity of environments was the relative positioning of the tumor with respect to the kidney. By manipulating the spatial relationship between the tumor and the kidney, we translated the trajectories of surgical tools in space, simulating scenarios where tumors were located in different regions of the kidney. These differences were carefully chosen to introduce variability in the trajectories of surgical tools and to create visually distinct environments for simulation purposes, so that the ViT-LSTM model can learn the tumor-kidney spatial relationships and generalize across different scenarios. Despite these variations, the position of the kidney remained consistent across all environments, centered at the origin of the SOFA environment.

Figure 1. 3D reconstruction of the Kidney tumor environment imported into SOFA. SOFA: Simulation open framework architecture.

Video demonstrations execution

To capture surgical demonstrations for the imitation learning dataset, we employed a manual execution approach conducted by a single engineer, as our primary focus was testing the model across different environments rather than evaluating variability across different experts in the same environment. This approach enabled us to assess the model’s adaptability and robustness in diverse scenarios. The demonstrations were made using specialized surgical tools within a simulated environment. The task of tumor removal from kidneys was executed using two distinct instruments: a cauter for the cutting phase and a grasper for grasping and removing the tumor post-incision. These instruments were controlled manually via a console controller, allowing for precise manipulation within the simulation environment. The integrated RGB-D camera is used as a stereoscope to capture videos of the demonstrations. The cauter, employed for the cutting phase, simulated the use of electro-cautery to incise tissue. This instrument was responsible for executing precise incisions along predetermined trajectories within the kidney tumor environment. Conversely, the grasper functioned as a versatile tool for grasping and removing the dissected tumor fragments from the surgical site. Together, these tools enabled the execution of a complete tumor removal procedure within the simulation. The execution of surgical demonstrations was facilitated by a modified version of the tissue dissection application developed within the SOFA_ENV framework, utilizing LapGym^[27]. This application was tailored to our specific use case, providing a specialized environment for executing kidney tumor removal procedures. By leveraging the capabilities of LapGym, a laparoscopic surgery simulation module within SOFA_ENV, we were able to recreate realistic surgical scenarios while enabling manual control of surgical instruments. The cutting procedure starts from the left side of the tumor and proceeds clockwise along the visible perimeter. Once the cutting edge has completed more than half of its trajectory and passed the gripping tool, positioned on the right side, the latter approaches the tumor to grasp it and establish one secure grip, holding it in place for the final stage of the procedure. Once all contact points are removed, the tumor is completely separated from the kidney and must be excised from the surrounding healthy tissue: the task is considered complete when the tumor is positioned at a safe distance from the organ. Some of these steps are shown in Figure 2. It is very important to carry out all simulations, always respecting the same guidelines to ensure a certain consistency between the dataset demonstrations: not only is it essential to carry out the task correctly, but it is equally crucial to obtain a homogeneous dataset so that, during the training phase, the learning agent can recognize visual and sequential patterns in the actions performed across different environments.

Figure 2. Some steps of the execution of tumor cutting and removal from the kidney executed by the expert. The cauter (tool on the left) cuts the tumor and the grasper (tool on the right) takes the tumor out of the kidney.

Dataset composition

The dataset for our imitation learning project is composed of multi-modal data captured during the execution of surgical demonstrations within the simulated kidney tumor environments. Each demonstration is recorded using an RGB-D camera, providing both visual and depth information essential for understanding the surgical scene. The primary component of the dataset consists of videos captured through RGB-D cameras, with each demonstration recorded as a sequence of frames. These videos offer a detailed visual representation of the surgical procedure, enabling the observation of tool manipulation and tissue interaction throughout the task execution. In addition to RGB video frames, depth information is also captured for every frame of each video using the same RGB-D camera. These depth data provide crucial spatial context, allowing for accurate estimation of the 3D geometry of the surgical scene. Depth information is invaluable for understanding the relative positions of surgical instruments, tissues, and anatomical structures within the environment. Alongside video and depth data, the dataset includes precise 3D pose information for both the cauter and the grasper instruments. The pose data comprises the XYZ position and unit quaternion for orientation of the tips of both tools at each time step of the trajectory. These tool poses serve as ground truth data for our imitation learning model, guiding the prediction of tool actions in subsequent time steps. Additionally, two more variables, one for each tool, are provided, describing the state of the tools: for the cauter, a boolean variable is considered with the value of 1 if the tool is cutting, 0 otherwise; for the grasper, the variable considered indicates the opening angle of the tool. In conclusion, the creation of the dataset described herein was necessitated by the absence of suitable resources for learning surgical tasks in the intended manner. Despite the growing interest in surgical automation and imitation learning, existing datasets fail to adequately address the specific requirements of our task. Real-world surgical datasets often lack the detail and precision required for effective imitation learning. Obtaining annotated data from actual surgical procedures is challenging due to ethical and logistical constraints. Additionally, real surgical environments may exhibit variability that is difficult to control and replicate consistently. A challenge for the future will be to have a good-quality dataset with real-world data. On the other side, while simulation platforms would offer a controlled and safe environment and a convenient, albeit preliminary, way for data collection, there are very few platforms where surgical environments can be simulated, and the existing datasets within these platforms often lack the nuanced details necessary for accurate imitation learning, such as realistic tissue deformation, instrument-tissue interaction dynamics, and the kinematics data of the tools exploited. Furthermore, the specific task of tumor removal from kidneys presents unique challenges that necessitate the creation of a tailored dataset.

Problem formulation

At its core, imitation learning seeks to train an agent to replicate expert actions based on observed demonstrations, without requiring explicit task specification or reward design. One common formulation of imitation learning is behavioral cloning from observation^[28], where an agent learns to imitate expert behavior by directly mapping observed states to corresponding actions. In this paradigm, the agent aims to replicate the demonstrated behavior without explicitly understanding the underlying task dynamics. Instead, it learns a mapping from state observations to actions through supervised learning techniques. Central to the success of imitation learning is the notion of policy learning. A policy represents the strategy or decision-making function employed by an agent to select actions based on its observations of the environment. In our particular context, we employ behavioral cloning to train an agent for performing surgical tasks, such as tumor removal from kidneys. At each step of the task execution, the agent must decide on an action based on the current state of the environment and past states. This involves finding a policy - a function of the current visual observation and the state of the surgical tools (e.g., kinematic data) - that dictates the agent’s actions to achieve the desired surgical outcome, i.e., the state the tools have to assume in the next time step. We can formulate our problem as follows: let $$ s_t $$ denote the current state of the environment at time step $$ t $$, and $$ a_t $$ represent the action taken by the agent in response to $$ s_t $$, and the agent learns a mapping from states to actions, represented by a policy $$ \pi(s) $$. This mapping is learned from a dataset of expert demonstrations $$ \{(s_1, a_1), (s_2, a_2), \ldots, (s_N, a_N)\} $$, where $$ N $$ is the number of demonstrations. The goal of policy learning is to find an optimal policy $$ \pi^* $$ that minimizes the discrepancy between the agent’s actions and the expert actions. In the case of supervised learning, this discrepancy is typically measured using a loss function $$ L(\pi(s), a) $$.

where $$ s_i $$ and $$ a_i $$ are the states and actions from the expert demonstrations, respectively.

In our context, the agent’s action at each time step is determined by a policy $$ \pi(s_t, \theta) $$, parameterized by $$ \theta $$. The policy is learned to mimic the expert behavior observed in the demonstrations, where $$ \theta $$ is optimized to minimize the loss between the agent’s actions and the expert actions, as follows

where $$ N $$ is the number of demonstrations. The state $$ s_t $$ is composed of the observation and the robot tools’ state (their kinematic data), denoted as $$ s_t = (o_t, \text{tools}_t) $$. Additionally, the parameter $$ \theta $$ depends on past states, reflecting the sequential nature of policy learning in imitation learning tasks.

Let $$ s_t $$ denote the state of the environment at time step $$ t $$, consisting of the observation and the predicted action at time $$ t $$. Additionally, let $$ w $$ represent the size of the window of past states, including both past observations and past predicted actions.

The action at time step $$ t+1 $$, denoted as $$ a_{t+1} $$, is obtained based on the application of the policy function $$ \pi(s_t, \theta) $$ to the current state $$ s_t $$, along with a window $$ w $$ of past states $$ s_{t-w+1}, s_{t-w+2}, \ldots, s_t $$:

where $$ \pi(s_t) $$ is the policy function that takes the current state $$ s_t $$ as input and predicts the action $$ a_{t+1} $$ to be taken at the next time step.

Each state $$ s_t $$ in the window is composed of both past observations and past predicted actions:

where $$ o_t $$ represents the observation at time step $$ t $$, and $$ a_t $$ represents the predicted action at time step $$ t $$.

Proposed architecture

The imitation learning model employed in our study leverages a hybrid architecture comprising a ViT^[29], pretrained on ImageNet, and a LSTM network as shown in Figure 3. This architecture is designed to effectively handle the multi-modal nature of the dataset, integrating both visual information from RGB-D frames and temporal sequences of tool poses. The model takes, as input, sequences of RGB-D frames along with past predictions iteratively to predict the actions that the surgical tools should take in the next time step, i.e., the poses every tool has to assume in the next step. By iteratively refining its predictions based on past observations and ground truth tool poses, the model learns to mimic the expert demonstrations and perform surgical tasks autonomously in new kidney tumor environments. The ViT serves as the backbone of our model, specifically configured with 16 patches and a resolution of 224 pixels, and is responsible for processing the visual information extracted from RGB-D frames. By employing self-attention mechanisms, the ViT can capture spatial relationships and long-range dependencies within the video data, enabling effective feature extraction from each frame. This allows the model to encode rich visual representations of the surgical scene.

Figure 3. Cycle to predict an action for a single step of the trajectory.

In our model, the LSTM network is a sequence-to-one architecture, where it takes a sequence of input features and predicts a single output value. It plays a crucial role in processing sequences of features extracted by the ViT: it is responsible for capturing temporal dependencies and learning the dynamics of the surgical task over time. We consider the trajectories of surgical tools as sequences of steps, where the action at time step t is influenced by both the preceding steps and the current state of the environment. This approach enables us to effectively model the long-term dynamic of the trajectories in surgical operations. We opted for an LSTM network instead of a traditional deep network because LSTMs are specifically designed to handle sequential data, such as the trajectories in our study. The LSTM’s ability to retain information over time and capture dependencies between time steps allows it to model the sequential nature of surgical operations more effectively. This approach resulted in better performance and more accurate predictions compared to using a standard deep network. Moreover, an attention mechanism is incorporated within the LSTM architecture to enhance the model’s ability to focus on relevant parts of the input sequence. The LSTM layer processes the input sequence, which consists of features extracted by the ViT at time step $$ t $$ and the previous $$ w $$ features, along with the corresponding predicted actions of the robot tools. This sequence is fed into the LSTM, allowing the model to capture long-term dependencies and temporal dynamics.

Following the LSTM layer, an attention mechanism is employed to dynamically weigh the importance of each input feature in the sequence. The attention mechanism computes attention weights using a linear layer applied to the LSTM outputs. These weights reflect the relevance of each feature in the sequence, allowing the model to focus more on informative features while attenuating the impact of less relevant ones. The attention weights computed by the linear layer are used to compute a weighted sum of the LSTM outputs, yielding an attended input representation. This attended input is then passed through another linear layer to produce the final output, which represents the action the robot tools are expected to take at the next time step. The integration of the attention mechanism within the LSTM architecture enhances the model’s performance by enabling it to selectively focus on relevant information in the input sequence, thereby improving its predictive capabilities.

Training of the model

During each training cycle, the model processes a single RGB-D frame and corresponding data from previous time steps to predict the actions of the surgical tool. The process unfolds as follows: (1) The RGB-D frame at time $$ t $$ is divided into patches and passed through the ViT and a multilayer perceptron (MLP) to extract high-dimensional features capturing spatial and temporal information; (2) The features extracted at time $$ t $$ are combined with features from previous time steps $$ t-1 $$ to $$ t-w $$, where $$ w $$ is the window size of past states considered in the input sequence. Additionally, past predicted actions from time $$ t $$ to $$ t-w $$ are included in the fusion process. This temporal fusion mechanism enables the model to leverage historical information and past predictions for context-aware action prediction; (3) The temporally fused features are passed into an Attention LSTM block, which comprises three main components: an LSTM layer, an attention mechanism, and a MLP. The LSTM layer processes the temporally fused features, capturing temporal dependencies and dynamics across the input sequence.The attention mechanism computes attention weights using a linear layer applied to the LSTM outputs. These weights reflect the relevance of each feature in the sequence, allowing the model to focus more on informative features while attenuating the impact of less relevant ones. The output of the attention mechanism is then passed through the MLP to produce the final prediction. The MLP outputs a 16-element vector comprising the xyz positions, unit quaternions of the two surgical tools, the cutting/not-cutting boolean variable for the cauter, and the opening angle of the grasper. This process is repeated for each point in the trajectory and for every trajectory in the dataset (i.e., every video of the dataset), allowing the model to learn and generalize across diverse surgical scenarios.

For training the model, we employed Adam optimizer, with an initial learning rate of $$ 10^{-5} $$, which is reduced if the loss does not improve for 5 consecutive epochs (best results obtained to $$ 10^{-6} $$). This adaptive learning rate schedule helps the model to converge more effectively and avoid getting stuck in local minima. We empirically determined that a window parameter of 10 elements in the input sequence provides a good balance between capturing temporal dependencies and computational efficiency. This window size allows the model to consider relevant past states and actions while avoiding excessive computational overhead. The model was trained for 50 epochs. Through experimentation, we observed that training beyond this point did not lead to significant improvements in performance, indicating that the model had converged to a stable solution. A custom loss function, which we call Temporal Loss, is introduced to address the challenge of mispredictions in later time steps. This loss function is designed to penalize errors more heavily as the predicted actions deviate further from the ground truth over time. By assigning higher weights to errors in later time steps, the Temporal Loss encourages the model to prioritize accurate predictions for actions in the immediate future while still considering the overall trajectory of the surgical task, since, due to the complex and dynamic nature of surgical procedures, the model may encounter difficulties in accurately predicting actions in later stages of the task. This is often exacerbated by accumulated errors and uncertainties as the task progresses. The dataset was split into training and testing sets using an 85/15 ratio. This division ensures that the model is trained on a sufficiently large portion of the data while still reserving a separate subset for evaluating its generalization performance on unseen examples.