feat: add DiscreteMaxEntropyDeepIRL (#7)

HokageM · web-flow · commit 67539d1c748d · 2023-11-17T12:25:34.000+01:00
* add DiscreteMaxEntropyDeepIRL

* update main
diff --git a/README.md b/README.md
@@ -28,15 +28,17 @@ pip install .
 # Usage
 
 ```commandline
-usage: irl [-h] [--version] [--training] [--testing] [--render]
+usage: irl [-h] [--version] [--training] [--testing] [--render] ALGORITHM
 
 Implementation of IRL algorithms
 
+positional arguments:
+  ALGORITHM   Currently supported training algorithm: [max-entropy, discrete-max-entropy-deep]
+
 options:
   -h, --help  show this help message and exit
   --version   show program's version number and exit
   --training  Enables training of model.
   --testing   Enables testing of previously created model.
   --render    Enables visualization of mountaincar.
-
 ```
diff --git a/src/irlwpython/ContinuousMaxEntropyDeepIRL.py b/src/irlwpython/ContinuousMaxEntropyDeepIRL.py
@@ -107,12 +107,20 @@ def get_demonstrations(self):
                 demonstrations[x][y][0] = state_idx
                 demonstrations[x][y][1] = raw_demo[x][y][2]
 
+        print(demonstrations)
         return demonstrations
 
+    def get_expert_state_frequencies(self):
+        raw_demo = np.load(file="expert_demo/expert_demo.npy")
+        expert_state_frequencies = []
+        return expert_state_frequencies
+
     def train(self):
         demonstrations = self.get_demonstrations()
         expert = self.expert_feature_expectations(demonstrations)
 
+        expert_state_frequencies = self.get_expert_state_frequencies()
+
         learner_feature_expectations = torch.zeros(self.state_dim, requires_grad=True)  # Add requires_grad=True
         episodes, scores = [], []
 
diff --git a/src/irlwpython/DiscreteMaxEntropyDeepIRL.py b/src/irlwpython/DiscreteMaxEntropyDeepIRL.py
@@ -0,0 +1,168 @@
+import gym
+import numpy as np
+import torch
+import torch.optim as optim
+import torch.nn as nn
+import matplotlib.pyplot as plt
+
+
+class ActorNetwork(nn.Module):
+    def __init__(self, num_inputs, num_output, hidden_size):
+        super(ActorNetwork, self).__init__()
+        self.fc1 = nn.Linear(num_inputs, hidden_size)
+        self.fc2 = nn.Linear(hidden_size, hidden_size)
+        self.fc3 = nn.Linear(hidden_size, num_output)
+
+    def forward(self, x):
+        x = nn.functional.relu(self.fc1(x))
+        x = nn.functional.relu(self.fc2(x))
+        return self.fc3(x)  # torch.nn.functional.softmax(self.fc3(x))
+
+
+class CriticNetwork(nn.Module):
+    def __init__(self, num_inputs, hidden_size):
+        super(CriticNetwork, self).__init__()
+        self.fc1 = nn.Linear(num_inputs, hidden_size)
+        self.fc2 = nn.Linear(hidden_size, hidden_size)
+        self.fc3 = nn.Linear(hidden_size, 1)
+
+        self.theta_layer = nn.Linear(hidden_size, 3)
+
+    def forward(self, x):
+        x_ = nn.functional.relu(self.fc1(x))
+        x_ = nn.functional.relu(self.fc2(x_))
+        theta_ = self.theta_layer(x_)
+        return self.fc3(x_) + torch.matmul(theta_, x)
+
+
+class DiscreteMaxEntropyDeepIRL:
+    def __init__(self, target, state_dim, action_dim, feature_matrix=None, learning_rate=0.001, gamma=0.99,
+                 num_epochs=1000):
+        self.feat_matrix = feature_matrix
+        self.one_feature = 20
+
+        self.target = target
+        self.state_dim = state_dim
+        self.action_dim = action_dim
+        self.learning_rate = learning_rate
+
+        self.gamma = gamma
+        self.num_epochs = num_epochs
+        self.actor_network = ActorNetwork(state_dim, action_dim, 100)
+        self.critic_network = CriticNetwork(state_dim + 1, 100)
+        self.optimizer_actor = optim.Adam(self.actor_network.parameters(), lr=learning_rate)
+        self.optimizer_critic = optim.Adam(self.critic_network.parameters(), lr=learning_rate)
+
+    def get_reward(self, state, action):
+        state_action = list(state) + list([action])
+        state_action = torch.Tensor(state_action)
+        return self.critic_network(state_action)
+
+    def expert_feature_expectations(self, demonstrations):
+        feature_expectations = torch.zeros(400)
+
+        for demonstration in demonstrations:
+            for state, _, _ in demonstration:
+                state_tensor = torch.tensor(state, dtype=torch.float32)
+                feature_expectations += state_tensor.squeeze()
+
+        feature_expectations /= demonstrations.shape[0]
+        return feature_expectations
+
+    def maxent_irl(self, expert, learner):
+        # Update critic network
+
+        self.optimizer_critic.zero_grad()
+
+        # Loss function for critic network
+        loss_critic = torch.nn.functional.mse_loss(learner, expert)
+        loss_critic.backward()
+
+        self.optimizer_critic.step()
+
+    def update_q_network(self, state_array, action, reward, next_state):
+        self.optimizer_actor.zero_grad()
+
+        state_tensor = torch.tensor(state_array, dtype=torch.float32)
+        next_state_tensor = torch.tensor(next_state, dtype=torch.float32)
+
+        q_values = self.actor_network(state_tensor)
+        q_1 = self.actor_network(state_tensor)[action]
+
+        q_2 = reward + self.gamma * max(self.actor_network(next_state_tensor))
+        next_q_values = reward + self.gamma * (q_2 - q_1)  # self.actor_network(next_state_tensor)
+
+        loss_actor = nn.functional.mse_loss(q_values, next_q_values)
+        loss_actor.backward()
+        self.optimizer_actor.step()
+
+    def train(self):
+        demonstrations = self.target.get_demonstrations()
+        expert = self.expert_feature_expectations(demonstrations)
+
+        learner_feature_expectations = torch.zeros(400, requires_grad=True)
+        episodes, scores = [], []
+
+        for episode in range(self.num_epochs):
+            state, info = self.target.env_reset()
+            score = 0
+
+            while True:
+                state_tensor = torch.tensor(state, dtype=torch.float32)
+
+                q_state = self.actor_network(state_tensor)
+                action = torch.argmax(q_state).item()
+                next_state, reward, done, _, _ = self.target.env_step(action)
+
+                # Actor update
+                irl_reward = self.get_reward(state, action)
+                self.update_q_network(state, action, irl_reward, next_state)
+
+                score += reward
+                state = next_state
+                if done:
+                    scores.append(score)
+                    episodes.append(episode)
+                    break
+
+            # Critic update
+            state_idx = state[0] + state[1] * self.one_feature
+            learner_feature_expectations = learner_feature_expectations + torch.Tensor(
+                self.feat_matrix[int(state_idx)])
+            learner = learner_feature_expectations / episode
+            self.maxent_irl(expert, learner)
+
+            if episode % 1 == 0:
+                score_avg = np.mean(scores)
+                print('{} episode score is {:.2f}'.format(episode, score_avg))
+                plt.plot(episodes, scores, 'b')
+                plt.savefig("./learning_curves/maxent_30000_network.png")
+
+        torch.save(self.q_network.state_dict(), "./results/maxent_30000_q_network.pth")
+
+    def test(self):
+        episodes, scores = [], []
+
+        for episode in range(10):
+            state = self.target.env_reset()
+            score = 0
+
+            while True:
+                self.target.env_render()
+                state_tensor = torch.tensor(state, dtype=torch.float32).unsqueeze(0)
+
+                action = torch.argmax(self.q_network(state_tensor)).item()
+                next_state, reward, done, _, _ = self.target.env_step(action)
+
+                score += reward
+                state = next_state
+
+                if done:
+                    scores.append(score)
+                    episodes.append(episode)
+                    plt.plot(episodes, scores, 'b')
+                    plt.savefig("./learning_curves/maxent_test_30000_network.png")
+                    break
+
+            if episode % 1 == 0:
+                print('{} episode score is {:.2f}'.format(episode, score))
diff --git a/src/irlwpython/MaxEntropyIRL.py b/src/irlwpython/MaxEntropyIRL.py
@@ -123,9 +123,7 @@ def train(self, theta_learning_rate):
                 self.update_q_table(state_idx, action, irl_reward, next_state_idx)
 
                 # State counting for densitiy
-                learner_feature_expectations += self.get_feature_matrix()[int(state_idx)]
-
-                print(reward, irl_reward)
+                learner_feature_expectations += self.feature_matrix[int(state_idx)]
 
                 score += reward
                 state = next_state
diff --git a/src/irlwpython/main.py b/src/irlwpython/main.py
@@ -3,13 +3,11 @@
 import numpy as np
 import sys
 
-from MountainCar import MountainCar
-from MaxEntropyIRL import MaxEntropyIRL
-from MaxEntropyDeepIRL import MaxEntropyDeepIRL
+from irlwpython.MountainCar import MountainCar
+from irlwpython.MaxEntropyIRL import MaxEntropyIRL
+from irlwpython.DiscreteMaxEntropyDeepIRL import DiscreteMaxEntropyDeepIRL
 
-#from irlwpython import __version__
-
-import gym
+from irlwpython import __version__
 
 __author__ = "HokageM"
 __copyright__ = "HokageM"
@@ -34,9 +32,10 @@ def parse_args(args):
     parser.add_argument(
         "--version",
         action="version",
-       # version=f"IRLwPython {__version__}",
+        version=f"IRLwPython {__version__}",
     )
-    parser.add_argument('--deep', action='store_true', help="Uses Max Entropy Deep IRL.")
+    parser.add_argument('algorithm', metavar='ALGORITHM', type=str,
+                        help='Currently supported training algorithm: [max-entropy, discrete-max-entropy-deep]')
     parser.add_argument('--training', action='store_true', help="Enables training of model.")
     parser.add_argument('--testing', action='store_true',
                         help="Enables testing of previously created model.")
@@ -86,25 +85,24 @@ def main(args):
     else:
         car = MountainCar(False, one_feature)
 
-    if args.deep:
-
-        # Create MountainCar environment
-        env = gym.make('MountainCar-v0', render_mode="human")
-        state_dim = env.observation_space.shape[0]
-        action_dim = env.action_space.n
+    if args.algorithm == "discrete-max-entropy-deep" and args.training:
+        state_dim = 2
 
         # Run MaxEnt Deep IRL using MountainCar environment
-        maxent_deep_irl_agent = MaxEntropyDeepIRL(env, state_dim, action_dim)
+        maxent_deep_irl_agent = DiscreteMaxEntropyDeepIRL(car, state_dim, n_actions, feature_matrix)
         maxent_deep_irl_agent.train()
-        maxent_deep_irl_agent.test()
+        # maxent_deep_irl_agent.test()
+
+    if args.algorithm == "discrete-max-entropy-deep" and args.testing:
+        pass
 
-    if args.training:
+    if args.algorithm == "max-entropy" and args.training:
         q_table = np.zeros((n_states, n_actions))
         trainer = MaxEntropyIRL(car, feature_matrix, one_feature, q_table, q_learning_rate, gamma, n_states, theta)
         trainer.train(theta_learning_rate)
 
-    if args.testing:
-        q_table = np.load(file="./results/maxent_q_table.npy")  # (400, 3)
+    if args.algorithm == "max-entropy" and args.testing:
+        q_table = np.load(file="./results/maxent_q_table.npy")
         trainer = MaxEntropyIRL(car, feature_matrix, one_feature, q_table, q_learning_rate, gamma, n_states, theta)
         trainer.test()
 
