Skip to content

Commit

Permalink
Clean up FRC 2022 shooter example coordinate systems (#580)
Browse files Browse the repository at this point in the history
  • Loading branch information
calcmogul authored Jun 29, 2024
1 parent 403fcd1 commit e2501a5
Show file tree
Hide file tree
Showing 2 changed files with 101 additions and 86 deletions.
95 changes: 49 additions & 46 deletions examples/FRC2022Shooter/main.py
Original file line number Diff line number Diff line change
Expand Up @@ -16,32 +16,29 @@
import numpy as np
from numpy.linalg import norm

field_width = 8.2296 # 27 ft
field_length = 16.4592 # 54 ft
field_width = 8.2296 # 27 ft -> m
field_length = 16.4592 # 54 ft -> m
target_wrt_field = np.array(
[[field_length / 2.0], [field_width / 2.0], [2.64], [0.0], [0.0], [0.0]]
)
target_radius = 0.61
g = 9.806 # m/s²


def main():
# Robot initial velocity
robot_initial_v_x = 0.2 # ft/s
robot_initial_v_y = -0.2 # ft/s
robot_initial_v_z = 0.0 # ft/s
def lerp(a, b, t):
return a + t * (b - a)

max_launch_velocity = 10.0

shooter = np.array([[field_length / 4.0], [field_width / 4.0], [1.2]])
shooter_x = shooter[0, 0]
shooter_y = shooter[1, 0]
shooter_z = shooter[2, 0]
def main():
# Robot initial state
robot_wrt_field = np.array(
[[field_length / 4.0], [field_width / 4.0], [0.0], [1.524], [-1.524], [0.0]]
)

target = np.array([[field_length / 2.0], [field_width / 2.0], [2.64]])
target_x = target[0, 0]
target_y = target[1, 0]
target_z = target[2, 0]
target_radius = 0.61
max_launch_velocity = 10.0 # m/s

def lerp(a, b, t):
return a + t * (b - a)
shooter_wrt_robot = np.array([[0.0], [0.0], [1.2], [0.0], [0.0], [0.0]])
shooter_wrt_field = robot_wrt_field + shooter_wrt_robot

problem = OptimizationProblem()

Expand All @@ -52,6 +49,8 @@ def lerp(a, b, t):
T.set_value(1)
dt = T / N

# Ball state in field frame
#
# [x position]
# [y position]
# [z position]
Expand All @@ -60,42 +59,46 @@ def lerp(a, b, t):
# [z velocity]
X = problem.decision_variable(6, N)

p = X[:3, :]
p_x = X[0, :]
p_y = X[1, :]
p_z = X[2, :]

v = X[3:, :]
v_x = X[3, :]
v_y = X[4, :]
v_z = X[5, :]

# Position initial guess is linear interpolation between start and end position
for k in range(N):
p_x[k].set_value(lerp(shooter_x, target_x, k / N))
p_y[k].set_value(lerp(shooter_y, target_y, k / N))
p_z[k].set_value(lerp(shooter_z, target_z, k / N))
p_x[k].set_value(lerp(shooter_wrt_field[0, 0], target_wrt_field[0, 0], k / N))
p_y[k].set_value(lerp(shooter_wrt_field[1, 0], target_wrt_field[1, 0], k / N))
p_z[k].set_value(lerp(shooter_wrt_field[2, 0], target_wrt_field[2, 0], k / N))

# Velocity initial guess is max launch velocity toward goal
uvec_shooter_to_target = target - shooter
uvec_shooter_to_target = target_wrt_field[:3, :] - shooter_wrt_field[:3, :]
uvec_shooter_to_target /= norm(uvec_shooter_to_target)
for k in range(N):
v_x[k].set_value(max_launch_velocity * uvec_shooter_to_target[0, 0])
v_y[k].set_value(max_launch_velocity * uvec_shooter_to_target[1, 0])
v_z[k].set_value(max_launch_velocity * uvec_shooter_to_target[2, 0])
v[:, k].set_value(
robot_wrt_field[3:, :] + max_launch_velocity * uvec_shooter_to_target
)

# Shooter initial position
problem.subject_to(X[:3, 0] == shooter)
problem.subject_to(p[:, 0] == shooter_wrt_field[:3, 0])

# Require initial launch velocity is below max
#
# √{v_x² + v_y² + v_z²) ≤ vₘₐₓ
# v_x² + v_y² + v_z² ≤ vₘₐₓ²
problem.subject_to(
v_x[0] ** 2 + v_y[0] ** 2 + v_z[0] ** 2 <= max_launch_velocity**2
(v_x[0] - robot_wrt_field[3, 0]) ** 2
+ (v_y[0] - robot_wrt_field[4, 0]) ** 2
+ (v_z[0] - robot_wrt_field[5, 0]) ** 2
<= max_launch_velocity**2
)

# Require final position is in center of target circle
problem.subject_to(p_x[-1] == target_x)
problem.subject_to(p_y[-1] == target_y)
problem.subject_to(p_z[-1] == target_z)
problem.subject_to(p[:, -1] == target_wrt_field[:3, :])

# Require the final velocity is down
problem.subject_to(v_z[-1] < 0.0)
Expand All @@ -115,9 +118,9 @@ def f(x):
m = 2.0 # kg
a_D = lambda v: 0.5 * rho * v**2 * C_D * A / m

v_x = x[3, 0] + robot_initial_v_x
v_y = x[4, 0] + robot_initial_v_y
v_z = x[5, 0] + robot_initial_v_z
v_x = x[3, 0]
v_y = x[4, 0]
v_z = x[5, 0]
return VariableMatrix(
[[v_x], [v_y], [v_z], [-a_D(v_x)], [-a_D(v_y)], [-g - a_D(v_z)]]
)
Expand All @@ -140,16 +143,16 @@ def f(x):
problem.solve(diagnostics=True)

# Initial velocity vector
v = X[3:, 0].value()
v0 = X[3:, 0].value() - robot_wrt_field[3:, :]

launch_velocity = norm(v)
launch_velocity = norm(v0)
print(f"Launch velocity = {launch_velocity:.03f} m/s")

pitch = math.atan2(v[2, 0], math.hypot(v[0, 0], v[1, 0]))
print(f"Pitch = {pitch * 180.0 / math.pi:.03f}°")
pitch = math.atan2(v0[2, 0], math.hypot(v0[0, 0], v0[1, 0]))
print(f"Pitch = {np.rad2deg(pitch):.03f}°")

yaw = math.atan2(v[1, 0], v[0, 0])
print(f"Yaw = {yaw * 180.0 / math.pi:.03f}°")
yaw = math.atan2(v0[1, 0], v0[0, 0])
print(f"Yaw = {np.rad2deg(yaw):.03f}°")

print(f"Total time = {T.value():.03f} s")
print(f"dt = {dt.value() * 1e3:.03f} ms")
Expand All @@ -175,19 +178,19 @@ def plot_wireframe(ax, f, x_range, y_range, color):

# Target
ax.plot(
target_x,
target_y,
target_z,
target_wrt_field[0, 0],
target_wrt_field[1, 0],
target_wrt_field[2, 0],
color="black",
marker="x",
)
xs = []
ys = []
zs = []
for angle in np.arange(0.0, 2.0 * math.pi, 0.1):
xs.append(target_x + target_radius * math.cos(angle))
ys.append(target_y + target_radius * math.sin(angle))
zs.append(target_z)
xs.append(target_wrt_field[0, 0] + target_radius * math.cos(angle))
ys.append(target_wrt_field[1, 0] + target_radius * math.sin(angle))
zs.append(target_wrt_field[2, 0])
ax.plot(xs, ys, zs, color="black")

# Trajectory
Expand Down
92 changes: 52 additions & 40 deletions examples/FRC2022Shooter/src/Main.cpp
Original file line number Diff line number Diff line change
Expand Up @@ -12,31 +12,33 @@
// This program finds the optimal initial launch velocity and launch angle for
// the 2022 FRC game's target.

constexpr double field_width = 8.2296; // 27 ft
constexpr double field_length = 16.4592; // 54 ft
constexpr double g = 9.806; // m/s²
namespace slp = sleipnir;

using Eigen::Vector3d;
using Vector6d = Eigen::Vector<double, 6>;

constexpr double field_width = 8.2296; // 27 ft -> m
constexpr double field_length = 16.4592; // 54 ft -> m
constexpr Vector6d target_wrt_field{
{field_length / 2.0}, {field_width / 2.0}, {2.64}, {0.0}, {0.0}, {0.0}};
constexpr double g = 9.806; // m/s²

#ifndef RUNNING_TESTS
int main() {
// Robot initial velocity
double robot_initial_v_x = 0.2; // ft/s
double robot_initial_v_y = -0.2; // ft/s
double robot_initial_v_z = 0.0; // ft/s
// Robot initial state
constexpr Vector6d robot_wrt_field{{field_length / 4.0},
{field_width / 4.0},
{0.0},
{1.524},
{-1.524},
{0.0}};

constexpr double max_launch_velocity = 10.0;

Eigen::Vector3d shooter{{field_length / 4.0}, {field_width / 4.0}, {1.2}};
const auto& shooter_x = shooter(0, 0);
const auto& shooter_y = shooter(1, 0);
const auto& shooter_z = shooter(2, 0);

constexpr Eigen::Vector3d target{
{field_length / 2.0}, {field_width / 2.0}, {2.64}};
const auto& target_x = target(0, 0);
const auto& target_y = target(1, 0);
const auto& target_z = target(2, 0);
Vector6d shooter_wrt_robot{{0.0}, {0.0}, {1.2}, {0.0}, {0.0}, {0.0}};
Vector6d shooter_wrt_field = robot_wrt_field + shooter_wrt_robot;

sleipnir::OptimizationProblem problem;
slp::OptimizationProblem problem;

// Set up duration decision variables
constexpr int N = 10;
Expand All @@ -45,6 +47,8 @@ int main() {
T.SetValue(1);
auto dt = T / N;

// Ball state in field frame
//
// [x position]
// [y position]
// [z position]
Expand All @@ -53,48 +57,55 @@ int main() {
// [z velocity]
auto X = problem.DecisionVariable(6, N);

auto p = X.Block(0, 0, 3, N);
auto p_x = X.Row(0);
auto p_y = X.Row(1);
auto p_z = X.Row(2);

auto v = X.Block(3, 0, 3, N);
auto v_x = X.Row(3);
auto v_y = X.Row(4);
auto v_z = X.Row(5);

// Position initial guess is linear interpolation between start and end
// position
for (int k = 0; k < N; ++k) {
p_x(k).SetValue(std::lerp(shooter_x, target_x, static_cast<double>(k) / N));
p_y(k).SetValue(std::lerp(shooter_y, target_y, static_cast<double>(k) / N));
p_z(k).SetValue(std::lerp(shooter_z, target_z, static_cast<double>(k) / N));
p_x(k).SetValue(std::lerp(shooter_wrt_field(0), target_wrt_field(0),
static_cast<double>(k) / N));
p_y(k).SetValue(std::lerp(shooter_wrt_field(1), target_wrt_field(1),
static_cast<double>(k) / N));
p_z(k).SetValue(std::lerp(shooter_wrt_field(2), target_wrt_field(2),
static_cast<double>(k) / N));
}

// Velocity initial guess is max launch velocity toward goal
Eigen::Vector3d uvec_shooter_to_target = (target - shooter).normalized();
Vector3d uvec_shooter_to_target =
(target_wrt_field.block(0, 0, 3, 1) - shooter_wrt_field.block(0, 0, 3, 1))
.normalized();
for (int k = 0; k < N; ++k) {
v_x(k).SetValue(max_launch_velocity * uvec_shooter_to_target(0, 0));
v_y(k).SetValue(max_launch_velocity * uvec_shooter_to_target(1, 0));
v_z(k).SetValue(max_launch_velocity * uvec_shooter_to_target(2, 0));
v.Col(k).SetValue(robot_wrt_field.block(3, 0, 3, 1) +
max_launch_velocity * uvec_shooter_to_target);
}

// Shooter initial position
problem.SubjectTo(X.Block(0, 0, 3, 1) == shooter);
problem.SubjectTo(p.Col(0) == shooter_wrt_field.block(0, 0, 3, 1));

// Require initial launch velocity is below max
//
// √{v_x² + v_y² + v_z²) ≤ vₘₐₓ
// v_x² + v_y² + v_z² ≤ vₘₐₓ²
problem.SubjectTo(v_x(0) * v_x(0) + v_y(0) * v_y(0) + v_z(0) * v_z(0) <=
problem.SubjectTo(slp::pow(v_x(0) - robot_wrt_field(3), 2) +
slp::pow(v_y(0) - robot_wrt_field(4), 2) +
slp::pow(v_z(0) - robot_wrt_field(5), 2) <=
max_launch_velocity * max_launch_velocity);

// Require final position is in center of target circle
problem.SubjectTo(p_x(N - 1) == target_x);
problem.SubjectTo(p_y(N - 1) == target_y);
problem.SubjectTo(p_z(N - 1) == target_z);
problem.SubjectTo(p.Col(N - 1) == target_wrt_field.block(0, 0, 3, 1));

// Require the final velocity is down
problem.SubjectTo(v_z(N - 1) < 0.0);

auto f = [&](sleipnir::VariableMatrix x) {
auto f = [](slp::VariableMatrix x) {
// x' = x'
// y' = y'
// z' = z'
Expand All @@ -109,11 +120,11 @@ int main() {
constexpr double m = 2.0; // kg
auto a_D = [](auto v) { return 0.5 * rho * v * v * C_D * A / m; };

auto v_x = x(3, 0) + robot_initial_v_x;
auto v_y = x(4, 0) + robot_initial_v_y;
auto v_z = x(5, 0) + robot_initial_v_z;
return sleipnir::VariableMatrix{{v_x}, {v_y}, {v_z},
{-a_D(v_x)}, {-a_D(v_y)}, {-g - a_D(v_z)}};
auto v_x = x(3, 0);
auto v_y = x(4, 0);
auto v_z = x(5, 0);
return slp::VariableMatrix{{v_x}, {v_y}, {v_z},
{-a_D(v_x)}, {-a_D(v_y)}, {-g - a_D(v_z)}};
};

// Dynamics constraints - RK4 integration
Expand All @@ -135,15 +146,16 @@ int main() {
problem.Solve({.diagnostics = true});

// Initial velocity vector
Eigen::Vector3d v = X.Block(3, 0, 3, 1).Value();
Eigen::Vector3d v0 =
X.Block(3, 0, 3, 1).Value() - robot_wrt_field.block(3, 0, 3, 1);

double launch_velocity = v.norm();
double launch_velocity = v0.norm();
std::println("Launch velocity = {:.03} ms", launch_velocity);

double pitch = std::atan2(v(2), std::hypot(v(0), v(1)));
double pitch = std::atan2(v0(2), std::hypot(v0(0), v0(1)));
std::println("Pitch = {:.03}°", pitch * 180.0 / std::numbers::pi);

double yaw = std::atan2(v(1), v(0));
double yaw = std::atan2(v0(1), v0(0));
std::println("Yaw = {:.03}°", yaw * 180.0 / std::numbers::pi);

std::println("Total time = {:.03} s", T.Value());
Expand Down

0 comments on commit e2501a5

Please sign in to comment.