Initial commit

.gitattributes

@@ -0,0 +1,2 @@
+# Auto detect text files and perform LF normalization
+* text=auto

DDPG.py

@@ -0,0 +1,115 @@
+import torch as T
+import torch.nn.functional as F
+import numpy as np
+from networks import ActorNetwork, CriticNetwork
+from buffer import ReplayBuffer
+
+device = T.device("cuda:0" if T.cuda.is_available() else "cpu")
+
+
+class DDPG:
+    def __init__(self, alpha, beta, state_dim, action_dim, actor_fc1_dim,
+                 actor_fc2_dim, critic_fc1_dim, critic_fc2_dim, ckpt_dir,
+                 gamma=0.99, tau=0.005, action_noise=0.1, max_size=1000000,
+                 batch_size=256):
+        self.gamma = gamma
+        self.tau = tau
+        self.action_noise = action_noise
+        self.checkpoint_dir = ckpt_dir
+
+        self.actor = ActorNetwork(alpha=alpha, state_dim=state_dim, action_dim=action_dim,
+                                  fc1_dim=actor_fc1_dim, fc2_dim=actor_fc2_dim)
+        self.target_actor = ActorNetwork(alpha=alpha, state_dim=state_dim, action_dim=action_dim,
+                                         fc1_dim=actor_fc1_dim, fc2_dim=actor_fc2_dim)
+        self.critic = CriticNetwork(beta=beta, state_dim=state_dim, action_dim=action_dim,
+                                    fc1_dim=critic_fc1_dim, fc2_dim=critic_fc2_dim)
+        self.target_critic = CriticNetwork(beta=beta, state_dim=state_dim, action_dim=action_dim,
+                                           fc1_dim=critic_fc1_dim, fc2_dim=critic_fc2_dim)
+
+        self.memory = ReplayBuffer(max_size=max_size, state_dim=state_dim, action_dim=action_dim,
+                                   batch_size=batch_size)
+
+        self.update_network_parameters(tau=1.0)
+
+    def update_network_parameters(self, tau=None):
+        if tau is None:
+            tau = self.tau
+
+        for actor_params, target_actor_params in zip(self.actor.parameters(),
+                                                     self.target_actor.parameters()):
+            target_actor_params.data.copy_(tau * actor_params + (1 - tau) * target_actor_params)
+
+        for critic_params, target_critic_params in zip(self.critic.parameters(),
+                                                       self.target_critic.parameters()):
+            target_critic_params.data.copy_(tau * critic_params + (1 - tau) * target_critic_params)
+
+    def remember(self, state, action, reward, state_, done):
+        self.memory.store_transition(state, action, reward, state_, done)
+
+    def choose_action(self, observation, train=True):
+        self.actor.eval()
+        state = T.tensor([observation], dtype=T.float).to(device)
+        action = self.actor.forward(state).squeeze()
+
+        if train:
+            noise = T.tensor(np.random.normal(loc=0.0, scale=self.action_noise),
+                             dtype=T.float).to(device)
+            action = T.clamp(action+noise, -1, 1)
+        self.actor.train()
+
+        return action.detach().cpu().numpy()
+
+    def learn(self):
+        if not self.memory.ready():
+            return
+
+        states, actions, reward, states_, terminals = self.memory.sample_buffer()
+        states_tensor = T.tensor(states, dtype=T.float).to(device)
+        actions_tensor = T.tensor(actions, dtype=T.float).to(device)
+        rewards_tensor = T.tensor(reward, dtype=T.float).to(device)
+        next_states_tensor = T.tensor(states_, dtype=T.float).to(device)
+        terminals_tensor = T.tensor(terminals).to(device)
+
+        with T.no_grad():
+            next_actions_tensor = self.target_actor.forward(next_states_tensor)
+            q_ = self.target_critic.forward(next_states_tensor, next_actions_tensor).view(-1)
+            q_[terminals_tensor] = 0.0
+            target = rewards_tensor + self.gamma * q_
+        q = self.critic.forward(states_tensor, actions_tensor).view(-1)
+
+        critic_loss = F.mse_loss(q, target.detach())
+        self.critic.optimizer.zero_grad()
+        critic_loss.backward()
+        self.critic.optimizer.step()
+
+        new_actions_tensor = self.actor.forward(states_tensor)
+        actor_loss = -T.mean(self.critic(states_tensor, new_actions_tensor))
+        self.actor.optimizer.zero_grad()
+        actor_loss.backward()
+        self.actor.optimizer.step()
+
+        self.update_network_parameters()
+
+    def save_models(self, episode):
+        self.actor.save_checkpoint(self.checkpoint_dir + 'Actor/DDPG_actor_{}.pth'.format(episode))
+        print('Saving actor network successfully!')
+        self.target_actor.save_checkpoint(self.checkpoint_dir +
+                                          'Target_actor/DDPG_target_actor_{}.pth'.format(episode))
+        print('Saving target_actor network successfully!')
+        self.critic.save_checkpoint(self.checkpoint_dir + 'Critic/DDPG_critic_{}'.format(episode))
+        print('Saving critic network successfully!')
+        self.target_critic.save_checkpoint(self.checkpoint_dir +
+                                           'Target_critic/DDPG_target_critic_{}'.format(episode))
+        print('Saving target critic network successfully!')
+
+    def load_models(self, episode):
+        self.actor.load_checkpoint(self.checkpoint_dir + 'Actor/DDPG_actor_{}.pth'.format(episode))
+        print('Loading actor network successfully!')
+        self.target_actor.load_checkpoint(self.checkpoint_dir +
+                                          'Target_actor/DDPG_target_actor_{}.pth'.format(episode))
+        print('Loading target_actor network successfully!')
+        self.critic.load_checkpoint(self.checkpoint_dir + 'Critic/DDPG_critic_{}'.format(episode))
+        print('Loading critic network successfully!')
+        self.target_critic.load_checkpoint(self.checkpoint_dir +
+                                           'Target_critic/DDPG_target_critic_{}'.format(episode))
+        print('Loading target critic network successfully!')

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2022 indigoLovee
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in all
+copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
+SOFTWARE.

