To solve this problem, I use “rejection sampling”. In a nutshell, this means we have a way to generate “probably good” boards, and then afterwards we filter through those to see which ones are “actually good”. The “probably good” boards form a proposal distribution, and our decision criteria for whether a board is “actually good” is an acceptance ratio.

More mathematically, we have a proposal distribution \(q(y)\), from which we can easily sample and estimate probability densities. Then, using a target distribution \(p(y)\) from which we cannot easily sample but can easily estimate probability densities, we determine whether \(q(y)\) overrepresents a certain sample; if so, by how much; and then reject with a certain probability accordingly. We first sample from the proposal distribution \(y \sim q(y)\), and then \(u \sim \operatorname{Unif}(0, 1)\). Then, we compute \(r = \frac{p(y)}{cq(y)}\), where \(p(y)\) is the target distribution and \(c\) is a constant such that \(r \leq 1\) for all \(y\). If \(u < r\), we accept \(y\) as a sample from \(p(x)\). We can also write this in terms of log probabilities. We have \(\log r = \log p(y) - \log q(y) - \log c\). If \(\log u < \log r\), we accept \(y\) as a sample from \(p(x)\). In this case, because \(\log p(x) = -\infty\) for all \(x\) that don’t satisfy the constraints, we can simply reject any values that do not satisfy the constraints, and accept the values that do.

Rejection Sampling with Uniform Proposal Distribution

We note that no two queens can be in the same row. Therefore, there must be exactly one queen in each row. So, we can represent the board as a list (e.g. \([2, 0, 3, 1]\) for a \(4 \times 4\) board), where the position in the list represents which row the queen is in, and the value at that position represents the column. We can randomly generate such boards uniformly by taking random permutations of \([0, 1, 2, 3]\). To add rejection sampling, we simply check if the resulting board after the random permutation has any attacking queens, and if not, accept it. The proportion of accepted samples to proposal samples is called the acceptance ratio, which I report here.

import numpy as np


def count_conflicts(board):
    has_conflict = [False] * len(board)
    for queen_index in range(len(board)):
        for other_queen_index in range(queen_index + 1, len(board)):
            if board[queen_index] == board[other_queen_index]:
                has_conflict[queen_index] = True
                has_conflict[other_queen_index] = True

            if abs(board[queen_index] - board[other_queen_index]) == abs(
                queen_index - other_queen_index
            ):
                has_conflict[queen_index] = True
                has_conflict[other_queen_index] = True

    return sum(has_conflict)


def is_satisfied(board):
    return count_conflicts(board) == 0


def energy_function(board):
    return np.exp(-count_conflicts(board))


def propose_initial_board_uniform(board_size):
    return np.random.permutation(board_size)

Acceptance Ratio with Increasing Board Size

As we see here, increasing the board size results in an exponentially decreasing acceptance ratio. This suggests that we need to improve the proposal distribution. The better of a match the proposal distribution has, the highest the acceptance ratio. If the proposal distribution matches the target distribution, the acceptance ratio is 1.

X_uniform_proposal = []
y_uniform_proposal = []

for board_size in [4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]:
    freqs = {}
    nshots = 100000
    # Now let's sample a couple of boards.
    samples = []
    for _ in range(nshots):
        board = propose_initial_board_uniform(board_size=board_size)
        if is_satisfied(board):
            # for the uniform distribution, E(x) = 0 everywhere, and logc can just be 0, so we can accept all samples.
            samples.append(board)
            board_str = ','.join(str(item) for item in board)
            freqs[board_str] = freqs.get(board_str, 0) + 1

    print(board_size, sum(freqs.values())/nshots, "acceptance ratio")

    X_uniform_proposal.append(board_size)
    y_uniform_proposal.append(sum(freqs.values())/nshots)

import matplotlib.pyplot as plt

plt.plot(X_uniform_proposal, y_uniform_proposal)
plt.yscale("log")
plt.show()

4 0.08289 acceptance ratio
5 0.0831 acceptance ratio
6 0.00584 acceptance ratio
7 0.00854 acceptance ratio
8 0.00239 acceptance ratio
9 0.00093 acceptance ratio
10 0.00013 acceptance ratio
11 0.00013 acceptance ratio
12 4e-05 acceptance ratio
13 1e-05 acceptance ratio
14 0.0 acceptance ratio
15 1e-05 acceptance ratio

Improving the Proposal Distribution

A nice idea is to use MCMC to sample from a superior proposal distribution, which might represent a sort of “surrogate objective”. For example, reducing the number of conflicts on the board. But wait - how can we do this? Aren’t the distributions unnormalized?

