Upload 14 files

.gitattributes

@@ -33,3 +33,5 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+figures/neurips_logo.png filter=lfs diff=lfs merge=lfs -text
+figures/plot.png filter=lfs diff=lfs merge=lfs -text

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) 2023 jangwonkim-cocel
+
+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

@@ -1,3 +1,67 @@
----
-license: mit
----
+<div align="center">
+  <h1>Belief Projection-Based Q-Learning</h1>
+  <a href="https://www.python.org/">
+    <img src="https://img.shields.io/badge/Python-3.8-blue?logo=python&style=flat-square" alt="Python Badge"/>
+  </a>
+  &nbsp;&nbsp;
+  <a href="https://pytorch.org/">
+    <img src="https://img.shields.io/badge/PyTorch-2.1.0-%23EE4C2C?logo=pytorch&style=flat-square" alt="PyTorch Badge"/>
+  </a>
+  &nbsp;&nbsp;
+  <a href="https://proceedings.neurips.cc/paper_files/paper/2023/hash/0252a434b18962c94910c07cd9a7fecc-Abstract-Conference.html">
+    <img src="https://img.shields.io/badge/NeurIPS%202023-Paper-%23007ACC?style=flat-square" alt="NeurIPS 2023 Badge"/>
+  </a>
+    <br/><br/>
+  <img src="./figures/neurips_logo.png" width="200px" style="margin: 0 10px;"/>
+</div>
+
+## [NeurIPS 2023] Official Implementation of Belief Projection-Based Q-Learning (BPQL)
+This repository contains the PyTorch implementation of **BPQL** introduced in the paper: **_Belief Projection-Based Reinforcement Learning for Environments with Delayed Feedback_** by Jangwon Kim et al., presented at Advances in Neural Information Processing Systems (NeurIPS), 2023.
+
+
+## 📄 Paper Link
+>You can see the paper here: https://proceedings.neurips.cc/paper_files/paper/2023/hash/0252a434b18962c94910c07cd9a7fecc-Abstract-Conference.html
+
+
+## 🚀 Achieves S.O.T.A. Performance, Yet Very Simple to Implement
+
+* **Supports both observation delay, action delay, and their combination**
+* **Performance Plot ⬇️**
+<p align="center">
+  <img src="./figures/plot.png" alt="BPQL Performance Plot" width="600"/>
+</p>
+
+
+## ▶️ How to Run?
+### Option 1: Run the script file 
+```
+>chmod +x run.sh
+>./run.sh
+```
+
+### Option 2: Run main.py with arguments
+```
+python main.py --env-name HalfCheetah-v3 --random-seed 2023 --obs-delayed-steps 5 --act-delayed-steps 4 --max-step 1000000
+```
+---
+
+## ✅Test Environment
+```
+python == 3.8.10
+gym == 0.26.2
+mujoco_py == 2.1.2.14
+pytorch == 2.1.0
+numpy == 1.24.3
+```
+
+## 📚 Citation Example
+```
+@inproceedings{kim2023cocel,
+   author = {Kim, Jangwon and Kim, Hangyeol and Kang, Jiwook and Baek, Jongchan and Han, Soohee},
+   booktitle = {Advances in Neural Information Processing Systems},
+   pages = {678--696},
+   title = {Belief Projection-Based Reinforcement Learning for Environments with Delayed Feedback},
+   volume = {36},
+   year = {2023}
+}
+```

bpql.py

