SAFER: Data-Efficient and Safe Reinforcement Learning via Skill Acquisition

To address these criteria, we a latent variable called the safety variable ${\bm{c}}\in\mathcal{C}$ that encodes safety context about the environment, i.e., $f_{\phi}:\mathcal{Z}\times\mathcal{C}\times\mathcal{S}\rightarrow\mathcal{A}.$ This construction encodes information beyond the current state ${\bm{s}}$ to help SAFER model complex per task safety dynamics. For example, the safety variable could encode the locations of people or animals while a robot performs household tasks. Because we do not assume the task variable $\mathcal{C}$ is provided, we infer it from a network.

In order to train the prior $f_{\phi}$ and posterior over the safety variable, we adopt a variational inference (VI) approach. We jointly train an invertible conditional normalizing flow $f_{\phi}$ [42] as the prior $f_{\phi}$ and posterior over the safety variable using VI. At each state $s\in\mathcal{S}$ and safety variable $c\in\mathcal{C}$, the flow model $f_{\phi}$ maps a unit Normal abstract action $z\in\mathcal{Z}$ (i.e., samples $z=f^{-1}_{\phi}({\bm{a}}|{\bm{s}},{\bm{c}})$ of the inverse flow model follow the distribution $p_{\mathcal{Z}}(\cdot):=\mathcal{N}(0,I)$) onto the action space $\mathcal{A}$ of safe behaviors, and thus, the corresponding prior action distribution is given by

The flow model is a good choice for the mapping $f_{\phi}$ because it allows computing exact log likelihoods. Further, it yields a mapping such that actions taken in the abstract action space ${\bm{z}}\in\mathcal{Z}$ can easily be transformed into useful ones ${\bm{a}}=f_{\phi}({\bm{z}};{\bm{s}};{\bm{c}})$. However, since VI approximates the lower bound of maximum likelihood, it does not explicitly enforce the safety requirements in the safety variable ${\bm{c}}$. To overcome this issue, we encode safety to ${\bm{c}}$ by formulating the learning problem as a chance constrained optimization [43] problem.

Chance Constrained Optimization Formally, our objective arises from optimizing a neural network to infer the posterior over the safety variable $\mathcal{C}$ using amortized variational inference [44]. In particular, we parameterize the posterior over the safety variable as $q_{\rho}({\bm{c}}\,|\,\Lambda)$, where ${\bm{c}}$ is the safety variable, and $\Lambda$ is information from which to infer the variable. We set $\Lambda$ as a sliding window of states, such that if ${\bm{s}}_{t}$ is the current state at time $t$ and $w$ is the window size, then the information is given by $\Lambda=\left[{\bm{s}}_{t},{\bm{s}}_{t-1},...,{\bm{s}}_{t-w}\right]$. We infer the safety variable from the sliding window of states $\Lambda$ because we expect $\Lambda$ to contain useful information concerning safe learning. For example, in a robotics setting where the observations are images, previous states may contain useful information concerning the locations of objects to avoid, which may be unobserved in the current state. We write the evidence-lower bound (ELBO) of our model as

where ${\bm{a}}$ is the safe action (i.e, $\omega({\bm{a}},{\bm{s}})=0$) and $p({\bm{c}})$ is a prior over the safety variable ${\bm{c}}$. To ensure that SAFER only samples unsafe actions with low probability, we add a chance constraint about the likelihood of unsafe actions [45] to the ELBO optimization,

where the constraint states that with probability $\xi$ with the safety variable $c$ drawn from $\mathcal{C}$ the distribution of the corresponding unsafe actions (i.e., $\omega({\bm{a}}_{\textrm{unsafe}},{\bm{s}})=1$) is always less than the safety threshold $\epsilon$. Intuitively, this objective enforces that the safety variable makes safe actions as likely as possible while minimizing the probability of unsafe actions.

We denote this objective as the SAFER Contrastive Objective. Further, this objective function has an intuitive interpretation. The first two terms act as a contrastive loss that encourages safe actions (high likelihood) while discourages unsafe ones (low likelihood). Together with the final term, the variable ${\bm{c}}$ is forced to contain useful information about safety. Thus the objective satisfies our goals, allowing for the inference of safety constraints through the task variable and discouraging unsafe behaviors. Finally, since SAFER can increase the likelihood of any safe behaviors, the final criteria that the offline primitive discovery technique can accelerate downstream policy learning will be met by using safe and successful trajectory data during SAFER training.

Parametization Choices To parameterize the SAFER action mapping $f_{\phi}$, we use the Real NVP conditional normalizing flow, proposed by Dinh et al. [42], due to it being highly expressive and allowing exact log-likelihood calculations. Next, we parameterize the posterior distribution $q_{\rho}(c|\Lambda)$ over the safety variable as a diagonal Gaussian to compute the KL efficiently while enabling an expressive latent space. We use a transformer architecture to model the sequential dependency between Gaussian safety variable ${\bm{c}}$ and the window of previous states $\Lambda$ [46]. Finally, because the state space is an image pixel space, we also encode each observation to a vector using a CNN. An overview of the architecture is given in Figure 2.

