On-Robot Bayesian Reinforcement Learning for POMDPs

Bayesian RL involves maintaining a probability distribution over an augmented state, including both the environment dynamics and the environment state, as proposed by the BA-POMDP framework [7]. State-of-the-art BRL methods [9, 10, 6] for partially observable domains have improved on the BA-POMDP by incorporating the Partially Observable Monte Carlo Planning (POMCP) [11], an efficient online planner for large POMDPs. While earlier methods [9, 10] are only applicable to small discrete POMDPs or ones with factorizable dynamics, Bayes-adaptive deep dropout RL (BADDr) [6] was introduced to tackle more complex domains. BADDr leverages the representation power of dropout networks [12] to capture priors, making BRL scalable to larger domains while being considerably more sample-efficient than pure RL methods. In this paper, we also use dropout networks to construct (parts of) a specialized dynamics model, leveraging assumptions about the factored dynamics and full observability of the robot state.

Recently, many methods have performed learning directly on physical hardware. In [13], multiple robot workers are used to collect data for online training of a robot arm to open a door autonomously. [14] used human feedback to learn manipulation tasks in 1-4 hours. Meanwhile, [15] combined RL and contrastive learning [16] to learn manipulation tasks online from pixels. In addition, [17, 18] used symmetry-aware neural networks to encode domain symmetries to perform online learning manipulation tasks within several hours, while [19, 20] conducted online learning for locomotion tasks on legged robots. Nevertheless, unlike ours, none of these works addresses partially observable domains or employs a Bayesian approach directly to a physical system.

A POMDP [21] is formally defined by the tuple $(\mathcal{S},\mathcal{A},\Omega,T,O,R,p_{s_{0}})$, where $\mathcal{S},\mathcal{A},\Omega$ are respectively the state, action, and observation spaces; $T:\mathcal{S}\times\mathcal{A}\rightarrow\Delta(\mathcal{S}),O:\mathcal{S}\times\mathcal{A}\rightarrow\Delta(\Omega)$, and $R:\mathcal{S}\times\mathcal{A}\rightarrow\mathbb{R}$ are respectively the transition, observation, and reward functions; $p_{s_{0}}=\Delta(\mathcal{S})$ is the prior about the initial state $s_{0}$. The goal is to find a sequence of actions to maximize the discounted return defined as $\mathbb{E}[\sum_{t=0}^{\infty}\gamma^{t}r_{t}]$ [22] with some discount factor $\gamma\in[0,1]$. An action $a$ taken in state $s$ will result in a state $s^{\prime}$, sampled from the transition function $T(s,a,s^{\prime})=p(s^{\prime}\mid s,a)$. In a POMDP, the agent can only observe $o\in\Omega$, indirectly related to $s^{\prime}$ via the observation function $O(s^{\prime},a,o)=p(o\mid s^{\prime},a)$. This implies that the agent might need to use the entire action-observation history, which grows with the episode length, to act optimally. Alternatively, acting optimally can rely on a belief (posterior probability distribution) $b\in\Delta(\mathcal{S})$ over possible states, where $b(s)$ denotes the probability that the environment’s true state is $s$. Given a new action $a$ and observation $o$, a new belief $b^{\prime}$ can be calculated as:

Equation 1 indicates that explicitly belief tracking (and therefore planning) must rely on access to the known dynamics $T(s,a,s^{\prime})$ of the system. When these are unknown, we must turn to learning-based approaches instead. This work adopts the Bayesian perspective.

BRL assumes that, instead of access to the system’s dynamics $\mathcal{D}$, we have a parameterized prior probability distribution $p(\mathcal{D};\theta)$. On a high level, our agent will maintain a belief $\bar{b}\in\Delta(S,\Theta)$ over the dynamics and the state and pick actions that optimize future return with respect to this belief. Technically, this Bayesian framework constructs a “belief” POMDP: an augmented POMDP of which the state space (and dynamics) include parameters that describe the dynamics of the underlying system. Even if the dynamics of the original system are unknown, it is now again possible to plan and track beliefs in the larger POMDP instead.

Since the GBA-POMDP is a POMDP, it can be solved through belief tracking and planning. Unfortunately, the state space is very large, and we need to reach for approximation techniques instead, which will be covered next.

