Added reward_scale and reward_shift to complex env wrappers and rectified it in the toy env

Author: RaghuSpaceRajan

experiments/a3c_lstm_tune_hps.py renamed to experiments/a3c_lstm_tune_hps.txt

@@ -20,6 +20,8 @@ class GymEnvWrapper(gym.Env):
         transition noise (for discrete environments)
         reward delay
         reward noise
+        reward scale
+        reward shift
         image_transforms
 
     The wrapper is pretty general and can be applied to any Gym Environment. The environment should be instantiated and passed as the 1st argument to the __init__ method of this class. If using this wrapper with Atari, additional keys may be added specifying either atari_preprocessing = True or wrap_deepmind_ray = True. These would use the AtariPreprocessing wrapper from OpenAI Gym or wrap_deepmind() wrapper from Ray Rllib.
@@ -92,6 +94,17 @@ def __init__(self, env, **config):
         else:
             self.reward_noise = None
 
+        if "reward_scale" not in config:
+            self.reward_scale = 1.0
+        else:
+            self.reward_scale = config["reward_scale"]
+
+        if "reward_shift" not in config:
+            self.reward_shift = 0.0
+        else:
+            self.reward_shift = config["reward_shift"]
+
+
         if "image_transforms" not in config:
             self.image_transforms = False
         else:
@@ -360,6 +373,8 @@ def step(self, action):
         self.total_abs_noise_in_reward_episode += np.abs(noise_in_reward)
         self.total_reward_episode += reward
         reward += noise_in_reward
+        reward *= self.reward_scale
+        reward += self.reward_shift
 
         return next_state, reward, done, info
 

experiments/a3c_tune_hps.py renamed to experiments/a3c_tune_hps.txt

@@ -1877,12 +1877,12 @@ def reward_function(self, state, action):
                     if list(new_relevant_state) == self.target_point:
                         reward += 1.0
 
-        reward *= self.reward_scale
         noise_in_reward = self.reward_noise(self.np_random) if self.reward_noise else 0
         # #random ###TODO Would be better to parameterise this in terms of state, action and time_step as well. Would need to change implementation to have a queue for the rewards achieved and then pick the reward that was generated delay timesteps ago.
         self.total_abs_noise_in_reward_episode += np.abs(noise_in_reward)
         self.total_reward_episode += reward
         reward += noise_in_reward
+        reward *= self.reward_scale
         reward += self.reward_shift
         return reward
 

experiments/dqn_tune_hps.py renamed to experiments/dqn_tune_hps.txt

@@ -68,6 +68,108 @@ def test_r_delay(self):
         print("total_reward:", total_reward)
         aew.reset()
 
+    def test_r_shift(self):
+        """ """
+        print("\033[32;1;4mTEST_REWARD_SHIFT\033[0m")
+        config = {
+            "AtariEnv": {
+                "game": "beam_rider",  # "breakout",
+                "obs_type": "image",
+                "frameskip": 1,
+            },
+            "reward_shift": 1,
+            # "GymEnvWrapper": {
+            "atari_preprocessing": True,
+            "frame_skip": 4,
+            "grayscale_obs": False,
+            "state_space_type": "discrete",
+            "action_space_type": "discrete",
+            "seed": 0,
+            # },
+            # 'seed': 0, #seed
+        }
+
+        # config["log_filename"] = log_filename
+
+        from gym.envs.atari import AtariEnv
+
+        ae = AtariEnv(**{"game": "beam_rider", "obs_type": "image", "frameskip": 1})
+        aew = GymEnvWrapper(ae, **config)
+        ob = aew.reset()
+        print("observation_space.shape:", ob.shape)
+        # print(ob)
+        total_reward = 0.0
+        for i in range(200):
+            act = aew.action_space.sample()
+            next_state, reward, done, info = aew.step(act)
+            print("step, reward, done, act:", i, reward, done, act)
+            if i == 153 or i == 158:
+                assert reward == 45.0, (
+                    "Shifted reward in step: "
+                    + str(i)
+                    + " should have been 45.0."
+                )
+            if i == 154 or i == 160:
+                assert reward == 1.0, (
+                    "Shifted reward in step: "
+                    + str(i)
+                    + " should have been 1.0."
+                )
+            total_reward += reward
+        print("total_reward:", total_reward)
+        aew.reset()
+
+    def test_r_scale(self):
+        """ """
+        print("\033[32;1;4mTEST_REWARD_SHIFT\033[0m")
+        config = {
+            "AtariEnv": {
+                "game": "beam_rider",  # "breakout",
+                "obs_type": "image",
+                "frameskip": 1,
+            },
+            "reward_scale": 2,
+            # "GymEnvWrapper": {
+            "atari_preprocessing": True,
+            "frame_skip": 4,
+            "grayscale_obs": False,
+            "state_space_type": "discrete",
+            "action_space_type": "discrete",
+            "seed": 0,
+            # },
+            # 'seed': 0, #seed
+        }
+
+        # config["log_filename"] = log_filename
+
+        from gym.envs.atari import AtariEnv
+
+        ae = AtariEnv(**{"game": "beam_rider", "obs_type": "image", "frameskip": 1})
+        aew = GymEnvWrapper(ae, **config)
+        ob = aew.reset()
+        print("observation_space.shape:", ob.shape)
+        # print(ob)
+        total_reward = 0.0
+        for i in range(200):
+            act = aew.action_space.sample()
+            next_state, reward, done, info = aew.step(act)
+            print("step, reward, done, act:", i, reward, done, act)
+            if i == 153 or i == 158:
+                assert reward == 88.0, (
+                    "Scaled reward in step: "
+                    + str(i)
+                    + " should have been 88.0."
+                )
+            if i == 154 or i == 160:
+                assert reward == 0.0, (
+                    "Scaled reward in step: "
+                    + str(i)
+                    + " should have been 0.0."
+                )
+            total_reward += reward
+        print("total_reward:", total_reward)
+        aew.reset()
+
     def test_r_delay_ray_frame_stack(self):
         """
         Uses wrap_deepmind_ray to frame stack Atari