@@ -0,0 +1,99 @@
+import torch
+import torch.nn.functional as F
+from replay_memory import ReplayMemory
+from network import Twin_Q_net, GaussianPolicy
+from temporary_buffer import TemporaryBuffer
+from utils import hard_update, soft_update
+
+
+class BPQLAgent:  # SAC for the base learning algorithm
+    def __init__(self, args, state_dim, action_dim, action_bound, action_space, device):
+        self.args = args
+
+        self.state_dim = state_dim
+        self.action_dim = action_dim
+        self.action_bound = action_bound
+
+        self.device = device
+        self.replay_memory = ReplayMemory(args.obs_delayed_steps + args.act_delayed_steps, state_dim, action_dim, device, args.buffer_size)
+        self.temporary_buffer = TemporaryBuffer(args.obs_delayed_steps + args.act_delayed_steps)
+        self.eval_temporary_buffer = TemporaryBuffer(args.obs_delayed_steps + args.act_delayed_steps)
+        self.batch_size = args.batch_size
+
+        self.gamma = args.gamma
+        self.tau = args.tau
+
+        self.actor = GaussianPolicy(args, args.obs_delayed_steps + args.act_delayed_steps, state_dim, action_dim, action_bound, args.hidden_dims, F.relu, device).to(device)
+        self.critic = Twin_Q_net(state_dim, action_dim, device, args.hidden_dims).to(device)  # Network for the beta Q-values.
+        self.target_critic = Twin_Q_net(state_dim, action_dim, device, args.hidden_dims).to(device)
+
+        self.actor_optimizer = torch.optim.Adam(self.actor.parameters(), lr=args.actor_lr)
+        self.critic_optimizer = torch.optim.Adam(self.critic.parameters(), lr=args.critic_lr)
+
+        # Automated Entropy Adjustment for Maximum Entropy RL
+        if args.automating_temperature is True:
+            self.target_entropy = -torch.prod(torch.Tensor(action_space.shape)).to(device)
+            self.log_alpha = torch.zeros(1, requires_grad=True, device=device)
+            self.alpha_optimizer = torch.optim.Adam([self.log_alpha], lr=args.temperature_lr)
+        else:
+            self.log_alpha = torch.log(torch.tensor(args.temperature, device=device, dtype=torch.float32))
+
+        hard_update(self.critic, self.target_critic)
+
+    def get_action(self, state, evaluation=True):
+        with torch.no_grad():
+            if evaluation:
+                _, _, action = self.actor.sample(state)
+            else:
+                action, _, _ = self.actor.sample(state)
+        return action.cpu().numpy()[0]
+
+    def train_actor(self, augmented_states, states, train_alpha=True):
+        self.actor_optimizer.zero_grad()
+        actions, log_pis, _ = self.actor.sample(augmented_states)
+        q_values_A, q_values_B = self.critic(states, actions)
+        q_values = torch.min(q_values_A, q_values_B)
+
+        actor_loss = (self.log_alpha.exp().detach() * log_pis - q_values).mean()
+        actor_loss.backward()
+        self.actor_optimizer.step()
+
+        if train_alpha:
+            self.alpha_optimizer.zero_grad()
+            alpha_loss = -(self.log_alpha.exp() * (log_pis + self.target_entropy).detach()).mean()
+            alpha_loss.backward()
+            self.alpha_optimizer.step()
+        else:
+            alpha_loss = torch.tensor(0.)
+
+        return actor_loss.item(), alpha_loss.item()
+
+    def train_critic(self, actions, rewards, next_augmented_states, dones,  states, next_states):
+        self.critic_optimizer.zero_grad()
+        with torch.no_grad():
+            next_actions, next_log_pis, _ = self.actor.sample(next_augmented_states)
+            next_q_values_A, next_q_values_B = self.target_critic(next_states, next_actions)
+            next_q_values = torch.min(next_q_values_A, next_q_values_B) - self.log_alpha.exp() * next_log_pis
+            target_q_values = rewards + (1 - dones) * self.gamma * next_q_values
+
+        q_values_A, q_values_B = self.critic(states, actions)
+        critic_loss = ((q_values_A - target_q_values)**2).mean() + ((q_values_B - target_q_values)**2).mean()
+
+        critic_loss.backward()
+        self.critic_optimizer.step()
+
+        return critic_loss.item() # 2 * Squared-Loss = (2*|TD-error|^2)
+
+    def train(self):
+        augmented_states, actions, rewards, next_augmented_states, dones, states, next_states = self.replay_memory.sample(self.batch_size)
+
+        critic_loss = self.train_critic(actions, rewards, next_augmented_states, dones, states, next_states)
+        if self.args.automating_temperature is True:
+            actor_loss, log_alpha_loss = self.train_actor(augmented_states, states, train_alpha=True)
+        else:
+            actor_loss, log_alpha_loss = self.train_actor(augmented_states, states, train_alpha=False)
+
+        soft_update(self.critic, self.target_critic, self.tau)
+
+        return critic_loss, actor_loss, log_alpha_loss
+

