import numpy as np from cliport.utils import utils from cliport.agents.transporter import TransporterAgent from cliport.models.streams.one_stream_attention_lang_fusion import OneStreamAttentionLangFusion from cliport.models.streams.one_stream_transport_lang_fusion import OneStreamTransportLangFusion from cliport.models.streams.two_stream_attention_lang_fusion import TwoStreamAttentionLangFusion from cliport.models.streams.two_stream_transport_lang_fusion import TwoStreamTransportLangFusion, TwoStreamTransportLangFusionLatReduce, TwoStreamTransportLangFusionLatPretrained18 from cliport.models.streams.two_stream_attention_lang_fusion import TwoStreamAttentionLangFusionLat, TwoStreamAttentionLangFusionLatReduce from cliport.models.streams.two_stream_transport_lang_fusion import TwoStreamTransportLangFusionLatReduceOneStream from cliport.models.streams.two_stream_transport_lang_fusion import TwoStreamTransportLangFusionLat import torch import time class TwoStreamClipLingUNetTransporterAgent(TransporterAgent): def __init__(self, name, cfg, train_ds, test_ds): super().__init__(name, cfg, train_ds, test_ds) def _build_model(self): stream_one_fcn = 'plain_resnet' stream_two_fcn = 'clip_lingunet' self.attention = TwoStreamAttentionLangFusion( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=1, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) self.transport = TwoStreamTransportLangFusion( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=self.n_rotations, crop_size=self.crop_size, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) def attn_forward(self, inp, softmax=True): inp_img = inp['inp_img'] if type(inp_img) is not torch.Tensor: inp_img = torch.from_numpy(inp_img).to('cuda').float().contiguous() lang_goal = inp['lang_goal'] out = self.attention.forward(inp_img.float(), lang_goal, softmax=softmax) return out def attn_training_step(self, frame, backprop=True, compute_err=False): inp_img = frame['img'] if type(inp_img) is not torch.Tensor: inp_img = torch.from_numpy(inp_img).to('cuda').float() p0, p0_theta = frame['p0'], frame['p0_theta'] lang_goal = frame['lang_goal'] inp = {'inp_img': inp_img, 'lang_goal': lang_goal} out = self.attn_forward(inp, softmax=False) return self.attn_criterion(backprop, compute_err, inp, out, p0, p0_theta) def trans_forward(self, inp, softmax=True): inp_img = inp['inp_img'] if type(inp_img) is not torch.Tensor: inp_img = torch.from_numpy(inp_img).to('cuda').float() p0 = inp['p0'] lang_goal = inp['lang_goal'] out = self.transport.forward(inp_img.float(), p0, lang_goal, softmax=softmax) return out def transport_training_step(self, frame, backprop=True, compute_err=False): inp_img = frame['img'] p0 = frame['p0'] p1, p1_theta = frame['p1'], frame['p1_theta'] lang_goal = frame['lang_goal'] inp = {'inp_img': inp_img, 'p0': p0, 'lang_goal': lang_goal} out = self.trans_forward(inp, softmax=False) err, loss = self.transport_criterion(backprop, compute_err, inp, out, p0, p1, p1_theta) return loss, err def act(self, obs, info, goal=None): # pylint: disable=unused-argument """Run inference and return best action given visual observations.""" # Get heightmap from RGB-D images. img = self.test_ds.get_image(obs) lang_goal = info['lang_goal'] # Attention model forward pass. pick_inp = {'inp_img': img, 'lang_goal': lang_goal} pick_conf = self.attn_forward(pick_inp) pick_conf = pick_conf[0].permute(1, 2, 0).detach().cpu().numpy() # argmax = np.argmax(pick_conf) # import IPython; IPython.embed() argmax = np.unravel_index(argmax, shape=pick_conf.shape) p0_pix = argmax[:2] p0_theta = argmax[2] * (2 * np.pi / pick_conf.shape[2]) # Transport model forward pass. place_inp = {'inp_img': img, 'p0': p0_pix, 'lang_goal': lang_goal} place_conf = self.trans_forward(place_inp) place_conf = place_conf.squeeze().permute(1, 2, 0) place_conf = place_conf.detach().cpu().numpy() argmax = np.argmax(place_conf) argmax = np.unravel_index(argmax, shape=place_conf.shape) p1_pix = argmax[:2] p1_theta = argmax[2] * (2 * np.pi / place_conf.shape[2]) # Pixels to end effector poses. hmap = img[:, :, 3] p0_xyz = utils.pix_to_xyz(p0_pix, hmap, self.bounds, self.pix_size) p1_xyz = utils.pix_to_xyz(p1_pix, hmap, self.bounds, self.pix_size) p0_xyzw = utils.eulerXYZ_to_quatXYZW((0, 0, -p0_theta)) p1_xyzw = utils.eulerXYZ_to_quatXYZW((0, 0, -p1_theta)) return { 'pose0': (np.asarray(p0_xyz), np.asarray(p0_xyzw)), 'pose1': (np.asarray(p1_xyz), np.asarray(p1_xyzw)), 'pick': [p0_pix[0], p0_pix[1], p0_theta], 'place': [p1_pix[0], p1_pix[1], p1_theta], } class TwoStreamClipFilmLingUNetLatTransporterAgent(TwoStreamClipLingUNetTransporterAgent): def __init__(self, name, cfg, train_ds, test_ds): super().__init__(name, cfg, train_ds, test_ds) def _build_model(self): stream_one_fcn = 'plain_resnet_lat' stream_two_fcn = 'clip_film_lingunet_lat' self.attention = TwoStreamAttentionLangFusionLat( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=1, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) self.transport = TwoStreamTransportLangFusionLat( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=self.n_rotations, crop_size=self.crop_size, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) class TwoStreamClipLingUNetLatTransporterAgent(TwoStreamClipLingUNetTransporterAgent): # This is our model def __init__(self, name, cfg, train_ds, test_ds): super().__init__(name, cfg, train_ds, test_ds) def _build_model(self): stream_one_fcn = 'plain_resnet_lat' stream_two_fcn = 'clip_lingunet_lat' self.attention = TwoStreamAttentionLangFusionLat( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=1, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) self.transport = TwoStreamTransportLangFusionLat( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=self.n_rotations, crop_size=self.crop_size, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) class TwoStreamClipLingUNetLatTransporterAgentReduce(TwoStreamClipLingUNetTransporterAgent): # This is our model def __init__(self, name, cfg, train_ds, test_ds): super().__init__(name, cfg, train_ds, test_ds) def _build_model(self): stream_one_fcn = 'plain_resnet_lat' stream_two_fcn = 'clip_lingunet_lat' self.attention = TwoStreamAttentionLangFusionLat( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=1, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) self.transport = TwoStreamTransportLangFusionLatReduce( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=self.n_rotations, crop_size=self.crop_size, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) class TwoStreamClipLingUNetLatTransporterAgentReduceOneStream(TwoStreamClipLingUNetTransporterAgent): # This is our model def __init__(self, name, cfg, train_ds, test_ds): super().__init__(name, cfg, train_ds, test_ds) def _build_model(self): stream_one_fcn = 'plain_resnet_lat' stream_two_fcn = 'clip_lingunet_lat' self.attention = TwoStreamAttentionLangFusionLatReduce( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=1, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) self.transport = TwoStreamTransportLangFusionLatReduceOneStream( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=self.n_rotations, crop_size=self.crop_size, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) class TwoStreamClipLingUNetLatTransporterAgentReducePretrained(TwoStreamClipLingUNetTransporterAgent): # This is our model def __init__(self, name, cfg, train_ds, test_ds): super().__init__(name, cfg, train_ds, test_ds) def _build_model(self): stream_one_fcn = 'plain_resnet_lat' stream_two_fcn = 'clip_lingunet_lat' self.attention = TwoStreamAttentionLangFusionLat( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=1, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) self.transport = TwoStreamTransportLangFusionLatPretrained18( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=self.n_rotations, crop_size=self.crop_size, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) class TwoStreamRN50BertLingUNetTransporterAgent(TwoStreamClipLingUNetTransporterAgent): def __init__(self, name, cfg, train_ds, test_ds): super().__init__(name, cfg, train_ds, test_ds) def _build_model(self): stream_one_fcn = 'plain_resnet' stream_two_fcn = 'rn50_bert_lingunet' self.attention = TwoStreamAttentionLangFusion( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=1, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) self.transport = TwoStreamTransportLangFusion( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=self.n_rotations, crop_size=self.crop_size, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) class TwoStreamUntrainedRN50BertLingUNetTransporterAgent(TwoStreamClipLingUNetTransporterAgent): def __init__(self, name, cfg, train_ds, test_ds): super().__init__(name, cfg, train_ds, test_ds) def _build_model(self): stream_one_fcn = 'plain_resnet' stream_two_fcn = 'untrained_rn50_bert_lingunet' self.attention = TwoStreamAttentionLangFusion( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=1, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) self.transport = TwoStreamTransportLangFusion( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=self.n_rotations, crop_size=self.crop_size, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) class TwoStreamRN50BertLingUNetLatTransporterAgent(TwoStreamClipLingUNetTransporterAgent): def __init__(self, name, cfg, train_ds, test_ds): super().__init__(name, cfg, train_ds, test_ds) def _build_model(self): stream_one_fcn = 'plain_resnet_lat' stream_two_fcn = 'rn50_bert_lingunet_lat' self.attention = TwoStreamAttentionLangFusionLat( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=1, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) self.transport = TwoStreamTransportLangFusionLat( stream_fcn=(stream_one_fcn, stream_two_fcn), in_shape=self.in_shape, n_rotations=self.n_rotations, crop_size=self.crop_size, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) class OriginalTransporterLangFusionAgent(TwoStreamClipLingUNetTransporterAgent): def __init__(self, name, cfg, train_ds, test_ds): super().__init__(name, cfg, train_ds, test_ds) def _build_model(self): stream_fcn = 'plain_resnet_lang' self.attention = OneStreamAttentionLangFusion( stream_fcn=(stream_fcn, None), in_shape=self.in_shape, n_rotations=1, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) self.transport = OneStreamTransportLangFusion( stream_fcn=(stream_fcn, None), in_shape=self.in_shape, n_rotations=self.n_rotations, crop_size=self.crop_size, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) class ClipLingUNetTransporterAgent(TwoStreamClipLingUNetTransporterAgent): def __init__(self, name, cfg, train_ds, test_ds): super().__init__(name, cfg, train_ds, test_ds) def _build_model(self): stream_fcn = 'clip_lingunet' self.attention = OneStreamAttentionLangFusion( stream_fcn=(stream_fcn, None), in_shape=self.in_shape, n_rotations=1, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, ) self.transport = OneStreamTransportLangFusion( stream_fcn=(stream_fcn, None), in_shape=self.in_shape, n_rotations=self.n_rotations, crop_size=self.crop_size, preprocess=utils.preprocess, cfg=self.cfg, device=self.device_type, )