Fix gym state inverted ang_vel

mtheall · Oct 30, 2023 · 826ad12 · 826ad12
1 parent 7116bce
commit 826ad12
Show file tree

Hide file tree

Showing 3 changed files with 61 additions and 13 deletions.
diff --git a/python-mtheall/Arena.cpp b/python-mtheall/Arena.cpp
@@ -200,17 +200,10 @@ void assign (RocketSim::Python::PyArrayRef &array_, unsigned row_, unsigned col_
 	array_ (row_, col_ + 2) = value_.z ();
 }
 
-void assign (RocketSim::Python::PyArrayRef &array_,
-    unsigned col_,
-    btVector3 const &value_,
-    bool angular_ = false) noexcept
+void assign (RocketSim::Python::PyArrayRef &array_, unsigned col_, btVector3 const &value_) noexcept
 {
 	assign (array_, 0, col_, value_);
-
-	if (!angular_)
-		assign (array_, 1, col_, btVector3 (-value_.x (), -value_.y (), value_.z ()));
-	else
-		assign (array_, 1, col_, value_);
+	assign (array_, 1, col_, btVector3 (-value_.x (), -value_.y (), value_.z ()));
 }
 
 void assign (RocketSim::Python::PyArrayRef &array_, unsigned col_, btQuaternion const &value_) noexcept
@@ -257,13 +250,13 @@ void assign (RocketSim::Python::PyArrayRef &array_,
 	{
 		assign (array_, col_ + 0, btVector3 (state_->pos));
 		assign (array_, col_ + 7, btVector3 (state_->vel));
-		assign (array_, col_ + 10, btVector3 (state_->angVel), true);
+		assign (array_, col_ + 10, btVector3 (state_->angVel));
 	}
 	else
 	{
 		assign (array_, col_ + 0, rigidBody_.getWorldTransform ().getOrigin () * BT_TO_UU);
 		assign (array_, col_ + 7, rigidBody_.getLinearVelocity () * BT_TO_UU);
-		assign (array_, col_ + 10, rigidBody_.getAngularVelocity (), true);
+		assign (array_, col_ + 10, rigidBody_.getAngularVelocity ());
 	}
 
 	assign (array_, col_ + 3, rigidBody_.getOrientation ());

diff --git a/python-mtheall/regression_test.py b/python-mtheall/regression_test.py
@@ -144,5 +144,60 @@ def test_hoops_boost_pads(self):
     arena = rs.Arena(rs.GameMode.HOOPS)
     arena.get_boost_pads()
 
+  def test_gym_state_inverse(self):
+    arena = rs.Arena(rs.GameMode.SOCCAR)
+
+    car_a = arena.add_car(rs.Team.BLUE)
+    car_b = arena.add_car(rs.Team.ORANGE)
+
+    # cars should spawn mirror of each other
+    arena.reset_kickoff()
+
+    for i in range(100):
+      controls = rs.CarControls(throttle=1.0, yaw=1.0, pitch=1.0, roll=1.0, jump=True)
+      car_a.set_controls(controls)
+      car_b.set_controls(controls)
+      arena.step()
+
+      gym_state = arena.get_gym_state()
+
+      # pos
+      self.assertTrue(abs(gym_state[3][0][11] - gym_state[4][1][11]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][12] - gym_state[4][1][12]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][13] - gym_state[4][1][13]) < 1e-3)
+
+      # quat
+      q1 = gym_state[3][0][14:18]
+      q2 = gym_state[4][1][14:18]
+      t1 = np.amax(np.absolute(q1 + q2))
+      t2 = np.amax(np.absolute(q1 - q2))
+      self.assertTrue(min(t1, t2) < 1e-6)
+
+      # vel
+      self.assertTrue(abs(gym_state[3][0][18] - gym_state[4][1][18]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][19] - gym_state[4][1][19]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][20] - gym_state[4][1][20]) < 1e-3)
+
+      # ang_vel
+      self.assertTrue(abs(gym_state[3][0][21] - gym_state[4][1][21]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][22] - gym_state[4][1][22]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][23] - gym_state[4][1][23]) < 1e-3)
+
+      # mat3
+      self.assertTrue(abs(gym_state[3][0][24] - gym_state[4][1][24]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][25] - gym_state[4][1][25]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][26] - gym_state[4][1][26]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][27] - gym_state[4][1][27]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][28] - gym_state[4][1][28]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][29] - gym_state[4][1][29]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][30] - gym_state[4][1][30]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][31] - gym_state[4][1][31]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][32] - gym_state[4][1][32]) < 1e-3)
+
+      # pyr
+      self.assertTrue(abs(gym_state[3][0][33] - gym_state[4][1][33]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][34] - gym_state[4][1][34]) < 1e-3)
+      self.assertTrue(abs(gym_state[3][0][35] - gym_state[4][1][35]) < 1e-3)
+
 if __name__ == "__main__":
   unittest.main()
diff --git a/python-mtheall/unit_test.py b/python-mtheall/unit_test.py
@@ -1727,7 +1727,7 @@ def check_pyr(pyr: np.ndarray) -> bool:
           self.assertEqual(gym_state[0:3], invert_vector(ball_state.pos.as_numpy()))
           #self.assertEqual(gym_state[3:7], ball_state.quat.as_numpy())
           self.assertEqual(gym_state[7:10], invert_vector(ball_state.vel.as_numpy()))
-          self.assertEqual(gym_state[10:13], ball_state.ang_vel.as_numpy())
+          self.assertEqual(gym_state[10:13], invert_vector(ball_state.ang_vel.as_numpy()))
           self.assertEqual(invert_mat3(load_mat3(gym_state[13:22])).transpose(), ball_state.rot_mat.as_numpy())
           #self.assertEqual(gym_state[22:25], ball_state.pyr.as_numpy())
 
@@ -1799,7 +1799,7 @@ def check_pyr(pyr: np.ndarray) -> bool:
             gym_state = state[i][1]
             self.assertEqual(gym_state[11:14], invert_vector(car_state.pos.as_numpy()))
             self.assertEqual(gym_state[18:21], invert_vector(car_state.vel.as_numpy()))
-            self.assertEqual(gym_state[21:24], car_state.ang_vel.as_numpy())
+            self.assertEqual(gym_state[21:24], invert_vector(car_state.ang_vel.as_numpy()))
 
             if car_state.is_demoed:
               continue