README.md

@@ -0,0 +1,2 @@
+# DDPG
+ DDPG in Pytorch

buffer.py

@@ -0,0 +1,41 @@
+import numpy as np
+
+
+class ReplayBuffer:
+    def __init__(self, max_size, state_dim, action_dim, batch_size):
+        self.mem_size = max_size
+        self.batch_size = batch_size
+        self.mem_cnt = 0
+
+        self.state_memory = np.zeros((self.mem_size, state_dim))
+        self.action_memory = np.zeros((self.mem_size, action_dim))
+        self.reward_memory = np.zeros((self.mem_size, ))
+        self.next_state_memory = np.zeros((self.mem_size, state_dim))
+        self.terminal_memory = np.zeros((self.mem_size, ), dtype=np.bool)
+
+    def store_transition(self, state, action, reward, state_, done):
+        mem_idx = self.mem_cnt % self.mem_size
+
+        self.state_memory[mem_idx] = state
+        self.action_memory[mem_idx] = action
+        self.reward_memory[mem_idx] = reward
+        self.next_state_memory[mem_idx] = state_
+        self.terminal_memory[mem_idx] = done
+
+        self.mem_cnt += 1
+
+    def sample_buffer(self):
+        mem_len = min(self.mem_size, self.mem_cnt)
+        batch = np.random.choice(mem_len, self.batch_size, replace=False)
+
+        states = self.state_memory[batch]
+        actions = self.action_memory[batch]
+        rewards = self.reward_memory[batch]
+        states_ = self.next_state_memory[batch]
+        terminals = self.terminal_memory[batch]
+
+        return states, actions, rewards, states_, terminals
+
+    def ready(self):
+        return self.mem_cnt >= self.batch_size
+

networks.py

@@ -0,0 +1,72 @@
+import torch as T
+import torch.nn as nn
+import torch.optim as optim
+
+device = T.device("cuda:0" if T.cuda.is_available() else "cpu")
+
+
+def weight_init(m):
+    if isinstance(m, nn.Linear):
+        nn.init.xavier_normal_(m.weight)
+        if m.bias is not None:
+            nn.init.constant_(m.bias, 0.0)
+    elif isinstance(m, nn.BatchNorm1d):
+        nn.init.constant_(m.weight, 1.0)
+        nn.init.constant_(m.bias, 0.0)
+
+
+class ActorNetwork(nn.Module):
+    def __init__(self, alpha, state_dim, action_dim, fc1_dim, fc2_dim):
+        super(ActorNetwork, self).__init__()
+        self.fc1 = nn.Linear(state_dim, fc1_dim)
+        self.ln1 = nn.LayerNorm(fc1_dim)
+        self.fc2 = nn.Linear(fc1_dim, fc2_dim)
+        self.ln2 = nn.LayerNorm(fc2_dim)
+        self.action = nn.Linear(fc2_dim, action_dim)
+
+        self.optimizer = optim.Adam(self.parameters(), lr=alpha)
+        self.apply(weight_init)
+        self.to(device)
+
+    def forward(self, state):
+        x = T.relu(self.ln1(self.fc1(state)))
+        x = T.relu(self.ln2(self.fc2(x)))
+        action = T.tanh(self.action(x))
+
+        return action
+
+    def save_checkpoint(self, checkpoint_file):
+        T.save(self.state_dict(), checkpoint_file)
+
+    def load_checkpoint(self, checkpoint_file):
+        self.load_state_dict(T.load(checkpoint_file))
+
+
+class CriticNetwork(nn.Module):
+    def __init__(self, beta, state_dim, action_dim, fc1_dim, fc2_dim):
+        super(CriticNetwork, self).__init__()
+        self.fc1 = nn.Linear(state_dim, fc1_dim)
+        self.ln1 = nn.LayerNorm(fc1_dim)
+        self.fc2 = nn.Linear(fc1_dim, fc2_dim)
+        self.ln2 = nn.LayerNorm(fc2_dim)
+        self.fc3 = nn.Linear(action_dim, fc2_dim)
+        self.q = nn.Linear(fc2_dim, 1)
+
+        self.optimizer = optim.Adam(self.parameters(), lr=beta, weight_decay=0.001)
+        self.apply(weight_init)
+        self.to(device)
+
+    def forward(self, state, action):
+        x_s = T.relu(self.ln1(self.fc1(state)))
+        x_s = self.ln2(self.fc2(x_s))
+        x_a = self.fc3(action)
+        x = T.relu(x_s + x_a)
+        q = self.q(x)
+
+        return q
+
+    def save_checkpoint(self, checkpoint_file):
+        T.save(self.state_dict(), checkpoint_file)
+
+    def load_checkpoint(self, checkpoint_file):
+        self.load_state_dict(T.load(checkpoint_file))

