Merge changes I4c9f0071,Iaa5afeac,I7dbea619 into main
* changes:
Add flashbax support
Fix documentation bug in dlqr
Fix NaNs in physics
diff --git a/frc971/control_loops/python/controls.py b/frc971/control_loops/python/controls.py
index 4cb10a2..7c7c310 100644
--- a/frc971/control_loops/python/controls.py
+++ b/frc971/control_loops/python/controls.py
@@ -85,11 +85,11 @@
J = sum(0, inf, x.T * Q * x + u.T * R * u)
"""
- # P = (A.T * P * A) - (A.T * P * B * numpy.linalg.inv(R + B.T * P *B) * (A.T * P.T * B).T + Q
- # 0.5 * X.T * P * X -> optimal cost to infinity
+ # P = (A.T * P * A) - (A.T * P * B * numpy.linalg.inv(R + B.T * P * B) * (A.T * P.T * B).T + Q
+ # X.T * P * X -> optimal cost to infinity
P = scipy.linalg.solve_discrete_are(a=A, b=B, q=Q, r=R)
- F = numpy.linalg.inv(R + B.T * P * B) * B.T * P * A
+ F = numpy.linalg.inv(R + B.T @ P @ B) @ B.T @ P @ A
if optimal_cost_function:
return F, P
else:
diff --git a/frc971/control_loops/swerve/jax_dynamics.py b/frc971/control_loops/swerve/jax_dynamics.py
index 8b05ef4..76a00d5 100644
--- a/frc971/control_loops/swerve/jax_dynamics.py
+++ b/frc971/control_loops/swerve/jax_dynamics.py
@@ -115,14 +115,14 @@
def softsign(x, gain):
- return -2 / (1 + jax.numpy.exp(gain * x)) + 1
+ return 1 - 2.0 * jax.nn.sigmoid(-gain * x)
def soft_atan2(y, x):
kMaxLogGain = 1.0 / 0.05
kAbsLogGain = 1.0 / 0.01
- softabs_x = x * (1.0 - 2.0 / (1 + jax.numpy.exp(kAbsLogGain * x)))
+ softabs_x = x * softsign(x, kAbsLogGain)
return jax.numpy.arctan2(
y,
@@ -272,7 +272,7 @@
Fwx = (coefficients.Ktd / (coefficients.Gd * coefficients.rw)) * Id
Fwy = coefficients.Cy * slip_angle
- softsign_velocity = softsign(wheel_ground_velocity[0], 100)
+ softsign_velocity = softsign(wheel_ground_velocity[0], 100.0)
Ms = -Fwy * (
(softsign_velocity * coefficients.contact_patch_length / 3.0) +
diff --git a/tools/python/requirements.lock.txt b/tools/python/requirements.lock.txt
index 7faf08a..93171a9 100644
--- a/tools/python/requirements.lock.txt
+++ b/tools/python/requirements.lock.txt
@@ -192,7 +192,9 @@
chex==0.1.86 \
--hash=sha256:251c20821092323a3d9c28e1cf80e4a58180978bec368f531949bd9847eee568 \
--hash=sha256:e8b0f96330eba4144659e1617c0f7a57b161e8cbb021e55c6d5056c7378091d1
- # via optax
+ # via
+ # flashbax
+ # optax
click==8.1.7 \
--hash=sha256:ae74fb96c20a0277a1d615f1e4d73c8414f5a98db8b799a7931d1582f3390c28 \
--hash=sha256:ca9853ad459e787e2192211578cc907e7594e294c7ccc834310722b41b9ca6de
@@ -338,6 +340,10 @@
# optax
# orbax-checkpoint
# tensorflow-datasets
+flashbax==0.1.2 \
+ --hash=sha256:ac50b2580808ce63787da0ae240db14986e3404ade98a33335e6d7a5efe84859 \
+ --hash=sha256:b566ac5a78975b3e0a0a404a8844a26aa45e9cacfaad2829dcbcac2ffb3d5f7a
+ # via -r tools/python/requirements.txt
flask==3.0.3 \
--hash=sha256:34e815dfaa43340d1d15a5c3a02b8476004037eb4840b34910c6e21679d288f3 \
--hash=sha256:ceb27b0af3823ea2737928a4d99d125a06175b8512c445cbd9a9ce200ef76842
@@ -352,6 +358,7 @@
# via
# -r tools/python/requirements.txt
# clu
+ # flashbax
fonttools==4.53.1 \
--hash=sha256:02569e9a810f9d11f4ae82c391ebc6fb5730d95a0657d24d754ed7763fb2d122 \
--hash=sha256:0679a30b59d74b6242909945429dbddb08496935b82f91ea9bf6ad240ec23397 \
@@ -527,6 +534,7 @@
# -r tools/python/requirements.txt
# chex
# clu
+ # flashbax
# flax
# optax
# orbax-checkpoint
@@ -568,9 +576,14 @@
# via
# chex
# clu
+ # flashbax
# jax
# optax
# orbax-checkpoint
+jaxtyping==0.2.34 \
+ --hash=sha256:2f81fb6d1586e497a6ea2d28c06dcab37b108a096cbb36ea3fe4fa2e1c1f32e5 \
+ --hash=sha256:eed9a3458ec8726c84ea5457ebde53c964f65d2c22c0ec40d0555ae3fed5bbaf
+ # via -r tools/python/requirements.txt
jinja2==3.1.4 \
--hash=sha256:4a3aee7acbbe7303aede8e9648d13b8bf88a429282aa6122a993f0ac800cb369 \
--hash=sha256:bc5dd2abb727a5319567b7a813e6a2e7318c39f4f487cfe6c89c6f9c7d25197d
@@ -969,6 +982,7 @@
# clu
# contourpy
# etils
+ # flashbax
# flax
# h5py
# jax
@@ -1698,6 +1712,7 @@
--hash=sha256:cc315029f49c0f294f0721462c221e0ef4c15360a526cc34392ac81565fd63b8 \
--hash=sha256:f47597209ce11228cfe6b94999f582788aac5571e85c3e8dcaa43b1f07660589
# via
+ # flashbax
# flax
# orbax-checkpoint
termcolor==2.4.0 \
@@ -1735,6 +1750,10 @@
--hash=sha256:90279a3770753eafc9194a0364852159802111925aa30eb3f9d85b0e805ac7cd \
--hash=sha256:e1020aef2e5096702d8a025ac7d16b1577279c9d63f8375b63083e9a5f0fcbad
# via tensorflow-datasets
+typeguard==2.13.3 \
+ --hash=sha256:00edaa8da3a133674796cf5ea87d9f4b4c367d77476e185e80251cc13dfbb8c4 \
+ --hash=sha256:5e3e3be01e887e7eafae5af63d1f36c849aaa94e3a0112097312aabfa16284f1
+ # via jaxtyping
typing-extensions==4.12.2 \
--hash=sha256:04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d \
--hash=sha256:1a7ead55c7e559dd4dee8856e3a88b41225abfe1ce8df57b7c13915fe121ffb8
@@ -1742,6 +1761,7 @@
# chex
# clu
# etils
+ # flashbax
# flax
# optree
# orbax-checkpoint
diff --git a/tools/python/requirements.txt b/tools/python/requirements.txt
index e7f23c3..2b806f5 100644
--- a/tools/python/requirements.txt
+++ b/tools/python/requirements.txt
@@ -25,9 +25,13 @@
bokeh
tabulate
+flask
+
casadi>=3.6.6
+# ML libraries
jax[cuda12]
+jaxtyping
optax
flax
clu
@@ -38,4 +42,5 @@
tensorflow
tensorflow_datasets
-flask
+# Experience buffer for reinforcement learning
+flashbax
diff --git a/tools/python/whl_overrides.json b/tools/python/whl_overrides.json
index c3d9a89..586a9d3 100644
--- a/tools/python/whl_overrides.json
+++ b/tools/python/whl_overrides.json
@@ -63,6 +63,10 @@
"sha256": "6dc882d355e1e98a5d1a148d6323679dc47c9a5792939b9de72615aa4737eb0b",
"url": "https://software.frc971.org/Build-Dependencies/wheelhouse/etils-1.5.2-py3-none-any.whl"
},
+ "flashbax==0.1.2": {
+ "sha256": "ac50b2580808ce63787da0ae240db14986e3404ade98a33335e6d7a5efe84859",
+ "url": "https://software.frc971.org/Build-Dependencies/wheelhouse/flashbax-0.1.2-py3-none-any.whl"
+ },
"flask==3.0.3": {
"sha256": "34e815dfaa43340d1d15a5c3a02b8476004037eb4840b34910c6e21679d288f3",
"url": "https://software.frc971.org/Build-Dependencies/wheelhouse/flask-3.0.3-py3-none-any.whl"
@@ -143,6 +147,10 @@
"sha256": "f74a6b0e09df4b5e2ee399ebb9f0e01190e26e84ccb0a758fadb516415c07f18",
"url": "https://software.frc971.org/Build-Dependencies/wheelhouse/jaxlib-0.4.30-cp39-cp39-manylinux2014_x86_64.whl"
},
+ "jaxtyping==0.2.34": {
+ "sha256": "2f81fb6d1586e497a6ea2d28c06dcab37b108a096cbb36ea3fe4fa2e1c1f32e5",
+ "url": "https://software.frc971.org/Build-Dependencies/wheelhouse/jaxtyping-0.2.34-py3-none-any.whl"
+ },
"jinja2==3.1.4": {
"sha256": "bc5dd2abb727a5319567b7a813e6a2e7318c39f4f487cfe6c89c6f9c7d25197d",
"url": "https://software.frc971.org/Build-Dependencies/wheelhouse/jinja2-3.1.4-py3-none-any.whl"
@@ -443,6 +451,10 @@
"sha256": "90279a3770753eafc9194a0364852159802111925aa30eb3f9d85b0e805ac7cd",
"url": "https://software.frc971.org/Build-Dependencies/wheelhouse/tqdm-4.66.5-py3-none-any.whl"
},
+ "typeguard==2.13.3": {
+ "sha256": "5e3e3be01e887e7eafae5af63d1f36c849aaa94e3a0112097312aabfa16284f1",
+ "url": "https://software.frc971.org/Build-Dependencies/wheelhouse/typeguard-2.13.3-py3-none-any.whl"
+ },
"typing_extensions==4.12.2": {
"sha256": "04e5ca0351e0f3f85c6853954072df659d0d13fac324d0072316b67d7794700d",
"url": "https://software.frc971.org/Build-Dependencies/wheelhouse/typing_extensions-4.12.2-py3-none-any.whl"
diff --git a/y2018/control_loops/python/extended_lqr.py b/y2018/control_loops/python/extended_lqr.py
index e19ee8f..a2a31af 100755
--- a/y2018/control_loops/python/extended_lqr.py
+++ b/y2018/control_loops/python/extended_lqr.py
@@ -85,6 +85,7 @@
Returns:
numpy.matrix(1, 1), The quadratic cost of being at X
"""
+ print('TODO(austin): S -> X.T * S * X, need to fix if we care')
return 0.5 * X.T * self.S * X
def cost(self, X, U):