LS 3 : Latent Space Safe Sets for Long-Horizon Visuomotor Control of Sparse Reward Iterative Tasks
Abstract:
Reinforcement learning (RL) has shown impressive success in exploring high-dimensional environments to learn complex tasks, but can often exhibit unsafe behaviors and require extensive environment interaction when exploration is unconstrained.
A promising strategy for learning in dynamically uncertain environments is requiring that the agent can robustly return to learned safe sets, where task success (and therefore safety) can be guaranteed.
While this approach has been successful in low-dimensions, enforcing this constraint in environments with visual observations is exceedingly challenging.
We present a novel continuous representation for safe sets by framing it as a binary classiﬁcation problem in a learned latent space, which ﬂexibly scales to image observations.
We then present a new algorithm, Latent Space Safe Sets (LS 3 ), which uses this representation for long-horizon tasks with sparse rewards.
We evaluate LS 3 on 4 domains, including a challenging sequential pushing task in simulation and a physical cable routing task.
We ﬁnd that LS 3 can use prior task successes to restrict exploration and learn more efﬁciently than prior algorithms while satisfying constraints.
See https://tinyurl.com/latent-ss for code and supplementary material.
Introduction
Visual planning over learned forward dynamics models is a popular area of research in robotic control from images [1, 2, 3, 4, 5, 6, 7], as it enables closed-loop, model-based control for tasks where the state of the system is not directly observable or difﬁcult to analytically model, such as the conﬁguration of a sheet of fabric or segment of cable.
These methods learn predictive models over either images or a learned latent space, which can then be used to plan actions which maximize some task reward.
While these approaches have signiﬁcant promise, there are several open challenges in learning policies from visual observations.
First, reward speciﬁcation is particularly challenging for visuomotor control tasks, because high-dimensional observations often do not expose the necessary features required to design dense, informative reward functions [8], especially for long-horizon tasks.
Second, while many prior reinforcement learning methods have been successfully applied to image-based control tasks [9, 10, 11, 12, 13], learning policies from image observations often requires extensive exploration due to the high dimensionality of the observation space and the difﬁculties in reward speciﬁcation, making safe and efﬁcient learning exceedingly challenging.
One promising strategy for efﬁciently learning safe control policies is to learn a safe set [14, 15], which captures the set of states from which the agent is known to behave safely, which is often reformulated as the set of states where it has previously completed the task.
When used to restrict exploration, this safe set can be used to enable highly efﬁcient and safe learning [14, 16, 17], as exploration is restricted to states in which the agent is conﬁdent in task success.
However, while these safe sets can give rise to algorithms with a number of appealing theoretical properties such as convergence to a goal set, constraint satisfaction, and iterative improvement [14, 16, 18], using them for controller design for practical problems requires developing continuous approximations at the expense of maintaining theoretical guarantees [17].
This choice of continuous approximation is a key element in determining the applications to which these safe sets can be used for control.
[Figure]
Prior works have presented approaches which collect a discrete safe set of states from previously successful trajectories and represent a continuous relaxation of this set by constructing a convex hull of these states [14] or via kernel density estimation with a tophat kernel function [17].
While these approaches have been successful for control tasks with low-dimensional states, extending them to high-dimensional observations presents two key challenges: (1) scalability: these prior methods cannot be efﬁciently applied when the number of observations in prior successful trajectories is large, as querying safe set inclusion scales linearly with number of samples it contains and (2) representation capacity: both of these prior approaches do not scale well to high dimensional observations and are limited in the space of continuous sets that they can efﬁciently represent.
Applying these ideas to visuomotor control is even more challenging, since images do not directly expose details about the system state or dynamics that are typically needed for formal controller analysis [14, 16, 19].
This work makes several contributions.
First, we introduce a scalable continuous approximation method which makes it possible to leverage safe sets for visuomotor policy learning.
The key idea is to reframe the safe set approximation as a binary classiﬁcation problem in a learned latent space, where the objective is to distinguish states from successful trajectories from those in unsuccessful trajectories.
Second, we present Latent Space Safe Sets (LS 3 ), a model-based RL algorithm which encourages the agent to maintain plans back to regions in which it is conﬁdent in task completion, even when learning in high dimensional spaces.
This constraint makes it possible to deﬁne a control strategy to (1) improve safely by encouraging consistent task completion (and therefore avoid unsafe behavior) and (2) learn efﬁciently since the agent only explores promising states in the immediate neighborhood of those in which it was previously successful.
Third, we present simulation experiments onvisuomotor control tasks which suggest that LS 3 can learn to improve upon demonstrations more safely and efﬁciently than prior algorithms.
Fourth, we conduct physical experiments on a visionbased cable routing task which suggest that LS 3 can learn more efﬁciently than prior algorithms while consistently completing the task and satisfying constraints during learning.
Related Work 2.1 Safe, Iterative Learning Control
In iterative learning control (ILC), the agent tracks a reference trajectory and uses data from controller rollouts to reﬁne tracking performance [20].
Rosolia et al.
[21], Rosolia and Borrelli [18, 14] present a new class of algorithms, known as Learning Model Predictive Control (LMPC), which are referencefree and instead iteratively improve upon the performance of an initial feasible trajectory.
To achieve this, Rosolia et al.
[21], Rosolia and Borrelli [18, 14] use data from controller rollouts to learn a safe set and value function, with which recursive feasibility, stability, and local optimality can be guaranteed given a known, deterministic nonlinear system or stochastic linear system under certain regularity assumptions.
However, a core challenge with these algorithms is that they assume known system dynamics, and cannot be applied to high-dimensional control problems.
Thananjeyan et al.
[17] extends the LMPC framework to higher dimensional settings in which system dynamics are unknown and must be learned, but the visuomotor control setting introduces a number of new challenges as learned system dynamics, safe sets, and value functions must ﬂexibly scale to visual inputs.
Richards et al.
[15] designs expressive safe sets for ﬁxed policies using neural network classiﬁers with Lyapunov constraints.
In contrast, LS 3 constructs a safe set for an improving policy by optimizing a task cost function instead of uniformly expanding across the state space.
2.2 Model Based Reinforcement Learning
There has been signiﬁcant recent progress in algorithms which combine ideas from model-based planning and control with deep learning [22, 23, 24, 25, 26, 27].
These algorithms are gaining popularity in the robotics community as they enable leaning complex policies from data while maintaining some of the sample efﬁciency and safety beneﬁts of classical model-based control techniques.
However, these algorithms typically require hand-engineered dense cost functions for task speciﬁcation, which can often be difﬁcult to provide, especially in high-dimensional spaces.
This motivates leveraging demonstrations (possibly suboptimal) to provide an initial signal regarding desirable agent behavior.
There has been some prior work on leveraging demonstrations in modelbased algorithms such as Quinlan and Khatib [28] and Ichnowski et al.
[29], which use modelbased control with known dynamics to reﬁne initially suboptimal motion plans, and Fu et al.
[24], which uses demonstrations to seed a learned dynamics model for fast online adaptation using iLQR [24].
Thananjeyan et al.
[17], Zhu et al.
[30] present ILC algorithms which rapidly improve upon suboptimal demonstrations when system dynamics are unknown.
However, these algorithms either require knowledge of system dynamics [28, 29] or are limited to low-dimensional state spaces [24, 17, 30] and cannot be ﬂexibly applied to visuomotor control tasks.
2.3 Reinforcement Learning from Pixels
Reinforcement learning and model-based planning from visual observations is gaining signiﬁcant recent interest as RGB images provide an easily available observation space for robot learning [1, 31].
Recent work has proposed a number of model-free and model-based algorithms that have seen success in laboratory settings in a number of robotic tasks when learning from visual observations [32, 33, 10, 34, 12, 13, 1, 35, 31].
However, two core issues that prevent application of many RL algorithms in practice, inefﬁcient exploration and safety, are signiﬁcantly exacerbated when learning from high-dimensional visual observations in which the space of possible behaviors is very large and the features required to determine whether the robot is safe are not readily exposed.
There has been signiﬁcant prior work on addressing inefﬁciencies in exploration for visuomotor control such as latent space planning [2, 31, 35] and goal-conditioned reinforcement learning [13, 10].
However, safe reinforcement learning for visuomotor tasks has received substantially less attention.
Thananjeyan et al.
[36] and Kahn et al.
[37] present reinforcement learning algorithms which estimate the likelihood of constraint violations to avoid them [36] or reduce the robot’s velocity [37].
Unlike these algorithms, which focus on presenting methods to avoid violating user-speciﬁed constraints, LS 3 additionally provides consistent task completion during learning by limiting exploration to the neighborhood of prior task successes.
This difference makes LS 3 less susceptible to the challenges of unconstrained exploration present in standard model-free reinforcement learning algorithms.
Problem Statement
We consider an agent interacting in a ﬁnite horizon goal-conditioned Markov Decision Processes (MDP) which can be described with the tuple M =
( S , G , A , P ( ·|· , · ) , R ( · , · ) , µ, T ) .
S and A are the state and action spaces, P : S ×
A × S → [0 , 1] maps a state and action to a probability distribution over subsequent states, R : S ×
A × S → R is the reward function, µ is the initial state distribution ( s 0 ∼ µ ), and T is the time horizon.
In this work, the agent is only provided with
RGB image observations s t ∈ R W × H × 3 + = S , where W and H are the image width and height in pixels, respectively.
We consider iterative tasks, where the agent must reach a ﬁxed goal set G ⊆ S as efﬁciently as possible and the support of µ is small.
While there are a number of possible choices of reward functions that would encourage fast convergence to G , providing shaped reward functions can be exceedingly challenging, especially when learning from high dimensional observations.
Thus, as in Thananjeyan et al.
[17], we consider a sparse reward function that only indicates task completion:
R ( s, a, s (cid:48) )
=
if s (cid:48) ∈ G and − 1 otherwise.
To incorporate constraints, we augment M with an extra constraint indicator function C : S → { 0 , 1 } which indicates whether a state satisﬁes user-speciﬁed state-space constraints, such as avoiding known obstacles.
This is consistent with the modiﬁed CMDP formulation used in [36].
We assume that R and C can be evaluated on the current state of the system, but may be approximated using prior data for use during planning.
We make this assumption because
[Figure]
in practice we plan over predicted future states, which may not be predicted at sufﬁciently high ﬁdelity to expose the necessary information to directly evaluate R and C during planning.
Given a policy π : S → A , we deﬁne its expected total return in M as R π = E π,µ,P [ (cid:80) t R ( s t , a t )] .
Furthermore, we deﬁne P πC ( s ) as the probability of future constraint violation (within time horizon T ) under policy π from state s .
The objective is to maximize the expected return R π while maintaining a constraint violation probability lower than δ C .
This can be written formally as follows:
[Equation]
We assume that the agent is provided with an ofﬂine dataset D of transitions in the environment of which some subset D constraint (cid:40) D are constraint violating and some subset D success (cid:40) D appear in successful demonstrations from a suboptimal supervisor.
As in [36], D constraint contains examples of constraint violating behaviors (for example from prior runs of different policies or collected under human supervision) so that the agent can learn about states which violate user-speciﬁed constraints.
Latent Space Safe Sets (LS 3 )
We describe how LS 3 uses demonstrations and online interaction to safely learn iteratively improving policies.
Section 4.1 describes how we learn a low-dimensional latent representation of image observations to facilitate efﬁcient model-based planning.
To enable this planning, we learn a probabilistic forward dynamics model as in [26] in the learned latent space and models to estimate whether plans will likely complete the task (Section 4.2) and to estimate future rewards and constraint violations (Section 4.3) from predicted trajectories.
In Section 4.4, we discuss how these components are synthesized in LS 3 .
Dataset D is expanded using online rollouts of LS 3 and used to update all latent space models (Sections 4.2 and 4.3) after every K rollouts.
See Algorithm 1 and the supplement for further details on training procedures and data collection.
Require: ofﬂine dataset D , number of updates U
4.1 Learning a Latent Space for Planning
Learning compressed representations of images has been a popular approach in vision based control to facilitate efﬁcient algorithms for planning and control which can reason about lower dimensional inputs [2, 35, 6, 38, 39, 31].
To learn such a representation, we train a β -variational autoencoder [40] on states in D to map states to a probability distribution over a d -dimensional latent space Z .
The resulting encoder network f enc ( z | s ) is then used to sample latent vectors z t ∼ f enc ( z t | s t ) to train a forward dynamics model, value function, reward estimator, constraint classiﬁer, safe set, and combine these elements to deﬁne a policy for model-based planning.
Motivated by Laskin et al.
[41], during training we augment inputs to the encoder with random cropping, which we found to be helpful in learning representations that are useful for planning.
For all environments we use a latent dimension of d = 32 , as in [2] and found that varying d did not signiﬁcantly affect performance.
4.2 Latent Safe Sets for Model-Based Control
LS 3 learns a binary classiﬁer for latent states to learn a latent space safe set that represents states from which the agent has high conﬁdence in task completion based on prior experience.
Because the agent can reach the goal from these states, they are safe: the agent can avoid constraint violations by simply completing the task as before.
While classical algorithms use known dynamics to construct safe sets, we approximate this set using successful trajectories from prior iterations.
At each iteration j , the algorithm collects K trajectories in the environment.
We then deﬁne the sampled safe set at iteration j , S j , as the set of states from which the agent has successfully navigated to G in iterations 0 through j of training, where demonstrations trajectories are those collected at iteration 0 .
We refer to the dataset collecting all these states as D success .
This discrete set is difﬁcult to plan to with continuous-valued state distributions so we leverage data from D success (data in the sampled safe set), data from D \ D success (data outside the sampled safe set), and the learned encoder from Section 4.1 to learn a continuous relaxation of this set in latent space (the latent safe set).
We train a neural network with a binary cross-entropy loss to learn a binary classiﬁer f S ( · ) that predicts the probability of a state s t with encoding z t being in S j .
To mitigate the negative bias that appears when trajectories that start in safe regions fail, we utilize the intuition that if a state s t +1 ∈ S j then it is likely that s t is also safe.
To do this, rather than just predict 1 S j , we train f S with a recursive objective to predict max( 1 S j , γ S f S ( s t +1 )) .
The relaxed latent safe set is parameterized by the superlevel sets of f S , where the level δ S is adaptively set during execution:
S j Z = { z t | f S ( · )( z t ) ≥ δ S } .
4.3 Reward and Constraint Estimation
In this work, we deﬁne rewards based on whether the agent has reached a state s ∈ G , but we need rewards that are deﬁned on predictions from the dynamics, which may not correspond to valid real images.
To address this, we train a classiﬁer f G : Z → { 0 , 1 } to map the encoding of a state to whether the state is contained in G using terminal states in D success (which are known to be in G ) and other states in D .
However, in the temporally-extended, sparse reward tasks we consider, reward prediction alone is insufﬁcient because rewards only indicate whether the agent is in the goal set, and thus provide no signal on task progress unless the agent can plan to the goal set.
To address this, as in prior MPC-literature [17, 16, 14, 8], we train a recursively-deﬁned value function (details in the supplement).
Similar to the reward function, we use the encoder (Section 4.1) to train a classiﬁer f C : Z → [0 , 1] with data of constraint violating states from D constraint and the constraint satisfying states in D \ D constraint to map the encoding of a state to the probability of constraint violation.
4.4 Model-Based Planning with LS
LS 3 aims to maximize total rewards attained in the environment while limiting constraint violation probability within some threshold δ C (equation 1).
We optimize an approximation of this objective over an H -step receding horizon with model-predictive control.
Precisely, LS 3 solves the following optimization problem to generate an action to execute at timestep t :
[Equation]
In this problem, the expectations and probabilities are taken with respect to the learned, probabilistic dynamics model f dyn ( z t +1 | z t , a t ) .
The optimization problem is solved approximately using the crossentropy method (CEM) [42] which is a popular optimizer in model-based RL [43, 17, 16, 44, 36].
The objective function is the expected sum of future rewards if the agent executes a t : t + H − 1 and then subsequently executes π (equation 2).
First, the current state s t is encoded to z t (equation 3).
Then, for a candidate sequence of actions a t : t + H − 1 , an H -step latent trajectory { z t +1 , . . .
, z t + H } is sampled from the learned dynamics f dyn (equation 4).
LS 3 constrains exploration using two chance constraints: (1) the terminal latent state in the plan must fall in the safe set (equation 5) and (2) all latent states in the trajectory must satisfy user-speciﬁed state-space constraints (equation 6).
Z C is the set of all latent states such that the corresponding observation is constraint violating.
The optimizer estimates constraint satisfaction probabilities for a candidate action sequence by simulating it repeatedly over f dyn .
The ﬁrst chance constraint ensures the agent maintains the ability to return to safe states where it knows how to do the task within H steps if necessary.
Because the agent replans at each timestep, the agent need not return to the safe set: during training, the safe set expands, enabling further exploration.
In practice, we set δ S for the safe set classiﬁer f S adaptively as described in the supplement.
The second chance constraint encourages constraint violation probability of no more than δ C .
After solving the optimization problem, the agent executes the ﬁrst action in the plan: π ( z t ) =
a t where a t is the ﬁrst element of a ∗ t : t + H − 1 , observes a new state, and replans.
Experiments
We evaluate LS 3 on 3 robotic control tasks in simulation and a physical cable routing task on the da Vinci Research Kit (dVRK) [45].
Safe RL is of particular interest for surgical robots such as the dVRK due to its delicate structure, motivating safety, and relatively imprecise controls [17, 46], motivating closed-loop control.
We study whether LS 3 can learn more safely and efﬁciently than algorithms that do not structure exploration based on prior task successes.
5.1 Comparisons
We evaluate LS 3 in comparison to prior algorithms that behavior clone suboptimal demonstrations before exploring online ( SACfD ) [47] or leverage ofﬂine reinforcement learning to learn a policy using all ofﬂine data before updating the policy online ( AWAC ) [48].
For both of these comparisons we enforce constraints via a tuned reward penalty of λ for constraint violations as in [49].
We also implement a version of SACfD with a learned recovery policy ( SACfD+RRL ) using the Recovery
RL algorithm [36] to use prior constraint violating data to try to avoid constraint violating states.
We then compare LS 3 to an ablated version without the safe set constraint (just binary classiﬁcation (BC)) in equation 5 ( LS 3 ( − Safe Set) ) to evaluate if the safe set promotes consistent task completion and stable learning.
Finally, we compare LS 3 to an ablated version of the safe set classiﬁer (Section 4.2) without a recursive objective, where the classiﬁer is just trained to predict 1 S j ( LS 3 (BC SS) ).
See the supplement for details on hyperparameters and ofﬂine data used for LS 3 and prior algorithms.
5.2 Evaluation Metrics
For each algorithm on each domain, we aggregate statistics over random seeds (10 for simulation experiments, 3 for the physical experiment), reporting the mean and standard error across the seeds.
[Figure]
We present learning curves that show the total sum reward for each training trajectory to study how efﬁciently LS 3 and the comparisons learn each task.
Because all tasks use the sparse task completion based rewards deﬁned in Section 3, the total reward for a trajectory is the time to reach the goal set, where more negative rewards correspond to slower convergence to G .
Thus, for a task with task horizon T , a total reward greater than − T implies successful task completion.
The state is frozen in place upon constraint violation until the task horizon elapses.
We also report task success and constraint satisfaction rates for LS 3 and comparisons during learning to study (1) the degree to which task completion inﬂuences sample efﬁciency and (2) how safely different algorithms explore.
LScollects K = 10 trajectories in between training phases on simulated tasks and K = 5 in between training phases for physical tasks, while the SACfD and AWAC comparisons update their parameters after each timestep.
This presents a metric in terms of the amount of data collected across algorithms.
5.3 Domains
In simulation, we evaluate LS 3 on 3 vision-based continuous control domains that are illustrated in Figure 3.
We evaluate LS 3 and comparisons on a constrained visual navigation task (Pointmass
Navigation) where the agent navigates from a ﬁxed start state to a ﬁxed goal set while avoiding a large central obstacle.
We study this domain to gain intuition and visualize the learned value function, goal/constraint indicators, and safe set in Figure 2.
We then study a constrained image-based reaching task (Reacher) based on [50], where the objective is to navigate the end effector of a 2-link planar robotic arm to a yellow goal position without the end-effector entering a red stay out zone.
We then study a challenging sequential image-based robotic pushing domain (Sequential Pushing), in which the objective is to push each of the 3 blocks forward on the table without pushing them to either side and causing them to fall out of the workspace.
Finally, we evaluate LS 3 with an image-based physical experiment on the da Vinci Research Kit (dVRK) [51] (Figure 3), where the objective is to guide the endpoint of a cable to a goal region without letting the cable or end effector collide with an obstacle.
The Pointmass Navigation and Reaching domains have a task horizon of T = 100 while the Sequential Pushing domain and physical experiment have task horizons of T = 150 and T =respectively.
See the supplement for more details on all domains.
5.4 Simulation Results
We ﬁnd that LS 3 is able to learn more stably and efﬁciently than all comparisons across all simulated domains while converging to similar performance within 250 trajectories collected online (Figure 4).
LS 3 is able to consistently complete the task during learning, while the comparisons, which do not learn a safe set to structure exploration based on prior successes, exhibit much less stable learning.
Additionally, in Table 1 and Table 2, we report the task success rate and constraint violation rate of all algorithms during training.
We ﬁnd that LS 3 achieves a signiﬁcantly higher task success rate than comparisons on all tasks.
We also ﬁnd that LS 3 violates constraints less often than comparisons on the
Reacher task, but violates constraints more often than SACfD and SACfD+RRL on the other domains.
This is because SACfD and SACfD+RRL spend much less time in the neighborhood of constraint violating states during training due to their lower task success rates.
Because they do not efﬁciently learn to perform the tasks, they do not violate constraints as often.
We ﬁnd that the AWAC comparison achieves very low task performance.
While AWAC is designed for ofﬂine reinforcement learning, to the best of our knowledge, it has not been previously evaluated on long-horizon, image-based tasks as in this paper, which we hypothesize are very challenging for it.
[Table] [Table]
[Figure] We ﬁnd LS 3 has a lower success rate when the safe set constraint is removed (LS 3 ( − Safe Set)) as expected.
The safe set is particularly important in the sequential pushing task, and LS 3 ( − Safe Set) has a much lower task completion rate than LS 3 .
LS 3 without the recursive classiﬁcation objective from Section 4.2 (LS 3 (BC SS)) has similar performance to LS 3 on the navigation environment, but learns substantially more slowly on the Reacher environment and performs signiﬁcantly worse than LS 3 on the more challenging Pushing environment as the learned safe set is unable to exploit temporal structure to distinguish safe states from unsafe states.
See the supplement for details on experimental parameters and ofﬂine data used for LS 3 and comparisons and ablations studying the effect of the planning horizon and threshold used to deﬁne the safe set.
5.5 Physical Results
In physical experiments, we compare LS 3 to SACfD and SACfD+RRL (Figure 5) on the physical cable routing task illustrated in Figure 3.
We ﬁnd LS 3 quickly outperforms the suboptimal demonstrations while succeeding at the task signiﬁcantly more often than both comparisons, which are unable to learn the task and also violate constraints more than LS 3 .
We hypothesize that the difﬁculty of reasoning about cable collisions and deformation from images makes it challenging for prior algorithms to make sufﬁcient task progress as they do not use prior successes to structure exploration.
See the supplement for details on experimental parameters and ofﬂine data used for LS 3 and comparisons.
Discussion and Future Work
We present LS 3 , a scalable algorithm for safe and efﬁcient policy learning for visuomotor tasks.
LS 3 structures exploration by learning a safe set in a learned latent space, which captures the set of states from which the agent is conﬁdent in task completion.
LS 3 then ensures that the agent can plan back to states in the safe set, encouraging consistent task completion during learning.
Experiments suggest that LS 3 can safely and efﬁciently learn 4 visuomotor control tasks, including a challenging sequential pushing task in simulation and a cable routing task on a physical robot.
In future work, we are excited to explore further physical evaluation of LS 3 on safety critical visuomotor control tasks and applications to systems with dynamic constraints on velocity or acceleration.