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/y2018/control_loops/python/arm_trajectory.py b/y2018/control_loops/python/arm_trajectory.py
index 9db7717..8b1ce47 100755
--- a/y2018/control_loops/python/arm_trajectory.py
+++ b/y2018/control_loops/python/arm_trajectory.py
@@ -1,4 +1,4 @@
-#!/usr/bin/python3
+#!/usr/bin/env python3
import copy
import numpy
@@ -617,7 +617,7 @@
self.distance_array = copy.copy(distance_array)
self.max_dvelocity_unfiltered = self.curvature_trajectory_pass(
dynamics, alpha_unitizer, distance_array, vmax)
- print 'Finished curvature pass'
+ print('Finished curvature pass')
self.max_dvelocity_unfiltered[0] = 0.0
self.max_dvelocity_unfiltered[-1] = 0.0
@@ -626,12 +626,12 @@
self.max_dvelocity_unfiltered, dynamics, alpha_unitizer,
distance_array, vmax)
self.max_dvelocity = self.max_dvelocity_back_pass
- print 'Finished backwards pass'
+ print('Finished backwards pass')
self.max_dvelocity_forward_pass = self.forward_trajectory_pass(
self.max_dvelocity_back_pass, dynamics, alpha_unitizer,
distance_array, vmax)
- print 'Finished forwards pass'
+ print('Finished forwards pass')
def interpolate_velocity(self, d, d0, d1, v0, v1):
if v0 + v1 > 0:
@@ -718,7 +718,7 @@
dynamics = Dynamics(dt)
trajectory = Trajectory(path_step_size)
- print 'Initialized path'
+ print('Initialized path')
distance_array = numpy.linspace(
0.0, trajectory.length(),
@@ -769,7 +769,7 @@
io0 += alpha[0, 0] * dd
io1 += alpha[1, 0] * dd
- print 'Iterated through path'
+ print('Iterated through path')
# Bounds on the accelerations of the two DOFs.
# We'll draw an oval to represent the actual bounds here.
@@ -786,7 +786,7 @@
distance_array,
vmax=vmax)
- print 'Computed trajectory'
+ print('Computed trajectory')
# Now, we can get acceleration, velocity, and position as a function of distance.
# Acceleration is best effort (derived from the velocities), but velocity
@@ -839,7 +839,7 @@
sim_dt = dt
- print 'Starting simulation'
+ print('Starting simulation')
# Now, we can start following the trajectory!
for t in numpy.arange(0.0, 1.0, sim_dt):
if goal_distance == trajectory.length():
@@ -914,9 +914,9 @@
# starting point.
if (numpy.abs(U) > vmax).any():
# Saturated. Let's do a binary search.
- print "Saturated."
+ print("Saturated.")
if (goal_distance - last_goal_distance) < 1e-8:
- print "Not bothering to move"
+ print("Not bothering to move")
# Avoid the divide by 0 when interpolating. Just don't move
# since we are saturated.
fraction_along_path = 0.0
@@ -936,7 +936,8 @@
fraction_along_path -= step_size
else:
fraction_along_path += step_size
- print "Fraction", fraction_along_path, "at", goal_distance, "rad,", t, "sec", goal_velocity
+ print("Fraction", fraction_along_path, "at", goal_distance, "rad,",
+ t, "sec", goal_velocity)
goal_distance = (
(goal_distance - last_goal_distance) * fraction_along_path +
@@ -986,7 +987,7 @@
goal_distance = trajectory.length()
goal_velocity = 0.0
- print 'Finished simulation'
+ print('Finished simulation')
pylab.figure()
pylab.title("Trajecotry")
pylab.plot(theta0_array, theta1_array, label="desired path")
diff --git a/y2018/control_loops/python/extended_lqr.py b/y2018/control_loops/python/extended_lqr.py
index 9a4705b..e19ee8f 100755
--- a/y2018/control_loops/python/extended_lqr.py
+++ b/y2018/control_loops/python/extended_lqr.py
@@ -66,14 +66,14 @@
final_B = numerical_jacobian_u(self.dynamics.discrete_dynamics,
numpy.matrix(numpy.zeros((4, 1))),
numpy.matrix(numpy.zeros((2, 1))))
- print 'Final A', final_A
- print 'Final B', final_B
+ print('Final A', final_A)
+ print('Final B', final_B)
K, self.S = controls.dlqr(final_A,
final_B,
self.Q,
self.R,
optimal_cost_function=True)
- print 'Final eig:', numpy.linalg.eig(final_A - final_B * K)
+ print('Final eig:', numpy.linalg.eig(final_A - final_B * K))
def final_cost(self, X, U):
"""Computes the final cost of being at X
@@ -111,8 +111,9 @@
numpy.matrix(self.num_states, self.num_states)
"""
zero_U = numpy.matrix(numpy.zeros((self.num_inputs, 1)))
- print 'S', self.S
- print 'Q_final', numerical_jacobian_x_x(self.final_cost, X_hat, zero_U)
+ print('S', self.S)
+ print('Q_final', numerical_jacobian_x_x(self.final_cost, X_hat,
+ zero_U))
return numerical_jacobian_x_x(self.final_cost, X_hat, zero_U)
def estimate_partial_cost_partial_x_final(self, X_hat):
@@ -448,8 +449,8 @@
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)
self.s_bar_t[1] = -self.s_t[1] - (S_bar_stiff +
self.S_t[1]) * optimal_x_1
self.s_scalar_bar_t[1] = 0.5 * (optimal_u_1.T * TotalS_1 * optimal_u_1 \
@@ -458,20 +459,20 @@
- optimal_x_1.T * (self.s_bar_t[1] + self.s_t[1]) \
- self.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 * self.S_t[0] * x_hat_initial) \
- + x_hat_initial.T * self.s_t[0] + self.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 * self.S_t[0] * x_hat_initial) \
+ + x_hat_initial.T * self.s_t[0] + self.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' % (self.L_t[0], )).replace('\n', '\n ')
print('l[ 0]: %s' % (self.l_t[0], )).replace('\n', '\n ')
@@ -482,14 +483,14 @@
# TODO(austin): optimal_x_1 is x_hat
x_hat = -numpy.linalg.solve((self.S_t[1] + S_bar_stiff),
(self.s_t[1] + self.s_bar_t[1]))
- print 'new xhat', x_hat
+ print('new xhat', x_hat)
self.S_bar_t[1] = S_bar_stiff
self.last_x_hat_t[1] = x_hat
for t in range(1, l):
- print 't forward', t
+ print('t forward', t)
u_t = self.L_t[t] * x_hat + self.l_t[t]
x_hat_next = self.dynamics.discrete_dynamics(x_hat, u_t)
@@ -499,8 +500,8 @@
self.dynamics.inverse_discrete_dynamics, x_hat_next, 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,
self.L_t[t],
@@ -509,7 +510,7 @@
t,
self.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 ')
@@ -561,20 +562,20 @@
* (self.s_t[l] + self.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 = self.L_bar_t[t + 1] * x_hat + self.l_bar_t[t + 1]
x_hat_prev = self.dynamics.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, self.L_bar_t[t + 1])).replace(
'\n', '\n ')
print('l_bar[%2d]: %s' % (t + 1, self.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(self.dynamics.discrete_dynamics,