From 29646ad43c85e5329931c487c5ca6e96f5461ef5 Mon Sep 17 00:00:00 2001 From: Luis Pineda Date: Tue, 29 Mar 2022 14:07:40 -0400 Subject: [PATCH] Added batching support to tactile pushing example. (#132) * Added batching support to tactile pushing example. * Changed outer loss so that it uses SE2.local(). * Added more options for the inner optimizer. * Added options to control backward mode and step size.a * Added results logging. * Added cfg values for lr and early stopping. --- examples/configs/tactile_pose_estimation.yaml | 23 ++- examples/tactile_pose_estimation.py | 150 ++++++++++++++---- .../examples/tactile_pose_estimation/misc.py | 96 ++++++++--- .../tactile_pose_estimation/models.py | 10 +- .../tactile_pose_estimation/pose_estimator.py | 51 ++++-- 5 files changed, 249 insertions(+), 81 deletions(-) diff --git a/examples/configs/tactile_pose_estimation.yaml b/examples/configs/tactile_pose_estimation.yaml index 4fb7e9722..74d0d90cd 100644 --- a/examples/configs/tactile_pose_estimation.yaml +++ b/examples/configs/tactile_pose_estimation.yaml @@ -1,10 +1,20 @@ seed: 0 +save_all: true dataset_name: "rectangle-pushing-corners-keypoints" sdf_name: "rect" -episode: 0 +episode_length: 100 max_steps: 100 +max_episodes: 1 + +inner_optim: + max_iters: 3 + optimizer: GaussNewton + reg_w: 1e-4 + backward_mode: IMPLICIT + step_size: 1.0 + keep_step_size: true # 0: disc, 1: rect-edges, 2: rect-corners, 3: ellip class_label: 2 @@ -24,13 +34,12 @@ tactile_cost: train: # options: "weights_only" or "weights_and_measurement_nn" - mode: "weights_only" - - batch_size: 1 - num_batches: 1 + mode: "weights_and_measurement_nn" - num_epochs: 100 - eps_tracking_loss: 1e-5 + batch_size: 4 + num_epochs: 50 + lr: 1e-3 # 5.0 for weights_only + eps_tracking_loss: 1e-10 options: vis_traj: True diff --git a/examples/tactile_pose_estimation.py b/examples/tactile_pose_estimation.py index 4e464579c..99f88e5c1 100644 --- a/examples/tactile_pose_estimation.py +++ b/examples/tactile_pose_estimation.py @@ -2,20 +2,26 @@ # # This source code is licensed under the MIT license found in the # LICENSE file in the root directory of this source tree. - +import logging +import os import pathlib import random +import time +from typing import Dict import hydra import matplotlib.pyplot as plt import numpy as np import torch import torch.nn as nn -import torch.nn.functional as F import torch.optim as optim +import theseus as th import theseus.utils.examples as theg +# Logger +logger = logging.getLogger(__name__) + # To run this example, you will need a tactile pushing dataset available at # https://dl.fbaipublicfiles.com/theseus/tactile_pushing_data.tar.gz # @@ -61,10 +67,41 @@ # 2021 (https://arxiv.org/abs/1705.10664) +def pack_batch_results( + theseus_outputs: Dict[str, torch.Tensor], + qsp_state_dict: torch.Tensor, + mfb_state_dict: torch.Tensor, + meas_state_dict: torch.Tensor, + info: th.optimizer.OptimizerInfo, + loss_value: float, + total_time: float, +) -> Dict: + def _clone(t_): + return t_.detach().cpu().clone() + + return { + "theseus_outputs": dict((s, _clone(t)) for s, t in theseus_outputs.items()), + "qsp_state_dict": qsp_state_dict, + "mfb_state_dict": mfb_state_dict, + "meas_state_dict": meas_state_dict, + "err_history": info.err_history, # type: ignore + "loss": loss_value, + "total_time": total_time, + } + + def run_learning_loop(cfg): + root_path = pathlib.Path(os.getcwd()) dataset_path = EXP_PATH / "datasets" / f"{cfg.dataset_name}.json" sdf_path = EXP_PATH / "sdfs" / f"{cfg.sdf_name}.json" - dataset = theg.TactilePushingDataset(dataset_path, sdf_path, cfg.episode, device) + dataset = theg.TactilePushingDataset( + dataset_path, + sdf_path, + cfg.episode_length, + cfg.train.batch_size, + cfg.max_episodes, + device, + ) # -------------------------------------------------------------------- # # Create pose estimator (which wraps a TheseusLayer) @@ -77,6 +114,10 @@ def run_learning_loop(cfg): step_window_moving_frame=cfg.tactile_cost.step_win_mf, rectangle_shape=(cfg.shape.rect_len_x, cfg.shape.rect_len_y), device=device, + optimizer_cls=getattr(th, cfg.inner_optim.optimizer), + max_iterations=cfg.inner_optim.max_iters, + step_size=cfg.inner_optim.step_size, + regularization_w=cfg.inner_optim.reg_w, ) time_steps = pose_estimator.time_steps @@ -94,29 +135,33 @@ def run_learning_loop(cfg): qsp_model, mf_between_model, learnable_params, - hyperparameters, ) = theg.create_tactile_models( cfg.train.mode, device, measurements_model_path=measurements_model_path ) - eps_tracking_loss = hyperparameters.get( - "eps_tracking_loss", cfg.train.eps_tracking_loss - ) - outer_optim = optim.Adam(learnable_params, lr=hyperparameters["learning_rate"]) + eps_tracking_loss = cfg.train.eps_tracking_loss + outer_optim = optim.Adam(learnable_params, lr=cfg.train.lr) # -------------------------------------------------------------------- # # Main learning loop # -------------------------------------------------------------------- # # Use theseus_layer in an outer learning loop to learn different cost # function parameters: - measurements = dataset.get_measurements( - cfg.train.batch_size, cfg.train.num_batches, time_steps - ) - obj_poses_gt = dataset.obj_poses[0:time_steps, :].clone().requires_grad_(True) - eff_poses_gt = dataset.eff_poses[0:time_steps, :].clone().requires_grad_(True) - theseus_inputs = {} - for _ in range(cfg.train.num_epochs): + measurements = dataset.get_measurements(time_steps) + results = {} + for epoch in range(cfg.train.num_epochs): + results[epoch] = {} + logger.info(" ********************* EPOCH {epoch} *********************") losses = [] + image_idx = 0 for batch_idx, batch in enumerate(measurements): + pose_and_motion_batch = dataset.get_start_pose_and_motion_for_batch( + batch_idx, time_steps + ) # x_y_theta format + pose_estimator.update_start_pose_and_motion_from_batch( + pose_and_motion_batch + ) + theseus_inputs = {} + # Updates the above with measurement factor data theg.update_tactile_pushing_inputs( dataset=dataset, batch=batch, @@ -129,28 +174,49 @@ def run_learning_loop(cfg): theseus_inputs=theseus_inputs, ) - theseus_inputs, _ = pose_estimator.forward( - theseus_inputs, optimizer_kwargs={"verbose": True} + start_time = time.time_ns() + theseus_outputs, info = pose_estimator.forward( + theseus_inputs, + optimizer_kwargs={ + "verbose": True, + "track_err_history": True, + "backward_mode": getattr( + th.BackwardMode, cfg.inner_optim.backward_mode + ), + "__keep_final_step_size__": cfg.inner_optim.keep_step_size, + }, ) + end_time = time.time_ns() obj_poses_opt, eff_poses_opt = theg.get_tactile_poses_from_values( - values=theseus_inputs, time_steps=time_steps + values=theseus_outputs, time_steps=time_steps + ) + obj_poses_gt, eff_poses_gt = dataset.get_gt_data_for_batch( + batch_idx, time_steps ) - loss = F.mse_loss(obj_poses_opt[batch_idx, :], obj_poses_gt) + se2_opt = th.SE2(x_y_theta=obj_poses_opt.view(-1, 3)) + se2_gt = th.SE2(x_y_theta=obj_poses_gt.view(-1, 3)) + loss = se2_opt.local(se2_gt).norm() loss.backward() nn.utils.clip_grad_norm_(qsp_model.parameters(), 100, norm_type=2) nn.utils.clip_grad_norm_(mf_between_model.parameters(), 100, norm_type=2) + nn.utils.clip_grad_norm_(measurements_model.parameters(), 100, norm_type=2) with torch.no_grad(): for name, param in qsp_model.named_parameters(): - print(name, param.data) + logger.info(f"{name} {param.data}") for name, param in mf_between_model.named_parameters(): - print(name, param.data) + logger.info(f"{name} {param.data}") - print(" grad qsp", qsp_model.param.grad.norm().item()) - print(" grad mfb", mf_between_model.param.grad.norm().item()) + def _print_grad(msg_, param_): + logger.info(f"{msg_} {param_.grad.norm().item()}") + + _print_grad(" grad qsp", qsp_model.param) + _print_grad(" grad mfb", mf_between_model.param) + _print_grad(" grad nn_weight", measurements_model.fc1.weight) + _print_grad(" grad nn_bias", measurements_model.fc1.bias) outer_optim.step() @@ -162,17 +228,35 @@ def run_learning_loop(cfg): losses.append(loss.item()) - if cfg.options.vis_traj: - theg.visualize_tactile_push2d( - obj_poses=obj_poses_opt[0, :], - eff_poses=eff_poses_opt[0, :], - obj_poses_gt=obj_poses_gt, - eff_poses_gt=eff_poses_gt, - rect_len_x=cfg.shape.rect_len_x, - rect_len_y=cfg.shape.rect_len_y, - ) + if cfg.save_all: + results[epoch][batch_idx] = pack_batch_results( + theseus_outputs, + qsp_model.state_dict(), + mf_between_model.state_dict(), + measurements_model.state_dict(), + info, + loss.item(), + end_time - start_time, + ) + torch.save(results, root_path / "results.pt") + + if cfg.options.vis_traj: + for i in range(len(obj_poses_gt)): + save_dir = root_path / f"img_{image_idx}" + save_dir.mkdir(parents=True, exist_ok=True) + save_fname = save_dir / f"epoch{epoch}.png" + theg.visualize_tactile_push2d( + obj_poses=obj_poses_opt[i], + eff_poses=eff_poses_opt[i], + obj_poses_gt=obj_poses_gt[i], + eff_poses_gt=eff_poses_gt[i], + rect_len_x=cfg.shape.rect_len_x, + rect_len_y=cfg.shape.rect_len_y, + save_fname=save_fname, + ) + image_idx += 1 - print(f"AVG. LOSS: {np.mean(losses)}") + logger.info(f"AVG. LOSS: {np.mean(losses)}") if np.mean(losses) < eps_tracking_loss: break diff --git a/theseus/utils/examples/tactile_pose_estimation/misc.py b/theseus/utils/examples/tactile_pose_estimation/misc.py index 11aa0ff8c..e81c072e9 100644 --- a/theseus/utils/examples/tactile_pose_estimation/misc.py +++ b/theseus/utils/examples/tactile_pose_estimation/misc.py @@ -1,4 +1,4 @@ -from typing import Dict, List, Tuple +from typing import Dict, List, Optional, Tuple import matplotlib.patches as mpatches import matplotlib.pyplot as plt @@ -15,11 +15,14 @@ def __init__( self, data_fname: str, sdf_fname: str, - episode: int, + episode_length: int, + batch_size: int, + max_episodes: int, device: torch.device, ): + batch_size = min(batch_size, max_episodes) data = TactilePushingDataset._load_dataset_from_file( - data_fname, episode, device + data_fname, episode_length, max_episodes, device ) ( self.sdf_data_tensor, @@ -32,13 +35,21 @@ def __init__( self.obj_poses = data["obj_poses"] self.contact_episode = data["contact_episode"] self.contact_flag = data["contact_flag"] + self.dataset_size: int = -1 for key, val in data.items(): setattr(self, key, val) + if self.dataset_size == -1: + self.dataset_size = val.shape[0] + else: + assert self.dataset_size == val.shape[0] + self.batch_size = batch_size @staticmethod def _load_dataset_from_file( - filename: str, episode: int, device: torch.device + filename: str, episode_length: int, max_episodes: int, device: torch.device ) -> Dict[str, torch.Tensor]: + + # Load all episode data with open(filename) as f: import json @@ -61,11 +72,30 @@ def _load_dataset_from_file( data_from_file["contact_flag"], device=device ) - data = {} - ds_idxs = torch.nonzero(dataset_all["contact_episode"] == episode).squeeze() - for key, val in dataset_all.items(): - data[key] = val[ds_idxs] - return data + # Read all episodes and filter those with length less than desired + episode_indices = [ + idx.item() for idx in dataset_all["contact_episode"].unique() + ] + data: Dict[str, List[torch.Tensor]] = dict( + [(k, []) for k in dataset_all.keys()] + ) + + for i, episode in enumerate(episode_indices): + if i == max_episodes: + break + ds_idxs = torch.nonzero(dataset_all["contact_episode"] == episode).squeeze() + if len(ds_idxs) < episode_length: + continue + ds_idxs = ds_idxs[:episode_length] + for key, val in dataset_all.items(): + data[key].append(val[ds_idxs]) + + # Stack all episode data into single tensors + data_tensors = {} + for key in data: + data_tensors[key] = torch.stack(data[key]) + print(f"Read {len(data_tensors[key])} episodes of length {episode_length}.") + return data_tensors @staticmethod def _load_tactile_sdf_from_file( @@ -97,19 +127,41 @@ def _load_tactile_sdf_from_file( return sdf_data_tensor, cell_size, origin def get_measurements( - self, batch_size: int, num_batches: int, time_steps: int + self, time_steps: int ) -> List[Tuple[torch.Tensor, torch.Tensor, torch.Tensor]]: - def _process(tensor: torch.Tensor) -> torch.Tensor: - return tensor.unsqueeze(0).repeat(batch_size, 1, 1) - - return [ - ( - _process(self.img_feats[0:time_steps]), - _process(self.eff_poses[0:time_steps]), - _process(self.obj_poses[0:time_steps]), + num_batches = (self.dataset_size - 1) // self.batch_size + 1 + batches = [] + for batch_idx in range(num_batches): + start = batch_idx * self.batch_size + end = min(start + self.batch_size, self.dataset_size) + batches.append( + ( + self.img_feats[start:end, 0:time_steps], + self.eff_poses[start:end, 0:time_steps], + self.obj_poses[start:end, 0:time_steps], + ) ) - for _ in range(num_batches) - ] + return batches + + def get_start_pose_and_motion_for_batch( + self, batch_idx: int, time_steps: int + ) -> Dict[str, torch.Tensor]: + pose_and_motion_batch = {} + start = batch_idx * self.batch_size + end = min(start + self.batch_size, self.dataset_size) + pose_and_motion_batch["obj_start_pose"] = self.obj_poses[start:end, 0] + for i in range(time_steps): + pose_and_motion_batch[f"motion_capture_{i}"] = self.eff_poses[start:end, i] + return pose_and_motion_batch + + def get_gt_data_for_batch( + self, batch_idx: int, time_steps: int + ) -> Tuple[torch.Tensor, torch.Tensor]: + start = batch_idx * self.batch_size + end = min(start + self.batch_size, self.dataset_size) + obj_poses_gt = self.obj_poses[start:end, :time_steps, :].clone() + eff_poses_gt = self.eff_poses[start:end, :time_steps, :].clone() + return obj_poses_gt, eff_poses_gt # ----------------------------------------------------------------------------------- # @@ -192,6 +244,7 @@ def visualize_tactile_push2d( eff_poses_gt: torch.Tensor, rect_len_x: float, rect_len_y: float, + save_fname: Optional[str] = None, ): plt.cla() @@ -211,3 +264,6 @@ def visualize_tactile_push2d( plt.show() plt.pause(1e-9) + + if save_fname is not None: + plt.savefig(save_fname) diff --git a/theseus/utils/examples/tactile_pose_estimation/models.py b/theseus/utils/examples/tactile_pose_estimation/models.py index d17a5ddf0..2037221cb 100644 --- a/theseus/utils/examples/tactile_pose_estimation/models.py +++ b/theseus/utils/examples/tactile_pose_estimation/models.py @@ -65,8 +65,7 @@ def create_tactile_models( model_type: str, device: torch.device, measurements_model_path: Optional[pathlib.Path] = None, -) -> Tuple[nn.Module, nn.Module, nn.Module, List[nn.Parameter], Dict[str, float]]: - hyperparams = {} +) -> Tuple[nn.Module, nn.Module, nn.Module, List[nn.Parameter]]: if model_type == "weights_only": qsp_model = TactileWeightModel( device, wt_init=torch.tensor([[50.0, 50.0, 50.0]]) @@ -76,7 +75,6 @@ def create_tactile_models( ) measurements_model = None - hyperparams["learning_rate"] = 5.0 learnable_params = list(qsp_model.parameters()) + list( mf_between_model.parameters() ) @@ -93,8 +91,6 @@ def create_tactile_models( ) measurements_model.to(device) - hyperparams["learning_rate"] = 1.0e-3 - hyperparams["eps_tracking_loss"] = 5.0e-4 # early stopping learnable_params = ( list(measurements_model.parameters()) + list(qsp_model.parameters()) @@ -108,7 +104,6 @@ def create_tactile_models( qsp_model, mf_between_model, learnable_params, - hyperparams, ) @@ -275,8 +270,9 @@ def update_tactile_pushing_inputs( time_steps: int, theseus_inputs: Dict[str, torch.Tensor], ): + batch_size = batch[0].shape[0] theseus_inputs["sdf_data"] = ( - (dataset.sdf_data_tensor.data).repeat(cfg.train.batch_size, 1, 1).to(device) + (dataset.sdf_data_tensor.data).repeat(batch_size, 1, 1).to(device) ) theseus_inputs.update( diff --git a/theseus/utils/examples/tactile_pose_estimation/pose_estimator.py b/theseus/utils/examples/tactile_pose_estimation/pose_estimator.py index 8e320bc5c..829e692fc 100644 --- a/theseus/utils/examples/tactile_pose_estimation/pose_estimator.py +++ b/theseus/utils/examples/tactile_pose_estimation/pose_estimator.py @@ -1,9 +1,10 @@ -from typing import Tuple +from typing import Dict, List, Optional, Tuple, Type import numpy as np import torch import theseus as th +from theseus.optimizer.nonlinear.levenberg_marquardt import LevenbergMarquardt from .misc import TactilePushingDataset @@ -18,9 +19,14 @@ def __init__( step_window_moving_frame: int, rectangle_shape: Tuple[float, float], device: torch.device, + optimizer_cls: Optional[Type[th.NonlinearLeastSquares]] = LevenbergMarquardt, + max_iterations: int = 3, + step_size: float = 1.0, + regularization_w: float = 0.0, ): self.dataset = dataset - self.time_steps = np.minimum(max_steps, len(self.dataset.obj_poses)) + # obj_poses is shape (batch_size, episode_length, 3) + self.time_steps = np.minimum(max_steps, self.dataset.obj_poses.shape[1]) # -------------------------------------------------------------------- # # Creating optimization variables @@ -40,18 +46,13 @@ def __init__( # - nn_measurements: tactile measurement prediction from image features # - sdf_data, sdf_cell_size, sdf_origin: signed distance field data, # cell_size and origin - obj_start_pose = th.SE2( - x_y_theta=self.dataset.obj_poses[0].unsqueeze(0), name="obj_start_pose" - ) + obj_start_pose = th.SE2(name="obj_start_pose") + self.obj_start_pose = obj_start_pose - motion_captures = [] + motion_captures: List[th.SE2] = [] for i in range(self.time_steps): - motion_captures.append( - th.SE2( - x_y_theta=self.dataset.eff_poses[i].unsqueeze(0), - name=f"motion_capture_{i}", - ) - ) + motion_captures.append(th.SE2(name=f"motion_capture_{i}")) + self.motion_captures = motion_captures nn_measurements = [] for i in range(min_window_moving_frame, self.time_steps): @@ -175,16 +176,38 @@ def __init__( ) ) + if regularization_w > 0.0: + reg_w = th.ScaleCostWeight(np.sqrt(regularization_w)) + reg_w.to(dtype=torch.double) + identity_se2 = th.SE2(name="identity") + for pose_list in [obj_poses, eff_poses]: + for pose in pose_list: + objective.add( + th.eb.VariableDifference( + pose, reg_w, identity_se2, name=f"reg_{pose.name}" + ) + ) + # -------------------------------------------------------------------- # # Creating TheseusLayer # -------------------------------------------------------------------- # # Wrap the objective and inner-loop optimizer into a `TheseusLayer`. # Inner-loop optimizer here is the Levenberg-Marquardt nonlinear optimizer # coupled with a dense linear solver based on Cholesky decomposition. - nl_optimizer = th.LevenbergMarquardt( - objective, th.CholeskyDenseSolver, max_iterations=3 + nl_optimizer = optimizer_cls( + objective, + th.CholeskyDenseSolver, + max_iterations=max_iterations, + step_size=step_size, ) self.theseus_layer = th.TheseusLayer(nl_optimizer) self.theseus_layer.to(device=device, dtype=torch.double) self.forward = self.theseus_layer.forward + + # This method updates the start pose and motion catpure variables with the + # xytheta data coming from the batch + def update_start_pose_and_motion_from_batch(self, batch: Dict[str, torch.Tensor]): + self.obj_start_pose.update_from_x_y_theta(batch[self.obj_start_pose.name]) + for motion_capture_var in self.motion_captures: + motion_capture_var.update_from_x_y_theta(batch[motion_capture_var.name])