def __init__(self,
input_dim,
hidden_dim,
num_objects,
output_size,
act_fn='relu'):
super(DecoderCNNMedium, self).__init__()
width, height = output_size[1] // 5, output_size[2] // 5
output_dim = width * height
self.fc1 = nn.Linear(input_dim, hidden_dim)
self.fc2 = nn.Linear(hidden_dim, hidden_dim)
self.fc3 = nn.Linear(hidden_dim, output_dim)
self.ln = nn.LayerNorm(hidden_dim)
self.deconv1 = nn.ConvTranspose2d(num_objects,
hidden_dim,
kernel_size=5,
stride=5)
self.deconv2 = nn.ConvTranspose2d(hidden_dim,
output_size[0],
kernel_size=9,
padding=4)
self.ln1 = nn.BatchNorm2d(hidden_dim)
self.input_dim = input_dim
self.num_objects = num_objects
self.map_size = output_size[0], width, height
self.act1 = util.get_act_fn(act_fn)
self.act2 = util.get_act_fn(act_fn)
self.act3 = util.get_act_fn(act_fn)
def __init__(self,
input_dim,
hidden_dim,
action_dim,
num_objects,
residual=True,
act_fn='relu'):
super(TransitionGNN, self).__init__()
self.input_dim = input_dim
self.hidden_dim = hidden_dim
self.num_objects = num_objects
self.action_dim = action_dim
self.residual = residual
self.edge_mlp = nn.Sequential(nn.Linear(input_dim * 2, hidden_dim),
util.get_act_fn(act_fn),
nn.Linear(hidden_dim, hidden_dim),
nn.LayerNorm(hidden_dim),
util.get_act_fn(act_fn),
nn.Linear(hidden_dim, hidden_dim))
node_input_dim = hidden_dim + input_dim + self.action_dim
self.node_mlp = nn.Sequential(nn.Linear(node_input_dim, hidden_dim),
util.get_act_fn(act_fn),
nn.Linear(hidden_dim, hidden_dim),
nn.LayerNorm(hidden_dim),
util.get_act_fn(act_fn),
nn.Linear(hidden_dim, input_dim))
self.edge_list = None
self.batch_size = 0
def __init__(self,
latent_dim,
output_dim,
hidden_channels,
num_layers,
img_dims,
act_fn='elu'):
super(BroadcastDecoder, self).__init__()
self.latent_dim = latent_dim
self.width = img_dims[0]
self.height = img_dims[1]
self.num_layers = num_layers
mods = [
nn.Conv2d(latent_dim + 2, hidden_channels, 3),
util.get_act_fn(act_fn)
]
for _ in range(num_layers - 1):
mods.extend([
nn.Conv2d(hidden_channels, hidden_channels, 3),
util.get_act_fn(act_fn)
])
# 1x1 conv for output layer
mods.append(nn.Conv2d(hidden_channels, output_dim, 1))
self.seq = nn.Sequential(*mods)
def __init__(self,
latent_dim,
output_dim,
hidden_channels,
num_layers,
img_dim,
act_fn='elu'):
super(GenesisBroadcastDecoder, self).__init__()
broad_dim = img_dim + 2 * num_layers
mods = [
BroadcastLayer(broad_dim),
nn.Conv2d(latent_dim + 2, hidden_channels, 3),
util.get_act_fn(act_fn)
]
for _ in range(num_layers - 1):
mods.extend([
nn.Conv2d(hidden_channels, hidden_channels, 3),
util.get_act_fn(act_fn)
])
# 1x1 conv for output layer
mods.append(nn.Conv2d(hidden_channels, output_dim, 1))
self.seq = nn.Sequential(*mods)
def __init__(self,
input_dim,
hidden_dim,
num_objects,
act_fn='sigmoid',
act_fn_hid='relu'):
super(EncoderCNNSmall, self).__init__()
self.cnn1 = nn.Conv2d(input_dim, hidden_dim, (10, 10), stride=10)
self.cnn2 = nn.Conv2d(hidden_dim, num_objects, (1, 1), stride=1)
self.ln1 = nn.BatchNorm2d(hidden_dim)
self.act1 = util.get_act_fn(act_fn_hid)
self.act2 = util.get_act_fn(act_fn)
def __init__(self,
input_dim,
hidden_dim,
num_objects,
act_fn='sigmoid',
act_fn_hid='leaky_relu'):
super(EncoderCNNMedium, self).__init__()
self.cnn1 = nn.Conv2d(input_dim, hidden_dim, (9, 9), padding=4)
self.act1 = util.get_act_fn(act_fn_hid)
self.ln1 = nn.BatchNorm2d(hidden_dim)
self.cnn2 = nn.Conv2d(hidden_dim, num_objects, (5, 5), stride=5)
self.act2 = util.get_act_fn(act_fn)
def __init__(self,
input_dim,
hidden_dim,
num_objects,
output_size,
act_fn='relu'):
super(DecoderMLP, self).__init__()
self.fc1 = nn.Linear(input_dim + num_objects, hidden_dim)
self.fc2 = nn.Linear(hidden_dim, hidden_dim)
self.fc3 = nn.Linear(hidden_dim, np.prod(output_size))
self.input_dim = input_dim
self.num_objects = num_objects
self.output_size = output_size
self.act1 = util.get_act_fn(act_fn)
self.act2 = util.get_act_fn(act_fn)
def __init__(self,
input_dim,
output_dim,
hidden_dim,
num_objects,
act_fn='relu'):
super(EncoderMLP, self).__init__()
self.num_objects = num_objects
self.input_dim = input_dim
self.fc1 = nn.Linear(self.input_dim, hidden_dim)
self.fc2 = nn.Linear(hidden_dim, hidden_dim)
self.fc3 = nn.Linear(hidden_dim, output_dim)
self.ln = nn.LayerNorm(hidden_dim)
self.act1 = util.get_act_fn(act_fn)
self.act2 = util.get_act_fn(act_fn)
def __init__(self,
input_dim,
hidden_dim,
num_objects,
act_fn='sigmoid',
act_fn_hid='relu'):
super(EncoderCNNSeg, self).__init__()
self.cnn1 = nn.Conv2d(6, 32, (8, 8), stride=4)
self.act1 = util.get_act_fn(act_fn_hid)
# self.ln1 = nn.BatchNorm2d(hidden_dim)
self.cnn2 = nn.Conv2d(32, 64, (4, 4), stride=2)
self.act2 = util.get_act_fn(act_fn_hid)
# self.ln2 = nn.BatchNorm2d(hidden_dim)
self.cnn3 = nn.Conv2d(64, 64, (3, 3), stride=1)
self.act3 = util.get_act_fn(act_fn_hid)
# self.ln3 = nn.BatchNorm2d(hidden_dim)
self.fc1 = nn.Flatten(1, -1)
self.act4 = util.get_act_fn(act_fn_hid)
self.fc2 = nn.Linear(256, 1323)
self.act5 = util.get_act_fn(act_fn_hid)
self.act6 = util.get_act_fn(act_fn)
self.cnn4 = nn.Conv2d(3, 3, (1, 1))
self.upsample1 = nn.UpsamplingBilinear2d(32)
self.upsample2 = nn.UpsamplingBilinear2d(50)
def __init__(self,
input_dim,
hidden_dim,
num_objects,
act_fn='relu',
act_fn_hid='relu'):
super(EncoderCNNLarge, self).__init__()
self.cnn1 = nn.Conv2d(input_dim, hidden_dim, (3, 3), padding=1)
self.act1 = util.get_act_fn(act_fn_hid)
self.ln1 = nn.BatchNorm2d(hidden_dim)
self.cnn2 = nn.Conv2d(hidden_dim, hidden_dim, (3, 3), padding=1)
self.act2 = util.get_act_fn(act_fn_hid)
self.ln2 = nn.BatchNorm2d(hidden_dim)
self.cnn3 = nn.Conv2d(hidden_dim, hidden_dim, (3, 3), padding=1)
self.act3 = util.get_act_fn(act_fn_hid)
self.ln3 = nn.BatchNorm2d(hidden_dim)
self.cnn4 = nn.Conv2d(hidden_dim, num_objects, (1, 1), padding=0)
self.act4 = util.get_act_fn(act_fn)
def __init__(self,
input_dim,
hidden_dim,
action_dim,
num_objects,
ignore_action=False,
copy_action=False,
act_fn='relu'):
super(TransitionGNN, self).__init__()
self.input_dim = input_dim
self.hidden_dim = hidden_dim
self.num_objects = num_objects
self.ignore_action = ignore_action
self.copy_action = copy_action
if self.ignore_action:
self.action_dim = 0
else:
self.action_dim = action_dim
self.edge_mlp = nn.Sequential(nn.Linear(input_dim * 2, hidden_dim),
util.get_act_fn(act_fn),
nn.Linear(hidden_dim, hidden_dim),
nn.LayerNorm(hidden_dim),
util.get_act_fn(act_fn),
nn.Linear(hidden_dim, hidden_dim))
node_input_dim = hidden_dim + input_dim + self.action_dim
self.node_mlp = nn.Sequential(nn.Linear(node_input_dim, hidden_dim),
util.get_act_fn(act_fn),
nn.Linear(hidden_dim, hidden_dim),
nn.LayerNorm(hidden_dim),
util.get_act_fn(act_fn),
nn.Linear(hidden_dim, input_dim))
self.edge_list = None
self.batch_size = 0