diff --git a/agent/DQN_agent.py b/agent/DQN_agent.py index 199604b..80c3a16 100644 --- a/agent/DQN_agent.py +++ b/agent/DQN_agent.py @@ -46,10 +46,11 @@ class DQNAgent: self.history_buffer.pop(0) self.history_buffer.append(next_state) next_state = np.vstack(self.history_buffer) + total_reward += np.sum(reward * self.config.reward_weight) + reward = self.config.reward_shift_fn(reward) if not deterministic: self.replay.feed([state, action, reward, next_state, int(done)]) self.total_steps += 1 - total_reward += np.sum(reward * self.config.reward_weight) steps += 1 state = next_state if done: diff --git a/async_worker/actor_critic.py b/async_worker/actor_critic.py index 9ad9417..e821da6 100644 --- a/async_worker/actor_critic.py +++ b/async_worker/actor_critic.py @@ -33,6 +33,7 @@ class AdvantageActorCritic: steps += 1 total_reward += reward + reward = config.reward_shift_fn(reward) if deterministic: if terminal: diff --git a/async_worker/n_step_q.py b/async_worker/n_step_q.py index 3a5899b..47c4cfb 100644 --- a/async_worker/n_step_q.py +++ b/async_worker/n_step_q.py @@ -34,6 +34,7 @@ class NStepQLearning: steps += 1 total_reward += reward + reward = config.reward_shift_fn(reward) if deterministic: if terminal: diff --git a/async_worker/one_step_q.py b/async_worker/one_step_q.py index d190753..a37bc4c 100644 --- a/async_worker/one_step_q.py +++ b/async_worker/one_step_q.py @@ -34,6 +34,7 @@ class OneStepQLearning: steps += 1 total_reward += reward + reward = config.reward_shift_fn(reward) if deterministic: if terminal: diff --git a/async_worker/one_step_sarsa.py b/async_worker/one_step_sarsa.py index d763867..1434108 100644 --- a/async_worker/one_step_sarsa.py +++ b/async_worker/one_step_sarsa.py @@ -37,6 +37,7 @@ class OneStepSarsa: steps += 1 total_reward += reward + reward = config.reward_shift_fn(reward) if deterministic: if terminal: diff --git a/component/atari_wrapper.py b/component/atari_wrapper.py index ac16b9e..acd8e90 100644 --- a/component/atari_wrapper.py +++ b/component/atari_wrapper.py @@ -176,11 +176,6 @@ class ProcessFrame(gym.Wrapper): def _reset(self): return self.process_fn(self.env.reset()) -class ClippedRewardsWrapper(gym.Wrapper): - def _step(self, action): - obs, reward, done, info = self.env.step(action) - return obs, np.sign(reward), done, info - class NormalizeFrame(gym.Wrapper): def __init__(self, env=None): super(NormalizeFrame, self).__init__(env) diff --git a/component/task.py b/component/task.py index d52d009..8079d34 100644 --- a/component/task.py +++ b/component/task.py @@ -70,8 +70,7 @@ class PixelAtari(BasicTask): env = MaxAndSkipEnv(env, skip=frame_skip) if 'FIRE' in env.unwrapped.get_action_meanings(): env = FireResetEnv(env) - env = ProcessFrame(env, frame_size) - self.env = ClippedRewardsWrapper(env) + self.env = ProcessFrame(env, frame_size) self.action_dim = self.env.action_space.n def normalize_state(self, state): diff --git a/main.py b/main.py index 9b74e82..26ba78c 100644 --- a/main.py +++ b/main.py @@ -79,6 +79,7 @@ def dqn_pixel_atari(name): # config.network_fn = lambda optimizer_fn: DuelingNatureConvNet(config.history_length, n_actions, optimizer_fn) config.policy_fn = lambda: GreedyPolicy(epsilon=1.0, final_step=1000000, min_epsilon=0.1) config.replay_fn = lambda: Replay(memory_size=1000000, batch_size=32, dtype=np.uint8) + config.reward_shift_fn = lambda r: np.sign(r) config.discount = 0.99 config.target_network_update_freq = 10000 config.max_episode_length = 0 @@ -104,6 +105,7 @@ def async_pixel_atari(name): # config.worker = OneStepSarsa # config.worker = NStepQLearning config.worker = OneStepQLearning + config.reward_shift_fn = lambda r: np.sign(r) config.discount = 0.99 config.target_network_update_freq = 10000 config.max_episode_length = 10000 @@ -123,6 +125,7 @@ def a3c_pixel_atari(name): config.optimizer_fn = lambda params: torch.optim.Adam(params, lr=0.0001) config.network_fn = lambda: OpenAIActorCriticConvNet( config.history_length, task.env.action_space.n, LSTM=True) + config.reward_shift_fn = lambda r: np.sign(r) config.policy_fn = SamplePolicy config.worker = AdvantageActorCritic config.discount = 0.99 diff --git a/utils/config.py b/utils/config.py index 132a866..e7fa540 100644 --- a/utils/config.py +++ b/utils/config.py @@ -37,6 +37,7 @@ class Config: self.noise_decay_interval = 0 self.target_network_mix = 0.001 self.action_shift_fn = lambda a: a + self.reward_shift_fn = lambda r: r self.reward_weight = 1 self.hybrid_reward = False self.target_type = self.q_target