log/temp.md

@@ -0,0 +1 @@
+temp.

main.py

@@ -0,0 +1,56 @@
+import argparse
+import torch
+from bpql import BPQLAgent
+from trainer import Trainer
+from utils import set_seed, make_delayed_env
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--env-name', default='HalfCheetah-v3', type=str)
+
+    parser.add_argument('--obs-delayed-steps', default=4, type=int)  # Delayed timesteps (Observation, Reward)
+    parser.add_argument('--act-delayed-steps', default=5, type=int)  # Delayed timesteps (Action)
+
+    parser.add_argument('--random-seed', default=-1, type=int)
+    parser.add_argument('--eval_flag', default=True, type=bool)
+    parser.add_argument('--eval-freq', default=5000, type=int)
+    parser.add_argument('--eval-episode', default=5, type=int)
+    parser.add_argument('--automating-temperature', default=True, type=bool)
+    parser.add_argument('--temperature', default=0.2, type=float)
+    parser.add_argument('--start-step', default=10000, type=int)
+    parser.add_argument('--max-step', default=1000000, type=int)
+    parser.add_argument('--update_after', default=1000, type=int)
+    parser.add_argument('--hidden-dims', default=(256, 256))
+    parser.add_argument('--batch-size', default=256, type=int)
+    parser.add_argument('--buffer-size', default=1000000, type=int)
+    parser.add_argument('--update-every', default=50, type=int)
+    parser.add_argument('--log_std_bound', default=[-20, 2])
+    parser.add_argument('--gamma', default=0.99, type=float)
+    parser.add_argument('--actor-lr', default=3e-4, type=float)
+    parser.add_argument('--critic-lr', default=3e-4, type=float)
+    parser.add_argument('--temperature-lr', default=3e-4, type=float)
+    parser.add_argument('--tau', default=0.005, type=float)
+    parser.add_argument('--show-loss', default=False, type=bool)
+    args = parser.parse_args()
+
+    # Set Device
+    device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+    # Set Seed
+    random_seed = set_seed(args.random_seed)
+
+    # Create Delayed Environment
+    env, eval_env = make_delayed_env(args, random_seed, obs_delayed_steps=args.obs_delayed_steps, act_delayed_steps=args.act_delayed_steps)
+
+    state_dim = env.observation_space.shape[0]
+    action_dim = env.action_space.shape[0]
+    action_bound = [env.action_space.low[0], env.action_space.high[0]]
+
+    print(f"Environment: {args.env_name}, Obs. Delayed Steps: {args.obs_delayed_steps}, Act. Delayed Steps: {args.act_delayed_steps}, Random Seed: {args.random_seed}", "\n")
+
+    # Create Agent
+    agent = BPQLAgent(args, state_dim, action_dim, action_bound, env.action_space, device)
+
+    # Create Trainer & Train
+    trainer = Trainer(env, eval_env, agent, args)
+    trainer.train()

network.py

