Repository: WPI-MMR/gym_solo

Issue:
A reward for getting the robot from flat on the ground to the home position
It would be easier to create a reward that is geared towards getting the robot from flat on the ground to the home position.
The home position would be when all the legs are straight and the body is parallel to the ground. In theory, this is a much easier task as compared to getting the robot to completely stand up on two legs while being flat on the ground. 

Discussion:

Patch:
diff --git a/.gitignore b/.gitignore
index 81bd9dc..3327ae7 100644
--- a/.gitignore
+++ b/.gitignore
@@ -2,6 +2,7 @@
 .vscode/
 venv-*/
 *.egg-info/
+*.swp
 
 # Byte-compiled / optimized / DLL files
 __pycache__/
diff --git a/examples/solo8_vanilla/observation_dump.py b/examples/solo8_vanilla/observation_dump.py
index 0873f1d..6a00fee 100644
--- a/examples/solo8_vanilla/observation_dump.py
+++ b/examples/solo8_vanilla/observation_dump.py
@@ -11,6 +11,7 @@ import gym_solo
 from gym_solo.envs import solo8v2vanilla
 from gym_solo.core import obs
 from gym_solo.core import rewards
+from gym_solo.core import termination as terms
 
 
 if __name__ == '__main__':
@@ -19,6 +20,8 @@ if __name__ == '__main__':
 
   env.obs_factory.register_observation(obs.TorsoIMU(env.robot))
   env.reward_factory.register_reward(1,rewards.UprightReward(env.robot))
+  env.termination_factory.register_termination(
+    terms.TimeBasedTermination(100000000))
 
   try:
     print("""\n
diff --git a/gym_solo/core/rewards.py b/gym_solo/core/rewards.py
index 7d696bb..44073c9 100644
--- a/gym_solo/core/rewards.py
+++ b/gym_solo/core/rewards.py
@@ -105,3 +105,42 @@ class UprightReward(Reward):
     _, 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
+
+
+class HomePositionReward(Reward):
+  """Rewards the robot for being in the home position. Currently, this rewards
+  the robot for being orientated properly and being as tall as possible. This
+  will require some experimentation--if the robot ends up jumping to maximize
+  the reward, this might need to be modified to a more intelligent height 
+  reward.
+  """
+  # Found via experimentation--should be close enough to prevent any 
+  # non-trivial errors
+  _quad_standing_height = 0.3
+  _max_angle = np.pi
+
+  def __init__(self, robot_id: int):
+    """Create a new HomePositionReward.
+
+    Args:
+      robot_id (int): The PyBullet body id of the robot
+    """
+    self._robot_id = robot_id
+
+  def compute(self) -> float:
+    """Compute the HomePositionReward for the current state. 
+
+    Returns:
+      float: A real-valued in [0, 1], where 1 is in the home position.
+    """
+    (unused_x, unused_y, z), quat = p.getBasePositionAndOrientation(
+      self._robot_id)
+    theta_x, theta_y, unused_z = np.array(p.getEulerFromQuaternion(quat))
+
+    x_reward = self._max_angle - abs(theta_x)
+    y_reward = self._max_angle - abs(theta_y)
+
+    orientation_reward = (x_reward + y_reward) / (2 * self._max_angle)
+    height_reward = z / self._quad_standing_height
+    
+    return 0.25 * orientation_reward + 0.75 * height_reward
\ No newline at end of file


Test patch:
diff --git a/gym_solo/core/test_rewards_factory.py b/gym_solo/core/test_rewards.py
similarity index 56%
rename from gym_solo/core/test_rewards_factory.py
rename to gym_solo/core/test_rewards.py
index b855bd3..15c4b3a 100644
--- a/gym_solo/core/test_rewards_factory.py
+++ b/gym_solo/core/test_rewards.py
@@ -55,5 +55,43 @@ class TestUprightReward(unittest.TestCase):
     self.assertEqual(reward.compute(), expected_reward)
 
 
+class TestHomePositionReward(unittest.TestCase):
+  def test_init(self):
+    robot_id = 0
+    r = rewards.HomePositionReward(robot_id)
+    self.assertEqual(robot_id, r._robot_id)
+
+  @mock.patch('pybullet.getBasePositionAndOrientation')
+  @mock.patch('pybullet.getEulerFromQuaternion')
+  def test_computation(self, mock_euler, mock_orien):
+    r = rewards.HomePositionReward(0)
+
+    def _mocked_reward(pos, orien):
+      mock_orien.return_value = pos, None
+      mock_euler.return_value = orien
+      return r.compute()
+
+    # Starting position
+    start_reward = _mocked_reward((0, 0, 0), (0, 0, 0))
+
+    # Rotated about the z-axis
+    z_rot_reward = _mocked_reward((0, 0, 0), (0, 0, .5))
+
+    # Semi standing up at an angle, 
+    standing_skewed_reward = _mocked_reward((0, 0, .15), (0, np.pi / 4, 0))
+
+    # Semi standing up completely flat (should be better than skewed)
+    standing_straight_reward = _mocked_reward((0, 0, .15), (0, 0, 0))
+
+    # Home position
+    home_reward = _mocked_reward(
+      (0, 0, rewards.HomePositionReward._quad_standing_height), (0, 0, 0))
+    
+    self.assertEqual(start_reward, z_rot_reward)
+    self.assertLess(start_reward, standing_straight_reward)
+    self.assertLess(standing_skewed_reward, standing_straight_reward)
+    self.assertLess(standing_straight_reward, home_reward)
+    
+
 if __name__ == '__main__':
   unittest.main()
\ No newline at end of file