For a simple proposal distribution, we can use the energy function \(E(y) = \operatorname{CountConflicts}(y)\), with the corresponding proposal distribution \(q(y) \propto e^{-E(y)}\). Although we cannot sample from this proposal distribution directly, we can indirectly sample from it using Markov Chain Monte Carlo. There are a few details about Markov Chain Monte Carlo that I will omit here. But the premise is that you take a starting board, \(x_t\), and change it slightly to see if there is a reduction in the number of conflicting queens. If there is, then you accept the change. If there isn’t, then you conditionally accept the change, according to an exponential decay based on the number of mismatched queens. The decision for whether to “accept” the changed board is made with probability \(A(x', x)\), where \(x\) is the source board, and \(x'\) is the changed board. The range of boards you can get by “changing” your current board is called a proposal distribution, \(g(\cdot \\mid x)\).

In the general case, the acceptance function is:

If your proposal distribution is symmetric (proposing a change and proposing “undoing” the change have the same probability for all \(x\) and \(x'\)), then the acceptance function becomes:

This is because \(g(x \mid x') = g(x' \mid x)\) in that case.

Bringing it back to our problem, let’s say that the proposal distribution is a mutation of some sorts. Like switching index \(i\) and index \(j\) of our list representation. There are \(\frac{n(n-1)}{2}\) ways to make this switch, and each is sampled from uniformly randomly. Therefore, the proposal distribution is symmetric, and we can use the second simplified acceptance ratio.

Now, we wanted to use this to sample from \(q(x) \sim e^{-E(x)}\). Then \(\log A(x', x) = \max(0, \log q(x') - \log q(x)) = \max(0, -(E(x') - E(x)))\). In practical terms, this means we accept any mutation that reduces the number of conflicts, and accept mutations that increase the number of conflicts with exponential hesitancy.

def propose_initial_board_mcmc(E, x0, proposal_distribution, steps):
    # E is the energy function.
    x = [x0]
    for _ in range(steps):
        next = proposal_distribution(x[-1])
        logA = -(E(next) - E(x[-1]))
        if logA >= 0 or np.random.rand() < np.exp(logA):
            x.append(next)
    return x[-1]


def propose_next_board_by_swap(board):
    new_board = board.copy()
    i, j = np.random.choice(len(board), 2, replace=False)
    new_board[i], new_board[j] = new_board[j], new_board[i]
    return new_board

X_mcmc_proposal = []
y_mcmc_proposal = []

for board_size in [4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]:
    freqs = {}
    nshots = 5000
    # Now let's sample a couple of boards.
    samples = []
    for _ in range(nshots):
        x0 = propose_initial_board_uniform(board_size=board_size)
        board = propose_initial_board_mcmc(
            count_conflicts,
            x0,
            propose_next_board_by_swap,
            steps=20,
        )
        if is_satisfied(board):
            # for the uniform distribution, E(x) = 0 everywhere, and logc can just be 0, so we can accept all samples.
            samples.append(board)
            board_str = ",".join(str(item) for item in board)
            freqs[board_str] = freqs.get(board_str, 0) + 1

    print(board_size, sum(freqs.values()) / nshots, "acceptance ratio")

    X_mcmc_proposal.append(board_size)
    y_mcmc_proposal.append(sum(freqs.values()) / nshots)

import matplotlib.pyplot as plt

plt.plot(X_mcmc_proposal, y_mcmc_proposal)
plt.yscale("log")
plt.show()

4 0.5934 acceptance ratio
5 0.7486 acceptance ratio
6 0.083 acceptance ratio
7 0.1806 acceptance ratio
8 0.0816 acceptance ratio
9 0.0558 acceptance ratio
10 0.0168 acceptance ratio
11 0.0064 acceptance ratio
12 0.005 acceptance ratio
13 0.004 acceptance ratio
14 0.0018 acceptance ratio
15 0.0012 acceptance ratio

Comparison

It looks like this approach is much better! In both cases, we gave the opportunity for \(100000\) “proposals” or “evaluations”. (For the second, we ran for \(5000\) samples which each took \(20\) Monte Carlo steps to generate). Let’s compare the number of boards we get.

num_boards_mcmc_proposal = [v * 5000 for v in y_mcmc_proposal]
num_boards_uniform_proposal = [v * 100000 for v in y_uniform_proposal]

plt.plot(X_mcmc_proposal, num_boards_mcmc_proposal, label="MCMC")
plt.plot(X_uniform_proposal, num_boards_uniform_proposal, label="Uniform")
plt.yscale("log")
plt.legend()
plt.show()

As we can see, using MCMC to initialize drastically reduces the number of samples we need to take. The curve is still exponential (which makes sense, as there are an increasing number of constraints), but the slope has been weakened, allowing us to sample larger boards as if they were a smaller size (here, seemingly as if they were 25% smaller or so). But can we do better? MCMC does not take into account the gradient. Basically, it doesn’t even know that you are better off moving conflicting queens than ones that are already safe. Langevin Dynamics might be a way to fix this. Here, I’m sort of just trying things out and seeing what works.

Continuous Board Representations

It’s possible that my current optimization space is not great. Maybe it’s too easy to get caught in local minima. What if we use a continuous version instead, where we have a tensor \(x \in \mathbb{R}^{n \times n}\)? We can interpret this as \(x_{ij} =\) log probability of the queen in row \(i\) being in column \(j\). Then the energy function can still be the number of conflicts, but we could try to guide each step using the gradient.

Of course, the actual number of conflicts is a discontinuous function of \(x\). We can try to make a continuous approximation, though. The way I’ll try this is by marking the argmax of each row as a “de facto” queen. Then, we can count the number of “expected conflicts” by looking at the squares that this “de facto” queen attacks, and summing up the probabilities. We can see that for the discrete case, this is the same as simply counting the number of pairs of queens that attack each other.

The energy function then becomes the following. Let \(\mathcal{G}\) represent the “de facto” queens; i.e. \(\{(i, j) : i \in \{1, \ldots, n\}, j : \forall k \in \{1, \ldots, n\}, k \neq j : x_{ij} > x_{ik} \}\), and let \(\operatorname{Conflicts}(q)\) represent the board squares that a queen at position \(q\) attacks. Let’s let \(\ell(q)\) represent the log probability that a queen is the selected one in its row.

Noting that conflict sets are symmetric; i.e. if \(q_b \in \operatorname{Conflicts}(q_a)\), then \(q_a \in \operatorname{Conflicts}(q_b)\), we can rewrite this as a sum over all of the board squares:

Then, we can calculate the gradient:

This intuitively says that the gradient with respect to logprobs is proportional to the probability of a queen being here times the number of conflicts with argmax queens.

# here, the board is represented by log probabilities.
def continuous_energy_function(board: np.ndarray):
    grad_E = np.zeros_like(board)
    E = 0

    for de_facto_row in range(board.shape[0]):
        # de facto queen
        de_facto_column = np.argmax(board[de_facto_row])
        for j in range(board.shape[0]):
            if de_facto_row == j:
                continue

            # Conflict at row occupied column in row j
            grad_E[j, de_facto_column] += np.exp(board[j, de_facto_column])
            E += np.exp(board[j, de_facto_column])

            # Conflict at diagonal locations
            occupied_column_offset = abs(de_facto_row - j)
            if de_facto_column - occupied_column_offset >= 0:
                grad_E[j, de_facto_column - occupied_column_offset] += np.exp(
                    board[j, de_facto_column - occupied_column_offset]
                )
                E += np.exp(board[j, de_facto_column - occupied_column_offset])

            if de_facto_column + occupied_column_offset < board.shape[1]:
                grad_E[j, de_facto_column + occupied_column_offset] += np.exp(
                    board[j, de_facto_column + occupied_column_offset]
                )
                E += np.exp(board[j, de_facto_column + occupied_column_offset])

    return E, grad_E

Sanity-Checking our Energy function

Let’s make a simple board with 4 queens, and visualize the gradient of the energy function. We’ll intentionally misplace one of the queens, and plot the gradient of the energy function, \(\nabla_x E(x)\). This misplaced queen conflicts with 3 other queens. Each of the other queens has only a single conflict. Therefore, decreasing the probability of the misplaced queen will proportionally decrease the energy 3 times as much as decreasing the probability of any of the other queens.

board = np.array([
    [-100,    0, -100, -100],
    [-100, -100, -100,    0],
    [   0, -100, -100, -100],
    [-100,    0, -100, -100], # misplaced queen.
])

E, grad_E = continuous_energy_function(board)

# plot the gradient of the energy function.

plt.imshow(grad_E)
plt.colorbar()
plt.show()

Projected Gradient Descent for Optimization with Langevin Dynamics

Here, we use the above energy function and corresponding gradient to perform Langevin Dynamics sampling over “good” N-queens boards. Concretely, at each step, we “should” do (see Eq.6 of Yang Song’s blog):

Where \(\epsilon_t\) is a variance schedule. However I empirically found that just using no coefficient of \(\epsilon_t\) led to better results. Additionally, because we are working with log probabilities, we must project to a simplex, as taking arbitrary gradient steps in log space may cause the probabilities to no longer sum to 1.

def project_to_simplex(x):
    sums = np.sum(np.exp(x), axis=1)
    x -= np.log(sums)[:, None]
    return x

def sample_langevin(x0, noise_schedule, differentiable_energy_function):
    x = [x0]
    for noise_level in noise_schedule:
        xnext = project_to_simplex(
            x[-1]
            - differentiable_energy_function(x[-1])[1]
            + np.sqrt(2 * noise_level) * np.random.randn(*x[-1].shape)
        )
        x.append(xnext)

    return x

Visualizing an Optimization Path

I had to run this cell a few times to get a nice timeline… So it is not guaranteed. However, we see that it is capable of converging to a valid N-queens board when it gets close enough to the local minimum in the energy function.

nsteps = 100
board_size = 6
steps = sample_langevin(
    project_to_simplex(np.random.randn(board_size, board_size)),
    np.exp(np.linspace(0, np.log(0.001), nsteps + 1)),
    # np.linspace(1, 0.001, nsteps + 1),
    continuous_energy_function,
)
for step in steps[::nsteps//4]:
    plt.imshow(np.exp(step))
    # set value range to 0, 1
    plt.clim(0, 1)
    plt.colorbar()
    plt.show()

Seeing if this results in an even better acceptance ratio

Now, we use Langevin Dynamics to sample from our desired proposal distribution \(q(y) \sim e^{-\operatorname{Conflicts}(y)}\), and see if it creates an improvement over Monte Carlo sampling.

Note that we did forgo the previous prior that all the queens’ columns had to be unique! When we go to check the constraints, before we were guaranteed that each of the queens belonged to different columns. Now, when taking the argmax of each row, we could potentially get two queens in the same column. Is this prior essential?

One of the primary reasons I wanted to create a continuous extension of this problem (beyond the fact that it gives us a sort of gradient) is to allow the proposed samples to “wormhole” through valleys of low density.

def propose_initial_board_langevin(energy_function, nsteps, board_size):
    steps = sample_langevin(
        project_to_simplex(np.random.randn(board_size, board_size)),
        np.exp(np.linspace(0, np.log(0.001), nsteps + 1)),
        # np.linspace(1, 0.001, nsteps + 1),
        energy_function,
    )
    final = steps[-1]

    # plt.imshow(final)
    # plt.show()

    # Convert this to a board with the same representation as before.
    board = []
    for row in final:
        board.append(np.argmax(row))
    
    return board

X_langevin_proposal = []
y_langevin_proposal = []

for board_size in [4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15]:
    freqs = {}
    nshots = 5000
    # Now let's sample a couple of boards.
    samples = []
    for _ in range(nshots):
        board = propose_initial_board_langevin(
            energy_function=continuous_energy_function, nsteps=20, board_size=board_size
        )
        # print(board)
        if is_satisfied(board):
            # for the uniform distribution, E(x) = 0 everywhere, and logc can just be 0, so we can accept all samples.
            samples.append(board)
            board_str = ",".join(str(item) for item in board)
            freqs[board_str] = freqs.get(board_str, 0) + 1

    print(board_size, sum(freqs.values()) / nshots, "acceptance ratio")

    X_langevin_proposal.append(board_size)
    y_langevin_proposal.append(sum(freqs.values()) / nshots)

import matplotlib.pyplot as plt

plt.plot(X_langevin_proposal, y_langevin_proposal)
plt.yscale("log")
plt.show()

4 0.643 acceptance ratio
5 0.8308 acceptance ratio
6 0.1306 acceptance ratio
7 0.294 acceptance ratio
8 0.1512 acceptance ratio
9 0.114 acceptance ratio
10 0.043 acceptance ratio
11 0.0254 acceptance ratio
12 0.0218 acceptance ratio
13 0.0134 acceptance ratio
14 0.0086 acceptance ratio
15 0.0068 acceptance ratio

Comparison to Other Approaches

The acceptance probabilities seem really great compared to the Monte Carlo version, EVEN without explicitly adding the constraint that all queens belong to different columns. This suggests that continuous representations for optimization could possibly be even more efficient in higher dimensions than discrete counterparts.

Maybe the next thing to consider would be, how well can we learn a cost function simply from samples? If we trained a model purely with example N-queens solutions, would it be able to infer the cost function?

Perhaps something concerning to note is the runtime. The Langevin approach takes quite a bit longer to run, possibly because we need to take gradient steps on large arrays.

One interesting thing to note is that the gap between the uniform and other approaches increases as the dimensionality of the problem increases.

num_boards_langevin_proposal = [v * 5000 for v in y_langevin_proposal]
num_boards_mcmc_proposal = [v * 5000 for v in y_mcmc_proposal]
num_boards_uniform_proposal = [v * 100000 for v in y_uniform_proposal]

plt.plot(X_langevin_proposal, num_boards_langevin_proposal, label="Langevin")
plt.plot(X_mcmc_proposal, num_boards_mcmc_proposal, label="MCMC")
plt.plot(X_uniform_proposal, num_boards_uniform_proposal, label="Uniform")
plt.yscale("log")
plt.legend()
plt.show()

Closing

Anyway, that’s all I had planned to write for today. It would be interesting to continue exploring whether continuous extensions of discrete problems can allow for better optimization, and more broadly explore the space of Monte Carlo sampling more deeply.

Early Imitation Learning

Date: Late December

Joint Velocity Prediction

Problem Formulation. Given a set of demonstrations for a task, can we learn to perform that task? We formulate this problem as “next keypoint prediction", in which given a state \(s\), we aim to predict the action \(a\) (that is, an \(SE(3)\) end-effector position and rotation) that accurately imitates the demonstrations towards completing the task.

Experimental Setup. I used CoppeliaSim and a small set of tasks from the RLBench repository [@rlbench]. Because it was time-consuming to generate new demonstrations online, I used a set of prepopulated task demonstrations from the PerAct repository [@peract]. The training objective was Huber loss (a type of mean-squared error meant to be robust to outliers) trained on joint velocities (\(\dot q\)): \(\mathcal{L}_{huber} = \begin{cases} \frac{1}{2}(\dot q - \hat {\dot q})^2 & \text{if } |\dot q - \hat {\dot q}| \leq \delta \\ |\dot q - \hat {\dot q}| & \text{if } |\dot q - \hat {\dot q}| > \delta \end{cases}\) I experimented with two image encoders: the vision transformer backbone from CLIP [@clip], and ResNet18 [@resnet] (pretrained on ImageNet). One thing I did not realize at the time, but think would have been interesting to compare, would be to use the vision transformer backbone for CLIP and to compare it to the ResNet backbone for CLIP. I believe this would have elucidated, in a more controlled manner, whether convolutional neural networks or vision transformers have the advantage in robotic manipulation. My intuition at this point tells me that vision transformers are better for “where" inferences, and that convolutional models are better for “what" inferences.

Results. I found that CLIP is a strong feature extractor compared to the ResNet. In particular, it was able to predict joint velocities with relatively high efficacy. I now realize that this is potentially because of the ability for the CLIP vision transformer to utilize positional embeddings, whereas ResNet models are translation-invariant. As an inductive bias in object detection, this is great, but when action outcomes can be dependent on position within an image, position-sensitive models may be better.

Keypoint-Conditioned Imitation Learning

The next iteration on this was to do keypoint-conditioned imitation learning, rather than directly predicting joint velocities. The hypothesized advantage is that keypoints are simpler to predict, which should make learning easier.

Experimental Setup. Here, I adopt a scheme similar to RVT [@rvt] (Robotic View Transformer) and PerAct [@peract] (Perceiver-Actor). Given a set of demonstrations – particularly, a set of observations (RGB images) and robot states (end-effector positions and rotations) – I extract keyframes, and then create a small dataset in which observations are mapped to the “next keyframe".

Keyframe Extraction is done with a simple heuristic: Whether the gripper state changes, or whether there is a change in velocity. Keyframe extraction is simple for demonstrations that occur in simulation, because motion is piecewise linear. It is also possible to extract keyframes from human demonstrations [@waypointslxs].

Simple Experiment: Vision Transformer Token Head. Aforementioned RGB observations (images multiple input directions) have future keypoint locations (members of \(SE(3)\)) projected onto them according to camera matrices. I used a pretrained vision transformer (CLIP) trained with cross-entropy loss to identify future keypoint locations.

Aside: Vision Transformers: Vision transformers [@vit] (ViTs) tokenize images by (1) (optional preprocessing step, not used everywhere) resizing them to \(224 \times 224\) and potentially scaling pixel values, and (2) breaking up the preprocessed image into tokens (often, these are \(14 \times 14\) or \(16 \times 16\) chunks of pixels). These raw tokens (which can be thought of as the word embeddings before a language-based transformer) are encoded through the transformer layers, resulting in contextualized features. For my approach, I simply take an overhead input image and train a classification head for whether a certain token contains the future keypoint.

Aside (Camera Projection Matrices): Given a Cartesian keypoint location \($[\vec x] = [x, y, z, 1]^\top\)$ in robot frame (with homogeneous coordinates), then for each camera, there exists a camera matrix \(M \in \mathbb{R}^{3 \times 4}\) such that \([\vec y] = M [\vec x] \in \mathbb{RP}^2\) is a pixel location. Basically, multiply by the matrix, and turn robot coordinates into pixel locations. This matrix, if you are trying to self-teach, is finicky to construct.

Simple Experiment: Voxelization. Following the Robotic View Transformer from Nvidia, I took point clouds and applied a voxelization step followed by a rendering step. I created my own voxelizer to do this because I was having problems with pytorch3d. I rendered “virtual views" along the orthogonal coordinate axes, and applied the aforementioned keypoint prediction pipeline. This resulted in decent-ish results. However, I was never able to quantify them strongly, nor test them in closed loop.

Aside: Voxelization: To perform voxelization, I defined workspace boundaries (\(\vec x_{min}\) and \(\vec x_{max}\) in robot coordinate system) and discretized the coordinates into buckets. Each 3-tuple of coordinates formed a voxel, which was assigned an occupancy level (\(\alpha\)) and a color (\(c\)). Then, I created a voxel renderer along each axis, through a rendering technique very similar to NeRF: creating rays parallel to the view axis (using an orthographic projection, so all such rays were parallel), and marching them forward. As an example, if the workspace along the \(z\) axis ranges from \(z_{near}\) to \(z_{far}\), I calculate the final pixel value in the image as \(\hat c(x, y) \gets \alpha_n c(x, y, z_n) + (1 - \alpha_n) \hat c(x, y)\) iterating from \(z_n = z_{far}\) to \(z_n = z_{near}\), with \(\hat c(x, y)\) representing the final pixel value for the \(z\) axis image. \(\hat c(x, y)\) is initialized as \((0, 0, 0)\) to give images a black background.

TAKEAWAY: Try things in closed-loop. Keep going until they work!

In-Context Learning

Date: Early January

One motif that I was interested in (since the beginning of the lab) was in-context learning. My vision was to create a robot for which you could explain tasks in natural language, and have the robot ultimately perform them for you. However, I think I did not realize how challenging this would be, because I did not stop to think about what pieces would be necessary.

Background Reading. There was an object-centric prompting mechanism [@vima] and meta-learning [@adaptiveagents] set of papers which were cool cool. I do not believe that I understood them deeply. I also designed an approach to use a temporal (i.e., multiple input-image) vision transformer to predict future actions conditioned on previous actions. Finally, I considered that it was important to train policies to run on the actual robot. I made a couple of setups:

1. Pick Same Object. Sequences of two images. The first would have a certain item circled. The second would have the items shuffled, and the target label would be the item that was circled.

2. Squares. I made a prompt: observe the <color> square, and continue the pattern.

3. Left, Right, Red, Blue. A flag would be set in the model input, prompting it to change its selection behavior between left/right side of an image, or red/blue side of an image.

It took a surprising amount of compute for the model to learn which one to use with in-context instructions. It was easy to train the model to make predictions for single tasks, though. Some of the things I learned were:

1. Learning rate is important in transformers. Dynamic learning rate schedules can be important as well. For example, without a linear warmup, some transformer models will simply never converge, or it would take forever to work. After the linear warmup, some form of annealing is good.

2. Without pretraining, in-context learning is fairly difficult.

Using a Real-Life Robot

Date: Early February

AprilTags

I wanted to try working on a real-life robot. I started by working with a pair of Azure Kinect RGBD cameras, and an AprilTag. To interact with cameras, I use the pyk4a library on GitHub, which requires the proprietary libk4a package (which is unfortunately only available in binary form). I used the apriltag library on GitHub, which provides the pixel locations of each AprilTag’s corners. Using the built-in camera intrinsic parameters, distortion coefficients, and known coordinates of the AprilTag in robot frame, it is possible to calculate a perspective-and-point (PnP) transform, which essentially determines the extrinsic matrix of the camera. Note that this is susceptible to noisy measurements of AprilTag coordinates, or inaccurate AprilTag placement. I (much later) go on to develop a more sophisticated and robust calibration system, which results in strong results after just 1-2 minutes of calibration per camera. The ability to place point clouds in robot frame allows for point clouds from different perspectives to be fused.

Using Offline Data

I tried to use some existing data we had collected. Some notes for future work:

1. Ensure timestamps are correct: In addition to being unsynchronized, I ran into the surprising issue that dependent on the run, images and robot control commands were collected with trial-dependent time offsets. This made it a bit difficult to do anything with the offline data.

2. Use a standardized setup: It was also difficult to create correct mappings from robot state space to pixel space.

At this point I tried to label 3D keypoints visually. That is, given two camera views, I specified target locations by manually selecting the target location as 2D coordinates. As a result of my camera calibration, I was able to have a robot move to a precise 3D location, simply using 2D information. This was an important moment because I realized that 2D image information could be highly useful for robotic policies.

Video Diffusion Models

Date: Late February to Early March

I then decided that large pretrained diffusion models would be good for robotic manipulation. I recognized that many problems in robotic control, similar to large language models with transformers, can be reduced to the form “generate a video of a robot performing X action", conditioned on the robot’s previous observations. There were several signals that would indicate positive performance in this direction. (1) Diffusion Policy [@diffusionpolicy] was demonstrating impressive multimodal (“multi-solution") planning capabilities in robot state space. (2) Image inpainting and video prediction models seemed to be doing well as well [@videolanguageplanning; @susie]. (3) Sergey Levine and Chelsea Finn created a patent for image prediction-based methods for robotic planning (although I think other works had done it earlier?). This may appear true, at least superficially.

I made a presentation of some existing video diffusion methods (view here). Still think I should have thought more logically / at a higher level about this one though, even if the theory was technically there. It’s important to consider the compute and talent budgets that other labs have to execute on more ambitious ideas.

Regardless, I decided to just try this out. I learned how to train large models on multiple A100 GPUs, but I think I could have saved a lot of time lol by focusing on more direct parts of the project and saving this for future work. I learned about EDM preconditioning strategies for diffusion models [@edm], which essentially dictate the loss function weighting at each training diffusion timestep, and the tradeoff between predicting \(x_0\) and predicting \(\epsilon\). I also rederived the forward diffusion reparameterization trick.

For training, I used a pretrained Stable Video Diffusion (SVD) [@stablevideodiffusion] model, and modified an existing HuggingFace training script for the image-based Stable Diffusion. This was harder than I anticipated, because the model architectures were different, and it was my first time training a diffusion model in such a way. In fact, this was so difficult that in order to check my sanity, I followed a guide from nn.labml.ai and trained/implemented my own diffusion model on MNIST. Yang Song and Lilian Weng’s blogposts is a great place to learn how diffusion models work.

In terms of learning about diffusion, I learned about the importance of the signal-to-noise ratio schedule: It should go from \(1\) to \(0\) [@fixtheschedule], this is much easier to tune correctly than a \(\beta\) schedule; besides, you can derive a \(\beta\) schedule from a signal-to-noise ratio, and I think the signal-to-noise ratio is more interpretable. Offset noise is also important to be able to sample very bright or very dark images from diffusion models. I also learned about the importance of the sampling step. You can get significantly better results with slightly under-optimized diffusion models simply by applying more backward inference steps. From a standpoint of annealed Langevin Dynamics, this allows the model to make more corrections at a given noise level, which seems to stabilize things. This all sort of ties into a great underlying theme about the success of diffusion models: They allow ML models to specify cost functions / local gradients (the score function), which are significantly easier to learn than direct samples from the probability distribution. The point is, it is easier to detect mistakes than to generate a satisfying sample directly.

This whole principle that diffusion models are optimizers allows for striking out-of-the-box capabilities. One example is that of optical illusions created by diffusion models [@diffusionillusions].

Anyway. I’m getting sidetracked. Here’s more training details about the video diffusion model. I trained my models on the RT-1 dataset from Google [@rt1]. I formatted the dataset by taking \(25\)-frame chunks (which is what SVD was finetuned for), using no text conditioning (I just wanted to get video diffusion working whatsoever). The results are conditioned on the initial image, the features of which are concatenated to the noisy latents of all future frames to guide generation. I used the HuggingFace Accelerate library, which includes advanced techniques for distributing workloads to multiple GPUs, for training. Overall, I used 10k samples and trained for 2 epochs. The result was a model capable of generating relatively decent results for the compute budget I had. I initially experimented with LoRA to save memory, but I found that the model did not improve enough if I did that.

One important detail that I had not noticed before during training was the impact of the VAE. Misconfiguring the VAE = death for your training. I did not realize that the VAE was trained with a scaling parameter, and during my finetuning, I at one point was training without regard for that. And thus, I deviated too far from the pretraining distribution, resulting in subpar performance for a bit. Once this was fixed, video generations looked significantly better.

Diffused Image from My Model

Loss Curve for Diffusion Training

Considering the lower-level details in retrospect, there are many reasons this should not have worked:

1. Compute limitations. Pretrained diffusion models require massive amounts of compute, and I was not a researcher on model efficiency, I was a researcher in robotics. Additionally, it is unclear whether video diffusion models even should be the standard for motion planning and prediction. Sure, it makes sense to make prediction in image space in some sense, but there are also many more practical approaches that should work equally well before getting to a point where that would make a difference.

2. Implementation limitations. Even with a trained diffusion model, performing evaluations on such a model would be difficult, when considering how it would look in material form. This would require some form of visual servoing, or otherwise, some model to encode the action space (serving as the dynamics model as a function of actions), and some model to extract meaningful reward information about future states (a reward model).

2D Diffusion Models

Date: March

Note: I found that concurrent work actually implemented what I was thinking about doing here. This paper is called Render and Diffuse, and it came from the Dyson Robot Learning Lab. [@renderanddiffuse].

One valid question I had, after seeing that much robotic planning could be performed in 2D space, was to use diffusion over images. This is already sort of what Diffusion Policy does (in retrospect). There were many theoretical reasons for this though:

1. Access to internet-scale pretraining. There are a ton of image-language pretrained transformers that are capable of a wide range of downstream tasks. Additionally, such transformers demonstrate great scene understanding.

2. Compute efficiency. Video diffusion models are extremely costly. Performing diffusion in action space instead is much better.

3. Analogies to visual servoing. Humans primarily measure motor errors visually. Therefore, it should seem natural for an agent to be able to correct its movements given an input image and proprioception, and in the best case, this reduces to a deep-learned version of visual servoing.

There are some issues though with this approach:

1. Simpler prior work: At a high level, this approach purports to have exactly the same benefits as the well-established Diffusion Policy [@diffusionpolicy]. Therefore, I should have worked with that.

2. Insufficiency of 2D corrective terms: My model outputted 2D correction terms, and my hope was that by optimizing the score function across all views simultaneously, I would eventually reach a mode of the joint probability distribution \(\pi(x, y, z | o)\). However, this is simply not true. A maximizer of \(P(x, y)\), \(P(y, z)\), \(P(z, x)\) is not necessarily a maximizer of \(P(x, y, z)\); there is too much ambiguity.

I created a noise-conditioned score matching model, where conditional on the noise level, each output token of a vision transformer would include a direction to move to the nearest mode. The training loss was as follows. For a token at position \((x, y)\), and true keypoint position at \((x_0, y_0)\), predicting the noise \(\hat \epsilon\) and quaternion map \(\hat q\), the loss function was as follows: \(\mathcal{L} = \underbrace{\frac{1}{\sigma^2}((x - x_0) - \hat \epsilon)^2}_{\text{translation noise}} + \underbrace{(\hat q - q_0)^2(e^{-(x - x_0)^2/\sigma^2})}_{\text{quaternion field}}\) Where \(\sigma\) is the standard deviation used by noise conditioning. The quaternion loss was used to enable prediction of end-effector rotations at multiple modes; broadly encouraging nearby quaternion predictions to be the same, but not penalizing distant quaternion predictions for one mode when they were correct for a closer mode. The result was a model that seemed capable of predicting keypoints. Below I have some images:

2D Keypoint Diffusion

Rotation Angle Prediction

I was actually able to get virtual plan rendering working, which was deceptively difficult! This required loading the URDF files and a 3D model into a 3D rendering environment, and using the quaternions and translation vectors to calculate where the end-effector should land in the scene. I don’t think I ever ended up doing anything with this though.

Virtual Plan Rendering

Language Models for Planning

After I realized that it was mathematically impossible for my formulation of 2D diffusion to work (optimizing from multiple views), I decided to pivot towards using language models as planners. The primary reason for this is their impressive zero-shot capability: Language models are great for code generation, and have an impressive prior over general world knowledge. With a bit of fine-tuning, they should perform excellently for robotic control tasks. One line of work on language models for robotic planning, particularly for long-horizon robotic planning, envisions the use of a language model that processes observations about the environment via external APIs called via tool use [@innermonologue] or code generation [@codeaspolicies; @voxposer; @progprompt].

There are also some interesting open questions for language models:

1. Closed-loop control: Language models generate code policies for multiple timesteps. Executing these code policies without checking for task execution errors isn’t a great idea. What ways are there to use VLMs as “error checkers"?

2. Instruction memory: Can we encode instructions and use language models to “recall" desired behavior in the future? What ways are there to implicitly create this type of memory?

3. Visual grounding: When (vision) language models reason, it is hard to connect their reasoning to the images they refer to. However, if we could reason over grounded objects, we could unlock significantly more advanced planning capabilities. This is still an open question and something I am interested in exploring further.

I explored the first and third points by creating an agentic LLM planner. Particularly, I created a notion of an “event stream", in which observations, actions, rationales, and reflections would enter a shared space for future reference by the language model.

There was also a ton of engineering effort required to make this into a reality. First, there was camera calibration to be able to fuse RGBD point clouds in robot space. Then, there was camera-robot calibration, which adjusted the resulting detections. Finally, there was grasp synthesis, and execution on the robot control system.

Object Detection and Selection

Plan generation. Based on previous actions taken, the current instruction, and the current image observation, the VLM generates a plan for how to interact with the objects in scene.

Code generation. Based on the reasoning generated by the VLM, I then use a separate prompt to generate code to control the robot. I expose a small set of APIs:

1. Object detection: This returns objects with 3D point clouds in robot frame, conditioned on a target text label.

2. Relative motion: This allows the robot to translate or rotate its end-effector by a calculated amount, in meters. This is useful, for example, for lifting items up simply by performing an upward translation.

3. Absolute motion: This allows motion with respect to anchor coordinate locations. One example for why this might be useful is if the robot arm needs to move above a certain item to pick it up.

4. Gripper control: This allows the robot to open or close its gripper. Additional logic adjusts the grasp axis to perform a force closure, through a rough (and non-robust) heuristic I made.

Object selection. I used an OwlV2 [@owlv2] model (which is great for open-vocabulary object detection) to locate objects of a specific type. To choose specific objects based on fine-grained attributes (like color, text, etc.), I create image crops centered around each of the object detections, and ask the VLM to describe the physical attributes of those objects. Then, these descriptions are compiled into a summary (a list of lines Object 1: <Description>\nObject 2: <Description>\n ...), before the VLM makes its final decision for which object to select, or none of the objects if none match.

Object point cloud extraction. Once an object has been located in a 2D image, we must further ground it in a 3D location. This is done in two steps. The first step is to use the Segment Anything Model [@sam] to generate an image segmentation. Because images are in RGBD, this image segmentation automatically generates a set of object points in robot state space. However, these object points only represent one side of the object, and grasp planning with this will not work. To fix this, I use a camera on the other side of the robot, as depicted below.

Camera Layout

Camera-Robot Cross Calibration. To create accurate calibrations of images to robot state, I wrote a script that (1) allowed a user to place a block with an AprilTag on it in the robot’s gripper, and (2) moved the AprilTag block to predefined locations in a 3D cubic grid. I then logged robot state information directly from the robot (rather than using the waypoint as a ground truth) due to control error. Using the correspondences between image pixel locations from the center of the AprilTag, and the known positions in robot coordinate space, I was able to use OpenCV’s solvePnP function to create a more accurate camera extrinsic calibration than the original AprilTag approach described earlier.

Robot Control. Robotic control was done with Polymetis, a framework that allows specialized control policies to be deployed at a high frequency. It includes a barebones min-jerk motion planner as well. Control is done with PID over joint positions/velocities in a tight loop rather than model predictive control. I found that this mostly worked out of the box except for two major challenges:

1. Flange frame miscalibration: For some reason, all actions the robot took were offset by \(10.34\) centimeters downward, and with a rotation of \(\pi/4\) radians around the z axis. Other people experienced this too (here).

2. Improper control coefficients: This was much easier to discover and debug. I found that the robot sagged under gravity at times, and was otherwise susceptible to slight pushes; to fix this, I increased the proportional gain in the PID.

I also used a very simple motion planning heuristic to reduce the number of collisions with items on the table. First, whenever the robot’s Z value was anticipated to move downward during a motion, I would first make the horizontal motion, and then the downward motion. If a movement involved upward motion, I would make the upward motion first, and then the horizontal motion. The overall result from this is that most of the horizontal motion occurs at higher Z levels, resulting in fewer collisions for basic tabletop tasks.

Robot Completing a Task

Grasp Generation Heuristic. When the robot attempts to grasp a target object, it uses a grasp generation heuristic I created. This voxelizes the object point cloud, rotates the voxels on their side, and rotates around the (former) z-axis, looking for appropriate planes to grasp along. This plane estimation approach is highly non-robust, and simply uses the slope of the voxelization at the top and bottom of the object to calculate the angles (\(\alpha\)) that dictate the amount of friction generated. Hole-filling or other techniques would make it better. A score is calculated based on these angles of approach which determines whether there is enough friction to perform the grasp. This is based on the Fast Force Closure algorithm by Van-Duc Nguyen [@fastforceclosure]. I found that this approach was capable of correcting gripper rotation angles for objects rotated at 45 degrees, enabling more successful grasps.

Object Memory for Language Models

I performed some additional experiments to create a sort of “visual object memory" for language models. For example, when looking at a snickers bar, a robot should be able to recognize that it contains peanuts, and thus not recommend it to someone with a peanut allergy. In retrospect, this set of experiments was hard to test, and it was not clear there was high demand for this sort of feature compared to other, high-priority components. To implement a non-verbal form of object memory, I used the following technique:

CLIP heatmap of the scene: I used the MaskCLIP reparameterization trick [@maskclip] to create text-aligned, dense CLIP features. The trick works by replacing the final transformer layer in the CLIP ViT: instead of applying self-attention, we perform a value projection on each token, and project from image embedding space to the multimodal text-image embedding space. To enable CLIP heatmaps to be created from images of arbitrary aspect ratios, I use positional encoding interpolation from the original ViT paper and DINO [@vit; @dino], which I borrow from the Distilled Feature Features [@featurefields] GitHub repository. Essentially, this reshapes the \(16 \times 16\) grid of positional encodings built into CLIP and applies a bilinear interpolation to any arbitrary grid size desired. Finally, I take the token-level CLIP features and interpolate them into pixel space.

Registration of object features: Given an object segmentation obtained via SAM, I then segment the pixel-level CLIP features to create a positive set (inside the segmentation) and negative set (outside the segmentation). I also augment these sets with Gaussian blurs, slight rotations, and background removal. Finally, I reweight the classes to 25% positive features, and 75% negative features, before applying a Support Vector Machine [@svm]. I found it to be crucial to perform class reweighting and data augmentation to increase the robustness of the classification. I used SVMs because of their ability to output well-calibrated class probabilities compared to a naïve dot product. Another reason for using SVMs is that they can use any kernel for classification, be it a linear function (as I used) or radial basis functions (which represent a mixture of Gaussians).

Object tracking: To further augment object features, I used XMem [@xmem], a method aimed at generating continual object segmentations in video. This seemed highly capable and generally worked for tracking objects as they rotated in a scene to enable fast collection of additional features.

Instance segmentation: After creating a feature detector, I needed to be able to reconstruct accurate object segmentations. Given a new RGB image of a scene, I extracted dense token-level CLIP features in the same way as before. I used a threshold to identify positive features and create a mask. Then, I encoded the input image with SAM and prompted it with randomly-selected points within the mask. Any masks that had sufficient duplication with other masks (e.g., IOU > 0.7) would be removed, and any masks which seemed to be false positives (e.g., positive features only made a small fraction of the resulting mask) were also discarded.

Result: Below is an image of one-shot object segmentation. In this case, the model is given a small set of annotations for a certain type of food packaging, and is able to recognize the target object regardless of distractors in the scene.

One-shot object segmentation

Alas, after this, I did not perform a closed-loop test of whether this could be used to modify a robot’s behavior.

Visual Grounding for Language Models

Date: May

I realized rather quickly for VLMs that a really important question is whether we can have VLMs output spatial locations rather than just language. That is, whether the objects of their reasoning can be grounded in the input image. One way to do this is to first generate all objects in a scene, and then detailed descriptions for each of these objects, and finally, select the optimal object out of that stack.

Can we have language models which reason about objects in a grounded manner? One compelling ideal outcome is to have the same reasoning capability as textual chain-of-thought, but for scenes. To experiment with this, I use the test task of generating referring expressions and detailed object captions given bounding boxes. Compared to other approaches, though, which might require reencoding an image crop for every object in a scene, allows one to directly decode object captions prompted by bounding boxes. In addition to improved computational speed, we can reason about objects within bounding boxes in terms of the surrounding context, which is impossible to do when just describing an object based on its cropped image.

Vision-Language Models Effectively Extract Object Attributes. I found that when using excellent vision-language models (such as GPT-4V), accuracy increases significantly for tasks requiring references to specific object attributes. For example, whereas a simple image embedding might capture the general type of an object, generating object captions with a vision-language model (which have more tokens allocated to represent larger and potentially more complex objects) allows for more granular detail, even compared to approaches that claim to “unleash extraordinary object grounding", like Set-of-Mark prompting [@setofmark]. However, in my opinion, it is unclear still whether these results are due to superior attribute extraction or due to the choice of vision model used (in our case, GPT-4V, which has who knows how many parameters, compared to OwlV2, which has a couple hundred million parameters).

A set of visually-similar objects

Pretraining with Object Classification. I initially trained the multimodal projection layer and sequence modeling components of PaliGemma with object classification. Particularly, I prompted the model to classify an object corresponding to a bounding box (represented as four <locXXXX> tokens in xyxy format), and set the target to be the text representation of the MSCOCO class name. When applied to the train2014 split of MSCOCO, this provides approximately 600000 training examples, spanning 80000 images. This provides high object detection density per image, forcing models to respect bounding box information to generate correct labels.

Finetuning with Referring Expressions: OCID-Ref and RefCOCO. I found that OCID-Ref [@ocidref] seemed to be a good dataset for referring expression comprehension. Importantly, it had a high density of objects per image, making it a great way to train object captioning models with strong respect for spatial locations. The RefCOCO datasets [@refcoco] were also great for this. I fine-tuned PaliGemma to generate referring expressions for objects given bounding boxes. I then ran an experiment to see whether objects could be correctly disambiguated based on their referring expressions.

Finetuning with Descriptive Captions: Synthetic Data Approach. I generated a small dataset of \(5000\) object-centric captions for finetuning, using the Phi-3 Vision [@phi3] model from Microsoft to extract object attributes for a large set of object detections in the MSCOCO dataset. I found that including the class of object in the prompt was crucial to generate attributes correctly, as otherwise, models would misclassify objects and describe corresponding features incorrectly due to the lack of surrounding context. Additionally, I prompted Phi-3 to describe physical attributes of the images only, as it would otherwise hallucinate features that correspond to the target object that were not explicitly visible, or misclassify visible features. I finetuned PaliGemma on this dataset for \(4\) epochs.

A detailed object caption

A referring expression generated from an object bounding box

Left: No object pretraining; Right: Object pretraining

Conclusion

I learned a lot from this experience. In retrospect, much of the research that was performed here was close to being a rich project, and simply required more pointed investigation and testing.

Acknowledgements

I would like to thank Mohammad Samin Yasar for his continual support of my work, even during the week I returned to Charlottesville to try to get something working, and for teaching me how to use the robots in the lab.

References

Black, Kevin, Mitsuhiko Nakamoto, Pranav Atreya, Homer Walke, Chelsea Finn, Aviral Kumar, and Sergey Levine. 2023. “Zero-Shot Robotic Manipulation with Pretrained Image-Editing Diffusion Models.” https://arxiv.org/abs/2310.10639.

Blattmann, Andreas, Tim Dockhorn, Sumith Kulal, Daniel Mendelevitch, Maciej Kilian, Dominik Lorenz, Yam Levi, et al. 2023. “Stable Video Diffusion: Scaling Latent Video Diffusion Models to Large Datasets.” https://arxiv.org/abs/2311.15127.

Brohan, Anthony, Noah Brown, Justice Carbajal, Yevgen Chebotar, Joseph Dabis, Chelsea Finn, Keerthana Gopalakrishnan, et al. 2023. “RT-1: Robotics Transformer for Real-World Control at Scale.” https://arxiv.org/abs/2212.06817.

Burgert, Ryan, Xiang Li, Abe Leite, Kanchana Ranasinghe, and Michael S. Ryoo. 2023. “Diffusion Illusions: Hiding Images in Plain Sight.” https://arxiv.org/abs/2312.03817.

Caron, Mathilde, Hugo Touvron, Ishan Misra, Hervé Jégou, Julien Mairal, Piotr Bojanowski, and Armand Joulin. 2021. “Emerging Properties in Self-Supervised Vision Transformers.” https://arxiv.org/abs/2104.14294.

Cheng, Ho Kei, and Alexander G. Schwing. 2022. “XMem: Long-Term Video Object Segmentation with an Atkinson-Shiffrin Memory Model.” https://arxiv.org/abs/2207.07115.

Chi, Cheng, Zhenjia Xu, Siyuan Feng, Eric Cousineau, Yilun Du, Benjamin Burchfiel, Russ Tedrake, and Shuran Song. 2024. “Diffusion Policy: Visuomotor Policy Learning via Action Diffusion.” https://arxiv.org/abs/2303.04137.

Dong, Xiaoyi, Jianmin Bao, Yinglin Zheng, Ting Zhang, Dongdong Chen, Hao Yang, Ming Zeng, et al. 2023. “MaskCLIP: Masked Self-Distillation Advances Contrastive Language-Image Pretraining.” https://arxiv.org/abs/2208.12262.

Dosovitskiy, Alexey, Lucas Beyer, Alexander Kolesnikov, Dirk Weissenborn, Xiaohua Zhai, Thomas Unterthiner, Mostafa Dehghani, et al. 2021. “An Image Is Worth 16x16 Words: Transformers for Image Recognition at Scale.” https://arxiv.org/abs/2010.11929.

Du, Yilun, Mengjiao Yang, Pete Florence, Fei Xia, Ayzaan Wahid, Brian Ichter, Pierre Sermanet, et al. 2023. “Video Language Planning.” https://arxiv.org/abs/2310.10625.

Goyal, Ankit, Jie Xu, Yijie Guo, Valts Blukis, Yu-Wei Chao, and Dieter Fox. 2023. “RVT: Robotic View Transformer for 3D Object Manipulation.” https://arxiv.org/abs/2306.14896.

He, Kaiming, Xiangyu Zhang, Shaoqing Ren, and Jian Sun. 2015. “Deep Residual Learning for Image Recognition.” https://arxiv.org/abs/1512.03385.

Hearst, M. A., S. T. Dumais, E. Osuna, J. Platt, and B. Scholkopf. 1998. “Support Vector Machines.” IEEE Intelligent Systems and Their Applications 13 (4): 18–28. https://doi.org/10.1109/5254.708428.

Huang, Wenlong, Chen Wang, Ruohan Zhang, Yunzhu Li, Jiajun Wu, and Li Fei-Fei. 2023. “VoxPoser: Composable 3D Value Maps for Robotic Manipulation with Language Models.” https://arxiv.org/abs/2307.05973.

Huang, Wenlong, Fei Xia, Ted Xiao, Harris Chan, Jacky Liang, Pete Florence, Andy Zeng, et al. 2022. “Inner Monologue: Embodied Reasoning Through Planning with Language Models.” https://arxiv.org/abs/2207.05608.

James, Stephen, Zicong Ma, David Rovick Arrojo, and Andrew J. Davison. 2020. “RLBench: The Robot Learning Benchmark & Learning Environment.” IEEE Robotics and Automation Letters.

Jiang, Yunfan, Agrim Gupta, Zichen Zhang, Guanzhi Wang, Yongqiang Dou, Yanjun Chen, Li Fei-Fei, Anima Anandkumar, Yuke Zhu, and Linxi Fan. 2023. “VIMA: General Robot Manipulation with Multimodal Prompts.” https://arxiv.org/abs/2210.03094.

Karras, Tero, Miika Aittala, Timo Aila, and Samuli Laine. 2022. “Elucidating the Design Space of Diffusion-Based Generative Models.” https://arxiv.org/abs/2206.00364.

Kazemzadeh, Sahar, Vicente Ordonez, Mark Matten, and Tamara Berg. 2014. “ReferItGame: Referring to Objects in Photographs of Natural Scenes.” In Proceedings of the 2014 Conference on Empirical Methods in Natural Language Processing (EMNLP), edited by Alessandro Moschitti, Bo Pang, and Walter Daelemans, 787–98. Doha, Qatar: Association for Computational Linguistics. https://doi.org/10.3115/v1/D14-1086.

Kirillov, Alexander, Eric Mintun, Nikhila Ravi, Hanzi Mao, Chloe Rolland, Laura Gustafson, Tete Xiao, et al. 2023. “Segment Anything.” https://arxiv.org/abs/2304.02643.

Liang, Jacky, Wenlong Huang, Fei Xia, Peng Xu, Karol Hausman, Brian Ichter, Pete Florence, and Andy Zeng. 2023. “Code as Policies: Language Model Programs for Embodied Control.” https://arxiv.org/abs/2209.07753.

Lin, Shanchuan, Bingchen Liu, Jiashi Li, and Xiao Yang. 2024. “Common Diffusion Noise Schedules and Sample Steps Are Flawed.” https://arxiv.org/abs/2305.08891.

Microsoft. 2024. “Phi-3 Technical Report: A Highly Capable Language Model Locally on Your Phone.” https://arxiv.org/abs/2404.14219.

Minderer, Matthias, Alexey Gritsenko, and Neil Houlsby. 2023. “Scaling Open-Vocabulary Object Detection.” https://arxiv.org/abs/2306.09683.

Nguyen, Van-Duc. 1988. “Constructing Force- Closure Grasps.” The International Journal of Robotics Research 7 (3): 3–16. https://doi.org/10.1177/027836498800700301.

Radford, Alec, Jong Wook Kim, Chris Hallacy, Aditya Ramesh, Gabriel Goh, Sandhini Agarwal, Girish Sastry, et al. 2021. “Learning Transferable Visual Models from Natural Language Supervision.” https://arxiv.org/abs/2103.00020.

Shen, William, Ge Yang, Alan Yu, Jansen Wong, Leslie Pack Kaelbling, and Phillip Isola. 2023. “Distilled Feature Fields Enable Few-Shot Language-Guided Manipulation.” https://arxiv.org/abs/2308.07931.

Shi, Lucy Xiaoyang, Archit Sharma, Tony Z. Zhao, and Chelsea Finn. 2023. “Waypoint-Based Imitation Learning for Robotic Manipulation.” https://arxiv.org/abs/2307.14326.

Shridhar, Mohit, Lucas Manuelli, and Dieter Fox. 2022. “Perceiver-Actor: A Multi-Task Transformer for Robotic Manipulation.” In Proceedings of the 6th Conference on Robot Learning (CoRL).

Singh, Ishika, Valts Blukis, Arsalan Mousavian, Ankit Goyal, Danfei Xu, Jonathan Tremblay, Dieter Fox, Jesse Thomason, and Animesh Garg. 2022. “ProgPrompt: Generating Situated Robot Task Plans Using Large Language Models.” https://arxiv.org/abs/2209.11302.

Team, Adaptive Agent, Jakob Bauer, Kate Baumli, Satinder Baveja, Feryal Behbahani, Avishkar Bhoopchand, Nathalie Bradley-Schmieg, et al. 2023. “Human-Timescale Adaptation in an Open-Ended Task Space.” https://arxiv.org/abs/2301.07608.

Vosylius, Vitalis, Younggyo Seo, Jafar Uruç, and Stephen James. 2024. “Render and Diffuse: Aligning Image and Action Spaces for Diffusion-Based Behaviour Cloning.” https://arxiv.org/abs/2405.18196.

Wang, Ke-Jyun, Yun-Hsuan Liu, Hung-Ting Su, Jen-Wei Wang, Yu-Siang Wang, Winston H. Hsu, and Wen-Chin Chen. 2021. “OCID-Ref: A 3D Robotic Dataset with Embodied Language for Clutter Scene Grounding.” https://arxiv.org/abs/2103.07679.

Yang, Jianwei, Hao Zhang, Feng Li, Xueyan Zou, Chunyuan Li, and Jianfeng Gao. 2023. “Set-of-Mark Prompting Unleashes Extraordinary Visual Grounding in GPT-4V.” https://arxiv.org/abs/2310.11441.

Context: Object Detection for Autonomous Racing

Our autonomous vehicle is supplied by the Indy Autonomous Challenge, which includes several state-of-the-art Luminar LiDAR sensors as well as six cameras (two front-facing, normal FOV cameras, and four wide-angle to capture the side and back views).

Each camera sends detections to a YOLOv5 model, which provides bounding boxes for the opponent vehicle. When given bounding boxes for the opponent that correspond across multiple views, we can calculate a 3D position. Because all cars used in this race have the same dimensions (the frames are provided by Dallara, an Italian auto manufacturer), we can use the size of the vehicle to estimate depth.

To calculate the direction, we can use our knowledge of the camera’s intrinsic/extrinsic matrices and distortion parameters. The cameras on this car follow a Brown-Conrady distortion model, and although I have no idea how to derive how that works, I do know that I can use OpenCV to undistort the images for me! So I just use that, along with the distortion parameters provided from the manufacturer. We can calculate the direction of the opponent vehicle in map frame by taking the undistorted 2D locations on the camera’s embedding plan (at \(z = 1\)) and then rotating them to map frame. Given a point \([x, y, 1]\) on the camera embedding plane, we can undistort this to world frame simply by applying the camera’s rotation matrix \(R\) (more information in the Camera Background: Pinhole Camera Model section below). Then, we can scale this by the estimated depth of the vehicle, and add the camera’s translation vector \(\hat t = (x, y, z)\) to get the final 3D position in map frame.

We train these models using the standard pipeline from Ultralytics. We provide 2D bounding box labels for a large dataset of images collected from real-life historical data: an MIT-KAIST race from LVMS in a prior year, a UVA-MIT race from LVMS 2024, and a Euroracing-TUM race.

We are using Grounding Dino to generate labels. This is a great foundation model, although we cannot use it in the race for two reasons: (1) it is slow; (2) there are some false positives (e.g. non-Indy cars). YOLO (or similar) models provide an opportunity to create distilled versions of foundation models with the added benefits of high speed and a more curated dataset.

We have a simple method for eliminating false positives from human supervision. We iterate through the video, frame-by-frame, automatically labelling new bounding boxes as correct if they are within a certain pixel distance of the previous detection. However, when a false positive appears, it is generally a high distance from our current car detection, and we automatically discard those. Whenever there is ambiguity (for example, the opponent enters the frame, and we cannot determine whether the label is accurate or not automatically), we defer to human supervision. This does not happen often. The end result is that we can take advantage of the general capability of the foundation model to detect cars, while simultaneously curating with human feedback.

Result

The resulting model is \(14\) Mb, which is great and can run at high Hz. Additionally, it has significantly higher recall than the original foundation model. We have integrated this into our object detection pipeline, and using the simple monocular depth estimation equation from before, we can estimate 3D location:

Additionally, we find that the resulting model is excellent even in harsh lighting conditions. The following image grid consists only of images that were not in the training dataset, because their labels were not detected successfully by the foundation model:

Camera Background: Pinhole Camera Model

Feel free to skip this section if you already know how the pinhole camera model works!

Cameras are determined by a position/orientation (extrinsic) and a focal length (intrinsic). Intrinsic means no matter how you move the camera around, they will stay the same. Extrinsic means when you move the camera, they change.

The pinhole camera model maps 3D locations to images. This is done via matrix multiplication. Think of matrix multiplication as rotating the world. We rotate the world so that (0, 0, -1) faces the camera. Then x and y in this rotated world correspond to the same x and y directions as an image. We use the intrinsic matrix to convert x and y in these rotated coordinates into pixel values (e.g. 1 meter in world coordinates becomes, say, 240 pixels). There is also a notion of translation; just think of this as subtracting coordinates so that (0, 0, 0) in the translated world is at the center of the camera.

Extrinsics

An extrinsic camera matrix follows the following format: [R t] where R is a 3x3 rotation matrix and t is a 3x1 translation vector.

The extrinsic matrix is something that takes a point in rear_axle_middle_ground frame and converts it to “camera” frame. The main difference between these (besides the angle/perspective of the camera) is the difference between the axes.

For the car, “forward” is X, “right” is Y, and “up” is Z. For the camera, though, “forward” is Z, “right” is X, and “up” is Y (or -Y, if you go top->bottom in images).

Let’s say we want to design a matrix to convert a point in “world coordinates” to the coordinates from the point of view of the camera.

The first step is to translate. Let’s say the input point is (x, y, z), and the camera’s location is (tx, ty, tz). We can write this as a (4D) matrix multiplication in the following way:

[1 0 0 -tx][x]   [x - tx]
[0 1 0 -ty][y] = [y - ty]
[0 0 1 -tz][z]   [z - tz]
[0 0 0   1][1]   [1     ]

Consider this from the perspective of the camera. A point at (tx, ty, tz) must go to the point (0, 0, 0) from the perspective of the camera. So this is the first part.

translation = np.array([
    [1, 0, 0, -x],
    [0, 1, 0, -y],
    [0, 0, 1, -z],
    [0, 0, 0, 1],
])

The next step is to rotate. Imagine the camera is pointing straight ahead.

For a rotation matrix (which is orthonormal), every row represents what that part of the matrix is “detecting”. To explain more clearly, consider this matrix:

    [0 1 0 0][x]   [y]
    [0 0 1 0][y] = [z]
    [1 0 0 0][z]   [x]
    [0 0 0 1][1]   [1]

Recall: For the car, “forward” is X, “right” is Y, and “up” is Z. For the camera, though, “forward” is Z, “right” is X, and “up” is Y (or -Y, if you go top->bottom in images).

This means that for the output of this matrix,

• Row 1 (the “x” row) is detecting Y values, by dot-producting [0 1 0 0] with [x y z 1]
• Row 2 (the “y” row) is detecting Z values, by dot-producting [0 0 1 0] with [x y z 1]
• Row 3 (the “z” row) is detecting Z values, by dot-producting [1 0 0 0] with [x y z 1]

This makes sense if you consider a matrix multiplication onto a vector as dotting each row of the matrix with the vector to produce a new vector.

Another thing this means is that if we invert the rotation matrix (which, by virtue of being orthonormal, amounts to a matrix transpose [making the rows -> columns and vice versa]), then multiplying by a point in “camera” coordinates outputs how the point would look in rear_axle_middle_ground coordinates.

To create the final camera extrinsics matrix, we multiply the rotation matrix by the translation matrix. This creates a matrix that has the same effect of (1) translating the world to the center of the camera, and (2) rotates the world so the coordinate axes line up with the axes of the image.

I wonder how we can meaningfully code for orientation and scale in images through neural networks. My goal is to maximize the number of explainable features that neural networks generate.

I experiment with the CelebA dataset and finetune a pretrained model (from ImageNet) to disambiguate between faces. The goal is that the model learns to output a fine-tuned vector space that distinguishes between faces. I could use a VAE but that takes a while; instead I will estimate mean and variance from a decent batch size and use KL divergence on that.

Download Dataset

We’ll be using the CelebA dataset for this. I’ll assume you have a Kaggle account and can download with that. To save me from the hassle of kaggle.json, I’m using environment variables.

# First, download CelebA dataset.
import os
import getpass

os.environ['KAGGLE_USERNAME'] = getpass.getpass("Username")
os.environ['KAGGLE_KEY'] = getpass.getpass("API Key")

!kaggle datasets download -d jessicali9530/celeba-dataset -p CelebA

Username··········
API Key··········
Downloading celeba-dataset.zip to CelebA
100% 1.33G/1.33G [00:12<00:00, 230MB/s]
100% 1.33G/1.33G [00:12<00:00, 117MB/s]

# Our images unzip to `./img_align_celeba/img_align_celeba` (relative to the current directory; not the CelebA folder for some reason.)
!unzip CelebA/celeba-dataset.zip -d data > /dev/null

Load dataset

Loading images from a folder is a common process, so Torchvision has a way for us to do this automatically. To test that everything works, I set up a quick visualization. It renders the first 16 images as a 4x4 grid using matplotlib.

import matplotlib.pyplot as plt
import torch
import torchvision
import torchvision.transforms as T
from torch.utils.data import DataLoader
from torchvision.datasets import ImageFolder

def visualize_4x4(images: torch.Tensor, title=None):
    # Shape: [N, C, H, W]

    # Make a grid of images
    grid = torchvision.utils.make_grid(images, nrow=4).numpy()

    # Render
    plt.figure(figsize=(4, 4))
    plt.title(title or "Celebrities!")
    plt.imshow(grid.transpose((1, 2, 0)))
    plt.axis("off")
    plt.show()

# Now, train a contrastive model. Use CelebA dataset.
# dataset[i] = (PIL.Image.Image, <dummy class label>)
dataset = ImageFolder(root="./data", transform=T.Compose([
    T.ToTensor(),
    T.Resize((224, 224)),
]))

visualize_4x4([dataset[i][0] for i in range(16)])

/usr/local/lib/python3.10/dist-packages/torchvision/transforms/functional.py:1603: UserWarning: The default value of the antialias parameter of all the resizing transforms (Resize(), RandomResizedCrop(), etc.) will change from None to True in v0.17, in order to be consistent across the PIL and Tensor backends. To suppress this warning, directly pass antialias=True (recommended, future default), antialias=None (current default, which means False for Tensors and True for PIL), or antialias=False (only works on Tensors - PIL will still use antialiasing). This also applies if you are using the inference transforms from the models weights: update the call to weights.transforms(antialias=True).
  warnings.warn(

Training

We’ll use a resnet18 pretrained model which is pretty standard, and it should have decent feature representations already, which should easily be finetuned to this domain.

from torchvision.models import resnet

# Decide whether to use the GPU. This speeds up training a *LOT* if it's available.
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

model = torch.hub.load('pytorch/vision:v0.10.0', 'resnet18', pretrained=True)
# Remove the fully-connected layer because we're interested in embeddings.
model.fc = torch.nn.Identity()
model = model.to(device)

Using cache found in /root/.cache/torch/hub/pytorch_vision_v0.10.0
/usr/local/lib/python3.10/dist-packages/torchvision/models/_utils.py:208: UserWarning: The parameter 'pretrained' is deprecated since 0.13 and may be removed in the future, please use 'weights' instead.
  warnings.warn(
/usr/local/lib/python3.10/dist-packages/torchvision/models/_utils.py:223: UserWarning: Arguments other than a weight enum or `None` for 'weights' are deprecated since 0.13 and may be removed in the future. The current behavior is equivalent to passing `weights=ResNet18_Weights.IMAGENET1K_V1`. You can also use `weights=ResNet18_Weights.DEFAULT` to get the most up-to-date weights.
  warnings.warn(msg)

Training the model

Because this is a smaller model, its capacity is strained for distinguishing between faces. This model was trained to detect objects of a certain range of classes (in ImageNet), and its internal/latent representation is a continuous vector. Classifications are made by having class-specific vectors that are dotted with the latent vector to create scores. In this case, we want to minimize the similarity between vectors belonging to different faces; so we use a contrastive loss. This forces the model to output dissimilar vectors for dissimilar faces. By inspecting the mapping the model creates as a result (which we assume will create similar vectors for similar faces, as the model is forced to compress some of the information it sees), we can try to extract salient facial features.

The loss function is as follows. Let $X$ be the set of input images. \(y \leftarrow \text{Model}(X) \\ \text{Loss}(y) := L_{ce} + L_{orient}\)

Where Cross Entropy is defined as \(L_{ce}(logprobs, target) := -\sum_{i=0}^{k} target_i logprobs_i\)

Where $logprobs$ is the log-Softmax output of the prediction vector.

To train the model to incorporate implicit orientation, I train it to reduce the mean-squared-error between the encoding of the input image and the negative of the encoding of the flipped image. \(L_{orient}(y_{normal}, y_{flipped}) := L_{mse}(y_{normal}, -y_{flipped})\)

import tqdm

batch_size = 128
dataloader = DataLoader(dataset, batch_size=batch_size, shuffle=True)

optim = torch.optim.Adam(model.parameters(), lr=1e-3)

N_EPOCHS = 1
for epoch in range(N_EPOCHS):
  with tqdm.tqdm(dataloader, desc="Training...") as pbar:
    for (images, _) in pbar:
      images = images.to(device)
      # Train model to disambiguate between these images.
      # Use embeddings as features. Instead of retraining
      # all of EfficientNet, add a linear layer at the end
      # to get embeddings. Rotate vector around first dimension.
      # 180º rotated images should have negative embeddings of each other.
      images_rotated = torch.rot90(images, 2, (2, 3))
      # [N, embedding_dim]
      embeddings = model(images)
      embeddings_rotated = model(images_rotated)

      # ORIENTATION LOSS
      loss_orientation = torch.nn.functional.mse_loss(embeddings, -embeddings_rotated)

      # CONTRASTIVE LOSS
      # Dot each of these embeddings with each other and treat these as logits.
      # Then use Cross Entropy loss.
      embeddings_normalized = torch.nn.functional.normalize(embeddings, dim=1)
      scores = embeddings @ embeddings.T
      loss_contrastive = torch.nn.functional.cross_entropy(scores, torch.arange(len(scores), device=device))

      loss = loss_contrastive + loss_orientation

      optim.zero_grad()
      loss.backward()
      optim.step()
      pbar.set_postfix({"loss": loss.item(), "contrastive_loss": loss_contrastive.item(), "orientation_loss": loss_orientation.item()})

Training...: 100%|██████████| 1583/1583 [34:38<00:00,  1.31s/it, loss=0.0944, contrastive_loss=0.00796, orientation_loss=0.0864]

Save our model.

This took like half an hour. I trained on the whole of CelebFace while writing the below content. Let’s please save this model lol.

torch.save(model.state_dict(), "resnet_18_celebface_state_dict.pt")

Create inferences.

We did a single pass over this dataset. Let’s create vectors for each of the faces in the dataset now.

results = []
model.eval()

# Notably, we *DO NOT SHUFFLE*!
dataloader_eval = DataLoader(dataset, batch_size=batch_size, shuffle=False)

with torch.no_grad():
  with tqdm.tqdm(dataloader_eval, desc="Generating output vectors...") as pbar:
    for (images, _) in pbar:
      images = images.to(device)
      embeddings = model(images)
      results.append(embeddings.to('cpu'))

# Took about 17m50s for 1583 batches (202599 input faces)

# Generate a single tensor
face_vectors = torch.cat(results, dim=0)

# See how much data we got :0
# Each vector has 512 dimensions.
print(face_vectors.shape)

# Store this tensor
torch.save(face_vectors, "face_vectors.pt")

torch.Size([202599, 512])

After training for a bit, we hopefully have a model that is strong at distinguishing different faces! (I only trained on 1/3 of the dataset, but it’s probably fine… the dataset is big and this is just for experiment.)

I want to see what the most salient features are for this model, too. How can we do this?

A page from Stanford CS231n shares a nice overview of ways to evaluate CNNs. One blog post shares a suite of ways to visualize convolutional neural network predictions. We can:

• Visualize maximally-activating patches (i.e., see what inputs for the receptive field of a given output cause the highest activation). We do this by inputting a bunch of images and see which ones are the most strongly correlated with certain features.
• Occlusion and gradient-based experiments are useful for identifying what contributed the most to an output. Examples include SHAP, GradCAM, and Integrated Gradients.
  • We can look at the gradient of a classification w.r.t a specific part of the input image, or we can occlude parts of the image and see how our classification changes.
• Deconv/Guided backprop are ways to synthetically generate maximally-activating patches. Basically, we are causing the model to hallucinate a given response, and inspecting what input it’s “imagining”. One popular example is DeepDream. See Peeking Inside ConvNets for a page with nice examples.
  • On a technical level, what we’re doing is using the same technique for optimizing model weights against a loss function to optimize an image input against a class activation. We start with an image and perform gradient descent, using the class activation as the objective (so $loss = -activation$).
• Embeddings: We can use methods like t-SNE to reduce the dimensionality of input images. Essentially, output vectors (which can be in a massive $512$-dimension vector space, etc.) are converted to 2D vectors, in a way that keeps neighboring vectors in the larger space close together in the 2D space. This gives a nice at-a-glance way to visualize what images the network perceives as similar to each other.

With respect to language models, Anthropic introduces a Sparse Autoencoder, which takes the set of output features (which form a densely-packed vector space of real-valued dimensions) and outputs a large set of disrete features. This is similar to work done by OpenAI to discover multimodal neurons.

• The sparse autoencoder is used specifically in language models. The goal of the sparse autoencoder is to discover features with monosemanticity. Because the inner workings of language models are not regularized at all, the intermediate vector space can look like a mess, as long as the optimizer knows how to improve predictions given more input. Although it makes optimization simply “work”, it’s bad for us because we don’t know what anything means. Additionally, because we can assume that most words don’t have all possible semantic meanings at once, some sort of compression is happening, resulting in an overpopulated basis. What this means is that $1024$ binary features might be stored in a $512$-dimensional real-valued vector. Why is this bad? Because we can’t point to any one feature and say, “this is the poetry detector”. What the sparse autoencoder does is take the range of internal representations the model creates across a wide variety of inputs, and translates it to a vector space where each dimension of the vector corresponds to a different feature.

I realized I forgot to store the output vectors of the model. It’s fine; using torch.no_grad() for non-training tasks makes compute take less time. Also, this is just for experimentation.

The question is now, will the model identify my face separately from someone elses? (And what features might the model be using?)

Testing Methods

We’ll now try various attribution methods to understand the features of our model. I’ll use the sparse autoencoder proposed by Anthropic for this. The method resembles a standard linear encoder and decoder, with the difference that decoder bias is subtracted from the input before being encoded. \(\overline{x} = x - b_d \\ f = \text{ReLU}(W_e\overline{x} + b_e) \\ \hat{x} = W_df+b_d \\ L = \frac{1}{|X|} \sum_{x \in X} \|x - \hat{x}\|_2^2 + \lambda\|\textbf{f}\|_1\)

The loss function $L$ decomposes into:

• A reconstruction loss (input vs. reconstructed input; MSE)
• A regularization loss (L1 norm of sparse features)

class SparseAutoencoder(torch.nn.Module):
  def __init__(self, input_feature_dim, sparse_feature_dim):
    super().__init__() # necessary for any torch.nn.Module subclass

    self.weight_encoder = torch.nn.Parameter(torch.zeros((input_feature_dim, sparse_feature_dim), dtype=torch.float32))
    self.weight_decoder = torch.nn.Parameter(torch.zeros((sparse_feature_dim, input_feature_dim), dtype=torch.float32))
    self.bias_decoder = torch.nn.Parameter(torch.zeros(input_feature_dim, dtype=torch.float32))
    self.bias_encoder = torch.nn.Parameter(torch.zeros(sparse_feature_dim, dtype=torch.float32))
    # initialize weight matrices
    torch.nn.init.normal_(self.weight_encoder)
    torch.nn.init.normal_(self.weight_decoder)

  def encode(self, input_features):
    # (n, input_feature_dim)
    xbar = input_features - self.bias_decoder
    # (n, sparse_feature_dim)
    f = torch.nn.functional.relu(xbar @ self.weight_encoder + self.bias_encoder)
    return f

  def decode(self, sparse_features):
    # (n, sparse_feature_dim)
    xhat = sparse_features @ self.weight_decoder + self.bias_decoder
    return xhat

def sparse_autoencoder_loss(model, input_features, l1_lambda):
  # Assume input is batched: (n, input_feature_dim)
  f = model.encode(input_features)
  x_hat = model.decode(f)
  # L_reconstruction = torch.nn.functional.mse_loss(input_features, x_hat)
  L_reconstruction = torch.norm(input_features - x_hat, 2)
  L_complexity = l1_lambda * torch.norm(f, 1)
  return L_reconstruction, L_complexity

Detecting Important Features

Here, we use the Sparse Autoencoder to detect monosemantic features. (I could have added an L1 loss to the feature vectors during training at minimal overhead…)

We first load the face vectors we saved. (My runtime crashed, lol.)

import PIL.Image

import torch

face_vectors = torch.load("face_vectors.pt")

sparse_feature_dim = 512
sae = SparseAutoencoder(input_feature_dim=512, sparse_feature_dim=sparse_feature_dim).to('cuda')
sae_optim = torch.optim.Adam(sae.parameters(), lr=1e-5, weight_decay=1e-4)

import tqdm

# l_recon = 30 or so.
# Takes like 1 minute train this many epochs, because of CUDA and the small size of the dataset.
for epoch in range(30):
  i = 0
  sae_batch_size = 512
  N = len(face_vectors)
  order = torch.randperm(N)
  with tqdm.tqdm(total=N, desc="Training sparse autoencoder...") as pbar:
    while i < N:
      batch = face_vectors[order[i:i + 512]].to('cuda')
      # L_reconstruction, L_complexity = sparse_autoencoder_loss(sae, torch.nn.functional.normalize(batch) * torch.sqrt(torch.tensor(512.0)), l1_lambda=0.1)
      L_reconstruction, L_complexity = sparse_autoencoder_loss(sae, torch.nn.functional.normalize(batch), l1_lambda=0.1)
      loss_sae = L_reconstruction + L_complexity
      sae_optim.zero_grad()
      loss_sae.backward()
      sae_optim.step()
      pbar.set_postfix({"loss_sae": loss_sae.item(), "recon": L_reconstruction.item(), "compl": L_complexity.item()})
      pbar.update(sae_batch_size)
      i += sae_batch_size

# Save our sparse autoencoder.
torch.save(sae.state_dict(), "sparse_autoencoder_weights.pt")

Generate features for each image

Now that we’re done training, let’s infer the feature set for each image.

# Let's sort images by their activation of certain neurons.
sparse_features = []

with torch.no_grad():
  i = 0
  sae_batch_size = 512
  N = len(face_vectors)
  with tqdm.tqdm(total=N, desc="Extracting sparse features...") as pbar:
    while i < N:
      batch = face_vectors[i:i + 512].to('cuda')
      sparse_features.append(sae.encode(batch).to('cpu'))
      i += sae_batch_size
      pbar.update(sae_batch_size)

feats = torch.cat(sparse_features, dim=0)

Extracting sparse features...: 202752it [00:00, 419656.01it/s]                            

# Save our sparse features.
torch.save(feats, "face_vectors_sparse.pt")

Inspecting Results

The sparse autoencoder has now been trained. Now, let’s look at which images maximally activate each feature. I’ll normalize the features using torch.nn.functional.normalize before sorting so we can look at which images have the strongest relative activation for each feature.

Result Without Sparse Autoencoder

Let’s see if the initial face vector (i.e., before being encoded) has interpretable feature representations.

feats_dense_normed = torch.nn.functional.normalize(face_vectors)

for feature_id in range(16):
  # Sort images by feature `feature_id`.
  order = torch.argsort(feats_dense_normed[:, feature_id], descending=True)
  visualize_4x4(
      [dataset[order[i]][0] for i in range(16)],
      title="Maximally-Activating Examples",
  )

/usr/local/lib/python3.10/dist-packages/torchvision/transforms/functional.py:1603: UserWarning: The default value of the antialias parameter of all the resizing transforms (Resize(), RandomResizedCrop(), etc.) will change from None to True in v0.17, in order to be consistent across the PIL and Tensor backends. To suppress this warning, directly pass antialias=True (recommended, future default), antialias=None (current default, which means False for Tensors and True for PIL), or antialias=False (only works on Tensors - PIL will still use antialiasing). This also applies if you are using the inference transforms from the models weights: update the call to weights.transforms(antialias=True).
  warnings.warn(

Result With Sparse Autoencoder

feats_normed = torch.nn.functional.normalize(feats)

for feature_id in range(16):
  # Sort images by feature `feature_id`.
  order = torch.argsort(feats_normed[:, feature_id], descending=True)
  visualize_4x4(
      [dataset[order[i]][0] for i in range(16)],
      title="Maximally-Activating Examples",
  )

/usr/local/lib/python3.10/dist-packages/torchvision/transforms/functional.py:1603: UserWarning: The default value of the antialias parameter of all the resizing transforms (Resize(), RandomResizedCrop(), etc.) will change from None to True in v0.17, in order to be consistent across the PIL and Tensor backends. To suppress this warning, directly pass antialias=True (recommended, future default), antialias=None (current default, which means False for Tensors and True for PIL), or antialias=False (only works on Tensors - PIL will still use antialiasing). This also applies if you are using the inference transforms from the models weights: update the call to weights.transforms(antialias=True).
  warnings.warn(

Unbiased Blind Interpretability Test

To see whether there is a measurable difference between the feature encoders, I’ll take a few features I haven’t seen before and rate each feature for its explainability. Then, I’ll see whether I categorized sparse autoencoder features or regular features as being more or less interpretable.

import random
import time

test_features = [*range(16, 32)]
orders = []

for feature_id in test_features:
  # Sort images by feature `feature_id`.
  order_dense = torch.argsort(feats_dense_normed[:, feature_id], descending=True)
  order_sparse = torch.argsort(feats_normed[:, feature_id], descending=True)
  orders.append((order_dense, "dense:" + str(feature_id)))
  orders.append((order_sparse, "sparse:" + str(feature_id)))

random.shuffle(orders)
results = []
for (order, feature_id) in orders:
  visualize_4x4(
      [dataset[order[i]][0] for i in range(16)],
      title="Maximally-Activating Examples",
  )
  time.sleep(0.5)
  rating = int(input("How visually-consistent do images for this feature appear? (1 - 7)\n"))
  results.append((feature_id, rating))

How visually-consistent do images for this feature appear? (1 - 7)
5

…

How visually-consistent do images for this feature appear? (1 - 7)
5

# Let's see how the ratings turned out.

ratings = {"sparse": [], "dense": []}
for (feature_id, rating) in results:
  vector_type = feature_id.split(":")[0]
  ratings[vector_type].append(rating)

# Create a figure and axis
fig, ax = plt.subplots()

# Create histograms
n, bins, patches = ax.hist([ratings['sparse'], ratings['dense']], bins=[1, 2, 3, 4, 5, 6, 7], label=["Sparse Features", "Dense Features"], alpha=0.5)

# Calculate the positions for centered ticks
tick_positions = 0.5 * (bins[:-1] + bins[1:])
tick_labels = [str(int(tick)) for tick in tick_positions]

# Set the x-tick positions and labels
ax.set_xticks(tick_positions)
ax.set_xticklabels(tick_labels)

# Set labels and legend
ax.set_xlabel("Interpretability Rating 1-7")
ax.set_title("Human-Annotated Blind Interpretability Ratings")
ax.legend()

# Show the plot
plt.show()

# Let's get the mean and standard deviation of these ratings through a Pandas summary.
# We'll also do a t-test to see if sparse features are truly more explainable than dense features.
import pandas as pd
from scipy.stats import ttest_ind

display(pd.DataFrame(ratings).describe())

ttest_result = ttest_ind(ratings['sparse'], ratings['dense'])
print(f"t-stat: {ttest_result.statistic:.4f}, pvalue: {ttest_result.pvalue:.4f}")

t-stat: 2.2486, pvalue: 0.0320

Assessment of Results

We can see that some features correspond to the background color of the images (because that was useful when the model was trained to disambiguate between images). We also see that some features correspond specifically to women with blond hair, or images of celebrities in front of specifically dark / specifically light backgrounds. Some neurons are less interpretable and may be combinations of more subtle features discovered during the disambiguation step.

Overall, this analysis provides insight into the types of images that neural networks perceive as similar to each other, and what components each feature vector is made of.

The sparse autoencoder is a lightweight and useful way to uncover monosemantic features, as backed by a t-test. This is simply based on the incorporation of an L1-loss on the sparse feature set.

Ignore

This is me verifying that the ImageDataset loaded images in alphabetical order, for reproducibility after closing this notebook.

visualize_4x4([dataset[267][0]] * 16)

/usr/local/lib/python3.10/dist-packages/torchvision/transforms/functional.py:1603: UserWarning: The default value of the antialias parameter of all the resizing transforms (Resize(), RandomResizedCrop(), etc.) will change from None to True in v0.17, in order to be consistent across the PIL and Tensor backends. To suppress this warning, directly pass antialias=True (recommended, future default), antialias=None (current default, which means False for Tensors and True for PIL), or antialias=False (only works on Tensors - PIL will still use antialiasing). This also applies if you are using the inference transforms from the models weights: update the call to weights.transforms(antialias=True).
  warnings.warn(

import PIL.Image
PIL.Image.open("data/img_align_celeba/img_align_celeba/000268.jpg")

Sound in Augmented Reality

My goal was to create a portable device to visually display which objects were creating sounds, and I accomplished it, measuring the angular accuracy to be within 15 degrees.

Initial Research and Development

I started by researching how humans evolved to locate sounds to find a source of inspiration for a way to develop a computerized algorithm. My main finding was that humans used the time difference of arrival (TDOA) of sounds reaching each ear to synthesize where sounds were coming from. TDOA is based on the fact that sounds travel at a constant speed, which means that the time delay between two microphones’ detection of a sound is proportional to the two microphones’ distance from the sound. My initial plan was to use an array of microphones as “ears” and compare when any given sound arrived at each microphone.

I also gathered an understanding of how sounds are represented to computers as waveforms. Microphones record data by detecting vibrations in the air and converting these continuous sensor values to discrete binary values (such as from 0 to 65535, -32768 to 32767, or 0 to 256.) The microphone captures the current state of vibration of the air as a frame. By putting a lot of these frames together, microphones create an image of the wave that reached it. When there’s a sudden sound, there’s a sudden vibration in the air that’s captured by the image of the wave (hereafter referred to as the “waveform”.)

It would have been helpful if I had researched already-existing computerized algorithms for sound localization, because I spent a little over a month trying to apply the time difference of arrival concept myself before discovering the widely-used SRP-PHAT algorithm. I recommend that anyone developing an algorithm for their senior research project implement an existing one before trying to design their own.

Time Difference of Arrival

My first attempt at an algorithm was to estimate the time difference of arrival by detecting the time differences between sharp jumps in microphone signals. I used the AV16.3 dataset, which contained recordings of people speaking around an array of eight microphones, to experiment with this algorithm’s potential accuracy. I found out that sharp jumps in a microphone signal can be ambiguous and lead to inconsistent time delays. The following is a sample of a recording from the AV16.3 dataset.

Machine Learning-based Approach

I then tried to create an ML-based algorithm. I found a paper from Amazon’s voice assistant division about how machine learning was used to locate the direction of a user’s voice – they used a convolutional neural network which outputted a “coarse location” where they expected a sound to be vaguely present, and a “fine location” which specified the exact location within that course region. The paper can be found here. I implemented the paper to the best of my ability, but did not have enough data to train with. I tried to simplify the problem by only outputting the coarse location rather than the fine location (which is essentially predicting which microphone detected the sound first.) I achieved a cross-entropy loss of 0.39. Even though it’s better than guessing, it was not particularly accurate.

I realized that I had not verified that the dataset was balanced. I rebalanced the dataset by creating seven identical copies with the sounds being rotated between the microphones. Upon doing this, my loss jumped. This indicated that the previous minor success was due to an imbalanced dataset.

This part of the project taught me to read the training procedure of the paper carefully. I found out that the Amazon paper trained their model on synthetic data and fine-tuned it on AV16.3, rather than only relying on real-life data. This was important because synthetic data allowed them to collect more data than would have been feasible within a short time frame in real life. I recommend looking into synthetic data when training your own models. I spent a lot of time writing boilerplate code that was very specific to the dataset I was using, which resulted in a lot of effort spent on code that was unreusable.

In retrospect, I probably had the time to generate synthetic data and train an accurate ML model, which would have led to some interesting research on the potential of ML for this problem. However, I still found ways to locate sounds effectively that I was able to integrate into my device.

Steered-Response Power Phase Transform (SRP-PHAT)

I still had the AV16.3 dataset code from before, so I decided to test it with a library called ODAS (“Open EmbeddeD Audition System”) (link here). This library was written entirely in C and had significant gaps in its documentation, such as how to customize the audio localization process. To figure out how it worked, I looked at the code itself, and reasoned what parameters I should use with the dataset. I don’t recommend using this library because I later found out that the same functionality can be implemented in Python with ~300 lines of code, and the ability to do sound localization “out of the box” doesn’t outweigh the difficulties in setting it up. The algorithms used by ODAS are an extension of an algorithm called Steered-Response Power Phase Transform (SRP-PHAT) (original paper here).

Instead of calculating when each sound arrives and triangulating the location of the source, SRP-PHAT divides space into a grid and calculates how likely a sound was to have arrived from there. Calculating this likelihood is based on the concept of time delay. If, hypothetically, there were a sound at a location, we would know how long it would take to reach each microphone. This is because we choose the location to check, and therefore we know the distance from that position to each of the microphones. The time delay that we would expect is the distance divided by the speed of sound. Therefore, we expect similar features to appear in the audio signals of each microphone at slightly different times.

These two diagrams demonstrate how it works. We set up two microphones at a distance d apart. If we are testing a location to the right of the microphone setup, then we expect the sound to reach the right microphone (d/Speed of Sound) seconds earlier than the sound reaches the left microphone. Therefore, whatever sound occurs at time t = 0 on the right microphone would appear on the left microphone at time t = d/Speed of Sound.

Example of phase offset during SRP-PHAT:

The left microphone’s signal is shifted backward

The black box shows the regions of interest (ROI) for each microphone. We can line up the ROIs by shifting the waveforms from each microphone. Now, we want a way to measure how well the sound from the right microphone matches with the sound in the left microphone signal’s ROI. To measure this correlation, we could consider using cross-correlation.

Cross-correlation takes the average of the product of the intensities between each microphone. Therefore, if both microphones have high intensity or both microphones have low intensity at the same time (i.e., their signals are strongly correlated), the cross-correlation will be high. Cross-correlation works for pairs of microphones, but when we use a microphone array, we will have more than two microphones, and we still want to be able to measure how well the signals between the microphones line up.

This is why we use beamforming instead of cross-correlation. The beamforming algorithm shifts the ROIs of each microphone (a “phase transform”) based on their expected time delays (“steered response power” towards a particular location) and adds them together. This produces a new waveform, which is the sum of several other waveforms. Wherever all the microphones have a high intensity, the intensity of the new waveform will be high. However, if some microphones have high intensity (e.g. 5) and some have low intensity (e.g. -5) then the new waveform will have an intensity that cancels out. Therefore, beamforming allows us to capture the underlying correlation between several microphone signals. The overall quality of a given sound location is calculated by taking the decibel level of the added-together sound. This is the method that ODAS uses for sound localization.

Code to perform beamforming for one grid square xyz:

# For each microphone:
for mic in range(num_waveforms):
   # Calculate the expected time delay
   delay = int(get_time_delay_of_arrival_in_seconds(
       mic_locations[mic], xyz) * sample_rate)

   # If delay < 0, then we should be playing future sounds earlier.
   # If delay > 0, then we should be playing past sounds later.
   # Shift the waveform accordingly, and add to the summed wave.
   if delay < 0:
       summed_waveforms[:-delay] += waveform[mic][delay:]
   else:  # Include delay == 0
       summed_waveforms[delay:] += waveform[mic][:-delay]

average_beam = summed_waveforms / num_waveforms

First, the time delay for each microphone is calculated. Then, the signals from each microphone are shifted by that delay and summed together.

All the components of SRP-PHAT have now been covered. Steered Response Power refers to the act of choosing ROIs that are based on the time delay to each microphone, and the Phase Transform refers to the shifting of the waveforms to make the ROIs line up.

ODAS applies this algorithm over a sparse set of points at a distance of 1 meter from the microphone array, resulting in a hemisphere-like point cloud. The points with the highest intensity have a neighborhood of more precise points sampled. Then, the points in this cloud with the highest intensity are output as JSON.

Making the ODAS configuration work involved writing the geometry of the microphone array used in the AV16.3 dataset to the ODAS config file. To compile ODAS, I needed to install CMake and build ODAS with a few commands: First, creating a “build/” folder, then doing “cd build/”, then initializing CMake with “cmake ..”, and finally building the project with “make install”. This outputted two binary files: “odaslive” and “odasserver.” I only used odaslive, which is a command-line program which could take a config file, read sound data from an audio file, and print the sound locations to standard output or write them to a TCP connection (a socket.) These sound locations are written as JSON objects separated by newlines. In order to compile ODAS, I used the Windows Subsystem for Linux (which simulates Linux on the Windows filesystem) because ODAS uses some low-level functions specific to the Linux operating system. I was also unable to recompile ODAS when I upgraded my Windows computer to a Macbook, because some libraries were not yet supported by Apple’s new chip architecture. I elaborate more on this later, but the result was that I needed to recreate the ODAS library from scratch.

ODAS also came with a visualizer, which required installing Node.js and running the program with “npx electron .” in the directory containing it. The visualizer showed that for the AV16.3 dataset, audio source location could be located along a single axis, but not further than that. Around this point of my development, the hardware arrived. The GitHub repository for the visualizer can be found at https://github.com/introlab/odas_web.

Hardware

I wanted to experiment with real data, so I found the MATRIX Voice microphone array board. I submitted a hardware order on October 12, and it arrived on November 3. I would highly recommend submitting hardware orders earlier than this.

When the MATRIX board finally arrived, I connected it to a Raspberry Pi. Reading any data of any kind did not work. I followed official demonstrations from their website, using the same Raspberry Pi version (v3) and running the same code, but no recordings had any audio features. The board also came with thermometers and accelerometers, and the data from those was “0.” It did not seem to be a problem with the board’s connectivity, because writing data to the board’s built-in LEDs worked.

I tried using other utilities like arecord (a command line tool) and Audacity to record audio data, but none of them were able to detect the board. I factory-reset the Pi and repeated the same process without success. I also purchased another MATRIX board and followed the same steps, along with a factory reset, but was unable to read any data. I considered purchasing another microphone array, so I submitted a hardware order for the ReSpeaker 6-mic circular array kit.

Because my goal was to work with real-world data, I wrote a test algorithm with an Arduino in the two weeks it took for the new board to arrive. I set up an oscilloscope in my basement to read data from Arduino microphones and compare the signals. I used a speaker to play sounds from several locations and exported the recorded signals from the oscilloscope to process them on my computer. The algorithm in this test was to calculate the expected time delay for each microphone and use the cross-correlation between pairs of microphones to calculate a relative intensity level. I implemented this in xcorrmap/xcorrmap.py.

The ReSpeaker board was functional, and I could read its data with a library called PyAudio. There were several helpful demos on the ReSpeaker Wiki:

https://wiki.seeedstudio.com/ReSpeaker_6-Mic_Circular_Array_kit_for_Raspberry_Pi.

Both the ReSpeaker and MATRIX Voice boards connected to the Pi through the General-Purpose Input/Output (GPIO) pins. This meant that the array “stacked” on top of the Pi.

I revisited ODAS, the sound localization library from before. Because the configuration files allowed the output format to be a “socket”, I installed ODAS on the Raspberry Pi and had it output sound locations to an ODAS visualizer on the same device. I soon realized, however, that ODAS was not fast enough to run on the Pi in real time. To achieve this, I offloaded the processing to my own computer, which meant I would need to set up a connection over the network.

The ODAS visualizer. Each blue dot indicates a high likelihood that a sound is coming from that direction.

Reading Raw Audio Data

I needed a way to record sound data and send it over the network, so I wrote a Python program that read data with PyAudio and established a socket server that I could use to send data to outside servers. This involved creating a TCP server based on the specifications from the Python documentation: https://docs.python.org/3/library/socket.html.

Sending raw sound data required me to learn how sound data is stored. In this case, each “frame” of sound from a single waveform is stored as two bytes: a signed 16-bit integer. There are 8 channels on the ReSpeaker 6-Mic Array, two of which are used for audio playback and can be ignored. When reading from a device with PyAudio, the corresponding waveforms for all channels – whether for recording or playback – are returned. Data is read in chunks, where the parallel streams of data from each channel are serialized in the format “[channel 1 bytes][channel 2 bytes][…][channel 8 bytes]”. The two bytes for each channel store one form of a digitized version of the sound wave. Code for this is in the main() function of vox_server.py.

Sometimes, the channels for the ReSpeaker array would go out of order. Let’s say that the original channel order was (1, 2, 3, 4, 5, 6, 7, 8). Sometimes, the data read under channels (1, 2, 3, 4, 5, 6, 7, 8) would actually be the data corresponding to channels (3, 4, 5, 6, 7, 8, 1, 2), or some other cyclic permutation by 2, 4, or 6 channels. I was able to deal with this because channels 7 and 8 are reserved for playback, so they are always silent when recording audio. At runtime, I would determine which channels had received a sound, and which had remained silent, and shift the order of the channels such that the silent ones were at indexes 7 and 8. After testing (once I had implemented the sound localization algorithm), I found that the order of the input channels was always preserved if the indexes of the output channels were known (channel_permutation.py).

Cloud Data Processing

Once I was able to read and send raw audio data, I needed to be able to consume it on the processing node. Optimally, I would be able to send sound data directly to a server running ODAS, but ODAS had no support for streaming data from a network connection. So, I learned how to use sockets from the Linux documentation and implemented the feature myself. It was very difficult to navigate ODAS’s codebase, as it was written in C, but the knowledge from working with C++ in Computer Vision at TJ transferred well. I made a new kind of sound input called “socket”, which connected to a device given its IP address and inserted the received data into a queue to be further processed. Because this feature integrated with the rest of ODAS, I was able to use the ODAS visualizer to test whether it was receiving the sounds correctly. Fortunately, it was – during testing, the localizations shown in the ODAS visualizer were consistent with where I held a sound source. I now had a way to locate sounds in real time!

Once I verified that this was working, I submitted a pull request to the ODAS GitHub repository. The library authors reviewed the code and approved it. (I was pleasantly surprised that something from this project made it into a real library!)

Visualizing the Data

My overall goal was to create a method to view the indicators in augmented reality, such that indicators would align visually with the objects making the sound. I initially looked at devices that would enable me to create an augmented reality headset with an iPhone by having the screen reflect from a heads-up visor, but I instead prioritized creating a web client because it required less overhead and allowed me to test my algorithm sooner.

Web Client with Reactive 3D Rendering

I first attempted to render the sound locations as 3D indicators in a VR-style environment. I chose to do this because I was comfortable with libraries that I knew I would be able to accomplish this with (React and THREE.js), the web is supported by almost every platform (including mobile phones), and the web is a very convenient way to create user interfaces without significant overhead code. React is a Javascript library that makes dynamic web applications easier to create and less bug-prone by handling all user interface changes behind-the-scenes. It allows the developer to specify the layout of the UI in a syntax that resembles HTML called JSX, which can include dynamic components that automatically respond to changes in the underlying data. I created a React component to render sound locations as spherical indicators in a 3D scene, and I took advantage of the library automatically updating the scene.

However, the direct creation of TCP sockets is not supported by most major web browsers; they only support a protocol called WebSocket. I solved this problem by creating a server that supported both WebSocket and TCP sockets and translated between the two.

Sound data from each microphone was first serialized by PyAudio and sent in chunks to ODAS. ODAS then outputted a list of sound locations to a broadcasting server, which translated them into the WebSocket protocol. Then, they were read by the web client and rendered onto a 3D VR scene by THREE.js and React.
THREE.js includes strong support for virtual reality. When using a mobile device, I was able to control the rotation of the camera based on the orientation of the device with DeviceOrientationControls from an extension of THREE.js called [drei](http://github.com/pmndrs/drei). This enabled me to create a more consistent sense of perspective. I used the library [react-three-fiber](https://github.com/pmndrs/react-three-fiber) to extend React’s dynamic UI support to THREE.js scenes, and allow me to quickly prototype VR environments by using an HTML-like syntax. Each sound location was rendered as a pink sphere, with the opacity being based on the intensity of the sound detected at that location.

<Canvas
    style=
    camera={camera}
    ref={canvasRef}
>
    <Grid />

    <Localizations localizations={localizations.items} />

    <Gate active={mobile}>
        <DeviceOrientationControls camera={camera} />
    </Gate>
    <Gate active={!mobile}>
        <OrbitControls camera={camera} />
    </Gate>

    <pointLight position={[10, 10, 10]} />
    <ambientLight />
</Canvas>

This code sets up a 3D scene in THREE.js

To use `DeviceOrientationControls`, Safari requires the server to provide an SSL certificate. I found it most convenient to use a [CloudFlare tunnel](https://developers.cloudflare.com/cloudflare-one/tutorials/single-command/) at this part of the development stage, because it allowed me to host the server locally and have CloudFlare manage the SSL for me. However, this came at the expense of some of the user experience, because I now needed to type in the IP of the broadcasting server manually.

Algorithmic Improvements

Non-Maximum Suppression

The output from the SRP-PHAT algorithm is a cloud of points, with each point having a corresponding intensity. I performed a threshold to ensure that only points with a relative intensity level above 25% were rendered. While this filtered out random noise, there were still clusters of points where the sound localization algorithm was uncertain of the true maximum location.

To make the localizations look cleaner, I implemented a non-maximum suppression algorithm. This algorithm reduces clusters of points into the member that has the maximum intensity. It is more effective than calculating the overall maximum because it allows for multiple sound sources. This simplification is made by iterating through the points and filtering out points that are not the most intense of the points within a certain Euclidean distance r. Therefore, local maxima are kept in, and all local maxima enforced to be at least r away from each other. The algorithm is computationally inexpensive enough to run in the browser in real time.

Before and after applying non-maximum suppression.

The following is a code snippet that performs non-maximum suppression.

def get_non_maximum_suppression_indexes(locations, energies, deduplication_radius):
   unsuppressed_indexes = []

   for index in range(len(locations)):
       location = locations[index]
       energy = energies[index]

       for other_index in range(len(locations)):
           if index == other_index:
               continue

           other_location = locations[other_index]
           other_energy = energies[other_index]

           if _distance(location, other_location) <= deduplication_radius:
               if other_energy > energy or (other_energy == energy and other_index > index):
                   break
       else:
           # Runs if the loop was not broken out of
           unsuppressed_indexes.append(index)

   return unsuppressed_indexes

Frequency-Based Filtering and Ambient Noise Removal

Due to ambient noise in the syslab, high-pitched sounds like my voice would only be detected when they were louder than the fans in the background. To counteract this effect, I created a way to isolate sounds within certain frequency ranges and prevent them from interfering with the detection of other sounds. I converted the waveforms from each microphone into 16 component waveforms, each representing different frequency ranges. For example, one waveform represents all sounds below 125 Hz; the next represents all sounds between 125 Hz and 250 Hz; and so on. Then, I pass these waveforms to the same beamforming algorithm that was used for unfiltered audio.

To create these filtered signals, I modified the “frequency domain” representation of the sound. While sounds can be represented as vibrations in the air as a function of time (the “time domain”), they can also be represented as combinations of several sinusoidal waves at different magnitudes and phases.

Visual representation of converting between the time domain and frequency domain

Transforming into the frequency domain is done with the Fast Fourier Transform (FFT), and transforming into the time domain is done with the Inverse Fast Fourier Transform (IFFT), both of which can be done with the functions stft and istft from the Python library librosa.

# Create a spectrogram.
# This is done by taking the fft of the last fft_window_length frames, skipping
# fft_hop_length frames between steps.
fft_hop_length = 128
fft_frequency_count = 2048
fft_window_length = fft_frequency_count
# Perform sliding Fourier Transform
spectrogram = librosa.stft(
    waveform,
    n_fft=fft_frequency_count,
    hop_length=fft_hop_length,
    win_length=fft_window_length
)

To isolate sounds within a given frequency range for each microphone’s audio data, I transformed their waveforms into their frequency domain representations, set the magnitudes of the frequencies outside of the desired range to 0, and transformed the sounds back to the time domain.

for bucket_number in range(len(frequency_buckets)):
    # filter_frequencies sets the magnitudes of the frequencies between start_frequency_index and end_frequency_index to 0.
    start_frequency_index = int(
        frequency_buckets[bucket_number][0] * 1025 / 2048)
    end_frequency_index = int(
        frequency_buckets[bucket_number][1] * 1025 / 2048)
    masked_spectrogram = filter_frequencies(
        spectrogram, start_frequency_index, end_frequency_index)

    # Normalize magnitudes to only consider phase offsets of each frequency
    # Not necessary but improves accuracy
    masked_spectrogram_phases = masked_spectrogram / \
        (np.abs(spectrogram) + 0.00001)

    # Reconstructs the time domain signal from the frequency domain signal
    filtered_signal = librosa.istft(
        masked_spectrogram_phases,
        n_fft=fft_frequency_count,
        hop_length=fft_hop_length,
        win_length=fft_window_length
    )

Now that I was able to locate sounds between 0 and 125 Hz, 125 Hz and 250 Hz, and 14 other ranges of frequencies, I could represent sounds differently depending on what their frequency was. First, I could set custom thresholds. I added a step at the beginning of processing that would measure the ambient noise level of each frequency range. Then, the intensity requirement for a sound to be shown on-screen would be calculated based on that. Therefore, frequency ranges with more ambient noise, like 0 to 125 Hz, can be thresholded more strictly than frequency ranges with quieter sounds like a person’s voice. Another benefit was that if there were a high-pitched sound and a low-pitched sound at the same time, they would be more easily isolated and located independently. This is important when a microphone array only has six microphones, compared to the commercial standards of 96 or more.

I chose to replace ODAS by implementing my own version of the SRP-PHAT algorithm in Python (written in test_spectrogram.py). Part of this was inspired by the fact that I had recently upgraded my computer to a new Macbook, and the Macbook had a chip architecture that was unsupported by some libraries that were needed by it. I tried to compile those libraries directly for the architecture, but after a few days of attempts, I found it would be faster to implement it independently. ODAS also only supported one microphone array per process, and I wanted to perform localization on the 16 frequency ranges separately. To integrate this algorithm to my pipeline, I made a Python socket client (vox_client.py) which would connect to the Pi and read a stream of microphone data. Whenever the socket client received microphone data, it would be added to a queue where a separate processing thread would evaluate it. After evaluation, the detected sound locations would be added to another queue, where another thread would send it to the broadcasting server. I divided work into threads because some components used blocking I/O, and the program would become too laggy if it waited for blocking I/O operations to complete before starting to process the most recent sound data.

Evaluation

Upon implementing these improvements, I needed a way to measure the accuracy of the overall algorithm. First, I marked speaker locations at 45-degree increments around the microphone array.

Then, I added a capture option to the web app which would store any sound locations it received to a variable for five seconds before exporting them as a JSON file. Finally, I placed a speaker at each of the locations and captured five seconds of sound localizations from that location. I loaded the sound locations into a post-processing script written in Python, and plotted the sound locations with the matplotlib library. Markers of the same color represent localizations received during the same capture. Each capture was taken with the speaker being at a different location, and they appear to not have significant overlap.

To evaluate these results more rigorously, I took mean and standard deviation of the azimuthal angle (angle around the Z axis) and compared them to the expected angles. I found that the measurements were always within 15 degrees of the true sound source, and that the standard deviations were usually less than 10 degrees.

Integrating the Camera Feed

Up to this point, I had created a sound “camera” that could locate sounds and render them in a 3D web app. However, the 3D web app lacked a sense of perspective, the locations of sound markers relative to the physical world was subjective. To fix this, I needed to add a visual camera to the back of the device. Unfortunately, because of the difficulty in sending a specific video feed to a web app, I was not able to use the web app I built for this representation. Instead, I created a script that ran locally on the Pi, with the intention that it could be used with an LCD display taped to the front of the device and I could be more efficient by avoiding sending video over the network. I implemented the same algorithms of non-maximum suppression and frequency-based filtering on this new script.

I used an Arducam 5MP Raspberry Pi camera I had at home and taped it to the face of the microphone array. I read from the camera with the picamera Python library and rendered a video feed with OpenCV’s imshow. Then, I set up a socket server to send microphone data to the processing node and read sound localizations. Because I was no longer using neither a web app nor ODAS, I also would no longer need a broadcasting server to translate between the TCP and WebSocket protocol, and my overall processing pipeline was vastly simplified to just a server (the Pi) and a client (the processing node.)

To render the sounds in this new environment, I could no longer use the 3D VR libraries from before, so I represented sounds as circles of varying radius and color, which I drew with OpenCV’s circle function. The color denoted the frequency range, with colors with a deeper blue identifying lower-pitched sounds and colors with a deeper red identifying higher-pitched sounds, and the radius denoted the intensity.

Sound frequency range vs. Color of indicator

Each sound was tagged with a 3D location satisfying x^2 + y^2 + z^2 = 1; z ≥ 0, and I calculated where to render them in 2D with a perspective projection, setting the plane of the camera at z=-1.

Photo of indicators (colored circles) rendered near sound sources.

Because the video feed, networking, and audio streaming were all done from the same program, I split each task into separate threads to ensure that lag from one process didn’t interfere with the others. The server (Raspberry Pi) had the following threads:

• A thread to listen for connections from the processing node
• A thread to read camera data with picamera
• A thread to render annotated camera frames with OpenCV
• A thread to read sound location data from the processing node
• A thread to read sound data with PyAudio and send it to the processing node

Some of these threads, like the thread reading sound location data and the thread rendering sound location data, needed to communicate, and I did this through global variables.

Making It Portable

Making the device portable required (1) removing all dependence on external power, and (2) creating a wireless way to display what was going on on the Pi.

I first submitted a hardware order for a Raspberry Pi Lithium-Polymer (Li-Po) battery chip, which could attach to the opposite side of the Pi from the microphone array. Setting this up simply involved connecting an included Li-Po battery with the regulator chip and using a micro-USB cable to connect the regulator chip to the Pi’s power supply port.

Because all the video rendering was now performed on the Pi, and I did not have an LCD screen to display what was running on the Pi, I needed to render the Pi’s content externally. I submitted a hardware order for a longer HDMI cable, but Dr. Gabor recommended that I look for ways to connect to the device wirelessly. I found out that I can stream data from the Raspberry Pi with the VNC protocol – a standard for screencasting and remote control. I enabled VNC in the Raspberry Pi settings and connected the Pi and my computer to the same Wi-Fi network. Then, I could use the Pi remotely from my computer and project its content during a presentation. This requires knowing the IP address of the Pi beforehand.

Final Device

From top to bottom:

• The Arducam 5MP camera. It is connected to the Pi under the microphone array through a ribbon cable.
• The ReSpeaker 6-mic array. There is a microphone at each corner of this hexagonal chip. It connects to the Pi through the GPIO pins.
• The Raspberry Pi.
• The Lithium-Polymer battery chip.

It is still a prototype, and while it can’t fit in a pocket, it can easily be used as a handheld device.

Further Directions for Exploration

This project was largely successful, but there are several avenues that I would encourage students in the future to pursue!

1. Sound classification. If each detected sound could be annotated with a semantic meaning, maybe with a machine learning algorithm, this could be an interesting project.
2. Augmented reality. The current device requires looking at a video feed on a small screen. If the device could be integrated with an augmented reality headset, the sounds would automatically be overlaid with the environment, no camera feed would be needed, and sounds would be significantly more intuitively understood.
3. Speech Recognition. This would tie in very well with the augmented reality track: performing speech recognition on separated sounds could enable the creation of live augmented reality subtitles. This could be especially useful for real-time translation applications.

What I Would Have Done Differently

• I could have saved time and conducted research more deeply if I had been more persistent with analyzing the time difference of arrival, machine learning, and beamforming algorithms in the Initial Research and Development phase.
• I would have benefitted from a clearer long-term goal to aid my decisions and prevent myself from spending too much time off-course, such as in Visualizing the Data and I invested a lot of time into a web client with reactive 3D rendering that I did not end up using. Alternatively, I could have looked for a way to send video from the Pi to the web client. Either way, I ended up implementing core functionality twice, and I could have spent this time instead on improving other parts of the project.

Advice for Future Developers

• The Raspberry Pi is a great way to prototype applications that require hardware unavailable on a phone or laptop. However, it comes with a tradeoff of being harder to set up. It will be handy to have backups of both the software and hardware setup on the Pi, and sometimes, the best solution is to factory reset it. The most common method of using a Raspberry Pi is with a keyboard, mouse, and monitor connected over HDMI, but for streaming data from the Raspberry Pi wirelessly, it can be very helpful to investigate VNC software, like VNC Viewer.
• If you’re using Windows, try the Windows Subsystem for Linux (WSL). Linux is much more developer-friendly because of the strong ecosystem of tools and libraries available. Otherwise, you may need to figure out on your own how to compile a given library.