commit f4d819c122b5332d89835249008552b0e1619fd2
author justinT21 <jjturcot@gmail.com> Mon Oct 21 13:51:58 2024 -0700
committer Justin Turcotte <jjturcot@gmail.com> Sat Oct 26 15:12:02 2024 -0700
tree 0457dc633922a3f5bbc4fcb32d50b64a133cba73
parent 96c0f25f69a04f4cfb76441500b89433ee2ae859

Fix goal not being passed to gradient from data

Signed-off-by: justinT21 <jjturcot@gmail.com>
Change-Id: I6a66c96b347f1cd23a04d66bac0f59a0ee12db36

frc971/control_loops/swerve/velocity_controller/train.py

diff --git a/frc971/control_loops/swerve/velocity_controller/train.py b/frc971/control_loops/swerve/velocity_controller/train.py
index b1ad9a8..cf54bc4 100644
--- a/frc971/control_loops/swerve/velocity_controller/train.py
+++ b/frc971/control_loops/swerve/velocity_controller/train.py
@@ -216,8 +216,9 @@
                        data: ArrayLike):
     """Computes the Soft Actor-Critic loss for alpha."""
     observations1 = data['observations1']
+    R = data['goals']
     pi, logp_pi, _, _ = jax.lax.stop_gradient(
-        state.pi_apply(rng=rng, params=params, observation=observations1))
+        state.pi_apply(rng=rng, params=params, R=R, observation=observations1))
 
     return (-jax.numpy.exp(params['logalpha']) *
             (logp_pi + state.target_entropy)).mean()