state_estimation_2d.py

# Copyright (c) Meta Platforms, Inc. and affiliates.
#
# This source code is licensed under the MIT license found in the
# LICENSE file in the root directory of this source tree.

from typing import List

import matplotlib.pyplot as plt
import torch
import torch.nn as nn
import torch.nn.functional as F

import theseus as th

device = "cpu"
torch.manual_seed(0)
path_length = 50
state_size = 2
batch_size = 4
learning_method = "leo"  # "default", "leo"

vis_flag = True
plt.ion()


# --------------------------------------------------- #
# --------------------- Utilities ------------------- #
# --------------------------------------------------- #
def plot_path(optimizer_path, groundtruth_path):
    plt.cla()
    plt.gca().axis("equal")

    plt.xlim(-250, 250)
    plt.ylim(-100, 400)

    batch_idx = 0
    plt.plot(
        optimizer_path[batch_idx, :, 0],
        optimizer_path[batch_idx, :, 1],
        linewidth=2,
        linestyle="-",
        color="tab:orange",
        label="optimizer",
    )
    plt.plot(
        groundtruth_path[batch_idx, :, 0],
        groundtruth_path[batch_idx, :, 1],
        linewidth=2,
        linestyle="-",
        color="tab:green",
        label="groundtruth",
    )

    plt.show()
    plt.pause(1e-12)


def generate_path_data(
    batch_size_,
    num_measurements_,
    generator=None,
):
    vel_ = torch.ones(batch_size, 2)
    path_ = [torch.zeros(batch_size, 2)]
    for _ in range(1, num_measurements_):
        new_state_ = path_[-1] + vel_
        path_.append(new_state_)
        vel_ += 0.75 * torch.randn(batch_size, 2, generator=generator)
    return path_


class SimpleNN(nn.Module):
    def __init__(self, in_size, out_size, hid_size=30, use_offset=False):
        super().__init__()
        self.fc = nn.Sequential(
            nn.Linear(in_size, hid_size),
            nn.ReLU(),
            nn.Linear(hid_size, hid_size),
            nn.ReLU(),
            nn.Linear(hid_size, out_size),
        )

    def forward(self, state_):
        return self.fc(state_)


def run_model(
    mode_,
    cost_weights_model_,
    current_inputs_,
    path_length_,
    print_stuff=False,
):
    weights_ = get_weights_dict_from_model(
        mode_,
        cost_weights_model_,
        current_inputs_,
        path_length_,
        print_stuff=print_stuff,
    )
    theseus_inputs_ = {}
    theseus_inputs_.update(current_inputs_)
    theseus_inputs_.update(weights_)

    return theseus_inputs_


def get_weights_dict_from_model(
    mode_, cost_weights_model_, values_, path_length_, print_stuff=False
):
    if mode_ == "constant":
        weights_dict = {}
        unique_weight_ = cost_weights_model_()
        for i in range(path_length_):
            weights_dict[f"scale_gps_{i}"] = unique_weight_[0, 0].view(1, 1)
            if i < path_length_ - 1:
                weights_dict[f"scale_between_{i}"] = unique_weight_[0, 1].view(1, 1)
    else:
        weights_dict = {}
        all_states_ = []
        # will compute weight for all cost weights in the path
        for i in range(path_length_):
            all_states_.append(values_[f"pose_{i}"])
        model_input_ = torch.cat(all_states_, dim=0)

        weights_ = cost_weights_model_(model_input_)
        for i in range(path_length_):
            weights_dict[f"scale_gps_{i}"] = weights_[i - 1, 0].view(1, 1)
            if i < path_length_ - 1:
                weights_dict[f"scale_between_{i}"] = weights_[i - 1, 1].view(1, 1)

    if print_stuff:
        with torch.no_grad():
            print("scale5", weights_dict["scale_gps_5"].item())
            print("scale45", weights_dict["scale_gps_45"].item())
            print("btwn5", weights_dict["scale_between_5"].item())
            print("btwn45", weights_dict["scale_between_45"].item())

    return weights_dict