@@ -0,0 +1,119 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torch.distributions import Normal
+from utils import weight_init
+
+
+class Twin_Q_net(nn.Module):
+    def __init__(self, state_dim, action_dim, device, hidden_dims=(256, 256), activation_fc=F.relu):
+        super(Twin_Q_net, self).__init__()
+        self.device = device
+
+        self.activation_fc = activation_fc
+
+        self.input_layer_A = nn.Linear(state_dim + action_dim, hidden_dims[0])
+        self.hidden_layers_A = nn.ModuleList()
+        for i in range(len(hidden_dims)-1):
+            hidden_layer_A = nn.Linear(hidden_dims[i], hidden_dims[i+1])
+            self.hidden_layers_A.append(hidden_layer_A)
+        self.output_layer_A = nn.Linear(hidden_dims[-1], 1)
+
+        self.input_layer_B = nn.Linear(state_dim + action_dim, hidden_dims[0])
+        self.hidden_layers_B = nn.ModuleList()
+        for i in range(len(hidden_dims)-1):
+            hidden_layer_B = nn.Linear(hidden_dims[i], hidden_dims[i+1])
+            self.hidden_layers_B.append(hidden_layer_B)
+        self.output_layer_B = nn.Linear(hidden_dims[-1], 1)
+        self.apply(weight_init)
+
+    def _format(self, state, action):
+        x, u = state, action
+
+        if not isinstance(x, torch.Tensor):
+            x = torch.tensor(x, device=self.device, dtype=torch.float32)
+            x = x.unsqueeze(0)
+
+        if not isinstance(u, torch.Tensor):
+            u = torch.tensor(u, device=self.device, dtype=torch.float32)
+            u = u.unsqueeze(0)
+
+        return x, u
+
+    def forward(self, state, action):
+        x, u = self._format(state, action)
+        x = torch.cat([x, u], dim=1)
+
+        x_A = self.activation_fc(self.input_layer_A(x))
+        for i, hidden_layer_A in enumerate(self.hidden_layers_A):
+            x_A = self.activation_fc(hidden_layer_A(x_A))
+        x_A = self.output_layer_A(x_A)
+
+        x_B = self.activation_fc(self.input_layer_B(x))
+        for i, hidden_layer_B in enumerate(self.hidden_layers_B):
+            x_B = self.activation_fc(hidden_layer_B(x_B))
+        x_B = self.output_layer_B(x_B)
+
+        return x_A, x_B
+
+
+class GaussianPolicy(nn.Module):
+    def __init__(self, args, delayed_steps, state_dim, action_dim, action_bound,
+                 hidden_dims=(256, 256), activation_fc=F.relu, device='cuda'):
+        super(GaussianPolicy, self).__init__()
+        self.device = device
+
+        self.log_std_min = args.log_std_bound[0]
+        self.log_std_max = args.log_std_bound[1]
+
+        self.activation_fc = activation_fc
+
+        self.input_layer = nn.Linear(state_dim + delayed_steps * action_dim, hidden_dims[0])
+        self.hidden_layers = nn.ModuleList()
+        for i in range(len(hidden_dims)-1):
+            hidden_layer = nn.Linear(hidden_dims[i], hidden_dims[i+1])
+            self.hidden_layers.append(hidden_layer)
+
+        self.mean_layer = nn.Linear(hidden_dims[-1], action_dim)
+        self.log_std_layer = nn.Linear(hidden_dims[-1], action_dim)
+
+        self.action_rescale = torch.as_tensor((action_bound[1] - action_bound[0]) / 2., dtype=torch.float32)
+        self.action_rescale_bias = torch.as_tensor((action_bound[1] + action_bound[0]) / 2., dtype=torch.float32)
+
+        self.apply(weight_init)
+
+    def _format(self, state):
+        x = state
+        if not isinstance(x, torch.Tensor):
+            x = torch.tensor(x, device=self.device, dtype=torch.float32)
+            x = x.unsqueeze(0)
+        return x
+
+    def forward(self, state):
+        x = self._format(state)
+        x = self.activation_fc(self.input_layer(x))
+        for i, hidden_layer in enumerate(self.hidden_layers):
+            x = self.activation_fc(hidden_layer(x))
+        mean = self.mean_layer(x)
+        log_std = self.log_std_layer(x)
+        log_std = torch.clamp(log_std, self.log_std_min, self.log_std_max)
+        return mean, log_std
+
+    def sample(self, state):
+        mean, log_std = self.forward(state)
+        distribution = Normal(mean, log_std.exp())
+
+        unbounded_action = distribution.rsample()
+
+        bounded_action = torch.tanh(unbounded_action)
+        action = bounded_action * self.action_rescale + self.action_rescale_bias
+
+        log_prob = distribution.log_prob(unbounded_action) - torch.log(self.action_rescale *
+                                                                       (1 - bounded_action.pow(2).clamp(0, 1)) + 1e-6)
+        log_prob = log_prob.sum(dim=1, keepdim=True)
+        mean = torch.tanh(mean) * self.action_rescale + self.action_rescale_bias
+        return action, log_prob, mean
+
+
+
+

replay_memory.py

