<jupyter_start><jupyter_text>Notebook PurposePrototype using probabilistic programming in NumPyro for control synthesis with STL specifications.<jupyter_code>import torch

import numpyro
import numpyro.distributions as dist
from numpyro.infer import MCMC, NUTS

import numpy as np
import jax.numpy as jnp
from jax import random
# Let's define a simple model for a LTI dynamical system. No STL yet, but we just want to measure
# the distance between the output and setpoint after 5 seconds.
def lti_stability_model(damping=None, actuation=None, final_location=None):
    # The drift and actuation matrices A and B are random, but
    # are centered around a damped single integrator
    damping = numpyro.sample("damping", dist.Uniform(np.array(0.05), np.array(0.15)), obs=damping)
    actuation = numpyro.sample("actuation", dist.Uniform(np.array(0.8), np.array(1.2)), obs=actuation)

    A = np.zeros((damping.shape[0], 2, 2))
    A[:, 0, 1] = 1.0
    A[:, 1, 1] = -damping
    B = np.zeros((damping.shape[0], 2, 1))
    B[:, 1, 0] = actuation

    # Define a Gausian prior over the control parameters
    K = numpyro.sample("K", dist.Normal(np.ones((2, 1)), 5.0 * np.ones(2)))

    print(A.shape)
    print(B.shape)
    print(K.shape)

    # Define the closed-loop transfer matrix
    Acl = A - np.matmul(B, K)

    # Simulate for 5 seconds
    dt = 0.1
    T = 5.0
    n_steps = T // dt
    x = np.zeros((n_steps, 2, 1))
    x[0, 0, 1] = 1.0  # start at location = 1.0, velocity = 0.0
    for i in range(1, n_steps):
        x[i] = Acl @ x[i - 1]

    # Observe the final location
    observation_noise = 0.01
    numpyro.sample("observed_location", dist.Normal(x[-1, 0, 0], observation_noise), obs=final_location)
# Get an RNG key to work with
rng_key = random.PRNGKey(0)
rng_key, rng_subkey = random.split(rng_key)

# Get some damping values
N_samples = 10
damping = random.uniform(rng_subkey, (N_samples,), minval=0.05, maxval=0.15)
rng_key, rng_subkey = random.split(rng_key)
actuation = random.uniform(rng_subkey, (N_samples,), minval=0.8, maxval=1.2)
rng_key, rng_subkey = random.split(rng_key)
final_location = 0.01 * random.normal(rng_subkey, (N_samples,))

# Run NUTS.
kernel = NUTS(lti_stability_model)
num_samples = 2000
mcmc = MCMC(kernel, num_warmup=1000, num_samples=num_samples)
mcmc.run(
    rng_subkey, damping=damping, actuation=actuation, final_location=final_location,
)
mcmc.print_summary()
samples_1 = mcmc.get_samples()<jupyter_output>(10, 2, 2)
(10, 2, 1)
(2, 2)
{'fraction_non_alphanumeric': 0.12615508236239453, 'fraction_numerical': 0.035757332261952594, 'mean_word_length': 4.050709939148073, 'pattern_counts': {'":': 0, '<': 4, '<?xml version=': 0, '>': 4, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '28468', 'license': 'BSD-2-Clause', 'n_tokens_mistral': 1002, 'n_tokens_neox': 927, 'n_words': 253, 'path': 'prototyping/NumPyro Experiments.ipynb', 'repo_name': 'MIT-REALM/architect', 'size': '2489'}