def get_initial_inputs(gps_targets_):
    inputs_ = {}
    for i, _ in enumerate(gps_targets_):
        inputs_[f"pose_{i}"] = gps_targets_[i] + torch.randn(1)
    return inputs_


def get_path_from_values(batch_size_, values_, path_length_):
    path = torch.empty(batch_size_, path_length_, 2, device=device)
    for i in range(path_length_):
        path[:, i, :2] = values_[f"pose_{i}"]
    return path


def get_values_from_path(path_):
    """
    :param path_: tensor of dim batch_size_ x path_length_ x 2
    :return: values: dict of (x,y) pos values
    """
    [batch_size_, path_length_, dim] = path_.shape
    values = {}
    for i in range(path_length_):
        values[f"pose_{i}"] = path_[:, i, :2]
    return values


def get_average_sample_cost(x_samples, cost_weights_model, objective, mode_):
    cost_opt = None
    n_samples = x_samples.shape[-1]
    for sidx in range(0, n_samples):
        x_sample_vals = get_values_from_path(
            x_samples[:, :, sidx].reshape(x_samples.shape[0], -1, 2)
        )
        theseus_inputs = run_model(
            mode_,
            cost_weights_model,
            x_sample_vals,
            path_length,
            print_stuff=False,
        )
        objective.update(theseus_inputs)
        if cost_opt is not None:
            cost_opt = cost_opt + torch.sum(objective.error(), dim=1)
        else:
            cost_opt = torch.sum(objective.error(), dim=1)
    cost_opt = cost_opt / n_samples

    return cost_opt