Training It is necessary to use the reparameterization trick to compute gradients across the objective in Equation 5 [47]. Second, optimizing Equation 5 involves minimizing an unbounded log-likelihood in the second term of the objective. This term can lead to numerical instabilities when $p_{\phi}({\bm{a}}_{\textrm{unsafe}}|{\bm{s}},{\bm{c}})$ is too small. To overcome these issues, we use gradient clipping and freeze this term if it starts to diverge. Psuedo code of the procedure to train SAFER is provided in Appendix F in Algorithm 2 and hyperparameter details are provided in Appendix D.

When using SAFER on a safe RL task, the goal is to accelerate safe learning by leveraging the mapping $f_{\phi}$ in the hierarchical policy $\mu_{\psi}({\bm{s}},{\bm{c}})=\int_{\bm{z}}f_{\phi}({\bm{z}};{\bm{s}};{\bm{c}})d\pi_{\theta}({\bm{z}}|{\bm{s}})$ where the policy parameters of the mapping $\phi$ are fixed and the parameters $\theta$ need to be optimized (Psuedo code of the procedure is provided in Algorithm 1). The policy $\pi_{\theta}({\bm{z}}|{\bm{s}})$ can be learned by any standard RL methods (e.g., SAC [48]) that produces continuous actions. To leverage SAFER at inference time, at each timestep $t$ the RL policy takes an action in the abstract action space ${\bm{z}}_{t}\sim\pi_{\theta}({\bm{z}}|{\bm{s}}={\bm{s}}_{t})$. Using the sliding window of states $\Lambda$, the safety variable posterior computes the distribution over the safety variable $c_{t}$.²²2If there are insufficient states to compute a task window of size $w$ (e.g., at the beginning of the rollout), we pad the available states with $0$’s in order to construct a window of $w$ states. Because a single safety variable value $c_{t}$ is required, we fix it at its mean, $\textrm{E}[{\bm{c}}_{t}]=\int{\bm{c}}\,dq_{\rho}({\bm{c}}|\Lambda_{t})$. Finally, SAFER computes the action ${\bm{a}}_{t}=f_{\phi}({\bm{z}}_{t};{\bm{s}}_{t};\textrm{E}[{\bm{c}}_{t}])$, the action is taken the environment, and the reward $r({\bm{s}}_{t},{\bm{a}}_{t})$ and safety violations $\omega({\bm{s}}_{t},{\bm{a}}_{t})$ are returned. The action ${\bm{z}}_{t}$ and reward $r_{t}$ are added to the replay buffer for subsequent RL training.

As a result, we can construct a latent variable bound around the mean of $\mathcal{Z}$ as the actions $f_{\phi}({\bm{z}};{\bm{c}};{\bm{s}})$ that are more likely to be safe and successful are closer to the mean. Because unsafe actions under the SAFER prior have lower likelihood and our assumption ensures that safe actions exists, there must exist a finite latent variable bound $\eta$ that contains all safe actions. Consequently, with such an $\eta$ from Proposition 4.2, any agent $\pi_{\theta}$ that is trained under the SAFER prior and has a bounded abstraction action output ${\bm{z}}\in(-\eta,\eta)$ is safe. The full proof details are provided in Appendix B.

In practice, we use an offline data set with safe (${\bm{s}}$, ${\bm{a}}$) and unsafe (${\bm{s}}$, ${\bm{a}}_{\text{unsafe}}$) state-action pairs to determine the value of $\eta$ that ensures safety. Also, it is acceptable to fix a range ($-\eta$, $\eta$) that includes a small number of unsafe actions (e.g., at most $1-\epsilon$ portion of all actions in data is unsafe) to avoid an overly tight bound. We optimize the real-valued $\eta>0$ by a numerical gradient-free approach. First, we initialize $\eta=\eta_{0}$ with a large constant and use that to generate the corresponding SAFER abstract actions ${\bm{z}}$ w.r.t. the offline data, whose latent variable is bounded in $(-\eta_{0},\eta_{0})$. We sub-sample this SAFER-action bootstrapped dataset to construct a refined one that only has at most $1-\epsilon$ portion of unsafe actions. Since the normalizing flow in SAFER is invertible, for every (${\bm{s}}$, ${\bm{a}}$)-pair in this data computing the corresponding latent action value is straightforward. This allows us to estimate $\eta_{1}$, which is the maximum latent action in this dataset. We repeat this procedure until convergence. If the offline dataset contains sufficiently diverse state-action data that covers most situations encountered by SAFER, we expect the above safety threshold to be generalizable [49]. We denote this procedure computing the SAFER safety assurances and provide pseudo code in Algorithm 3 in the Appendix.

