ravens/agents/conv_mlp.py
def act(self, obs, gt_act, info):
        """Run inference and return best action given visual observations."""
        self.regression_model.set_batch_size(1)

        act = {'camera_config': self.camera_config, 'primitive': None}
        if not obs:
            return act

        # Get heightmap from RGB-D images.
        colormap, heightmap = self.get_heightmap(obs, self.camera_config)

        # Concatenate color with depth images.
        input_image = np.concatenate((colormap,
                                      heightmap[..., None],
                                      heightmap[..., None],
                                      heightmap[..., None]), axis=2)[None, ...]

        # or just use rgb
        #input_image = colormap[None, ...]

        # Regression
        prediction = self.regression_model.forward(input_image)

        if self.USE_MDN:
            mdn_prediction = prediction
            pi, mu, var = mdn_prediction
            #prediction = mdn_utils.pick_max_mean(pi, mu, var)
            prediction = mdn_utils.sample_from_pdf(pi, mu, var)
            prediction = prediction[:, 0, :]

        prediction = prediction[0]

        p0_position = np.hstack((prediction[0:2], 0.02))
        p1_position = np.hstack((prediction[3:5], 0.02))

        p0_rotation = utils.get_pybullet_quaternion_from_rot((0, 0, -prediction[2]*self.THETA_SCALE))
        p1_rotation = utils.get_pybullet_quaternion_from_rot((0, 0, -prediction[5]*self.THETA_SCALE))

        act['primitive'] = 'pick_place'
        if self.task == 'sweeping':
            act['primitive'] = 'sweep'
        elif self.task == 'pushing':
            act['primitive'] = 'push'
        params = {'pose0': (p0_position, p0_rotation),
                  'pose1': (p1_position, p1_rotation)}
        act['params'] = params
        self.regression_model.set_batch_size(self.batch_size)

        return act
Run inference and return best action given visual observations.
Run inference and return best action given visual observations.
Run inference and return best action given visual observations .
def act(self, obs, gt_act, info):
        self.regression_model.set_batch_size(1)
        act = {'camera_config': self.camera_config, 'primitive': None}
        if not obs:
            return act
        colormap, heightmap = self.get_heightmap(obs, self.camera_config)
        input_image = np.concatenate((colormap,
                                      heightmap[..., None],
                                      heightmap[..., None],
                                      heightmap[..., None]), axis=2)[None, ...]
        prediction = self.regression_model.forward(input_image)
        if self.USE_MDN:
            mdn_prediction = prediction
            pi, mu, var = mdn_prediction
            prediction = mdn_utils.sample_from_pdf(pi, mu, var)
            prediction = prediction[:, 0, :]
        prediction = prediction[0]
        p0_position = np.hstack((prediction[0:2], 0.02))
        p1_position = np.hstack((prediction[3:5], 0.02))
        p0_rotation = utils.get_pybullet_quaternion_from_rot((0, 0, -prediction[2]*self.THETA_SCALE))
        p1_rotation = utils.get_pybullet_quaternion_from_rot((0, 0, -prediction[5]*self.THETA_SCALE))
        act['primitive'] = 'pick_place'
        if self.task == 'sweeping':
            act['primitive'] = 'sweep'
        elif self.task == 'pushing':
            act['primitive'] = 'push'
        params = {'pose0': (p0_position, p0_rotation),
                  'pose1': (p1_position, p1_rotation)}
        act['params'] = params
        self.regression_model.set_batch_size(self.batch_size)
        return act
def act ( self , obs , gt_act , info ) : self . regression_model . set_batch_size ( 1 ) act = { 'camera_config' : self . camera_config , 'primitive' : None } if not obs : return act colormap , heightmap = self . get_heightmap ( obs , self . camera_config ) input_image = np . concatenate ( ( colormap , heightmap [ ... , None ] , heightmap [ ... , None ] , heightmap [ ... , None ] ) , axis = 2 ) [ None , ... ] prediction = self . regression_model . forward ( input_image ) if self . USE_MDN : mdn_prediction = prediction pi , mu , var = mdn_prediction prediction = mdn_utils . sample_from_pdf ( pi , mu , var ) prediction = prediction [ : , 0 , : ] prediction = prediction [ 0 ] p0_position = np . hstack ( ( prediction [ 0 : 2 ] , 0.02 ) ) p1_position = np . hstack ( ( prediction [ 3 : 5 ] , 0.02 ) ) p0_rotation = utils . get_pybullet_quaternion_from_rot ( ( 0 , 0 , - prediction [ 2 ] * self . THETA_SCALE ) ) p1_rotation = utils . get_pybullet_quaternion_from_rot ( ( 0 , 0 , - prediction [ 5 ] * self . THETA_SCALE ) ) act [ 'primitive' ] = 'pick_place' if self . task == 'sweeping' : act [ 'primitive' ] = 'sweep' elif self . task == 'pushing' : act [ 'primitive' ] = 'push' params = { 'pose0' : ( p0_position , p0_rotation ) , 'pose1' : ( p1_position , p1_rotation ) } act [ 'params' ] = params self . regression_model . set_batch_size ( self . batch_size ) return act
Run inference and return best action given visual observations.
Run inference and return best action given visual observations .
"""Run inference and return best action given visual observations.""" # Get heightmap from RGB-D images. # Concatenate color with depth images. # or just use rgb #input_image = colormap[None, ...] # Regression #prediction = mdn_utils.pick_max_mean(pi, mu, var)
self
obs
gt_act
info