LS 3 : Latent Space Safe Sets for Long-Horizon Visuomotor Control of Sparse Reward Iterative Tasks
Anonymous Author(s)
AfﬁliationAddress email
Abstract:
Reinforcement learning (RL) has shown impressive success in exploring high-dimensional environments to learn complex tasks, but can often exhibitunsafe behaviors and require extensive environment interaction when explorationis unconstrained.
A promising strategy for learning in dynamically uncertainenvironments is requiring that the agent can robustly return to learned safe sets,where task success (and therefore safety) can be guaranteed.
While this approachhas been successful in low-dimensions, enforcing this constraint in environmentswith visual observation spaces is exceedingly challenging.
We present a novelcontinuous representation for safe sets framed as a binary classiﬁcation problemin a learned latent space, which ﬂexibly scales to high-dimensional image observations.
We then present a new algorithm, Latent Space Safe Sets (LS 3 ), whichuses this representation for long-horizon control.
We evaluate LS 3 on 4 domains,including a challenging sequential pushing task in simulation and a physical cablerouting task.
We ﬁnd that LS 3 can use prior task successes to restrict explorationand learn more efﬁciently than prior algorithms while satisfying constraints.
Seehttps://tinyurl.com/latent-safe-sets for supplementary material.
Introduction
Visual planning over learned forward dynamics models is a popular area of research in roboticcontrol from images [1, 2, 3, 4, 5, 6, 7], as it enables closed-loop, model-based control for taskswhere the state of the system is not directly observable or difﬁcult to analytically model, such as theconﬁguration of a sheet of fabric or segment of cable.
These methods learn predictive models overeither images or a learned latent space, which can then be used to plan actions which maximize sometask reward.
While these approaches have signiﬁcant promise, there are several open challenges inlearning policies from visual observations.
First, reward speciﬁcation is particularly challenging forvisuomotor control tasks, because high-dimensional observations often do not expose the necessaryfeatures required to design dense, informative reward functions [8], especially for long-horizontasks.
Second, while many prior reinforcement learning methods have been successfully applied toimage-based control tasks [9, 10, 11, 12, 13], learning policies from image observations often requiresextensive exploration due to the high dimensionality of the observation space and the difﬁculties inreward 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 oftenreformulated as the set of states where it has previously completed the task.
When used to restrictexploration, this safe set can be used to enable highly efﬁcient and safe learning [16, 17, 14], asexploration is restricted to states in which the agent is conﬁdent in task success.
However, whilethese safe sets can give rise to algorithms with a number of appealing theoretical properties such asconvergence to a goal set, constraint satisfaction, and iterative improvement [14, 18, 16], using themfor controller design for practical problems requires developing continuous approximations at theexpense of maintaining theoretical guarantees [17].
This choice of continuous approximation is a keyelement 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 previouslysuccessful trajectories and represent a continuous relaxation of this set by constructing a convex hullof these states [14] or via kernel density estimation with a tophat kernel function [17].
While theseapproaches have been successful for control tasks with low-dimensional states, extending them tohigh-dimensional observations presents two key challenges: (1) scalability: these prior methodscannot 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) representationcapacity: both of these prior approaches do not scale well to high dimensional observations and arelimited in the space of continuous sets that they can efﬁciently represent.
Applying these ideas tovisuomotor control is even more challenging, since images do not directly expose details about thesystem state or dynamics that are typically needed for formal controller analysis [16, 14, 19].
This work makes several contributions.
First, we introduce a scalable continuous approximationmethod which makes it possible to leverage safe sets for visuomotor policy learning.
The key idea is toreframe the safe set approximation as a binary classiﬁcation problem in a learned latent space, wherethe 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 encouragesthe agent to maintain plans back to regions in which it is conﬁdent in task completion, even whenlearning 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 neighborhoodof those in which it was previously successful.
Third, we present simulation experiments on 3visuomotor control tasks which suggest that LS 3 can learn to improve upon demonstrations moresafely 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 algorithmswhile 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 controllerrollouts to reﬁne tracking performance [20].
Rosolia et al.
[21], Rosolia and Borrelli [18, 14] presenta 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 achievethis, Rosolia et al.
[21], Rosolia and Borrelli [18, 14] use data from controller rollouts to learna safe set and value function, with which recursive feasibility, stability, and local optimality canbe guaranteed given a known, deterministic nonlinear system or stochastic linear system undercertain regularity assumptions.
However, a core challenge with these algorithms is that they assumeknown system dynamics, and cannot be applied to high-dimensional control problems. Thananjeyanet al.
[17] extends the LMPC framework to higher dimensional settings in which system dynamicsare 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 visualinputs.
Richards et al. [15] designs expressive safe sets for ﬁxed policies using neural networkclassiﬁers with Lyapunov constraints.
In contrast, LS 3 constructs a safe set for an improving policyby 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-basedplanning and control with deep learning [22, 23, 24, 25, 26, 27].
These algorithms are gainingpopularity in the robotics community as they enable leaning complex policies from data whilemaintaining some of the sample efﬁciency and safety beneﬁts of classical model-based controltechniques.
However, these algorithms typically require hand-engineered dense cost functions fortask 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 regardingdesirable 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 usingiLQR
Thananjeyan et al.
[17], Zhu et al.
[30] present ILC algorithms which rapidly improveupon suboptimal demonstrations when system dynamics are unknown.
However, these algorithmseither require knowledge of system dynamics [28, 29] or are limited to low-dimensional statespaces [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ﬁcantrecent 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 successin 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 algorithmsin practice, inefﬁcient exploration and safety, are signiﬁcantly exacerbated when learning fromhigh-dimensional visual observations in which the space of possible behaviors is very large and thefeatures required to determine whether the robot is safe are not readily exposed.
There has beensigniﬁcant prior work on addressing inefﬁciencies in exploration for visuomotor control such aslatent space planning [2, 31, 35] and goal-conditioned reinforcement learning [13, 10].
However, safereinforcement learning for visuomotor tasks has received substantially less attention.
Thananjeyanet al.
[36] and Kahn et al.
[37] present reinforcement learning algorithms which estimate the likelihoodof 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 additionallyprovides consistent task completion during learning by limiting exploration to the neighborhood ofprior task successes.
This difference makes LS 3 less susceptible to the challenges of unconstrainedexploration 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 arethe state and action spaces, P : S × A × S → [0 , 1] maps a state and action to a probabilitydistribution over subsequent states, R : S ×
A × S → R is the reward function, µ is the initialstate 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 inpixels, respectively.
We consider iterative tasks, where the agent must reach a ﬁxed goal set G ⊆ S asefﬁciently as possible and the support of µ is small.
While there are a number of possible choices ofreward functions that would encourage fast convergence to G , providing shaped reward functions canbe exceedingly challenging, especially when learning from high dimensional observations.
Thus, asin 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 extraconstraint indicator function C : S → { 0 , 1 } which indicates whether a state satisﬁes user-speciﬁedstate-space constraints, such as avoiding known obstacles.
This is consistent with the modiﬁed CMDPformulation 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 maintaininga 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 ofwhich some subset D constraint (cid:40) D are constraint violating and some subset D success (cid:40) D appear insuccessful demonstrations from a suboptimal supervisor.
As in [36], D constraint contains examples ofconstraint violating behaviors (for example from prior runs of different policies or collected underhuman 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 improvingpolicies.
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 probabilisticforward dynamics model as in [26] in the learned latent space and models to estimate whether planswill 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 aresynthesized in LS 3 .
Dataset D is expanded using online rollouts of LS 3 and used to update all latentspace models (Sections 4.2 and 4.3) after every K rollouts.
See Algorithm 1 and the supplement forfurther 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 controlto 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 .
Theresulting encoder network f enc ( z | s ) is then used to sample latent vectors z t ∼ f enc ( z t | s t ) to train aforward dynamics model, value function, reward estimator, constraint classiﬁer, safe set, and combinethese elements to deﬁne a policy for model-based planning.
Motivated by Laskin et al.
[41], duringtraining we augment inputs to the encoder with random cropping, which we found to be helpful inlearning representations that are useful for planning.
For all environments we use a latent dimensionof 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 statesfrom which the agent has high conﬁdence in task completion based on prior experience.
Because theagent can reach the goal from these states, they are safe: the agent can avoid constraint violations bysimply completing the task as before.
While classical algorithms use known dynamics to constructsafe sets, we approximate this set using successful trajectories from prior iterations.
At each iterationj , the algorithm collects K trajectories in the environment.
We then deﬁne the sampled safe set atiteration j , S j , as the set of states from which the agent has successfully navigated to G in iterations0 through j of training, where demonstrations trajectories are those collected at iteration 0 .
Werefer to the dataset collecting all these states as D success .
This discrete set is difﬁcult to plan to withcontinuous-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 tolearn a continuous relaxation of this set in latent space (the latent safe set).
We train a neural networkwith a binary cross-entropy loss to learn a binary classiﬁer f S ( · ) that predicts the probability of astate s t with encoding z t being in S j .
To mitigate the negative bias that appears when trajectoriesthat start in safe regions fail, we utilize the intuition that if a state s t +1 ∈ S j then it is likely thats t is also safe.
To do this, rather than just predict 1 S j , we train f S with a recursive objective topredict max( 1 S j , γ S f S ( s t +1 )) .
The relaxed latent safe set is parameterized by the superlevel sets off 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 needrewards that are deﬁned on predictions from the dynamics, which may not correspond to valid realimages.
To address this, we train a classiﬁer f G : Z → { 0 , 1 } to map the encoding of a state towhether the state is contained in G using terminal states in D success (which are known to be in G ) andother states in D .
However, in the temporally-extended, sparse reward tasks we consider, rewardprediction alone is insufﬁcient because rewards only indicate whether the agent is in the goal set, andthus provide no signal on task progress unless the agent can plan to the goal set.
To address this, asin prior MPC-literature [17, 16, 14, 8], we train a recursively-deﬁned value function (details in thesupplement).
Similar to the reward function, we use the encoder (Section 4.1) to train a classiﬁerf C : Z → [0 , 1] with data of constraint violating states from D constraint and the constraint satisfyingstates in D \ D constraint to map the encoding of a state to the probability of constraint violation.
Model-Based Planning with LS
LS 3 aims to maximize total rewards attained in the environment while limiting constraint violationprobability within some threshold δ C (equation 1).
We optimize an approximation of this objectiveover an H -step receding horizon with model-predictive control.
Precisely, LS 3 solves the followingoptimization 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, probabilisticdynamics 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 andthen 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 } issampled from the learned dynamics f dyn (equation 4).
LS 3 constrains exploration using two chanceconstraints: (1) the terminal latent state in the plan must fall in the safe set (equation 5) and (2) alllatent states in the trajectory must satisfy user-speciﬁed state-space constraints (equation 6).
Z Cis the set of all latent states such that the corresponding observation is constraint violating.
Theoptimizer estimates constraint satisfaction probabilities for a candidate action sequence by simulatingit repeatedly over f dyn .
The ﬁrst chance constraint ensures the agent maintains the ability to return tosafe states where it knows how to do the task within H steps if necessary.
Because the agent replans ateach timestep, the agent need not return to the safe set: during training, the safe set expands, enablingfurther exploration.
In practice, we set δ S for the safe set classiﬁer f S adaptively as described inthe supplement.
The second chance constraint encourages constraint violation probability of nomore 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 theda Vinci Research Kit (dVRK) [45].
Safe RL is of particular interest for surgical robots such asthe 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 thanalgorithms 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 demonstrationsbefore exploring online ( SACfD ) [47] or leverage ofﬂine reinforcement learning to learn a policyusing all ofﬂine data before updating the policy online ( AWAC ) [48].
For both of these comparisonswe enforce constraints via a tuned reward penalty of λ for constraint violations as in [49].
We alsoimplement 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.
Wethen 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 andstable learning.
Finally, we compare LS 3 to a 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) ).
Seethe 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 simulationexperiments, 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 howefﬁciently LS 3 and the comparisons learn each task.
Because all tasks use the sparse task completionbased rewards deﬁned in Section 3, the total reward for a trajectory is the time to reach the goalset, where more negative rewards correspond to slower convergence to G .
Thus, for a task with taskhorizon T , a total reward greater than − T implies successful task completion.
The state is frozenin place upon constraint violation until the task horizon elapses.
We also report task success andconstraint satisfaction rates for LS 3 and comparisons during learning to study (1) the degree to whichtask completion inﬂuences sample efﬁciency and (2) how safely different algorithms explore.
LS 3collects K = 10 trajectories in between training phases on simulated tasks and K = 5 in betweentraining phases for physical tasks, while the SACfD and AWAC comparisons update their parametersafter 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 illustratedin 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 alarge 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 reachingtask (Reacher) based on [50], where the objective is to navigate the end effector of a 2-link planarrobotic arm to a yellow goal position without the end-effector entering a red stay out zone.
We thenstudy a challenging sequential image-based robotic pushing domain (Sequential Pushing), in whichthe objective is to push each of the 3 blocks forward on the table without pushing them to eitherside and causing them to fall out of the workspace.
Finally, we evaluate LS 3 with an image-basedphysical experiment on the da Vinci Research Kit (dVRK) [51] (Figure 3), where the objective is toguide the endpoint of a cable to a goal region without letting the cable or end effector collide with anobstacle.
The Pointmass Navigation and Reaching domains have a task horizon of T = 100 whilethe Sequential Pushing domain and physical experiment have task horizons of T = 150 and T = 50respectively.
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 simulateddomains 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 notlearn 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 ofall algorithms during training.
We ﬁnd that LS 3 achieves a signiﬁcantly higher task success rate thancomparisons 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 constraintviolating states during training due to their lower task success rates.
Because they do not efﬁcientlylearn to perform the tasks, they do not violate constraints as often.
We ﬁnd that the AWAC comparisonachieves very low task performance.
While AWAC is designed for ofﬂine reinforcement learning, tothe best of our knowledge, it has not been previously evaluated on long-horizon, image-based tasksas 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)) asexpected.
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 objectivefrom 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 worsethan LS 3 on the more challenging Pushing environment as the learned safe set is unable to exploittemporal structure to distinguish safe states from unsafe states.
See the supplement for details onexperimental parameters and ofﬂine data used for LS 3 and comparisons and ablations studying theeffect 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 cablerouting task illustrated in Figure 3.
We ﬁnd LS 3 quickly outperforms the suboptimal demonstrationswhile succeeding at the task signiﬁcantly more often than both comparisons, which are unable to learnthe task and also violate constraints more than LS 3 .
We hypothesize that the difﬁculty of reasoningabout cable collisions and deformation from images makes it challenging for prior algorithms tomake sufﬁcient task progress as they do not use prior successes to structure exploration.
See thesupplement 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 ofstates from which the agent is conﬁdent in task completion.
LS 3 then ensures that the agent can planback to states in the safe set, encouraging consistent task completion during learning.
Experimentssuggest that LS 3 can safely and efﬁciently learn 4 visuomotor control tasks, including a challengingsequential pushing task in simulation and a cable routing task on a physical robot.
In future work, weare excited to explore further physical evaluation of LS 3 on safety critical visuomotor control tasksand applications to systems with dynamic constraints on velocity or acceleration.  Appendix In Appendices 7.1 and 7.2 we discuss algorithmic details and implementation/hyperparameter detailsrespectively for LS 3 and all comparisons.
We then provide full details regarding each of the experimental domains and how data is collected in these domains in Appendix 7.3.
In Appendix 7.4, weprovide results from an additional experiment studying the task success rate of LS 3 and comparisonsevolves as training progresses.
Finally, in Appendix 7.5 we perform sensitivity experiments andablations.
7.1 Algorithm Details
In this section, we provide implementation details and additional background information for LS 3and comparison algorithms.
7.1.1 Latent Space Safe Sets (LS 3 )
We now discuss additional details for each of the components of LS 3 , including network architectures,training data, and loss functions.
Variational Autoencoders: We scale all image inputs to a size of (64 , 64 , 3) before feeding themto the β -VAE, which uses a convolutional neural network for f enc and a transpose convolutionalneural network for f dec .
We use the encoder and decoder from Hafner et al.
[2], but modify the secondconvolutional layer in the encoder to have a stride of 3 rather than 2.
As is standard for β -VAEs, wetrain with a mean-squared error loss combined with a KL-divergence loss.
For a particular observations t ∈ S the loss is
[Equation]
where z t ∼ f enc ( z t | s t ) is modeled using the reparameterization trick.
Probabilistic Dynamics: As in Chua et al.
[43] we train a probabilistic ensemble of neural networksto learn dynamics.
Each network has two hidden layers with 128 hidden units.
We train these networkswith a maximum log-likelihood objective, so for two particular latent states z t , z t +1 ∈ Z and thecorresponding action a t ∈
A the loss is as follows for dynamics model f dyn ,θ with parameter θ :
[Equation]
When using the learned dynamics model f dyn for planning, we use the TS-1 method from Chua et al.[43].
Value Functions: As discussed in Section 4.3, we train an ensemble of recursively deﬁned valuefunctions to predict long term reward.
We represent these functions using fully connected neuralnetworks with 3 hidden layers with 256 hidden units.
Similarly to [17], we use separate trainingobjectives during ofﬂine and online training.
During ofﬂine training, we train the function to predictactual discounted cost-to-go on all trajectories in D .
Hence, for a latent vector z t , the loss duringofﬂine training is given as follows where V has parameter θ :
[Equation]
Then in online training we also store a target network V π (cid:48) and use it to calculate a temporal difference(TD-1) error,
[Equation]
(10) where θ (cid:48) are the parameters of a lagged target network and π (cid:48) is the policy at the timestep at which θ (cid:48)was set.
We update the target network every 100 updates.
In each of these equations, γ is a discountfactor (we use γ = 0 . 99 ).
Because all episodes end by hitting a time horizon, we found it wasbeneﬁcial to remove the mask multiplier usually used with TD-1 error losses.
For all simulated experiments we update value functions using only data collected by the suboptimaldemonstrator or collected online, ignoring ofﬂine data collected with random interactions or ofﬂinedemonstrations of constraint violating behavior.
Constraint and Goal Estimators: We represent constraint indicator f C :
Z → { 0 , 1 } with a neuralnetwork with 3 hidden layers and 256 hidden units for each layer with a binary cross entropy loss withtransitions from D constraint for unsafe examples and the constraint satisfying states in D \ D constraint assafe examples.
Similarly, we represent the goal estimator f G : Z → { 0 , 1 } with a neural networkwith 3 hidden layers and 256 hidden units.
This estimator is also trained with a binary cross entropyloss with positive examples from D success and negative examples sampled from all datasets.
Forthe constraint estimator and goal indicator, training data is sampled uniformly from a replay buffercontaining D success , D rand and D constraint .
Safe Set: The safe set classiﬁer f S ( · ) is represented with neural network with 3 hidden layers and256 hidden units.
We train the safe set classiﬁer to predict
[Equation]
using a binary cross entropy loss, where 1 S j ( s t ) is an indicator function indicating whether s t is partof a successful trajectory.
Training data is sampled uniformly from a replay buffer containing all of
D .
Similar to deep value function learning literature [52, 53, 17], the safe set is trained to solve theabove equation by ﬁxed point iteration: the safe set is used to construct its own targets, which arethen used to update the safe set before using the updated safe set to construct new targets.
Cross Entropy Method: We use the cross entropy method to solve the optimization problem inequation 2.
We build on the implementation of the cross entropy method provided in [54], whichworks by sampling n candidate action sequences from a diagonal Gaussian distribution, simulating eachone n particle times over the learned dynamics, and reﬁtting the parameters of the Gaussian on the n elitetrajectories with the highest score under equation 2 where constraints are implemented by assigninglarge negative rewards to trajectories which violate either the safe set constraint or user-speciﬁedconstraints.
This process is repeated for n cem iters to iteratively reﬁne the set of sampled trajectoriesto optimize equation 2.
To improve the optimizer’s efﬁciency on tasks where subsequent actionsare often correlated, we sample a proportion (1 − p random ) of the optimizer’s candidates at the ﬁrstiteration from the distribution it learned when planning the last action.
To avoid local minima, wesample a proportion p random uniformly from the action space.
See Chua et al.
[43] for more details onthe cross entropy method as applied to planning over neural network dynamics models.
As mentioned in Section 4.4, we set δ S for the safe set classiﬁer f S adaptively by checking whetherthere exists at least one plan which satisﬁes the safe set constraint at each CEM iteration.
If no suchplan exists, we multiply δ S by 0 .
8 and re-initialize the optimizer at the ﬁrst CEM iteration with thenew value of δ S .
This ensures that δ S is set such that planning back to the safe set is possible.
Weinitialize δ S = 0 .
7.1.2 Soft Actor-Critic from Demonstrations (SACfD)
We utilize the implementation of the Soft Actor Critic algorithm from [55] and initialize the actorand critic from demonstrations but keep all other hyperparameters the same as the default in theprovided implementation.
We create a new dataset D demos (cid:40) D using only data from the suboptimaldemonstrator, and use the data from D demos to behavior clone the actor and initialize the critic usingofﬂine bellman backups.
We use the same mean-squared loss function for behavior cloning as for thebehavior clone policy, but only train the mean of the SAC policy.
Precisely, we use the followingloss for some policy π with parameter θ : L ( θ, D demos )
=
(cid:80)
τ i ∈D demos (cid:80)
Tt =
1 || µ θ ( s it ) − a it || 2 where s itand a it are the state and action at timestep t of trajectory
τ i and π ( ·| s t ) ∼ N ( µ θ ( s t ) , σ φ ( s t )) .
Wealso experimented with training the SAC critic on all data provided to LS 3 in D but found that thishurt performance.
We use the architecture from [55] and update neural network weights using an
Adam optimizer with a learning rate of 3 × 10 − 4 .
The only hyperparameter for SACfD that wetuned across environments was the reward penalty λ which was imposed upon constraint violations.
For all simulation experiments, we evaluated λ ∈ {− 1 , − 3 , − 5 , − 10 , − 20 } and report the highestperforming value.
Accordingly, we use λ = − 3 for all experiments except the reacher task, for whichwe used λ = − 1 .
We observed that higher values of λ resulted in worse task performance withoutsigniﬁcant increase in constraint satisfaction.
We hypothesize that since the agent is frozen in theenvironment upon constraint violations, the resulting loss of rewards from this is sufﬁcient to enable
SACfD to avoid constraint violations.
[Table]
7.1.3 Soft Actor-Critic from Demonstrations with Learned Recovery Zones (SACfD+RRL)
We build on the implementation of the Recovery RL algorithm [36] provided in [56].
We train thesafety critic on all ofﬂine data from D .
Recovery RL uses SACfD as its task policy optimizationalgorithm, and introduces two new hyperparameters: ( γ risk , (cid:15) risk ) .
For each of the simulationenvironments, we evaluated SACfD+RRL across 3-4 ( γ risk , (cid:15) risk ) settings and reported resultsfrom the highest performing run.
Accordingly, for the navigation environment, we use: ( γ risk =0 .
95 , (cid:15) risk = 0 .
For the reacher environment, we use ( γ risk = 0 .
55 , (cid:15) risk = 0 . 7) , and we use( γ risk = 0 .
75 , (cid:15) risk = 0 .
for the sequential pushing environment.
For the cable routing environment,we use ( γ risk = 0 .
55 , (cid:15) risk = 0 . 7) .
4 Advantage Weighted Actor-Critic (AWAC)
To provide a comparison to state of the art ofﬂine reinforcement learning algorithms, we evaluate
AWAC [48] on the experimental domains in this work.
We use the implementation of AWACfrom [57].
For all simulation experiments, we evaluated λ ∈ {− 1 , − 3 , − 5 , − 10 , − 20 } and reportthe highest performing value.
Accordingly, we use λ = − 1 for all experiments.
We used the defaultsettings from [57] for all other hyperparameters.
7.2 LS 3 Implementation Details
In Table 3, we present the hyperparameters used to train and run LS 3 .
We present details for theconstraint thresholds δ C and δ S .
We also present the planning horizon H and VAE KL regularizationweight β .
We present the number of particles sampled over the probabilistic latent dynamicsmodel for a ﬁxed action sequence n particles , which is used to provide an estimated probability ofconstraint satisfaction and expected rewards.
For the cross entropy method, we sample n candidateaction sequences at each iteration, take the best n elite action sequences and then reﬁt the samplingdistribution.
This process iterates n cem iters times.
We also report the latent space dimension d ,whether frame stacking is used as input, training batch size, and discount factor γ .
Finally, we presentvalues of the safe set bellman coefﬁcient γ S .
For all domains, we scale RGB observations to a size of(64 , 64 , 3) .
For all modules we use the Adam optimizer with a learning rate of 1 × 10 − 4 , except fordynamics which use a learning rate of 1 × 10 − 3 .
7.3 Experimental Domain Details
7.3.1 Navigation
The visual navigation domain has 2-D single integrator dynamics with additive Gaussian noise sampled from N (0 , σ 2 I 2 )
where σ = 0 .
The start position is (30 , 75) and goal set is B 2 ((150 , 75) , 3) ,where B 2 ( c, r ) is a Euclidean ball centered at c with radius r .
The demonstrations are created byguiding the agent north for 20 timesteps, east for 40 timesteps, and directly towards the goal until theepisode terminates.
This tuned controller ensures that demonstrations avoid the obstacle and alsoreach the goal set, but they are very suboptimal.
To collect demonstrations of constraint violating behavior, we randomly sample starting points throughout the environment, move in a random direction for 15 time steps, and then move directly towards the obstacle.
We do not collect additional data for
D rand in this environment.
We collect 50 demonstrations of successful behaviors and 50 trajectoriescontaining constraint violating behaviors.
7.3.2 Reacher
The reacher domain is built on the reacher domain provided in the DeepMind Control Suite from [50].
The robot is represented with a planar 2-link arm and the agent supplies torques to each of the 2 joints.
Because velocity is not observable from a single frame, algorithms are provided with several stackedframes as input.
The start position of the end-effector is ﬁxed and the objective is to navigate the endeffector to a ﬁxed goal set on the top left of the workspace without allowing the end effector to entera large red stay-out zone.
To collect data from D constraint we randomly sample starting states in theenvironment, and then use a PID controller to move towards the constraint.
To sample random datathat will require the agent to model velocity for accurate prediction, we start trajectories at randomplaces in the environment, and then sample each action from a normal distribution centered aroundthe previous action, a t +1 ∼ N ( a t , σ 2 I ) for σ 2 = 0 .
We collect 50 demonstrations of successfulbehavior, 50 trajectories containing constraint violations and 100 short (length 20 ) trajectories orrandom data.
7.3.3 Sequential Pushing
This sequential pushing environment is implemented in MuJoCo [58], and the robot can specify adesired planar displacement a = (∆ x, ∆ y ) for the end effector position.
The goal is to push all 3blocks backwards by at least some displacement on the table, but constraints are violated if blocksare pushed backwards off of the table.
For the sequential pushing environment, demonstrations arecreated by guiding the end effector to the center of each block and then moving the end effector ina straight line at a low velocity until the block is in the goal set.
This same process is repeated foreach of the 3 blocks.
Data of constraint violations and random transitions for D constraint and D rand arecollected by randomly switching between a policy that moves towards the blocks and a policy thatrandomly samples from the action space.
We collect 500 demonstrations of successful behavior and300 trajectories of random and/or constraint violating behavior.
7.3.4 Physical Cable Routing
This task starts with the robot grasping one endpoint of the red cable, and it can make (∆ x, ∆ y )motions with its end effector.
The goal is to guide the red cable to intersect with the green goalset while avoiding the blue obstacle.
The ground-truth goal and obstacle checks are performedwith color masking.
LS 3 and all baselines are provided with a segmentation mask of the cable asinput.
The demonstrator generates trajectories D success by moving the end effector well over theobstacle and to the right before executing a straight line trajectory to the goal set.
This ensuresthat it avoids the obstacle as there is signiﬁcant margin to the obstacle, but the demonstrations maynot be optimal trajectories for the task.
Random trajectories D rand are collected by following ademonstrator trajectory for some random amount of time and then sampling from the action spaceuntil the episode hits the time horizon.
We collect 420 demonstrations of successful behavior and150 random trajectories.
We use data augmentation to increase the size of the dataset used to trainf enc and f dec , taking the images in D and creating an expanded dataset by adding randomly sampledafﬁne translations and perspective shifts, until |D VAE | > 100000 .
7.4 Additional Results
We additionally study how the task success rate of LS 3 and comparisons evolves as training progressesin Figure 6.
Precisely, we checkpoint each policy after each training trajectory and evaluate it over 10rollouts for each of the 10 random seeds (100 total trials per datapoint).
We ﬁnd that LS 3 achievesa much higher task success rate than comparisons early on in training, and maintains a higher tasksuccess rate throughout the course of training on all simulation domains.
7.5 Sensitivity Experiments
Key hyperparameters in LS 3 are the constraint threshold δ C and safe set threshold δ S , which controlwhether the agent decides predicted states are constraint violating or in the safe set respectively.
Weablate these parameters for the Sequential Pushing environment in Figures 7 and 9.
We ﬁnd thatlower values of δ C made the agent less likely to violate constraints as expected.
Additionally, we ﬁndthat higher values of δ S helped constrain exploration more effectively, but too high of a threshold led to poor performance suffered as the agent exploited local maxima in the safe set estimation.
Finally,we ablate the planning horizon H for LS 3 and ﬁnd that when H is too high, Latent Space Safe Sets(LS 3 )can explore too aggressively away from the safe set, leading to poor performance.
When His lower, LS 3 explores much more stably, but if it is too low (ie.
H = 1 ), LS 3 is eventually unableto explore signiﬁcantly new plans, while slightly increasing H (ie.
H = 3 ) allows for continuousimprovement in performance.
[Figure]
[Figure]
No Constraints =0.
=0.
[Figure]
=
[Figure]
[Figure]
=
[Figure]
[Figure]
=
[Figure]
[Figure]
[Figure]
Demonstrations
H =
H =
H =
H =
H =