SAFE-CHEM: Uncertainty-Aware Policy Switching for Robust Robotic Chemistry

SAFE-CHEM: Uncertainty-Aware Policy Switching for Robust Robotic Chemistry

Abstract Summary:

The paper addresses the deployment of autonomous robotic systems in chemistry laboratories, which accelerates experimental workflows and provides foundational data for AI-driven scientific discovery. Despite the success of data-driven methods in acquiring dexterous skills, safety remains a primary barrier to deployment in high-risk domains such as early-stage materials chemistry experiments. Learning-based policies frequently struggle to distinguish between safe and unsafe actions, leading to overconfident extrapolation and potentially catastrophic failures. To mitigate these risks, the authors propose SAFE-CHEM, an uncertainty-aware framework for robust, learning-based robotic chemists. The approach leverages an ensemble of recurrent neural network-based imitation learning policies to quantify epistemic uncertainty online through the variance of action predictions. By characterising the success-conditioned density of this variance using kernel density estimation, they introduce a hybrid control architecture that autonomously switches from the learned policy to a deterministic, rule-based backup controller when uncertainty exceeds a calibrated safety threshold. The framework is evaluated across three fundamental laboratory manipulation tasks, where empirical results demonstrate that this hybrid strategy improves overall task success rates and reduces critical safety violations compared to traditional single-policy baselines. Finally, the practical viability is demonstrated through zero-shot sim-to-real transfer onto a physical Franka Production 3 robot manipulator.

Introduction Summary:

The integration of robotics and AI has the potential to accelerate traditional human-driven chemistry experiments towards autonomous and more efficient workflows. Robotic chemists have been demonstrated globally across diverse workflows, but most remain restricted to rigid, non-reconfigurable workspaces where changing to novel experiments incurs substantial downtime. Recent works have pivoted towards learning-based methods, specifically imitation learning (IL) and reinforcement learning (RL), allowing manipulators to acquire complex, dexterous skills such as powder weighing, pouring of fluids, or sample scraping. Although IL offers a sample-efficient alternative to RL, its application to chemistry remains underexplored. The widespread adoption of robotic systems in human-centric chemistry laboratories remains contingent on overcoming critical safety hurdles, such as chemical spillage and reactive collision avoidance. A critical challenge in chemistry laboratory-based IL is the distributional shift that occurs when an agent encounters states outside the expert manifold. Because vanilla IL models often fail to distinguish between epistemic uncertainty (lack of data) and aleatoric uncertainty (inherent noise), they are prone to overconfident extrapolation, potentially leading to catastrophic failures without any internal mechanism to trigger a safety intervention or seek human input. SAFE-CHEM treats failure prevention as a real-time uncertainty-quantification problem, utilising an ensemble of neural network policies to monitor epistemic drift during inference and trigger a switching mechanism that transitions control back to a fail-safe state whenever the model’s confidence falls below a critical threshold.

Contributions:

(1) An ensemble-based imitation learning architecture specifically tailored for chemistry laboratory robotic manipulation, enhancing task robustness by capturing distribution of uncertainty across task rollouts to provide a more resilient policy than traditional single-model approaches.

(2) An uncertainty-aware, policy-agnostic safety framework that monitors variance across ensemble predictions to autonomously identify high-risk states and trigger a switching mechanism, effectively balancing task success with the frequency of autonomous intervention.

(3) An empirical quantification of the relationship between policy confidence and task success, providing a systematic analysis of how safety thresholds can be tuned based on the specific risk profile of a laboratory task.

(4) A zero-shot real-world demonstration of the framework onto a physical robotic manipulator.

Related Work Summary:

The adoption of robotic systems in chemistry labs is evolving towards adaptive, modular systems capable of learning new laboratory skills. Current approaches fall into two categories: RL or more recently IL. For instance, Xu et al. demonstrated the first application of imitation learning in a materials laboratory by training a behaviour cloning with a recurrent neural network (BC-RNN) policy on 50 kinesthetic demonstrations to achieve autonomous cathode painting. Suzuki et al. proposed TVF-DiT, a compact imitation learning framework that aligns vision and vision-language foundation models with a diffusion transformer-based action expert. These prior works rely on a single policy, whereas SAFE-CHEM introduces an ensemble-based policy that monitors action uncertainty across task manifolds. For uncertainty estimation, Romer et al. proposed FIPER using random network distillation to identify OOD observations, but such visual-heavy embeddings can be confounded by optical challenges of real-world chemistry labs. FAIL-Detect employs conformal prediction to calibrate safety bands but is sensitive to temporal drift. Zhu et al. proposed a heterogeneity uncertainty quantification method for autonomous navigation, distinguishing between aleatoric and epistemic uncertainties, but primarily for low-dimensional action spaces. SAFE-CHEM distinguishes itself by treating failure prevention as a pre-emptive uncertainty-quantification problem, monitoring epistemic drift to trigger a safety intervention before an irreversible event occurs, without requiring any failure modes to be shown to the robot during training.

