From Pixels to UI Actions: Learning to Follow Instructions via Graphical User Interfaces

We consider an environment with states $\mathcal{S}$ and actions $\mathcal{A}$. The reward function, $r(s)$, returns a scalar corresponding to the reward given for transitioning to state $s\in\mathcal{S}$, and is described below. MiniWob++ environments are randomly generated, but transitions are deterministic within an environment generated by a particular random seed. The transition function, $f(s,a)$, returns the state resulting from taking action $a\in\mathcal{A}$ in state $s\in\mathcal{S}$.

Rather than using the raw reward directly provided by the MiniWob++ environment, we consider a surrogate reward: $r(s)=\alpha_{s}+r^{t}(s)$, where $\alpha_{s}$ provides a small negative reward that encourages shorter trajectories without unnecessary actions. $r^{t}(s)$ is the raw reward from the MiniWob++ environment if $s$ is a terminal state and the raw reward is $>0.8$, or $0$ otherwise. We use $\alpha_{S}=-\frac{1}{30}$. As all tasks can be completed within 30 steps, this is small enough to ensure a positive reward is possible for all tasks. Additionally, the penalty is small enough such that in practice the agent should not be incentivized to sacrifice raw reward to reduce the number of steps taken.

The value function $v^{\pi}(s)$ for a given policy $\pi$ is the expected future rewards from state $s$ if actions are selected according to policy $\pi$. The optimal value function, $v^{*}(s)$, is the expected future rewards if optimal actions are chosen. We attempt to learn an approximation of this function, $\hat{v}_{\phi}(s)\approx v^{*}(s)$, parameterized as a Pix2Struct-initialized model with parameters $\phi$, which we refer to as the value network. The model is trained on transitions from the human demonstrations, which demonstrate close to optimal behavior in many cases. For every state in the human demonstrations, we compute the actual future rewards for the given episode, according to the surrogate reward. We map these future rewards to discrete bins and represent them as integers in the Pix2Struct decoder. At inference time, we approximate the mean of the distribution over these discrete bins by considering the top-$n$ predictions from the model using beam search (with $n=3$), weighted proportional to their respective probabilities.

For consistency with prior work, we will refer to the Pix2Struct model tuned to generate actions (i.e. Pix2Act) as the policy network, with parameters $\theta$. The greedy policy $\pi_{\theta}(s)$ selects the action $a$ with the highest approximate probability $p_{\theta}(a|s)$ in the top-$k$ beam (see §B.1), subject to the conditions described in $\lx@sectionsign\ref{sec:agent}$.

We can use lookahead search to implement a policy, $\pi^{*}_{\theta}(s)$, which leverages interactions with the environment ($f(s,a)$ and $r(s)$) to select actions in a more optimal way than the greedy policy $\pi_{\theta}(s)$. Both the policy network and value network are used to constrain and prioritize the search.

MCTS performs $K$ rounds of traversing a search tree with nodes corresponding to states, and edges corresponding to actions. Due to the computational cost of the policy and value networks, we use a modest number of rounds, $K=16$, for our experiments. The search tree is initialized with a single root node for state $s$. Each round starts at $s$ and traverses the tree. At each step $t$ of a given round, an action $a_{t}$ is selected for state $s_{t}$, where $a_{t}=\max_{a}Q(s_{t},a)+U(s_{t},a)$. $Q(s_{t},a)$ is an average reward over all rounds that have traversed the associated edge. It is based on actual accumulated rewards during tree traversal and the value estimates of leaf states (described below). $U(s_{t},a)=c*p_{\theta}(a|s)*\frac{\sqrt{N(s_{t})}}{1+n(s_{t},a)}$ is a term that encourages exploration, where $n(s_{t},a)$ is the number of times action $a$ has been selected from state $s_{t}$, $N(s_{t})$ is the total number of times state $s_{t}$ has been visited, and $c$ is a scalar hyperparameter that we set to $0.1$. Following  Silver et al. [2017], we use the policy network to bias this exploration term. To constrain the search, we only consider the top-$k$ actions according to the policy network, where $k=8$ in our experiments.

If we select an action $a_{t}$ for state $s_{t}$ which has never been previously selected from $s_{t}$, then the simulation ends and we add a new leaf state, $s_{L}=f(s_{t},a)$, to the search tree. If $s_{L}$ is not a terminal state, then we estimate its value (i.e. future returns) using both the value network and a rollout with the greedy policy. Specifically, following Silver et al. [2017], we estimate its value as $\lambda*\hat{v}_{\phi}(s_{L})+(1-\lambda)*v^{\pi_{\theta}}(s_{L})$ where $v^{\pi_{\theta}}(s_{L})$ is equal to the actual returns from following the policy $\pi_{\theta}$ starting at $s_{L}$ for a maximum of $20$ steps, with actual returns clipped to a minimum value of $0$. Is there $\lambda$ is a mixing parameter that we set to $0.1$. For challenging environments, rollouts may be unlikely to find a terminal state with positive reward, and in such cases rollouts may not be very informative. On the other hand, the value network can provide poor value estimates for certain states, especially if they are not well represented in the human demonstrations. By combining both methods we aim to provide a better approximation of the value of leaf states. Returns are propagated up the tree to each parent $s^{\prime}$ to update $Q(s^{\prime},a)$. As $Q(s_{L},a)$ is undefined prior to selecting $a$ from $s_{L}$ for the first time, we initialize $Q(s_{L},a)$ for each action to be equal to the initial value estimate of $s_{L}$ plus $\alpha_{s}$.

To understand the impact of rollouts and value estimates using the value network, in Table 3 we compare mean scores over 12 challenging MiniWob++ tasks for different values of $\lambda$: 0 (rollout only), 0.1 (both rollout and value network), and 1 (value network only). We also include the mean score using the greedy policy for reference. These results use the policy network and value network trained on the human demonstrations. The results show that using a combination of rollouts and the value network gives the best results. The value network is primarily useful for challenging tasks that require longer trajectories, such as number-checkboxes, relative to using rollouts only.

Once we have completed $K$ rounds, $\pi^{*}_{\theta}(s)$ selects the most visited action $a$ for state $s$, and we begin the process again at the subsequent state. We reuse the search tree for subsequent time steps for efficiency, so we require only $K-n(s,a)$ additional rounds for the subsequent state.

We can sample trajectories with $\pi^{*}_{\theta}$, then update $\theta$ by training $\pi_{\theta}(s)$ to approximate $\pi^{*}_{\theta}(s)$ for each $s$ in the sampled trajectories. This then also improves $\pi^{*}_{\theta}(s)$, as $\theta$ informs how the search space is constrained and prioritized. Therefore, we can continue to iteratively improve $\pi_{\theta}(s)$. To produce these trajectories, we randomly sample MiniWob++ tasks and seeds, and select actions according to $\pi^{*}_{\theta}$. We then filter trajectories where the raw reward is $<0.8$. We then tune $\theta$ on these new trajectories. For simplicity, we keep the value network (i.e. $\phi$) fixed.

We initially found that tuning on trajectories from MCTS could be unstable, leading to an early loss spike. To resolve this, we slightly decreased the learning rate (from $1e-3$ to $5e-4$) and increased the number of warmup steps (from $1000$ to $4000$) relative to the hyperparameters used for behavioral cloning.