<jupyter_start><jupyter_text>The state transition equation is the following:$$\def \vx {\boldsymbol{\color{Plum}{x}}}\def \vu {\boldsymbol{\color{orange}{u}}}\dot{\vx} = f(\vx, \vu) \quad\left\{\begin{array}{l}\dot{x} = s \cos \theta \\\dot{y} = s \sin \theta \\\dot{\theta} = \frac{s}{L} \tan \phi \\\dot{s} = a\end{array}\right. \quad\vx = (x\;y\;\theta\;s) \quad\vu = (\phi\;a)$$<jupyter_code>def f(x, u, t=None):
    """
    Kinematic model for tricycle
    ẋ(t) = f[x(t), u(t), t]
    x: states (x, y, θ, s)
    u: control
    t: time
    f: kinematic model
    ẋ = dx/dt
    x' = x + f(x, u, t) * dt
    """
    L = 1  # m
    x, y, θ, s = x
    
    ϕ, a = u
    f = torch.zeros(4)
    f[0] = s * torch.cos(θ)
    f[1] = s * torch.sin(θ)
    f[2] = s / L * torch.tan(ϕ)
    f[3] = a
    return f
# Manual driving
x = torch.tensor((0, 0, 0, 1),dtype=torch.float32)
# Optimal action from back propagation
u = torch.tensor([
    [0.1280, 0.0182],
    [0.0957, 0.0131],
    [0.0637, 0.0085],
    [0.0318, 0.0043],
    [0.0000, 0.0000]
])
# Brake
u = torch.ones(10, 2) * -0.1
u[:, 0] = 0
# S
u = torch.zeros(10, 2)
u[:5, 0] = 0.2
u[5:, 0] = -0.2
# Straight
# u = torch.zeros(10, 2)

dt = 1  # s
trajectory = [x.clone()]
for t in range(10):
    x += f(x, u[t]) * dt
    print(x)
    trajectory.append(x.clone())
τ = torch.stack(trajectory)

# plt.plot(0,0,'gx', markersize=20, markeredgewidth=5)
# plt.plot(5,1,'rx', markersize=20, markeredgewidth=5)

# Costs definition
# x: states (x, y, θ, s)
def vanilla_cost(state, target):
    x_x, x_y = target
    return (state[-1][0] - x_x).pow(2) + (state[-1][1] - x_y).pow(2)

def cost_with_target_s(state, target):
    x_x, x_y = target
    return (state[-1][0] - x_x).pow(2) + (state[-1][1] - x_y).pow(2) + (state[-1][-1]).pow(2)

def cost_sum_distances(state, target):
    x_x, x_y = target
    dists = ((state[:, 0] - x_x).pow(2) + (state[:, 1] - x_y).pow(2)).pow(0.5)
    return dists.mean()

def cost_sum_square_distances(state, target):
    x_x, x_y = target
    dists = ((state[:, 0] - x_x).pow(2) + (state[:, 1] - x_y).pow(2))
    return dists.mean()

def cost_logsumexp(state, target):
    x_x, x_y = target
    dists = ((state[:, 0] - x_x).pow(2) + (state[:, 1] - x_y).pow(2))#.pow(0.5)
    return -1 * torch.logsumexp(-1 * dists, dim=0)
# Path planning
def path_planning_with_cost(x_x, x_y, s, T, epochs, stepsize, cost_f, ax=None, ax_lims=None, debug=False):
    """
    Path planning for tricycle
    x_x: x component of target postion vector
    x_y: y component of target postion vector
    s: initial speed
    T: time steps
    epochs: number of epochs for back propagation
    stepsize: stepsize for back propagation
    cost_f: cost funciton that takes the trajectory and the tuple (x, y) - target.
    ax: axis to plot the trajectory
    """

    u = nn.Parameter(torch.zeros(T, 2))
    optimizer = SGD((u,), lr=stepsize)
    dt = 1  # s
    costs = []
    for epoch in range(epochs):
        x = [torch.tensor((0, 0, 0, s),dtype=torch.float32)]
        for t in range(1, T+1):
            x.append(x[-1] + f(x[-1], u[t-1]) * dt)
        x_t = torch.stack(x)
        τ = torch.stack(x).detach()
        cost = cost_f(x_t, (x_x, x_y))
        costs.append(cost.item())
        optimizer.zero_grad()
        cost.backward()
        optimizer.step()
        if debug: 
            print(u.data)

path_planning_with_cost(x_x=5, x_y=1, s=1, T=5, epochs=100, stepsize=0.01, cost_f=vanilla_cost, debug=True)<jupyter_output>tensor([[0.0800, 0.0000],
        [0.0600, 0.0000],
        [0.0400, 0.0000],
        [0.0200, 0.0000],
        [0.0000, 0.0000]])
tensor([[0.1114, 0.0109],
        [0.0833, 0.0078],
        [0.0554, 0.0051],
        [0.0277, 0.0026],
        [0.0000, 0.0000]])
tensor([[0.1230, 0.0159],
        [0.0920, 0.0115],
        [0.0612, 0.0075],
        [0.0306, 0.0037],
        [0.0000, 0.0000]])
tensor([[0.1268, 0.0176],
        [0.0949, 0.0127],
        [0.0631, 0.0083],
        [0.0315, 0.0042],
        [0.0000, 0.0000]])
tensor([[0.1280, 0.0182],
        [0.0957, 0.0131],
        [0.0637, 0.0085],
        [0.0318, 0.0043],
        [0.0000, 0.0000]])
tensor([[0.1283, 0.0183],
        [0.0960, 0.0132],
        [0.0639, 0.0086],
        [0.0319, 0.0043],
        [0.0000, 0.0000]])
tensor([[0.1284, 0.0184],
        [0.0961, 0.0132],
        [0.0639, 0.0086],
        [0.0319, 0.0043],
        [0.0000, 0.0000]])
tensor([[0.1285, 0.0184],
        [0.0961, 0.0132],
        [0.0640, 0.0086],
     [...]
{'fraction_non_alphanumeric': 0.208955223880597, 'fraction_numerical': 0.11851191802183114, 'mean_word_length': 2.365817091454273, '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': '615819', 'license': 'Apache-2.0', 'n_tokens_mistral': 2271, 'n_tokens_neox': 1986, 'n_words': 432, 'path': 'Optimal control, MPC with pytorch.ipynb', 'repo_name': 'JiaheXu/Pytorch_tutorial', 'size': '4489'}