Add cagct env scripts

2026-09-12 12:31:54 +08:00 · 2019-04-21 12:31:37 +09:00
parent a78f2bddb7
commit 8816c42a7d
10 changed files with 2772 additions and 13 deletions
@@ -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
@@ -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)
+
@@ -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