def test_output_time(self): env = OneCharMemory(n=2, num_steps=3, output_time=True) self.assertEqual(env.observation_space.flat_dim, 3 + 4) obs = env.reset() action = np.zeros(3) time = np.zeros(4) time[0] = 1 self.assertNpEqual(obs[-4:], time) next_ob, reward, terminal, _ = env.step(action) time[0] = 0 time[1] = 1 self.assertNpEqual(next_ob[-4:], time) next_ob, reward, terminal, _ = env.step(action) time[1] = 0 time[2] = 1 self.assertNpEqual(next_ob[-4:], time) next_ob, reward, terminal, _ = env.step(action) time[2] = 0 time[3] = 1 self.assertNpEqual(next_ob[-4:], time) self.assertTrue(terminal) obs = env.reset() time[3] = 0 time[0] = 1 self.assertNpEqual(obs[-4:], time)
def test_episode_length_one(self): env = OneCharMemory(n=3, num_steps=1, reward_for_remembering=0.5) action = env.reset() next_ob, reward, terminal, _ = env.step(action) self.assertTrue(terminal) self.assertNpArraysEqual(next_ob, np.array([1, 0, 0, 0])) self.assertAlmostEqual(reward, 0.5)
def test_episode_length_is_right(self): env = OneCharMemory(num_steps=4) action = env.reset() for _ in range(3): _, _, terminal, _ = env.step(action) self.assertFalse(terminal) _, _, terminal, _ = env.step(action) self.assertTrue(terminal)
def test_reward_for_optimal_input_is_correct(self): env = OneCharMemory(n=3, num_steps=4, reward_for_remembering=.25) init_obs = env.reset() action = np.zeros((4, 1)) action[0] = 1 for _ in range(3): next_ob, reward, terminal, _ = env.step(action) self.assertAlmostEqual(reward, 0.) next_ob, reward, terminal, _ = env.step(init_obs) self.assertAlmostEqual(reward, 0.25)
def test_reward_for_wrong_input_is_correct(self): env = OneCharMemory(n=3, num_steps=4) init_obs = env.reset() action = init_obs for _ in range(3): next_ob, reward, terminal, _ = env.step(action) self.assertTrue(reward < 0) based_reward = reward action = np.zeros((4,)) next_ob, reward, terminal, _ = env.step(action) self.assertTrue(reward < 0)
def test_output_target(self): env = OneCharMemory(n=2, num_steps=5, output_target_number=True) self.assertEqual(env.observation_space.flat_dim, 4) obs = env.reset() target_number = np.argmax(obs[:3]) action = np.zeros(3) next_ob, reward, terminal, _ = env.step(action) expected_obs = np.hstack((obs[:3], [target_number])) expected_next_obs = np.array([1, 0, 0, target_number]) self.assertNpArraysEqual(obs, expected_obs) self.assertNpArraysEqual(next_ob, expected_next_obs) next_ob, reward, terminal, _ = env.step(action) self.assertNpArraysEqual(next_ob, expected_next_obs)
def test_reward_for_wrong_input_is_correct(self): env = OneCharMemory(n=3, num_steps=4, reward_for_remembering=10) init_obs = env.reset() action = init_obs for _ in range(3): next_ob, reward, terminal, _ = env.step(action) self.assertEqual(reward, 0) action = np.zeros((4, )) if init_obs[2]: action[3] = 1 else: action[2] = 1 next_ob, reward, terminal, _ = env.step(action) self.assertEqual(reward, 0)
def test_init_state_is_one_hot(self): env = OneCharMemory(n=3, num_steps=4) init_state = env.reset() self.assertEqual(init_state.shape, (4,))