BoxGPT: Embodied Data Collection

Mathematical formulation

From model uncertainty to robot actions

BoxGPT is a simplification of a general problem in embodied data collection. We are learning a generative model, denoted by \(p(y\mid x;\theta)\). Here, \(x\) is the state of the robot, \(y\) is the variable of interest and the variable of measurement, and \(\theta\) is the parameter of the model. For example, \(y\) could represent images in a vision model.

Learning from data

The robot moves across the space to collect data. Each data point is denoted by \(d_i=(x_i,y_i)\), where \(i\) is the index of the data point. The dataset is denoted by \(D=\{d_i\}_{i=1}^{N}\).

The learning process, or model regression, infers a distribution of the model parameters from the dataset. We denote this distribution by \(q(\theta\mid D)\). Maximum likelihood estimation can be considered a special case in which \(q(\theta\mid D)\) is a Dirac delta function.

Given an inferred distribution of model parameters, the posterior predictive distribution at a particular robot state \(x\) is

\[ \bar p(y\mid x,D) = \mathbb{E}_{\theta\sim q(\theta\mid D)}\!\left[p(y\mid x;\theta)\right]. \]

We quantify the predictive uncertainty at \(x\) using an information function \(h(x)\):

\[ h(x)=\mathbb{H}_{y}\!\left[\bar p(y\mid x,D)\right], \]

where \(\mathbb{E}\) denotes expectation and \(\mathbb{H}\) denotes entropy. The information function quantifies how uncertain the model is at a particular robot state. The expected information that a measurement at that state provides about the model parameters is given by the conditional mutual information

\[ \mathbb{I}(y;\theta\mid x,D) = h(x)-\mathbb{E}_{\theta\sim q(\theta\mid D)} \!\left[\mathbb{H}_{y}\!\left[p(y\mid x;\theta)\right]\right]. \]

In BoxGPT, every rectangle predicts the binary measurement deterministically. The second term is therefore zero, making predictive entropy and expected information gain equivalent for this problem.

Data collection as a control problem

The problem, however, is how to leverage the information function to generate robot actions in order to acquire new data. A robot that collects data faces unique limitations. The data is inherently collected sequentially, and the robot's motion is constrained by its dynamics. BoxGPT is designed to highlight and examine these challenges. It serves as a testbed for prototyping data-collection policies for embodied agents.

BoxGPT simplifies the general problem formulation to make model regression tractable. In fact, model regression is analytical, leaving the challenge entirely to the data-collection policy rather than the capability of the model. Modern AI paradigms tend to overemphasize the importance of model architecture and overlook the impact of data quality. More importantly, they often overlook the impact of the data-collection process itself on learning performance. BoxGPT is designed to expose the importance of decision-making in data collection, especially for embodied agents.

The BoxGPT model

More specifically, in BoxGPT, the measurement \(y\in\{\text{Positive},\text{Negative}\}\) is binary. The predictive model is represented as a Bernoulli distribution:

\[ p(y\mid x;\theta)= \begin{cases} g(x;\theta), & \text{if }y\text{ is Positive},\\ 1-g(x;\theta), & \text{if }y\text{ is Negative}. \end{cases} \]

The parameter \(\theta\) fully characterizes a rectangle through its width, height, and center coordinates. The function \(g(x;\theta)\) indicates whether the robot state falls within the rectangle characterized by \(\theta\):

\[ g(x;\theta)= \begin{cases} 1, & x\in\operatorname{Rectangle}(\theta),\\ 0, & \text{otherwise}. \end{cases} \]

Analytical model regression

Following this simplification, model regression becomes analytical. Given a dataset \(D=\{(x_i,y_i)\}_{i=1}^{N}\), define \(\Theta(D)\) as the set of parameters that characterize rectangles containing all positive signals and none of the negative signals. The inferred distribution \(q(\theta\mid D)\) is uniform over this consistent set:

\[ q(\theta\mid D)\propto \mathbf{1}\!\left[\theta\in\Theta(D)\right]. \]

In practice, BoxGPT represents this distribution using candidate rectangles sampled from the consistent set. These candidate rectangles are shown in red in the interactive demonstration.

An intuitive uncertainty distribution

Following this simplification, the information function also has an intuitive explanation. The probability of a positive measurement at \(x\) is the fraction of inferred rectangles that include \(x\):

\[ p_{+}(x) =p(y=\text{Positive}\mid x,D) =\mathbb{P}_{\theta\sim q(\theta\mid D)} \!\left[x\in\operatorname{Rectangle}(\theta)\right]. \]

The posterior predictive distribution is therefore Bernoulli with positive probability \(p_{+}(x)\). The information function is its binary entropy:

\[ h(x)=-p_{+}(x)\log p_{+}(x) -\bigl(1-p_{+}(x)\bigr)\log\bigl(1-p_{+}(x)\bigr). \]

A state \(x\) has the highest uncertainty when exactly half of the inferred rectangles include it while the other half do not.

BoxGPT

Collect data and update the model

Code and notebooks

Random walk