The previous section described the GBA-POMDP as a recipe that constructs a belief POMDP from a prior $\theta_{0}$ and parameter update $\mathcal{U}$. BADDr [6] is a realization of GBA-POMDP that uses dropout neural networks [12] to represent the dynamics. In particular, it assumes that the prior over the dynamics is captured by a parameter set $w\in\mathcal{W}$ (e.g., $w$ can be neural network weights) and chooses the update function to be stochastic gradient descent with dropout (SGD) [23], which can be interpreted as an approximation of Bayesian inference over neural network parameters [24]. Consequently, BADDr is a POMDP with augmented state space $\bar{\mathcal{S}}=\mathcal{S}\times\mathcal{W}$ and dynamics:

where $\mathcal{L}((s,a),(s^{\prime},o);w)=-\log p(s^{\prime},o\mid s,a;w)$ is the cross-entropy loss between the predicted and true next state and observation.

Typical assumptions include:

• •

  There is a high-fidelity simulator of the state of the robot

• •

  The (internal) state of the robot is fully observable

• •

  Some high-level system behavior is unknown

• •

  The model of the sensors is somewhat known

Expert knowledge in robotic tasks typically includes a physical simulator of the robot in which we can simulate actions. Similarly, the internal of the robot state is observable in most systems. What is likely unknown, and over which we may only have a distribution, is some high-level element of the dynamics of the environment in which the robot will be deployed. In this work, this is the human behavior in a human-robot setting. Other assumptions may apply for different applications, but we hope that most classes of assumptions have been covered here and that it will be relatively straightforward to adapt to future requirements.

We observe that, without loss of generalization, the state can be divided into internal (robot $s_{r}$) and external (here “human’s” $s_{h}$) parts: $s=(s_{r},s_{h})$ (see fig. 1). Then we denote $T_{r}(s_{r},a,s^{\prime}_{r})$ as the accurate but low-level simulator in which we can accurately predict the agent’s internal state by unrolling a controller. On the other hand, the dynamics of the environment $T_{h}(s,a,s^{\prime}_{h})$ are unknown before deployment.

We make a similar distinction in the observation and divide it into an internal (robot) $o_{r}$ and external (human) $o_{h}$ part $o=(o_{r},o_{h})$. We assume that the internal state $o_{r}$ (can be computed from the state through a given function $O_{r}$: $s^{\prime}\rightarrow o_{r}$) is fully observable (note that a similar assumption is formalized in the mixed observability MDP [8]). The observation function over the variable $o_{h}$, $O_{h}(s^{\prime},a,o_{h})$, may or may not be given depending on our understanding of the robot’s sensors.

The GBA-POMDP does not provide the tools to capture the expert knowledge we discussed above or how to do further inference over. Instead, we derive a more sophisticated framework starting from the assumptions stated above.

Here we leverage the shape of the posterior derived above to define a Bayesian framework for robotics. In particular, we use the factorization of the posterior computation to construct a model with a factorization of the parameter update rule $\mathcal{U}$ — one for each factor.

The transition posterior computation, for example, is defined by an update rule for the distribution over the physics simulator and one for the distribution over the human dynamics. The update rule that corresponds to the Kronecker-delta function is simply its identity $\mathcal{U}(T_{r},s,a,s^{\prime},o)=T_{r}$. The posterior parameters of the human dynamics depend on the choice of the prior model $p(T_{h})$, which in this case will be neural networks $w_{hT}$ with the corresponding dropout SGD update rule $\mathcal{U}(w_{hT},s,a,s^{\prime},o)$ as in eq. 4. For example, eq. 5 can be (re-)written with a factored parameter update rule:

and we apply the same approach to the observation function.

The prior of known components (e.g., $T_{r}$ and $O_{r}$) are, as discussed before, encoded by Kronecker-delta distributions. Priors over other parts of the dynamics, such as the human’s dynamics $T_{h}$ and observations $O_{h}$, are a design choice, and here are assumed to be represented by some network parameters $(w_{hT},w_{hO})$.

Starting with the problem definition of a general robotic system, we have developed a belief POMDP that captures typical expert prior knowledge and allows for Bayesian inference over unknown parts of the dynamics. To solve the resulting problem, however, we require efficient action selection and belief approximation. In particular, the belief space is far too large to either apply naive planning or exact inference. Instead, we propose to combine: a specialized POMCP algorithm, a particle filtering approximation technique, and a targeted prior construction.

