Repository: WPI-MMR/gym_solo

Issue:
Create Rewards Factory
Similar to the observation factory. The reward factory should be able to take in `Reward` objects, similar to an `Observation` object, evaluate the rewards for the state, and combine them.

I'm thinking of making the final reward a linear combination of the registered rewards. With that in mind, consider the following example:

```python
r1 = Reward()
r2 = Reward()
r3 = Reward()

rf = RewardFactory()

r1.register_reward(-1, r1) 
r1.register_reward(.1, r2)
r1.register_reward(.9, r3)
```

Notice that `register_reward()` has two args: `weight: float` and `reward: gym_solo.core.rewards.Reward`. Thus, the final reward would evaluate to: `-r1() + 0.1 * r2() + 0.9 * r3()`. Figure that if you need functionality more elaborate than a linear combination, you should be offloading that processing into a `Reward` class.

@mahajanrevant thoughts on this? lmk if you think we need anything stronger than linear combinations in the `RewardFactory`.

Comments:

Patch:
diff --git a/gym_solo/core/rewards.py b/gym_solo/core/rewards.py
new file mode 100644
index 0000000..d364f98
--- /dev/null
+++ b/gym_solo/core/rewards.py
@@ -0,0 +1,63 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass
+from typing import List
+
+from gym_solo import solo_types
+
+
+class Reward(ABC):  
+  @abstractmethod
+  def compute(self) -> solo_types.reward:
+    """Compute the reward for the current state.
+
+    Returns:
+      solo_types.reward: The reward evalulated at the current state.
+    """
+    pass
+
+
+@dataclass
+class _WeightedReward:
+  reward: Reward
+  weight: float
+
+
+class RewardFactory:
+  """A factory to dynamically create rewards.
+  
+  Note that this factory is currently implemented to combined rewards via
+  a linear combination. For example, if the user wanted to register rewards
+  r1, r2, and r3, the final reward would be r1 + r2 + r3. 
+  
+  Obviously, you can add coefficients to the rewards, and that functionality 
+  is further explained in register_reward() below. 
+  
+  If you need more functionality then a linear combination (exponential 
+  temporal decay), then it's probably in your best interest to implement that
+  in a custom Reward.
+  """
+  def __init__(self):
+    """Create a new RewardFactory."""
+    self._rewards: List[_WeightedReward] = []
+
+  def register_reward(self, weight: float, reward: Reward):
+    """Register a reward to be computed per state.
+
+    Args:
+      weight (float): The weight to be applied to this reward when it is 
+        is combined linearly with the other rewards. The domain for this
+        value is (-∞, ∞).
+      reward (Reward): A Reward object which .compute() will be called on at
+        reward computation time.
+    """
+    self._rewards.append(_WeightedReward(reward=reward, weight=weight))
+
+  def get_reward(self) -> float:
+    """Evaluate the current state and get the combined reward.
+
+    Returns:
+      float: The reward from the current state. Note that this reward is a 
+      combination of multiple atomic sub-rewards, as explained by the 
+      strategies earlier.
+    """
+    return sum(wr.weight * wr.reward.compute() for wr in self._rewards)
\ No newline at end of file
diff --git a/gym_solo/solo_types.py b/gym_solo/solo_types.py
index 85dc7ae..365d7e8 100644
--- a/gym_solo/solo_types.py
+++ b/gym_solo/solo_types.py
@@ -3,4 +3,7 @@ from typing import List
 import numpy as np
 
 # A state observation
-obs = np.ndarray
\ No newline at end of file
+obs = np.ndarray
+
+# A reward after a step
+reward = float
\ No newline at end of file


Test patch:
diff --git a/gym_solo/core/test_rewards_factory.py b/gym_solo/core/test_rewards_factory.py
new file mode 100644
index 0000000..fd70bea
--- /dev/null
+++ b/gym_solo/core/test_rewards_factory.py
@@ -0,0 +1,35 @@
+import unittest
+from gym_solo.core import rewards
+
+from parameterized import parameterized
+
+
+class TestReward(rewards.Reward):
+  def __init__(self, return_value):
+    self._return_value = return_value
+
+  def compute(self):
+    return self._return_value
+
+
+class TestRewardsFactory(unittest.TestCase):
+  def test_empty(self):
+    rf = rewards.RewardFactory()
+    self.assertListEqual(rf._rewards, [])
+
+  @parameterized.expand([
+    ('single', {1: 2.5}, 2.5),
+    ('two_happy', {1: 1, 2: 2}, 5),
+    ('0-weight', {0: 1, 2: 2}, 4),
+    ('negative-weight', {-1: 1, 2: 2}, 3),
+    ('three', {1: 1, 2: 2, 3: 3}, 14),
+  ])
+  def test_register_and_compute(self, name, rewards_dict, expected_reward):
+    rf = rewards.RewardFactory()
+    for weight, reward in rewards_dict.items():
+      rf.register_reward(weight, TestReward(reward))
+    self.assertEqual(rf.get_reward(), expected_reward)
+
+
+if __name__ == '__main__':
+  unittest.main()
\ No newline at end of file