@@ -0,0 +1,44 @@
+import torch
+import numpy as np
+
+
+class ReplayMemory:
+    def __init__(self, delayed_steps, state_dim, action_dim, device, capacity=1e6):
+        self.device = device
+        self.capacity = int(capacity)
+        self.size = 0
+        self.position = 0
+
+        self.augmented_state_buffer = np.empty(shape=(self.capacity, state_dim + action_dim * delayed_steps), dtype=np.float32)
+        self.action_buffer = np.empty(shape=(self.capacity, action_dim), dtype=np.float32)
+        self.reward_buffer = np.empty(shape=(self.capacity, 1), dtype=np.float32)
+        self.next_augmented_state_buffer = np.empty(shape=(self.capacity, state_dim + action_dim * delayed_steps), dtype=np.float32)
+        self.done_buffer = np.empty(shape=(self.capacity, 1), dtype=np.float32)
+        self.state_buffer = np.empty(shape=(self.capacity, state_dim), dtype=np.float32)
+        self.next_state_buffer = np.empty(shape=(self.capacity, state_dim), dtype=np.float32)
+
+    def push(self, augmented_state, state, action, reward, next_augmented_state, next_state, done):
+        self.size = min(self.size + 1, self.capacity)
+
+        self.augmented_state_buffer[self.position] = augmented_state
+        self.action_buffer[self.position] = action
+        self.reward_buffer[self.position] = reward
+        self.next_augmented_state_buffer[self.position] = next_augmented_state
+        self.done_buffer[self.position] = done
+        self.state_buffer[self.position] = state
+        self.next_state_buffer[self.position] = next_state
+
+        self.position = (self.position + 1) % self.capacity
+
+    def sample(self, batch_size):
+        idxs = np.random.randint(0, self.size, size=batch_size)
+
+        augmented_states = torch.FloatTensor(self.augmented_state_buffer[idxs]).to(self.device)
+        actions = torch.FloatTensor(self.action_buffer[idxs]).to(self.device)
+        rewards = torch.FloatTensor(self.reward_buffer[idxs]).to(self.device)
+        next_augmented_states = torch.FloatTensor(self.next_augmented_state_buffer[idxs]).to(self.device)
+        dones = torch.FloatTensor(self.done_buffer[idxs]).to(self.device)
+        states = torch.FloatTensor(self.state_buffer[idxs]).to(self.device)
+        next_states = torch.FloatTensor(self.next_state_buffer[idxs]).to(self.device)
+
+        return augmented_states, actions, rewards, next_augmented_states, dones, states, next_states

run.sh

@@ -0,0 +1,15 @@
+# ----Test Environment----
+# python == 3.8.10
+# gym == 0.26.2
+# mujoco_py == 2.1.2.14
+# pytorch == 2.1.0
+# numpy == 1.24.3
+#-------------------------
+
+#!/bin/bash
+python main.py \
+--env-name "HalfCheetah-v3" \
+--random-seed 2023 \
+--obs-delayed-steps 5 \
+--act-delayed-steps 4 \
+--max-step 1000000

temporary_buffer.py

@@ -0,0 +1,34 @@
+import numpy as np
+from collections import deque
+
+
+class TemporaryBuffer:
+    def __init__(self, delayed_steps):
+        self.d = delayed_steps
+        self.states = deque(maxlen=delayed_steps + 2)
+        self.actions = deque(maxlen=2 * delayed_steps + 1)
+
+    def clear(self):
+        self.states.clear()
+        self.actions.clear()
+
+    def get_augmented_state(self, last_observed_state, first_action_idx):
+        aug_state = np.concatenate([last_observed_state, self.actions[first_action_idx]])
+        for i in range(first_action_idx + 1, first_action_idx + self.d):
+            aug_state = np.concatenate([aug_state, self.actions[i]])
+        return aug_state
+
+    def get_tuple(self):
+        assert len(self.states) == self.d + 2 and len(self.actions) == 2 * self.d + 1
+
+        aug_s = self.get_augmented_state(self.states[0], 0)
+        s = self.states[-2]
+        a = self.actions[self.d]
+
+        next_aug_s = self.get_augmented_state(self.states[1], 1)
+        next_s = self.states[-1]
+
+        self.states.popleft()
+        self.actions.popleft()
+        return aug_s, s, a, next_aug_s, next_s
+

trainer.py