Methodology Summary:

The framework uses an ensemble of deep neural networks to capture the predictive distribution across multiple models trained for the same task, aggregating their outputs into a single, robust action. This enables estimation of epistemic uncertainty to identify OOD scenarios. A threshold-triggered switching mechanism transitions control from the learned policy to a deterministic, rule-based fail-safe system whenever uncertainty exceeds a predefined limit.

For imitation learning, the robot manipulation task is formalised as a discrete-time Markov decision process. Behaviour cloning is used to learn a policy directly from expert demonstrations, minimising the expected mean squared error between expert actions and policy predictions. A recurrent neural network architecture is employed to model temporal correlations, maintaining an internal hidden state that serves as a compressed latent representation of interaction history.

For uncertainty quantification, an ensemble of M neural networks is employed, each independently initialised and trained on the expert dataset. Each member generates an action prediction, and the final control signal is the predictive mean. The ensemble captures temporal divergence of recurrent latent states through variance, hypothesising that inter-model variance increases proportionally with prediction error as the robot encounters states further from the training manifold.

For the uncertainty-aware switching mechanism, the design is predicated on stochastic divergence of the action ensemble between nominal and failure states. By modelling the density of action variances observed during successful task executions using kernel density estimation with a Gaussian kernel via Scott’s Rule, component-specific confidence thresholds are calibrated at specified quantiles. During online execution, a sliding window peak detection algorithm ensures that a switch to the safety backup controller occurs only during sustained periods of high variance. If more than two action components are flagged, the system transitions to the backup policy.

Experimental Evaluation Summary:

All simulation environments were developed within IsaacLab, run on a system with an AMD Ryzen Threadripper 7970X processor, 128GB RAM, and an NVIDIA GeForce RTX 5090 GPU. A Franka Production 3 robotic manipulator was used in both simulated and real-world configurations. The learning framework employed an ensemble of BC-RNN models from the RoboMimic framework, using BC as the policy backbone due to its low computational overhead and high inference frequency. The network observes a low-dimensional state vector O ∈ R29 encompassing proprioceptive data and environmental states. Temporal dependencies are processed using a two-layer LSTM architecture with a hidden dimension of 400, trained with a sequence length of 16. Data generation began with 10 human demonstrations collected in simulation using a SpaceMouse, expanded using MimicGen to produce 4,000 valid trajectories per task. The continuous action space is modelled using a Gaussian mixture model with 10 modes. A deep ensemble of 15 distinct BC-RNN models is used, each initialised with a unique random seed and trained on a randomly sampled 25% subset of the total demonstration dataset. Models are trained using the Adam optimiser with an initial learning rate of 1e−4 and a batch size of 100 for 1000 epochs.

The experimental evaluation is conducted across three manipulation tasks: lift (picking up an object and holding it in a desired position for visual inspection), pick and place (picking up an object from an instrument and placing it on another), and insertion (picking up a vial and placing it inside a rack, a high-precision, contact-rich skill). In lift and pick and place tasks, the robot grasps objects from the side rather than from the top, adding complexity as the robot works closer to joint limits. Domain randomisation is introduced with object start and goal positions randomly assigned within a 0.2m × 0.2m region.

Experiment I: Ensemble Policy Evaluation

Task success rate and safety incident rate (defined as any instance where an object is dropped or tilted beyond 45°) are measured. The ensemble of policies runs in an environment with a step size of 0.01s (100Hz) and an action repeat of two, yielding an effective policy control frequency of 50Hz. Results show that employing an ensemble of policies significantly enhances success rates across all tasks. For the lift task, diminishing returns were observed beyond n=10. The pick and place task required n=15 to achieve a significant performance leap. For the insertion task, the policy showed no improvement beyond n=10, and performance declined significantly for n=15, suggesting larger ensembles may produce overly conservative action predictions leading to timeouts. Specific success rates: Lift with n=10 achieved 93.0% ± 1.7%; Pick and Place with n=15 achieved 83.0% ± 2.7%; Insertion with n=10 achieved 55.0% ± 9.8%.

