Trajectory optimization

[jlmaccal]

I'm trying to solve a trajectory optimization problem using gpytorch and botorch.

Imagine I have a trajectory of d-dimensional points x_1, x_2, ..., x_n.

I have a cost function that is evaluated on adjacent points c(x_i, x_{i+1}) = z, where z is a scalar. This cost function is stationary in time, so the same function is used to evaluate c(x_1, x_2) as for any other c(x_i, x_{i+1}). I want to model this function using a GP, which seems straightforward.

What I am having difficulty with is the next part. What I want to do now is to compute the total cost for an entire trajectory as: c(x_1, x_2) + c(x_2, x_3), + .. C(x_{n-1}, x_n). This is also a GP as it's a sum of correlated gaussian random variables. The mean will be the sum of the means of the individual steps and the variance will be the sum over the entire covariance matrix.

However, it isn't clear to me how to best implement this model. Ideally, I would like the resulting model to be treated as a GP, so that I can use analytic acquisition functions. I would also like to make full use of batching for efficiency. Any pointers would be appreciated.

[Balandat]

Interesting use case. We have ScalarizedObjective uses scalarize_posterior under the hood, which is really what you want to do. However, this setup is meant for scalarizing outcomes of multi-output models, rather than scalarizing a batch posterior.

My initial idea to do this would be the following:

1. Use a standard single-output model to model the stationary cost over R^2d. You can batch-evaluate this model using an input X of shape n-1 x 1 x 2d, where X[i, 0, :] = (x_i, x_{i+1}). This will produce a batched posterior of shape n-1 x 1 x 1.
2. Define a lightweight wrapper around this model that applies the scalarization from scalarize_posterior across the batch dimension (rather than the output dimension)

That way you should be able to directly use the analytic acquisition functions with such wrapper model. To achieve 2, you could either translate scalarize_posterior to work on the batch dims, or you could grab the batched Posterior object from the base model, flip the batch (first) and output (last) dimension of the mean and covariance, and then apply scalarize_posterior verbatim.

Details need to be worked out of course, but I think such approach should work without too much pain.

[jlmaccal]

Thanks, I'll check out scalarize_posterior.

My current thinking is that I would have an "outer" model that evaluates entire trajectories and an "inner" model that evaluates the cost of the steps. It would work something like this:

1. The outer model takes as input a batch of trajectories with shape b x 1 x nd. It restructures this have shape b x (n-1) x 2d. (I think the n-1 segments from each trajectory should be inferred jointly, as they all come from the same realization of the step cost function.)
2. The inner model takes b x (n-1) x 2d as input and produces a b x (n-1) x 1 batch of posteriors.
3. The outer model takes the b x (n-1) x 1 and transforms it to a b x 1 x 1 batch of posteriors.

For 2, I haven't actually experimented much yet, so it's not clear to me how BoTorch handles posteriors with t-batching. Do functions return a list of posteriors or a posterior with an additional batch dimension?

For 3, the output mean is the sum of the n-1 means along the q dimension. The output variance is the sum of the (n-1) * (n-1) covariance entries along the q dimension.

At the moment, it's not clear to me which methods of Model I need to implement. I want to be able to optimize one of the analytical acquisition functions, like EI. Is it sufficient to only implement posterior, or will I need other methods as well?

Sorry for all of the questions. I'm trying to do something non-standard without having really used the library before, which is always difficult.

[saitcakmak]

For 2, I haven't actually experimented much yet, so it's not clear to me how BoTorch handles posteriors with t-batching. Do functions return a list of posteriors or a posterior with an additional batch dimension?

I can answer that. It will return a posterior with the same batch dimension, where the posterior is independent across t-batches. So, for batch_shape x q x d input, the posterior would have a shape batch_shape x q x 1, and this is equivalent to calling posterior for each q x d with a for loop over t-batches (but much faster!).

At the moment, it's not clear to me which methods of Model I need to implement. I want to be able to optimize one of the analytical acquisition functions, like EI. Is it sufficient to only implement posterior, or will I need other methods as well?

If you're building on top of the standard models (which I think is the best way to go, at least for starters), I don't think you need to implement a new Model. The easiest way would probably be to transform the posterior, using a modified ScalarizedObjective wrapper as Max suggested above.
For defining a new model, you'd typically define the __init__ and forward methods, where forward returns the prior mean & covariance, and let the existing posteriors handle the rest.

Going over your 3-step outline above, I'd let the acquisition function or some other wrapper handle the conversion in step 1 (since it's just reshaping the tensor, a custom InputTransform may also work), define a GP over the 2d input space and use it's posterior for step 2, and use a scalarization over this q x 1 posterior to handle step 3 (goes back to ScalarizedObjective). Hope this helps!