@@ -0,0 +1,165 @@
+import numpy as np
+
+from utils import log_to_txt
+
+
+class Trainer:
+    def __init__(self, env, eval_env, agent, args):
+        self.args = args
+        self.agent = agent
+
+        self.delayed_env = env
+        self.eval_delayed_env = eval_env
+
+        self.start_step = args.start_step
+        self.update_after = args.update_after
+        self.max_step = args.max_step
+        self.batch_size = args.batch_size
+        self.update_every = args.update_every
+
+        self.eval_flag = args.eval_flag
+        self.eval_episode = args.eval_episode
+        self.eval_freq = args.eval_freq
+
+        self.episode = 0
+        self.total_step = 0
+        self.local_step = 0
+        self.eval_local_step = 0
+        self.eval_num = 0
+        self.finish_flag = False
+
+        self.total_delayed_steps = args.obs_delayed_steps + self.args.act_delayed_steps
+
+    def train(self):
+        # The train process starts here.
+        while not self.finish_flag:
+            self.episode += 1
+            self.local_step = 0
+
+            # Initialize the delayed environment & the temporal buffer
+            self.delayed_env.reset()
+            self.agent.temporary_buffer.clear()
+            done = False
+
+            # Episode starts here.
+            while not done:
+                self.local_step += 1
+                self.total_step += 1
+
+                if self.local_step < self.total_delayed_steps:  # if t < d
+                    action = np.zeros_like(self.delayed_env.action_space.sample())  # Select the 'no-op' action
+                    _, _, _, _ = self.delayed_env.step(action)
+
+                    self.agent.temporary_buffer.actions.append(action)
+                elif self.local_step == self.total_delayed_steps:  # if t == d
+                    if self.total_step < self.start_step:
+                        action = self.delayed_env.action_space.sample()
+                    else:
+                        action = np.zeros_like(self.delayed_env.action_space.sample())  # Select the 'no-op' action
+
+                    next_observed_state, _, _, _ = self.delayed_env.step(action)
+                    #                s(1)       <-     Env: a(d)
+                    self.agent.temporary_buffer.actions.append(action)  # Put a(d) to the temporary buffer
+                    self.agent.temporary_buffer.states.append(next_observed_state)  # Put s(1) to the temporary buffer
+                else:  # if t > d
+                    last_observed_state = self.agent.temporary_buffer.states[-1]
+                    first_action_idx = len(self.agent.temporary_buffer.actions) - self.total_delayed_steps
+
+                    # Get the augmented state(t)
+                    augmented_state = self.agent.temporary_buffer.get_augmented_state(last_observed_state, first_action_idx)
+
+                    if self.total_step < self.start_step:
+                        action = self.delayed_env.action_space.sample()
+                    else:
+                        action = self.agent.get_action(augmented_state, evaluation=False)
+                        # a(t) <- policy: augmented_state(t)
+                    next_observed_state, reward, done, info = self.delayed_env.step(action)
+                    #          s(t+1-d),  r(t-d)      <-      Env: a(t)
+                    true_done = 0.0 if self.local_step == self.delayed_env._max_episode_steps + self.args.obs_delayed_steps else float(done)
+
+                    self.agent.temporary_buffer.actions.append(action)  # Put a(t) to the temporary buffer
+                    self.agent.temporary_buffer.states.append(next_observed_state)  # Put s(t+1-d) to the temporary buffer
+
+                    if self.local_step > 2 * self.total_delayed_steps:  # if t > 2d
+                        augmented_s, s, a, next_augmented_s, next_s = self.agent.temporary_buffer.get_tuple()
+                        #  aug_s(t-d),  s(t-d),  a(t-d),  aug_s(t+1-d),  s(t+1-d)  <- Temporal Buffer
+                        self.agent.replay_memory.push(augmented_s, s, a, reward, next_augmented_s, next_s, true_done)
+                        #  Store (aug_s(t-d), s(t-d), a(t-d), r(t-d), aug_s(t+1-d), s(t+1-d)) in the replay memory.
+
+                # Update parameters
+                if self.agent.replay_memory.size >= self.batch_size and self.total_step >= self.update_after and \
+                        self.total_step % self.update_every == 0:
+                    total_actor_loss = 0
+                    total_critic_loss = 0
+                    total_log_alpha_loss = 0
+                    for i in range(self.update_every):
+                        # Train the policy and the beta Q-network (critic).
+                        critic_loss, actor_loss, log_alpha_loss = self.agent.train()
+                        total_critic_loss += critic_loss
+                        total_actor_loss += actor_loss
+                        total_log_alpha_loss += log_alpha_loss
+
+                    # Print the loss.
+                    if self.args.show_loss:
+                        print("Loss  |  Actor loss {:.3f}  |  Critic loss {:.3f}  |  Log-alpha loss {:.3f}"
+                              .format(total_actor_loss / self.update_every, total_critic_loss / self.update_every,
+                                      total_log_alpha_loss / self.update_every))
+
+                # Evaluate.
+                if self.eval_flag and self.total_step % self.eval_freq == 0:
+                    self.evaluate()
+
+                # Raise finish flag.
+                if self.total_step == self.max_step:
+                    self.finish_flag = True
+
+    def evaluate(self):
+        # Evaluate process
+        self.eval_num += 1
+        reward_list = []
+
+        for epi in range(self.eval_episode):
+            episode_reward = 0
+            self.eval_delayed_env.reset()
+            self.agent.eval_temporary_buffer.clear()
+            done = False
+            self.eval_local_step = 0
+
+            while not done:
+                self.eval_local_step += 1
+                if self.eval_local_step < self.total_delayed_steps:
+                    action = np.zeros_like(self.delayed_env.action_space.sample())
+                    _, _, _, _ = self.eval_delayed_env.step(action)
+                    self.agent.eval_temporary_buffer.actions.append(action)
+                elif self.eval_local_step == self.total_delayed_steps:
+                    action = np.zeros_like(self.eval_delayed_env.action_space.sample())
+                    next_observed_state, _, _, _ = self.eval_delayed_env.step(action)
+                    self.agent.eval_temporary_buffer.actions.append(action)
+                    self.agent.eval_temporary_buffer.states.append(next_observed_state)
+                else:
+                    last_observed_state = self.agent.eval_temporary_buffer.states[-1]
+                    first_action_idx = len(self.agent.eval_temporary_buffer.actions) - self.total_delayed_steps
+                    augmented_state = self.agent.eval_temporary_buffer.get_augmented_state(last_observed_state,
+                                                                                          first_action_idx)
+                    action = self.agent.get_action(augmented_state, evaluation=True)
+                    next_observed_state, reward, done, _ = self.eval_delayed_env.step(action)
+                    self.agent.eval_temporary_buffer.actions.append(action)
+                    self.agent.eval_temporary_buffer.states.append(next_observed_state)
+                    episode_reward += reward
+
+            reward_list.append(episode_reward)
+
+        log_to_txt(self.args.env_name, self.args.random_seed, self.total_step, sum(reward_list) / len(reward_list))
+        print("Eval  |  Total Steps {}  |  Episodes {}  |  Average Reward {:.2f}  |  Max reward {:.2f}  |  "
+              "Min reward {:.2f}".format(self.total_step, self.episode, sum(reward_list) / len(reward_list),
+                                          max(reward_list), min(reward_list)))
+
+
+
+
+
+
+
+
+
+