Experiment II: Calibrating the Switching Logic

Trajectory data was collected from ensemble policy rollouts without intervention, classified as successful or failed. KDE applied to aggregated variance data showed a distinct shift toward lower variances for successful trajectories compared to failed ones, validating the assumption that the ensemble policy maintains higher consistency during confident, successful actions. Three confidence thresholds were investigated: conservative (q=0.90), moderate (q=0.95), and relaxed (q=0.99). A sliding window size w ∈ [8, 10] and grid search over npeaks ∈ 1,..., 10 was conducted. An npeaks value of 8 consistently provided the highest F1-score across all tasks when using the conservative θ0.90 limit. This configuration was selected for final implementation to maximise failure detection while minimising unnecessary interruptions.

Experiment III: Confidence Level Intervention

A parallelised, finite-state machine backup controller was implemented utilising NVIDIA Warp for parallelisation across N simulation environments directly on the GPU. The hybrid uncertainty-aware framework consistently improved success rates, with the hybrid system outperforming the non-intervened ensemble in every instance. For the lift task, the best performing model used an ensemble of 10 policies with the most relaxed confidence level (θ0.99), achieving an average success rate of 99.3%. For the pick and place task, the highest success rate was observed with an ensemble of n=5 and the moderate confidence interval (θ0.95), achieving 98.3% ± 1.2%. For the most challenging vial insertion task, the best task success rates were achieved with the most conservative threshold (θ0.90) for all ensemble sizes, with n=15 achieving 68.0% ± 1.7%. Notably, instances were observed where the learned policy resulted in unrecoverable states that failed to trigger the switching mechanism, such as a vial slightly misplaced in the rack remaining improperly seated without exceeding the calibrated uncertainty threshold.

Regarding safety failure rates, for the lift task, the ensemble of 10 policies did not cause any safety violations. For the pick and place task, the hybrid approach reduced safety violations, with the best result of 1.7% ± 1.2% at θ0.95 with n=5. For the insertion task, the ensemble of policies reduced the safety failure rate from 62.3% ± 5.1% (single policy) to 21.0% ± 7.0% (n=10), though the hybrid approach did not further reduce this rate, as manipulation failures predominantly occurred when the robot attempted to insert the vial at an angular offset, leading to irreversible environmental states.

Real-World Deployment Summary:

Experiments were conducted using a 7-DoF FP3 robotic manipulator, focusing on the pick and place task. The system used a zero-shot sim-to-real configuration, where inference was performed by the ensemble policy within the Isaac Lab environment. At each control step, inferred joint position targets were transmitted via a low-latency UDP connection to a custom ROS 2 bridge node, which translated simulated targets into ROS 2 messages for a joint position controller. MoveIt 2 provided the interface to low-level joint tracking controllers. The successful sim-to-real transferred policies and uncertainty-aware switching logic demonstrated high-fidelity generalisation to physical hardware without iterative real-world fine-tuning or domain adaptation.

Conclusion Summary:

SAFE-CHEM uses an ensemble of BC-RNN policies to quantify epistemic uncertainty online through ensemble action variance. By calibrating confidence thresholds using kernel density estimation, the framework introduces an uncertainty-driven switching mechanism that transitions control to a deterministic, rule-based backup controller before a critical failure can occur. Empirical evaluations across three representative chemistry manipulation skills demonstrated that the hybrid framework significantly improves overall task success rates. Crucially, the ensemble of policies reduced the occurrence of catastrophic safety violations, such as dropped glassware or potential chemical spills from tilted beakers, compared to standard single-policy baselines. The practical viability was demonstrated through successful zero-shot sim-to-real transfer onto a physical FP3 robotic manipulator. The method is agnostic to policy, task, and environment, and does not require an exhaustive list of all possible failure modes. Future work will explore extending the uncertainty-quantification framework to incorporate multi-modal sensory inputs, such as force-torque feedback and visual-based policies, with a focus on long-horizon multi-skill experiments.

Improvements to AI Systems:

1. Uncertainty-Aware Action Selection: Implement an ensemble of recurrent neural network policies (e.g., 10–15 models) that output a distribution of actions. Use the variance across ensemble predictions as a real-time measure of epistemic uncertainty, and compute the mean action as the primary control signal. This replaces single-policy extrapolation with a consensus-based approach that inherently flags out-of-distribution states.

