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 = utils.get_act_fn(act_fn)
self.act2 = utils.get_act_fn(act_fn)
self.act3 = utils.get_act_fn(act_fn)
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 = utils.get_act_fn(act_fn_hid)
self.act2 = utils.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 = utils.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 = utils.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 = utils.get_act_fn(act_fn)
self.act2 = utils.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 = utils.get_act_fn(act_fn)
self.act2 = utils.get_act_fn(act_fn)
def __init__(self,
input_dim,
hidden_dim,
num_objects,
act_fn='sigmoid',
act_fn_hid='relu'):
super(EncoderCNNLarge, self).__init__()
self.cnn1 = nn.Conv2d(input_dim, hidden_dim, (3, 3), padding=1)
self.act1 = utils.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 = utils.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 = utils.get_act_fn(act_fn_hid)
self.ln3 = nn.BatchNorm2d(hidden_dim)
self.cnn4 = nn.Conv2d(hidden_dim, num_objects, (3, 3), padding=1)
self.act4 = utils.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),
utils.get_act_fn(act_fn),
nn.Linear(hidden_dim, hidden_dim),
nn.LayerNorm(hidden_dim),
utils.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),
utils.get_act_fn(act_fn),
nn.Linear(hidden_dim, hidden_dim),
nn.LayerNorm(hidden_dim),
utils.get_act_fn(act_fn),
nn.Linear(hidden_dim, input_dim))
self.edge_list = None
self.batch_size = 0