First, we train networks on the priors we have over the unknown dynamic terms (e.g., $\{w_{hT}\}\approx p(T_{h})$). Then we initialize our belief $b_{0}$ by sampling particles $(s,w_{hT},w_{hO})$ from our joint prior ($s\sim p_{s_{0}}$, and $(w_{hT},w_{hO})\sim\{w_{hT},w_{hO}\}$). At each time step, our planner builds a look-ahead tree, as POMCP does, but where each simulation starts with sampling an augmented state $\sim b$ and using dynamics in eq. 7 to generate trajectories. For belief tracking, we use a new type of rejection sampling. This algorithm samples particles $(s,w_{hT},w_{hO})$, proposes the next augmented states (again with dynamics in eq. 7) and accepts those that generated the perceived observation.

This approach, first, can do approximate belief tracking through particle filtering in our specialized Bayes model and, second, pick actions efficiently with respect to our most up-to-date understanding of the environment. The result is a quick and feasible online algorithm that manages to both be sample efficient and capable of exploiting expert knowledge directly into the initial belief in a principled manner.

A Turtlebot must deliver $T$ tools to a human in an unknown order to complete an assembly task (fig. 2a). Turtlebot can carry multiple tools in its basket at once, and the tools are stored in a tool room while the human works in a work room. Each time the human worker receives a correct tool, he needs a fixed number of timesteps to use it before needing another. When a correct tool is delivered, the human’s working step (located in the top right corner) increases by one.

The navigation and tool pickup transitions are known a-priori $(s_{r}{=}x_{coord}\sim T_{r}(s_{r},a)$, while the human behavior $(s_{h}{=}(t_{carry},t_{need})\sim T_{h}(s,a)$ is not.

In this domain, the Turtlebot must deliver tools to two human workers (Human-0 and Human-1) in separate rooms (fig. 2b). Importantly, given the same tool, one of the workers works faster than the other. The table now holds two identical sets of tools, each containing three tools, and each human requires one set to complete the task. The tool order is the same for both workers and assumed to be known, i.e., $\text{Tool }0\rightarrow 1\rightarrow 2$.

In OTD, the transition prior assumes a uniform distribution over $T!$ possible delivery orders for $T$ tools. Assuming a correct observation model, we represent these priors using $T!$ parameter sets $(w_{hT},w_{hO})$. To generate training data for training a parameter set, we pair a random tool order with the correct observation model to sample the dynamics. We then use a random policy to interact with the resulting dynamics to generate the data.

In STD, we consider a set $S$ consisting of $|S|$ possible working speeds. Because the speeds are different (one human worker works faster than the other worker), $|S|^{2}-|S|$ possible different speed combinations are possible. Over these combinations, we assume a uniform transition prior, which is defined using $|S|^{2}-|S|$ parameter sets $(w_{hT},w_{hO})$. Similar to OTD, to generate training data for training a parameter set, we pair a random speed combination with a correct observation model and the correct tool order to create the dynamics and then use a random policy to interact.

In OTD, our agent nearly reaches the performance of POMCP running on the true POMDP after three episodes. During the first episode of OTD, the agent relies solely on the initial prior, as the dynamics parameters have not yet been updated. The observed behavior of the agent is to gather all three tools and then deliver them. While this strategy yields a reasonable outcome, it is sub-optimal as the worker only requires a tool once they have finished using the current one. In the second episode, the agent performs better by only taking two tools in two consecutive actions, delivering them, and then returning later for the third tool. In the third episode, it learns to retrieve one tool at a time and deliver it until all tools have been supplied, which is optimal. In order to see the robot in action, please see our video.

STD seems more challenging as the agent needs more episodes to perform well. Specifically, fig. 4(b) indicates that our agent can reach the performance of POMCP within ten episodes by delivering all tools in the correct order for both tested working speed combinations. The final policy involves taking two tools of the same type (e.g., tool 0) and delivering them one by one, starting with the faster worker and then moving on to the slower worker. This strategy allows the agent to leverage the close distance between the two work rooms. Please see our video for the learned policy.

To ablate our agent, we perform experiments in simulation with the OTD task with three tools.