unimodals.robotics package

Submodules

unimodals.robotics.decoders module

Implements decoders for robotics tasks.

class unimodals.robotics.decoders.ContactDecoder(z_dim, deterministic, head=1)

Bases: Module

Decodes everything, given some input.

__init__(z_dim, deterministic, head=1)

Initialize ContactDecoder Module.

Parameters:

  • z_dim (float) – Z dimension size

  • deterministic (float) – Whether input parameters are deterministic or sampled from some distribution given prior mu and var.

  • head (int) – Output dimension of head.

forward(input)

Apply ContactDecoder Module to Layer Input.

Parameters:

input (torch.Tensor) – Layer Input

Returns:

Decoded Output

Return type:

torch.Tensor

class unimodals.robotics.decoders.EeDeltaDecoder(z_dim, action_dim, initailize_weights=True)

Bases: Module

Implements an EE Delta Decoder.

__init__(z_dim, action_dim, initailize_weights=True)

Initialize EeDeltaDecoder Module.

Parameters:

  • z_dim (float) – Z dimension size

  • alpha (float) – Alpha to multiply proprio input by.

  • initialize_weights (bool, optional) – Whether to initialize weights or not. Defaults to True.

forward(mm_act_feat)

Apply EeDeltaDecoder Module to EE Delta.

Parameters:

mm_act_feat (torch.Tensor) – EE Delta

Returns:

Decoded Output

Return type:

torch.Tensor

class unimodals.robotics.decoders.OpticalFlowDecoder(z_dim, initailize_weights=True)

Bases: Module

Implements optical flow and optical flow mask decoder.

__init__(z_dim, initailize_weights=True)

Initialize OpticalFlowDecoder Module.

Parameters:

  • z_dim (float) – Z dimension size

  • alpha (float) – Alpha to multiply proprio input by.

  • initialize_weights (bool, optional) – Whether to initialize weights or not. Defaults to True.

forward(tiled_feat, img_out_convs)

Predicts the optical flow and optical flow mask.

Parameters:

  • tiled_feat – action conditioned z (output of fusion + action network)

  • img_out_convs – outputs of the image encoders (skip connections)

unimodals.robotics.encoders module

Implements encoders for robotics tasks.

class unimodals.robotics.encoders.ActionEncoder(action_dim)

Bases: Module

Implements an action-encoder module.

__init__(action_dim)

Instantiate ActionEncoder module.

Parameters:

action_dim (int) – Dimension of action.

forward(action)

Apply action encoder to action input.

Parameters:

action (torch.Tensor optional) – Action input

Returns:

Encoded output

Return type:

torch.Tensor

class unimodals.robotics.encoders.DepthEncoder(z_dim, alpha, initialize_weights=True)

Bases: Module

Implements a simplified depth-encoder module.

Sourced from Making Sense of Vision and Touch.

__init__(z_dim, alpha, initialize_weights=True)

Initialize DepthEncoder Module.

Parameters:

  • z_dim (float) – Z dimension size

  • alpha (float) – Alpha to multiply input by.

  • initialize_weights (bool, optional) – Whether to initialize weights or not. Defaults to True.

forward(depth_in)

Apply encoder to depth input.

Parameters:

depth_in (torch.Tensor) – Depth input

Returns:

Output of encoder, Output of encoder after each convolution.

Return type:

tuple(torch.Tensor, torch.Tensor)

class unimodals.robotics.encoders.ForceEncoder(z_dim, alpha, initialize_weights=True)

Bases: Module

Implements force encoder module.

Sourced from selfsupervised code.

__init__(z_dim, alpha, initialize_weights=True)

Initialize ForceEncoder Module.

Parameters:

  • z_dim (float) – Z dimension size

  • alpha (float) – Alpha to multiply proprio input by.

  • initialize_weights (bool, optional) – Whether to initialize weights or not. Defaults to True.

forward(force)

Apply ForceEncoder to Force Input.

Parameters:

force (torch.Tensor) – Force Input

Returns:

Encoded Output

Return type:

torch.Tensor

class unimodals.robotics.encoders.ImageEncoder(z_dim, alpha, initialize_weights=True)

Bases: Module

Implements image encoder module.

Sourced from Making Sense of Vision and Touch.

__init__(z_dim, alpha, initialize_weights=True)

Initialize ImageEncoder Module.

Parameters:

  • z_dim (float) – Z dimension size

  • alpha (float) – Alpha to multiply input by.

  • initialize_weights (bool, optional) – Whether to initialize weights or not. Defaults to True.

forward(vis_in)

Apply encoder to image input.

Parameters:

vis_in (torch.Tensor) – Image input

Returns:

Output of encoder, Output of encoder after each convolution.

Return type:

