Add cagct env scripts

scripts/algorithms/fd/td3_agent.py

@@ -39,8 +39,9 @@ class Agent(TD3Agent):
             # load demo replay memory
             # TODO: should make new demo to set protocol 2
             #       e.g. pickle.dump(your_object, your_file, protocol=2)
-            with open(self.args.demo_path, "rb") as f:
-                demos = pickle.load(f)
+#            with open(self.args.demo_path, "rb") as f:
+#                demos = pickle.load(f)
+            demos = []
 
             if self.use_n_step:
                 demos, demos_n_step = common_utils.get_n_step_info_from_demo(
@@ -91,20 +92,17 @@ class Agent(TD3Agent):
         )
         next_actions = (self.actor_target(next_states) + clipped_noise).clamp(-1.0, 1.0)
 
-        target_values1 = self.critic1_target(
-            torch.cat((next_states, next_actions), dim=-1)
-        )
-        target_values2 = self.critic2_target(
-            torch.cat((next_states, next_actions), dim=-1)
-        )
+        target_values1 = self.critic1_target(next_states, next_actions)
+        target_values2 = self.critic2_target(next_states, next_actions)
+
         target_values = torch.min(target_values1, target_values2)
         target_values = rewards + (gamma * target_values * masks).detach()
 
         # train critic
-        values1 = self.critic1(torch.cat((states, actions), dim=-1))
+        values1 = self.critic1(next_states, next_actions)
         critic1_loss_element_wise = (values1 - target_values.detach()).pow(2)
 
-        values2 = self.critic2(torch.cat((states, actions), dim=-1))
+        values2 = self.critic2(next_states, next_actions)
         critic2_loss_element_wise = (values2 - target_values.detach()).pow(2)
 
         return critic1_loss_element_wise, critic2_loss_element_wise
@@ -115,6 +113,12 @@ class Agent(TD3Agent):
         states, actions, rewards, next_states, dones, weights, indices, eps_d = (
             experiences
         )
