def __init__(self, input_shape, num_rotations, crop_size, preprocess): # pylint: disable=g-doc-args
"""Inits transport module with separate goal FCN.
Assumes the presence of a goal image, that cropping is done after the
query, that per-pixel loss is not used, and SE(2) grasping.
"""
self.num_rotations = num_rotations
self.crop_size = crop_size # crop size must be N*16 (e.g. 96)
self.preprocess = preprocess
self.lr = 1e-5
self.pad_size = int(self.crop_size / 2)
self.padding = np.zeros((3, 2), dtype=int)
self.padding[:2, :] = self.pad_size
input_shape = np.array(input_shape)
input_shape[0:2] += self.pad_size * 2
input_shape = tuple(input_shape)
# Output dimension (i.e., number of channels) of 3.
self.odim = output_dim = 3
# 3 fully convolutional ResNets. Third one is for the goal.
in0, out0 = ResNet43_8s(input_shape, output_dim, prefix='s0_')
in1, out1 = ResNet43_8s(input_shape, output_dim, prefix='s1_')
in2, out2 = ResNet43_8s(input_shape, output_dim, prefix='s2_')
self.model = tf.keras.Model(
inputs=[in0, in1, in2], outputs=[out0, out1, out2])
self.optim = tf.keras.optimizers.Adam(learning_rate=self.lr)
self.metric = tf.keras.metrics.Mean(name='transport_loss')
def __init__(self, input_shape, d_action, use_mdn):
del use_mdn
self.batch_size = 4
self.input_shape = input_shape
self.d_action = d_action
channel_depth_dim = 16
in0, out0 = ResNet43_8s(self.input_shape,
channel_depth_dim,
prefix="s0_",
cutoff_early=False,
include_batchnorm=True)
# out0 = tf.nn.avg_pool(out0, ksize=(1,4,4,1), strides=(1,4,4,1),
# padding="SAME", data_format="NHWC")
out0 = compute_spatial_soft_argmax(out0, self.batch_size, 320, 160,
16) # shape (B, C, 2)
out0 = tf.keras.layers.Flatten()(out0)
out0 = tf.keras.layers.Dense(128,
kernel_initializer="normal",
bias_initializer="normal",
activation="relu")(out0)
out0 = tf.keras.layers.Dense(128,
kernel_initializer="normal",
bias_initializer="normal",
activation="relu")(out0)
out0 = tf.keras.layers.Dense(self.d_action,
kernel_initializer="normal",
bias_initializer="normal")(out0)
self.model = tf.keras.Model(inputs=[in0], outputs=[out0])
def __init__(self, input_shape, preprocess):
self.preprocess = preprocess
RESNET = False
if RESNET:
output_dim = 6
in0, out0 = ResNet43_8s(input_shape, output_dim, prefix='s0_')
out0 = tf.nn.avg_pool(out0,
ksize=(1, 4, 4, 1),
strides=(1, 4, 4, 1),
padding="SAME",
data_format="NHWC")
out0 = tf.keras.layers.Flatten()(out0)
out0 = tf.keras.layers.Dense(128,
kernel_initializer="normal",
bias_initializer="normal",
activation='relu')(out0)
out0 = tf.keras.layers.Dense(3,
kernel_initializer="normal",
bias_initializer="normal")(out0)
self.model = tf.keras.Model(inputs=[in0], outputs=[out0])
else:
self.model = ConvMLP(d_action=3)
self.optim = tf.keras.optimizers.Adam()
self.metric = tf.keras.metrics.Mean(name='regression_loss')
self.loss_criterion = tf.keras.losses.MeanSquaredError()
def __init__(self,
input_shape,
descriptor_dim,
num_rotations,
preprocess,
lite=False):
self.preprocess = preprocess
self.num_rotations = num_rotations
self.descriptor_dim = descriptor_dim
max_dim = np.max(input_shape[:2])
self.padding = np.zeros((3, 2), dtype=int)
pad = (max_dim - np.array(input_shape[:2])) / 2
self.padding[:2] = pad.reshape(2, 1)
input_shape = np.array(input_shape)
input_shape += np.sum(self.padding, axis=1)
input_shape = tuple(input_shape)
# Initialize fully convolutional Residual Network with 43 layers and
# 8-stride (3 2x2 max pools and 3 2x bilinear upsampling)
if lite:
d_in, d_out = ResNet36_4s(input_shape, self.descriptor_dim)
else:
d_in, d_out = ResNet43_8s(input_shape, self.descriptor_dim)
self.model = tf.keras.models.Model(inputs=[d_in], outputs=[d_out])
self.optim = tf.keras.optimizers.Adam(learning_rate=1e-5)
self.metric = tf.keras.metrics.Mean(name='attention_loss')
def __init__(self, input_shape, num_rotations, crop_size, preprocess):
self.num_rotations = num_rotations
self.crop_size = crop_size # crop size must be N*16 (e.g. 96)
self.preprocess = preprocess
self.pad_size = int(self.crop_size / 2)
self.padding = np.zeros((3, 2), dtype=int)
self.padding[:2, :] = self.pad_size
input_shape = np.array(input_shape)
input_shape[0:2] += self.pad_size * 2
input_shape = tuple(input_shape)
self.odim = output_dim = 3
# 3 fully convolutional ResNets. Third one is for the goal.
in0, out0 = ResNet43_8s(input_shape, output_dim, prefix='s0_')
in1, out1 = ResNet43_8s(input_shape, output_dim, prefix='s1_')
in2, out2 = ResNet43_8s(input_shape, output_dim, prefix='s2_')
self.model = tf.keras.Model(inputs=[in0, in1, in2],
outputs=[out0, out1, out2])
self.optim = tf.keras.optimizers.Adam(learning_rate=1e-4)
self.metric = tf.keras.metrics.Mean(name='transport_loss')
def __init__(self,
input_shape,
num_rotations,
crop_size,
preprocess,
per_pixel_loss=False,
crop_bef_q=True,
use_goal_image=False):
self.num_rotations = num_rotations
self.crop_size = crop_size # crop size must be N*16 (e.g. 96)
self.preprocess = preprocess
self.per_pixel_loss = per_pixel_loss
self.crop_bef_q = crop_bef_q
self.use_goal_image = use_goal_image
self.pad_size = int(self.crop_size / 2)
self.padding = np.zeros((3, 2), dtype=int)
self.padding[:2, :] = self.pad_size
input_shape = np.array(input_shape)
input_shape[0:2] += self.pad_size * 2
input_shape = tuple(input_shape)
kernel_shape = (self.crop_size, self.crop_size, input_shape[2])
output_dim = 6 if self.per_pixel_loss else 3
# 2 fully convolutional ResNets [Daniel: I think 43 layers and stride 8]
in0, out0 = ResNet43_8s(input_shape, output_dim, prefix='s0_')
if self.crop_bef_q:
in1, out1 = ResNet43_8s(kernel_shape, 3, prefix='s1_')
else:
in1, out1 = ResNet43_8s(input_shape, output_dim, prefix='s1_')
self.model = tf.keras.Model(inputs=[in0, in1], outputs=[out0, out1])
self.optim = tf.keras.optimizers.Adam(learning_rate=1e-4)
self.metric = tf.keras.metrics.Mean(name='transport_loss')
def __init__(self,
input_shape,
num_rotations,
crop_size,
preprocess,
per_pixel_loss=False,
six_dof=False,
crop_bef_q=True,
use_goal_image=False):
"""Defines the transport module for determining placing.
Args:
input_shape: Shape of original, stacked color and heightmap images
before padding, usually (320,160,6).
num_rotations: Number of rotations considered, which determines the
number of challens of the output.
crop_size: Size of crop for Transporter Network.
preprocess: Method to subtract mean and divide by standard dev.
per_pixel_loss: True if training using a per-pixel loss. Default is
False (which performs better in CoRL submission experiments) and
means that a softmax is applied on all pixels, and the designated
pixel (from demonstration data) is used as the correct class,
passing gradients via every pixel, and avoiding the need for
negative data samples.
six_dof: True if using the 6 DoF extension of Transporters.
crop_bef_q: True if cropping the input before passing it through the
query FCN, meaning that the FCN takes input of shape (H,W) which
is `kernel_shape`, and has output of the same `kernel_shape` (H,W)
dims.
use_goal_image: True only if we're using goal-conditioning, and if
so, this is the naive version where we stack the current and goal
images and pass them through this Transport module.
"""
self.num_rotations = num_rotations
self.crop_size = crop_size # crop size must be N*16 (e.g. 96)
self.preprocess = preprocess
self.per_pixel_loss = per_pixel_loss
self.six_dof = six_dof
self.crop_bef_q = crop_bef_q
self.lr = 1e-5
self.pad_size = int(self.crop_size / 2)
self.padding = np.zeros((3, 2), dtype=int)
self.padding[:2, :] = self.pad_size
input_shape = np.array(input_shape)
input_shape[0:2] += self.pad_size * 2
input_shape = tuple(input_shape)
kernel_shape = (self.crop_size, self.crop_size, input_shape[2])
if not self.per_pixel_loss and not self.six_dof:
output_dim = 3
kernel_dim = 3
elif self.per_pixel_loss and not self.six_dof:
output_dim = 6
kernel_dim = 3
elif self.six_dof:
output_dim = 24
kernel_dim = 24
self.regress_loss = tf.keras.losses.Huber()
self.z_regressor = Regressor()
self.roll_regressor = Regressor()
self.pitch_regressor = Regressor()
else:
raise ValueError("I don't support this config!")
# 2 fully convolutional ResNets with 57 layers and 16-stride
if not self.six_dof:
in0, out0 = ResNet43_8s(input_shape, output_dim, prefix="s0_")
if self.crop_bef_q:
# Passing in kernels: (64,64,6) --> (64,64,3)
in1, out1 = ResNet43_8s(kernel_shape, kernel_dim, prefix="s1_")
else:
# Passing in original images: (384,224,6) --> (394,224,3)
in1, out1 = ResNet43_8s(input_shape, output_dim, prefix="s1_")
else:
in0, out0 = ResNet43_8s(input_shape, output_dim, prefix="s0_")
# early cutoff just so it all fits on GPU.
in1, out1 = ResNet43_8s(kernel_shape,
kernel_dim,
prefix="s1_",
cutoff_early=True)
self.model = tf.keras.Model(inputs=[in0, in1], outputs=[out0, out1])
self.optim = tf.keras.optimizers.Adam(learning_rate=self.lr)
self.metric = tf.keras.metrics.Mean(name="transport_loss")
if self.six_dof:
self.z_metric = tf.keras.metrics.Mean(name="z_loss")
self.roll_metric = tf.keras.metrics.Mean(name="roll_loss")
self.pitch_metric = tf.keras.metrics.Mean(name="pitch_loss")
# For visualization
self.feature_visualize = False
if self.feature_visualize:
self.fig, self.ax = plt.subplots(5, 1)
self.write_visualize = False
self.plot_interval = 20
self.iters = 0