To evaluate SAFER, we introduce a suite of safety-critical robotic grasping tasks that are inspired by the game Operation³³3https://en.wikipedia.org/wiki/OperationGame.

Safety-critical Robotic Grasping Tasks Based on the game Operation, whose goal is to extract objects from different sized containers without touching the container, we construct a set of $40$ grasping tasks, each consisting of a container and object defined in PyBullet [50]. We collect data from these tasks to train SAFER and use $6$ of the more complex tasks for evaluation. In our tasks, the objects are randomly selected from ones available in PyBullet package, and the containers are generated to fit the objects, whose dimensions (heights and widths) are generated randomly. Our agent controls a $5$DoF robotic arm and gripper. The agent receives positive reward ($r(s,a)=1$) when it extracts the object from the box and a negative reward ($r(s,a)=-1$) at every time step while the task is incomplete. The agent incurs a safety violation ($\omega(s,a)=1$) if the arm touches the box (examples of safe/unsafe trajectories in Figure 3, examples of the tasks are in Figure 10). The states are $48\times 48$ pixel image observations of the scene collected from a fixed camera.

Offline Data Collection To generate the offline data for the SAFER training algorithm, for each robot grasping task we use the scripted policy from Singh et al. [1] to collect trajectories with a total of $1,000,000$ steps. The scripted policy controls the robotic arm to grasp the object generally by minimizing the absolute distance between objects and the robot. To obtain more diverse/exploratory trajectories, one also adds random actuation noise to the policy. After collecting the trajectories, for each state-action pair $({\bm{s}},{\bm{a}})$ in the dataset we provide labels for i) safety violation $\omega(s,a)\in\{0,1\}$, and ii) whether the pair $({\bm{s}},{\bm{a}})$ is part of a successful rollout (i.e., $({\bm{s}},{\bm{a}})$ such that $\mathbb{E}[r({\bm{s}}_{T},{\bm{a}}_{T})|\mu_{\text{data}},{\bm{s}}_{0}={\bm{s}},{\bm{a}}_{0}={\bm{a}}]=1$, where $T$ is the trajectory length random variable). To create the state window $\Lambda$ for SAFER training, for each $({\bm{s}},{\bm{a}})$ in the data buffer we save the previous $w$ states. One can utilize these labels to categorize safe versus unsafe data to train SAFER.

Baseline Comparisons To demonstrate the improved safety performance of SAFER over existing offline primitive learning techniques, we compare against baseline methods that leverage offline data to accelerate learning, including PARROT [1], a contextual version of PARROT (Context. PAR) that uses a latent variable to help accelerate learning, Prior Explore (a method that samples from SAFER to help with data collection during training) and RL from scratch using SAC. See Appendix E for more details. Last, because our setting requires learning safety constraints entirely from labeled offline data, we do not compare against methods that require online safety constraint functions, such as many existing safe RL methods which use a constraint function during training.

Effectiveness of RL training with SAFER In Table 1 we compare SAFER with the baseline methods both in terms of cumulative safety violations and success rate. Note, here we use the underlying reward and safety violation functions for each task to evaluate performance. We choose a SAFER policy primitive with a safety assurance upper bound that guarantees at most $15\%$ unsafe actions, which empirically maintains a good balance between performance and safety. For each downstream task, we then train the RL agent $\pi_{\theta}$ with SAC for only $50,000$ steps because we are more interested to evaluate the power of the primitive learning algorithm. Overall, we see that SAFER has the lowest cumulative safety violations, indicating that it is the most effective method in promoting safe policy learning. Interestingly, SAFER also consistently outperforms other methods in policy performance. The strong success rates of SAFER are potentially due to the fact that discouraging unsafe behaviors may indeed help refining the space of useful behaviors, thus improves policy learning.

Safety Assurance Calibration We evaluate whether the safe abstract action bound of SAFER computed in Section 4.4 is well calibrated, i.e., the empirical percent of unsafe actions should be less than the upper bound. To study this, we compute the $\mathcal{Z}$-action bound $(-\eta,\eta)$ corresponding to an upper bound of $0\%$, $15\%$, $30\%$ and $45\%$ unsafe actions for SAFER. We compute the percentage of unsafe actions by randomly sampling actions from SAFER on each evaluation task and report the results in Figure 4, showing that the SAFER bounds are indeed well calibrated.

Impact of latent safety variable We train SAFER on Tasks $2$ and $5$ using the contrastive objective in Equation 5 but without the safety variable. In this case, the success rate never exceeds $10\%$ and the safety violations are quite high (see Appendix C for the Task 2 results). In contrast, the safety variable in SAFER has at least a $60\%$ success rate on both tasks ( Table 1). This result suggests that the latent safety variable is crucial for success and safety.