utils.py

@@ -0,0 +1,82 @@
+import numpy as np
+import random
+import torch
+import torch.nn as nn
+from wrapper import DelayedEnv
+
+
+def weight_init(m):
+    """Custom weight init for Conv2D and Linear layers.
+        Reference: https://github.com/MishaLaskin/rad/blob/master/curl_sac.py"""
+
+    if isinstance(m, nn.Linear):
+        nn.init.orthogonal_(m.weight.data)
+        m.bias.data.fill_(0.0)
+    elif isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
+        # delta-orthogonal init from https://arxiv.org/pdf/1806.05393.pdf
+        assert m.weight.size(2) == m.weight.size(3)
+        m.weight.data.fill_(0.0)
+        m.bias.data.fill_(0.0)
+        mid = m.weight.size(2) // 2
+        gain = nn.init.calculate_gain('relu')
+        nn.init.orthogonal_(m.weight.data[:, :, mid, mid], gain)
+
+
+def hard_update(network, target_network):
+    with torch.no_grad():
+        for param, target_param in zip(network.parameters(), target_network.parameters()):
+            target_param.data.copy_(param.data)
+
+
+def soft_update(network, target_network, tau):
+    with torch.no_grad():
+        for param, target_param in zip(network.parameters(), target_network.parameters()):
+            target_param.data.copy_(target_param.data * (1.0 - tau) + param.data * tau)
+
+
+def set_seed(random_seed):
+    if random_seed <= 0:
+        random_seed = np.random.randint(1, 9999)
+    else:
+        random_seed = random_seed
+
+    torch.manual_seed(random_seed)
+    np.random.seed(random_seed)
+    random.seed(random_seed)
+
+    return random_seed
+
+
+def make_env(env_name, random_seed):
+    import gym
+    # openai gym
+    env = gym.make(env_name)
+    env.seed(random_seed)
+    env.action_space.seed(random_seed)
+
+    eval_env = gym.make(env_name)
+    eval_env.seed(random_seed)
+    eval_env.action_space.seed(random_seed)
+
+    return env, eval_env
+
+
+def make_delayed_env(args, random_seed, obs_delayed_steps, act_delayed_steps):
+    import gym
+    # openai gym
+    env_name = args.env_name
+
+    env = gym.make(env_name)
+    delayed_env = DelayedEnv(env, seed=random_seed, obs_delayed_steps=obs_delayed_steps, act_delayed_steps=act_delayed_steps)
+
+    eval_env = gym.make(env_name)
+    eval_delayed_env = DelayedEnv(eval_env, seed=random_seed, obs_delayed_steps=obs_delayed_steps, act_delayed_steps=act_delayed_steps)
+
+    return delayed_env, eval_delayed_env
+
+def log_to_txt(env_name, random_seed, total_step, result):
+    seed = '(' + str(random_seed) + ')'
+    f = open('./log/' + env_name + '_seed' + seed + '.txt', 'a')
+    log = str(total_step) + ' ' + str(result) + '\n'
+    f.write(log)
+    f.close()