# ------------------------------------------------------------- #
# --------------------------- Learning ------------------------ #
# ------------------------------------------------------------- #
def run_learning(mode_, path_data_, gps_targets_, measurements_):
    # first input is scale for GPS costs, second is scale for Between costs
    if mode_ == "constant":
        model_params = nn.Parameter(10 * torch.rand(1, 2, device=device))

        def cost_weights_model():
            return model_params * torch.ones(1)

        model_optimizer = torch.optim.Adam([model_params], lr=5e-2)
    else:
        cost_weights_model = SimpleNN(state_size, 2, hid_size=100, use_offset=False).to(
            device
        )
        model_optimizer = torch.optim.Adam(
            cost_weights_model.parameters(),
            lr=7e-5,
        )

    # GPS and Between cost weights
    gps_cost_weights = []
    between_cost_weights = []
    for i in range(path_length):
        gps_cost_weights.append(
            th.ScaleCostWeight(th.Variable(torch.ones(1, 1), name=f"scale_gps_{i}"))
        )
        if i < path_length - 1:
            between_cost_weights.append(
                th.ScaleCostWeight(
                    th.Variable(torch.ones(1, 1), name=f"scale_between_{i}")
                )
            )

    # ## Variables
    poses = []
    for i in range(path_length):
        poses.append(th.Point2(name=f"pose_{i}"))

    # ## Cost functions
    cost_functions: List[th.CostFunction] = []

    # ### GPS and between cost functions
    for i in range(path_length):
        cost_functions.append(
            th.Difference(
                poses[i],
                th.Point2(tensor=gps_targets_[i]),
                gps_cost_weights[i],
                name=f"gps_{i}",
            )
        )
        if i < path_length - 1:
            cost_functions.append(
                (
                    th.Between(
                        poses[i],
                        poses[i + 1],
                        th.Point2(tensor=measurements_[i]),
                        between_cost_weights[i],
                        name=f"between_{i}",
                    )
                )
            )

    # # Create Theseus layer and initial values for variables
    objective = th.Objective()
    for cost_function in cost_functions:
        objective.add(cost_function)
    optimizer = th.GaussNewton(
        objective,
        th.CholeskyDenseSolver,
        max_iterations=1,
        step_size=0.9,
    )
    state_estimator = th.TheseusLayer(optimizer)
    state_estimator.to(device)

    # ## Learning loop
    best_loss = 1000.0
    inner_loop_iters = 3
    groundtruth_path = torch.stack(path_data_).permute(1, 0, 2)
    best_solution = None
    losses = []
    for epoch in range(500):
        model_optimizer.zero_grad()

        theseus_inputs = get_initial_inputs(gps_targets_)
        theseus_inputs = run_model(
            mode_,
            cost_weights_model,
            theseus_inputs,
            path_length,
            print_stuff=False,
        )
        objective.update(theseus_inputs)
        with torch.no_grad():
            if epoch % 10 == 0:
                print("Initial error:", objective.error_metric().mean().item())

        for i in range(inner_loop_iters):
            theseus_inputs, _ = state_estimator.forward(
                theseus_inputs,
                optimizer_kwargs={
                    "track_best_solution": True,
                    "verbose": epoch % 10 == 0,
                },
            )
            theseus_inputs = run_model(
                mode_,
                cost_weights_model,
                theseus_inputs,
                path_length,
                print_stuff=epoch % 10 == 0 and i == 0,
            )

        optimizer_path = get_path_from_values(
            objective.batch_size, theseus_inputs, path_length
        )
        mse_loss = F.mse_loss(optimizer_path, groundtruth_path)

        # LEO (Sodhi et al., https://arxiv.org/abs/2108.02274) is a method to learn
        # models end-to-end within second-order optimizers. The main difference is that
        # instead of unrolling the optimizer and minimizing the MSE tracking loss,
        # it uses a NLL energy-based loss that does not backpropagate through the optimizer.
        if learning_method == "leo":
            x_samples = state_estimator.compute_samples(
                optimizer.linear_solver, n_samples=10, temperature=1.0
            )  # batch_size x n_vars x n_samples
            # When x_samples is None, this defaults to a perceptron loss
            # using the mean trajectory solution from the optimizer.
            if x_samples is None:
                x_opt_dict = {key: val.detach() for key, val in theseus_inputs.items()}
                x_samples = get_path_from_values(
                    objective.batch_size, x_opt_dict, path_length
                )
                x_samples = x_samples.reshape(x_samples.shape[0], -1).unsqueeze(
                    -1
                )  # batch_size x n_vars x 1
            cost_opt = get_average_sample_cost(
                x_samples, cost_weights_model, objective, mode_
            )
            x_gt = get_values_from_path(groundtruth_path)
            theseus_inputs_gt = run_model(
                mode_,
                cost_weights_model,
                x_gt,
                path_length,
                print_stuff=False,
            )
            objective.update(theseus_inputs_gt)
            cost_gt = torch.sum(objective.error(), dim=1)
            loss = cost_gt - cost_opt
        else:
            loss = mse_loss

        loss = torch.mean(loss, dim=0)
        loss.backward()
        model_optimizer.step()

        loss_value = loss.item()
        losses.append(loss_value)
        if loss_value < best_loss:
            best_loss = loss_value
            best_solution = optimizer_path.detach()

        if epoch % 10 == 0:
            if vis_flag:
                plot_path(
                    optimizer_path.detach().cpu().numpy(),
                    groundtruth_path.detach().cpu().numpy(),
                )
            print("Loss: ", loss.item())
            print("MSE error: ", mse_loss.item())
            print(f" ---------------- END EPOCH {epoch} -------------- ")

    return best_solution, losses


path_data = generate_path_data(batch_size, 50)
gps_targets = []
measurements = []
for i in range(path_length):
    gps_noise = 0.075 * path_data[i][1].abs() * torch.randn(batch_size, 2)
    gps_target = (path_data[i] + gps_noise).view(batch_size, 2)
    gps_targets.append(gps_target)

    if i < path_length - 1:
        measurement = (path_data[i + 1] - path_data[i]).view(batch_size, 2)
        measurement_noise = 0.005 * torch.randn(batch_size, 2).view(batch_size, 2)
        measurements.append(measurement + measurement_noise)

constant_solution, constant_losses = run_learning(
    "constant", path_data, gps_targets, measurements
)