Merge pull request #21 from pockerman/add_q_learning_algorithm

Add q learning algorithm

Author: pockerman

src/algorithms/anonymity_a2c_ray.py

@@ -6,7 +6,7 @@
 import ray.rllib.agents.a3c as a3c
 from ray.tune.logger import pretty_print
 from ray.rllib.env.env_context import EnvContext
-from src.spaces.environment import TimeStep, StepType
+from src.spaces.discrete_state_environment import TimeStep, StepType
 from src.spaces.observation_space import ObsSpace
 
 

src/algorithms/q_learning.py

@@ -11,8 +11,11 @@
 Env = TypeVar('Env')
 Policy = TypeVar('Policy')
 
-class QLearnConfig(object):
 
+class QLearnConfig(object):
+    """
+    Configuration  for Q-learning
+    """
     def __init__(self):
         self.gamma: float = 1.0
         self.alpha: float = 0.1
@@ -21,16 +24,15 @@ def __init__(self):
 
 
 class QLearning(WithMaxActionMixin):
+    """
+    Q-learning algorithm implementation
+    """
 
     def __init__(self, algo_config: QLearnConfig):
         super(QLearning, self).__init__()
         self.q_table = {}
         self.config = algo_config
 
-        # monitor performance
-        self.total_rewards: np.array = None
-        self.iterations_per_episode = []
-
     @property
     def name(self) -> str:
         return "QLearn"
@@ -40,8 +42,8 @@ def actions_before_training(self, env: Env, **options):
         if self.config.policy is None:
             raise InvalidParamValue(param_name="policy", param_value="None")
 
-        for state in range(env.observation_space.n):
-            for action in range(env.action_space.n):
+        for state in range(1, env.n_states):
+            for action in range(env.n_actions):
                 self.q_table[state, action] = 0.0
 
     def actions_after_episode_ends(self, **options):
@@ -57,8 +59,9 @@ def actions_after_episode_ends(self, **options):
     def train(self, env: Env, **options) -> tuple:
 
         # episode score
-        episode_score = 0  # initialize score
+        episode_score = 0
         counter = 0
+        total_distortion = 0
 
         time_step = env.reset()
         state = time_step.observation
@@ -70,24 +73,28 @@ def train(self, env: Env, **options) -> tuple:
 
             action = env.get_action(action_idx)
 
+            if action.action_type.name == "GENERALIZE" and action.column_name == "salary":
+                print("Attempt to generalize salary")
+            else:
+                print(action.action_type.name, " on ", action.column_name)
+
             # take action A, observe R, S'
             next_time_step = env.step(action)
             next_state = next_time_step.observation
             reward = next_time_step.reward
 
-            next_state_id = next_state.state_id if next_state is not None else None
-
             # add reward to agent's score
-            episode_score += next_time_step.reward
-            self._update_Q_table(state=state.state_id, action=action_idx, reward=reward,
-                                 next_state=next_state_id, n_actions=env.action_space.n)
+            episode_score += reward
+            self._update_Q_table(state=state, action=action_idx, reward=reward,
+                                 next_state=next_state, n_actions=env.n_actions)
             state = next_state  # S <- S'
             counter += 1
+            total_distortion += next_time_step.info["total_distortion"]
 
             if next_time_step.last():
                 break
 
