Add flax and friends for training neural networks
This updates our formatter, which changed what it expects very slightly.
Update all the pieces of code to match
optax is for the actual optimizers, jax is the underlying accelerated
linear algebra library, tensorflow is for loading datasets and exporting
models.
Change-Id: Ic4c3b425cda74267e1d0ad1615c42452cbefab8a
Signed-off-by: Austin Schuh <austin.linux@gmail.com>
diff --git a/y2014/control_loops/python/extended_lqr.py b/y2014/control_loops/python/extended_lqr.py
index 699dd30..dca57eb 100755
--- a/y2014/control_loops/python/extended_lqr.py
+++ b/y2014/control_loops/python/extended_lqr.py
@@ -294,18 +294,18 @@
S_bar_1_eigh_eigenvalues_stiff[i] = max(
S_bar_1_eigh_eigenvalues_stiff) * 1.0
- #print 'eigh eigenvalues of S bar', S_bar_1_eigh_eigenvalues
- #print 'eigh eigenvectors of S bar', S_bar_1_eigh_eigenvectors.T
+ #print('eigh eigenvalues of S bar', S_bar_1_eigh_eigenvalues)
+ #print('eigh eigenvectors of S bar', S_bar_1_eigh_eigenvectors.T)
- #print 'S bar eig recreate', S_bar_1_eigh_eigenvectors * numpy.matrix(numpy.diag(S_bar_1_eigh_eigenvalues)) * S_bar_1_eigh_eigenvectors.T
- #print 'S bar eig recreate error', (S_bar_1_eigh_eigenvectors * numpy.matrix(numpy.diag(S_bar_1_eigh_eigenvalues)) * S_bar_1_eigh_eigenvectors.T - S_bar_t[1])
+ #print('S bar eig recreate', S_bar_1_eigh_eigenvectors * numpy.matrix(numpy.diag(S_bar_1_eigh_eigenvalues)) * S_bar_1_eigh_eigenvectors.T)
+ #print('S bar eig recreate error', (S_bar_1_eigh_eigenvectors * numpy.matrix(numpy.diag(S_bar_1_eigh_eigenvalues)) * S_bar_1_eigh_eigenvectors.T - S_bar_t[1]))
S_bar_stiff = S_bar_1_eigh_eigenvectors * numpy.matrix(
numpy.diag(
S_bar_1_eigh_eigenvalues_stiff)) * S_bar_1_eigh_eigenvectors.T
- print 'Min u', -numpy.linalg.solve(TotalS_1, Totals_1)
- print 'Min x_hat', optimal_x_1
+ print('Min u', -numpy.linalg.solve(TotalS_1, Totals_1))
+ print('Min x_hat', optimal_x_1)
s_bar_t[1] = -s_t[1] - (S_bar_stiff + S_t[1]) * optimal_x_1
s_scalar_bar_t[1] = 0.5 * (
optimal_u_1.T * TotalS_1 * optimal_u_1 - optimal_x_1.T *
@@ -313,21 +313,22 @@
optimal_x_1) + optimal_u_1.T * Totals_1 - optimal_x_1.T * (
s_bar_t[1] + s_t[1]) - s_scalar_t[1] + Totals_scalar_1
- print 'optimal_u_1', optimal_u_1
- print 'TotalS_1', TotalS_1
- print 'Totals_1', Totals_1
- print 'Totals_scalar_1', Totals_scalar_1
- print 'overall cost 1', 0.5 * (
- optimal_u_1.T * TotalS_1 *
- optimal_u_1) + optimal_u_1.T * Totals_1 + Totals_scalar_1
- print 'overall cost 0', 0.5 * (x_hat_initial.T * S_t[0] * x_hat_initial
- ) + x_hat_initial.T * s_t[0] + s_scalar_t[0]
+ print('optimal_u_1', optimal_u_1)
+ print('TotalS_1', TotalS_1)
+ print('Totals_1', Totals_1)
+ print('Totals_scalar_1', Totals_scalar_1)
+ print(
+ 'overall cost 1', 0.5 * (optimal_u_1.T * TotalS_1 * optimal_u_1) +
+ optimal_u_1.T * Totals_1 + Totals_scalar_1)
+ print(
+ 'overall cost 0', 0.5 * (x_hat_initial.T * S_t[0] * x_hat_initial) +
+ x_hat_initial.T * s_t[0] + s_scalar_t[0])
- print 't forward 0'
- print 'x_hat_initial[ 0]: %s' % (x_hat_initial)
- print 'x_hat[%2d]: %s' % (0, x_hat.T)
- print 'x_hat_next[%2d]: %s' % (0, x_hat_next.T)
- print 'u[%2d]: %s' % (0, u_t.T)
+ print('t forward 0')
+ print('x_hat_initial[ 0]: %s' % (x_hat_initial))
+ print('x_hat[%2d]: %s' % (0, x_hat.T))
+ print('x_hat_next[%2d]: %s' % (0, x_hat_next.T))
+ print('u[%2d]: %s' % (0, u_t.T))
print('L[ 0]: %s' % (L_t[0], )).replace('\n', '\n ')
print('l[ 0]: %s' % (l_t[0], )).replace('\n', '\n ')
@@ -337,14 +338,14 @@
# TODO(austin): optimal_x_1 is x_hat
x_hat = -numpy.linalg.solve((S_t[1] + S_bar_stiff), (s_t[1] + s_bar_t[1]))
- print 'new xhat', x_hat
+ print('new xhat', x_hat)
S_bar_t[1] = S_bar_stiff
last_x_hat_t[1] = x_hat
for t in range(1, l):
- print 't forward', t
+ print('t forward', t)
u_t = L_t[t] * x_hat + l_t[t]
x_hat_next = discrete_dynamics(x_hat, u_t)
@@ -354,8 +355,8 @@
u_t)
c_bar_t = x_hat - A_bar_t * x_hat_next - B_bar_t * u_t
- print 'x_hat[%2d]: %s' % (t, x_hat.T)
- print 'x_hat_next[%2d]: %s' % (t, x_hat_next.T)
+ print('x_hat[%2d]: %s' % (t, x_hat.T))
+ print('x_hat_next[%2d]: %s' % (t, x_hat_next.T))
print('L[%2d]: %s' % (
t,
L_t[t],
@@ -364,7 +365,7 @@
t,
l_t[t],
)).replace('\n', '\n ')
- print 'u[%2d]: %s' % (t, u_t.T)
+ print('u[%2d]: %s' % (t, u_t.T))
print('A_bar_t[%2d]: %s' % (t, A_bar_t)).replace(
'\n', '\n ')
@@ -414,19 +415,19 @@
x_hat = -numpy.linalg.inv(S_t[l] + S_bar_t[l]) * (s_t[l] + s_bar_t[l])
for t in reversed(range(l)):
- print 't backward', t
+ print('t backward', t)
# TODO(austin): I don't think we can use L_t like this here.
# I think we are off by 1 somewhere...
u_t = L_bar_t[t + 1] * x_hat + l_bar_t[t + 1]
x_hat_prev = inverse_discrete_dynamics(x_hat, u_t)
- print 'x_hat[%2d]: %s' % (t, x_hat.T)
- print 'x_hat_prev[%2d]: %s' % (t, x_hat_prev.T)
+ print('x_hat[%2d]: %s' % (t, x_hat.T))
+ print('x_hat_prev[%2d]: %s' % (t, x_hat_prev.T))
print('L_bar[%2d]: %s' % (t + 1, L_bar_t[t + 1])).replace(
'\n', '\n ')
print('l_bar[%2d]: %s' % (t + 1, l_bar_t[t + 1])).replace(
'\n', '\n ')
- print 'u[%2d]: %s' % (t, u_t.T)
+ print('u[%2d]: %s' % (t, u_t.T))
# Now compute the linearized A, B, and C
# Start by doing it numerically, and then optimize.
A_t = numerical_jacobian_x(discrete_dynamics, x_hat_prev, u_t)