test.py

@@ -0,0 +1,41 @@
+import gym
+import imageio
+import argparse
+from DDPG import DDPG
+from utils import scale_action
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--filename', type=str, default='./output_images/LunarLander.gif')
+parser.add_argument('--checkpoint_dir', type=str, default='./checkpoints/DDPG/')
+parser.add_argument('--save_video', type=bool, default=True)
+parser.add_argument('--fps', type=int, default=30)
+parser.add_argument('--render', type=bool, default=True)
+
+args = parser.parse_args()
+
+
+def main():
+    env = gym.make('LunarLanderContinuous-v2')
+    agent = DDPG(alpha=0.0003, beta=0.0003, state_dim=env.observation_space.shape[0],
+                 action_dim=env.action_space.shape[0], actor_fc1_dim=400, actor_fc2_dim=300,
+                 critic_fc1_dim=400, critic_fc2_dim=300, ckpt_dir=args.checkpoint_dir,
+                 batch_size=256)
+    agent.load_models(1000)
+    video = imageio.get_writer(args.filename, fps=args.fps)
+
+    done = False
+    observation = env.reset()
+    while not done:
+        if args.render:
+            env.render()
+        action = agent.choose_action(observation, train=True)
+        action_ = scale_action(action.copy(), env.action_space.high, env.action_space.low)
+        observation_, reward, done, info = env.step(action_)
+        observation = observation_
+        if args.save_video:
+            video.append_data(env.render(mode='rgb_array'))
+
+
+if __name__ == '__main__':
+    main()
+

train.py

@@ -0,0 +1,53 @@
+import gym
+import numpy as np
+import argparse
+from DDPG import DDPG
+from utils import create_directory, plot_learning_curve, scale_action
+
+parser = argparse.ArgumentParser("DDPG parameters")
+parser.add_argument('--max_episodes', type=int, default=1000)
+parser.add_argument('--checkpoint_dir', type=str, default='./checkpoints/DDPG/')
+parser.add_argument('--figure_file', type=str, default='./output_images/reward.png')
+
+args = parser.parse_args()
+
+
+def main():
+    env = gym.make('LunarLanderContinuous-v2')
+    agent = DDPG(alpha=0.0003, beta=0.0003, state_dim=env.observation_space.shape[0],
+                 action_dim=env.action_space.shape[0], actor_fc1_dim=400, actor_fc2_dim=300,
+                 critic_fc1_dim=400, critic_fc2_dim=300, ckpt_dir=args.checkpoint_dir,
+                 batch_size=256)
+    create_directory(args.checkpoint_dir,
+                     sub_paths=['Actor', 'Target_actor', 'Critic', 'Target_critic'])
+
+    reward_history = []
+    avg_reward_history = []
+    for episode in range(args.max_episodes):
+        done = False
+        total_reward = 0
+        observation = env.reset()
+        while not done:
+            action = agent.choose_action(observation, train=True)
+            action_ = scale_action(action.copy(), env.action_space.high, env.action_space.low)
+            observation_, reward, done, info = env.step(action_)
+            agent.remember(observation, action, reward, observation_, done)
+            agent.learn()
+            total_reward += reward
+            observation = observation_
+
+        reward_history.append(total_reward)
+        avg_reward = np.mean(reward_history[-100:])
+        avg_reward_history.append(avg_reward)
+        print('Ep: {} Reward: {:.1f} AvgReward: {:.1f}'.format(episode+1, total_reward, avg_reward))
+
+        if (episode + 1) % 200 == 0:
+            agent.save_models(episode+1)
+
+    episodes = [i+1 for i in range(args.max_episodes)]
+    plot_learning_curve(episodes, avg_reward_history, title='AvgReward',
+                        ylabel='reward', figure_file=args.figure_file)
+
+
+if __name__ == '__main__':
+    main()