-        return episode_score, counter
+        return episode_score, total_distortion, counter
 
     def _update_Q_table(self, state: int, action: int, n_actions: int, reward: float, next_state: int = None) -> None:
         """

src/algorithms/trainer.py

@@ -17,10 +17,26 @@ def __init__(self, env: Env,  agent: Agent, configuration: dir) -> None:
         self.agent = agent
         self.configuration = configuration
         # monitor performance
-        self.total_rewards: np.array = None
+        self.total_rewards: np.array = np.zeros(configuration['n_episodes'])
         self.iterations_per_episode = []
+        self.total_distortions = []
+
+    def avg_rewards(self) -> np.array:
+        """
+        Returns the average reward per episode
+        :return:
+        """
+        avg = np.zeros(self.configuration['n_episodes'])
+
+        for i in range(self.total_rewards.shape[0]):
+            avg[i] = self.total_rewards[i] / self.iterations_per_episode[i]
+        return avg
 
     def actions_before_training(self):
+        """
+        Any actions to perform before training begins
+        :return:
+        """
         self.total_rewards: np.array = np.zeros(self.configuration['n_episodes'])
         self.iterations_per_episode = []
 
@@ -29,27 +45,32 @@ def actions_before_training(self):
     def actions_after_episode_ends(self, **options):
         self.agent.actions_after_episode_ends(**options)
 
+        if options["episode_idx"] % self.configuration['output_msg_frequency'] == 0:
+            if self.env.config.distorted_set_path is not None:
+                self.env.save_current_dataset(options["episode_idx"])
+
     def train(self):
 
         print("{0} Training agent {1}".format(INFO, self.agent.name))
         self.actions_before_training()
 
         for episode in range(0, self.configuration["n_episodes"]):
-            print("INFO: Episode {0}/{1}".format(episode, self.configuration["n_episodes"]))
+            print("{0} On episode {1}/{2}".format(INFO, episode, self.configuration["n_episodes"]))
 
             # reset the environment
             ignore = self.env.reset()
 
             # train for a number of iterations
-            episode_score, n_itrs = self.agent.train(self.env)
+            episode_score, total_distortion, n_itrs = self.agent.train(self.env)
 
-            if episode % self.configuration['output_msg_frequency'] == 0:
-                print("{0}: On episode {1} training finished with  "
-                      "{2} iterations. Total reward={3}".format(INFO, episode, n_itrs, episode_score))
+            print("{0} Episode score={1}, episode total distortion {2}".format(INFO, episode_score, total_distortion / n_itrs))
+
+            #if episode % self.configuration['output_msg_frequency'] == 0:
+            print("{0} Episode finished after {1} iterations".format(INFO, n_itrs))
 
             self.iterations_per_episode.append(n_itrs)
             self.total_rewards[episode] = episode_score
-
+            self.total_distortions.append(total_distortion)
             self.actions_after_episode_ends(**{"episode_idx": episode})
 
         print("{0} Training finished for agent {1}".format(INFO, self.agent.name))

src/apps/qlearning_on_mock.py

@@ -1,8 +1,11 @@
+import matplotlib.pyplot as plt
+import numpy as np
+
 from src.algorithms.q_learning import QLearning, QLearnConfig
 from src.algorithms.trainer import Trainer
-from src.utils.string_distance_calculator import DistanceType
-from src.spaces.actions import ActionSuppress, ActionIdentity, ActionGeneralize, ActionTransform
-from src.spaces.environment import Environment, EnvConfig
+from src.utils.string_distance_calculator import StringDistanceType
+from src.spaces.actions import ActionSuppress, ActionIdentity, ActionStringGeneralize, ActionTransform
+from src.spaces.discrete_state_environment import Environment, EnvConfig
 from src.spaces.action_space import ActionSpace
 from src.datasets.datasets_loaders import MockSubjectsLoader
 from src.utils.reward_manager import RewardManager
@@ -11,45 +14,74 @@
 from src.utils.numeric_distance_type import NumericDistanceType
 
 
+def plot_running_avg(avg_rewards):
+
+    running_avg = np.empty(avg_rewards.shape[0])
+    for t in range(avg_rewards.shape[0]):
+        running_avg[t] = np.mean(avg_rewards[max(0, t-100) : (t+1)])
+    plt.plot(running_avg)
+    plt.xlabel("Number of episodes")
+    plt.ylabel("Reward")
+    plt.title("Running average")
+    plt.show()
+
+def get_ethinicity_hierarchies():
+
+    ethnicity_hierarchy = SerialHierarchy()
+    ethnicity_hierarchy.add("Mixed White/Asian", values=["Mixed", '*'])
+    ethnicity_hierarchy.add("Chinese", values=["Asian", '*'])
+    ethnicity_hierarchy.add("Indian", values=["Asian", '*'])
+    ethnicity_hierarchy.add("Mixed White/Black African", values=["Mixed", '*'])
+    ethnicity_hierarchy.add("Black African", values=["Black", '*'])
+    ethnicity_hierarchy.add("Asian other", values=["Asian", "*"])
+    ethnicity_hierarchy.add("Black other", values=["Black", "*"])
+    ethnicity_hierarchy.add("Mixed White/Black Caribbean", values=["Mixed", "*"])
+    ethnicity_hierarchy.add("Mixed other", values=["Mixed", "*"])
+    ethnicity_hierarchy.add("Arab", values=["Asian", "*"])
+    ethnicity_hierarchy.add("White Irish", values=["White", "*"])
+    ethnicity_hierarchy.add("Not stated", values=["Not stated", "*"])
+    ethnicity_hierarchy.add("White Gypsy/Traveller", values=["White", "*"])
+    ethnicity_hierarchy.add("White British", values=["White", "*"])
+    ethnicity_hierarchy.add("Bangladeshi", values=["Asian", "*"])
+    ethnicity_hierarchy.add("White other", values=["White", "*"])
+    ethnicity_hierarchy.add("Black Caribbean", values=["Black", "*"])
+    ethnicity_hierarchy.add("Pakistani", values=["Asian", "*"])
+
+    return ethnicity_hierarchy
+
+
 if __name__ == '__main__':
 
     EPS = 1.0
     GAMMA = 0.99
     ALPHA = 0.1
+    N_EPISODES = 100
 
     # load the dataset
     ds = MockSubjectsLoader()
 
+    # generalization table for the ethnicity column
+    ethinicity_table = get_ethinicity_hierarchies()
+
     # specify the action space. We need to establish how these actions
     # are performed
-    action_space = ActionSpace(n=4)
-
-    generalization_table = {"Mixed White/Asian": SerialHierarchy(values=["Mixed", ]),
-                            "Chinese": SerialHierarchy(values=["Asian", ]),
-                            "Indian": SerialHierarchy(values=["Asian", ]),
-                            "Mixed White/Black African": SerialHierarchy(values=["Mixed", ]),
-                            "Black African": SerialHierarchy(values=["Black", ]),
-                            "Asian other": SerialHierarchy(values=["Asian", ]),
-                            "Black other": SerialHierarchy(values=["Black", ]),
-                            "Mixed White/Black Caribbean": SerialHierarchy(values=["Mixed", ]),
-                            "Mixed other": SerialHierarchy(values=["Mixed", ]),
-                            "Arab": SerialHierarchy(values=["Asian", ]),
-                            "White Irish": SerialHierarchy(values=["White", ]),
-                            "Not stated": SerialHierarchy(values=["Not stated"]),
-                            "White Gypsy/Traveller": SerialHierarchy(values=["White", ]),
-                            "White British": SerialHierarchy(values=["White", ]),
-                            "Bangladeshi": SerialHierarchy(values=["Asian", ]),
-                            "White other": SerialHierarchy(values=["White", ]),
-                            "Black Caribbean": SerialHierarchy(values=["Black", ]),
-                            "Pakistani": SerialHierarchy(values=["Asian", ])}
-
+    action_space = ActionSpace(n=5)
     action_space.add_many(ActionSuppress(column_name="gender", suppress_table={"F": SerialHierarchy(values=['*', ]),
                                                                                'M': SerialHierarchy(values=['*', ])}),
-                          ActionIdentity(column_name="salary"), ActionIdentity(column_name="education"),
-                          ActionGeneralize(column_name="ethnicity", generalization_table=generalization_table))
-
+                          ActionIdentity(column_name="salary"),
+                          ActionIdentity(column_name="education"),
+                          ActionStringGeneralize(column_name="ethnicity", generalization_table=ethinicity_table),
+                          ActionSuppress(column_name="preventative_treatment",
+                                         suppress_table={"No":  SerialHierarchy(values=['Maybe', '*']),
+                                                         'Yes': SerialHierarchy(values=['Maybe', '*']),
+                                                         "NA":  SerialHierarchy(values=['Maybe', '*']),
+                                                         "Maybe": SerialHierarchy(values=['*', '*'])
+                                                         }))
+
+    # average distirtion
     average_distortion_constraint = {"salary": [0.0, 0.0, 0.0], "education": [0.0, 0.0, 0.0],
-                                     "ethnicity": [3.0, 1.0, -1.0], "gender": [4.0, 1.0, -1.0]}
+                                     "ethnicity": [3.0, 1.0, -1.0], "gender": [4.0, 1.0, -1.0],
+                                     "preventative_treatment": [4.0, 1.0, -1.0]}
 
     # specify the reward manager to use
     reward_manager = RewardManager(average_distortion_constraint=average_distortion_constraint)
@@ -66,20 +98,36 @@
     env = Environment(env_config=env_config)
 
     # initialize text distances
-    env.initialize_text_distances(distance_type=DistanceType.COSINE)
+    env.initialize_text_distances(distance_type=StringDistanceType.COSINE)
 
     algo_config = QLearnConfig()
-    algo_config.n_itrs_per_episode = 1000
+    algo_config.n_itrs_per_episode = 10
     algo_config.gamma = 0.99
     algo_config.alpha = 0.1
     algo_config.policy = EpsilonGreedyPolicy(eps=EPS, env=env,
                                              decay_op=EpsilonDecreaseOption.INVERSE_STEP)
 
     agent = QLearning(algo_config=algo_config)
 
-    configuration = {"n_episodes": 10, "output_msg_frequency": 100}
+    configuration = {"n_episodes": N_EPISODES, "output_msg_frequency": 10}
 
     # create a trainer to train the A2C agent
     trainer = Trainer(env=env, agent=agent, configuration=configuration)
 
     trainer.train()
+
+    # get the state space
+    state_space = env.state_space
+
+    for state in state_space:
+        print("Column {0} history {1}".format(state, state_space[state].history))
+
+    total_reward = trainer.total_rewards
+    episodes = [episode for episode in range(N_EPISODES)]
+
+    plt.plot(episodes, total_reward)
+    plt.xlabel("Episodes")
+    plt.ylabel("Reward")
+    plt.show()
+
+