+        # monkey spanner
+        states = states.view(-1, 1, 40, 40, 40)
+        next_states = next_states.view(-1, 1, 40, 40, 40)
+        actions = actions.view(-1, 1)
+        rewards = rewards.view(-1, 1)
+        experiences = states, actions, rewards, next_states, dones, weights, indices, eps_d
 
         gamma = self.hyper_params["GAMMA"]
         critic1_loss_element_wise, critic2_loss_element_wise = self._get_critic_loss(
@@ -149,9 +153,7 @@ class Agent(TD3Agent):
         if self.episode_steps % self.hyper_params["POLICY_UPDATE_FREQ"] == 0:
             # train actor
             actions = self.actor(states)
-            actor_loss_element_wise = -self.critic1(
-                torch.cat((states, actions), dim=-1)
-            )
+            actor_loss_element_wise = -self.critic1(states, actions)
             actor_loss = torch.mean(actor_loss_element_wise * weights)
             self.actor_optim.zero_grad()
             actor_loss.backward()

scripts/config/agent/cagct/network.py

@@ -0,0 +1,66 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .utils import *
+
+
+device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+
+
+class Actor(nn.Module):
+
+    def __init__(self, in_channel, action_dim, max_action):
+        super(Actor, self).__init__()
+
+        self.conv1 = nn.Conv3d(in_channel, 3, (3, 3, 3), 2, 1)
+        self.conv2 = nn.Conv3d(3, 128, (3, 3, 3), 2, 1)
+        self.conv3 = nn.Conv3d(128, 256, (3, 3, 3), 2, 1)
+        self.avg_pool4 = nn.AvgPool3d(5)
+        self.fc5 = nn.Linear(256, 512)
+        self.fc5.weight.data.uniform_(-3e-3, 3e-3)
+        self.fc5.bias.data.uniform_(-3e-3, 3e-3)
+        self.fc6 = nn.Linear(512, action_dim)
+        self.fc6.weight.data.uniform_(-3e-3, 3e-3)
+        self.fc6.bias.data.uniform_(-3e-3, 3e-3)
+
+        self.max_action = max_action
+
+    def forward(self, x):
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = F.relu(self.conv3(x))
+        x = self.avg_pool4(x)
+        x = F.relu(self.fc5(x.view(x.size(0), -1)))
+        x = self.fc6(x)
+        x = self.max_action * torch.tanh(x)
+        return x.squeeze()
+
+
+class Critic(nn.Module):
+
+    def __init__(self, in_channel, action_dim):
+        super(Critic, self).__init__()
+
+        self.conv1 = nn.Conv3d(in_channel, 3, (3, 3, 3), 2, 1)
+        self.conv2 = nn.Conv3d(3, 128, (3, 3, 3), 2, 1)
+        self.conv3 = nn.Conv3d(128, 256, (3, 3, 3), 2, 1)
+        self.avg_pool4 = nn.AvgPool3d(5)
+
+        self.fc5 = nn.Linear(256 + action_dim, 512)
+        self.fc5.weight.data.uniform_(-3e-3, 3e-3)
+        self.fc5.bias.data.uniform_(-3e-3, 3e-3)
+        self.fc6 = nn.Linear(512, 1)
+        self.fc6.weight.data.uniform_(-3e-3, 3e-3)
+        self.fc6.bias.data.uniform_(-3e-3, 3e-3)
+
+    def forward(self, x, u):
+        x = F.relu(self.conv1(x))
+        x = F.relu(self.conv2(x))
+        x = F.relu(self.conv3(x))
+        x = self.avg_pool4(x)
+        xu = torch.cat([x.view(x.size(0), -1), u], 1)
+
+        x = F.relu(self.fc5(xu))
+        x = self.fc6(x)
+        return x

scripts/config/agent/cagct/td3fd.py

@@ -0,0 +1,113 @@
+# -*- coding: utf-8 -*-
+"""Run module for SACfD on LunarLanderContinuous-v2.
+
+- Author: Seungjae Ryan Lee
+- Contact: seungjaeryanlee@gmail.com
+"""
+
+import torch
+import torch.optim as optim
+
+from algorithms.common.networks.mlp import MLP
+from algorithms.common.noise import GaussianNoise
+from algorithms.fd.td3_agent import Agent
+from .network import Actor, Critic
+
+
+device = torch.device("cuda:0" if torch.cuda.is_available() else "cpu")
+
+# hyper parameters
+# TODO Tune hyperparameters on LunarLander-v2
+
+# hyper parameters
+hyper_params = {
+    "N_STEP": 1,
+    "GAMMA": 0.99,
+    "TAU": 5e-3,
+    "BUFFER_SIZE": int(3e4),
+    "BATCH_SIZE": 100,
+    "LR_ACTOR": 1e-3,
+    "LR_CRITIC": 1e-3,
+    "EXPLORATION_NOISE": 0.1,
+    "TARGET_POLICY_NOISE": 0.2,
+    "TARGET_POLICY_NOISE_CLIP": 0.5,
+    "POLICY_UPDATE_FREQ": 2,
+    "INITIAL_RANDOM_ACTIONS": 1000,
+    "PRETRAIN_STEP": 1000,
+    "MULTIPLE_LEARN": 2,  # multiple learning updates
+    "LAMBDA1": 1.0,  # N-step return weight
+    "LAMBDA2": 1e-5,  # l2 regularization weight
+    "LAMBDA3": 1.0,  # actor loss contribution of prior weight
+    "PER_ALPHA": 0.3,
+    "PER_BETA": 1.0,
+    "PER_EPS": 1e-6,
+    "PER_EPS_DEMO": 1.0,
+    "NETWORK": {"ACTOR_HIDDEN_SIZES": [400, 300], "CRITIC_HIDDEN_SIZES": [400, 300]},
+}
+
+
+def get(env, args):
+    """Run training or test.
+
+    Args:
+        env (gym.Env): openAI Gym environment with continuous action space
+        args (argparse.Namespace): arguments including training settings
+
+    """
+    state_dim = 1
+    action_dim = env.action_space.shape[0]
+
+    # create actor
+    actor = Actor(state_dim, action_dim, 1).to(device)
+    actor_target = Actor(state_dim, action_dim, 1).to(device)
+    actor_target.load_state_dict(actor.state_dict())
+
+    # create critic1
+    critic1 = Critic(state_dim, action_dim).to(device)
+
+    critic1_target = Critic(state_dim, action_dim).to(device)
+    critic1_target.load_state_dict(critic1.state_dict())
+
+    # create critic2
+    critic2 = Critic(state_dim, action_dim).to(device)
+
+    critic2_target = Critic(state_dim, action_dim).to(device)
+    critic2_target.load_state_dict(critic2.state_dict())
+
+    # concat critic parameters to use one optim
+    critic_parameters = list(critic1.parameters()) + list(critic2.parameters())
+
+    # create optimizer
+    actor_optim = optim.Adam(
+        actor.parameters(),
+        lr=hyper_params["LR_ACTOR"],
+        weight_decay=hyper_params["LAMBDA2"],
+    )
+
+    critic_optim = optim.Adam(
+        critic_parameters,
+        lr=hyper_params["LR_CRITIC"],
+        weight_decay=hyper_params["LAMBDA2"],
+    )
+
+    # noise
+    exploration_noise = GaussianNoise(
+        action_dim,
+        min_sigma=hyper_params["EXPLORATION_NOISE"],
+        max_sigma=hyper_params["EXPLORATION_NOISE"],
+    )
+
+    target_policy_noise = GaussianNoise(
+        action_dim,
+        min_sigma=hyper_params["TARGET_POLICY_NOISE"],
+        max_sigma=hyper_params["TARGET_POLICY_NOISE"],
+    )
+
+    # make tuples to create an agent
+    models = (actor, actor_target, critic1, critic1_target, critic2, critic2_target)
+    optims = (actor_optim, critic_optim)
+    noises = (exploration_noise, target_policy_noise)
+
+    # create an agent
+    return Agent(env, args, hyper_params, models, optims, noises)
+

scripts/config/agent/cagct/utils.py

@@ -0,0 +1,158 @@
+from copy import deepcopy
+from math import pi, ceil
+
+import scipy.ndimage
+from scipy.spatial import distance
+import numpy as np
+import matplotlib.pyplot as plt
+import trimesh
+from trimesh import voxel
+
+
+def load_stl(stl_path):
+    return trimesh.load(stl_path)
+
+
+def stl_to_arr(stl_mesh, pad=True):
+    arr = voxel.VoxelMesh(mesh=stl_mesh, pitch=1).matrix_solid
+    if pad:
+        max_edge_dist = get_max_bbox_edge(stl_mesh)
+        arr = center_constant_pad(arr, max_edge_dist)
+    return arr
+
+
+def generate_label(stl_mesh):
+    random_angle = np.random.choice(np.arange(-30, 30), 3, replace=True)
+    stl_mesh_random_rotated = rotate_stl(
+        stl_mesh, random_angle, axes=(1, 1, 1))
+    label_3d_arr = stl_to_arr(stl_mesh_random_rotated)
+    label_2d_arr = project_3d_arr_to_2d_arr(label_3d_arr)
+    return label_3d_arr, label_2d_arr, random_angle
+
+
+def get_max_bbox_edge(stl_mesh):
+    vertices = stl_mesh.bounding_box.vertices
+    edge_dists = distance.cdist(vertices, vertices, 'euclidean')
+    max_edge_dist = int(edge_dists.max())  # temporary disable
+    return 40
+
+
+def center_constant_pad(arr, target_shape):
+    x, y, z = arr.shape
+    x_num_pad_left = (target_shape - x) // 2
+    x_num_pad_right = ceil((target_shape - x) / 2)
+    y_num_pad_left = (target_shape - y) // 2
+    y_num_pad_right = ceil((target_shape - y) / 2)
+    z_num_pad_left = (target_shape - z) // 2
+    z_num_pad_right = ceil((target_shape - z) / 2)
+    return np.pad(arr, ((x_num_pad_left, x_num_pad_right), (y_num_pad_left,
+        y_num_pad_right), (z_num_pad_left, z_num_pad_right)), 'constant')
+
+
+def rotate_stl(stl_mesh, angles, axes=(1, 1, 1)):
+    stl_mesh_rotated = deepcopy(stl_mesh)
+    rads = [i * pi / 180 for i in angles]
+    x_rad, y_rad, z_rad = rads
+    axis_x, axis_y, axis_z = axes
+
+    if axis_x:
+        stl_mesh_rotated.apply_transform(
+            trimesh.transformations.rotation_matrix(x_rad,
+            (1, 0, 0)))
+    if axis_y:
+        stl_mesh_rotated.apply_transform(
+            trimesh.transformations.rotation_matrix(y_rad,
+            (0, 1, 0)))
+    if axis_z:
+        stl_mesh_rotated.apply_transform(
+            trimesh.transformations.rotation_matrix(z_rad,
+            (0, 0, 1)))
+    return stl_mesh_rotated
+
+
+def load_3d_img(npy_path):
+    arr_3d = np.load(npy_path)
+    assert arr_3d.ndim == 3
+    return arr_3d
+
+
+def rotate_3d_arr(arr, rotate=(0, 0, 0)):
+    x_angle, y_angle, z_angle = rotate
+    arr_x_rotated = scipy.ndimage.interpolation.rotate(
+        arr, x_angle, mode="nearest", axes=(1, 2), reshape=False
+    )
+    arr_y_rotated = scipy.ndimage.interpolation.rotate(
+        arr_x_rotated, y_angle, mode="nearest", axes=(0, 2), reshape=False
+    )
+    arr_z_rotated = scipy.ndimage.interpolation.rotate(
+        arr_y_rotated, z_angle, mode="nearest", axes=(0, 1), reshape=False
+    )
+    return arr_z_rotated
+
+
+def project_3d_arr_to_2d_arr(arr_3d, axis=-1):
+    arr_2d = arr_3d.max(axis)
+    return arr_2d
+
+
+def get_iou(pred_img, target_img):
+    union = (pred_img + target_img).sum()
+    num_pred_ones = (pred_img).sum()
+    num_target_ones = (target_img).sum()
+    intersection = num_pred_ones + num_target_ones - union
+    iou = intersection / num_target_ones
+    return iou
+
+
+def plot_3d_arr(arr_3d):
+    fig = plt.figure()
+    ax = fig.gca(projection="3d")
+    ax.voxels(arr_3d, edgecolor="k")
+    plt.show()
+
+
+def plot_2d_img(arr_2d):
+    plt.figure()
+    plt.imshow(arr_2d)
+    plt.show()
+
+
+def plot_stl_mesh(stl_mesh):
+    stl_mesh.show()
+
+
+def unnormalize_action(action):
+    return [i * 30 for i in action]
+
+
+class ReplayBuffer(object):
+    def __init__(self, max_size=1e6):
+        self.storage = []
+        self.max_size = max_size
+        self.ptr = 0
+
+    def add(self, data):
+        if len(self.storage) == self.max_size:
+            self.storage[int(self.ptr)] = data
+            self.ptr = (self.ptr + 1) % self.max_size
+        else:
+            self.storage.append(data)
+
+    def sample(self, batch_size):
+        ind = np.random.randint(0, len(self.storage), size=batch_size)
+        x, y, u, r, d = [], [], [], [], []
+
+        for i in ind:
+            X, Y, U, R, D = self.storage[i]
+            x.append(np.array(X, copy=False))
+            y.append(np.array(Y, copy=False))
+            u.append(np.array(U, copy=False))
+            r.append(np.array(R, copy=False))
+            d.append(np.array(D, copy=False))
+
+        x = np.array(x).reshape(-1, 1, *X.shape)
+        y = np.array(y).reshape(-1, 1, *Y.shape)
+        u = np.array(u)
+        r = np.array(r).reshape(-1, 1)
+        d = np.array(d).reshape(-1, 1)
+        return x, y, u, r, d

scripts/config/environment/cagct.py

@@ -0,0 +1,11 @@
+import numpy as np
+
+
+env_name = "CAGCTRegistratorEnv"
+stl_path = "./iSight_clamp.stl"
+stl_resize_ratio = 0.3
+max_action = 1
+action_low = np.array([-1, -1, -1])
+action_high = np.array([1, 1, 1])
+max_episode_steps = 50
+succeed_iou = 0.9

scripts/config/environment/utils.py

@@ -0,0 +1,158 @@
+from copy import deepcopy
+from math import pi, ceil
+
+import scipy.ndimage
+from scipy.spatial import distance
+import numpy as np
+import matplotlib.pyplot as plt
+import trimesh
+from trimesh import voxel
+
+
+def load_stl(stl_path):
+    return trimesh.load(stl_path)
+
+
+def stl_to_arr(stl_mesh, pad=True):
+    arr = voxel.VoxelMesh(mesh=stl_mesh, pitch=1).matrix_solid
+    if pad:
+        max_edge_dist = get_max_bbox_edge(stl_mesh)
+        arr = center_constant_pad(arr, max_edge_dist)
+    return arr
+
+
+def generate_label(stl_mesh):
+    random_angle = np.random.choice(np.arange(-30, 30), 3, replace=True)
+    stl_mesh_random_rotated = rotate_stl(
+        stl_mesh, random_angle, axes=(1, 1, 1))
+    label_3d_arr = stl_to_arr(stl_mesh_random_rotated)
+    label_2d_arr = project_3d_arr_to_2d_arr(label_3d_arr)
+    return label_3d_arr, label_2d_arr, random_angle
+
+
+def get_max_bbox_edge(stl_mesh):
+    vertices = stl_mesh.bounding_box.vertices
+    edge_dists = distance.cdist(vertices, vertices, 'euclidean')
+    max_edge_dist = int(edge_dists.max())  # temporary disable
+    return 40
+
+
+def center_constant_pad(arr, target_shape):
+    x, y, z = arr.shape
+    x_num_pad_left = (target_shape - x) // 2
+    x_num_pad_right = ceil((target_shape - x) / 2)
+    y_num_pad_left = (target_shape - y) // 2
+    y_num_pad_right = ceil((target_shape - y) / 2)
+    z_num_pad_left = (target_shape - z) // 2
+    z_num_pad_right = ceil((target_shape - z) / 2)
+    return np.pad(arr, ((x_num_pad_left, x_num_pad_right), (y_num_pad_left,
+        y_num_pad_right), (z_num_pad_left, z_num_pad_right)), 'constant')
+
+
+def rotate_stl(stl_mesh, angles, axes=(1, 1, 1)):
+    stl_mesh_rotated = deepcopy(stl_mesh)
+    rads = [i * pi / 180 for i in angles]
+    x_rad, y_rad, z_rad = rads
+    axis_x, axis_y, axis_z = axes
+
+    if axis_x:
+        stl_mesh_rotated.apply_transform(
+            trimesh.transformations.rotation_matrix(x_rad,
+            (1, 0, 0)))
+    if axis_y:
+        stl_mesh_rotated.apply_transform(
+            trimesh.transformations.rotation_matrix(y_rad,
+            (0, 1, 0)))
+    if axis_z:
+        stl_mesh_rotated.apply_transform(
+            trimesh.transformations.rotation_matrix(z_rad,
+            (0, 0, 1)))
+    return stl_mesh_rotated
+
+
+def load_3d_img(npy_path):
+    arr_3d = np.load(npy_path)
+    assert arr_3d.ndim == 3
+    return arr_3d
+
+
+def rotate_3d_arr(arr, rotate=(0, 0, 0)):
+    x_angle, y_angle, z_angle = rotate
+    arr_x_rotated = scipy.ndimage.interpolation.rotate(
+        arr, x_angle, mode="nearest", axes=(1, 2), reshape=False
+    )
+    arr_y_rotated = scipy.ndimage.interpolation.rotate(
+        arr_x_rotated, y_angle, mode="nearest", axes=(0, 2), reshape=False
+    )
+    arr_z_rotated = scipy.ndimage.interpolation.rotate(
+        arr_y_rotated, z_angle, mode="nearest", axes=(0, 1), reshape=False
+    )
+    return arr_z_rotated
+
+
+def project_3d_arr_to_2d_arr(arr_3d, axis=-1):
+    arr_2d = arr_3d.max(axis)
+    return arr_2d
+
+
+def get_iou(pred_img, target_img):
+    union = (pred_img + target_img).sum()
+    num_pred_ones = (pred_img).sum()
+    num_target_ones = (target_img).sum()
+    intersection = num_pred_ones + num_target_ones - union
+    iou = intersection / num_target_ones
+    return iou
+
+
+def plot_3d_arr(arr_3d):
+    fig = plt.figure()
+    ax = fig.gca(projection="3d")
+    ax.voxels(arr_3d, edgecolor="k")
+    plt.show()
+
+
+def plot_2d_img(arr_2d):
+    plt.figure()
+    plt.imshow(arr_2d)
+    plt.show()
+
+
+def plot_stl_mesh(stl_mesh):
+    stl_mesh.show()
+
+
+def unnormalize_action(action):
+    return [i * 30 for i in action]
+
+
+class ReplayBuffer(object):
+    def __init__(self, max_size=1e6):
+        self.storage = []
+        self.max_size = max_size
+        self.ptr = 0
+
+    def add(self, data):
+        if len(self.storage) == self.max_size:
+            self.storage[int(self.ptr)] = data
+            self.ptr = (self.ptr + 1) % self.max_size
+        else:
+            self.storage.append(data)
+
+    def sample(self, batch_size):
+        ind = np.random.randint(0, len(self.storage), size=batch_size)
+        x, y, u, r, d = [], [], [], [], []
+
+        for i in ind:
+            X, Y, U, R, D = self.storage[i]
+            x.append(np.array(X, copy=False))
+            y.append(np.array(Y, copy=False))
+            u.append(np.array(U, copy=False))
+            r.append(np.array(R, copy=False))
+            d.append(np.array(D, copy=False))
+
+        x = np.array(x).reshape(-1, 1, *X.shape)
+        y = np.array(y).reshape(-1, 1, *Y.shape)
+        u = np.array(u)
+        r = np.array(r).reshape(-1, 1)
+        d = np.array(d).reshape(-1, 1)
+        return x, y, u, r, d

scripts/envs/cagct/cagct.py

@@ -0,0 +1,54 @@
+import gym
+from trimesh import voxel
+
+from .utils import *
+
+
+class CAGCTRegistratorEnv(gym.Env):
+
+    def __init__(self, cfg):
+        self.env_name = cfg.env_name
+        self.cfg = cfg
+        self.action_space = gym.spaces.Box(low=self.cfg.action_low,
+            high=self.cfg.action_high, dtype=np.float32)
+        self._max_episode_steps = self.cfg.max_episode_steps
+        self.episode_steps = 0
+
+    def compute_reward(self):
+        state_arr = stl_to_arr(self.state_stl)
+        projected_state_arr = project_3d_arr_to_2d_arr(state_arr)
+        self.iou = get_iou(projected_state_arr, self.label_2d_arr)
+        reward = self.iou - 1
+        return reward
+
+    def reset(self):
+        self.state_stl = load_stl(self.cfg.stl_path)
+        # resize stl
+        self.state_stl.apply_scale(self.cfg.stl_resize_ratio)
+        _, self.label_2d_arr, _ = generate_label(self.state_stl)
+        return self.get_observation()
+
+    def get_observation(self):
+        state_3d_arr = stl_to_arr(self.state_stl).astype('float')
+        return state_3d_arr.reshape(1, 1, *state_3d_arr.shape)
+
+    def step(self, action):
+        done = False
+        succeed = None
+        self.episode_steps += 1
+
+        action = unnormalize_action(action)
+        self.state_stl = rotate_stl(self.state_stl, action)
+        obs = self.get_observation()
+        reward = self.compute_reward()
+        if self.iou > self.cfg.succeed_iou:
+            done = True
+            succeed = True
+        if self.episode_steps == self._max_episode_steps:
+            done = True
+            self.episode_steps = 0
+        return obs, reward, done, succeed
+
+    def render(self):
+        pass
+

scripts/envs/cagct/utils.py

@@ -0,0 +1,158 @@
+from copy import deepcopy
+from math import pi, ceil
+
+import scipy.ndimage
+from scipy.spatial import distance
+import numpy as np
+import matplotlib.pyplot as plt
+import trimesh
+from trimesh import voxel
+
+
+def load_stl(stl_path):
+    return trimesh.load(stl_path)
+
+
+def stl_to_arr(stl_mesh, pad=True):
+    arr = voxel.VoxelMesh(mesh=stl_mesh, pitch=1).matrix_solid
+    if pad:
+        max_edge_dist = get_max_bbox_edge(stl_mesh)
+        arr = center_constant_pad(arr, max_edge_dist)
+    return arr
+
+
+def generate_label(stl_mesh):
+    random_angle = np.random.choice(np.arange(-30, 30), 3, replace=True)
+    stl_mesh_random_rotated = rotate_stl(
+        stl_mesh, random_angle, axes=(1, 1, 1))
+    label_3d_arr = stl_to_arr(stl_mesh_random_rotated)
+    label_2d_arr = project_3d_arr_to_2d_arr(label_3d_arr)
+    return label_3d_arr, label_2d_arr, random_angle
+
+
+def get_max_bbox_edge(stl_mesh):
+    vertices = stl_mesh.bounding_box.vertices
+    edge_dists = distance.cdist(vertices, vertices, 'euclidean')
+    max_edge_dist = int(edge_dists.max())  # temporary disable
+    return 40
+
+
+def center_constant_pad(arr, target_shape):
+    x, y, z = arr.shape
+    x_num_pad_left = (target_shape - x) // 2
+    x_num_pad_right = ceil((target_shape - x) / 2)
+    y_num_pad_left = (target_shape - y) // 2
+    y_num_pad_right = ceil((target_shape - y) / 2)
+    z_num_pad_left = (target_shape - z) // 2
+    z_num_pad_right = ceil((target_shape - z) / 2)
+    return np.pad(arr, ((x_num_pad_left, x_num_pad_right), (y_num_pad_left,
+        y_num_pad_right), (z_num_pad_left, z_num_pad_right)), 'constant')
+
+
+def rotate_stl(stl_mesh, angles, axes=(1, 1, 1)):
+    stl_mesh_rotated = deepcopy(stl_mesh)
+    rads = [i * pi / 180 for i in angles]
+    x_rad, y_rad, z_rad = rads
+    axis_x, axis_y, axis_z = axes
+
+    if axis_x:
+        stl_mesh_rotated.apply_transform(
+            trimesh.transformations.rotation_matrix(x_rad,
+            (1, 0, 0)))
+    if axis_y:
+        stl_mesh_rotated.apply_transform(
+            trimesh.transformations.rotation_matrix(y_rad,
+            (0, 1, 0)))
+    if axis_z:
+        stl_mesh_rotated.apply_transform(
+            trimesh.transformations.rotation_matrix(z_rad,
+            (0, 0, 1)))
+    return stl_mesh_rotated
+
+
+def load_3d_img(npy_path):
+    arr_3d = np.load(npy_path)
+    assert arr_3d.ndim == 3
+    return arr_3d
+
+
+def rotate_3d_arr(arr, rotate=(0, 0, 0)):
+    x_angle, y_angle, z_angle = rotate
+    arr_x_rotated = scipy.ndimage.interpolation.rotate(
+        arr, x_angle, mode="nearest", axes=(1, 2), reshape=False
+    )
+    arr_y_rotated = scipy.ndimage.interpolation.rotate(
+        arr_x_rotated, y_angle, mode="nearest", axes=(0, 2), reshape=False
+    )
+    arr_z_rotated = scipy.ndimage.interpolation.rotate(
+        arr_y_rotated, z_angle, mode="nearest", axes=(0, 1), reshape=False
+    )
+    return arr_z_rotated
+
+
+def project_3d_arr_to_2d_arr(arr_3d, axis=-1):
+    arr_2d = arr_3d.max(axis)
+    return arr_2d
+
+
+def get_iou(pred_img, target_img):
+    union = (pred_img + target_img).sum()
+    num_pred_ones = (pred_img).sum()
+    num_target_ones = (target_img).sum()
+    intersection = num_pred_ones + num_target_ones - union
+    iou = intersection / num_target_ones
+    return iou
+
+
+def plot_3d_arr(arr_3d):
+    fig = plt.figure()
+    ax = fig.gca(projection="3d")
+    ax.voxels(arr_3d, edgecolor="k")
+    plt.show()
+
+
+def plot_2d_img(arr_2d):
+    plt.figure()
+    plt.imshow(arr_2d)
+    plt.show()
+
+
+def plot_stl_mesh(stl_mesh):
+    stl_mesh.show()
+
+
+def unnormalize_action(action):
+    return [i * 30 for i in action]
+
+
+class ReplayBuffer(object):
+    def __init__(self, max_size=1e6):
+        self.storage = []
+        self.max_size = max_size
+        self.ptr = 0
+
+    def add(self, data):
+        if len(self.storage) == self.max_size:
+            self.storage[int(self.ptr)] = data
+            self.ptr = (self.ptr + 1) % self.max_size
+        else:
+            self.storage.append(data)
+
+    def sample(self, batch_size):
+        ind = np.random.randint(0, len(self.storage), size=batch_size)
+        x, y, u, r, d = [], [], [], [], []
+
+        for i in ind:
+            X, Y, U, R, D = self.storage[i]
+            x.append(np.array(X, copy=False))
+            y.append(np.array(Y, copy=False))
+            u.append(np.array(U, copy=False))
+            r.append(np.array(R, copy=False))
+            d.append(np.array(D, copy=False))
+
+        x = np.array(x).reshape(-1, 1, *X.shape)
+        y = np.array(y).reshape(-1, 1, *Y.shape)
+        u = np.array(u)
+        r = np.array(r).reshape(-1, 1)
+        d = np.array(d).reshape(-1, 1)
+        return x, y, u, r, d

scripts/run_cagct.py

@@ -0,0 +1,77 @@
+#! /usr/bin/env python
+
+# -*- coding: utf-8 -*-
+"""Train or test algorithms on OpenManipulator Reacher-v0 on Gazebo.
+
+- Author: Kh Kim
+- Contact: kh.kim@medipixel.io
+"""
+
+import argparse
+import importlib
+
+import algorithms.common.helper_functions as common_utils
+from config.environment import cagct as env_cfg
+from envs.cagct.cagct import CAGCTRegistratorEnv
+
+# configurations
+parser = argparse.ArgumentParser(description="Pytorch RL algorithms")
+parser.add_argument(
+    "--seed", type=int, default=777, help="random seed for reproducibility"
+)
+parser.add_argument("--algo", type=str, default="td3fd", help="choose an algorithm")
+parser.add_argument(
+    "--test", dest="test", action="store_true", help="test mode (no training)"
+)
+parser.add_argument(
+    "--load-from", type=str, help="load the saved model and optimizer at the beginning"
+)
+parser.add_argument(
+    "--off-render", dest="render", action="store_false", help="turn off rendering"
+)
+parser.add_argument(
+    "--render-after",
+    type=int,
+    default=0,
+    help="start rendering after the input number of episode",
+)
+parser.add_argument("--log", dest="log", action="store_true", help="turn on logging")
+parser.add_argument("--save-period", type=int, default=200, help="save model period")
+parser.add_argument("--episode-num", type=int, default=20000, help="total episode num")
+parser.add_argument(
+    "--max-episode-steps", type=int, default=-1, help="max episode step"
+)
+parser.add_argument(
+    "--demo-path", type=str, default="data/reacher_demo.pkl", help="demonstration path"
+)
+
+parser.set_defaults(test=False)
+parser.set_defaults(load_from=None)
+parser.set_defaults(render=True)
+parser.set_defaults(log=False)
+args = parser.parse_args()
+
+
+def main():
+    """Main."""
+    # env initialization
+    env = CAGCTRegistratorEnv(env_cfg)
+
+    # set a random seed
+    common_utils.set_random_seed(args.seed, env)
+
+    # agent initialization
+    module_path = "config.agent.cagct." + args.algo
+    agent = importlib.import_module(module_path)
+    agent = agent.get(env, args)
+
+    # run
+    if args.test:
+        agent.test()
+    else:
+        agent.train()
+
+
+if __name__ == "__main__":
+    main()
+