wrapper.py

@@ -0,0 +1,65 @@
+from collections import deque
+import gym
+import numpy as np
+
+
+class DelayedEnv(gym.Wrapper):
+    def __init__(self, env, seed, obs_delayed_steps, act_delayed_steps):
+        super(DelayedEnv, self).__init__(env)
+        assert obs_delayed_steps + act_delayed_steps > 0
+        self.env.action_space.seed(seed)
+
+        self.observation_space = self.env.observation_space
+        self.action_space = self.env.action_space
+
+        self._max_episode_steps = self.env._max_episode_steps
+
+        self.obs_buffer = deque(maxlen=obs_delayed_steps)
+        self.reward_buffer = deque(maxlen=obs_delayed_steps)
+        self.done_buffer = deque(maxlen=obs_delayed_steps)
+
+        self.action_buffer = deque(maxlen=act_delayed_steps)
+
+        self.obs_delayed_steps = obs_delayed_steps
+        self.act_delayed_steps = act_delayed_steps
+
+    def reset(self):
+        for _ in range(self.act_delayed_steps):
+            self.action_buffer.append(np.zeros_like(self.env.action_space.sample()))
+
+        init_state, _ = self.env.reset()
+        for _ in range(self.obs_delayed_steps):
+            self.obs_buffer.append(init_state)
+            self.reward_buffer.append(0)
+            self.done_buffer.append(False)
+        return init_state
+
+    def step(self, action):
+        if self.act_delayed_steps > 0:
+            delayed_action = self.action_buffer.popleft()
+            self.action_buffer.append(action)
+        else:
+            delayed_action = action
+
+        current_obs, current_reward, current_terminated, current_truncated, _ = self.env.step(delayed_action)
+        current_done = current_terminated or current_truncated
+
+        if self.obs_delayed_steps > 0:
+            delayed_obs = self.obs_buffer.popleft()
+            delayed_reward = self.reward_buffer.popleft()
+            delayed_done = self.done_buffer.popleft()
+
+            self.obs_buffer.append(current_obs)
+            self.reward_buffer.append(current_reward)
+            self.done_buffer.append(current_done)
+        else:
+            delayed_obs = current_obs
+            delayed_reward = current_reward
+            delayed_done = current_done
+
+        return delayed_obs, delayed_reward, delayed_done, {'current_obs': current_obs, 'current_reward': current_reward,
+                                                           'current_done': current_done}
+
+
+
+