Repository: WPI-MMR/gym_solo

Issue:
Create Upright Reward
At the very least, we know that our reward will need to be somewhat correlated with how upright the robot is. I'm reaching out to a roommate rn to verify the math, but I feel like the cosine similarity between the torso IMU vector and straight upright should basically be it, correct? @andrew103 and @mahajanrevant. Choosing cosine similarity over inner product because it's automatically normalized to be in [-1, 1]


Discussion:

Patch:
diff --git a/gym_solo/core/rewards.py b/gym_solo/core/rewards.py
index 35284e6..7d696bb 100644
--- a/gym_solo/core/rewards.py
+++ b/gym_solo/core/rewards.py
@@ -2,6 +2,9 @@ from abc import ABC, abstractmethod
 from dataclasses import dataclass
 from typing import List
 
+import numpy as np
+import pybullet as p
+
 from gym_solo import solo_types
 
 
@@ -66,4 +69,39 @@ class RewardFactory:
     if not self._rewards:
       raise ValueError('Need to register at least one reward instance')
 
-    return sum(wr.weight * wr.reward.compute() for wr in self._rewards)
\ No newline at end of file
+    return sum(wr.weight * wr.reward.compute() for wr in self._rewards)
+
+
+class UprightReward(Reward):
+  """A reward for being fully upright. Note that this is technically only
+  designed for the Solo8v2Vanilla and should be extended in the future to 
+  support more models.
+
+  The reward's is in the interval [-1, 1], where 1 indicates that the torso
+  is upright, while -1 means that it is upside down. Observe that this means
+  that the reward is orientation-specific.
+  """
+  _fully_upright = -np.pi / 2
+
+  def __init__(self, robot_id: int):
+    """Create a new UprightReward.
+
+    Args:
+      robot_id ([int]): The PyBullet body id of the robot
+    """
+    self._robot_id = robot_id
+
+  def compute(self) -> float:
+    """Compute the UprightReward for the current state. 
+    
+    As this reward function is specifically designed for the solo8v2vanilla 
+    model, this means that the model just needs to be rotated -pi/2 radians 
+    about the y axis.
+
+    Returns:
+      float: A real-valued number in [-1, 1], where 1 means perfectly upright 
+      whilst -1 means that the robot is literally upside down. 
+    """
+    _, quat = p.getBasePositionAndOrientation(self._robot_id)
+    unused_x, y, unused_z = np.array(p.getEulerFromQuaternion(quat))
+    return self._fully_upright * y / self._fully_upright ** 2


Test patch:
diff --git a/gym_solo/core/test_rewards_factory.py b/gym_solo/core/test_rewards_factory.py
index c29bc5a..8c6f420 100644
--- a/gym_solo/core/test_rewards_factory.py
+++ b/gym_solo/core/test_rewards_factory.py
@@ -2,6 +2,9 @@ import unittest
 from gym_solo.core import rewards
 
 from parameterized import parameterized
+from unittest import mock
+
+import numpy as np
 
 
 class TestReward(rewards.Reward):
@@ -34,5 +37,30 @@ class TestRewardsFactory(unittest.TestCase):
     self.assertEqual(rf.get_reward(), expected_reward)
 
 
+class TestUprightReward(unittest.TestCase):
+  def test_init(self):
+    robot_id = 0
+    r = rewards.UprightReward(robot_id)
+    self.assertEqual(robot_id, r._robot_id)
+
+  @parameterized.expand([
+    ('flat', (0, 0, 0), 0),
+    ('upright', (0, 90, 0), -1.),
+    ('upside down', (0, -90, 0), 1.),
+    ('y dependence', (-45, 90, -90), -1.),
+  ])
+  @mock.patch('pybullet.getBasePositionAndOrientation')
+  @mock.patch('pybullet.getEulerFromQuaternion')
+  def test_computation(self, name, orien, expected_reward, mock_euler,
+                       mock_orien):
+    mock_orien.return_value = None, None
+
+    orien_radians = tuple(i * np.pi / 180 for i in orien)
+    mock_euler.return_value = orien_radians
+
+    reward = rewards.UprightReward(None)
+    self.assertEqual(reward.compute(), expected_reward)
+
+
 if __name__ == '__main__':
   unittest.main()
\ No newline at end of file