# Copyright 2016 The TensorFlow Authors All Rights Reserved. # # Licensed under the Apache License, Version 2.0 (the "License"); # you may not use this file except in compliance with the License. # You may obtain a copy of the License at # # http://www.apache.org/licenses/LICENSE-2.0 # # Unless required by applicable law or agreed to in writing, software # distributed under the License is distributed on an "AS IS" BASIS, # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== r"""Navidation Environment. Includes the following classes along with some helper functions. Building: Loads buildings, computes traversibility, exposes functionality for rendering images. GridWorld: Base class which implements functionality for moving an agent on a grid world. NavigationEnv: Base class which generates navigation problems on a grid world. VisualNavigationEnv: Builds upon NavigationEnv and Building to provide interface that is used externally to train the agent. MeshMapper: Class used for distilling the model, testing the mapper. BuildingMultiplexer: Wrapper class that instantiates a VisualNavigationEnv for each building and multiplexes between them as needed. """ import numpy as np import os import re import matplotlib.pyplot as plt import graph_tool as gt import graph_tool.topology from tensorflow.python.platform import gfile import logging import src.file_utils as fu import src.utils as utils import src.graph_utils as gu import src.map_utils as mu import src.depth_utils as du import render.swiftshader_renderer as sru from render.swiftshader_renderer import SwiftshaderRenderer import cv2 label_nodes_with_class = gu.label_nodes_with_class label_nodes_with_class_geodesic = gu.label_nodes_with_class_geodesic get_distance_node_list = gu.get_distance_node_list convert_to_graph_tool = gu.convert_to_graph_tool generate_graph = gu.generate_graph get_hardness_distribution = gu.get_hardness_distribution rng_next_goal_rejection_sampling = gu.rng_next_goal_rejection_sampling rng_next_goal = gu.rng_next_goal rng_room_to_room = gu.rng_room_to_room rng_target_dist_field = gu.rng_target_dist_field compute_traversibility = mu.compute_traversibility make_map = mu.make_map resize_maps = mu.resize_maps pick_largest_cc = mu.pick_largest_cc get_graph_origin_loc = mu.get_graph_origin_loc generate_egocentric_maps = mu.generate_egocentric_maps generate_goal_images = mu.generate_goal_images get_map_to_predict = mu.get_map_to_predict bin_points = du.bin_points make_geocentric = du.make_geocentric get_point_cloud_from_z = du.get_point_cloud_from_z get_camera_matrix = du.get_camera_matrix def _get_semantic_maps(folder_name, building_name, map, flip): # Load file from the cache. file_name = '{:s}_{:d}_{:d}_{:d}_{:d}_{:d}_{:d}.pkl' file_name = file_name.format(building_name, map.size[0], map.size[1], map.origin[0], map.origin[1], map.resolution, flip) file_name = os.path.join(folder_name, file_name) logging.info('Loading semantic maps from %s.', file_name) if fu.exists(file_name): a = utils.load_variables(file_name) maps = a['maps'] #HxWx#C cats = a['cats'] else: logging.error('file_name: %s not found.', file_name) maps = None cats = None return maps, cats def _select_classes(all_maps, all_cats, cats_to_use): inds = [] for c in cats_to_use: ind = all_cats.index(c) inds.append(ind) out_maps = all_maps[:,:,inds] return out_maps def _get_room_dimensions(file_name, resolution, origin, flip=False): if fu.exists(file_name): a = utils.load_variables(file_name)['room_dimension'] names = a.keys() dims = np.concatenate(a.values(), axis=0).reshape((-1,6)) ind = np.argsort(names) dims = dims[ind,:] names = [names[x] for x in ind] if flip: dims_new = dims*1 dims_new[:,1] = -dims[:,4] dims_new[:,4] = -dims[:,1] dims = dims_new*1 dims = dims*100. dims[:,0] = dims[:,0] - origin[0] dims[:,1] = dims[:,1] - origin[1] dims[:,3] = dims[:,3] - origin[0] dims[:,4] = dims[:,4] - origin[1] dims = dims / resolution out = {'names': names, 'dims': dims} else: out = None return out def _filter_rooms(room_dims, room_regex): pattern = re.compile(room_regex) ind = [] for i, name in enumerate(room_dims['names']): if pattern.match(name): ind.append(i) new_room_dims = {} new_room_dims['names'] = [room_dims['names'][i] for i in ind] new_room_dims['dims'] = room_dims['dims'][ind,:]*1 return new_room_dims def _label_nodes_with_room_id(xyt, room_dims): # Label the room with the ID into things. node_room_id = -1*np.ones((xyt.shape[0], 1)) dims = room_dims['dims'] for x, name in enumerate(room_dims['names']): all_ = np.concatenate((xyt[:,[0]] >= dims[x,0], xyt[:,[0]] <= dims[x,3], xyt[:,[1]] >= dims[x,1], xyt[:,[1]] <= dims[x,4]), axis=1) node_room_id[np.all(all_, axis=1), 0] = x return node_room_id def get_path_ids(start_node_id, end_node_id, pred_map): id = start_node_id path = [id] while id != end_node_id: id = pred_map[id] path.append(id) return path def image_pre(images, modalities): # Assumes images are ...xHxWxC. # We always assume images are RGB followed by Depth. if 'depth' in modalities: d = images[...,-1][...,np.newaxis]*1. d[d < 0.01] = np.NaN; isnan = np.isnan(d); d = 100./d; d[isnan] = 0.; images = np.concatenate((images[...,:-1], d, isnan), axis=images.ndim-1) if 'rgb' in modalities: images[...,:3] = images[...,:3]*1. - 128 return images def _get_relative_goal_loc(goal_loc, loc, theta): r = np.sqrt(np.sum(np.square(goal_loc - loc), axis=1)) t = np.arctan2(goal_loc[:,1] - loc[:,1], goal_loc[:,0] - loc[:,0]) t = t-theta[:,0] + np.pi/2 return np.expand_dims(r,axis=1), np.expand_dims(t, axis=1) def _gen_perturbs(rng, batch_size, num_steps, lr_flip, delta_angle, delta_xy, structured): perturbs = [] for i in range(batch_size): # Doing things one by one for each episode in this batch. This way this # remains replicatable even when we change the batch size. p = np.zeros((num_steps+1, 4)) if lr_flip: # Flip the whole trajectory. p[:,3] = rng.rand(1)-0.5 if delta_angle > 0: if structured: p[:,2] = (rng.rand(1)-0.5)* delta_angle else: p[:,2] = (rng.rand(p.shape[0])-0.5)* delta_angle if delta_xy > 0: if structured: p[:,:2] = (rng.rand(1, 2)-0.5)*delta_xy else: p[:,:2] = (rng.rand(p.shape[0], 2)-0.5)*delta_xy perturbs.append(p) return perturbs def get_multiplexer_class(args, task_number): assert(args.task_params.base_class == 'Building') logging.info('Returning BuildingMultiplexer') R = BuildingMultiplexer(args, task_number) return R class GridWorld(): def __init__(self): """Class members that will be assigned by any class that actually uses this class.""" self.restrict_to_largest_cc = None self.robot = None self.env = None self.category_list = None self.traversible = None def get_loc_axis(self, node, delta_theta, perturb=None): """Based on the node orientation returns X, and Y axis. Used to sample the map in egocentric coordinate frame. """ if type(node) == tuple: node = np.array([node]) if perturb is None: perturb = np.zeros((node.shape[0], 4)) xyt = self.to_actual_xyt_vec(node) x = xyt[:,[0]] + perturb[:,[0]] y = xyt[:,[1]] + perturb[:,[1]] t = xyt[:,[2]] + perturb[:,[2]] theta = t*delta_theta loc = np.concatenate((x,y), axis=1) x_axis = np.concatenate((np.cos(theta), np.sin(theta)), axis=1) y_axis = np.concatenate((np.cos(theta+np.pi/2.), np.sin(theta+np.pi/2.)), axis=1) # Flip the sampled map where need be. y_axis[np.where(perturb[:,3] > 0)[0], :] *= -1. return loc, x_axis, y_axis, theta def to_actual_xyt(self, pqr): """Converts from node to location on the map.""" (p, q, r) = pqr if self.task.n_ori == 6: out = (p - q * 0.5 + self.task.origin_loc[0], q * np.sqrt(3.) / 2. + self.task.origin_loc[1], r) elif self.task.n_ori == 4: out = (p + self.task.origin_loc[0], q + self.task.origin_loc[1], r) return out def to_actual_xyt_vec(self, pqr): """Converts from node array to location array on the map.""" p = pqr[:,0][:, np.newaxis] q = pqr[:,1][:, np.newaxis] r = pqr[:,2][:, np.newaxis] if self.task.n_ori == 6: out = np.concatenate((p - q * 0.5 + self.task.origin_loc[0], q * np.sqrt(3.) / 2. + self.task.origin_loc[1], r), axis=1) elif self.task.n_ori == 4: out = np.concatenate((p + self.task.origin_loc[0], q + self.task.origin_loc[1], r), axis=1) return out def raw_valid_fn_vec(self, xyt): """Returns if the given set of nodes is valid or not.""" height = self.traversible.shape[0] width = self.traversible.shape[1] x = np.round(xyt[:,[0]]).astype(np.int32) y = np.round(xyt[:,[1]]).astype(np.int32) is_inside = np.all(np.concatenate((x >= 0, y >= 0, x < width, y < height), axis=1), axis=1) x = np.minimum(np.maximum(x, 0), width-1) y = np.minimum(np.maximum(y, 0), height-1) ind = np.ravel_multi_index((y,x), self.traversible.shape) is_traversible = self.traversible.ravel()[ind] is_valid = np.all(np.concatenate((is_inside[:,np.newaxis], is_traversible), axis=1), axis=1) return is_valid def valid_fn_vec(self, pqr): """Returns if the given set of nodes is valid or not.""" xyt = self.to_actual_xyt_vec(np.array(pqr)) height = self.traversible.shape[0] width = self.traversible.shape[1] x = np.round(xyt[:,[0]]).astype(np.int32) y = np.round(xyt[:,[1]]).astype(np.int32) is_inside = np.all(np.concatenate((x >= 0, y >= 0, x < width, y < height), axis=1), axis=1) x = np.minimum(np.maximum(x, 0), width-1) y = np.minimum(np.maximum(y, 0), height-1) ind = np.ravel_multi_index((y,x), self.traversible.shape) is_traversible = self.traversible.ravel()[ind] is_valid = np.all(np.concatenate((is_inside[:,np.newaxis], is_traversible), axis=1), axis=1) return is_valid def get_feasible_actions(self, node_ids): """Returns the feasible set of actions from the current node.""" a = np.zeros((len(node_ids), self.task_params.num_actions), dtype=np.int32) gtG = self.task.gtG next_node = [] for i, c in enumerate(node_ids): neigh = gtG.vertex(c).out_neighbours() neigh_edge = gtG.vertex(c).out_edges() nn = {} for n, e in zip(neigh, neigh_edge): _ = gtG.ep['action'][e] a[i,_] = 1 nn[_] = int(n) next_node.append(nn) return a, next_node def take_action(self, current_node_ids, action): """Returns the new node after taking the action action. Stays at the current node if the action is invalid.""" actions, next_node_ids = self.get_feasible_actions(current_node_ids) new_node_ids = [] for i, (c,a) in enumerate(zip(current_node_ids, action)): if actions[i,a] == 1: new_node_ids.append(next_node_ids[i][a]) else: new_node_ids.append(c) return new_node_ids def set_r_obj(self, r_obj): """Sets the SwiftshaderRenderer object used for rendering.""" self.r_obj = r_obj class Building(GridWorld): def __init__(self, building_name, robot, env, category_list=None, small=False, flip=False, logdir=None, building_loader=None): self.restrict_to_largest_cc = True self.robot = robot self.env = env self.logdir = logdir # Load the building meta data. building = building_loader.load_building(building_name) if small: building['mesh_names'] = building['mesh_names'][:5] # New code. shapess = building_loader.load_building_meshes(building) if flip: for shapes in shapess: shapes.flip_shape() vs = [] for shapes in shapess: vs.append(shapes.get_vertices()[0]) vs = np.concatenate(vs, axis=0) map = make_map(env.padding, env.resolution, vertex=vs, sc=100.) map = compute_traversibility( map, robot.base, robot.height, robot.radius, env.valid_min, env.valid_max, env.num_point_threshold, shapess=shapess, sc=100., n_samples_per_face=env.n_samples_per_face) room_dims = _get_room_dimensions(building['room_dimension_file'], env.resolution, map.origin, flip=flip) class_maps, class_map_names = _get_semantic_maps( building['class_map_folder'], building_name, map, flip) self.class_maps = class_maps self.class_map_names = class_map_names self.building = building self.shapess = shapess self.map = map self.traversible = map.traversible*1 self.building_name = building_name self.room_dims = room_dims self.flipped = flip self.renderer_entitiy_ids = [] if self.restrict_to_largest_cc: self.traversible = pick_largest_cc(self.traversible) def load_building_into_scene(self): # Loads the scene. self.renderer_entitiy_ids += self.r_obj.load_shapes(self.shapess) # Free up memory, we dont need the mesh or the materials anymore. self.shapess = None def add_entity_at_nodes(self, nodes, height, shape): xyt = self.to_actual_xyt_vec(nodes) nxy = xyt[:,:2]*1. nxy = nxy * self.map.resolution nxy = nxy + self.map.origin Ts = np.concatenate((nxy, nxy[:,:1]), axis=1) Ts[:,2] = height; Ts = Ts / 100.; # Merge all the shapes into a single shape and add that shape. shape.replicate_shape(Ts) entity_ids = self.r_obj.load_shapes([shape]) self.renderer_entitiy_ids += entity_ids return entity_ids def add_shapes(self, shapes): scene = self.r_obj.viz.scene() for shape in shapes: scene.AddShape(shape) def add_materials(self, materials): scene = self.r_obj.viz.scene() for material in materials: scene.AddOrUpdateMaterial(material) def set_building_visibility(self, visibility): self.r_obj.set_entity_visible(self.renderer_entitiy_ids, visibility) def render_nodes(self, nodes, perturb=None, aux_delta_theta=0.): self.set_building_visibility(True) if perturb is None: perturb = np.zeros((len(nodes), 4)) imgs = [] r = 2 elevation_z = r * np.tan(np.deg2rad(self.robot.camera_elevation_degree)) for i in range(len(nodes)): xyt = self.to_actual_xyt(nodes[i]) lookat_theta = 3.0 * np.pi / 2.0 - (xyt[2]+perturb[i,2]+aux_delta_theta) * (self.task.delta_theta) nxy = np.array([xyt[0]+perturb[i,0], xyt[1]+perturb[i,1]]).reshape(1, -1) nxy = nxy * self.map.resolution nxy = nxy + self.map.origin camera_xyz = np.zeros((1, 3)) camera_xyz[...] = [nxy[0, 0], nxy[0, 1], self.robot.sensor_height] camera_xyz = camera_xyz / 100. lookat_xyz = np.array([-r * np.sin(lookat_theta), -r * np.cos(lookat_theta), elevation_z]) lookat_xyz = lookat_xyz + camera_xyz[0, :] self.r_obj.position_camera(camera_xyz[0, :].tolist(), lookat_xyz.tolist(), [0.0, 0.0, 1.0]) img = self.r_obj.render(take_screenshot=True, output_type=0) img = [x for x in img if x is not None] img = np.concatenate(img, axis=2).astype(np.float32) if perturb[i,3]>0: img = img[:,::-1,:] imgs.append(img) self.set_building_visibility(False) return imgs class MeshMapper(Building): def __init__(self, robot, env, task_params, building_name, category_list, flip, logdir=None, building_loader=None): Building.__init__(self, building_name, robot, env, category_list, small=task_params.toy_problem, flip=flip, logdir=logdir, building_loader=building_loader) self.task_params = task_params self.task = None self._preprocess_for_task(self.task_params.building_seed) def _preprocess_for_task(self, seed): if self.task is None or self.task.seed != seed: rng = np.random.RandomState(seed) origin_loc = get_graph_origin_loc(rng, self.traversible) self.task = utils.Foo(seed=seed, origin_loc=origin_loc, n_ori=self.task_params.n_ori) G = generate_graph(self.valid_fn_vec, self.task_params.step_size, self.task.n_ori, (0, 0, 0)) gtG, nodes, nodes_to_id = convert_to_graph_tool(G) self.task.gtG = gtG self.task.nodes = nodes self.task.delta_theta = 2.0*np.pi/(self.task.n_ori*1.) self.task.nodes_to_id = nodes_to_id logging.info('Building %s, #V=%d, #E=%d', self.building_name, self.task.nodes.shape[0], self.task.gtG.num_edges()) if self.logdir is not None: write_traversible = cv2.applyColorMap(self.traversible.astype(np.uint8)*255, cv2.COLORMAP_JET) img_path = os.path.join(self.logdir, '{:s}_{:d}_graph.png'.format(self.building_name, seed)) node_xyt = self.to_actual_xyt_vec(self.task.nodes) plt.set_cmap('jet'); fig, ax = utils.subplot(plt, (1,1), (12,12)) ax.plot(node_xyt[:,0], node_xyt[:,1], 'm.') ax.imshow(self.traversible, origin='lower'); ax.set_axis_off(); ax.axis('equal'); ax.set_title('{:s}, {:d}, {:d}'.format(self.building_name, self.task.nodes.shape[0], self.task.gtG.num_edges())) if self.room_dims is not None: for i, r in enumerate(self.room_dims['dims']*1): min_ = r[:3]*1 max_ = r[3:]*1 xmin, ymin, zmin = min_ xmax, ymax, zmax = max_ ax.plot([xmin, xmax, xmax, xmin, xmin], [ymin, ymin, ymax, ymax, ymin], 'g') with fu.fopen(img_path, 'w') as f: fig.savefig(f, bbox_inches='tight', transparent=True, pad_inches=0) plt.close(fig) def _gen_rng(self, rng): # instances is a list of list of node_ids. if self.task_params.move_type == 'circle': _, _, _, _, paths = rng_target_dist_field(self.task_params.batch_size, self.task.gtG, rng, 0, 1, compute_path=True) instances_ = paths instances = [] for instance_ in instances_: instance = instance_ for i in range(self.task_params.num_steps): instance.append(self.take_action([instance[-1]], [1])[0]) instances.append(instance) elif self.task_params.move_type == 'shortest_path': _, _, _, _, paths = rng_target_dist_field(self.task_params.batch_size, self.task.gtG, rng, self.task_params.num_steps, self.task_params.num_steps+1, compute_path=True) instances = paths elif self.task_params.move_type == 'circle+forward': _, _, _, _, paths = rng_target_dist_field(self.task_params.batch_size, self.task.gtG, rng, 0, 1, compute_path=True) instances_ = paths instances = [] for instance_ in instances_: instance = instance_ for i in range(self.task_params.n_ori-1): instance.append(self.take_action([instance[-1]], [1])[0]) while len(instance) <= self.task_params.num_steps: while self.take_action([instance[-1]], [3])[0] == instance[-1] and len(instance) <= self.task_params.num_steps: instance.append(self.take_action([instance[-1]], [2])[0]) if len(instance) <= self.task_params.num_steps: instance.append(self.take_action([instance[-1]], [3])[0]) instances.append(instance) # Do random perturbation if needed. perturbs = _gen_perturbs(rng, self.task_params.batch_size, self.task_params.num_steps, self.task_params.data_augment.lr_flip, self.task_params.data_augment.delta_angle, self.task_params.data_augment.delta_xy, self.task_params.data_augment.structured) return instances, perturbs def worker(self, instances, perturbs): # Output the images and the free space. # Make the instances be all the same length. for i in range(len(instances)): for j in range(self.task_params.num_steps - len(instances[i]) + 1): instances[i].append(instances[i][-1]) if perturbs[i].shape[0] < self.task_params.num_steps+1: p = np.zeros((self.task_params.num_steps+1, 4)) p[:perturbs[i].shape[0], :] = perturbs[i] p[perturbs[i].shape[0]:, :] = perturbs[i][-1,:] perturbs[i] = p instances_ = [] for instance in instances: instances_ = instances_ + instance perturbs_ = np.concatenate(perturbs, axis=0) instances_nodes = self.task.nodes[instances_,:] instances_nodes = [tuple(x) for x in instances_nodes] imgs_ = self.render_nodes(instances_nodes, perturbs_) imgs = []; next = 0; for instance in instances: img_i = [] for _ in instance: img_i.append(imgs_[next]) next = next+1 imgs.append(img_i) imgs = np.array(imgs) # Render out the maps in the egocentric view for all nodes and not just the # last node. all_nodes = [] for x in instances: all_nodes = all_nodes + x all_perturbs = np.concatenate(perturbs, axis=0) loc, x_axis, y_axis, theta = self.get_loc_axis( self.task.nodes[all_nodes, :]*1, delta_theta=self.task.delta_theta, perturb=all_perturbs) fss = None valids = None loc_on_map = None theta_on_map = None cum_fs = None cum_valid = None incremental_locs = None incremental_thetas = None if self.task_params.output_free_space: fss, valids = get_map_to_predict(loc, x_axis, y_axis, map=self.traversible*1., map_size=self.task_params.map_size) fss = np.array(fss) > 0.5 fss = np.reshape(fss, [self.task_params.batch_size, self.task_params.num_steps+1, self.task_params.map_size, self.task_params.map_size]) valids = np.reshape(np.array(valids), fss.shape) if self.task_params.output_transform_to_global_map: # Output the transform to the global map. loc_on_map = np.reshape(loc*1, [self.task_params.batch_size, self.task_params.num_steps+1, -1]) # Converting to location wrt to first location so that warping happens # properly. theta_on_map = np.reshape(theta*1, [self.task_params.batch_size, self.task_params.num_steps+1, -1]) if self.task_params.output_incremental_transform: # Output the transform to the global map. incremental_locs_ = np.reshape(loc*1, [self.task_params.batch_size, self.task_params.num_steps+1, -1]) incremental_locs_[:,1:,:] -= incremental_locs_[:,:-1,:] t0 = -np.pi/2+np.reshape(theta*1, [self.task_params.batch_size, self.task_params.num_steps+1, -1]) t = t0*1 incremental_locs = incremental_locs_*1 incremental_locs[:,:,0] = np.sum(incremental_locs_ * np.concatenate((np.cos(t), np.sin(t)), axis=-1), axis=-1) incremental_locs[:,:,1] = np.sum(incremental_locs_ * np.concatenate((np.cos(t+np.pi/2), np.sin(t+np.pi/2)), axis=-1), axis=-1) incremental_locs[:,0,:] = incremental_locs_[:,0,:] # print incremental_locs_[0,:,:], incremental_locs[0,:,:], t0[0,:,:] incremental_thetas = np.reshape(theta*1, [self.task_params.batch_size, self.task_params.num_steps+1, -1]) incremental_thetas[:,1:,:] += -incremental_thetas[:,:-1,:] if self.task_params.output_canonical_map: loc_ = loc[0::(self.task_params.num_steps+1), :] x_axis = np.zeros_like(loc_); x_axis[:,1] = 1 y_axis = np.zeros_like(loc_); y_axis[:,0] = -1 cum_fs, cum_valid = get_map_to_predict(loc_, x_axis, y_axis, map=self.traversible*1., map_size=self.task_params.map_size) cum_fs = np.array(cum_fs) > 0.5 cum_fs = np.reshape(cum_fs, [self.task_params.batch_size, 1, self.task_params.map_size, self.task_params.map_size]) cum_valid = np.reshape(np.array(cum_valid), cum_fs.shape) inputs = {'fs_maps': fss, 'valid_maps': valids, 'imgs': imgs, 'loc_on_map': loc_on_map, 'theta_on_map': theta_on_map, 'cum_fs_maps': cum_fs, 'cum_valid_maps': cum_valid, 'incremental_thetas': incremental_thetas, 'incremental_locs': incremental_locs} return inputs def pre(self, inputs): inputs['imgs'] = image_pre(inputs['imgs'], self.task_params.modalities) if inputs['loc_on_map'] is not None: inputs['loc_on_map'] = inputs['loc_on_map'] - inputs['loc_on_map'][:,[0],:] if inputs['theta_on_map'] is not None: inputs['theta_on_map'] = np.pi/2. - inputs['theta_on_map'] return inputs def _nav_env_reset_helper(type, rng, nodes, batch_size, gtG, max_dist, num_steps, num_goals, data_augment, **kwargs): """Generates and returns a new episode.""" max_compute = max_dist + 4*num_steps if type == 'general': start_node_ids, end_node_ids, dist, pred_map, paths = \ rng_target_dist_field(batch_size, gtG, rng, max_dist, max_compute, nodes=nodes, compute_path=False) target_class = None elif type == 'room_to_room_many': goal_node_ids = []; dists = []; node_room_ids = kwargs['node_room_ids'] # Sample the first one start_node_ids_, end_node_ids_, dist_, _, _ = rng_room_to_room( batch_size, gtG, rng, max_dist, max_compute, node_room_ids=node_room_ids, nodes=nodes) start_node_ids = start_node_ids_ goal_node_ids.append(end_node_ids_) dists.append(dist_) for n in range(num_goals-1): start_node_ids_, end_node_ids_, dist_, _, _ = rng_next_goal( goal_node_ids[n], batch_size, gtG, rng, max_dist, max_compute, node_room_ids=node_room_ids, nodes=nodes, dists_from_start_node=dists[n]) goal_node_ids.append(end_node_ids_) dists.append(dist_) target_class = None elif type == 'rng_rejection_sampling_many': num_goals = num_goals goal_node_ids = []; dists = []; n_ori = kwargs['n_ori'] step_size = kwargs['step_size'] min_dist = kwargs['min_dist'] sampling_distribution = kwargs['sampling_distribution'] target_distribution = kwargs['target_distribution'] rejection_sampling_M = kwargs['rejection_sampling_M'] distribution_bins = kwargs['distribution_bins'] for n in range(num_goals): if n == 0: input_nodes = None else: input_nodes = goal_node_ids[n-1] start_node_ids_, end_node_ids_, dist_, _, _, _, _ = rng_next_goal_rejection_sampling( input_nodes, batch_size, gtG, rng, max_dist, min_dist, max_compute, sampling_distribution, target_distribution, nodes, n_ori, step_size, distribution_bins, rejection_sampling_M) if n == 0: start_node_ids = start_node_ids_ goal_node_ids.append(end_node_ids_) dists.append(dist_) target_class = None elif type == 'room_to_room_back': num_goals = num_goals assert(num_goals == 2), 'num_goals must be 2.' goal_node_ids = []; dists = []; node_room_ids = kwargs['node_room_ids'] # Sample the first one. start_node_ids_, end_node_ids_, dist_, _, _ = rng_room_to_room( batch_size, gtG, rng, max_dist, max_compute, node_room_ids=node_room_ids, nodes=nodes) start_node_ids = start_node_ids_ goal_node_ids.append(end_node_ids_) dists.append(dist_) # Set second goal to be starting position, and compute distance to the start node. goal_node_ids.append(start_node_ids) dist = [] for i in range(batch_size): dist_ = gt.topology.shortest_distance( gt.GraphView(gtG, reversed=True), source=gtG.vertex(start_node_ids[i]), target=None) dist_ = np.array(dist_.get_array()) dist.append(dist_) dists.append(dist) target_class = None elif type[:14] == 'to_nearest_obj': # Generate an episode by sampling one of the target classes (with # probability proportional to the number of nodes in the world). # With the sampled class sample a node that is within some distance from # the sampled class. class_nodes = kwargs['class_nodes'] sampling = kwargs['sampling'] dist_to_class = kwargs['dist_to_class'] assert(num_goals == 1), 'Only supports a single goal.' ind = rng.choice(class_nodes.shape[0], size=batch_size) target_class = class_nodes[ind,1] start_node_ids = []; dists = []; goal_node_ids = []; for t in target_class: if sampling == 'uniform': max_dist = max_dist cnts = np.bincount(dist_to_class[t], minlength=max_dist+1)*1. cnts[max_dist+1:] = 0 p_each = 1./ cnts / (max_dist+1.) p_each[cnts == 0] = 0 p = p_each[dist_to_class[t]]*1.; p = p/np.sum(p) start_node_id = rng.choice(p.shape[0], size=1, p=p)[0] else: logging.fatal('Sampling not one of uniform.') start_node_ids.append(start_node_id) dists.append(dist_to_class[t]) # Dummy goal node, same as the start node, so that vis is better. goal_node_ids.append(start_node_id) dists = [dists] goal_node_ids = [goal_node_ids] return start_node_ids, goal_node_ids, dists, target_class class NavigationEnv(GridWorld, Building): """Wrapper around GridWorld which sets up navigation tasks. """ def _debug_save_hardness(self, seed): out_path = os.path.join(self.logdir, '{:s}_{:d}_hardness.png'.format(self.building_name, seed)) batch_size = 4000 rng = np.random.RandomState(0) start_node_ids, end_node_ids, dists, pred_maps, paths, hardnesss, gt_dists = \ rng_next_goal_rejection_sampling( None, batch_size, self.task.gtG, rng, self.task_params.max_dist, self.task_params.min_dist, self.task_params.max_dist, self.task.sampling_distribution, self.task.target_distribution, self.task.nodes, self.task_params.n_ori, self.task_params.step_size, self.task.distribution_bins, self.task.rejection_sampling_M) bins = self.task.distribution_bins n_bins = self.task.n_bins with plt.style.context('ggplot'): fig, axes = utils.subplot(plt, (1,2), (10,10)) ax = axes[0] _ = ax.hist(hardnesss, bins=bins, weights=np.ones_like(hardnesss)/len(hardnesss)) ax.plot(bins[:-1]+0.5/n_bins, self.task.target_distribution, 'g') ax.plot(bins[:-1]+0.5/n_bins, self.task.sampling_distribution, 'b') ax.grid('on') ax = axes[1] _ = ax.hist(gt_dists, bins=np.arange(self.task_params.max_dist+1)) ax.grid('on') ax.set_title('Mean: {:0.2f}, Median: {:0.2f}'.format(np.mean(gt_dists), np.median(gt_dists))) with fu.fopen(out_path, 'w') as f: fig.savefig(f, bbox_inches='tight', transparent=True, pad_inches=0) def _debug_save_map_nodes(self, seed): """Saves traversible space along with nodes generated on the graph. Takes the seed as input.""" img_path = os.path.join(self.logdir, '{:s}_{:d}_graph.png'.format(self.building_name, seed)) node_xyt = self.to_actual_xyt_vec(self.task.nodes) plt.set_cmap('jet'); fig, ax = utils.subplot(plt, (1,1), (12,12)) ax.plot(node_xyt[:,0], node_xyt[:,1], 'm.') ax.set_axis_off(); ax.axis('equal'); if self.room_dims is not None: for i, r in enumerate(self.room_dims['dims']*1): min_ = r[:3]*1 max_ = r[3:]*1 xmin, ymin, zmin = min_ xmax, ymax, zmax = max_ ax.plot([xmin, xmax, xmax, xmin, xmin], [ymin, ymin, ymax, ymax, ymin], 'g') ax.imshow(self.traversible, origin='lower'); with fu.fopen(img_path, 'w') as f: fig.savefig(f, bbox_inches='tight', transparent=True, pad_inches=0) def _debug_semantic_maps(self, seed): """Saves traversible space along with nodes generated on the graph. Takes the seed as input.""" for i, cls in enumerate(self.task_params.semantic_task.class_map_names): img_path = os.path.join(self.logdir, '{:s}_flip{:d}_{:s}_graph.png'.format(self.building_name, seed, cls)) maps = self.traversible*1. maps += 0.5*(self.task.class_maps_dilated[:,:,i]) write_traversible = (maps*1.+1.)/3.0 write_traversible = (write_traversible*255.).astype(np.uint8)[:,:,np.newaxis] write_traversible = write_traversible + np.zeros((1,1,3), dtype=np.uint8) fu.write_image(img_path, write_traversible[::-1,:,:]) def _preprocess_for_task(self, seed): """Sets up the task field for doing navigation on the grid world.""" if self.task is None or self.task.seed != seed: rng = np.random.RandomState(seed) origin_loc = get_graph_origin_loc(rng, self.traversible) self.task = utils.Foo(seed=seed, origin_loc=origin_loc, n_ori=self.task_params.n_ori) G = generate_graph(self.valid_fn_vec, self.task_params.step_size, self.task.n_ori, (0, 0, 0)) gtG, nodes, nodes_to_id = convert_to_graph_tool(G) self.task.gtG = gtG self.task.nodes = nodes self.task.delta_theta = 2.0*np.pi/(self.task.n_ori*1.) self.task.nodes_to_id = nodes_to_id logging.info('Building %s, #V=%d, #E=%d', self.building_name, self.task.nodes.shape[0], self.task.gtG.num_edges()) type = self.task_params.type if type == 'general': # Do nothing _ = None elif type == 'room_to_room_many' or type == 'room_to_room_back': if type == 'room_to_room_back': assert(self.task_params.num_goals == 2), 'num_goals must be 2.' self.room_dims = _filter_rooms(self.room_dims, self.task_params.room_regex) xyt = self.to_actual_xyt_vec(self.task.nodes) self.task.node_room_ids = _label_nodes_with_room_id(xyt, self.room_dims) self.task.reset_kwargs = {'node_room_ids': self.task.node_room_ids} elif type == 'rng_rejection_sampling_many': n_bins = 20 rejection_sampling_M = self.task_params.rejection_sampling_M min_dist = self.task_params.min_dist bins = np.arange(n_bins+1)/(n_bins*1.) target_d = np.zeros(n_bins); target_d[...] = 1./n_bins; sampling_d = get_hardness_distribution( self.task.gtG, self.task_params.max_dist, self.task_params.min_dist, np.random.RandomState(0), 4000, bins, self.task.nodes, self.task_params.n_ori, self.task_params.step_size) self.task.reset_kwargs = {'distribution_bins': bins, 'target_distribution': target_d, 'sampling_distribution': sampling_d, 'rejection_sampling_M': rejection_sampling_M, 'n_bins': n_bins, 'n_ori': self.task_params.n_ori, 'step_size': self.task_params.step_size, 'min_dist': self.task_params.min_dist} self.task.n_bins = n_bins self.task.distribution_bins = bins self.task.target_distribution = target_d self.task.sampling_distribution = sampling_d self.task.rejection_sampling_M = rejection_sampling_M if self.logdir is not None: self._debug_save_hardness(seed) elif type[:14] == 'to_nearest_obj': self.room_dims = _filter_rooms(self.room_dims, self.task_params.room_regex) xyt = self.to_actual_xyt_vec(self.task.nodes) self.class_maps = _select_classes(self.class_maps, self.class_map_names, self.task_params.semantic_task.class_map_names)*1 self.class_map_names = self.task_params.semantic_task.class_map_names nodes_xyt = self.to_actual_xyt_vec(np.array(self.task.nodes)) tt = utils.Timer(); tt.tic(); if self.task_params.type == 'to_nearest_obj_acc': self.task.class_maps_dilated, self.task.node_class_label = label_nodes_with_class_geodesic( nodes_xyt, self.class_maps, self.task_params.semantic_task.pix_distance+8, self.map.traversible, ff_cost=1., fo_cost=1., oo_cost=4., connectivity=8.) dists = [] for i in range(len(self.class_map_names)): class_nodes_ = np.where(self.task.node_class_label[:,i])[0] dists.append(get_distance_node_list(gtG, source_nodes=class_nodes_, direction='to')) self.task.dist_to_class = dists a_, b_ = np.where(self.task.node_class_label) self.task.class_nodes = np.concatenate((a_[:,np.newaxis], b_[:,np.newaxis]), axis=1) if self.logdir is not None: self._debug_semantic_maps(seed) self.task.reset_kwargs = {'sampling': self.task_params.semantic_task.sampling, 'class_nodes': self.task.class_nodes, 'dist_to_class': self.task.dist_to_class} if self.logdir is not None: self._debug_save_map_nodes(seed) def reset(self, rngs): rng = rngs[0]; rng_perturb = rngs[1]; nodes = self.task.nodes tp = self.task_params start_node_ids, goal_node_ids, dists, target_class = \ _nav_env_reset_helper(tp.type, rng, self.task.nodes, tp.batch_size, self.task.gtG, tp.max_dist, tp.num_steps, tp.num_goals, tp.data_augment, **(self.task.reset_kwargs)) start_nodes = [tuple(nodes[_,:]) for _ in start_node_ids] goal_nodes = [[tuple(nodes[_,:]) for _ in __] for __ in goal_node_ids] data_augment = tp.data_augment perturbs = _gen_perturbs(rng_perturb, tp.batch_size, (tp.num_steps+1)*tp.num_goals, data_augment.lr_flip, data_augment.delta_angle, data_augment.delta_xy, data_augment.structured) perturbs = np.array(perturbs) # batch x steps x 4 end_perturbs = perturbs[:,-(tp.num_goals):,:]*1 # fixed perturb for the goal. perturbs = perturbs[:,:-(tp.num_goals),:]*1 history = -np.ones((tp.batch_size, tp.num_steps*tp.num_goals), dtype=np.int32) self.episode = utils.Foo( start_nodes=start_nodes, start_node_ids=start_node_ids, goal_nodes=goal_nodes, goal_node_ids=goal_node_ids, dist_to_goal=dists, perturbs=perturbs, goal_perturbs=end_perturbs, history=history, target_class=target_class, history_frames=[]) return start_node_ids def take_action(self, current_node_ids, action, step_number): """In addition to returning the action, also returns the reward that the agent receives.""" goal_number = step_number / self.task_params.num_steps new_node_ids = GridWorld.take_action(self, current_node_ids, action) rewards = [] for i, n in enumerate(new_node_ids): reward = 0 if n == self.episode.goal_node_ids[goal_number][i]: reward = self.task_params.reward_at_goal reward = reward - self.task_params.reward_time_penalty rewards.append(reward) return new_node_ids, rewards def get_optimal_action(self, current_node_ids, step_number): """Returns the optimal action from the current node.""" goal_number = step_number / self.task_params.num_steps gtG = self.task.gtG a = np.zeros((len(current_node_ids), self.task_params.num_actions), dtype=np.int32) d_dict = self.episode.dist_to_goal[goal_number] for i, c in enumerate(current_node_ids): neigh = gtG.vertex(c).out_neighbours() neigh_edge = gtG.vertex(c).out_edges() ds = np.array([d_dict[i][int(x)] for x in neigh]) ds_min = np.min(ds) for i_, e in enumerate(neigh_edge): if ds[i_] == ds_min: _ = gtG.ep['action'][e] a[i, _] = 1 return a def get_targets(self, current_node_ids, step_number): """Returns the target actions from the current node.""" action = self.get_optimal_action(current_node_ids, step_number) action = np.expand_dims(action, axis=1) return vars(utils.Foo(action=action)) def get_targets_name(self): """Returns the list of names of the targets.""" return ['action'] def cleanup(self): self.episode = None class VisualNavigationEnv(NavigationEnv): """Class for doing visual navigation in environments. Functions for computing features on states, etc. """ def __init__(self, robot, env, task_params, category_list=None, building_name=None, flip=False, logdir=None, building_loader=None, r_obj=None): tt = utils.Timer() tt.tic() Building.__init__(self, building_name, robot, env, category_list, small=task_params.toy_problem, flip=flip, logdir=logdir, building_loader=building_loader) self.set_r_obj(r_obj) self.task_params = task_params self.task = None self.episode = None self._preprocess_for_task(self.task_params.building_seed) if hasattr(self.task_params, 'map_scales'): self.task.scaled_maps = resize_maps( self.traversible.astype(np.float32)*1, self.task_params.map_scales, self.task_params.map_resize_method) else: logging.fatal('VisualNavigationEnv does not support scale_f anymore.') self.task.readout_maps_scaled = resize_maps( self.traversible.astype(np.float32)*1, self.task_params.readout_maps_scales, self.task_params.map_resize_method) tt.toc(log_at=1, log_str='VisualNavigationEnv __init__: ') def get_weight(self): return self.task.nodes.shape[0] def get_common_data(self): goal_nodes = self.episode.goal_nodes start_nodes = self.episode.start_nodes perturbs = self.episode.perturbs goal_perturbs = self.episode.goal_perturbs target_class = self.episode.target_class goal_locs = []; rel_goal_locs = []; for i in range(len(goal_nodes)): end_nodes = goal_nodes[i] goal_loc, _, _, goal_theta = self.get_loc_axis( np.array(end_nodes), delta_theta=self.task.delta_theta, perturb=goal_perturbs[:,i,:]) # Compute the relative location to all goals from the starting location. loc, _, _, theta = self.get_loc_axis(np.array(start_nodes), delta_theta=self.task.delta_theta, perturb=perturbs[:,0,:]) r_goal, t_goal = _get_relative_goal_loc(goal_loc*1., loc, theta) rel_goal_loc = np.concatenate((r_goal*np.cos(t_goal), r_goal*np.sin(t_goal), np.cos(goal_theta-theta), np.sin(goal_theta-theta)), axis=1) rel_goal_locs.append(np.expand_dims(rel_goal_loc, axis=1)) goal_locs.append(np.expand_dims(goal_loc, axis=1)) map = self.traversible*1. maps = np.repeat(np.expand_dims(np.expand_dims(map, axis=0), axis=0), self.task_params.batch_size, axis=0)*1 if self.task_params.type[:14] == 'to_nearest_obj': for i in range(self.task_params.batch_size): maps[i,0,:,:] += 0.5*(self.task.class_maps_dilated[:,:,target_class[i]]) rel_goal_locs = np.concatenate(rel_goal_locs, axis=1) goal_locs = np.concatenate(goal_locs, axis=1) maps = np.expand_dims(maps, axis=-1) if self.task_params.type[:14] == 'to_nearest_obj': rel_goal_locs = np.zeros((self.task_params.batch_size, 1, len(self.task_params.semantic_task.class_map_names)), dtype=np.float32) goal_locs = np.zeros((self.task_params.batch_size, 1, 2), dtype=np.float32) for i in range(self.task_params.batch_size): t = target_class[i] rel_goal_locs[i,0,t] = 1. goal_locs[i,0,0] = t goal_locs[i,0,1] = np.NaN return vars(utils.Foo(orig_maps=maps, goal_loc=goal_locs, rel_goal_loc_at_start=rel_goal_locs)) def pre_common_data(self, inputs): return inputs def get_features(self, current_node_ids, step_number): task_params = self.task_params goal_number = step_number / self.task_params.num_steps end_nodes = self.task.nodes[self.episode.goal_node_ids[goal_number],:]*1 current_nodes = self.task.nodes[current_node_ids,:]*1 end_perturbs = self.episode.goal_perturbs[:,goal_number,:][:,np.newaxis,:] perturbs = self.episode.perturbs target_class = self.episode.target_class # Append to history. self.episode.history[:,step_number] = np.array(current_node_ids) # Render out the images from current node. outs = {} if self.task_params.outputs.images: imgs_all = [] imgs = self.render_nodes([tuple(x) for x in current_nodes], perturb=perturbs[:,step_number,:]) imgs_all.append(imgs) aux_delta_thetas = self.task_params.aux_delta_thetas for i in range(len(aux_delta_thetas)): imgs = self.render_nodes([tuple(x) for x in current_nodes], perturb=perturbs[:,step_number,:], aux_delta_theta=aux_delta_thetas[i]) imgs_all.append(imgs) imgs_all = np.array(imgs_all) # A x B x H x W x C imgs_all = np.transpose(imgs_all, axes=[1,0,2,3,4]) imgs_all = np.expand_dims(imgs_all, axis=1) # B x N x A x H x W x C if task_params.num_history_frames > 0: if step_number == 0: # Append the same frame 4 times for i in range(task_params.num_history_frames+1): self.episode.history_frames.insert(0, imgs_all*1.) self.episode.history_frames.insert(0, imgs_all) self.episode.history_frames.pop() imgs_all_with_history = np.concatenate(self.episode.history_frames, axis=2) else: imgs_all_with_history = imgs_all outs['imgs'] = imgs_all_with_history # B x N x A x H x W x C if self.task_params.outputs.node_ids: outs['node_ids'] = np.array(current_node_ids).reshape((-1,1,1)) outs['perturbs'] = np.expand_dims(perturbs[:,step_number, :]*1., axis=1) if self.task_params.outputs.analytical_counts: assert(self.task_params.modalities == ['depth']) d = image_pre(outs['imgs']*1., self.task_params.modalities) cm = get_camera_matrix(self.task_params.img_width, self.task_params.img_height, self.task_params.img_fov) XYZ = get_point_cloud_from_z(100./d[...,0], cm) XYZ = make_geocentric(XYZ*100., self.robot.sensor_height, self.robot.camera_elevation_degree) for i in range(len(self.task_params.analytical_counts.map_sizes)): non_linearity = self.task_params.analytical_counts.non_linearity[i] count, isvalid = bin_points(XYZ*1., map_size=self.task_params.analytical_counts.map_sizes[i], xy_resolution=self.task_params.analytical_counts.xy_resolution[i], z_bins=self.task_params.analytical_counts.z_bins[i]) assert(count.shape[2] == 1), 'only works for n_views equal to 1.' count = count[:,:,0,:,:,:] isvalid = isvalid[:,:,0,:,:,:] if non_linearity == 'none': None elif non_linearity == 'min10': count = np.minimum(count, 10.) elif non_linearity == 'sqrt': count = np.sqrt(count) else: logging.fatal('Undefined non_linearity.') outs['analytical_counts_{:d}'.format(i)] = count # Compute the goal location in the cordinate frame of the robot. if self.task_params.outputs.rel_goal_loc: if self.task_params.type[:14] != 'to_nearest_obj': loc, _, _, theta = self.get_loc_axis(current_nodes, delta_theta=self.task.delta_theta, perturb=perturbs[:,step_number,:]) goal_loc, _, _, goal_theta = self.get_loc_axis(end_nodes, delta_theta=self.task.delta_theta, perturb=end_perturbs[:,0,:]) r_goal, t_goal = _get_relative_goal_loc(goal_loc, loc, theta) rel_goal_loc = np.concatenate((r_goal*np.cos(t_goal), r_goal*np.sin(t_goal), np.cos(goal_theta-theta), np.sin(goal_theta-theta)), axis=1) outs['rel_goal_loc'] = np.expand_dims(rel_goal_loc, axis=1) elif self.task_params.type[:14] == 'to_nearest_obj': rel_goal_loc = np.zeros((self.task_params.batch_size, 1, len(self.task_params.semantic_task.class_map_names)), dtype=np.float32) for i in range(self.task_params.batch_size): t = target_class[i] rel_goal_loc[i,0,t] = 1. outs['rel_goal_loc'] = rel_goal_loc # Location on map to plot the trajectory during validation. if self.task_params.outputs.loc_on_map: loc, x_axis, y_axis, theta = self.get_loc_axis(current_nodes, delta_theta=self.task.delta_theta, perturb=perturbs[:,step_number,:]) outs['loc_on_map'] = np.expand_dims(loc, axis=1) # Compute gt_dist to goal if self.task_params.outputs.gt_dist_to_goal: gt_dist_to_goal = np.zeros((len(current_node_ids), 1), dtype=np.float32) for i, n in enumerate(current_node_ids): gt_dist_to_goal[i,0] = self.episode.dist_to_goal[goal_number][i][n] outs['gt_dist_to_goal'] = np.expand_dims(gt_dist_to_goal, axis=1) # Free space in front of you, map and goal as images. if self.task_params.outputs.ego_maps: loc, x_axis, y_axis, theta = self.get_loc_axis(current_nodes, delta_theta=self.task.delta_theta, perturb=perturbs[:,step_number,:]) maps = generate_egocentric_maps(self.task.scaled_maps, self.task_params.map_scales, self.task_params.map_crop_sizes, loc, x_axis, y_axis, theta) for i in range(len(self.task_params.map_scales)): outs['ego_maps_{:d}'.format(i)] = \ np.expand_dims(np.expand_dims(maps[i], axis=1), axis=-1) if self.task_params.outputs.readout_maps: loc, x_axis, y_axis, theta = self.get_loc_axis(current_nodes, delta_theta=self.task.delta_theta, perturb=perturbs[:,step_number,:]) maps = generate_egocentric_maps(self.task.readout_maps_scaled, self.task_params.readout_maps_scales, self.task_params.readout_maps_crop_sizes, loc, x_axis, y_axis, theta) for i in range(len(self.task_params.readout_maps_scales)): outs['readout_maps_{:d}'.format(i)] = \ np.expand_dims(np.expand_dims(maps[i], axis=1), axis=-1) # Images for the goal. if self.task_params.outputs.ego_goal_imgs: if self.task_params.type[:14] != 'to_nearest_obj': loc, x_axis, y_axis, theta = self.get_loc_axis(current_nodes, delta_theta=self.task.delta_theta, perturb=perturbs[:,step_number,:]) goal_loc, _, _, _ = self.get_loc_axis(end_nodes, delta_theta=self.task.delta_theta, perturb=end_perturbs[:,0,:]) rel_goal_orientation = np.mod( np.int32(current_nodes[:,2:] - end_nodes[:,2:]), self.task_params.n_ori) goal_dist, goal_theta = _get_relative_goal_loc(goal_loc, loc, theta) goals = generate_goal_images(self.task_params.map_scales, self.task_params.map_crop_sizes, self.task_params.n_ori, goal_dist, goal_theta, rel_goal_orientation) for i in range(len(self.task_params.map_scales)): outs['ego_goal_imgs_{:d}'.format(i)] = np.expand_dims(goals[i], axis=1) elif self.task_params.type[:14] == 'to_nearest_obj': for i in range(len(self.task_params.map_scales)): num_classes = len(self.task_params.semantic_task.class_map_names) outs['ego_goal_imgs_{:d}'.format(i)] = np.zeros((self.task_params.batch_size, 1, self.task_params.map_crop_sizes[i], self.task_params.map_crop_sizes[i], self.task_params.goal_channels)) for i in range(self.task_params.batch_size): t = target_class[i] for j in range(len(self.task_params.map_scales)): outs['ego_goal_imgs_{:d}'.format(j)][i,:,:,:,t] = 1. # Incremental locs and theta (for map warping), always in the original scale # of the map, the subequent steps in the tf code scale appropriately. # Scaling is done by just multiplying incremental_locs appropriately. if self.task_params.outputs.egomotion: if step_number == 0: # Zero Ego Motion incremental_locs = np.zeros((self.task_params.batch_size, 1, 2), dtype=np.float32) incremental_thetas = np.zeros((self.task_params.batch_size, 1, 1), dtype=np.float32) else: previous_nodes = self.task.nodes[self.episode.history[:,step_number-1], :]*1 loc, _, _, theta = self.get_loc_axis(current_nodes, delta_theta=self.task.delta_theta, perturb=perturbs[:,step_number,:]) previous_loc, _, _, previous_theta = self.get_loc_axis( previous_nodes, delta_theta=self.task.delta_theta, perturb=perturbs[:,step_number-1,:]) incremental_locs_ = np.reshape(loc-previous_loc, [self.task_params.batch_size, 1, -1]) t = -np.pi/2+np.reshape(theta*1, [self.task_params.batch_size, 1, -1]) incremental_locs = incremental_locs_*1 incremental_locs[:,:,0] = np.sum(incremental_locs_ * np.concatenate((np.cos(t), np.sin(t)), axis=-1), axis=-1) incremental_locs[:,:,1] = np.sum(incremental_locs_ * np.concatenate((np.cos(t+np.pi/2), np.sin(t+np.pi/2)), axis=-1), axis=-1) incremental_thetas = np.reshape(theta-previous_theta, [self.task_params.batch_size, 1, -1]) outs['incremental_locs'] = incremental_locs outs['incremental_thetas'] = incremental_thetas if self.task_params.outputs.visit_count: # Output the visit count for this state, how many times has the current # state been visited, and how far in the history was the last visit # (except this one) visit_count = np.zeros((self.task_params.batch_size, 1), dtype=np.int32) last_visit = -np.ones((self.task_params.batch_size, 1), dtype=np.int32) if step_number >= 1: h = self.episode.history[:,:(step_number)] visit_count[:,0] = np.sum(h == np.array(current_node_ids).reshape([-1,1]), axis=1) last_visit[:,0] = np.argmax(h[:,::-1] == np.array(current_node_ids).reshape([-1,1]), axis=1) + 1 last_visit[visit_count == 0] = -1 # -1 if not visited. outs['visit_count'] = np.expand_dims(visit_count, axis=1) outs['last_visit'] = np.expand_dims(last_visit, axis=1) return outs def get_features_name(self): f = [] if self.task_params.outputs.images: f.append('imgs') if self.task_params.outputs.rel_goal_loc: f.append('rel_goal_loc') if self.task_params.outputs.loc_on_map: f.append('loc_on_map') if self.task_params.outputs.gt_dist_to_goal: f.append('gt_dist_to_goal') if self.task_params.outputs.ego_maps: for i in range(len(self.task_params.map_scales)): f.append('ego_maps_{:d}'.format(i)) if self.task_params.outputs.readout_maps: for i in range(len(self.task_params.readout_maps_scales)): f.append('readout_maps_{:d}'.format(i)) if self.task_params.outputs.ego_goal_imgs: for i in range(len(self.task_params.map_scales)): f.append('ego_goal_imgs_{:d}'.format(i)) if self.task_params.outputs.egomotion: f.append('incremental_locs') f.append('incremental_thetas') if self.task_params.outputs.visit_count: f.append('visit_count') f.append('last_visit') if self.task_params.outputs.analytical_counts: for i in range(len(self.task_params.analytical_counts.map_sizes)): f.append('analytical_counts_{:d}'.format(i)) if self.task_params.outputs.node_ids: f.append('node_ids') f.append('perturbs') return f def pre_features(self, inputs): if self.task_params.outputs.images: inputs['imgs'] = image_pre(inputs['imgs'], self.task_params.modalities) return inputs class BuildingMultiplexer(): def __init__(self, args, task_number): params = vars(args) for k in params.keys(): setattr(self, k, params[k]) self.task_number = task_number self._pick_data(task_number) logging.info('Env Class: %s.', self.env_class) if self.task_params.task == 'planning': self._setup_planner() elif self.task_params.task == 'mapping': self._setup_mapper() elif self.task_params.task == 'map+plan': self._setup_mapper() else: logging.error('Undefined task: %s'.format(self.task_params.task)) def _pick_data(self, task_number): logging.error('Input Building Names: %s', self.building_names) self.flip = [np.mod(task_number / len(self.building_names), 2) == 1] id = np.mod(task_number, len(self.building_names)) self.building_names = [self.building_names[id]] self.task_params.building_seed = task_number logging.error('BuildingMultiplexer: Picked Building Name: %s', self.building_names) self.building_names = self.building_names[0].split('+') self.flip = [self.flip[0] for _ in self.building_names] logging.error('BuildingMultiplexer: Picked Building Name: %s', self.building_names) logging.error('BuildingMultiplexer: Flipping Buildings: %s', self.flip) logging.error('BuildingMultiplexer: Set building_seed: %d', self.task_params.building_seed) self.num_buildings = len(self.building_names) logging.error('BuildingMultiplexer: Num buildings: %d', self.num_buildings) def _setup_planner(self): # Load building env class. self.buildings = [] for i, building_name in enumerate(self.building_names): b = self.env_class(robot=self.robot, env=self.env, task_params=self.task_params, building_name=building_name, flip=self.flip[i], logdir=self.logdir, building_loader=self.dataset) self.buildings.append(b) def _setup_mapper(self): # Set up the renderer. cp = self.camera_param rgb_shader, d_shader = sru.get_shaders(cp.modalities) r_obj = SwiftshaderRenderer() r_obj.init_display(width=cp.width, height=cp.height, fov=cp.fov, z_near=cp.z_near, z_far=cp.z_far, rgb_shader=rgb_shader, d_shader=d_shader) self.r_obj = r_obj r_obj.clear_scene() # Load building env class. self.buildings = [] wt = [] for i, building_name in enumerate(self.building_names): b = self.env_class(robot=self.robot, env=self.env, task_params=self.task_params, building_name=building_name, flip=self.flip[i], logdir=self.logdir, building_loader=self.dataset, r_obj=r_obj) wt.append(b.get_weight()) b.load_building_into_scene() b.set_building_visibility(False) self.buildings.append(b) wt = np.array(wt).astype(np.float32) wt = wt / np.sum(wt+0.0001) self.building_sampling_weights = wt def sample_building(self, rng): if self.num_buildings == 1: building_id = rng.choice(range(len(self.building_names))) else: building_id = rng.choice(self.num_buildings, p=self.building_sampling_weights) b = self.buildings[building_id] instances = b._gen_rng(rng) self._building_id = building_id return self.buildings[building_id], instances def sample_env(self, rngs): rng = rngs[0]; if self.num_buildings == 1: building_id = rng.choice(range(len(self.building_names))) else: building_id = rng.choice(self.num_buildings, p=self.building_sampling_weights) return self.buildings[building_id] def pre(self, inputs): return self.buildings[self._building_id].pre(inputs) def __del__(self): self.r_obj.clear_scene() logging.error('Clearing scene.')