2. Adaptive Safety Switching via Calibrated Thresholds: Integrate a kernel density estimation (KDE) layer that models the variance distribution from successful task executions. Set a safety threshold at a specific quantile (e.g., 90th percentile) of this distribution. During inference, continuously monitor ensemble variance; if variance exceeds the threshold for a sustained period (e.g., using a sliding window of 8–10 timesteps with a peak detection algorithm), automatically switch control to a deterministic, rule-based backup controller (e.g., a finite-state machine) to prevent catastrophic failures.

3. Pre-Emptive Failure Prevention Without Training on Failures: Train the ensemble solely on expert demonstrations (imitation learning) without requiring any failure examples. The system detects high-risk states by measuring inter-model divergence, which increases as the robot moves away from the expert manifold. This enables the AI to intervene before an irreversible event (e.g., dropping a vial, spilling chemicals) occurs, rather than reacting after the fact.

4. Task-Specific Risk Profiling: Calibrate the safety threshold (conservative, moderate, or relaxed) based on the task’s risk profile. For high-precision, contact-rich tasks (e.g., insertion), use a conservative threshold (90th percentile) to maximize safety; for simpler tasks (e.g., lifting), use a relaxed threshold (99th percentile) to avoid unnecessary interruptions and improve success rates. This allows the AI to balance autonomy and safety dynamically.

5. Ensemble Size Optimization: Automatically select the optimal number of ensemble members per task. For tasks with low complexity, use fewer models (e.g., 5–10) to reduce computational overhead; for tasks requiring high robustness, increase to 15 models. Monitor performance gains and avoid oversizing, as larger ensembles can lead to overly conservative actions and timeouts (e.g., insertion task performance degrades at n=15).

6. Zero-Shot Sim-to-Real Transfer with Uncertainty Monitoring: Deploy the trained ensemble and switching logic directly onto physical robots without fine-tuning. Use low-latency communication (e.g., UDP) to transmit ensemble-averaged joint targets from simulation to real hardware, while retaining the uncertainty-based safety switch in the real-world loop. This enables safe, out-of-the-box deployment in unstructured environments.

What the Improved AI System Can Do:

• Operate autonomously in high-risk environments (e.g., chemistry labs) with minimal human supervision, by continuously assessing its own confidence and triggering a fail-safe backup when uncertain, reducing the risk of spills, broken glassware, or equipment damage.

• Achieve higher task success rates than single-policy baselines (e.g., lift success from 85% to 99.3%, pick-and-place from 70% to 98.3%) by leveraging ensemble consensus and pre-emptive intervention.

• Detect and avoid out-of-distribution states in real-time, even without prior exposure to failure scenarios, by monitoring epistemic uncertainty—this is critical for adapting to novel objects, positions, or environmental changes.

• Switch between learned and rule-based control seamlessly, ensuring that the system remains functional even when the learned policy is uncertain, without requiring manual reset or human intervention.

• Transfer learned skills from simulation to physical robots with zero-shot generalization, including uncertainty-aware safety logic, making it viable for rapid deployment in real laboratories without costly iterative tuning.

• Scale across diverse manipulation tasks (e.g., lifting, pick-and-place, insertion) by tuning safety thresholds and ensemble sizes per task, offering a flexible framework for various robotic chemistry workflows.

The deployment of autonomous robotic systems in chemistry laboratories is accelerating experimental workflows and providing the foundational data for AI-driven scientific discovery. However, despite the success of data-driven methods in acquiring dexterous skills, safety remains a primary barrier to their deployment in high-risk domains, such as early-stage materials chemistry experiments. Specifically, learning-based policies frequently struggle to distinguish between safe and unsafe actions, leading to overconfident extrapolation and potentially catastrophic failures. To mitigate these safety risks, we propose SAFE-CHEM, an uncertainty-aware framework designed for robust, learning-based robotic chemists. Our approach leverages an ensemble of recurrent neural network-based imitation learning policies to quantify epistemic uncertainty online through the variance of action predictions. By characterising the success-conditioned density of this variance using kernel density estimation, we introduce a hybrid control architecture that autonomously switches from the learned policy to a deterministic, rule-based backup controller when uncertainty exceeds a calibrated safety threshold. We evaluate SAFE-CHEM across three fundamental laboratory manipulation tasks, where our empirical results demonstrate that this hybrid strategy improves overall task success rates and reduces critical safety violations compared to traditional single-policy baselines. Finally, we demonstrate the practical viability of the framework through zero-shot sim-to-real transfer onto a physical Franka Production 3 robot manipulator.