tuple(torch.Tensor, torch.Tensor)

class unimodals.robotics.encoders.ProprioEncoder(z_dim, alpha, initialize_weights=True)

Bases: Module

Implements image encoder module.

Sourced from selfsupervised code.

__init__(z_dim, alpha, initialize_weights=True)

Initialize ProprioEncoder Module.

Parameters:

  • z_dim (float) – Z dimension size

  • alpha (float) – Alpha to multiply proprio input by.

  • initialize_weights (bool, optional) – Whether to initialize weights or not. Defaults to True.

forward(proprio)

Apply ProprioEncoder to Proprio Input.

Parameters:

proprio (torch.Tensor) – Proprio Input

Returns:

Encoded Output

Return type:

torch.Tensor

unimodals.robotics.layers module

Neural network layers.

class unimodals.robotics.layers.CausalConv1D(in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True)

Bases: Conv1d

Implements a causal 1D convolution.

__init__(in_channels, out_channels, kernel_size, stride=1, dilation=1, bias=True)

Initialize CausalConv1D Object.

Parameters:

  • in_channels (int) – Input Dimension

  • out_channels (int) – Output Dimension

  • kernel_size (int) – Kernel Size

  • stride (int, optional) – Stride Amount. Defaults to 1.

  • dilation (int, optional) – Dilation Amount. Defaults to 1.

  • bias (bool, optional) – Whether to add a bias after convolving. Defaults to True.

forward(x)

Apply CasualConv1D layer to layer input.

Parameters:

x (torch.Tensor) – Layer Input

Returns:

Layer Output

Return type:

torch.Tensor

class unimodals.robotics.layers.Flatten

Bases: Module

Implements .reshape(*,-1) for nn.Sequential usage.

__init__()

Initialize Flatten Object.

forward(x)

Apply Flatten to Input.

Parameters:

x (torch.Tensor) – Layer Input

Returns:

Layer Output

Return type:

torch.Tensor

class unimodals.robotics.layers.ResidualBlock(channels)

Bases: Module

Implements a simple residual block.

__init__(channels)

Initialize ResidualBlock object.

Parameters:

channels (int) – Number of input and hidden channels in block.

forward(x)

Apply ResidualBlock to Layer Input.

Parameters:

x (torch.Tensor) – Layer Input

Returns:

Layer Output

Return type:

torch.Tensor

class unimodals.robotics.layers.View(size)

Bases: Module

Extends .view() for nn.Sequential usage.

__init__(size)

Initialize View Object.

Parameters:

size (int) – Output Size

forward(tensor)

Apply View to Layer Input.

Parameters:

tensor (torch.Tensor) – Layer Input

Returns:

Layer Output

Return type:

torch.Tensor

unimodals.robotics.layers.conv2d(in_channels, out_channels, kernel_size=3, stride=1, dilation=1, bias=True)

Create nn.Module with same convolution with LeakyReLU, i.e. output shape equals input shape.

Parameters:

  • in_planes (int) – The number of input feature maps.

  • out_planes (int) – The number of output feature maps.

  • kernel_size (int) – The filter size.

  • dilation (int) – The filter dilation factor.

  • stride (int) – The filter stride.

unimodals.robotics.layers.crop_like(input, target)

Crop Tensor based on Another’s Size.

Parameters:

  • input (torch.Tensor) – Input Tensor

  • target (torch.Tensor) – Tensor to Compare and Crop To.

Returns:

Cropped Tensor

Return type:

torch.Tensor

unimodals.robotics.layers.deconv(in_planes, out_planes)

Create Deconvolution Layer.

Parameters:

  • in_planes (int) – Number of Input Channels.

  • out_planes (int) – Number of Output Channels.

Returns:

Deconvolution Object with Given Input/Output Channels.

Return type:

nn.Module

unimodals.robotics.layers.predict_flow(in_planes)

Create Optical Flow Prediction Layer.

Parameters:

in_planes (int) – Number of Input Channels.

Returns:

Convolution Object with Given Input/Output Channels.

Return type:

nn.Module

unimodals.robotics.models_utils module

Utility functions for robotics unimodals.

unimodals.robotics.models_utils.filter_depth(depth_image)

Get filter depth given a depth image.

Parameters:

depth_image (torch.Tensor) – Depth image.

Returns:

Output

Return type:

torch.Tensor

unimodals.robotics.models_utils.init_weights(modules)

Weight initialization from original SensorFusion Code.

unimodals.robotics.models_utils.rescaleImage(image, output_size=128, scale=0.00392156862745098)

Rescale the image in a sample to a given size.

Parameters:

output_size (tuple or int) – Desired output size. If tuple, output is matched to output_size. If int, smaller of image edges is matched to output_size keeping aspect ratio the same.

Module contents