def get_all_objtype_in_room(room_type: str):
'''
Get all Object Types in a certern room type: kitchen, living_room, bedroom, bathroom
'''
if room_type == "kitchen":
room_start_index = 0
elif room_type == "living_room":
room_start_index = 200
elif room_type == "bedroom":
room_start_index = 300
else:
room_start_index = 400
all_obj_type = []
controller = Controller(scene="FloorPlan1",
renderInstanceSegmentation=True,
width=1080,
height=1080)
for i in range(1, 31):
controller.reset(scene="FloorPlan" + str(room_start_index + i))
event = controller.step("Done")
all_obj = get_all_objtype_in_event(event)
for objtype in all_obj:
if objtype not in all_obj_type:
all_obj_type.append(objtype)
controller.stop()
return all_obj_type
def run_load(dataset_path, memory, conn):
init_logging()
controller = Controller(
x_display='0.%d' % global_config.actor_gpu,
visibilityDistance=actor_config.visibilityDistance,
renderDepthImage=global_config.depth)
dataset = ActiveDataset(dataset_path, memory, controller, conn)
dataset.process()
controller.stop()
class Ai2Thor():
def __init__(self):
self.visualize = False
self.verbose = False
self.save_imgs = True
self.do_orbslam = False
self.do_depth_noise = False
self.makevideo = True
# st()
# these are all map names
a = np.arange(1, 30)
b = np.arange(201, 231)
c = np.arange(301, 331)
d = np.arange(401, 431)
abcd = np.hstack((a, b, c, d))
mapnames = []
for i in list(abcd):
mapname = 'FloorPlan' + str(i)
mapnames.append(mapname)
random.shuffle(mapnames)
self.mapnames = mapnames
self.num_episodes = 1 #len(self.mapnames)
self.ignore_classes = []
# classes to save
self.include_classes = [
'ShowerDoor', 'Cabinet', 'CounterTop', 'Sink', 'Towel',
'HandTowel', 'TowelHolder', 'SoapBar', 'ToiletPaper',
'ToiletPaperHanger', 'HandTowelHolder', 'SoapBottle', 'GarbageCan',
'Candle', 'ScrubBrush', 'Plunger', 'SinkBasin', 'Cloth',
'SprayBottle', 'Toilet', 'Faucet', 'ShowerHead', 'Box', 'Bed',
'Book', 'DeskLamp', 'BasketBall', 'Pen', 'Pillow', 'Pencil',
'CellPhone', 'KeyChain', 'Painting', 'CreditCard', 'AlarmClock',
'CD', 'Laptop', 'Drawer', 'SideTable', 'Chair', 'Blinds', 'Desk',
'Curtains', 'Dresser', 'Watch', 'Television', 'WateringCan',
'Newspaper', 'FloorLamp', 'RemoteControl', 'HousePlant', 'Statue',
'Ottoman', 'ArmChair', 'Sofa', 'DogBed', 'BaseballBat',
'TennisRacket', 'VacuumCleaner', 'Mug', 'ShelvingUnit', 'Shelf',
'StoveBurner', 'Apple', 'Lettuce', 'Bottle', 'Egg', 'Microwave',
'CoffeeMachine', 'Fork', 'Fridge', 'WineBottle', 'Spatula',
'Bread', 'Tomato', 'Pan', 'Cup', 'Pot', 'SaltShaker', 'Potato',
'PepperShaker', 'ButterKnife', 'StoveKnob', 'Toaster',
'DishSponge', 'Spoon', 'Plate', 'Knife', 'DiningTable', 'Bowl',
'LaundryHamper', 'Vase', 'Stool', 'CoffeeTable', 'Poster',
'Bathtub', 'TissueBox', 'Footstool', 'BathtubBasin',
'ShowerCurtain', 'TVStand', 'Boots', 'RoomDecor', 'PaperTowelRoll',
'Ladle', 'Kettle', 'Safe', 'GarbageBag', 'TeddyBear',
'TableTopDecor', 'Dumbbell', 'Desktop', 'AluminumFoil', 'Window'
]
self.small_classes = []
self.rot_interval = 5.0
self.radius_max = 3.5 #3 #1.75
self.radius_min = 1.0 #1.25
self.num_flat_views = 3
self.num_any_views = 7
self.num_views = 25
self.obj_per_scene = 10
# self.origin_quaternion = np.quaternion(1, 0, 0, 0)
# self.origin_rot_vector = quaternion.as_rotation_vector(self.origin_quaternion)
self.homepath = f'/home/nel/gsarch/aithor/data/test'
# self.basepath = '/home/nel/gsarch/replica_traj_bed'
if not os.path.exists(self.homepath):
os.mkdir(self.homepath)
else:
val = input("Delete homepath? [y/n]: ")
if val == 'y':
import shutil
shutil.rmtree(self.homepath)
os.mkdir(self.homepath)
else:
print("ENDING")
assert (False)
self.W = 256
self.H = 256
self.fov = 90
hfov = float(self.fov) * np.pi / 180.
self.pix_T_camX = np.array([[
(self.W / 2.) * 1 / np.tan(hfov / 2.), 0., 0., 0.
], [0., (self.H / 2.) * 1 / np.tan(hfov / 2.), 0., 0.], [0., 0., 1, 0],
[0., 0., 0, 1]])
self.pix_T_camX[0, 2] = self.W / 2.
self.pix_T_camX[1, 2] = self.H / 2.
self.run_episodes()
def run_episodes(self):
self.ep_idx = 0
# self.objects = []
for episode in range(self.num_episodes):
print("STARTING EPISODE ", episode)
mapname = self.mapnames[episode]
print("MAPNAME=", mapname)
self.controller = Controller(
scene=mapname,
gridSize=0.25,
width=self.W,
height=self.H,
fieldOfView=self.fov,
renderObjectImage=True,
renderDepthImage=True,
)
self.basepath = self.homepath + f"/{mapname}_{episode}"
print("BASEPATH: ", self.basepath)
# self.basepath = f"/hdd/ayushj/habitat_data/{mapname}_{episode}"
if not os.path.exists(self.basepath):
os.mkdir(self.basepath)
self.run()
self.controller.stop()
time.sleep(1)
self.ep_idx += 1
def save_datapoint(self, observations, data_path, viewnum, flat_view):
if self.verbose:
print("Print Sensor States.", self.agent.state.sensor_states)
rgb = observations["color_sensor"]
semantic = observations["semantic_sensor"]
# st()
depth = observations["depth_sensor"]
agent_pos = observations["positions"]
agent_rot = observations["rotations"]
# Assuming all sensors have same extrinsics
color_sensor_pos = observations["positions"]
color_sensor_rot = observations["rotations"]
#print("POS ", agent_pos)
#print("ROT ", color_sensor_rot)
object_list = observations['object_list']
# print(viewnum, agent_pos)
# print(agent_rot)
if False:
plt.imshow(rgb)
plt_name = f'/home/nel/gsarch/aithor/data/test/img_mask{viewnum}.png'
plt.savefig(plt_name)
save_data = {
'flat_view': flat_view,
'objects_info': object_list,
'rgb_camX': rgb,
'depth_camX': depth,
'sensor_pos': color_sensor_pos,
'sensor_rot': color_sensor_rot
}
with open(os.path.join(data_path, str(viewnum) + ".p"), 'wb') as f:
pickle.dump(save_data, f)
f.close()
def quat_from_angle_axis(self, theta: float,
axis: np.ndarray) -> np.quaternion:
r"""Creates a quaternion from angle axis format
:param theta: The angle to rotate about the axis by
:param axis: The axis to rotate about
:return: The quaternion
"""
axis = axis.astype(np.float)
axis /= np.linalg.norm(axis)
return quaternion.from_rotation_vector(theta * axis)
def run2(self):
event = self.controller.step('GetReachablePositions')
for obj in event.metadata['objects']:
if obj['objectType'] not in self.objects:
self.objects.append(obj['objectType'])
def run(self):
event = self.controller.step('GetReachablePositions')
self.nav_pts = event.metadata['reachablePositions']
self.nav_pts = np.array([list(d.values()) for d in self.nav_pts])
objects = np.random.choice(event.metadata['objects'],
self.obj_per_scene,
replace=False)
# objects = np.random.shuffle(event.metadata['objects'])
# for obj in event.metadata['objects']: #objects:
# print(obj['name'])
# objects = objects[0]
for obj in objects:
print("Object is ", obj['objectType'])
# if obj['name'] in ['Microwave_b200e0bc']:
# print(obj['name'])
# else:
# continue
# print(obj['name'])
if obj['objectType'] not in self.include_classes:
print("Continuing... Invalid Object")
continue
# Calculate distance to object center
# obj_center = np.array(list(obj['position'].values()))
obj_center = np.array(
list(obj['axisAlignedBoundingBox']['center'].values()))
# obj_center = np.array(list(obj['position'].values()))
#print(obj_center)
obj_center = np.expand_dims(obj_center, axis=0)
#print(obj_center)
distances = np.sqrt(np.sum((self.nav_pts - obj_center)**2, axis=1))
# Get points with r_min < dist < r_max
valid_pts = self.nav_pts[np.where(
(distances > self.radius_min) * (distances < self.radius_max))]
# if not valid_pts:
# continue
# plot valid points that we happen to select
# self.plot_navigable_points(valid_pts)
# Bin points based on angles [vertical_angle (10 deg/bin), horizontal_angle (10 deg/bin)]
valid_pts_shift = valid_pts - obj_center
dz = valid_pts_shift[:, 2]
dx = valid_pts_shift[:, 0]
dy = valid_pts_shift[:, 1]
# Get yaw for binning
valid_yaw = np.degrees(np.arctan2(dz, dx))
# # pitch calculation
# dxdz_norm = np.sqrt((dx * dx) + (dz * dz))
# valid_pitch = np.degrees(np.arctan2(dy,dxdz_norm))
# binning yaw around object
# nbins = 18
nbins = 18
bins = np.linspace(-180, 180, nbins + 1)
bin_yaw = np.digitize(valid_yaw, bins)
num_valid_bins = np.unique(bin_yaw).size
# spawns_per_bin = int(self.num_views / num_valid_bins) + 2
if False:
import matplotlib.cm as cm
colors = iter(cm.rainbow(np.linspace(0, 1, nbins)))
plt.figure(2)
plt.clf()
print(np.unique(bin_yaw))
for bi in range(nbins):
cur_bi = np.where(bin_yaw == (bi + 1))
points = valid_pts[cur_bi]
x_sample = points[:, 0]
z_sample = points[:, 2]
plt.plot(z_sample, x_sample, 'o', color=next(colors))
plt.plot(self.nav_pts[:, 2],
self.nav_pts[:, 0],
'x',
color='red')
plt.plot(obj_center[:, 2],
obj_center[:, 0],
'x',
color='black')
plt_name = '/home/nel/gsarch/aithor/data/valid.png'
plt.savefig(plt_name)
if num_valid_bins == 0:
continue
spawns_per_bin = int(self.num_views / num_valid_bins) + 2
# print(f'spawns_per_bin: {spawns_per_bin}')
action = "do_nothing"
episodes = []
valid_pts_selected = []
cnt = 0
for b in range(nbins):
# get all angle indices in the current bin range
# st()
inds_bin_cur = np.where(
bin_yaw == (b + 1)) # bins start 1 so need +1
inds_bin_cur = list(inds_bin_cur[0])
if len(inds_bin_cur) == 0:
continue
for s in range(spawns_per_bin):
observations = {}
if len(inds_bin_cur) == 0:
continue
rand_ind = np.random.randint(0, len(inds_bin_cur))
s_ind = inds_bin_cur.pop(rand_ind)
# st()
# s_ind = np.random.choice(inds_bin_cur)
#s_ind = inds_bin_cur[0][0]
pos_s = valid_pts[s_ind]
valid_pts_selected.append(pos_s)
# event = self.controller.step('TeleportFull', x=pos_s[0], y=pos_s[1], z=pos_s[2], rotation=dict(x=0.0, y=180.0, z=0.0), horizon=0.0)
# agent_pos = list(event.metadata['agent']['position'].values())
# print("Agent rot " , event.metadata['agent']['rotation'])
# print("Agent pos " , event.metadata['agent']['position'])
# print("Object center ", obj['axisAlignedBoundingBox']['center'])
# add height from center of agent to camera
pos_s[1] = pos_s[1] + 0.675
# YAW calculation - rotate to object
agent_to_obj = np.squeeze(
obj_center) - pos_s # + np.array([0.0, 0.675, 0.0]))
agent_local_forward = np.array([0, 0, -1.0]) # y, z, x
# agent_local_forward = np.array([-1.0, 0, 0]) # y, z, x
flat_to_obj = np.array(
[agent_to_obj[0], 0.0, agent_to_obj[2]])
flat_dist_to_obj = np.linalg.norm(flat_to_obj)
flat_to_obj /= flat_dist_to_obj
det = (flat_to_obj[0] * agent_local_forward[2] -
agent_local_forward[0] * flat_to_obj[2])
turn_angle = math.atan2(
det, np.dot(agent_local_forward, flat_to_obj))
# quat_yaw = self.quat_from_angle_axis(turn_angle, np.array([0, 1.0, 0]))
# turn_yaw = np.degrees(quaternion.as_euler_angles(quat_yaw)[1])
turn_yaw = np.degrees(turn_angle)
# print("Turn Yaw=", turn_yaw)
# print("turn1 beg ", turn_yaw)
# agent_pos = list(event.metadata['agent']['position'].values())
# dist_to_origin = np.sqrt(agent_pos[0]**2 + agent_pos[2]**2)
# dist_to_object = np.sqrt((agent_pos[0] - obj_center[0,0])**2 + (agent_pos[2] - obj_center[0,2])**2)
# print("Agent rot " , event.metadata['agent']['rotation'])
# p0 = agent_pos
# p1 = np.squeeze(obj_center)
# C = np.cross(p0, p1)
# D = np.dot(p0, p1)
# NP0 = np.linalg.norm(p0)
# if ~np.all(C==0): # check for colinearity
# Z = np.array([[0, -C[2], C[1]], [C[2], 0, -C[0]], [-C[1], C[0], 0]])
# R = (np.eye(3) + Z + Z**2 * (1-D)/(np.linalg.norm(C)**2)) / NP0**2 # rotation matrix
# else:
# R = np.sign(D) * (np.linalg.norm(p1) / NP0) # orientation and scaling
# quat_rot = quaternion.from_rotation_matrix(R)
# turns = np.degrees(quaternion.as_euler_angles(quat_rot))
# turn_yaw = turns[1]
# print("turn1 ", turn_yaw)
# print("turn0 ", turns[0])
# print("turn2 ", turns[2])
# if dist_to_origin > dist_to_object:
# print('or>ob')
# turn_yaw2 = np.degrees(np.cos(dist_to_object/dist_to_origin))
# else:
# print('or>ob')
# turn_yaw2 = np.degrees(np.cos(dist_to_origin/dist_to_object))
# print("TURN YAW2, ", turn_yaw2)
# Calculate Pitch from head to object
turn_pitch = -np.degrees(
math.atan2(agent_to_obj[1], flat_dist_to_obj))
# movement = "LookUp" if turn_pitch>0 else "LookDown"
# event = controller.step(movement, degrees=np.abs(turn_pitch))
event = self.controller.step('TeleportFull',
x=pos_s[0],
y=pos_s[1],
z=pos_s[2],
rotation=dict(x=0.0,
y=180.0 +
int(turn_yaw),
z=0.0),
horizon=int(turn_pitch))
# movement = "RotateRight" if turn_yaw>0 else "RotateLeft"
# event = self.controller.step(action='RotateRight', rotation=int(np.abs(turn_yaw)))
# movement = "LookDown" if turn_pitch>0 else "LookUp"
# event = self.controller.step(movement, degrees=np.abs(turn_pitch))
# print("Agent rot " , event.metadata['agent']['rotation'])
# angle_ranges = np.arange(0, 360, 15)
# angles_test = np.ones_like(angle_ranges) * 15
# for i in angle_ranges:
# movement = "RotateRight"
# event = self.controller.step(movement, degrees=15.0)
# rgb = event.frame
# if True:
# plt.imshow(rgb)
# plt_name = f'/home/nel/gsarch/aithor/data/img{i}.png'.format(i)
# plt.savefig(plt_name)
# print(event.metadata['agent']['position'])
# print(event.metadata['agent']['rotation'])
rgb = event.frame
rotation_euler_radians = np.radians(
np.array([
event.metadata['agent']['cameraHorizon'],
event.metadata['agent']['rotation']['y'], 0.0
]))
observations["positions"] = np.array(
list(event.metadata['agent']['position'].values())
) + np.array([0.0, 0.675, 0.0])
# print(observations["positions"])
# print(pos_s)
# observations["rotations"] = quaternion.from_euler_angles(np.radians(np.array(list(event.metadata['agent']['rotation'].values()))))
observations["rotations"] = quaternion.from_euler_angles(
rotation_euler_radians)
# print(observations["positions"])
observations["color_sensor"] = rgb
observations["depth_sensor"] = event.depth_frame
observations[
"semantic_sensor"] = event.instance_segmentation_frame
if False:
plt.imshow(rgb)
plt_name = f'/home/nel/gsarch/aithor/data/test/img_true{s}{b}.png'
plt.savefig(plt_name)
# print("Processed image #", cnt, " for object ", obj['objectType'])
semantic = event.instance_segmentation_frame
object_id_to_color = event.object_id_to_color
color_to_object_id = event.color_to_object_id
obj_ids = np.unique(semantic.reshape(
-1, semantic.shape[2]),
axis=0)
instance_masks = event.instance_masks
instance_detections2D = event.instance_detections2D
obj_metadata_IDs = []
for obj_m in event.metadata['objects']: #objects:
obj_metadata_IDs.append(obj_m['objectId'])
object_list = []
for obj_idx in range(obj_ids.shape[0]):
try:
obj_color = tuple(obj_ids[obj_idx])
object_id = color_to_object_id[obj_color]
except:
# print("Skipping ", object_id)
continue
if object_id not in obj_metadata_IDs:
# print("Skipping ", object_id)
continue
obj_meta_index = obj_metadata_IDs.index(object_id)
obj_meta = event.metadata['objects'][obj_meta_index]
obj_category_name = obj_meta['objectType']
# continue if not visible or not in include classes
if obj_category_name not in self.include_classes or not obj_meta[
'visible']:
continue
obj_instance_mask = instance_masks[object_id]
obj_instance_detection2D = instance_detections2D[
object_id] # [start_x, start_y, end_x, end_y]
obj_instance_detection2D = np.array([
obj_instance_detection2D[1],
obj_instance_detection2D[0],
obj_instance_detection2D[3],
obj_instance_detection2D[2]
]) # ymin, xmin, ymax, xmax
if False:
print(object_id)
plt.imshow(obj_instance_mask)
plt_name = f'/home/nel/gsarch/aithor/data/test/img_mask{s}.png'
plt.savefig(plt_name)
obj_center_axisAligned = np.array(
list(obj_meta['axisAlignedBoundingBox']
['center'].values()))
obj_size_axisAligned = np.array(
list(obj_meta['axisAlignedBoundingBox']
['size'].values()))
# print(obj_category_name)
if self.verbose:
print("Saved class name is : ", obj_category_name)
obj_data = {
'instance_id': object_id,
'category_id': object_id,
'category_name': obj_category_name,
'bbox_center': obj_center_axisAligned,
'bbox_size': obj_size_axisAligned,
'mask_2d': obj_instance_mask,
'box_2d': obj_instance_detection2D
}
# object_list.append(obj_instance)
object_list.append(obj_data)
observations["object_list"] = object_list
# check if object visible (make sure agent is not behind a wall)
obj_id = obj['objectId']
obj_id_to_color = object_id_to_color[obj_id]
# if np.sum(obj_ids==object_id_to_color[obj_id]) > 0:
if self.verbose:
print("episode is valid......")
episodes.append(observations)
cnt += 1
if len(episodes) >= self.num_views:
print(f'num episodes: {len(episodes)}')
data_folder = obj['name']
data_path = os.path.join(self.basepath, data_folder)
print("Saving to ", data_path)
os.mkdir(data_path)
# flat_obs = np.random.choice(episodes, self.num_views, replace=False)
rand_inds = np.sort(
np.random.choice(len(episodes),
self.num_views,
replace=False))
bool_inds = np.zeros(len(episodes), dtype=bool)
bool_inds[rand_inds] = True
flat_obs = np.array(episodes)[bool_inds]
flat_obs = list(flat_obs)
viewnum = 0
for obs in flat_obs:
self.save_datapoint(obs, data_path, viewnum, True)
viewnum += 1
else:
print("Not enough episodes:", len(episodes))
class RoboThorEnvironment:
"""Wrapper for the robo2thor controller providing additional functionality
and bookkeeping.
See [here](https://ai2thor.allenai.org/robothor/documentation) for comprehensive
documentation on RoboTHOR.
# Attributes
controller : The AI2THOR controller.
config : The AI2THOR controller configuration
"""
def __init__(self, **kwargs):
self.config = dict(
rotateStepDegrees=30.0,
visibilityDistance=1.0,
gridSize=0.25,
agentType="stochastic",
continuousMode=True,
snapToGrid=False,
agentMode="bot",
width=640,
height=480,
)
recursive_update(self.config, {**kwargs, "agentMode": "bot"})
self.controller = Controller(**self.config)
self.known_good_locations: Dict[str, Any] = {
self.scene_name: copy.deepcopy(self.currently_reachable_points)
}
assert len(self.known_good_locations[self.scene_name]) > 10
# onames = [o['objectId'] for o in self.last_event.metadata['objects']]
# removed = []
# for oname in onames:
# if 'Painting' in oname:
# self.controller.step("RemoveFromScene", objectId=oname)
# removed.append(oname)
# get_logger().info("Removed {} Paintings from {}".format(len(removed), self.scene_name))
# get_logger().warning("init to scene {} in pos {}".format(self.scene_name, self.agent_state()))
# npoints = len(self.currently_reachable_points)
# assert npoints > 100, "only {} reachable points after init".format(npoints)
self.grids: Dict[str, Tuple[Dict[str, np.array], int, int, int,
int]] = {}
self.initialize_grid()
def initialize_grid_dimensions(
self, reachable_points: Collection[Dict[str, float]]
) -> Tuple[int, int, int, int]:
"""Computes bounding box for reachable points quantized with the
current gridSize."""
points = {(
round(p["x"] / self.config["gridSize"]),
round(p["z"] / self.config["gridSize"]),
): p
for p in reachable_points}
assert len(reachable_points) == len(points)
xmin, xmax = min([p[0] for p in points]), max([p[0] for p in points])
zmin, zmax = min([p[1] for p in points]), max([p[1] for p in points])
return xmin, xmax, zmin, zmax
def access_grid(self, target: str) -> float:
"""Returns the geodesic distance from the quantized location of the
agent in the current scene's grid to the target object of given
type."""
if target not in self.grids[self.scene_name][0]:
xmin, xmax, zmin, zmax = self.grids[self.scene_name][1:5]
nx = xmax - xmin + 1
nz = zmax - zmin + 1
self.grids[self.scene_name][0][target] = -2 * np.ones(
(nx, nz), dtype=np.float64)
p = self.quantized_agent_state()
if self.grids[self.scene_name][0][target][p[0], p[1]] < -1.5:
corners = self.path_corners(target)
dist = self.path_corners_to_dist(corners)
if dist == float("inf"):
dist = -1.0 # -1.0 for unreachable
self.grids[self.scene_name][0][target][p[0], p[1]] = dist
return dist
return self.grids[self.scene_name][0][target][p[0], p[1]]
def initialize_grid(self) -> None:
"""Initializes grid for current scene if not already initialized."""
if self.scene_name in self.grids:
return
self.grids[self.scene_name] = ({}, ) + self.initialize_grid_dimensions(
self.known_good_locations[self.scene_name]) # type: ignore
def object_reachable(self, object_type: str) -> bool:
"""Determines whether a path can be computed from the discretized
current agent location to the target object of given type."""
return (self.access_grid(object_type) > -0.5
) # -1.0 for unreachable, 0.0 for end point
def point_reachable(self, xyz: Dict[str, float]) -> bool:
"""Determines whether a path can be computed from the current agent
location to the target point."""
return self.dist_to_point(
xyz) > -0.5 # -1.0 for unreachable, 0.0 for end point
def path_corners(
self, target: Union[str, Dict[str,
float]]) -> List[Dict[str, float]]:
"""Returns an array with a sequence of xyz dictionaries objects
representing the corners of the shortest path to the object of given
type or end point location."""
pose = self.agent_state()
position = {k: pose[k] for k in ["x", "y", "z"]}
# get_logger().debug("initial pos in path corners {} target {}".format(pose, target))
try:
if isinstance(target, str):
path = metrics.get_shortest_path_to_object_type(
self.controller,
target,
position,
{**pose["rotation"]} if "rotation" in pose else None,
)
else:
path = metrics.get_shortest_path_to_point(
self.controller, position, target)
except ValueError:
get_logger().debug("No path to object {} from {} in {}".format(
target, position, self.scene_name))
path = []
finally:
if isinstance(target, str):
self.controller.step("TeleportFull", **pose)
# pass
new_pose = self.agent_state()
try:
assert abs(new_pose["x"] - pose["x"]) < 1e-5, "wrong x"
assert abs(new_pose["y"] - pose["y"]) < 1e-5, "wrong y"
assert abs(new_pose["z"] - pose["z"]) < 1e-5, "wrong z"
assert (abs(new_pose["rotation"]["x"] - pose["rotation"]["x"])
< 1e-5), "wrong rotation x"
assert (abs(new_pose["rotation"]["y"] - pose["rotation"]["y"])
< 1e-5), "wrong rotation y"
assert (abs(new_pose["rotation"]["z"] - pose["rotation"]["z"])
< 1e-5), "wrong rotation z"
assert (abs((new_pose["horizon"] % 360) -
(pose["horizon"] % 360)) <
1e-5), "wrong horizon {} vs {}".format(
(new_pose["horizon"] % 360),
(pose["horizon"] % 360))
except Exception:
# get_logger().error("new_pose {} old_pose {} in {}".format(new_pose, pose, self.scene_name))
pass
# if abs((new_pose['horizon'] % 360) - (pose['horizon'] % 360)) > 1e-5:
# get_logger().debug("wrong horizon {} vs {} after path to object {} from {} in {}".format((new_pose['horizon'] % 360), (pose['horizon'] % 360), target, position, self.scene_name))
# else:
# get_logger().debug("correct horizon {} vs {} after path to object {} from {} in {}".format((new_pose['horizon'] % 360), (pose['horizon'] % 360), target, position, self.scene_name))
# assert abs((new_pose['horizon'] % 360) - (pose['horizon'] % 360)) < 1e-5, "wrong horizon {} vs {}".format((new_pose['horizon'] % 360), (pose['horizon'] % 360))
# # TODO: the agent will continue with a random horizon from here on
# target_horizon = (pose['horizon'] % 360) - (360 if (pose['horizon'] % 360) >= 180 else 0)
# new_pose = self.agent_state()['horizon']
# update_horizon = (new_pose % 360) - (360 if (new_pose % 360) >= 180 else 0)
# cond = abs(target_horizon - update_horizon) > 1e-5
# nmovements = 0
# while cond:
# cond = abs(target_horizon - update_horizon) > 1e-5 and target_horizon > update_horizon
# while cond:
# self.controller.step("LookDown")
# old = update_horizon
# new_pose = self.agent_state()['horizon']
# update_horizon = (new_pose % 360) - (360 if (new_pose % 360) >= 180 else 0)
# get_logger().debug("LookDown horizon {} -> {} ({})".format(old, update_horizon, target_horizon))
# nmovements += 1
# cond = abs(target_horizon - update_horizon) > 1e-5 and target_horizon > update_horizon
#
# cond = abs(target_horizon - update_horizon) > 1e-5 and target_horizon < update_horizon
# while cond:
# self.controller.step("LookUp")
# old = update_horizon
# new_pose = self.agent_state()['horizon']
# update_horizon = (new_pose % 360) - (360 if (new_pose % 360) >= 180 else 0)
# get_logger().debug("LookUp horizon {} -> {} ({})".format(old, update_horizon, target_horizon))
# nmovements += 1
# cond = abs(target_horizon - update_horizon) > 1e-5 and target_horizon < update_horizon
#
# cond = abs(target_horizon - update_horizon) > 1e-5
# get_logger().debug("nmovements {}".format(nmovements))
# new_pose = self.agent_state()
# assert abs((new_pose['horizon'] % 360) - (pose['horizon'] % 360)) < 1e-5, "wrong horizon {} vs {}".format((new_pose['horizon'] % 360), (pose['horizon'] % 360))
# try:
# assert abs((new_pose['horizon'] % 360) - (pose['horizon'] % 360)) < 1e-5, "wrong horizon {} vs {}".format((new_pose['horizon'] % 360), (pose['horizon'] % 360))
# except Exception:
# get_logger().error("wrong horizon {} vs {}".format((new_pose['horizon'] % 360), (pose['horizon'] % 360)))
# self.controller.step("TeleportFull", **pose)
# assert abs(
# (new_pose['horizon'] % 360) - (pose['horizon'] % 360)) < 1e-5, "wrong horizon {} vs {} after teleport full".format(
# (new_pose['horizon'] % 360), (pose['horizon'] % 360))
# # get_logger().debug("initial pos in path corners {} current pos {} path {}".format(pose, self.agent_state(), path))
return path
def path_corners_to_dist(self, corners: Sequence[Dict[str,
float]]) -> float:
"""Computes the distance covered by the given path described by its
corners."""
if len(corners) == 0:
return float("inf")
sum = 0.0
for it in range(1, len(corners)):
sum += math.sqrt((corners[it]["x"] - corners[it - 1]["x"])**2 +
(corners[it]["z"] - corners[it - 1]["z"])**2)
return sum
def quantized_agent_state(
self,
xz_subsampling: int = 1,
rot_subsampling: int = 1) -> Tuple[int, int, int]:
"""Quantizes agent location (x, z) to a (subsampled) position in a
fixed size grid derived from the initial set of reachable points; and
rotation (around y axis) as a (subsampled) discretized angle given the
current `rotateStepDegrees`."""
pose = self.agent_state()
p = {k: float(pose[k]) for k in ["x", "y", "z"]}
xmin, xmax, zmin, zmax = self.grids[self.scene_name][1:5]
x = int(np.clip(round(p["x"] / self.config["gridSize"]), xmin, xmax))
z = int(np.clip(round(p["z"] / self.config["gridSize"]), zmin, zmax))
rs = self.config["rotateStepDegrees"] * rot_subsampling
shifted = pose["rotation"]["y"] + rs / 2
normalized = shifted % 360.0
r = int(round(normalized / rs))
return (x - xmin) // xz_subsampling, (z - zmin) // xz_subsampling, r
def dist_to_object(self, object_type: str) -> float:
"""Minimal geodesic distance to object of given type from agent's
current location.
It might return -1.0 for unreachable targets.
"""
return self.access_grid(object_type)
def dist_to_point(self, xyz: Dict[str, float]) -> float:
"""Minimal geodesic distance to end point from agent's current
location.
It might return -1.0 for unreachable targets.
"""
corners = self.path_corners(xyz)
dist = self.path_corners_to_dist(corners)
if dist == float("inf"):
dist = -1.0 # -1.0 for unreachable
return dist
def agent_state(self) -> Dict:
"""Return agent position, rotation and horizon."""
agent_meta = self.last_event.metadata["agent"]
return {
**{k: float(v)
for k, v in agent_meta["position"].items()},
"rotation":
{k: float(v)
for k, v in agent_meta["rotation"].items()},
"horizon": round(float(agent_meta["cameraHorizon"]), 1),
}
def teleport(self,
pose: Dict[str, float],
rotation: Dict[str, float],
horizon: float = 0.0):
e = self.controller.step(
"TeleportFull",
x=pose["x"],
y=pose["y"],
z=pose["z"],
rotation=rotation,
horizon=horizon,
)
return e.metadata["lastActionSuccess"]
def reset(self, scene_name: str = None) -> None:
"""Resets scene to a known initial state."""
if scene_name is not None and scene_name != self.scene_name:
self.controller.reset(scene_name)
assert self.last_action_success, "Could not reset to new scene"
if scene_name not in self.known_good_locations:
self.known_good_locations[scene_name] = copy.deepcopy(
self.currently_reachable_points)
assert len(self.known_good_locations[scene_name]) > 10
# onames = [o['objectId'] for o in self.last_event.metadata['objects']]
# removed = []
# for oname in onames:
# if 'Painting' in oname:
# self.controller.step("RemoveFromScene", objectId=oname)
# removed.append(oname)
# get_logger().info("Removed {} Paintings from {}".format(len(removed), scene_name))
# else:
# assert (
# self.scene_name in self.known_good_locations
# ), "Resetting scene without known good location"
# get_logger().warning("Resetting {} to {}".format(self.scene_name, self.known_good_locations[self.scene_name]))
# self.controller.step("TeleportFull", **self.known_good_locations[self.scene_name])
# assert self.last_action_success, "Could not reset to known good location"
# npoints = len(self.currently_reachable_points)
# assert npoints > 100, "only {} reachable points after reset".format(npoints)
self.initialize_grid()
def randomize_agent_location(
self,
seed: int = None,
partial_position: Optional[Dict[str, float]] = None
) -> Dict[str, Union[Dict[str, float], float]]:
"""Teleports the agent to a random reachable location in the scene."""
if partial_position is None:
partial_position = {}
k = 0
state: Optional[Dict] = None
while k == 0 or (not self.last_action_success and k < 10):
# self.reset()
state = {
**self.random_reachable_state(seed=seed),
**partial_position
}
# get_logger().debug("picked target location {}".format(state))
self.controller.step("TeleportFull", **state)
k += 1
if not self.last_action_success:
get_logger().warning((
"Randomize agent location in scene {} and current random state {}"
" with seed {} and partial position {} failed in "
"10 attempts. Forcing the action.").format(
self.scene_name, state, seed, partial_position))
self.controller.step("TeleportFull", **state,
force_action=True) # type: ignore
assert self.last_action_success, "Force action failed with {}".format(
state)
# get_logger().debug("location after teleport full {}".format(self.agent_state()))
# self.controller.step("TeleportFull", **self.agent_state()) # TODO only for debug
# get_logger().debug("location after re-teleport full {}".format(self.agent_state()))
return self.agent_state()
def random_reachable_state(
self,
seed: Optional[int] = None
) -> Dict[str, Union[Dict[str, float], float]]:
"""Returns a random reachable location in the scene."""
if seed is not None:
random.seed(seed)
# xyz = random.choice(self.currently_reachable_points)
assert len(self.known_good_locations[self.scene_name]) > 10
xyz = copy.deepcopy(
random.choice(self.known_good_locations[self.scene_name]))
rotation = random.choice(
np.arange(0.0, 360.0, self.config["rotateStepDegrees"]))
horizon = 0.0 # random.choice([0.0, 30.0, 330.0])
return {
**{k: float(v)
for k, v in xyz.items()},
"rotation": {
"x": 0.0,
"y": float(rotation),
"z": 0.0
},
"horizon": float(horizon),
}
def known_good_locations_list(self):
return self.known_good_locations[self.scene_name]
@property
def currently_reachable_points(self) -> List[Dict[str, float]]:
"""List of {"x": x, "y": y, "z": z} locations in the scene that are
currently reachable."""
self.controller.step(action="GetReachablePositions")
return self.last_action_return
@property
def scene_name(self) -> str:
"""Current ai2thor scene."""
return self.controller.last_event.metadata["sceneName"].replace(
"_physics", "")
@property
def current_frame(self) -> np.ndarray:
"""Returns rgb image corresponding to the agent's egocentric view."""
return self.controller.last_event.frame
@property
def current_depth(self) -> np.ndarray:
"""Returns depth image corresponding to the agent's egocentric view."""
return self.controller.last_event.depth_frame
@property
def last_event(self) -> ai2thor.server.Event:
"""Last event returned by the controller."""
return self.controller.last_event
@property
def last_action(self) -> str:
"""Last action, as a string, taken by the agent."""
return self.controller.last_event.metadata["lastAction"]
@property
def last_action_success(self) -> bool:
"""Was the last action taken by the agent a success?"""
return self.controller.last_event.metadata["lastActionSuccess"]
@property
def last_action_return(self) -> Any:
"""Get the value returned by the last action (if applicable).
For an example of an action that returns a value, see
`"GetReachablePositions"`.
"""
return self.controller.last_event.metadata["actionReturn"]
def step(self, action_dict: Dict) -> ai2thor.server.Event:
"""Take a step in the ai2thor environment."""
return self.controller.step(**action_dict)
def stop(self):
"""Stops the ai2thor controller."""
try:
self.controller.stop()
except Exception as e:
get_logger().warning(str(e))
def all_objects(self) -> List[Dict[str, Any]]:
"""Return all object metadata."""
return self.controller.last_event.metadata["objects"]
def all_objects_with_properties(
self, properties: Dict[str, Any]) -> List[Dict[str, Any]]:
"""Find all objects with the given properties."""
objects = []
for o in self.all_objects():
satisfies_all = True
for k, v in properties.items():
if o[k] != v:
satisfies_all = False
break
if satisfies_all:
objects.append(o)
return objects
def visible_objects(self) -> List[Dict[str, Any]]:
"""Return all visible objects."""
return self.all_objects_with_properties({"visible": True})
class IThorEnvironment(object):
"""Wrapper for the ai2thor controller providing additional functionality
and bookkeeping.
See [here](https://ai2thor.allenai.org/documentation/installation) for comprehensive
documentation on AI2-THOR.
# Attributes
controller : The ai2thor controller.
"""
def __init__(
self,
x_display: Optional[str] = None,
docker_enabled: bool = False,
local_thor_build: Optional[str] = None,
visibility_distance: float = VISIBILITY_DISTANCE,
fov: float = FOV,
player_screen_width: int = 300,
player_screen_height: int = 300,
quality: str = "Very Low",
restrict_to_initially_reachable_points: bool = False,
make_agents_visible: bool = True,
object_open_speed: float = 1.0,
simplify_physics: bool = False,
) -> None:
"""Initializer.
# Parameters
x_display : The x display into which to launch ai2thor (possibly necessarily if you are running on a server
without an attached display).
docker_enabled : Whether or not to run thor in a docker container (useful on a server without an attached
display so that you don't have to start an x display).
local_thor_build : The path to a local build of ai2thor. This is probably not necessary for your use case
and can be safely ignored.
visibility_distance : The distance (in meters) at which objects, in the viewport of the agent,
are considered visible by ai2thor and will have their "visible" flag be set to `True` in the metadata.
fov : The agent's camera's field of view.
player_screen_width : The width resolution (in pixels) of the images returned by ai2thor.
player_screen_height : The height resolution (in pixels) of the images returned by ai2thor.
quality : The quality at which to render. Possible quality settings can be found in
`ai2thor._quality_settings.QUALITY_SETTINGS`.
restrict_to_initially_reachable_points : Whether or not to restrict the agent to locations in ai2thor
that were found to be (initially) reachable by the agent (i.e. reachable by the agent after resetting
the scene). This can be useful if you want to ensure there are only a fixed set of locations where the
agent can go.
make_agents_visible : Whether or not the agent should be visible. Most noticable when there are multiple agents
or when quality settings are high so that the agent casts a shadow.
object_open_speed : How quickly objects should be opened. High speeds mean faster simulation but also mean
that opening objects have a lot of kinetic energy and can, possibly, knock other objects away.
simplify_physics : Whether or not to simplify physics when applicable. Currently this only simplies object
interactions when opening drawers (when simplified, objects within a drawer do not slide around on
their own when the drawer is opened or closed, instead they are effectively glued down).
"""
self._start_player_screen_width = player_screen_width
self._start_player_screen_height = player_screen_height
self._local_thor_build = local_thor_build
self.x_display = x_display
self.controller: Optional[Controller] = None
self._started = False
self._quality = quality
self._initially_reachable_points: Optional[List[Dict]] = None
self._initially_reachable_points_set: Optional[Set[Tuple[
float, float]]] = None
self._move_mag: Optional[float] = None
self._grid_size: Optional[float] = None
self._visibility_distance = visibility_distance
self._fov = fov
self.restrict_to_initially_reachable_points = (
restrict_to_initially_reachable_points)
self.make_agents_visible = make_agents_visible
self.object_open_speed = object_open_speed
self._always_return_visible_range = False
self.simplify_physics = simplify_physics
self.start(None)
# noinspection PyTypeHints
self.controller.docker_enabled = docker_enabled # type: ignore
@property
def scene_name(self) -> str:
"""Current ai2thor scene."""
return self.controller.last_event.metadata["sceneName"]
@property
def current_frame(self) -> np.ndarray:
"""Returns rgb image corresponding to the agent's egocentric view."""
return self.controller.last_event.frame
@property
def last_event(self) -> ai2thor.server.Event:
"""Last event returned by the controller."""
return self.controller.last_event
@property
def started(self) -> bool:
"""Has the ai2thor controller been started."""
return self._started
@property
def last_action(self) -> str:
"""Last action, as a string, taken by the agent."""
return self.controller.last_event.metadata["lastAction"]
@last_action.setter
def last_action(self, value: str) -> None:
"""Set the last action taken by the agent.
Doing this is rewriting history, be careful.
"""
self.controller.last_event.metadata["lastAction"] = value
@property
def last_action_success(self) -> bool:
"""Was the last action taken by the agent a success?"""
return self.controller.last_event.metadata["lastActionSuccess"]
@last_action_success.setter
def last_action_success(self, value: bool) -> None:
"""Set whether or not the last action taken by the agent was a success.
Doing this is rewriting history, be careful.
"""
self.controller.last_event.metadata["lastActionSuccess"] = value
@property
def last_action_return(self) -> Any:
"""Get the value returned by the last action (if applicable).
For an example of an action that returns a value, see
`"GetReachablePositions"`.
"""
return self.controller.last_event.metadata["actionReturn"]
@last_action_return.setter
def last_action_return(self, value: Any) -> None:
"""Set the value returned by the last action.
Doing this is rewriting history, be careful.
"""
self.controller.last_event.metadata["actionReturn"] = value
def start(
self,
scene_name: Optional[str],
move_mag: float = 0.25,
**kwargs,
) -> None:
"""Starts the ai2thor controller if it was previously stopped.
After starting, `reset` will be called with the scene name and move magnitude.
# Parameters
scene_name : The scene to load.
move_mag : The amount of distance the agent moves in a single `MoveAhead` step.
kwargs : additional kwargs, passed to reset.
"""
if self._started:
raise RuntimeError(
"Trying to start the environment but it is already started.")
self.controller = Controller(
x_display=self.x_display,
width=self._start_player_screen_width,
height=self._start_player_screen_height,
local_executable_path=self._local_thor_build,
quality=self._quality,
server_class=ai2thor.fifo_server.FifoServer,
)
if (
self._start_player_screen_height,
self._start_player_screen_width,
) != self.current_frame.shape[:2]:
self.controller.step({
"action": "ChangeResolution",
"x": self._start_player_screen_width,
"y": self._start_player_screen_height,
})
self._started = True
self.reset(scene_name=scene_name, move_mag=move_mag, **kwargs)
def stop(self) -> None:
"""Stops the ai2thor controller."""
try:
self.controller.stop()
except Exception as e:
get_logger().warning(str(e))
finally:
self._started = False
def reset(
self,
scene_name: Optional[str],
move_mag: float = 0.25,
**kwargs,
):
"""Resets the ai2thor in a new scene.
Resets ai2thor into a new scene and initializes the scene/agents with
prespecified settings (e.g. move magnitude).
# Parameters
scene_name : The scene to load.
move_mag : The amount of distance the agent moves in a single `MoveAhead` step.
kwargs : additional kwargs, passed to the controller "Initialize" action.
"""
self._move_mag = move_mag
self._grid_size = self._move_mag
if scene_name is None:
scene_name = self.controller.last_event.metadata["sceneName"]
self.controller.reset(scene_name)
self.controller.step({
"action": "Initialize",
"gridSize": self._grid_size,
"visibilityDistance": self._visibility_distance,
"fov": self._fov,
"makeAgentsVisible": self.make_agents_visible,
"alwaysReturnVisibleRange": self._always_return_visible_range,
**kwargs,
})
if self.object_open_speed != 1.0:
self.controller.step({
"action": "ChangeOpenSpeed",
"x": self.object_open_speed
})
self._initially_reachable_points = None
self._initially_reachable_points_set = None
self.controller.step({"action": "GetReachablePositions"})
if not self.controller.last_event.metadata["lastActionSuccess"]:
get_logger().warning(
"Error when getting reachable points: {}".format(
self.controller.last_event.metadata["errorMessage"]))
self._initially_reachable_points = self.last_action_return
def teleport_agent_to(
self,
x: float,
y: float,
z: float,
rotation: float,
horizon: float,
standing: Optional[bool] = None,
force_action: bool = False,
only_initially_reachable: Optional[bool] = None,
verbose=True,
ignore_y_diffs=False,
) -> None:
"""Helper function teleporting the agent to a given location."""
if standing is None:
standing = self.last_event.metadata.get(
"isStanding",
self.last_event.metadata["agent"].get("isStanding"))
original_location = self.get_agent_location()
target = {"x": x, "y": y, "z": z}
if only_initially_reachable is None:
only_initially_reachable = self.restrict_to_initially_reachable_points
if only_initially_reachable:
reachable_points = self.initially_reachable_points
reachable = False
for p in reachable_points:
if self.position_dist(target, p,
ignore_y=ignore_y_diffs) < 0.01:
reachable = True
break
if not reachable:
self.last_action = "TeleportFull"
self.last_event.metadata[
"errorMessage"] = "Target position was not initially reachable."
self.last_action_success = False
return
self.controller.step(
dict(
action="TeleportFull",
x=x,
y=y,
z=z,
rotation={
"x": 0.0,
"y": rotation,
"z": 0.0
},
horizon=horizon,
standing=standing,
forceAction=force_action,
))
if not self.last_action_success:
agent_location = self.get_agent_location()
rot_diff = (agent_location["rotation"] -
original_location["rotation"]) % 360
new_old_dist = self.position_dist(original_location,
agent_location,
ignore_y=ignore_y_diffs)
if (self.position_dist(original_location,
agent_location,
ignore_y=ignore_y_diffs) > 1e-2
or min(rot_diff, 360 - rot_diff) > 1):
get_logger().warning(
"Teleportation FAILED but agent still moved (position_dist {}, rot diff {})"
" (\nprevious location\n{}\ncurrent_location\n{}\n)".
format(new_old_dist, rot_diff, original_location,
agent_location))
return
if force_action:
assert self.last_action_success
return
agent_location = self.get_agent_location()
rot_diff = (agent_location["rotation"] - rotation) % 360
if (self.position_dist(agent_location, target, ignore_y=ignore_y_diffs)
> 1e-2 or min(rot_diff, 360 - rot_diff) > 1):
if only_initially_reachable:
self._snap_agent_to_initially_reachable(verbose=False)
if verbose:
get_logger().warning(
"Teleportation did not place agent"
" precisely where desired in scene {}"
" (\ndesired\n{}\nactual\n{}\n)"
" perhaps due to grid snapping."
" Action is considered failed but agent may have moved.".
format(
self.scene_name,
{
"x": x,
"y": y,
"z": z,
"rotation": rotation,
"standing": standing,
"horizon": horizon,
},
agent_location,
))
self.last_action_success = False
return
def random_reachable_state(self, seed: int = None) -> Dict:
"""Returns a random reachable location in the scene."""
if seed is not None:
random.seed(seed)
xyz = random.choice(self.currently_reachable_points)
rotation = random.choice([0, 90, 180, 270])
horizon = random.choice([0, 30, 60, 330])
state = copy.copy(xyz)
state["rotation"] = rotation
state["horizon"] = horizon
return state
def randomize_agent_location(
self,
seed: int = None,
partial_position: Optional[Dict[str, float]] = None) -> Dict:
"""Teleports the agent to a random reachable location in the scene."""
if partial_position is None:
partial_position = {}
k = 0
state: Optional[Dict] = None
while k == 0 or (not self.last_action_success and k < 10):
state = self.random_reachable_state(seed=seed)
self.teleport_agent_to(**{**state, **partial_position})
k += 1
if not self.last_action_success:
get_logger().warning(
("Randomize agent location in scene {}"
" with seed {} and partial position {} failed in "
"10 attempts. Forcing the action.").format(
self.scene_name, seed, partial_position))
self.teleport_agent_to(**{
**state,
**partial_position
},
force_action=True) # type: ignore
assert self.last_action_success
assert state is not None
return state
def object_pixels_in_frame(self,
object_id: str,
hide_all: bool = True,
hide_transparent: bool = False) -> np.ndarray:
"""Return an mask for a given object in the agent's current view.
# Parameters
object_id : The id of the object.
hide_all : Whether or not to hide all other objects in the scene before getting the mask.
hide_transparent : Whether or not partially transparent objects are considered to occlude the object.
# Returns
A numpy array of the mask.
"""
# Emphasizing an object turns it magenta and hides all other objects
# from view, we can find where the hand object is on the screen by
# emphasizing it and then scanning across the image for the magenta pixels.
if hide_all:
self.step({"action": "EmphasizeObject", "objectId": object_id})
else:
self.step({"action": "MaskObject", "objectId": object_id})
if hide_transparent:
self.step({"action": "HideTranslucentObjects"})
# noinspection PyShadowingBuiltins
filter = np.array([[[255, 0, 255]]])
object_pixels = 1 * np.all(self.current_frame == filter, axis=2)
if hide_all:
self.step({"action": "UnemphasizeAll"})
else:
self.step({"action": "UnmaskObject", "objectId": object_id})
if hide_transparent:
self.step({"action": "UnhideAllObjects"})
return object_pixels
def object_pixels_on_grid(
self,
object_id: str,
grid_shape: Tuple[int, int],
hide_all: bool = True,
hide_transparent: bool = False,
) -> np.ndarray:
"""Like `object_pixels_in_frame` but counts object pixels in a
partitioning of the image."""
def partition(n, num_parts):
m = n // num_parts
parts = [m] * num_parts
num_extra = n % num_parts
for k in range(num_extra):
parts[k] += 1
return parts
object_pixels = self.object_pixels_in_frame(
object_id=object_id,
hide_all=hide_all,
hide_transparent=hide_transparent)
# Divide the current frame into a grid and count the number
# of hand object pixels in each of the grid squares
sums_in_blocks: List[List] = []
frame_shape = self.current_frame.shape[:2]
row_inds = np.cumsum([0] + partition(frame_shape[0], grid_shape[0]))
col_inds = np.cumsum([0] + partition(frame_shape[1], grid_shape[1]))
for i in range(len(row_inds) - 1):
sums_in_blocks.append([])
for j in range(len(col_inds) - 1):
sums_in_blocks[i].append(
np.sum(object_pixels[row_inds[i]:row_inds[i + 1],
col_inds[j]:col_inds[j + 1]]))
return np.array(sums_in_blocks, dtype=np.float32)
def object_in_hand(self):
"""Object metadata for the object in the agent's hand."""
inv_objs = self.last_event.metadata["inventoryObjects"]
if len(inv_objs) == 0:
return None
elif len(inv_objs) == 1:
return self.get_object_by_id(
self.last_event.metadata["inventoryObjects"][0]["objectId"])
else:
raise AttributeError("Must be <= 1 inventory objects.")
@property
def initially_reachable_points(self) -> List[Dict[str, float]]:
"""List of {"x": x, "y": y, "z": z} locations in the scene that were
reachable after initially resetting."""
assert self._initially_reachable_points is not None
return copy.deepcopy(self._initially_reachable_points) # type:ignore
@property
def initially_reachable_points_set(self) -> Set[Tuple[float, float]]:
"""Set of (x,z) locations in the scene that were reachable after
initially resetting."""
if self._initially_reachable_points_set is None:
self._initially_reachable_points_set = set()
for p in self.initially_reachable_points:
self._initially_reachable_points_set.add(
self._agent_location_to_tuple(p))
return self._initially_reachable_points_set
@property
def currently_reachable_points(self) -> List[Dict[str, float]]:
"""List of {"x": x, "y": y, "z": z} locations in the scene that are
currently reachable."""
self.step({"action": "GetReachablePositions"})
return self.last_event.metadata["actionReturn"] # type:ignore
def get_agent_location(self) -> Dict[str, Union[float, bool]]:
"""Gets agent's location."""
metadata = self.controller.last_event.metadata
location = {
"x":
metadata["agent"]["position"]["x"],
"y":
metadata["agent"]["position"]["y"],
"z":
metadata["agent"]["position"]["z"],
"rotation":
metadata["agent"]["rotation"]["y"],
"horizon":
metadata["agent"]["cameraHorizon"],
"standing":
metadata.get("isStanding", metadata["agent"].get("isStanding")),
}
return location
@staticmethod
def _agent_location_to_tuple(p: Dict[str, float]) -> Tuple[float, float]:
return round(p["x"], 2), round(p["z"], 2)
def _snap_agent_to_initially_reachable(self, verbose=True):
agent_location = self.get_agent_location()
end_location_tuple = self._agent_location_to_tuple(agent_location)
if end_location_tuple in self.initially_reachable_points_set:
return
agent_x = agent_location["x"]
agent_z = agent_location["z"]
closest_reachable_points = list(self.initially_reachable_points_set)
closest_reachable_points = sorted(
closest_reachable_points,
key=lambda xz: abs(xz[0] - agent_x) + abs(xz[1] - agent_z),
)
# In rare cases end_location_tuple might be not considered to be in self.initially_reachable_points_set
# even when it is, here we check for such cases.
if (math.sqrt(((np.array(closest_reachable_points[0]) -
np.array(end_location_tuple))**2).sum()) < 1e-6):
return
saved_last_action = self.last_action
saved_last_action_success = self.last_action_success
saved_last_action_return = self.last_action_return
saved_error_message = self.last_event.metadata["errorMessage"]
# Thor behaves weirdly when the agent gets off of the grid and you
# try to teleport the agent back to the closest grid location. To
# get around this we first teleport the agent to random location
# and then back to where it should be.
for point in self.initially_reachable_points:
if abs(agent_x - point["x"]) > 0.1 or abs(agent_z -
point["z"]) > 0.1:
self.teleport_agent_to(
rotation=0,
horizon=30,
**point,
only_initially_reachable=False,
verbose=False,
)
if self.last_action_success:
break
for p in closest_reachable_points:
self.teleport_agent_to(
**{
**agent_location, "x": p[0],
"z": p[1]
},
only_initially_reachable=False,
verbose=False,
)
if self.last_action_success:
break
teleport_forced = False
if not self.last_action_success:
self.teleport_agent_to(
**{
**agent_location,
"x": closest_reachable_points[0][0],
"z": closest_reachable_points[0][1],
},
force_action=True,
only_initially_reachable=False,
verbose=False,
)
teleport_forced = True
self.last_action = saved_last_action
self.last_action_success = saved_last_action_success
self.last_action_return = saved_last_action_return
self.last_event.metadata["errorMessage"] = saved_error_message
new_agent_location = self.get_agent_location()
if verbose:
get_logger().warning((
"In {}, at location (x,z)=({},{}) which is not in the set "
"of initially reachable points;"
" attempting to correct this: agent teleported to (x,z)=({},{}).\n"
"Teleportation {} forced.").format(
self.scene_name,
agent_x,
agent_z,
new_agent_location["x"],
new_agent_location["z"],
"was" if teleport_forced else "wasn't",
))
def step(
self,
action_dict: Optional[Dict[str, Union[str, int, float, Dict]]] = None,
**kwargs: Union[str, int, float, Dict],
) -> ai2thor.server.Event:
"""Take a step in the ai2thor environment."""
if action_dict is None:
action_dict = dict()
action_dict.update(kwargs)
action = cast(str, action_dict["action"])
skip_render = "renderImage" in action_dict and not action_dict[
"renderImage"]
last_frame: Optional[np.ndarray] = None
if skip_render:
last_frame = self.current_frame
if self.simplify_physics:
action_dict["simplifyOPhysics"] = True
if "Move" in action and "Hand" not in action: # type: ignore
action_dict = {
**action_dict,
"moveMagnitude": self._move_mag,
} # type: ignore
start_location = self.get_agent_location()
sr = self.controller.step(action_dict)
if self.restrict_to_initially_reachable_points:
end_location_tuple = self._agent_location_to_tuple(
self.get_agent_location())
if end_location_tuple not in self.initially_reachable_points_set:
self.teleport_agent_to(**start_location,
force_action=True) # type: ignore
self.last_action = action
self.last_action_success = False
self.last_event.metadata[
"errorMessage"] = "Moved to location outside of initially reachable points."
elif "RandomizeHideSeekObjects" in action:
last_position = self.get_agent_location()
self.controller.step(action_dict)
metadata = self.last_event.metadata
if self.position_dist(last_position,
self.get_agent_location()) > 0.001:
self.teleport_agent_to(**last_position,
force_action=True) # type: ignore
get_logger().warning(
"In scene {}, after randomization of hide and seek objects, agent moved."
.format(self.scene_name))
sr = self.controller.step({"action": "GetReachablePositions"})
self._initially_reachable_points = self.controller.last_event.metadata[
"actionReturn"]
self._initially_reachable_points_set = None
self.last_action = action
self.last_action_success = metadata["lastActionSuccess"]
self.controller.last_event.metadata["actionReturn"] = []
elif "RotateUniverse" in action:
sr = self.controller.step(action_dict)
metadata = self.last_event.metadata
if metadata["lastActionSuccess"]:
sr = self.controller.step({"action": "GetReachablePositions"})
self._initially_reachable_points = self.controller.last_event.metadata[
"actionReturn"]
self._initially_reachable_points_set = None
self.last_action = action
self.last_action_success = metadata["lastActionSuccess"]
self.controller.last_event.metadata["actionReturn"] = []
else:
sr = self.controller.step(action_dict)
if self.restrict_to_initially_reachable_points:
self._snap_agent_to_initially_reachable()
if skip_render:
assert last_frame is not None
self.last_event.frame = last_frame
return sr
@staticmethod
def position_dist(
p0: Mapping[str, Any],
p1: Mapping[str, Any],
ignore_y: bool = False,
l1_dist: bool = False,
) -> float:
"""Distance between two points of the form {"x": x, "y":y, "z":z"}."""
if l1_dist:
return (abs(p0["x"] - p1["x"]) +
(0 if ignore_y else abs(p0["y"] - p1["y"])) +
abs(p0["z"] - p1["z"]))
else:
return math.sqrt((p0["x"] - p1["x"])**2 +
(0 if ignore_y else (p0["y"] - p1["y"])**2) +
(p0["z"] - p1["z"])**2)
@staticmethod
def rotation_dist(a: Dict[str, float], b: Dict[str, float]):
"""Distance between rotations."""
def deg_dist(d0: float, d1: float):
dist = (d0 - d1) % 360
return min(dist, 360 - dist)
return sum(deg_dist(a[k], b[k]) for k in ["x", "y", "z"])
@staticmethod
def angle_between_rotations(a: Dict[str, float], b: Dict[str, float]):
return np.abs(
(180 / (2 * math.pi)) *
(Rotation.from_euler("xyz", [a[k] for k in "xyz"], degrees=True) *
Rotation.from_euler("xyz", [b[k] for k in "xyz"],
degrees=True).inv()).as_rotvec()).sum()
def closest_object_with_properties(
self, properties: Dict[str, Any]) -> Optional[Dict[str, Any]]:
"""Find the object closest to the agent that has the given
properties."""
agent_pos = self.controller.last_event.metadata["agent"]["position"]
min_dist = float("inf")
closest = None
for o in self.all_objects():
satisfies_all = True
for k, v in properties.items():
if o[k] != v:
satisfies_all = False
break
if satisfies_all:
d = self.position_dist(agent_pos, o["position"])
if d < min_dist:
min_dist = d
closest = o
return closest
def closest_visible_object_of_type(
self, object_type: str) -> Optional[Dict[str, Any]]:
"""Find the object closest to the agent that is visible and has the
given type."""
properties = {"visible": True, "objectType": object_type}
return self.closest_object_with_properties(properties)
def closest_object_of_type(self,
object_type: str) -> Optional[Dict[str, Any]]:
"""Find the object closest to the agent that has the given type."""
properties = {"objectType": object_type}
return self.closest_object_with_properties(properties)
def closest_reachable_point_to_position(
self, position: Dict[str,
float]) -> Tuple[Dict[str, float], float]:
"""Of all reachable positions, find the one that is closest to the
given location."""
target = np.array([position["x"], position["z"]])
min_dist = float("inf")
closest_point = None
for pt in self.initially_reachable_points:
dist = np.linalg.norm(target - np.array([pt["x"], pt["z"]]))
if dist < min_dist:
closest_point = pt
min_dist = dist
if min_dist < 1e-3:
break
assert closest_point is not None
return closest_point, min_dist
@staticmethod
def _angle_from_to(a_from: float, a_to: float) -> float:
a_from = a_from % 360
a_to = a_to % 360
min_rot = min(a_from, a_to)
max_rot = max(a_from, a_to)
rot_across_0 = (360 - max_rot) + min_rot
rot_not_across_0 = max_rot - min_rot
rot_err = min(rot_across_0, rot_not_across_0)
if rot_across_0 == rot_err:
rot_err *= -1 if a_to > a_from else 1
else:
rot_err *= 1 if a_to > a_from else -1
return rot_err
def agent_xz_to_scene_xz(self, agent_xz: Dict[str,
float]) -> Dict[str, float]:
agent_pos = self.get_agent_location()
x_rel_agent = agent_xz["x"]
z_rel_agent = agent_xz["z"]
scene_x = agent_pos["x"]
scene_z = agent_pos["z"]
rotation = agent_pos["rotation"]
if abs(rotation) < 1e-5:
scene_x += x_rel_agent
scene_z += z_rel_agent
elif abs(rotation - 90) < 1e-5:
scene_x += z_rel_agent
scene_z += -x_rel_agent
elif abs(rotation - 180) < 1e-5:
scene_x += -x_rel_agent
scene_z += -z_rel_agent
elif abs(rotation - 270) < 1e-5:
scene_x += -z_rel_agent
scene_z += x_rel_agent
else:
raise Exception("Rotation must be one of 0, 90, 180, or 270.")
return {"x": scene_x, "z": scene_z}
def scene_xz_to_agent_xz(self, scene_xz: Dict[str,
float]) -> Dict[str, float]:
agent_pos = self.get_agent_location()
x_err = scene_xz["x"] - agent_pos["x"]
z_err = scene_xz["z"] - agent_pos["z"]
rotation = agent_pos["rotation"]
if abs(rotation) < 1e-5:
agent_x = x_err
agent_z = z_err
elif abs(rotation - 90) < 1e-5:
agent_x = -z_err
agent_z = x_err
elif abs(rotation - 180) < 1e-5:
agent_x = -x_err
agent_z = -z_err
elif abs(rotation - 270) < 1e-5:
agent_x = z_err
agent_z = -x_err
else:
raise Exception("Rotation must be one of 0, 90, 180, or 270.")
return {"x": agent_x, "z": agent_z}
def all_objects(self) -> List[Dict[str, Any]]:
"""Return all object metadata."""
return self.controller.last_event.metadata["objects"]
def all_objects_with_properties(
self, properties: Dict[str, Any]) -> List[Dict[str, Any]]:
"""Find all objects with the given properties."""
objects = []
for o in self.all_objects():
satisfies_all = True
for k, v in properties.items():
if o[k] != v:
satisfies_all = False
break
if satisfies_all:
objects.append(o)
return objects
def visible_objects(self) -> List[Dict[str, Any]]:
"""Return all visible objects."""
return self.all_objects_with_properties({"visible": True})
def get_object_by_id(self, object_id: str) -> Optional[Dict[str, Any]]:
for o in self.last_event.metadata["objects"]:
if o["objectId"] == object_id:
return o
return None
###
# Following is used for computing shortest paths between states
###
_CACHED_GRAPHS: Dict[str, nx.DiGraph] = {}
GRAPH_ACTIONS_SET = {
"LookUp", "LookDown", "RotateLeft", "RotateRight", "MoveAhead"
}
def reachable_points_with_rotations_and_horizons(self):
self.controller.step({"action": "GetReachablePositions"})
assert self.last_action_success
points_slim = self.last_event.metadata["actionReturn"]
points = []
for r in [0, 90, 180, 270]:
for horizon in [-30, 0, 30, 60]:
for p in points_slim:
p = copy.copy(p)
p["rotation"] = r
p["horizon"] = horizon
points.append(p)
return points
@staticmethod
def location_for_key(key, y_value=0.0):
x, z, rot, hor = key
loc = dict(x=x, y=y_value, z=z, rotation=rot, horizon=hor)
return loc
@staticmethod
def get_key(input_dict: Dict[str, Any]) -> Tuple[float, float, int, int]:
if "x" in input_dict:
x = input_dict["x"]
z = input_dict["z"]
rot = input_dict["rotation"]
hor = input_dict["horizon"]
else:
x = input_dict["position"]["x"]
z = input_dict["position"]["z"]
rot = input_dict["rotation"]["y"]
hor = input_dict["cameraHorizon"]
return (
round(x, 2),
round(z, 2),
round_to_factor(rot, 90) % 360,
round_to_factor(hor, 30) % 360,
)
def update_graph_with_failed_action(self, failed_action: str):
if (self.scene_name not in self._CACHED_GRAPHS
or failed_action not in self.GRAPH_ACTIONS_SET):
return
source_key = self.get_key(self.last_event.metadata["agent"])
self._check_contains_key(source_key)
edge_dict = self.graph[source_key]
to_remove_key = None
for target_key in self.graph[source_key]:
if edge_dict[target_key]["action"] == failed_action:
to_remove_key = target_key
break
if to_remove_key is not None:
self.graph.remove_edge(source_key, to_remove_key)
def _add_from_to_edge(
self,
g: nx.DiGraph,
s: Tuple[float, float, int, int],
t: Tuple[float, float, int, int],
):
def ae(x, y):
return abs(x - y) < 0.001
s_x, s_z, s_rot, s_hor = s
t_x, t_z, t_rot, t_hor = t
dist = round(math.sqrt((s_x - t_x)**2 + (s_z - t_z)**2), 2)
angle_dist = (round_to_factor(t_rot - s_rot, 90) % 360) // 90
horz_dist = (round_to_factor(t_hor - s_hor, 30) % 360) // 30
# If source and target differ by more than one action, continue
if sum(x != 0 for x in [dist, angle_dist, horz_dist]) != 1:
return
grid_size = self._grid_size
action = None
if angle_dist != 0:
if angle_dist == 1:
action = "RotateRight"
elif angle_dist == 3:
action = "RotateLeft"
elif horz_dist != 0:
if horz_dist == 11:
action = "LookUp"
elif horz_dist == 1:
action = "LookDown"
elif ae(dist, grid_size):
if ((s_rot == 0 and ae(t_z - s_z, grid_size))
or (s_rot == 90 and ae(t_x - s_x, grid_size))
or (s_rot == 180 and ae(t_z - s_z, -grid_size))
or (s_rot == 270 and ae(t_x - s_x, -grid_size))):
g.add_edge(s, t, action="MoveAhead")
if action is not None:
g.add_edge(s, t, action=action)
@functools.lru_cache(1)
def possible_neighbor_offsets(
self) -> Tuple[Tuple[float, float, int, int], ...]:
grid_size = round(self._grid_size, 2)
offsets = []
for rot_diff in [-90, 0, 90]:
for horz_diff in [-30, 0, 30, 60]:
for x_diff in [-grid_size, 0, grid_size]:
for z_diff in [-grid_size, 0, grid_size]:
if (rot_diff != 0) + (horz_diff != 0) + (
x_diff != 0) + (z_diff != 0) == 1:
offsets.append(
(x_diff, z_diff, rot_diff, horz_diff))
return tuple(offsets)
def _add_node_to_graph(self, graph: nx.DiGraph, s: Tuple[float, float, int,
int]):
if s in graph:
return
existing_nodes = set(graph.nodes())
graph.add_node(s)
for o in self.possible_neighbor_offsets():
t = (s[0] + o[0], s[1] + o[1], s[2] + o[2], s[3] + o[3])
if t in existing_nodes:
self._add_from_to_edge(graph, s, t)
self._add_from_to_edge(graph, t, s)
@property
def graph(self):
if self.scene_name not in self._CACHED_GRAPHS:
g = nx.DiGraph()
points = self.reachable_points_with_rotations_and_horizons()
for p in points:
self._add_node_to_graph(g, self.get_key(p))
self._CACHED_GRAPHS[self.scene_name] = g
return self._CACHED_GRAPHS[self.scene_name]
@graph.setter
def graph(self, g):
self._CACHED_GRAPHS[self.scene_name] = g
def _check_contains_key(self,
key: Tuple[float, float, int, int],
add_if_not=True):
if key not in self.graph:
get_logger().warning(
"{} was not in the graph for scene {}.".format(
key, self.scene_name))
if add_if_not:
self._add_node_to_graph(self.graph, key)
def shortest_state_path(self, source_state_key, goal_state_key):
self._check_contains_key(source_state_key)
self._check_contains_key(goal_state_key)
# noinspection PyBroadException
try:
path = nx.shortest_path(self.graph, source_state_key,
goal_state_key)
return path
except Exception as _:
return None
def action_transitioning_between_keys(self, s, t):
self._check_contains_key(s)
self._check_contains_key(t)
if self.graph.has_edge(s, t):
return self.graph.get_edge_data(s, t)["action"]
else:
return None
def shortest_path_next_state(self, source_state_key, goal_state_key):
self._check_contains_key(source_state_key)
self._check_contains_key(goal_state_key)
if source_state_key == goal_state_key:
raise RuntimeError(
"called next state on the same source and goal state")
state_path = self.shortest_state_path(source_state_key, goal_state_key)
return state_path[1]
def shortest_path_next_action(self, source_state_key, goal_state_key):
self._check_contains_key(source_state_key)
self._check_contains_key(goal_state_key)
next_state_key = self.shortest_path_next_state(source_state_key,
goal_state_key)
return self.graph.get_edge_data(source_state_key,
next_state_key)["action"]
def shortest_path_length(self, source_state_key, goal_state_key):
self._check_contains_key(source_state_key)
self._check_contains_key(goal_state_key)
try:
return nx.shortest_path_length(self.graph, source_state_key,
goal_state_key)
except nx.NetworkXNoPath as _:
return float("inf")
class RoboThorEnvironment:
"""Wrapper for the robo2thor controller providing additional functionality
and bookkeeping.
See [here](https://ai2thor.allenai.org/robothor/documentation) for comprehensive
documentation on RoboTHOR.
# Attributes
controller : The AI2-THOR controller.
config : The AI2-THOR controller configuration
"""
def __init__(self, all_metadata_available: bool = True, **kwargs):
self.config = dict(
rotateStepDegrees=30.0,
visibilityDistance=1.0,
gridSize=0.25,
continuousMode=True,
snapToGrid=False,
agentMode="locobot",
width=640,
height=480,
agentCount=1,
server_class=FifoServer,
)
if "agentCount" in kwargs:
assert kwargs["agentCount"] > 0
kwargs["agentMode"] = kwargs.get("agentMode", "locobot")
if kwargs["agentMode"] not in ["bot", "locobot"]:
warnings.warn(f"The RoboTHOR environment has not been tested using"
f" an agent of mode '{kwargs['agentMode']}'.")
recursive_update(self.config, kwargs)
self.controller = Controller(**self.config, )
self.all_metadata_available = all_metadata_available
self.scene_to_reachable_positions: Optional[Dict[str, Any]] = None
self.distance_cache: Optional[DynamicDistanceCache] = None
if self.all_metadata_available:
self.scene_to_reachable_positions = {
self.scene_name: copy.deepcopy(self.currently_reachable_points)
}
assert len(self.scene_to_reachable_positions[self.scene_name]) > 10
self.distance_cache = DynamicDistanceCache(rounding=1)
self.agent_count = self.config["agentCount"]
self._extra_teleport_kwargs: Dict[str, Any] = {
} # Used for backwards compatability with the teleport action
def initialize_grid_dimensions(
self, reachable_points: Collection[Dict[str, float]]
) -> Tuple[int, int, int, int]:
"""Computes bounding box for reachable points quantized with the
current gridSize."""
points = {(
round(p["x"] / self.config["gridSize"]),
round(p["z"] / self.config["gridSize"]),
): p
for p in reachable_points}
assert len(reachable_points) == len(points)
xmin, xmax = min([p[0] for p in points]), max([p[0] for p in points])
zmin, zmax = min([p[1] for p in points]), max([p[1] for p in points])
return xmin, xmax, zmin, zmax
def set_object_filter(self, object_ids: List[str]):
self.controller.step("SetObjectFilter",
objectIds=object_ids,
renderImage=False)
def reset_object_filter(self):
self.controller.step("ResetObjectFilter", renderImage=False)
def path_from_point_to_object_type(
self, point: Dict[str, float], object_type: str,
allowed_error: float) -> Optional[List[Dict[str, float]]]:
event = self.controller.step(
action="GetShortestPath",
objectType=object_type,
position=point,
allowedError=allowed_error,
)
if event.metadata["lastActionSuccess"]:
return event.metadata["actionReturn"]["corners"]
else:
get_logger().debug(
"Failed to find path for {} in {}. Start point {}, agent state {}."
.format(
object_type,
self.controller.last_event.metadata["sceneName"],
point,
self.agent_state(),
))
return None
def distance_from_point_to_object_type(self, point: Dict[str, float],
object_type: str,
allowed_error: float) -> float:
"""Minimal geodesic distance from a point to an object of the given
type.
It might return -1.0 for unreachable targets.
"""
path = self.path_from_point_to_object_type(point, object_type,
allowed_error)
if path:
# Because `allowed_error != 0` means that the path returned above might not start
# at `point`, we explicitly add any offset there is.
s_dist = math.sqrt((point["x"] - path[0]["x"])**2 +
(point["z"] - path[0]["z"])**2)
return metrics.path_distance(path) + s_dist
return -1.0
def distance_to_object_type(self,
object_type: str,
agent_id: int = 0) -> float:
"""Minimal geodesic distance to object of given type from agent's
current location.
It might return -1.0 for unreachable targets.
"""
assert 0 <= agent_id < self.agent_count
assert (
self.all_metadata_available
), "`distance_to_object_type` cannot be called when `self.all_metadata_available` is `False`."
def retry_dist(position: Dict[str, float], object_type: str):
allowed_error = 0.05
debug_log = ""
d = -1.0
while allowed_error < 2.5:
d = self.distance_from_point_to_object_type(
position, object_type, allowed_error)
if d < 0:
debug_log = (
f"In scene {self.scene_name}, could not find a path from {position} to {object_type} with"
f" {allowed_error} error tolerance. Increasing this tolerance to"
f" {2 * allowed_error} any trying again.")
allowed_error *= 2
else:
break
if d < 0:
get_logger().warning(
f"In scene {self.scene_name}, could not find a path from {position} to {object_type}"
f" with {allowed_error} error tolerance. Returning a distance of -1."
)
elif debug_log != "":
get_logger().debug(debug_log)
return d
return self.distance_cache.find_distance(
self.scene_name,
self.controller.last_event.events[agent_id].metadata["agent"]
["position"],
object_type,
retry_dist,
)
def path_from_point_to_point(
self, position: Dict[str, float], target: Dict[str, float],
allowedError: float) -> Optional[List[Dict[str, float]]]:
try:
return self.controller.step(
action="GetShortestPathToPoint",
position=position,
x=target["x"],
y=target["y"],
z=target["z"],
allowedError=allowedError,
).metadata["actionReturn"]["corners"]
except Exception:
get_logger().debug(
"Failed to find path for {} in {}. Start point {}, agent state {}."
.format(
target,
self.controller.last_event.metadata["sceneName"],
position,
self.agent_state(),
))
return None
def distance_from_point_to_point(self, position: Dict[str, float],
target: Dict[str, float],
allowed_error: float) -> float:
path = self.path_from_point_to_point(position, target, allowed_error)
if path:
# Because `allowed_error != 0` means that the path returned above might not start
# or end exactly at the position/target points, we explictly add any offset there is.
s_dist = math.sqrt((position["x"] - path[0]["x"])**2 +
(position["z"] - path[0]["z"])**2)
t_dist = math.sqrt((target["x"] - path[-1]["x"])**2 +
(target["z"] - path[-1]["z"])**2)
return metrics.path_distance(path) + s_dist + t_dist
return -1.0
def distance_to_point(self,
target: Dict[str, float],
agent_id: int = 0) -> float:
"""Minimal geodesic distance to end point from agent's current
location.
It might return -1.0 for unreachable targets.
"""
assert 0 <= agent_id < self.agent_count
assert (
self.all_metadata_available
), "`distance_to_object_type` cannot be called when `self.all_metadata_available` is `False`."
def retry_dist(position: Dict[str, float], target: Dict[str, float]):
allowed_error = 0.05
debug_log = ""
d = -1.0
while allowed_error < 2.5:
d = self.distance_from_point_to_point(position, target,
allowed_error)
if d < 0:
debug_log = (
f"In scene {self.scene_name}, could not find a path from {position} to {target} with"
f" {allowed_error} error tolerance. Increasing this tolerance to"
f" {2 * allowed_error} any trying again.")
allowed_error *= 2
else:
break
if d < 0:
get_logger().warning(
f"In scene {self.scene_name}, could not find a path from {position} to {target}"
f" with {allowed_error} error tolerance. Returning a distance of -1."
)
elif debug_log != "":
get_logger().debug(debug_log)
return d
return self.distance_cache.find_distance(
self.scene_name,
self.controller.last_event.events[agent_id].metadata["agent"]
["position"],
target,
retry_dist,
)
def agent_state(self, agent_id: int = 0) -> Dict:
"""Return agent position, rotation and horizon."""
assert 0 <= agent_id < self.agent_count
agent_meta = self.last_event.events[agent_id].metadata["agent"]
return {
**{k: float(v)
for k, v in agent_meta["position"].items()},
"rotation":
{k: float(v)
for k, v in agent_meta["rotation"].items()},
"horizon": round(float(agent_meta["cameraHorizon"]), 1),
}
def teleport(
self,
pose: Dict[str, float],
rotation: Dict[str, float],
horizon: float = 0.0,
agent_id: int = 0,
):
assert 0 <= agent_id < self.agent_count
try:
e = self.controller.step(
action="TeleportFull",
x=pose["x"],
y=pose["y"],
z=pose["z"],
rotation=rotation,
horizon=horizon,
agentId=agent_id,
**self._extra_teleport_kwargs,
)
except ValueError as e:
if len(self._extra_teleport_kwargs) == 0:
self._extra_teleport_kwargs["standing"] = True
else:
raise e
return self.teleport(pose=pose,
rotation=rotation,
horizon=horizon,
agent_id=agent_id)
return e.metadata["lastActionSuccess"]
def reset(self,
scene_name: str = None,
filtered_objects: Optional[List[str]] = None) -> None:
"""Resets scene to a known initial state."""
if scene_name is not None and scene_name != self.scene_name:
self.controller.reset(scene_name)
assert self.last_action_success, "Could not reset to new scene"
if (self.all_metadata_available
and scene_name not in self.scene_to_reachable_positions):
self.scene_to_reachable_positions[scene_name] = copy.deepcopy(
self.currently_reachable_points)
assert len(self.scene_to_reachable_positions[scene_name]) > 10
if filtered_objects:
self.set_object_filter(filtered_objects)
else:
self.reset_object_filter()
def random_reachable_state(
self,
seed: Optional[int] = None
) -> Dict[str, Union[Dict[str, float], float]]:
"""Returns a random reachable location in the scene."""
assert (
self.all_metadata_available
), "`random_reachable_state` cannot be called when `self.all_metadata_available` is `False`."
if seed is not None:
random.seed(seed)
# xyz = random.choice(self.currently_reachable_points)
assert len(self.scene_to_reachable_positions[self.scene_name]) > 10
xyz = copy.deepcopy(
random.choice(self.scene_to_reachable_positions[self.scene_name]))
rotation = random.choice(
np.arange(0.0, 360.0, self.config["rotateStepDegrees"]))
horizon = 0.0 # random.choice([0.0, 30.0, 330.0])
return {
**{k: float(v)
for k, v in xyz.items()},
"rotation": {
"x": 0.0,
"y": float(rotation),
"z": 0.0
},
"horizon": float(horizon),
}
def randomize_agent_location(
self,
seed: int = None,
partial_position: Optional[Dict[str, float]] = None,
agent_id: int = 0,
) -> Dict[str, Union[Dict[str, float], float]]:
"""Teleports the agent to a random reachable location in the scene."""
assert 0 <= agent_id < self.agent_count
if partial_position is None:
partial_position = {}
k = 0
state: Optional[Dict] = None
while k == 0 or (not self.last_action_success and k < 10):
# self.reset()
state = {
**self.random_reachable_state(seed=seed),
**partial_position
}
# get_logger().debug("picked target location {}".format(state))
self.controller.step("TeleportFull", **state, agentId=agent_id)
k += 1
if not self.last_action_success:
get_logger().warning((
"Randomize agent location in scene {} and current random state {}"
" with seed {} and partial position {} failed in "
"10 attempts. Forcing the action.").format(
self.scene_name, state, seed, partial_position))
self.controller.step("TeleportFull",
**state,
force_action=True,
agentId=agent_id) # type: ignore
assert self.last_action_success, "Force action failed with {}".format(
state)
# get_logger().debug("location after teleport full {}".format(self.agent_state()))
# self.controller.step("TeleportFull", **self.agent_state()) # TODO only for debug
# get_logger().debug("location after re-teleport full {}".format(self.agent_state()))
return self.agent_state(agent_id=agent_id)
def known_good_locations_list(self):
assert (
self.all_metadata_available
), "`known_good_locations_list` cannot be called when `self.all_metadata_available` is `False`."
return self.scene_to_reachable_positions[self.scene_name]
@property
def currently_reachable_points(self) -> List[Dict[str, float]]:
"""List of {"x": x, "y": y, "z": z} locations in the scene that are
currently reachable."""
self.controller.step(action="GetReachablePositions")
assert (
self.last_action_success
), f"Could not get reachable positions for reason {self.last_event.metadata['errorMessage']}."
return self.last_action_return
@property
def scene_name(self) -> str:
"""Current ai2thor scene."""
return self.controller.last_event.metadata["sceneName"].replace(
"_physics", "")
@property
def current_frame(self) -> np.ndarray:
"""Returns rgb image corresponding to the agent's egocentric view."""
return self.controller.last_event.frame
@property
def current_depth(self) -> np.ndarray:
"""Returns depth image corresponding to the agent's egocentric view."""
return self.controller.last_event.depth_frame
@property
def current_frames(self) -> List[np.ndarray]:
"""Returns rgb images corresponding to the agents' egocentric views."""
return [
self.controller.last_event.events[agent_id].frame
for agent_id in range(self.agent_count)
]
@property
def current_depths(self) -> List[np.ndarray]:
"""Returns depth images corresponding to the agents' egocentric
views."""
return [
self.controller.last_event.events[agent_id].depth_frame
for agent_id in range(self.agent_count)
]
@property
def last_event(self) -> ai2thor.server.Event:
"""Last event returned by the controller."""
return self.controller.last_event
@property
def last_action(self) -> str:
"""Last action, as a string, taken by the agent."""
return self.controller.last_event.metadata["lastAction"]
@property
def last_action_success(self) -> bool:
"""Was the last action taken by the agent a success?"""
return self.controller.last_event.metadata["lastActionSuccess"]
@property
def last_action_return(self) -> Any:
"""Get the value returned by the last action (if applicable).
For an example of an action that returns a value, see
`"GetReachablePositions"`.
"""
return self.controller.last_event.metadata["actionReturn"]
def step(
self,
action_dict: Optional[Dict[str, Union[str, int, float, Dict]]] = None,
**kwargs: Union[str, int, float, Dict],
) -> ai2thor.server.Event:
"""Take a step in the ai2thor environment."""
if action_dict is None:
action_dict = dict()
action_dict.update(kwargs)
return self.controller.step(**action_dict)
def stop(self):
"""Stops the ai2thor controller."""
try:
self.controller.stop()
except Exception as e:
get_logger().warning(str(e))
def all_objects(self) -> List[Dict[str, Any]]:
"""Return all object metadata."""
return self.controller.last_event.metadata["objects"]
def all_objects_with_properties(
self, properties: Dict[str, Any]) -> List[Dict[str, Any]]:
"""Find all objects with the given properties."""
objects = []
for o in self.all_objects():
satisfies_all = True
for k, v in properties.items():
if o[k] != v:
satisfies_all = False
break
if satisfies_all:
objects.append(o)
return objects
def visible_objects(self) -> List[Dict[str, Any]]:
"""Return all visible objects."""
return self.all_objects_with_properties({"visible": True})
class QueryAgentSAC():
def __init__(self, scene, target, room_object_types = G_livingroom_objtype):
# ai2thor
self.scene = scene
self.target = target
assert self.target in room_object_types
self.room_object_types = room_object_types
self.controller = Controller(scene=scene,
renderInstanceSegmentation=True,
width=1080,
height=1080)
self.target_type = target
# register query
FrameInfo.candidates = self.room_object_types
self.query_indicates = [True for _ in range(len(FrameInfo.candidates))]
# Keep track of qa history
self.keep_frame_map = False # weather to keep history to avoid duplicated query
self.history = deque(maxlen = 1000)
# self.frame_map = {} # (position and rotation) -> FrameInfo
# RL part
self.observation = None
self.last_action = None
self.episode_done = False
# reward
self.time_penalty = -0.01
self.action_fail_penalty = -0.01
self.first_seen = False
self.first_seen_reward = 1 # reward for finding the object initially
self.first_in_range = False
self.first_in_range_reward = 5.0 # object in interaction range
self.mission_success_reward = 10.0 # object in interaction range and done
# init event
self.event = None
self.step(5) # self.controller.step("Done")
self.observation = self.get_observation()
# training
self.use_gpu = True
# learning
self.batch_size = 4
self.learning_rate = 0.001
self.alpha = 1 # soc temparature
self.gamma = 0.95 # discount factor
self.soft_target_tau = 0.01
self.target_update_period = 1 # updatge target network frequency
# record
self.n_train_steps_total = 0
self.episode_steps = 0
self.episode_total_reward = 0
self.episode_policy_loss = []
self.episode_pf1_loss = []
self.episode_pf2_loss = []
# policy network
self.input_dim = len(self.observation)
self.hidden_dim = 64
self.action_dim = 6
self.policy = PolicyNetwork(self.input_dim, self.hidden_dim, self.action_dim)
self.qf1 = QNetwork(self.input_dim , self.action_dim, self.hidden_dim)
self.qf2 = QNetwork(self.input_dim , self.action_dim, self.hidden_dim)
self.target_qf1 = QNetwork(self.input_dim , self.action_dim, self.hidden_dim)
self.target_qf2 = QNetwork(self.input_dim , self.action_dim, self.hidden_dim)
if self.use_gpu:
self.policy = self.policy.cuda()
self.qf1 = self.qf1.cuda()
self.qf2 = self.qf2.cuda()
self.target_qf1 = self.target_qf1.cuda()
self.target_qf2 = self.target_qf2.cuda()
# loss
self.qf_criterion = nn.MSELoss()
self.update_target_networks()
self.policy_optimizer = optim.Adam(
self.policy.parameters(),
lr=self.learning_rate,
)
self.qf1_optimizer = optim.Adam(
self.qf1.parameters(),
lr=self.learning_rate,
)
self.qf2_optimizer = optim.Adam(
self.qf2.parameters(),
lr=self.learning_rate,
)
self.update_target_networks()
def update_target_networks(self):
soft_update_from_to(self.qf1, self.target_qf1, self.soft_target_tau)
soft_update_from_to(self.qf2, self.target_qf2, self.soft_target_tau)
def take_action(self, epsilon = 0.2, print_action = False):
current_state = self.observation
if np.random.rand() < epsilon:
action_code = np.random.randint(self.action_dim)
else:
current_state_tensor = torch.FloatTensor(current_state).unsqueeze(0)
if self.use_gpu:
current_state_tensor = current_state_tensor.to("cuda")
action, _ = self.policy.sample_action_with_prob(current_state_tensor)
action_code = torch.argmax(action, dim = -1)[0].item()
#print(action_code)
if print_action:
print("Agent action: ", action_code)
self.step(action_code)
next_observation = self.get_observation()
# calulate reward
reward = self.time_penalty
if not self.event.metadata["lastActionSuccess"]:
reward += self.action_fail_penalty
frame_info = FrameInfo(self.event)
if not self.first_seen:
for obj in frame_info.object_info:
if self.target_type == obj["objectType"]:
reward += self.first_seen_reward
self.first_seen = True
break
if not self.first_in_range:
for obj in frame_info.object_info:
if self.target_type == obj["objectType"] and obj["visible"] == True:
reward += self.first_in_range_reward
self.first_in_range = True
break
for obj in frame_info.object_info:
if self.target_type == obj["objectType"] and obj["visible"] == True:
if action_code == 5:
reward += self.mission_success_reward
self.episode_done = True
break
self.history.append([self.observation.copy(), action_code, reward, next_observation.copy(), self.episode_done])
self.observation = next_observation
self.episode_steps += 1
self.episode_total_reward += reward
def step(self, action_code:int):
self.last_action = action_code
self.event = self.controller.step(G_action_code2action[action_code])
def get_observation(self):
'''
Get state for RL
'''
state = []
frame_info = FrameInfo(self.event)
# target encode
target_encode = [0 for _ in range(len(self.room_object_types))]
target_index = self.room_object_types.index(self.target_type)
target_encode[target_index] = 1
state.extend(target_encode)
# agent state: last action success encode
last_action_success = 1 if self.event.metadata["lastActionSuccess"] else -1
state.append(last_action_success)
last_action_encode = [0] * len(G_action2code)
last_action_encode[self.last_action] = 1
state.extend(last_action_encode)
# agent head position
head_pose = round(self.event.metadata["agent"]['cameraHorizon']) // 30
state.append(head_pose)
# object state query
# if self.keep_history
frame_info = FrameInfo(self.event)
obj_query_state = frame_info.get_answer_array_for_all_candidates(self.query_indicates)
state.extend(obj_query_state)
return state
def sample_history(self, strategy="positive first"):
all_indexes = np.arange(len(self.history))
all_rewards = np.asarray([h[2] for h in self.history])
sample_prob = softmax(all_rewards)
sample_indexes = np.random.choice(all_indexes, size=self.batch_size, replace=True, p=sample_prob)
return [self.history[index] for index in sample_indexes]
def learn(self):
# sample history
# sample_list = random.sample((self.history), self.batch_size)
sample_list = self.sample_history()
s0 = [sample_list[i][0] for i in range(self.batch_size)]
a = [[1 if j == sample_list[i][1] else 0 for j in range(self.action_dim)] for i in range(self.batch_size)]
r = [[sample_list[i][2]] for i in range(self.batch_size)]
s1 = [sample_list[i][3] for i in range(self.batch_size)]
d = [[int(sample_list[i][4])] for i in range(self.batch_size)]
s0 = torch.FloatTensor(s0)
a = torch.LongTensor(a)
r = torch.FloatTensor(r)
s1 = torch.FloatTensor(s1)
d = torch.FloatTensor(d)
if self.use_gpu:
s0 = s0.to("cuda")
a = a.to("cuda")
r = r.to("cuda")
s1 = s1.to("cuda")
d = d.to("cuda")
"""
Policy loss
"""
new_obs_actions, log_pi = self.policy.sample_action_with_prob(s0)
log_pi = log_pi.unsqueeze(-1)
q_new_actions = torch.min(
self.qf1(s0, new_obs_actions),
self.qf2(s0, new_obs_actions),
)
policy_loss = (self.alpha*log_pi - q_new_actions).mean()
"""
QF Loss
"""
q1_pred = self.qf1(s0, a)
q2_pred = self.qf2(s0, a)
new_next_actions, new_log_pi = self.policy.sample_action_with_prob(s1)
new_log_pi = new_log_pi.unsqueeze(-1)
target_q_values = torch.min(self.target_qf1(s1, new_next_actions),self.target_qf2(s1, new_next_actions)) - self.alpha * new_log_pi
q_target = r + (1. - d) * self.gamma * target_q_values
qf1_loss = self.qf_criterion(q1_pred, q_target.detach())
qf2_loss = self.qf_criterion(q2_pred, q_target.detach())
#update parameters
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
self.qf1_optimizer.zero_grad()
qf1_loss.backward()
self.qf1_optimizer.step()
self.qf2_optimizer.zero_grad()
qf2_loss.backward()
self.qf2_optimizer.step()
self.n_train_steps_total += 1
if self.n_train_steps_total % self.target_update_period == 0:
self.update_target_networks()
# record
self.episode_policy_loss.append(policy_loss.item())
self.episode_pf1_loss.append(qf1_loss.item())
self.episode_pf2_loss.append(qf2_loss.item())
def reset_scene_and_target(self, scene: str, target: str):
self.controller.reset(scene=scene)
self.target_type = target
def reset_episode(self):
self.reset_scene_and_target(self.scene, self.target)
self.episode_done = False
self.episode_steps = 0
self.episode_total_reward = 0
self.episode_policy_loss.clear()
self.episode_pf1_loss.clear()
self.episode_pf2_loss.clear()
self.first_seen = False
self.first_in_range = False
def close(self):
self.controller.stop()
def save_model(self):
from datetime import datetime
now = datetime.now()
time_str = now.strftime("%H:%M:%S")
torch.save(self.qf1.state_dict(), "record/qf1_" + time_str + ".pth")
torch.save(self.qf2.state_dict(), "record/qf2_" + time_str + ".pth")
torch.save(self.qf2.state_dict(), "record/policy_" + time_str + ".pth")
class QueryAgentSIL():
def __init__(self, scene, target, room_object_types=G_livingroom_objtype):
# ai2thor
self.scene = scene
self.target = target
assert self.target in room_object_types
self.room_object_types = room_object_types
self.controller = Controller(scene=scene,
renderInstanceSegmentation=True,
width=1080,
height=1080)
self.target_type = target
# register query
FrameInfo.candidates = self.room_object_types
self.query_indicates = [True for _ in range(len(FrameInfo.candidates))]
# Keep track of qa history
self.keep_frame_map = False # weather to keep history to avoid duplicated query
self.replay_buffer = deque(maxlen=1000)
# self.frame_map = {} # (position and rotation) -> FrameInfo
# RL part
self.observation = None
self.last_action = None
self.episode_done = False
self.episode_history = []
# reward
self.time_penalty = -0.01
self.action_fail_penalty = -0.01
self.first_seen = False
self.first_seen_reward = 1 # reward for finding the object initially
self.first_in_range = False
self.first_in_range_reward = 5.0 # object in interaction range
self.mission_success_reward = 10.0 # object in interaction range and done
# init event
self.event = None
self.step(5) # self.controller.step("Done")
self.observation = self.get_observation()
# training
self.use_gpu = True
# learning
self.batch_size = 4
self.learning_rate = 0.001
self.alpha = 1 # soc temparature
self.gamma = 0.95 # discount factor
# record
self.n_train_steps_total = 0
self.episode_steps = 0
self.episode_total_reward = 0
# policy network
self.state_dim = len(self.observation)
self.hidden_dim = 64
self.action_dim = 6
self.policy = PolicyNetwork(self.state_dim, self.hidden_dim,
self.action_dim)
self.value_net = ValueNetwork(self.state_dim, self.action_dim)
if self.use_gpu:
self.policy = self.policy.cuda()
self.value_net = self.value_net.cuda()
# loss
self.vf_criterion = nn.MSELoss()
self.policy_optimizer = optim.Adam(
self.policy.parameters(),
lr=self.learning_rate,
)
self.value_optimizer = optim.Adam(
self.value_net.parameters(),
lr=self.learning_rate,
)
def step(self, action_code: int):
self.last_action = action_code
self.event = self.controller.step(G_action_code2action[action_code])
def get_observation(self):
'''
Get state for RL
'''
state = []
frame_info = FrameInfo(self.event)
# target encode
target_encode = [0 for _ in range(len(self.room_object_types))]
target_index = self.room_object_types.index(self.target_type)
target_encode[target_index] = 1
state.extend(target_encode)
# agent state: last action success encode
last_action_success = 1 if self.event.metadata[
"lastActionSuccess"] else -1
state.append(last_action_success)
last_action_encode = [0] * len(G_action2code)
last_action_encode[self.last_action] = 1
state.extend(last_action_encode)
# agent head position
head_pose = round(self.event.metadata["agent"]['cameraHorizon']) // 30
state.append(head_pose)
# object state query
# if self.keep_history
frame_info = FrameInfo(self.event)
obj_query_state = frame_info.get_answer_array_for_all_candidates(
self.query_indicates)
state.extend(obj_query_state)
return state
def take_action(self):
current_state = self.observation
current_state_tensor = torch.FloatTensor(current_state).unsqueeze(0)
if self.use_gpu:
current_state_tensor = current_state_tensor.to("cuda")
action, log_prob = self.policy.sample_action_with_prob(
current_state_tensor)
action_code = torch.argmax(action, dim=-1)[0].item()
#print(action_code)
self.step(action_code)
next_observation = self.get_observation()
# calulate reward
reward = self.time_penalty
if not self.event.metadata["lastActionSuccess"]:
reward += self.action_fail_penalty
frame_info = FrameInfo(self.event)
if not self.first_seen:
for obj in frame_info.object_info:
if self.target_type == obj["objectType"]:
reward += self.first_seen_reward
self.first_seen = True
break
if not self.first_in_range:
for obj in frame_info.object_info:
if self.target_type == obj["objectType"] and obj[
"visible"] == True:
reward += self.first_in_range_reward
self.first_in_range = True
break
for obj in frame_info.object_info:
if self.target_type == obj["objectType"] and obj["visible"] == True:
if action_code == 5:
reward += self.mission_success_reward
self.episode_done = True
break
self.episode_history.append([
self.observation.copy(), action_code, reward, log_probs[0],
self.episode_done
])
self.observation = next_observation
# for print record
self.episode_steps += 1
self.episode_total_reward += reward
def learn(self):
'''
On episode ends, learn somethings
'''
# A2C part
policy_loss = 0
value_loss = 0
R = 0
for i in reversed(range(len(self.episode_history))):
state = h[0]
action = h[1]
reward = h[2]
log_probs = h[3]
h = self.episode_history[i]
R = self.gamma * R + reward
s0 = torch.FloatTensor(state)
if self.use_gpu:
s0 = s0.to("cuda")
value_i = self.value_net(s0.unsqueeze(-1))[0]
advantage = R - values_i
value_loss += 0.5 * advantage**2
entropy = -torch.sum(torch.exp(log_probs) * log_probs)
policy_loss = policy_loss - log_probs[action] * advantage.detach(
) - self.alpha * entropy
#update parameters
self.policy_optimizer.zero_grad()
policy_loss.backward()
self.policy_optimizer.step()
self.value_optimizer.zero_grad()
value_loss.backward()
self.value_optimizer.step()
def reset_scene_and_target(self, scene: str, target: str):
self.controller.reset(scene=scene)
self.target_type = target
def reset_episode(self):
self.reset_scene_and_target(self.scene, self.target)
self.episode_done = False
self.episode_steps = 0
self.episode_total_reward = 0
self.episode_policy_loss.clear()
self.episode_pf1_loss.clear()
self.episode_pf2_loss.clear()
def close(self):
self.controller.stop()
def save_model(self):
from datetime import datetime
now = datetime.now()
time_str = now.strftime("%H:%M:%S")
torch.save(self.qf1.state_dict(), "record/qf1_" + time_str + ".pth")
torch.save(self.qf2.state_dict(), "record/qf2_" + time_str + ".pth")
torch.save(self.qf2.state_dict(), "record/policy_" + time_str + ".pth")
class drone_explorer(Thread):
def __init__(self, pid, cfg, queue, mode='train'):
super(drone_explorer, self).__init__()
### thread setting
self.pid = pid
set_random_seed(pid + cfg.framework.seed)
self.cfg = cfg
self.queue = queue
self.iters = 0
self.reset = False
self.kill = False
### data
self.objects = json.load(open(cfg.object_dir))
if mode == 'train':
self.trajectory = json.load(open(cfg.train.meta))
self.object_random = cfg.train.random_object
elif mode == 'val':
self.trajectory = json.load(open(cfg.val.meta))
self.object_random = cfg.val.random_object
elif mode == 'test':
self.trajectory = json.load(open(cfg.test.meta))
self.object_random = cfg.test.random_object
else:
raise NotImplementedError
self.trajectory_idx = np.arange(len(self.trajectory['scene']))
np.random.shuffle(self.trajectory_idx)
if pid == 0:
for epoch in range(self.cfg.train.num_epoch):
for ii in range(len(self.trajectory_idx)):
self.queue.put(self.trajectory_idx[ii])
self.force, self.angle_x, self.angle_y = 0, 0, 0
self.event = None
### THOR setting
self.port = cfg.thor.port + pid
self.x_display = '0.{}'.format(pid % cfg.framework.num_gpu + cfg.thor.x_display_offset)
self.restart_cond = cfg.thor.restart_cond
self.controller = None
### try it
self.restart()
def restart(self):
### reset the unity to avoid some latency issue ...
if isinstance(self.controller, type(None)):
self.controller = Controller(local_executable_path=PATH,
scene="FloorPlan201_physics",
x_display=self.x_display,
agentMode='drone',
fieldOfView=60,
port=self.port)
else:
self.controller.stop()
self.controller = Controller(local_executable_path=PATH,
scene="FloorPlan201_physics",
x_display=self.x_display,
agentMode='drone',
fieldOfView=60,
port=self.port)
_ = self.controller.reset(self.trajectory['scene'][0])
_ = self.controller.step(dict(action='ChangeAutoResetTimeScale', timeScale=self.cfg.thor.time_scale))
_ = self.controller.step(dict(action='ChangeFixedDeltaTime', fixedDeltaTime=self.cfg.thor.delta_time))
print('Thread:'+str(self.pid)+' ['+str(self.iters+1)+']: restart finish')
def run(self):
while not self.kill:
### check if it needs to reset or the agent is still exploring
if self.reset and not self.queue.empty():
### check if it reaches the restart condition
if self.iters % self.restart_cond == 0 and self.iters > 1:
self.restart()
### get meta data of the trajectory
try:
tidx = self.queue.get(timeout=1)
except Empty:
if self.kill:
break
else:
self.good_to_go = False
self.reset = False
continue
### have a loop to ensure the trajectory has a good start
self.good_to_go = False
while not self.good_to_go:
scene = self.trajectory['scene'][tidx]
object_name = self.trajectory['object'][tidx]
mass = self.objects[object_name][1]
drone_position = self.trajectory['drone_position'][tidx]
launcher_position = self.trajectory['launcher_position'][tidx]
force = self.trajectory['force'][tidx]
angle_y = self.trajectory['angle_y'][tidx]
angle_x = self.trajectory['angle_x'][tidx]
if object_name == "Glassbottle":
object_name = "Bottle"
# make sure the value is .2f
drone_position['x'] = np.round(drone_position['x'], 2)
drone_position['y'] = np.round(1.5, 2)
drone_position['z'] = np.round(drone_position['z'], 2)
launcher_position['x'] = np.round(launcher_position['x'], 2)
launcher_position['y'] = np.round(launcher_position['y'], 2)
launcher_position['z'] = np.round(launcher_position['z'], 2)
force = np.round(force, 2)
angle_y = np.round(angle_y, 2)
angle_x = np.round(angle_x, 2)
### set THOR
event = self.controller.reset(scene)
event = self.controller.step(dict(action='SpawnDroneLauncher', position=launcher_position))
event = self.controller.step(dict(action='FlyAssignStart',
position=drone_position,
x=launcher_position['x'],
y=launcher_position['y'],
z=launcher_position['z']))
event = self.controller.step(dict(action='Rotate', rotation=dict(x=0, y=0, z=0)))
event = self.controller.step(dict(action='ChangeAutoResetTimeScale',
timeScale=self.cfg.thor.time_scale))
event = self.controller.step(dict(action='ChangeFixedDeltaTime',
fixedDeltaTime=self.cfg.thor.delta_time))
if "noise_sigma" in self.cfg.agent:
event = self.controller.step(dict(action='ChangeDronePositionRandomNoiseSigma',
dronePositionRandomNoiseSigma=self.cfg.agent.noise_sigma))
### prepare to launch the object
position = event.metadata['agent']['position']
event = self.controller.step(dict(action = 'LaunchDroneObject',
moveMagnitude = force,
x = angle_x,
y = angle_y,
z = -1,
objectName=object_name,
objectRandom=self.object_random))
if np.round(event.metadata['currentTime'], 2) == 0.00:
event = self.controller.step(dict(action='Pass'))
if np.round(event.metadata['currentTime'], 2) == 0.02:
self.good_to_go = True
### some record
self.angle_x = angle_x
self.angle_y = angle_y
self.force = force
self.object_name = object_name
self.mass = mass
self.event = event
self.tidx = tidx
### thread setting
self.reset = False
self.iters += 1
### otherwise, just go to sleep
else:
time.sleep(0.1)
class Ai2Thor():
def __init__(self):
self.visualize = False
self.verbose = False
self.save_imgs = True
self.do_orbslam = False
self.do_depth_noise = False
self.makevideo = True
# st()
# these are all map names
a = np.arange(1, 30)
b = np.arange(201, 231)
c = np.arange(301, 331)
d = np.arange(401, 431)
abcd = np.hstack((a, b, c, d))
mapnames = []
for i in list(abcd):
mapname = 'FloorPlan' + str(i)
mapnames.append(mapname)
np.random.seed(1)
random.shuffle(mapnames)
self.mapnames = mapnames
self.num_episodes = len(self.mapnames)
self.ignore_classes = []
# classes to save
self.include_classes = [
'ShowerDoor', 'Cabinet', 'CounterTop', 'Sink', 'Towel',
'HandTowel', 'TowelHolder', 'SoapBar', 'ToiletPaper',
'ToiletPaperHanger', 'HandTowelHolder', 'SoapBottle', 'GarbageCan',
'Candle', 'ScrubBrush', 'Plunger', 'SinkBasin', 'Cloth',
'SprayBottle', 'Toilet', 'Faucet', 'ShowerHead', 'Box', 'Bed',
'Book', 'DeskLamp', 'BasketBall', 'Pen', 'Pillow', 'Pencil',
'CellPhone', 'KeyChain', 'Painting', 'CreditCard', 'AlarmClock',
'CD', 'Laptop', 'Drawer', 'SideTable', 'Chair', 'Blinds', 'Desk',
'Curtains', 'Dresser', 'Watch', 'Television', 'WateringCan',
'Newspaper', 'FloorLamp', 'RemoteControl', 'HousePlant', 'Statue',
'Ottoman', 'ArmChair', 'Sofa', 'DogBed', 'BaseballBat',
'TennisRacket', 'VacuumCleaner', 'Mug', 'ShelvingUnit', 'Shelf',
'StoveBurner', 'Apple', 'Lettuce', 'Bottle', 'Egg', 'Microwave',
'CoffeeMachine', 'Fork', 'Fridge', 'WineBottle', 'Spatula',
'Bread', 'Tomato', 'Pan', 'Cup', 'Pot', 'SaltShaker', 'Potato',
'PepperShaker', 'ButterKnife', 'StoveKnob', 'Toaster',
'DishSponge', 'Spoon', 'Plate', 'Knife', 'DiningTable', 'Bowl',
'LaundryHamper', 'Vase', 'Stool', 'CoffeeTable', 'Poster',
'Bathtub', 'TissueBox', 'Footstool', 'BathtubBasin',
'ShowerCurtain', 'TVStand', 'Boots', 'RoomDecor', 'PaperTowelRoll',
'Ladle', 'Kettle', 'Safe', 'GarbageBag', 'TeddyBear',
'TableTopDecor', 'Dumbbell', 'Desktop', 'AluminumFoil', 'Window'
]
self.small_classes = []
self.rot_interval = 5.0
self.radius_max = 3.5 #3 #1.75
self.radius_min = 1.0 #1.25
self.num_flat_views = 3
self.num_any_views = 7
self.num_views = 25
self.obj_per_scene = 5
# self.origin_quaternion = np.quaternion(1, 0, 0, 0)
# self.origin_rot_vector = quaternion.as_rotation_vector(self.origin_quaternion)
self.homepath = f'/home/sirdome/katefgroup/gsarch/ithor/data/test'
# self.basepath = '/home/nel/gsarch/replica_traj_bed'
if not os.path.exists(self.homepath):
os.mkdir(self.homepath)
else:
val = input("Delete homepath? [y/n]: ")
if val == 'y':
import shutil
shutil.rmtree(self.homepath)
os.mkdir(self.homepath)
else:
print("ENDING")
assert (False)
self.W = 256
self.H = 256
self.fov = 90
hfov = float(self.fov) * np.pi / 180.
self.pix_T_camX = np.array([[
(self.W / 2.) * 1 / np.tan(hfov / 2.), 0., 0., 0.
], [0., (self.H / 2.) * 1 / np.tan(hfov / 2.), 0., 0.], [0., 0., 1, 0],
[0., 0., 0, 1]])
self.pix_T_camX[0, 2] = self.W / 2.
self.pix_T_camX[1, 2] = self.H / 2.
self.run_episodes()
def run_episodes(self):
self.ep_idx = 0
# self.objects = []
for episode in range(self.num_episodes):
print("STARTING EPISODE ", episode)
mapname = self.mapnames[episode]
print("MAPNAME=", mapname)
self.controller = Controller(
scene=mapname,
gridSize=0.25,
width=self.W,
height=self.H,
fieldOfView=self.fov,
renderObjectImage=True,
renderDepthImage=True,
)
# self.controller.start()
self.basepath = self.homepath + f"/{mapname}_{episode}"
print("BASEPATH: ", self.basepath)
# self.basepath = f"/hdd/ayushj/habitat_data/{mapname}_{episode}"
if not os.path.exists(self.basepath):
os.mkdir(self.basepath)
self.run()
self.controller.stop()
time.sleep(1)
self.ep_idx += 1
def save_datapoint(self, observations, data_path, viewnum, flat_view):
if self.verbose:
print("Print Sensor States.", self.agent.state.sensor_states)
rgb = observations["color_sensor"]
semantic = observations["semantic_sensor"]
# st()
depth = observations["depth_sensor"]
agent_pos = observations["positions"]
agent_rot = observations["rotations"]
# Assuming all sensors have same extrinsics
color_sensor_pos = observations["positions"]
color_sensor_rot = observations["rotations"]
#print("POS ", agent_pos)
#print("ROT ", color_sensor_rot)
object_list = observations['object_list']
# print(viewnum, agent_pos)
# print(agent_rot)
if False:
plt.imshow(rgb)
plt_name = f'/home/nel/gsarch/aithor/data/test/img_mask{viewnum}.png'
plt.savefig(plt_name)
save_data = {
'flat_view': flat_view,
'objects_info': object_list,
'rgb_camX': rgb,
'depth_camX': depth,
'sensor_pos': color_sensor_pos,
'sensor_rot': color_sensor_rot
}
with open(os.path.join(data_path, str(viewnum) + ".p"), 'wb') as f:
pickle.dump(save_data, f)
f.close()
def quat_from_angle_axis(self, theta: float,
axis: np.ndarray) -> np.quaternion:
r"""Creates a quaternion from angle axis format
:param theta: The angle to rotate about the axis by
:param axis: The axis to rotate about
:return: The quaternion
"""
axis = axis.astype(np.float)
axis /= np.linalg.norm(axis)
return quaternion.from_rotation_vector(theta * axis)
def safe_inverse_single(self, a):
r, t = self.split_rt_single(a)
t = np.reshape(t, (3, 1))
r_transpose = r.T
inv = np.concatenate([r_transpose, -np.matmul(r_transpose, t)], 1)
bottom_row = a[3:4, :] # this is [0, 0, 0, 1]
# bottom_row = torch.tensor([0.,0.,0.,1.]).view(1,4)
inv = np.concatenate([inv, bottom_row], 0)
return inv
def split_rt_single(self, rt):
r = rt[:3, :3]
t = np.reshape(rt[:3, 3], 3)
return r, t
def eul2rotm(self, rx, ry, rz):
# inputs are shaped B
# this func is copied from matlab
# R = [ cy*cz sy*sx*cz-sz*cx sy*cx*cz+sz*sx
# cy*sz sy*sx*sz+cz*cx sy*cx*sz-cz*sx
# -sy cy*sx cy*cx]
# rx = np.expand_dims(rx, axis=1)
# ry = np.expand_dims(ry, axis=1)
# rz = np.expand_dims(rz, axis=1)
# st()
# these are B x 1
sinz = np.sin(rz)
siny = np.sin(ry)
sinx = np.sin(rx)
cosz = np.cos(rz)
cosy = np.cos(ry)
cosx = np.cos(rx)
r11 = cosy * cosz
r12 = sinx * siny * cosz - cosx * sinz
r13 = cosx * siny * cosz + sinx * sinz
r21 = cosy * sinz
r22 = sinx * siny * sinz + cosx * cosz
r23 = cosx * siny * sinz - sinx * cosz
r31 = -siny
r32 = sinx * cosy
r33 = cosx * cosy
r = np.array([[r11, r12, r13], [r21, r22, r23], [r31, r32, r33]])
# r1 = np.stack([r11,r12,r13],axis=2)
# r2 = np.stack([r21,r22,r23],axis=2)
# r3 = np.stack([r31,r32,r33],axis=2)
# r = np.concatenate([r1,r2,r3],axis=1)
return r
def rotm2eul(self, r):
# r is Bx3x3, or Bx4x4
r00 = r[0, 0]
r10 = r[1, 0]
r11 = r[1, 1]
r12 = r[1, 2]
r20 = r[2, 0]
r21 = r[2, 1]
r22 = r[2, 2]
## python guide:
# if sy > 1e-6: # singular
# x = math.atan2(R[2,1] , R[2,2])
# y = math.atan2(-R[2,0], sy)
# z = math.atan2(R[1,0], R[0,0])
# else:
# x = math.atan2(-R[1,2], R[1,1])
# y = math.atan2(-R[2,0], sy)
# z = 0
sy = np.sqrt(r00 * r00 + r10 * r10)
cond = (sy > 1e-6)
rx = np.where(cond, np.arctan2(r21, r22), np.arctan2(-r12, r11))
ry = np.where(cond, np.arctan2(-r20, sy), np.arctan2(-r20, sy))
rz = np.where(cond, np.arctan2(r10, r00), np.zeros_like(r20))
# rx = torch.atan2(r21, r22)
# ry = torch.atan2(-r20, sy)
# rz = torch.atan2(r10, r00)
# rx[cond] = torch.atan2(-r12, r11)
# ry[cond] = torch.atan2(-r20, sy)
# rz[cond] = 0.0
return rx, ry, rz
def generate_xyz_habitatCamXs(self, flat_obs):
pix_T_camX = self.pix_T_camX
xyz_camXs = []
for i in range(2): #depth_camXs.shape[0]):
K = pix_T_camX
xs, ys = np.meshgrid(
np.linspace(-1 * self.W / 2., 1 * self.W / 2., self.W),
np.linspace(1 * self.W / 2., -1 * self.W / 2., self.W))
depth = flat_obs[i]['depth_sensor'].reshape(1, self.W, self.W)
if i > 0:
rotation_X = flat_obs[i][
'rotations'] #quaternion.as_rotation_matrix(rot)
pos = flat_obs[i]['positions']
euler_rot_X = flat_obs[i]["rotations_euler"]
# euler_rot_X_rad = np.radians(flat_obs[i]["rotations_euler"])
origin_T_camX_4x4 = np.eye(4)
origin_T_camX_4x4[0:3, 0:3] = rotation_X
origin_T_camX_4x4[0:3, 3] = pos
origin_T_camX = rotation_X
camX_T_origin_4x4 = self.safe_inverse_single(origin_T_camX_4x4)
camX0_T_camX = np.matmul(camX0_T_origin, origin_T_camX)
camX0_T_camX_4x4 = np.matmul(camX0_T_origin_4x4,
origin_T_camX_4x4)
# camX0_T_camX = np.matmul(rotation_0, np.linalg.inv(origin_T_camX))
# camX0_T_camX_4x4 = np.matmul(origin_T_camX0_4x4, camX_T_origin_4x4)
# r = R.from_matrix(camX0_T_camX)
# print("CamX0_T_camX CHECK: ", r.as_euler('xyz', degrees=True))
# r = R.from_matrix(origin_T_camX_4x4[0:3,0:3])
# print("CamX0_T_camX CHECK: ", r.as_euler('xyz', degrees=True))
# r = R.from_matrix(rotation_X)
# rx = euler_rot_X_rad[0]
# ry = euler_rot_X_rad[1]
# rz = euler_rot_X_rad[2]
# rotm = self.eul2rotm(rx, ry, rz)
rx, ry, rz = self.rotm2eul(rotation_X)
print("EULER ACTUAL: ", euler_rot_X)
print("EULER OBTAINED: ", rx, ry, rz)
rx, ry, rz = self.rotm2eul(camX0_T_camX)
# print("origin_T_camx CHECK: ", r.as_euler('xyz', degrees=True))s
print("EULER SUBTRACT: ", euler_rot_0 - euler_rot_X)
print("EULER OBTAINED: ", rx, ry, rz)
# st()
elif i == 0:
rotation_0 = flat_obs[i][
'rotations'] #quaternion.as_rotation_matrix(rot)
pos = flat_obs[i]['positions']
euler_rot_0 = flat_obs[i]["rotations_euler"]
origin_T_camX0_4x4 = np.eye(4)
origin_T_camX0_4x4[0:3, 0:3] = rotation_0
origin_T_camX0_4x4[0:3, 3] = pos
camX0_T_origin_4x4 = self.safe_inverse_single(
origin_T_camX0_4x4)
camX0_T_origin = np.linalg.inv(rotation_0)
xs = xs.reshape(1, self.W, self.W)
ys = ys.reshape(1, self.W, self.W)
xys = np.vstack(
(xs * depth, ys * depth, -depth, np.ones(depth.shape)))
xys = xys.reshape(4, -1)
xy_c0 = np.matmul(np.linalg.inv(K), xys)
xyz_camX = xy_c0.T[:, :3]
xyz_camXs.append(xyz_camX)
if i == 1:
xyz = np.expand_dims(xyz_camX, axis=0)
B, N, _ = list(xyz.shape)
ones = np.ones_like(xyz[:, :, 0:1])
xyz1 = np.concatenate([xyz, ones], 2)
xyz1_t = np.transpose(xyz1, (0, 2, 1))
# this is B x 4 x N
xyz2_t = np.matmul(camX0_T_camX_4x4, xyz1_t)
xyz2 = np.transpose(xyz2_t, (0, 2, 1))
xyz2 = np.squeeze(xyz2[:, :, :3])
# xyz2 = xyz_camX
import matplotlib.pyplot as plt
from mpl_toolkits.mplot3d import Axes3D
fig = plt.figure()
# ax = fig.add_subplot(111, projection='3d')
xs = xyz2[:, 0]
ys = xyz2[:, 1]
zs = xyz2[:, 2]
# ax.scatter(xs, ys, zs)
plt.plot(xs, zs, 'x', color='red')
xyz3 = xyz_camXs[0]
xs = xyz3[:, 0]
ys = xyz3[:, 1]
zs = xyz3[:, 2]
# ax.scatter(xs, ys, zs)
plt.plot(xs, zs, 'o', color='blue')
plt_name = '/home/nel/gsarch/aithor/data/pointcloud.png'
plt.savefig(plt_name)
st()
return np.stack(xyz_camXs)
def run2(self):
event = self.controller.step('GetReachablePositions')
for obj in event.metadata['objects']:
if obj['objectType'] not in self.objects:
self.objects.append(obj['objectType'])
def run(self):
event = self.controller.step('GetReachablePositions')
self.nav_pts = event.metadata['reachablePositions']
self.nav_pts = np.array([list(d.values()) for d in self.nav_pts])
np.random.seed(1)
# objects = np.random.choice(event.metadata['objects'], self.obj_per_scene, replace=False)
objects = event.metadata['objects']
objects_inds = np.arange(len(event.metadata['objects']))
np.random.shuffle(objects_inds)
# objects = np.random.shuffle(event.metadata['objects'])
# for obj in event.metadata['objects']: #objects:
# print(obj['name'])
# objects = objects[0]
successes = 0
meta_obj_idx = 0
while successes < self.obj_per_scene and meta_obj_idx <= len(
event.metadata['objects']
) - 1: #obj in objects: #event.metadata['objects']: #objects:
obj = objects[objects_inds[meta_obj_idx]]
meta_obj_idx += 1
print("Object is ", obj['objectType'])
# if obj['name'] in ['Microwave_b200e0bc']:
# print(obj['name'])
# else:
# continue
# print(obj['name'])
if obj['objectType'] not in self.include_classes:
print("Continuing... Invalid Object")
continue
# Calculate distance to object center
obj_center = np.array(
list(obj['axisAlignedBoundingBox']['center'].values()))
obj_center = np.expand_dims(obj_center, axis=0)
distances = np.sqrt(np.sum((self.nav_pts - obj_center)**2, axis=1))
# Get points with r_min < dist < r_max
valid_pts = self.nav_pts[np.where(
(distances > self.radius_min) * (distances < self.radius_max))]
# Bin points based on angles [vertical_angle (10 deg/bin), horizontal_angle (10 deg/bin)]
valid_pts_shift = valid_pts - obj_center
dz = valid_pts_shift[:, 2]
dx = valid_pts_shift[:, 0]
dy = valid_pts_shift[:, 1]
# Get yaw for binning
valid_yaw = np.degrees(np.arctan2(dz, dx))
nbins = 18
bins = np.linspace(-180, 180, nbins + 1)
bin_yaw = np.digitize(valid_yaw, bins)
num_valid_bins = np.unique(bin_yaw).size
if False:
import matplotlib.cm as cm
colors = iter(cm.rainbow(np.linspace(0, 1, nbins)))
plt.figure(2)
plt.clf()
print(np.unique(bin_yaw))
for bi in range(nbins):
cur_bi = np.where(bin_yaw == (bi + 1))
points = valid_pts[cur_bi]
x_sample = points[:, 0]
z_sample = points[:, 2]
plt.plot(z_sample, x_sample, 'o', color=next(colors))
plt.plot(self.nav_pts[:, 2],
self.nav_pts[:, 0],
'x',
color='red')
plt.plot(obj_center[:, 2],
obj_center[:, 0],
'x',
color='black')
plt_name = '/home/nel/gsarch/aithor/data/valid.png'
plt.savefig(plt_name)
if num_valid_bins == 0:
continue
spawns_per_bin = int(self.num_views / num_valid_bins) + 2
# print(f'spawns_per_bin: {spawns_per_bin}')
action = "do_nothing"
episodes = []
valid_pts_selected = []
camXs_T_camX0_4x4 = []
camX0_T_camXs_4x4 = []
origin_T_camXs = []
origin_T_camXs_t = []
cnt = 0
for b in range(nbins):
# get all angle indices in the current bin range
inds_bin_cur = np.where(
bin_yaw == (b + 1)) # bins start 1 so need +1
inds_bin_cur = list(inds_bin_cur[0])
if len(inds_bin_cur) == 0:
continue
for s in range(spawns_per_bin):
observations = {}
if len(inds_bin_cur) == 0:
continue
rand_ind = np.random.randint(0, len(inds_bin_cur))
s_ind = inds_bin_cur.pop(rand_ind)
pos_s = valid_pts[s_ind]
valid_pts_selected.append(pos_s)
# add height from center of agent to camera
pos_s[1] = pos_s[1] + 0.675
# YAW calculation - rotate to object
agent_to_obj = np.squeeze(obj_center) - pos_s
agent_local_forward = np.array([0, 0, 1.0])
flat_to_obj = np.array(
[agent_to_obj[0], 0.0, agent_to_obj[2]])
flat_dist_to_obj = np.linalg.norm(flat_to_obj)
flat_to_obj /= flat_dist_to_obj
det = (flat_to_obj[0] * agent_local_forward[2] -
agent_local_forward[0] * flat_to_obj[2])
turn_angle = math.atan2(
det, np.dot(agent_local_forward, flat_to_obj))
turn_yaw = np.degrees(turn_angle)
turn_pitch = -np.degrees(
math.atan2(agent_to_obj[1], flat_dist_to_obj))
event = self.controller.step('TeleportFull',
x=pos_s[0],
y=pos_s[1],
z=pos_s[2],
rotation=dict(x=0.0,
y=int(turn_yaw),
z=0.0),
horizon=int(turn_pitch))
rgb = event.frame
eulers_xyz_rad = np.radians(
np.array([
event.metadata['agent']['cameraHorizon'],
event.metadata['agent']['rotation']['y'], 0.0
]))
rx = eulers_xyz_rad[0]
ry = eulers_xyz_rad[1]
rz = eulers_xyz_rad[2]
rotation_r_matrix = self.eul2rotm(-rx, -ry, rz)
agent_position = np.array(
list(event.metadata['agent']['position'].values())
) + np.array([0.0, 0.675, 0.0])
# need to invert since z is positive here by convention
agent_position[2] = -agent_position[2]
observations["positions"] = agent_position
observations["rotations"] = rotation_r_matrix
# rt_4x4 = np.eye(4)
# rt_4x4[0:3,0:3] = observations["rotations"]
# rt_4x4[0:3,3] = observations["positions"]
# rt_4x4_inv = self.safe_inverse_single(rt_4x4)
# r, t = self.split_rt_single(rt_4x4_inv)
# observations["positions"] = r
# observations["positions"] = t
# observations["rotations_euler"] = np.array([rx, ry, rz]) #rotation_r.as_euler('xyz', degrees=True)
observations["color_sensor"] = rgb
observations["depth_sensor"] = event.depth_frame
observations[
"semantic_sensor"] = event.instance_segmentation_frame
if False:
plt.imshow(rgb)
plt_name = f'/home/nel/gsarch/aithor/data/test/img_true{s}{b}.png'
plt.savefig(plt_name)
# print("Processed image #", cnt, " for object ", obj['objectType'])
semantic = event.instance_segmentation_frame
object_id_to_color = event.object_id_to_color
color_to_object_id = event.color_to_object_id
obj_ids = np.unique(semantic.reshape(
-1, semantic.shape[2]),
axis=0)
instance_masks = event.instance_masks
instance_detections2D = event.instance_detections2D
obj_metadata_IDs = []
for obj_m in event.metadata['objects']: #objects:
obj_metadata_IDs.append(obj_m['objectId'])
object_list = []
for obj_idx in range(obj_ids.shape[0]):
try:
obj_color = tuple(obj_ids[obj_idx])
object_id = color_to_object_id[obj_color]
except:
# print("Skipping ", object_id)
continue
if object_id not in obj_metadata_IDs:
# print("Skipping ", object_id)
continue
obj_meta_index = obj_metadata_IDs.index(object_id)
obj_meta = event.metadata['objects'][obj_meta_index]
obj_category_name = obj_meta['objectType']
# continue if not visible or not in include classes
if obj_category_name not in self.include_classes or not obj_meta[
'visible']:
continue
obj_instance_mask = instance_masks[object_id]
obj_instance_detection2D = instance_detections2D[
object_id] # [start_x, start_y, end_x, end_y]
obj_instance_detection2D = np.array([
obj_instance_detection2D[1],
obj_instance_detection2D[0],
obj_instance_detection2D[3],
obj_instance_detection2D[2]
]) # ymin, xmin, ymax, xmax
if False:
print(object_id)
plt.imshow(obj_instance_mask)
plt_name = f'/home/nel/gsarch/aithor/data/test/img_mask{s}.png'
plt.savefig(plt_name)
obj_center_axisAligned = np.array(
list(obj_meta['axisAlignedBoundingBox']
['center'].values()))
obj_center_axisAligned[2] = -obj_center_axisAligned[2]
obj_size_axisAligned = np.array(
list(obj_meta['axisAlignedBoundingBox']
['size'].values()))
# print(obj_category_name)
if self.verbose:
print("Saved class name is : ", obj_category_name)
obj_data = {
'instance_id': object_id,
'category_id': object_id,
'category_name': obj_category_name,
'bbox_center': obj_center_axisAligned,
'bbox_size': obj_size_axisAligned,
'mask_2d': obj_instance_mask,
'box_2d': obj_instance_detection2D
}
# object_list.append(obj_instance)
object_list.append(obj_data)
observations["object_list"] = object_list
# check if object visible (make sure agent is not behind a wall)
obj_id = obj['objectId']
obj_id_to_color = object_id_to_color[obj_id]
# if np.sum(obj_ids==object_id_to_color[obj_id]) > 0:
if self.verbose:
print("episode is valid......")
episodes.append(observations)
# print(cnt)
if False:
if cnt > 0:
origin_T_camX = episodes[cnt]["rotations"]
r = R.from_matrix(origin_T_camX)
print("EULER CHECK: ",
r.as_euler('xyz', degrees=True))
camX0_T_camX = np.matmul(camX0_T_origin,
origin_T_camX)
r = R.from_matrix(camX0_T_camX)
print("EULER CHECK: ",
r.as_euler('xyz', degrees=True))
# camX0_T_camX = np.matmul(camX0_T_origin, origin_T_camX)
r = R.from_matrix(camX0_T_origin)
print("EULER CHECK: ",
r.as_euler('xyz', degrees=True))
origin_T_camXs.append(origin_T_camX)
origin_T_camXs_t.append(episodes[cnt]["positions"])
origin_T_camX_4x4 = np.eye(4)
origin_T_camX_4x4[0:3, 0:3] = origin_T_camX
origin_T_camX_4x4[:3,
3] = episodes[cnt]["positions"]
camX0_T_camX_4x4 = np.matmul(
camX0_T_origin_4x4, origin_T_camX_4x4)
camX_T_camX0_4x4 = self.safe_inverse_single(
camX0_T_camX_4x4)
camXs_T_camX0_4x4.append(camX_T_camX0_4x4)
camX0_T_camXs_4x4.append(camX0_T_camX_4x4)
camX0_T_camX_quat = quaternion.from_rotation_matrix(
camX0_T_camX)
camX0_T_camX_eul = quaternion.as_euler_angles(
camX0_T_camX_quat)
camX0_T_camX_4x4 = self.safe_inverse_single(
camX_T_camX0_4x4)
origin_T_camX_4x4 = np.matmul(
origin_T_camX0_4x4, camX0_T_camX_4x4)
r_origin_T_camX, t_origin_T_camX, = self.split_rt_single(
origin_T_camX_4x4)
if self.verbose:
print(r_origin_T_camX)
print(origin_T_camX)
else:
origin_T_camX0 = episodes[0]["rotations"]
camX0_T_origin = np.linalg.inv(origin_T_camX0)
# camX0_T_origin = self.safe_inverse_single(origin_T_camX0)
origin_T_camXs.append(origin_T_camX0)
origin_T_camXs_t.append(episodes[0]["positions"])
origin_T_camX0_4x4 = np.eye(4)
origin_T_camX0_4x4[0:3, 0:3] = origin_T_camX0
origin_T_camX0_4x4[:3,
3] = episodes[0]["positions"]
camX0_T_origin_4x4 = self.safe_inverse_single(
origin_T_camX0_4x4)
camXs_T_camX0_4x4.append(np.eye(4))
camX0_T_camXs_4x4.append(np.eye(4))
origin_T_camX0_t = episodes[0]["positions"]
cnt += 1
if len(episodes) >= self.num_views:
print(f'num episodes: {len(episodes)}')
data_folder = obj['name']
data_path = os.path.join(self.basepath, data_folder)
print("Saving to ", data_path)
os.mkdir(data_path)
# flat_obs = np.random.choice(episodes, self.num_views, replace=False)
np.random.seed(1)
rand_inds = np.sort(
np.random.choice(len(episodes),
self.num_views,
replace=False))
bool_inds = np.zeros(len(episodes), dtype=bool)
bool_inds[rand_inds] = True
flat_obs = np.array(episodes)[bool_inds]
flat_obs = list(flat_obs)
viewnum = 0
if False:
self.generate_xyz_habitatCamXs(flat_obs)
for obs in flat_obs:
self.save_datapoint(obs, data_path, viewnum, True)
viewnum += 1
print("SUCCESS #", successes)
successes += 1
else:
print("Not enough episodes:", len(episodes))
class RoboThorEnvironment:
"""Wrapper for the robo2thor controller providing additional functionality
and bookkeeping.
See [here](https://ai2thor.allenai.org/robothor/documentation) for comprehensive
documentation on RoboTHOR.
# Attributes
controller : The AI2THOR controller.
config : The AI2THOR controller configuration
"""
def __init__(self, **kwargs):
self.config = dict(
rotateStepDegrees=30.0,
visibilityDistance=1.0,
gridSize=0.25,
agentType="stochastic",
continuousMode=True,
snapToGrid=False,
agentMode="bot",
width=640,
height=480,
)
recursive_update(self.config, {**kwargs, "agentMode": "bot"})
self.controller = Controller(
**self.config, server_class=ai2thor.fifo_server.FifoServer)
self.known_good_locations: Dict[str, Any] = {
self.scene_name: copy.deepcopy(self.currently_reachable_points)
}
self.distance_cache = DynamicDistanceCache(rounding=1)
assert len(self.known_good_locations[self.scene_name]) > 10
def initialize_grid_dimensions(
self, reachable_points: Collection[Dict[str, float]]
) -> Tuple[int, int, int, int]:
"""Computes bounding box for reachable points quantized with the
current gridSize."""
points = {(
round(p["x"] / self.config["gridSize"]),
round(p["z"] / self.config["gridSize"]),
): p
for p in reachable_points}
assert len(reachable_points) == len(points)
xmin, xmax = min([p[0] for p in points]), max([p[0] for p in points])
zmin, zmax = min([p[1] for p in points]), max([p[1] for p in points])
return xmin, xmax, zmin, zmax
def set_object_filter(self, object_ids: List[str]):
self.controller.step("SetObjectFilter",
objectIds=object_ids,
renderImage=False)
def reset_object_filter(self):
self.controller.step("ResetObjectFilter", renderImage=False)
def path_from_point_to_object_type(
self, point: Dict[str, float],
object_type: str) -> Optional[List[Dict[str, float]]]:
try:
return metrics.get_shortest_path_to_object_type(
self.controller, object_type, point)
except:
get_logger().debug(
"Failed to find path for {} in {}. Start point {}, agent state {}."
.format(
object_type,
self.controller.last_event.metadata["sceneName"],
point,
self.agent_state(),
))
return None
def distance_from_point_to_object_type(self, point: Dict[str, float],
object_type: str) -> float:
"""Minimal geodesic distance from a point to an object of the given
type.
It might return -1.0 for unreachable targets.
"""
path = self.path_from_point_to_object_type(point, object_type)
if path:
return metrics.path_distance(path)
return -1.0
def distance_to_object_type(self, object_type: str) -> float:
"""Minimal geodesic distance to object of given type from agent's
current location.
It might return -1.0 for unreachable targets.
"""
return self.distance_cache.find_distance(
self.controller.last_event.metadata["agent"]["position"],
object_type,
self.distance_from_point_to_object_type,
)
def path_from_point_to_point(
self, position: Dict[str, float],
target: Dict[str, float]) -> Optional[List[Dict[str, float]]]:
try:
return self.controller.step(
action="GetShortestPathToPoint",
position=position,
x=target["x"],
y=target["y"],
z=target["z"],
# renderImage=False
).metadata["actionReturn"]["corners"]
except:
get_logger().debug(
"Failed to find path for {} in {}. Start point {}, agent state {}."
.format(
target,
self.controller.last_event.metadata["sceneName"],
position,
self.agent_state(),
))
return None
def distance_from_point_to_point(self, position: Dict[str, float],
target: Dict[str, float]) -> float:
path = self.path_from_point_to_point(position, target)
if path:
return metrics.path_distance(path)
return -1.0
def distance_to_point(self, target: Dict[str, float]) -> float:
"""Minimal geodesic distance to end point from agent's current
location.
It might return -1.0 for unreachable targets.
"""
return self.distance_cache.find_distance(
self.controller.last_event.metadata["agent"]["position"],
target,
self.distance_from_point_to_point,
)
def agent_state(self) -> Dict:
"""Return agent position, rotation and horizon."""
agent_meta = self.last_event.metadata["agent"]
return {
**{k: float(v)
for k, v in agent_meta["position"].items()},
"rotation":
{k: float(v)
for k, v in agent_meta["rotation"].items()},
"horizon": round(float(agent_meta["cameraHorizon"]), 1),
}
def teleport(self,
pose: Dict[str, float],
rotation: Dict[str, float],
horizon: float = 0.0):
e = self.controller.step(
action="TeleportFull",
x=pose["x"],
y=pose["y"],
z=pose["z"],
rotation=rotation,
horizon=horizon,
)
return e.metadata["lastActionSuccess"]
def reset(self,
scene_name: str = None,
filtered_objects: Optional[List[str]] = None) -> None:
"""Resets scene to a known initial state."""
if scene_name is not None and scene_name != self.scene_name:
self.controller.reset(scene_name)
assert self.last_action_success, "Could not reset to new scene"
if scene_name not in self.known_good_locations:
self.known_good_locations[scene_name] = copy.deepcopy(
self.currently_reachable_points)
assert len(self.known_good_locations[scene_name]) > 10
if filtered_objects:
self.set_object_filter(filtered_objects)
else:
self.reset_object_filter()
def random_reachable_state(
self,
seed: Optional[int] = None
) -> Dict[str, Union[Dict[str, float], float]]:
"""Returns a random reachable location in the scene."""
if seed is not None:
random.seed(seed)
# xyz = random.choice(self.currently_reachable_points)
assert len(self.known_good_locations[self.scene_name]) > 10
xyz = copy.deepcopy(
random.choice(self.known_good_locations[self.scene_name]))
rotation = random.choice(
np.arange(0.0, 360.0, self.config["rotateStepDegrees"]))
horizon = 0.0 # random.choice([0.0, 30.0, 330.0])
return {
**{k: float(v)
for k, v in xyz.items()},
"rotation": {
"x": 0.0,
"y": float(rotation),
"z": 0.0
},
"horizon": float(horizon),
}
def randomize_agent_location(
self,
seed: int = None,
partial_position: Optional[Dict[str, float]] = None
) -> Dict[str, Union[Dict[str, float], float]]:
"""Teleports the agent to a random reachable location in the scene."""
if partial_position is None:
partial_position = {}
k = 0
state: Optional[Dict] = None
while k == 0 or (not self.last_action_success and k < 10):
# self.reset()
state = {
**self.random_reachable_state(seed=seed),
**partial_position
}
# get_logger().debug("picked target location {}".format(state))
self.controller.step("TeleportFull", **state)
k += 1
if not self.last_action_success:
get_logger().warning((
"Randomize agent location in scene {} and current random state {}"
" with seed {} and partial position {} failed in "
"10 attempts. Forcing the action.").format(
self.scene_name, state, seed, partial_position))
self.controller.step("TeleportFull", **state,
force_action=True) # type: ignore
assert self.last_action_success, "Force action failed with {}".format(
state)
# get_logger().debug("location after teleport full {}".format(self.agent_state()))
# self.controller.step("TeleportFull", **self.agent_state()) # TODO only for debug
# get_logger().debug("location after re-teleport full {}".format(self.agent_state()))
return self.agent_state()
def known_good_locations_list(self):
return self.known_good_locations[self.scene_name]
@property
def currently_reachable_points(self) -> List[Dict[str, float]]:
"""List of {"x": x, "y": y, "z": z} locations in the scene that are
currently reachable."""
self.controller.step(action="GetReachablePositions")
assert (
self.last_action_success
), f"Could not get reachable positions for reason {self.last_event.metadata['errorMessage']}."
return self.last_action_return
@property
def scene_name(self) -> str:
"""Current ai2thor scene."""
return self.controller.last_event.metadata["sceneName"].replace(
"_physics", "")
@property
def current_frame(self) -> np.ndarray:
"""Returns rgb image corresponding to the agent's egocentric view."""
return self.controller.last_event.frame
@property
def current_depth(self) -> np.ndarray:
"""Returns depth image corresponding to the agent's egocentric view."""
return self.controller.last_event.depth_frame
@property
def last_event(self) -> ai2thor.server.Event:
"""Last event returned by the controller."""
return self.controller.last_event
@property
def last_action(self) -> str:
"""Last action, as a string, taken by the agent."""
return self.controller.last_event.metadata["lastAction"]
@property
def last_action_success(self) -> bool:
"""Was the last action taken by the agent a success?"""
return self.controller.last_event.metadata["lastActionSuccess"]
@property
def last_action_return(self) -> Any:
"""Get the value returned by the last action (if applicable).
For an example of an action that returns a value, see
`"GetReachablePositions"`.
"""
return self.controller.last_event.metadata["actionReturn"]
def step(self, action_dict: Dict) -> ai2thor.server.Event:
"""Take a step in the ai2thor environment."""
return self.controller.step(**action_dict)
def stop(self):
"""Stops the ai2thor controller."""
try:
self.controller.stop()
except Exception as e:
get_logger().warning(str(e))
def all_objects(self) -> List[Dict[str, Any]]:
"""Return all object metadata."""
return self.controller.last_event.metadata["objects"]
def all_objects_with_properties(
self, properties: Dict[str, Any]) -> List[Dict[str, Any]]:
"""Find all objects with the given properties."""
objects = []
for o in self.all_objects():
satisfies_all = True
for k, v in properties.items():
if o[k] != v:
satisfies_all = False
break
if satisfies_all:
objects.append(o)
return objects
def visible_objects(self) -> List[Dict[str, Any]]:
"""Return all visible objects."""
return self.all_objects_with_properties({"visible": True})
class Ai2Thor():
def __init__(self):
self.visualize = False
self.verbose = False
self.save_imgs = True
self.plot_loss = True
# st()
# these are all map names
a = np.arange(1, 30)
b = np.arange(201, 231)
c = np.arange(301, 331)
d = np.arange(401, 431)
abcd = np.hstack((a, b, c, d))
mapnames = []
for i in list(abcd):
mapname = 'FloorPlan' + str(i)
mapnames.append(mapname)
train_len = int(0.9 * len(mapnames))
np.random.seed(1)
random.shuffle(mapnames)
self.mapnames_train = mapnames[:train_len]
self.mapnames_val = mapnames[train_len:]
# self.num_episodes = len(self.mapnames)
self.ignore_classes = []
# classes to save
self.include_classes = [
'ShowerDoor', 'Cabinet', 'CounterTop', 'Sink', 'Towel',
'HandTowel', 'TowelHolder', 'SoapBar', 'ToiletPaper',
'ToiletPaperHanger', 'HandTowelHolder', 'SoapBottle', 'GarbageCan',
'Candle', 'ScrubBrush', 'Plunger', 'SinkBasin', 'Cloth',
'SprayBottle', 'Toilet', 'Faucet', 'ShowerHead', 'Box', 'Bed',
'Book', 'DeskLamp', 'BasketBall', 'Pen', 'Pillow', 'Pencil',
'CellPhone', 'KeyChain', 'Painting', 'CreditCard', 'AlarmClock',
'CD', 'Laptop', 'Drawer', 'SideTable', 'Chair', 'Blinds', 'Desk',
'Curtains', 'Dresser', 'Watch', 'Television', 'WateringCan',
'Newspaper', 'FloorLamp', 'RemoteControl', 'HousePlant', 'Statue',
'Ottoman', 'ArmChair', 'Sofa', 'DogBed', 'BaseballBat',
'TennisRacket', 'VacuumCleaner', 'Mug', 'ShelvingUnit', 'Shelf',
'StoveBurner', 'Apple', 'Lettuce', 'Bottle', 'Egg', 'Microwave',
'CoffeeMachine', 'Fork', 'Fridge', 'WineBottle', 'Spatula',
'Bread', 'Tomato', 'Pan', 'Cup', 'Pot', 'SaltShaker', 'Potato',
'PepperShaker', 'ButterKnife', 'StoveKnob', 'Toaster',
'DishSponge', 'Spoon', 'Plate', 'Knife', 'DiningTable', 'Bowl',
'LaundryHamper', 'Vase', 'Stool', 'CoffeeTable', 'Poster',
'Bathtub', 'TissueBox', 'Footstool', 'BathtubBasin',
'ShowerCurtain', 'TVStand', 'Boots', 'RoomDecor', 'PaperTowelRoll',
'Ladle', 'Kettle', 'Safe', 'GarbageBag', 'TeddyBear',
'TableTopDecor', 'Dumbbell', 'Desktop', 'AluminumFoil', 'Window'
]
self.action_space = {
0: "MoveLeft",
1: "MoveRight",
2: "MoveAhead",
3: "MoveBack"
}
self.num_actions = len(self.action_space)
cfg_det = get_cfg()
cfg_det.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg_det.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.2 # set threshold for this model
cfg_det.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg_det.MODEL.DEVICE = 'cpu'
self.cfg_det = cfg_det
self.maskrcnn = DefaultPredictor(cfg_det)
self.conf_thresh_detect = 0.7 # for initially detecting a low confident object
self.conf_thresh_init = 0.8 # for after turning head toward object threshold
self.conf_thresh_end = 0.9 # if reach this then stop getting obs
self.BATCH_SIZE = 12
self.percentile = 70
self.max_iters = 100000
self.max_frames = 10
self.val_interval = 15
self.save_interval = 50
self.BATCH_SIZE = 1
self.percentile = 70
self.max_iters = 100000
self.max_frames = 1
self.val_interval = 1
self.save_interval = 1
self.small_classes = []
self.rot_interval = 5.0
self.radius_max = 3.5 #3 #1.75
self.radius_min = 1.0 #1.25
self.num_flat_views = 3
self.num_any_views = 7
self.num_views = 25
self.center_from_mask = False # get object centroid from maskrcnn (True) or gt (False)
self.obj_per_scene = 5
# self.origin_quaternion = np.quaternion(1, 0, 0, 0)
# self.origin_rot_vector = quaternion.as_rotation_vector(self.origin_quaternion)
# self.homepath = f'/home/nel/gsarch/aithor/data/test2'
self.homepath = '/home/sirdome/katefgroup/gsarch/ithor/data/test'
if not os.path.exists(self.homepath):
os.mkdir(self.homepath)
else:
val = input("Delete homepath? [y/n]: ")
if val == 'y':
import shutil
shutil.rmtree(self.homepath)
os.mkdir(self.homepath)
else:
print("ENDING")
assert (False)
self.log_freq = 1
self.log_dir = self.homepath + '/..' + '/log_cem' + '/aa'
if not os.path.exists(self.log_dir):
os.mkdir(self.log_dir)
MAX_QUEUE = 10 # flushes when this amount waiting
self.writer = SummaryWriter(self.log_dir,
max_queue=MAX_QUEUE,
flush_secs=60)
self.W = 256
self.H = 256
# self.fov = 90
# hfov = float(self.fov) * np.pi / 180.
# self.pix_T_camX = np.array([
# [(self.W/2.)*1 / np.tan(hfov / 2.), 0., 0., 0.],
# [0., (self.H/2.)*1 / np.tan(hfov / 2.), 0., 0.],
# [0., 0., 1, 0],
# [0., 0., 0, 1]])
# self.pix_T_camX[0,2] = self.W/2.
# self.pix_T_camX[1,2] = self.H/2.
self.fov = 90
self.camera_matrix = self.get_camera_matrix(self.W, self.H, self.fov)
self.K = self.get_habitat_pix_T_camX(self.fov)
self.init_network()
self.run_episodes()
def init_network(self):
self.cemnet = CEMNet(h1=32,
h2=64,
fc_dim=1024,
num_actions=self.num_actions)
self.loss = nn.CrossEntropyLoss()
self.optimizer = torch.optim.Adam(params=self.cemnet.parameters(),
lr=0.01)
def batch_iteration(self, mapnames, NN, BATCH_SIZE):
batch = {"rewards": [], "obs": [], "actions": []}
episode_rewards = 0.0
iter_idx = 0
while True:
mapname = np.random.choice(mapnames)
# self.basepath = self.homepath + f"/{mapname}_{episode}"
# print("BASEPATH: ", self.basepath)
# # self.basepath = f"/hdd/ayushj/habitat_data/{mapname}_{episode}"
# if not os.path.exists(self.basepath):
# os.mkdir(self.basepath)
self.controller = Controller(
scene=mapname,
gridSize=0.25,
width=self.W,
height=self.H,
fieldOfView=self.fov,
renderObjectImage=True,
renderDepthImage=True,
)
episode_rewards, obs, actions = self.run()
print("Total reward for train batch # ", iter_idx, " :",
episode_rewards)
self.controller.stop()
time.sleep(1)
if episode_rewards is None:
print("NO EPISODE REWARDS.. SKIPPING BATCH INSTANCE")
continue
batch["rewards"].append(episode_rewards)
batch["obs"].append(obs)
batch["actions"].append(actions)
iter_idx += 1
print(len(batch["rewards"]))
if len(batch["rewards"]) == BATCH_SIZE:
yield batch
iter_idx = 0
batch = {"rewards": [], "obs": [], "actions": []}
if len(batch["rewards"]) > BATCH_SIZE:
st()
def elite_batch(self, batch, percentile):
rewards = np.array(batch["rewards"])
obs = batch["obs"]
actions = batch["actions"]
rewards_mean = float(np.mean(rewards))
rewards_boundary = np.percentile(rewards, percentile)
print("Reward boundary: ", rewards_boundary)
rewards_mean = float(np.mean(rewards))
training_obs = []
training_actions = []
for idx in range(rewards.shape[0]):
reward_idx = rewards[idx]
if reward_idx < rewards_boundary:
continue
training_obs.extend(obs[idx])
training_actions.extend(actions[idx])
obs_tensor = torch.FloatTensor(training_obs).permute(0, 3, 1, 2)
act_tensor = torch.LongTensor(training_actions)
return obs_tensor, act_tensor, rewards_mean, rewards_boundary
def run_val(self, mapnames, BATCH_SIZE):
batch = {"rewards": [], "obs": [], "actions": []}
episode_rewards = 0.0
episode_steps = []
iter_idx = 0
while True:
mapname = np.random.choice(mapnames)
# self.basepath = self.homepath + f"/{mapname}_{episode}"
# print("BASEPATH: ", self.basepath)
# # self.basepath = f"/hdd/ayushj/habitat_data/{mapname}_{episode}"
# if not os.path.exists(self.basepath):
# os.mkdir(self.basepath)
self.controller = Controller(
scene=mapname,
gridSize=0.25,
width=self.W,
height=self.H,
fieldOfView=self.fov,
renderObjectImage=True,
renderDepthImage=True,
)
episode_rewards, obs, actions = self.run()
print("Total reward for val batch # ", iter_idx, " :",
episode_rewards)
self.controller.stop()
time.sleep(1)
if episode_rewards is None:
print("NO EPISODE REWARDS.. SKIPPING BATCH INSTANCE")
continue
batch["rewards"].append(episode_rewards)
batch["obs"].append(obs)
batch["actions"].append(actions)
train, labels, mean_rewards, boudary = self.elite_batch(
batch, self.percentile)
with torch.no_grad():
scores = self.cemnet(train, softmax=False)
loss_value = self.loss(scores, labels)
iter_idx += 1
if len(batch["rewards"]) == BATCH_SIZE:
break
return loss_value, mean_rewards
def run_episodes(self):
self.ep_idx = 0
# self.objects = []
train_loss = []
val_loss = []
train_iters = []
val_iters = []
mean_rewards_train = []
mean_rewards_val = []
if self.plot_loss:
# plt.figure(1) # loss
# plt.figure(2) # reward
fig, (ax1, ax2) = plt.subplots(1, 2)
for iteration, batch in enumerate(
self.batch_iteration(self.mapnames_train, self.cemnet,
self.BATCH_SIZE)):
print("ITERATION #", iteration)
self.summ_writer = utils.improc.Summ_writer(writer=self.writer,
global_step=iteration,
log_freq=self.log_freq,
fps=8,
just_gif=True)
train, labels, mean_rewards, boudary = self.elite_batch(
batch, self.percentile)
mean_rewards_train.append(mean_rewards)
self.optimizer.zero_grad()
scores = self.cemnet(train, softmax=False)
loss_value = self.loss(scores, labels)
train_loss.append(float(loss_value.clone().detach().cpu().numpy()))
train_iters.append(iteration)
back_v = loss_value.backward()
self.optimizer.step()
print('rewards mean = ', mean_rewards)
print('')
if iteration >= self.max_iters:
print("MAX ITERS REACHED")
self.writer.close()
break
# if self.plot_loss:
# plt.figure(1) # loss
# fig.clf()
# plt.figure(2) # reward
# plt.clf()
if iteration % self.val_interval == 0:
loss_val, mean_reward_v = self.run_val(self.mapnames_val,
self.BATCH_SIZE)
val_loss.append(float(loss_val.clone().detach().cpu().numpy()))
val_iters.append(iteration)
mean_rewards_val.append(mean_reward_v)
if self.plot_loss:
# plt.figure(1) # loss
ax1.plot(val_iters, val_loss, color='red')
# plt.figure(2) # reward
ax2.plot(val_iters, mean_rewards_val, color='red')
self.summ_writer.summ_scalar('unscaled_loss_val', loss_val)
self.summ_writer.summ_scalar('unscaled_mean_reward_val',
mean_reward_v)
if iteration % self.save_interval == 0:
PATH = self.homepath + f'/checkpoint{iteration}.tar'
torch.save(self.cemnet.state_dict(), PATH)
if self.plot_loss:
self.summ_writer.summ_scalar('unscaled_loss', loss_value)
self.summ_writer.summ_scalar('unscaled_mean_reward',
mean_rewards)
# plt.figure(1) # loss
ax1.plot(train_iters, train_loss, color='blue')
ax1.set(xlabel='iterations', ylabel='loss')
# ax1.xlabel('iterations')
# ax1.ylabel('loss')
# plt_name = '/home/nel/gsarch/aithor/data/test/loss.png'
# fig.savefig(plt_name)
# plt.figure(2) # reward
ax2.plot(train_iters, mean_rewards_train, color='blue')
ax2.set(xlabel='iterations', ylabel='reward')
# ax2.xlabel('iterations')
# ax2.ylabel('reward')
plt_name = os.path.join(self.homepath, 'train.png')
fig.savefig(plt_name)
# if mean_rewards > 500:
# print('Accomplished!')
# break
def save_datapoint(self, observations, data_path, viewnum, flat_view):
if self.verbose:
print("Print Sensor States.", self.agent.state.sensor_states)
rgb = observations["color_sensor"]
semantic = observations["semantic_sensor"]
# st()
depth = observations["depth_sensor"]
agent_pos = observations["positions"]
agent_rot = observations["rotations"]
# Assuming all sensors have same extrinsics
color_sensor_pos = observations["positions"]
color_sensor_rot = observations["rotations"]
#print("POS ", agent_pos)
#print("ROT ", color_sensor_rot)
object_list = observations['object_list']
# print(viewnum, agent_pos)
# print(agent_rot)
if False:
plt.imshow(rgb)
plt_name = f'/home/nel/gsarch/aithor/data/test/img_mask{viewnum}.png'
plt.savefig(plt_name)
save_data = {
'flat_view': flat_view,
'objects_info': object_list,
'rgb_camX': rgb,
'depth_camX': depth,
'sensor_pos': color_sensor_pos,
'sensor_rot': color_sensor_rot
}
with open(os.path.join(data_path, str(viewnum) + ".p"), 'wb') as f:
pickle.dump(save_data, f)
f.close()
def quat_from_angle_axis(self, theta: float,
axis: np.ndarray) -> np.quaternion:
r"""Creates a quaternion from angle axis format
:param theta: The angle to rotate about the axis by
:param axis: The axis to rotate about
:return: The quaternion
"""
axis = axis.astype(np.float)
axis /= np.linalg.norm(axis)
return quaternion.from_rotation_vector(theta * axis)
def safe_inverse_single(self, a):
r, t = self.split_rt_single(a)
t = np.reshape(t, (3, 1))
r_transpose = r.T
inv = np.concatenate([r_transpose, -np.matmul(r_transpose, t)], 1)
bottom_row = a[3:4, :] # this is [0, 0, 0, 1]
# bottom_row = torch.tensor([0.,0.,0.,1.]).view(1,4)
inv = np.concatenate([inv, bottom_row], 0)
return inv
def split_rt_single(self, rt):
r = rt[:3, :3]
t = np.reshape(rt[:3, 3], 3)
return r, t
def eul2rotm(self, rx, ry, rz):
# inputs are shaped B
# this func is copied from matlab
# R = [ cy*cz sy*sx*cz-sz*cx sy*cx*cz+sz*sx
# cy*sz sy*sx*sz+cz*cx sy*cx*sz-cz*sx
# -sy cy*sx cy*cx]
# rx = np.expand_dims(rx, axis=1)
# ry = np.expand_dims(ry, axis=1)
# rz = np.expand_dims(rz, axis=1)
# st()
# these are B x 1
sinz = np.sin(rz)
siny = np.sin(ry)
sinx = np.sin(rx)
cosz = np.cos(rz)
cosy = np.cos(ry)
cosx = np.cos(rx)
r11 = cosy * cosz
r12 = sinx * siny * cosz - cosx * sinz
r13 = cosx * siny * cosz + sinx * sinz
r21 = cosy * sinz
r22 = sinx * siny * sinz + cosx * cosz
r23 = cosx * siny * sinz - sinx * cosz
r31 = -siny
r32 = sinx * cosy
r33 = cosx * cosy
r = np.array([[r11, r12, r13], [r21, r22, r23], [r31, r32, r33]])
# r1 = np.stack([r11,r12,r13],axis=2)
# r2 = np.stack([r21,r22,r23],axis=2)
# r3 = np.stack([r31,r32,r33],axis=2)
# r = np.concatenate([r1,r2,r3],axis=1)
return r
def rotm2eul(self, r):
# r is Bx3x3, or Bx4x4
r00 = r[0, 0]
r10 = r[1, 0]
r11 = r[1, 1]
r12 = r[1, 2]
r20 = r[2, 0]
r21 = r[2, 1]
r22 = r[2, 2]
## python guide:
# if sy > 1e-6: # singular
# x = math.atan2(R[2,1] , R[2,2])
# y = math.atan2(-R[2,0], sy)
# z = math.atan2(R[1,0], R[0,0])
# else:
# x = math.atan2(-R[1,2], R[1,1])
# y = math.atan2(-R[2,0], sy)
# z = 0
sy = np.sqrt(r00 * r00 + r10 * r10)
cond = (sy > 1e-6)
rx = np.where(cond, np.arctan2(r21, r22), np.arctan2(-r12, r11))
ry = np.where(cond, np.arctan2(-r20, sy), np.arctan2(-r20, sy))
rz = np.where(cond, np.arctan2(r10, r00), np.zeros_like(r20))
# rx = torch.atan2(r21, r22)
# ry = torch.atan2(-r20, sy)
# rz = torch.atan2(r10, r00)
# rx[cond] = torch.atan2(-r12, r11)
# ry[cond] = torch.atan2(-r20, sy)
# rz[cond] = 0.0
return rx, ry, rz
def get_habitat_pix_T_camX(self, fov):
hfov = float(self.fov) * np.pi / 180.
pix_T_camX = np.array([[
(self.W / 2.) * 1 / np.tan(hfov / 2.), 0., 0., 0.
], [0., (self.H / 2.) * 1 / np.tan(hfov / 2.), 0., 0.], [0., 0., 1, 0],
[0., 0., 0, 1]])
return pix_T_camX
def get_camera_matrix(self, width, height, fov):
"""Returns a camera matrix from image size and fov."""
xc = (width - 1.) / 2.
zc = (height - 1.) / 2.
f = (width / 2.) / np.tan(np.deg2rad(fov / 2.))
camera_matrix = {'xc': xc, 'zc': zc, 'f': f}
camera_matrix = Namespace(**camera_matrix)
return camera_matrix
def run2(self):
event = self.controller.step('GetReachablePositions')
for obj in event.metadata['objects']:
if obj['objectType'] not in self.objects:
self.objects.append(obj['objectType'])
def get_detectron_conf_center_obj(self, im, frame=None):
im = Image.fromarray(im, mode="RGB")
im = cv2.cvtColor(np.asarray(im), cv2.COLOR_RGB2BGR)
# plt.imshow(im)
# plt.show()
outputs = self.maskrcnn(im)
v = Visualizer(im[:, :, ::-1],
MetadataCatalog.get(self.cfg_det.DATASETS.TRAIN[0]),
scale=1.2)
out = v.draw_instance_predictions(outputs['instances'].to("cpu"))
seg_im = out.get_image()
if False:
plt.figure(1)
plt.clf()
plt.imshow(seg_im)
plt_name = self.homepath + '/seg{frame}.png'
plt.savefig(plt_name)
# plt.show()
pred_masks = outputs['instances'].pred_masks
pred_scores = outputs['instances'].scores
pred_classes = outputs['instances'].pred_classes
len_pad = 5
W2_low = self.W // 2 - len_pad
W2_high = self.W // 2 + len_pad
H2_low = self.H // 2 - len_pad
H2_high = self.H // 2 + len_pad
ind_obj = None
max_overlap = 0
for idx in range(pred_masks.shape[0]):
pred_mask_cur = pred_masks[idx]
pred_masks_center = pred_mask_cur[W2_low:W2_high, H2_low:H2_high]
# print(torch.sum(pred_masks_center))
if torch.sum(pred_masks_center) > max_overlap:
ind_obj = idx
max_overlap = torch.sum(pred_masks_center)
if ind_obj is None:
return None
print("OBJ CLASS ID=",
int(pred_classes[ind_obj].detach().cpu().numpy()))
# pred_boxes = outputs['instances'].pred_boxes.tensor
# pred_classes = outputs['instances'].pred_classes
# pred_scores = outputs['instances'].scores
obj_score = pred_scores[ind_obj]
return obj_score
def detect_object_centroid(self, im, event):
im = Image.fromarray(im, mode="RGB")
im = cv2.cvtColor(np.asarray(im), cv2.COLOR_RGB2BGR)
# plt.imshow(im)
# plt.show()
outputs = self.maskrcnn(im)
v = Visualizer(im[:, :, ::-1],
MetadataCatalog.get(self.cfg_det.DATASETS.TRAIN[0]),
scale=1.2)
out = v.draw_instance_predictions(outputs['instances'].to("cpu"))
seg_im = out.get_image()
if False:
plt.figure(1)
plt.clf()
plt.imshow(seg_im)
plt_name = self.homepath + '/seg_init.png'
plt.savefig(plt_name)
# plt.show()
pred_masks = outputs['instances'].pred_masks
pred_boxes = outputs['instances'].pred_boxes.tensor
pred_classes = outputs['instances'].pred_classes
pred_scores = outputs['instances'].scores
# obj_ids = []
obj_catids = []
obj_scores = []
obj_masks = []
# obj_all_catids = []
# obj_all_scores = []
# obj_all_boxes = []
for segs in range(len(pred_masks)):
if pred_scores[segs] <= self.conf_thresh_detect:
# obj_ids.append(segs)
obj_catids.append(pred_classes[segs].item())
obj_scores.append(pred_scores[segs].item())
obj_masks.append(pred_masks[segs])
# obj_all_catids.append(pred_classes[segs].item())
# obj_all_scores.append(pred_scores[segs].item())
# y, x = torch.where(pred_masks[segs])
# pred_box = torch.Tensor([min(y), min(x), max(y), max(x)]) # ymin, xmin, ymax, xmaxs
# obj_all_boxes.append(pred_box)
# print("MASKS ", len(pred_masks))
# print("VALID ", len(obj_scores))
# print(obj_scores)
# print(pred_scores.shape)
eulers_xyz_rad = np.radians(
np.array([
event.metadata['agent']['cameraHorizon'],
event.metadata['agent']['rotation']['y'], 0.0
]))
rx = eulers_xyz_rad[0]
ry = eulers_xyz_rad[1]
rz = eulers_xyz_rad[2]
rotation_ = self.eul2rotm(-rx, -ry, rz)
translation_ = np.array(
list(event.metadata['agent']['position'].values())) + np.array(
[0.0, 0.675, 0.0])
# need to invert since z is positive here by convention
translation_[2] = -translation_[2]
T_world_cam = np.eye(4)
T_world_cam[0:3, 0:3] = rotation_
T_world_cam[0:3, 3] = translation_
if not obj_masks:
return None
elif self.center_from_mask:
# want an object not on the edges of the image
sum_interior = 0
while sum_interior == 0:
if len(obj_masks) == 0:
return None
random_int = np.random.randint(low=0, high=len(obj_masks))
obj_mask_focus = obj_masks.pop(random_int)
print("OBJECT ID INIT=", obj_catids[random_int])
sum_interior = torch.sum(obj_mask_focus[50:self.W - 50,
50:self.H - 50])
depth = event.depth_frame
xs, ys = np.meshgrid(np.linspace(-1 * 256 / 2., 1 * 256 / 2., 256),
np.linspace(1 * 256 / 2., -1 * 256 / 2., 256))
depth = depth.reshape(1, 256, 256)
xs = xs.reshape(1, 256, 256)
ys = ys.reshape(1, 256, 256)
xys = np.vstack(
(xs * depth, ys * depth, -depth, np.ones(depth.shape)))
xys = xys.reshape(4, -1)
xy_c0 = np.matmul(np.linalg.inv(self.K), xys)
xyz = xy_c0.T[:, :3].reshape(256, 256, 3)
xyz_obj_masked = xyz[obj_mask_focus]
xyz_obj_masked = np.matmul(
rotation_, xyz_obj_masked.T) + translation_.reshape(3, 1)
xyz_obj_mid = np.mean(xyz_obj_masked, axis=1)
xyz_obj_mid[2] = -xyz_obj_mid[2]
else:
# want an object not on the edges of the image
sum_interior = 0
while True:
if len(obj_masks) == 0:
return None
random_int = np.random.randint(low=0, high=len(obj_masks))
obj_mask_focus = obj_masks.pop(random_int)
print("OBJECT ID INIT=", obj_catids[random_int])
sum_interior = torch.sum(obj_mask_focus[50:self.W - 50,
50:self.H - 50])
if sum_interior < 500:
continue # exclude too small objects
pixel_locs_obj = np.where(obj_mask_focus)
x_mid = np.round(np.median(pixel_locs_obj[1]) / self.W, 4)
y_mid = np.round(np.median(pixel_locs_obj[0]) / self.H, 4)
if False:
plt.figure(1)
plt.clf()
plt.imshow(obj_mask_focus)
plt.plot(np.median(pixel_locs_obj[1]),
np.median(pixel_locs_obj[0]), 'x')
plt_name = self.homepath + '/seg_mask.png'
plt.savefig(plt_name)
event = self.controller.step('TouchThenApplyForce',
x=x_mid,
y=y_mid,
handDistance=1000000.0,
direction=dict(x=0.0,
y=0.0,
z=0.0),
moveMagnitude=0.0)
obj_focus_id = event.metadata['actionReturn']['objectId']
xyz_obj_mid = None
for o in event.metadata['objects']:
if o['objectId'] == obj_focus_id:
xyz_obj_mid = np.array(
list(o['axisAlignedBoundingBox']
['center'].values()))
if xyz_obj_mid is not None:
break
# if xyz_obj_mid is None:
# st()
print("MIDPOINT=", xyz_obj_mid)
return xyz_obj_mid
# semantic = event.instance_segmentation_frame
# object_id_to_color = event.object_id_to_color
# color_to_object_id = event.color_to_object_id
# obj_ids = np.unique(semantic.reshape(-1, semantic.shape[2]), axis=0)
# instance_masks = event.instance_masks
# instance_detections2D = event.instance_detections2D
# obj_metadata_IDs = []
# for obj_m in event.metadata['objects']: #objects:
# obj_metadata_IDs.append(obj_m['objectId'])
# object_list = []
# for obj_idx in range(obj_ids.shape[0]):
# try:
# obj_color = tuple(obj_ids[obj_idx])
# object_id = color_to_object_id[obj_color]
# except:
# # print("Skipping ", object_id)
# continue
# if object_id not in obj_metadata_IDs:
# # print("Skipping ", object_id)
# continue
# obj_meta_index = obj_metadata_IDs.index(object_id)
# obj_meta = event.metadata['objects'][obj_meta_index]
# obj_category_name = obj_meta['objectType']
# # continue if not visible or not in include classes
# if obj_category_name not in self.include_classes: # or not obj_meta['visible']:
# continue
# obj_instance_mask = instance_masks[object_id]
# obj_instance_detection2D = instance_detections2D[object_id] # [start_x, start_y, end_x, end_y]
# obj_instance_detection2D = np.array([obj_instance_detection2D[1], obj_instance_detection2D[0], obj_instance_detection2D[3], obj_instance_detection2D[2]]) # ymin, xmin, ymax, xmax
# if True:
# print(object_id)
# print(np.array(list(obj_meta['axisAlignedBoundingBox']['center'].values())))
# # plt.imshow(obj_instance_mask)
# # plt_name = f'/home/nel/gsarch/aithor/data/test/img_mask{s}.png'
# # plt.savefig(plt_name)
# obj_center_axisAligned = np.array(list(obj_meta['axisAlignedBoundingBox']['center'].values()))
# obj_center_axisAligned[2] = -obj_center_axisAligned[2]
# obj_size_axisAligned = np.array(list(obj_meta['axisAlignedBoundingBox']['size'].values()))
def get_rotation_to_obj(self, obj_center, pos_s):
# YAW calculation - rotate to object
agent_to_obj = np.squeeze(obj_center) - pos_s
agent_local_forward = np.array([0, 0, 1.0])
flat_to_obj = np.array([agent_to_obj[0], 0.0, agent_to_obj[2]])
flat_dist_to_obj = np.linalg.norm(flat_to_obj)
flat_to_obj /= flat_dist_to_obj
det = (flat_to_obj[0] * agent_local_forward[2] -
agent_local_forward[0] * flat_to_obj[2])
turn_angle = math.atan2(det, np.dot(agent_local_forward, flat_to_obj))
turn_yaw = np.degrees(turn_angle)
turn_pitch = -np.degrees(math.atan2(agent_to_obj[1], flat_dist_to_obj))
return turn_yaw, turn_pitch
def run(self):
event = self.controller.step('GetReachablePositions')
if not event.metadata['reachablePositions']:
# Different versions this is empty/full
event = self.controller.step(action='MoveAhead')
self.nav_pts = event.metadata['reachablePositions']
self.nav_pts = np.array([list(d.values()) for d in self.nav_pts])
np.random.seed(1)
# objects = np.random.choice(event.metadata['objects'], self.obj_per_scene, replace=False)
objects = event.metadata['objects']
objects_inds = np.arange(len(event.metadata['objects']))
np.random.shuffle(objects_inds)
# objects = np.random.shuffle(event.metadata['objects'])
# for obj in event.metadata['objects']: #objects:
# print(obj['name'])
# objects = objects[0]
successes = 0
meta_obj_idx = 0
while True: #successes < self.obj_per_scene and meta_obj_idx <= len(event.metadata['objects']) - 1:
if meta_obj_idx > len(event.metadata['objects']) - 1:
print("OUT OF OBJECT... RETURNING")
return None, None, None
obj = objects[objects_inds[meta_obj_idx]]
meta_obj_idx += 1
print("Center object is ", obj['objectType'])
# if obj['name'] in ['Microwave_b200e0bc']:
# print(obj['name'])
# else:
# continue
# print(obj['name'])
if obj['objectType'] not in self.include_classes:
print("Continuing... Invalid Object")
continue
# Calculate distance to object center
obj_center = np.array(
list(obj['axisAlignedBoundingBox']['center'].values()))
obj_center = np.expand_dims(obj_center, axis=0)
distances = np.sqrt(np.sum((self.nav_pts - obj_center)**2, axis=1))
# Get points with r_min < dist < r_max
valid_pts = self.nav_pts[np.where(
(distances > self.radius_min) * (distances < self.radius_max))]
# add height from center of agent to camera
rand_pos_int = np.random.randint(low=0, high=valid_pts.shape[0])
pos_s = valid_pts[rand_pos_int]
pos_s[1] = pos_s[1] + 0.675
turn_yaw, turn_pitch = self.get_rotation_to_obj(obj_center, pos_s)
event = self.controller.step('TeleportFull',
x=pos_s[0],
y=pos_s[1],
z=pos_s[2],
rotation=dict(x=0.0,
y=int(turn_yaw),
z=0.0),
horizon=int(turn_pitch))
rgb = event.frame
# get object center of a low confidence object
obj_center = self.detect_object_centroid(rgb, event)
if obj_center is None:
print("NO LOW CONFIDENCE OBJECTS... SKIPPING...")
continue
# initialize object in center of FOV
turn_yaw, turn_pitch = self.get_rotation_to_obj(obj_center, pos_s)
event = self.controller.step('TeleportFull',
x=pos_s[0],
y=pos_s[1],
z=pos_s[2],
rotation=dict(x=0.0,
y=int(turn_yaw),
z=0.0),
horizon=int(turn_pitch))
rgb = event.frame
init_conf = self.get_detectron_conf_center_obj(rgb)
if init_conf is None:
print("Nothing detected in the center... SKIPPING")
continue
conf_prev = init_conf
if init_conf > self.conf_thresh_init:
print("HIGH INITIAL CONFIDENCE... SKIPPING...")
continue
obs = []
actions = []
episode_rewards = 0.0
frame = 0
while True:
rgb_tensor = torch.FloatTensor([rgb]).permute(0, 3, 1, 2)
with torch.no_grad():
actions_probability = self.cemnet(rgb_tensor, softmax=True)
act_proba = actions_probability.data.numpy()[0]
action_ind = np.random.choice(len(act_proba), p=act_proba)
action = self.action_space[action_ind]
event = self.controller.step(action)
rgb = event.frame
conf_cur = self.get_detectron_conf_center_obj(rgb, frame)
if conf_cur is None:
reward = -1 #-0.2 # fixed negative reward for no detection
conf_cur = conf_prev
# conf_prev = conf_prev # use same conf
else:
# reward = (conf_cur - conf_prev)/(1 - init_conf) # normalize by intial confidence to account for differences in starting confidence
diff = conf_cur - conf_prev
if diff > 0:
reward = 1
elif diff == 0:
reward = 0
else:
reward = -1
conf_prev = conf_cur
episode_rewards += reward
obs.append(rgb)
actions.append(action_ind)
if conf_cur > self.conf_thresh_end:
print("CONFIDENCE THRESHOLD REACHED!")
print("End confidence: ", conf_cur)
break
if frame >= self.max_frames - 1:
print("MAX FRAMES REACHED")
print("End confidence: ", conf_cur)
break
frame += 1
return episode_rewards, obs, actions
class Actor:
"""
The basic THOR actor class that we use. This process can be controlled via pipes. The main functionality is to
provide RGBD views from THOR scenes, receive and perform pokes in these scenes, and provide the feedback.
The feedback can be in the form of raw images, or already processed (for performance). The agent can run in standard
mode (images are only extracted from iTHOR when the scene is at rest), or in video mode (images are extracted in
real time; very slow).
Additionally, this class can provide superpixel segmentations of the THOR scenes.
"""
def __init__(self, pipe, gpu, configs, ground_truth=0, run=True):
self.pipe = pipe
self.gc, self.ac = configs
self.ground_truth = ground_truth
self.directions = [
dict(x=a, y=b, z=c) for a in [-1, 1] for b in [1, 2]
for c in [-1, 1]
]
self.controller = Controller(
x_display='0.%d' % gpu,
visibilityDistance=self.ac.visibilityDistance,
renderDepthImage=self.gc.depth or ground_truth > 0,
renderClassImage=ground_truth > 0,
renderObjectImage=ground_truth > 0)
self.grid = [(x, y) for x in range(self.gc.grid_size)
for y in range(self.gc.grid_size)]
self.depth_correction = self._make_depth_correction(
self.gc.resolution, self.gc.resolution, 90)
self.kernel_size = (self.ac.check_change_kernel.shape[0] - 1) // 2
if run:
self.run()
@staticmethod
def _make_depth_correction(height, width, vertical_field_of_view):
focal_length_in_pixels = height / np.tan(
vertical_field_of_view / 2 * np.pi / 180) / 2
xs = (np.arange(height).astype(np.float32) - (height - 1) / 2)**2
ys = (np.arange(width).astype(np.float32) - (width - 1) / 2)**2
return np.sqrt(1 +
(xs[:, None] + ys[None, :]) / focal_length_in_pixels**2)
def run(self):
while True:
action = self.pipe.recv()
if action == 'stop':
break
elif type(action) == dict:
data = dict()
image, superpixel, metadata = self.set_scene(
action) if self.ac.use_dataset else self.set_scene0(action)
data['image'] = image
if self.gc.superpixels:
data['superpixels'] = superpixel
if self.ground_truth:
data['metadata'] = metadata
self.pipe.send(data)
if type(action) == list:
feedback = self.poke(action, superpixel)
self.pipe.send(feedback)
self.controller.stop()
self.pipe.send('stop')
def set_scene(self, scene_data):
if self.ac.video_mode:
self.controller.step(action='UnpausePhysicsAutoSim')
scene, position, rotation, horizon, seed = (scene_data['scene'],
scene_data['position'],
scene_data['rotation'],
scene_data['horizon'],
scene_data['seed'])
self.controller.reset(scene)
self.controller.step(action='InitialRandomSpawn',
seed=seed,
forceVisible=True,
numPlacementAttempts=5)
self.controller.step(action='MakeAllObjectsMoveable')
event = self.controller.step(action='TeleportFull',
x=position['x'],
y=position['y'],
z=position['z'],
rotation=rotation,
horizon=horizon)
image = event.frame
if self.gc.superpixels:
superpixel = felzenszwalb(
image, scale=200, sigma=.5,
min_size=200)[::self.gc.stride, ::self.gc.stride].astype(
np.int32)
else:
superpixel = None
if self.gc.depth:
depth = event.depth_frame
if self.gc.correct_depth:
depth = (
depth - .1
) * self.depth_correction # convert depth from camera plane to distance from camera
image = (image, depth)
if self.ground_truth:
ground_truth = self.compute_ground_truth(event)
metadata = (ground_truth, seed, position, rotation, horizon)
else:
metadata = None
return image, superpixel, metadata
def set_scene0(self, scene):
if self.ac.video_mode:
self.controller.step(action='UnpausePhysicsAutoSim')
self.controller.reset(scene['scene'])
seed = randint(0, 2**30)
self.controller.step(action='InitialRandomSpawn',
seed=seed,
forceVisible=True,
numPlacementAttempts=5)
self.controller.step(action='MakeAllObjectsMoveable')
event = self.controller.step(action='GetReachablePositions')
positions = deepcopy(event.metadata['reachablePositions'])
position, rotation, horizon = choice(positions), choice(
[0., 90., 180., 270.]), choice([-30., 0., 30., 60.])
event = self.controller.step(action='TeleportFull',
x=position['x'],
y=position['y'],
z=position['z'],
rotation=rotation,
horizon=horizon)
if self.gc.respawn_until_object:
contains_interactable_object = len([
o for o in event.metadata['objects']
if o['visible'] and (o['moveable'] or o['pickupable'])
]) > 0
while not contains_interactable_object:
position, rotation, horizon = choice(positions), choice([0., 90., 180., 270.]), \
choice([-30., 0., 30., 60.])
event = self.controller.step(action='TeleportFull',
x=position['x'],
y=position['y'],
z=position['z'],
rotation=rotation,
horizon=horizon)
contains_interactable_object = len([
o for o in event.metadata['objects']
if o['visible'] and (o['moveable'] or o['pickupable'])
]) > 0
image = event.frame
if self.gc.superpixels:
superpixel = felzenszwalb(
image, scale=200, sigma=.5,
min_size=200)[::self.gc.stride, ::self.gc.stride].astype(
np.int32)
else:
superpixel = None
if self.gc.depth:
depth = event.depth_frame
if self.gc.correct_depth:
depth = (
depth - .1
) * self.depth_correction # convert depth from camera plane to distance from camera
image = (image, depth)
if self.ground_truth:
ground_truth = self.compute_ground_truth(event)
metadata = (ground_truth, seed, position, rotation, horizon)
else:
metadata = None
return image, superpixel, metadata
def poke(self, action, superpixel):
feedback = []
if self.ac.video_mode:
im1 = self.controller.step(action='PausePhysicsAutoSim').frame
else:
im1 = self.controller.step(action='Pass').frame
for poke_point in action:
if self.ac.instance_only:
im2 = self.touch(poke_point['point'],
self.ac.force_buckets[-1])
poke_feedback = self.compute_feedback(im1, im2,
poke_point['point'],
superpixel)
elif self.ac.scaleable:
poke_feedback, im2 = self.touch_with_forces(
im1, poke_point, superpixel)
else:
poke_feedback, im2 = self.touch_with_forces_nonscaleable(
im1, poke_point, superpixel)
feedback.append(poke_feedback)
im1 = im2[-1]
return feedback
def touch_with_forces(self, im1, point_and_force, superpixel):
direction = choice(self.directions)
point, force = point_and_force['point'], point_and_force['force']
smaller_force = max(force - 1, 0)
im2 = self.touch(point, self.ac.force_buckets[smaller_force],
direction)
vis_feedback = self.compute_feedback(im1, im2, point, superpixel)
if self.get_score(vis_feedback, point) > 1.5:
return (vis_feedback, -1 if force > 0 else 0), im2
im1 = im2[-1]
im2 = self.touch(point, self.ac.force_buckets[force], direction)
vis_feedback = self.compute_feedback(im1, im2, point, superpixel)
if self.get_score(vis_feedback, point) > 1.5:
return (vis_feedback, 0), im2
if force < len(self.ac.force_buckets) - 1:
im1 = im2[-1]
im2 = self.touch(point, self.ac.force_buckets[-1], direction)
vis_feedback = self.compute_feedback(im1, im2, point, superpixel)
if self.get_score(vis_feedback, point) > 1.5:
return (vis_feedback, 1), im2
return (vis_feedback, 2), im2
def touch_with_forces_nonscaleable(self, im1, point_and_force, superpixel):
point = point_and_force['point']
smaller_force = 0
im2 = self.touch(point, self.ac.force_buckets[smaller_force])
vis_feedback = self.compute_feedback(im1, im2, point, superpixel)
if self.get_score(vis_feedback, point) > 1.5:
return (vis_feedback, 0), im2
im1 = im2[-1]
im2 = self.touch(point, self.ac.force_buckets[1])
vis_feedback = self.compute_feedback(im1, im2, point, superpixel)
if self.get_score(vis_feedback, point) > 1.5:
return (vis_feedback, 1), im2
im1 = im2[-1]
im2 = self.touch(point, self.ac.force_buckets[-1])
vis_feedback = self.compute_feedback(im1, im2, point, superpixel)
if self.get_score(vis_feedback, point) > 1.5:
return (vis_feedback, 1), im2
return (vis_feedback, 2), im2
def touch(self, point, force, direction=None):
y, x = point # x axis (=first axis) in numpy is y-axis (=second axis) in images (THOR)
if direction is None:
direction = choice(self.directions)
im2 = []
if self.ac.video_mode:
self.controller.step(
dict(action='TouchThenApplyForce',
x=x / self.gc.grid_size + 1 / 2 / self.gc.grid_size,
y=y / self.gc.grid_size + 1 / 2 / self.gc.grid_size,
direction=direction,
handDistance=self.ac.handDistance,
moveMagnitude=force))
im2.append(
self.controller.step(action='AdvancePhysicsStep',
timeStep=0.01).frame)
for _ in range(25):
im2.append(
self.controller.step(action='AdvancePhysicsStep',
timeStep=0.05).frame)
else:
im2.append(
self.controller.step(
dict(action='TouchThenApplyForce',
x=x / self.gc.grid_size + 1 / 2 / self.gc.grid_size,
y=y / self.gc.grid_size + 1 / 2 / self.gc.grid_size,
direction=direction,
handDistance=self.ac.handDistance,
moveMagnitude=force)).frame)
return im2
def compute_feedback(self, im1, im2, poke_point, superpixel):
if self.ac.raw_feedback:
return im1, im2
if self.ac.hsv:
im1, im2 = rgb2hsv(im1), [rgb2hsv(im) for im in im2]
else:
im1, im2 = im1.astype(
np.float32), [im.astype(np.float32) for im in im2]
if self.ac.video_mode:
diff = self.pca_diff(
[im1] + im2) if self.ac.pca else self.compute_mean_of([im1] +
im2)
else:
im2 = im2[0]
diff = [im1 - im2]
if self.ac.smooth_mask:
mask = self.make_smooth_mask(diff, [im1] + im2)
else:
mask = self.make_mask(diff)
if superpixel is not None and self.ac.superpixel_postprocessed_feedback:
if self.ac.smooth_mask:
raise ValueError
mask = self.smooth_mask_over_superpixels(mask, superpixel)
if self.ac.connectedness_postprocessed_feedback:
if self.ac.smooth_mask:
raise ValueError
mask = self.connected_component(mask, poke_point)
return mask
def make_mask(self, diff):
diff = diff[0]
scores = (diff**2).reshape(self.gc.grid_size, self.gc.stride,
self.gc.grid_size, self.gc.stride,
3).mean(axis=(1, 3, 4))
mask = scores > self.ac.pixel_change_threshold
return mask
def smooth_mask_over_superpixels(self, mask, superpixels):
smoothed_mask = np.zeros_like(mask)
superpixels = [superpixels == i for i in np.unique(superpixels)]
for superpixel in superpixels:
if mask[superpixel].sum() / superpixel.sum(
) > self.ac.superpixel_postprocessing_threshold:
smoothed_mask[superpixel] = True
return smoothed_mask
def connected_component(self, mask, poke_point):
x, y = poke_point
b = mask[x, y]
mask[x, y] = True
fat_mask = self.fatten(mask) if self.ac.fatten else mask
labels = label(fat_mask) * mask
i = labels[x, y]
labels[x, y] *= b
mask[x, y] = b
return labels == i
# Below: Functionality for computing videoPCA soft masks
def make_smooth_mask(self, diff, video):
with torch.no_grad():
diff = torch.from_numpy(diff.transpose(0, 3, 1, 2)).float()
diff = torch.sqrt((diff**2).sum(dim=1)).unsqueeze(1)
diff = torch.nn.functional.conv2d(
diff,
self.ac.kernel,
padding=(self.ac.kernel.shape[-1] - 1) // 2)
if (diff[0] > self.ac.soft_mask_threshold).sum() == 0:
return np.zeros(video[0].shape[:-1], dtype=np.float32)
diff = self.bn_torch(diff)
mask = diff.squeeze(1) > .5
mask_np = mask.numpy()
return self.color_histogram_soft_mask(mask_np, mask, video)
def color_histogram_soft_mask(self, mask, mask_torch, video):
if mask[0].sum() == 0 or (~mask[0]).sum() == 0:
return mask[0].astype(np.float32)
if not self.ac.hsv:
video = [rgb2hsv(im) for im in video]
video = self.combine_colors(video)
fg_histogram = np.histogram(video[mask],
bins=self.ac.num_bins,
density=True)[0]
bg_histogram = np.histogram(
video[~mask], bins=self.ac.num_bins, density=True)[0] + 1e-5
image = video[0].reshape(-1)
soft_mask = fg_histogram[image] / (fg_histogram[image] +
bg_histogram[image])
soft_mask = soft_mask.astype(np.float32).reshape(video[0].shape)
soft_mask = self.center_soft_masks(soft_mask, mask_torch[0])
return self.bn_np(self.hysteresis_threshold(self.bn_np(soft_mask)))
def combine_colors(self, video):
video = np.stack(video)
ret = np.zeros(video.shape[:-1], dtype=np.int)
ret += ((video[..., 1] * np.cos(2 * np.pi * video[..., 0]) + 1) / 2 *
(self.ac.colres1 - 1)).astype(np.int)
ret += ((video[..., 1] * np.sin(2 * np.pi * video[..., 0]) + 1) / 2 *
(self.ac.colres2 - 1)).astype(np.int) * self.ac.colres1
ret += (video[..., 2] * (self.ac.colres3 - 1)).astype(
np.int) * self.ac.colres1 * self.ac.colres2
return ret
def hysteresis_threshold(self, mask):
thresholding = self.fatten(
self.fatten(mask > self.ac.hyst_thresholds[0]))
thresholding *= mask > self.ac.hyst_thresholds[1]
return mask * thresholding
@staticmethod
def compute_mean_of(images):
mean = sum(images) / len(images)
return np.stack([im - mean for im in images])
def pca_diff(self, video):
video = np.stack(video)
bs = video.shape[0]
video = video.reshape(
(bs, self.gc.grid_size, self.gc.stride, self.gc.grid_size,
self.gc.stride, 3)).mean(axis=(2, 4))
video_shape = video[0].shape
video = video.reshape((bs, -1))
pca = PCA(n_components=self.ac.num_pca_components)
pca.fit(video)
reconstruction = pca.inverse_transform(pca.transform(video))
diff = video - reconstruction
return diff.reshape(*((bs, ) + video_shape))
def center_soft_masks(self, soft_mask, mask):
with torch.no_grad():
mask = mask.float().unsqueeze(0).unsqueeze(1)
mask = torch.nn.functional.conv2d(
mask,
self.ac.centering_kernel,
padding=(self.ac.centering_kernel.shape[-1] - 1) //
2).squeeze() > 1
return soft_mask * mask.numpy()
@staticmethod
def fatten(mask):
fat_mask = mask.copy()
fat_mask[:-1] = fat_mask[:-1] | mask[1:]
fat_mask[1:] = fat_mask[1:] | mask[:-1]
fat_mask[:, :-1] = fat_mask[:, :-1] | mask[:, 1:]
fat_mask[:, 1:] = fat_mask[:, 1:] | mask[:, :-1]
return fat_mask
@staticmethod
def bn_torch(array):
bs = array.shape[0]
unsqueeze = (bs, ) + (1, ) * (len(array.shape) - 1)
diffmin = array.view(bs, -1).min(1)[0].view(*unsqueeze)
diffmax = array.view(bs, -1).max(1)[0].view(*unsqueeze)
return (array - diffmin) / (diffmax - diffmin + 1e-6)
@staticmethod
def bn_np(array):
bs = array.shape[0]
unsqueeze = (bs, ) + (1, ) * (len(array.shape) - 1)
diffmin = array.reshape((bs, -1)).min(1).reshape(unsqueeze)
diffmax = array.reshape((bs, -1)).max(1).reshape(unsqueeze)
return (array - diffmin) / (diffmax - diffmin + 1e-6)
def compute_ground_truth(self, event):
depth = (event.depth_frame - 0.1) * self.depth_correction
reachable_pixels = depth < self.ac.handDistance
keys = [
o['objectId'] for o in event.metadata['objects']
if o['objectId'] in event.instance_masks.keys() and o['visible']
and (o['moveable'] or o['pickupable'])
and o['mass'] < self.ac.mass_threshold
]
masses_unf = [
self.round_mass(o['mass']) for o in event.metadata['objects']
if o['objectId'] in event.instance_masks.keys() and o['visible']
and (o['moveable'] or o['pickupable'])
and o['mass'] < self.ac.mass_threshold
]
masks_unf = [event.instance_masks[key] for key in keys]
masks, masses = [], []
for mask, mass in zip(masks_unf, masses_unf):
if self.ac.max_pixel_threshold > (
mask *
reachable_pixels).sum() > self.ac.min_pixel_threshold:
masks.append(mask)
masses.append(mass)
poking_points = [
np.stack(np.where(mask * reachable_pixels)) for mask in masks
]
poking_points = [[
self.round(*tuple(points[:, i]))
for i in sample(range(points.shape[1]),
k=min(points.shape[1], self.ac.max_poke_attempts))
] for points in poking_points]
return masks, poking_points, masses
def round(self, x, y):
return x // self.gc.stride, y // self.gc.stride
def round_mass(self, mass):
if mass < self.ac.mass_buckets[0]:
return 0
if mass < self.ac.mass_buckets[1]:
return 1
return 2
def get_score(self, mask, action):
x, y = action
dx1 = min(x, self.kernel_size)
dx2 = min(self.gc.grid_size - 1 - x, self.kernel_size) + 1
dy1 = min(y, self.kernel_size)
dy2 = min(self.gc.grid_size - 1 - y, self.kernel_size) + 1
x1, x2, y1, y2 = x - dx1, x + dx2, y - dy1, y + dy2
return (mask[x1:x2, y1:y2] * self.ac.check_change_kernel[
self.kernel_size - dx1:self.kernel_size + dx2,
self.kernel_size - dy1:self.kernel_size + dy2]).sum()
class Ai2Thor():
def __init__(self):
self.visualize = False
self.verbose = False
self.save_imgs = True
self.plot_loss = True
# st()
# these are all map names
a = np.arange(1, 30)
b = np.arange(201, 231)
c = np.arange(301, 331)
d = np.arange(401, 431)
abcd = np.hstack((a,b,c,d))
mapnames = []
for i in list(abcd):
mapname = 'FloorPlan' + str(i)
mapnames.append(mapname)
train_len = int(0.9 * len(mapnames))
random.shuffle(mapnames)
self.mapnames_train = mapnames[:train_len]
self.mapnames_val = mapnames[train_len:]
# self.num_episodes = len(self.mapnames)
self.ignore_classes = []
# classes to save
self.include_classes = [
'ShowerDoor', 'Cabinet', 'CounterTop', 'Sink', 'Towel', 'HandTowel', 'TowelHolder', 'SoapBar',
'ToiletPaper', 'ToiletPaperHanger', 'HandTowelHolder', 'SoapBottle', 'GarbageCan', 'Candle', 'ScrubBrush',
'Plunger', 'SinkBasin', 'Cloth', 'SprayBottle', 'Toilet', 'Faucet', 'ShowerHead', 'Box', 'Bed', 'Book',
'DeskLamp', 'BasketBall', 'Pen', 'Pillow', 'Pencil', 'CellPhone', 'KeyChain', 'Painting', 'CreditCard',
'AlarmClock', 'CD', 'Laptop', 'Drawer', 'SideTable', 'Chair', 'Blinds', 'Desk', 'Curtains', 'Dresser',
'Watch', 'Television', 'WateringCan', 'Newspaper', 'FloorLamp', 'RemoteControl', 'HousePlant', 'Statue',
'Ottoman', 'ArmChair', 'Sofa', 'DogBed', 'BaseballBat', 'TennisRacket', 'VacuumCleaner', 'Mug', 'ShelvingUnit',
'Shelf', 'StoveBurner', 'Apple', 'Lettuce', 'Bottle', 'Egg', 'Microwave', 'CoffeeMachine', 'Fork', 'Fridge',
'WineBottle', 'Spatula', 'Bread', 'Tomato', 'Pan', 'Cup', 'Pot', 'SaltShaker', 'Potato', 'PepperShaker',
'ButterKnife', 'StoveKnob', 'Toaster', 'DishSponge', 'Spoon', 'Plate', 'Knife', 'DiningTable', 'Bowl',
'LaundryHamper', 'Vase', 'Stool', 'CoffeeTable', 'Poster', 'Bathtub', 'TissueBox', 'Footstool', 'BathtubBasin',
'ShowerCurtain', 'TVStand', 'Boots', 'RoomDecor', 'PaperTowelRoll', 'Ladle', 'Kettle', 'Safe', 'GarbageBag', 'TeddyBear',
'TableTopDecor', 'Dumbbell', 'Desktop', 'AluminumFoil', 'Window']
self.include_classes_final = [
'Sink',
'Toilet', 'Bed', 'Book',
'CellPhone',
'AlarmClock', 'Laptop', 'Chair',
'Television', 'RemoteControl', 'HousePlant',
'Ottoman', 'ArmChair', 'Sofa', 'BaseballBat', 'TennisRacket', 'Mug',
'Apple', 'Bottle', 'Microwave', 'Fork', 'Fridge',
'WineBottle', 'Cup',
'ButterKnife', 'Toaster', 'Spoon', 'Knife', 'DiningTable', 'Bowl',
'Vase',
'TeddyBear', 'StoveKnob', 'StoveBurner',
]
# self.include_classes = [
# 'Sink',
# 'Toilet', 'Bed', 'Book',
# 'CellPhone',
# 'AlarmClock', 'Laptop', 'Chair',
# 'Television', 'RemoteControl', 'HousePlant',
# 'Ottoman', 'ArmChair', 'Sofa', 'BaseballBat', 'TennisRacket', 'Mug',
# 'Apple', 'Bottle', 'Microwave', 'Fork', 'Fridge',
# 'WineBottle', 'Cup',
# 'ButterKnife', 'Toaster', 'Spoon', 'Knife', 'DiningTable', 'Bowl',
# 'Vase',
# 'TeddyBear',
# ]
self.action_space = {0: "MoveLeft", 1: "MoveRight", 2: "MoveAhead", 3: "MoveBack", 4: "DoNothing"}
self.num_actions = len(self.action_space)
cfg_det = get_cfg()
cfg_det.merge_from_file(model_zoo.get_config_file("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg_det.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.1 # set threshold for this model
cfg_det.MODEL.WEIGHTS = model_zoo.get_checkpoint_url("COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg_det.MODEL.DEVICE='cuda'
self.cfg_det = cfg_det
self.maskrcnn = DefaultPredictor(cfg_det)
self.conf_thresh_detect = 0.7 # for initially detecting a low confident object
self.conf_thresh_init = 0.8 # for after turning head toward object threshold
self.conf_thresh_end = 0.9 # if reach this then stop getting obs
self.BATCH_SIZE = 50 #50 # frames (not episodes) - this is approximate - it could be higher
# self.percentile = 70
self.max_iters = 100000
self.max_frames = 10
self.val_interval = 10 #10 #10
self.save_interval = 50
# self.BATCH_SIZE = 2
# self.percentile = 70
# self.max_iters = 100000
# self.max_frames = 2
# self.val_interval = 1
# self.save_interval = 1
self.small_classes = []
self.rot_interval = 5.0
self.radius_max = 3.5 #3 #1.75
self.radius_min = 1.0 #1.25
self.num_flat_views = 3
self.num_any_views = 7
self.num_views = 25
self.center_from_mask = False # get object centroid from maskrcnn (True) or gt (False)
self.obj_per_scene = 5
mod = 'conf05'
# self.homepath = f'/home/nel/gsarch/aithor/data/test2'
self.homepath = '/home/sirdome/katefgroup/gsarch/ithor/data/' + mod
print(self.homepath)
if not os.path.exists(self.homepath):
os.mkdir(self.homepath)
else:
val = input("Delete homepath? [y/n]: ")
if val == 'y':
import shutil
shutil.rmtree(self.homepath)
os.mkdir(self.homepath)
else:
print("ENDING")
assert(False)
self.log_freq = 1
self.log_dir = self.homepath +'/..' + '/log_cem/' + mod
if not os.path.exists(self.log_dir):
os.mkdir(self.log_dir)
MAX_QUEUE = 10 # flushes when this amount waiting
self.writer = SummaryWriter(self.log_dir, max_queue=MAX_QUEUE, flush_secs=60)
self.W = 256
self.H = 256
self.fov = 90
self.utils = Utils(self.fov, self.W, self.H)
self.K = self.utils.get_habitat_pix_T_camX(self.fov)
self.camera_matrix = self.utils.get_camera_matrix(self.W, self.H, self.fov)
self.controller = Controller(
scene='FloorPlan30', # will change
gridSize=0.25,
width=self.W,
height=self.H,
fieldOfView= self.fov,
renderObjectImage=True,
renderDepthImage=True,
)
self.init_network()
self.run_episodes()
def init_network(self):
input_shape = np.array([3, self.W, self.H])
self.localpnet = LocalPNET(input_shape=input_shape, num_actions=self.num_actions).cuda()
self.loss = nn.CrossEntropyLoss()
self.optimizer = torch.optim.Adam(params=self.localpnet.parameters(),lr=0.00001)
def batch_iteration(self,mapnames,BATCH_SIZE):
batch = {"actions": [], "obs_all": [], "seg_ims": [], "conf_end_change": [], "conf_avg_change": [], "conf_median_change": []}
iter_idx = 0
total_loss = torch.tensor(0.0).cuda()
num_obs = 0
while True:
mapname = np.random.choice(mapnames)
# self.basepath = self.homepath + f"/{mapname}_{episode}"
# print("BASEPATH: ", self.basepath)
# # self.basepath = f"/hdd/ayushj/habitat_data/{mapname}_{episode}"
# if not os.path.exists(self.basepath):
# os.mkdir(self.basepath)
self.controller.reset(scene=mapname)
total_loss, obs, actions, seg_ims, confs = self.run("train", total_loss)
if obs is None:
print("NO EPISODE LOSS.. SKIPPING BATCH INSTANCE")
continue
num_obs += len(actions)
print("Total loss for train batch # ",iter_idx," :",total_loss)
confs = np.array(confs)
conf_end_change = confs[-1] - confs[0]
conf_avg_change = np.mean(np.diff(confs))
conf_median_change = np.median(np.diff(confs))
batch["actions"].append(actions)
# These are only used for plotting
batch["obs_all"].append(obs)
batch["seg_ims"].append(seg_ims)
batch["conf_end_change"].append(conf_end_change)
batch["conf_avg_change"].append(conf_avg_change)
batch["conf_median_change"].append(conf_median_change)
iter_idx += 1
# if len(batch["obs_all"]) == BATCH_SIZE:
if num_obs >= BATCH_SIZE:
print("NUM OBS IN BATCH=", num_obs)
# batch["total_loss"] = total_loss
print("Total loss for iter: ", total_loss)
return total_loss, batch, num_obs
# iter_idx = 0
# total_loss = torch.tensor(0.0).cuda()
# batch = {"actions": [], "obs_all": [], "seg_ims": [], "conf_end_change": [], "conf_avg_change": []}
# def elite_batch(self,batch,percentile):
# rewards = np.array(batch["rewards"])
# obs = batch["obs"]
# actions = batch["actions"]
# rewards_mean = float(np.mean(rewards))
# rewards_boundary = np.percentile(rewards,percentile)
# print("Reward boundary: ", rewards_boundary)
# rewards_mean = float(np.mean(rewards))
# training_obs = []
# training_actions = []
# for idx in range(rewards.shape[0]):
# reward_idx = rewards[idx]
# if reward_idx < rewards_boundary:
# continue
# training_obs.extend(obs[idx])
# training_actions.extend(actions[idx])
# obs_tensor = torch.FloatTensor(training_obs).permute(0, 3, 1, 2).cuda()
# act_tensor = torch.LongTensor(training_actions).cuda()
# return obs_tensor, act_tensor, rewards_mean, rewards_boundary
def run_val(self, mapnames, BATCH_SIZE, summ_writer=None):
# run validation every x steps
# batch = {"rewards": [], "obs": [], "actions": []}
episode_rewards = 0.0
seg_ims_batch = []
obs_ims_batch = []
iter_idx = 0
while True:
mapname = np.random.choice(mapnames)
self.controller.reset(scene=mapname)
_, obs, actions, seg_ims, confs = self.run("val", None)
# self.controller.stop()
# time.sleep(1)
if obs is None:
print("NO EPISODE REWARDS.. SKIPPING BATCH INSTANCE")
continue
seg_ims_batch.append(seg_ims)
obs_ims_batch.append(obs)
confs = np.array(confs)
conf_end_change = confs[-1] - confs[0]
conf_avg_change = np.mean(np.diff(confs))
conf_median_change = np.median(np.diff(confs))
print("val confidence change (end-start):", conf_end_change)
print("val average confidence difference between frames", conf_avg_change)
# batch["obs"].append(obs[1:]) # first obs is initial pos (for plotting)
# batch["actions"].append(actions)
iter_idx += 1
if len(obs_ims_batch) == 1: # only one for val
try:
if summ_writer is not None:
name = 'inputs_val/rgbs_original'
self.summ_writer.summ_imgs_aithor(name,obs_ims_batch, self.W, self.H, self.max_frames)
name = 'inputs_val/rgbs_maskrcnn'
self.summ_writer.summ_imgs_aithor(name,seg_ims_batch, self.W, self.H, self.max_frames)
except:
print("PLOTTING DIDNT WORK")
pass
break
return conf_end_change, conf_avg_change, conf_median_change
def run_episodes(self):
iteration = 0
while True:
iteration += 1
print("ITERATION #", iteration)
self.summ_writer = utils.improc.Summ_writer(
writer=self.writer,
global_step=iteration,
log_freq=self.log_freq,
fps=8,
just_gif=True)
total_loss, batch, num_obs = self.batch_iteration(self.mapnames_train,self.BATCH_SIZE)
self.optimizer.zero_grad()
total_loss.backward()
self.optimizer.step()
if iteration >= self.max_iters:
print("MAX ITERS REACHED")
self.writer.close()
break
if iteration % self.val_interval == 0:
conf_end_change, conf_avg_change, conf_median_change = self.run_val(self.mapnames_val, self.BATCH_SIZE, self.summ_writer)
if self.plot_loss:
self.summ_writer.summ_scalar('val_conf_end_change', conf_end_change)
self.summ_writer.summ_scalar('val_conf_avg_change', conf_avg_change)
self.summ_writer.summ_scalar('val_conf_median_change', conf_median_change)
if iteration % self.save_interval == 0:
PATH = self.homepath + f'/checkpoint{iteration}.tar'
torch.save(self.localpnet.state_dict(), PATH)
if self.plot_loss:
conf_end_change_t = np.mean(np.array(batch["conf_end_change"]))
conf_avg_change_t = np.mean(np.array(batch["conf_avg_change"]))
conf_median_change_t = np.mean(batch["conf_median_change"])
self.summ_writer.summ_scalar('train_conf_end_change_batchavg', conf_end_change_t)
self.summ_writer.summ_scalar('train_conf_avg_change_batchavg', conf_avg_change_t)
self.summ_writer.summ_scalar('train_conf_median_change_batchavg', conf_median_change_t)
self.summ_writer.summ_scalar('total_loss', total_loss)
## PLOTTING #############
try:
summ_writer = self.summ_writer
if summ_writer is not None and (iteration % self.val_interval == 0):
obs_ims_batch = batch["obs_all"]
seg_ims_batch = batch["seg_ims"]
name = 'inputs_train/rgbs_original'
self.summ_writer.summ_imgs_aithor(name,obs_ims_batch, self.W, self.H, self.max_frames)
name = 'inputs_train/rgbs_maskrcnn'
self.summ_writer.summ_imgs_aithor(name,seg_ims_batch, self.W, self.H, self.max_frames)
except:
print("PLOTTING DIDNT WORK")
pass
self.writer.close() # close tensorboard to flush
self.controller.stop()
time.sleep(10)
def run2(self):
event = self.controller.step('GetReachablePositions')
for obj in event.metadata['objects']:
if obj['objectType'] not in self.objects:
self.objects.append(obj['objectType'])
def get_detectron_conf_center_obj(self,im, obj_mask, frame=None):
im = Image.fromarray(im, mode="RGB")
im = cv2.cvtColor(np.asarray(im), cv2.COLOR_RGB2BGR)
outputs = self.maskrcnn(im)
pred_masks = outputs['instances'].pred_masks
pred_scores = outputs['instances'].scores
pred_classes = outputs['instances'].pred_classes
len_pad = 5
W2_low = self.W//2 - len_pad
W2_high = self.W//2 + len_pad
H2_low = self.H//2 - len_pad
H2_high = self.H//2 + len_pad
if False:
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(self.cfg_det.DATASETS.TRAIN[0]), scale=1.0)
out = v.draw_instance_predictions(outputs['instances'].to("cpu"))
seg_im = out.get_image()
plt.figure(1)
plt.clf()
plt.imshow(seg_im)
plt_name = self.homepath + f'/seg_all{frame}.png'
plt.savefig(plt_name)
seg_im[W2_low:W2_high, H2_low:H2_high,:] = 0.0
plt.figure(1)
plt.clf()
plt.imshow(seg_im)
plt_name = self.homepath + f'/seg_all_mask{frame}.png'
plt.savefig(plt_name)
ind_obj = None
# max_overlap = 0
sum_obj_mask = np.sum(obj_mask)
mask_sum_thresh = 7000
for idx in range(pred_masks.shape[0]):
pred_mask_cur = pred_masks[idx].detach().cpu().numpy()
pred_masks_center = pred_mask_cur[W2_low:W2_high, H2_low:H2_high]
sum_pred_mask_cur = np.sum(pred_mask_cur)
# print(torch.sum(pred_masks_center))
if np.sum(pred_masks_center) > 0:
if np.abs(sum_pred_mask_cur - sum_obj_mask) < mask_sum_thresh:
ind_obj = idx
mask_sum_thresh = np.abs(sum_pred_mask_cur - sum_obj_mask)
# max_overlap = torch.sum(pred_masks_center)
if ind_obj is None:
print("RETURNING NONE")
return None, None, None, None
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(self.cfg_det.DATASETS.TRAIN[0]), scale=1.0)
out = v.draw_instance_predictions(outputs['instances'][ind_obj].to("cpu"))
seg_im = out.get_image()
if False:
plt.figure(1)
plt.clf()
plt.imshow(seg_im)
plt_name = self.homepath + f'/seg{frame}.png'
plt.savefig(plt_name)
# print("OBJ CLASS ID=", int(pred_classes[ind_obj].detach().cpu().numpy()))
# pred_boxes = outputs['instances'].pred_boxes.tensor
# pred_classes = outputs['instances'].pred_classes
# pred_scores = outputs['instances'].scores
obj_score = float(pred_scores[ind_obj].detach().cpu().numpy())
obj_pred_classes = int(pred_classes[ind_obj].detach().cpu().numpy())
obj_pred_mask = pred_masks[ind_obj].detach().cpu().numpy()
return obj_score, obj_pred_classes, obj_pred_mask, seg_im
def detect_object_centroid(self, im, event):
im = Image.fromarray(im, mode="RGB")
im = cv2.cvtColor(np.asarray(im), cv2.COLOR_RGB2BGR)
outputs = self.maskrcnn(im)
v = Visualizer(im[:, :, ::-1], MetadataCatalog.get(self.cfg_det.DATASETS.TRAIN[0]), scale=1.2)
out = v.draw_instance_predictions(outputs['instances'].to("cpu"))
seg_im = out.get_image()
if False:
plt.figure(1)
plt.clf()
plt.imshow(seg_im)
plt_name = self.homepath + '/seg_init.png'
plt.savefig(plt_name)
pred_masks = outputs['instances'].pred_masks
pred_boxes = outputs['instances'].pred_boxes.tensor
pred_classes = outputs['instances'].pred_classes
pred_scores = outputs['instances'].scores
obj_catids = []
obj_scores = []
obj_masks = []
for segs in range(len(pred_masks)):
if pred_scores[segs] <= self.conf_thresh_detect:
obj_catids.append(pred_classes[segs].item())
obj_scores.append(pred_scores[segs].item())
obj_masks.append(pred_masks[segs])
eulers_xyz_rad = np.radians(np.array([event.metadata['agent']['cameraHorizon'], event.metadata['agent']['rotation']['y'], 0.0]))
rx = eulers_xyz_rad[0]
ry = eulers_xyz_rad[1]
rz = eulers_xyz_rad[2]
rotation_ = self.utils.eul2rotm(-rx, -ry, rz)
translation_ = np.array(list(event.metadata['agent']['position'].values())) + np.array([0.0, 0.675, 0.0])
# need to invert since z is positive here by convention
translation_[2] = -translation_[2]
T_world_cam = np.eye(4)
T_world_cam[0:3,0:3] = rotation_
T_world_cam[0:3,3] = translation_
if not obj_masks:
return None, None
elif self.center_from_mask:
# want an object not on the edges of the image
sum_interior = 0
while sum_interior==0:
if len(obj_masks)==0:
return None, None
random_int = np.random.randint(low=0, high=len(obj_masks))
obj_mask_focus = obj_masks.pop(random_int)
print("OBJECT ID INIT=", obj_catids[random_int])
sum_interior = torch.sum(obj_mask_focus[50:self.W-50, 50:self.H-50])
depth = event.depth_frame
xs, ys = np.meshgrid(np.linspace(-1*256/2.,1*256/2.,256), np.linspace(1*256/2.,-1*256/2., 256))
depth = depth.reshape(1,256,256)
xs = xs.reshape(1,256,256)
ys = ys.reshape(1,256,256)
xys = np.vstack((xs * depth , ys * depth, -depth, np.ones(depth.shape)))
xys = xys.reshape(4, -1)
xy_c0 = np.matmul(np.linalg.inv(self.K), xys)
xyz = xy_c0.T[:,:3].reshape(256,256,3)
xyz_obj_masked = xyz[obj_mask_focus]
xyz_obj_masked = np.matmul(rotation_, xyz_obj_masked.T) + translation_.reshape(3,1)
xyz_obj_mid = np.mean(xyz_obj_masked, axis=1)
xyz_obj_mid[2] = -xyz_obj_mid[2]
else:
# want an object not on the edges of the image
sum_interior = 0
while True:
if len(obj_masks)==0:
return None, None
random_int = np.random.randint(low=0, high=len(obj_masks))
obj_mask_focus = obj_masks.pop(random_int)
# print("OBJECT ID INIT=", obj_catids[random_int])
sum_interior = torch.sum(obj_mask_focus[50:self.W-50, 50:self.H-50])
if sum_interior < 500:
continue # exclude too small objects
pixel_locs_obj = np.where(obj_mask_focus.cpu().numpy())
x_mid = np.round(np.median(pixel_locs_obj[1])/self.W, 4)
y_mid = np.round(np.median(pixel_locs_obj[0])/self.H, 4)
if False:
plt.figure(1)
plt.clf()
plt.imshow(obj_mask_focus)
plt.plot(np.median(pixel_locs_obj[1]), np.median(pixel_locs_obj[0]), 'x')
plt_name = self.homepath + '/seg_mask.png'
plt.savefig(plt_name)
event = self.controller.step('TouchThenApplyForce', x=x_mid, y=y_mid, handDistance = 1000000.0, direction=dict(x=0.0, y=0.0, z=0.0), moveMagnitude = 0.0)
obj_focus_id = event.metadata['actionReturn']['objectId']
xyz_obj_mid = None
for o in event.metadata['objects']:
if o['objectId'] == obj_focus_id:
if o['objectType'] not in self.include_classes_final:
continue
xyz_obj_mid = np.array(list(o['axisAlignedBoundingBox']['center'].values()))
if xyz_obj_mid is not None:
break
print("MIDPOINT=", xyz_obj_mid)
return xyz_obj_mid, obj_mask_focus
def run(self, mode, total_loss, summ_writer=None):
event = self.controller.step('GetReachablePositions')
if not event.metadata['reachablePositions']:
# Different versions this is empty/full
event = self.controller.step(action='MoveAhead')
self.nav_pts = event.metadata['reachablePositions']
self.nav_pts = np.array([list(d.values()) for d in self.nav_pts])
# objects = np.random.choice(event.metadata['objects'], self.obj_per_scene, replace=False)
objects = event.metadata['objects']
objects_inds = np.arange(len(event.metadata['objects']))
np.random.shuffle(objects_inds)
# objects = np.random.shuffle(event.metadata['objects'])
# for obj in event.metadata['objects']: #objects:
# print(obj['name'])
# objects = objects[0]
successes = 0
meta_obj_idx = 0
num_obs = 0
while True: #successes < self.obj_per_scene and meta_obj_idx <= len(event.metadata['objects']) - 1:
if meta_obj_idx > len(event.metadata['objects']) - 1:
print("OUT OF OBJECT... RETURNING")
return total_loss, None, None, None, None
obj = objects[objects_inds[meta_obj_idx]]
meta_obj_idx += 1
print("Center object is ", obj['objectType'])
# if obj['name'] in ['Microwave_b200e0bc']:
# print(obj['name'])
# else:
# continue
# print(obj['name'])
if obj['objectType'] not in self.include_classes:
print("Continuing... Invalid Object")
continue
# Calculate distance to object center
obj_center = np.array(list(obj['axisAlignedBoundingBox']['center'].values()))
obj_center = np.expand_dims(obj_center, axis=0)
distances = np.sqrt(np.sum((self.nav_pts - obj_center)**2, axis=1))
# Get points with r_min < dist < r_max
valid_pts = self.nav_pts[np.where((distances > self.radius_min)*(distances<self.radius_max))]
# add height from center of agent to camera
rand_pos_int = np.random.randint(low=0, high=valid_pts.shape[0])
pos_s = valid_pts[rand_pos_int]
pos_s[1] = pos_s[1] + 0.675
turn_yaw, turn_pitch = self.utils.get_rotation_to_obj(obj_center, pos_s)
event = self.controller.step('TeleportFull', x=pos_s[0], y=pos_s[1], z=pos_s[2], rotation=dict(x=0.0, y=int(turn_yaw), z=0.0), horizon=int(turn_pitch))
rgb = event.frame
# get object center of a low confidence object
obj_center, obj_mask = self.detect_object_centroid(rgb, event)
if obj_center is None:
print("NO LOW CONFIDENCE OBJECTS... SKIPPING...")
continue
# initialize object in center of FOV
turn_yaw, turn_pitch = self.utils.get_rotation_to_obj(obj_center, pos_s)
if mode=="train":
pos_s_prev = pos_s
turn_yaw_prev = turn_yaw
turn_pitch_prev = turn_pitch
event = self.controller.step('TeleportFull', x=pos_s[0], y=pos_s[1], z=pos_s[2], rotation=dict(x=0.0, y=int(turn_yaw), z=0.0), horizon=int(turn_pitch))
rgb = event.frame
seg_ims = []
obs = []
init_conf, obj_pred_classes, obj_mask, seg_im = self.get_detectron_conf_center_obj(rgb, obj_mask.detach().cpu().numpy())
if init_conf is None:
print("Nothing detected in the center... SKIPPING")
continue
conf_cur = init_conf
conf_prev = init_conf
if init_conf > self.conf_thresh_init:
print("HIGH INITIAL CONFIDENCE... SKIPPING...")
continue
seg_ims.append(seg_im)
obs.append(rgb)
actions = []
confs = []
confs.append(conf_cur)
episode_rewards = 0.0
frame = 0
while True:
rgb_tensor = torch.FloatTensor([rgb]).permute(0, 3, 1, 2).cuda()
if mode=="train":
torch.set_grad_enabled(True)
action_ind, act_proba = self.localpnet(rgb_tensor)
elif mode=="val":
with torch.no_grad():
action_ind, act_proba = self.localpnet(rgb_tensor)
# action_ind, act_proba = actions_probability.data.cpu().numpy()[0]
# action_ind = np.random.choice(len(act_proba),p=act_proba)
action_ind = int(action_ind.detach().cpu().numpy())
action = self.action_space[action_ind]
print("ACTION=", action)
obs.append(rgb)
actions.append(action_ind)
# get best action in a confidence sense
if mode=="train":
best_action = 4 # "DoNothing"
best_conf = conf_prev
for action_idx in [0,1,2,3]:
action_t = self.action_space[action_idx]
event_t = self.controller.step(action_t)
agent_position_t = np.array(list(event_t.metadata['agent']['position'].values())) + np.array([0.0, 0.675, 0.0])
turn_yaw_t, turn_pitch_t = self.utils.get_rotation_to_obj(obj_center, agent_position_t)
event_t = self.controller.step('TeleportFull', x=agent_position_t[0], y=agent_position_t[1], z=agent_position_t[2], rotation=dict(x=0.0, y=int(turn_yaw_t), z=0.0), horizon=int(turn_pitch_t))
rgb_t = event_t.frame
conf_t, _, _, _ = self.get_detectron_conf_center_obj(rgb_t, obj_mask, frame)
if conf_t is None:
conf_t = best_conf - 1 # dont want no detection
if conf_t > best_conf:
best_action = action_idx
best_conf = conf_t
_ = self.controller.step('TeleportFull', x=pos_s_prev[0], y=pos_s_prev[1], z=pos_s_prev[2], rotation=dict(x=0.0, y=int(turn_yaw_prev), z=0.0), horizon=int(turn_pitch_prev))
best_action = torch.LongTensor([best_action]).cuda()
total_loss += self.loss(act_proba, best_action)
num_obs += 1
print("BEST ACTION=", self.action_space[int(best_action.detach().cpu().numpy())])
if not action=="DoNothing":
event = self.controller.step(action)
agent_position = np.array(list(event.metadata['agent']['position'].values())) + np.array([0.0, 0.675, 0.0])
turn_yaw, turn_pitch = self.utils.get_rotation_to_obj(obj_center, agent_position)
event = self.controller.step('TeleportFull', x=agent_position[0], y=agent_position[1], z=agent_position[2], rotation=dict(x=0.0, y=int(turn_yaw), z=0.0), horizon=int(turn_pitch))
else:
print("Do nothing reached")
print("End confidence: ", conf_prev)
break
if mode=="train":
pos_s_prev = agent_position
turn_yaw_prev = turn_yaw
turn_pitch_prev = turn_pitch
rgb = event.frame
conf_cur, obj_pred_classes, obj_mask_new, seg_im = self.get_detectron_conf_center_obj(rgb, obj_mask, frame)
seg_ims.append(seg_im)
if conf_cur is None:
conf_cur = conf_prev
seg_im = rgb
else:
obj_mask = obj_mask_new
if True:
plt.figure(1)
plt.clf()
plt.imshow(seg_im)
plt_name = self.homepath + f'/seg{frame}.png'
plt.savefig(plt_name)
confs.append(conf_cur)
conf_prev = conf_cur
if conf_cur > self.conf_thresh_end:
print("CONFIDENCE THRESHOLD REACHED!")
print("End confidence: ", conf_cur)
break
if frame >= self.max_frames - 1:
print("MAX FRAMES REACHED")
print("End confidence: ", conf_cur)
break
frame += 1
return total_loss, obs, actions, seg_ims, confs
class Ai2Thor():
def __init__(self):
self.visualize = False
self.verbose = False
self.save_imgs = True
self.plot_loss = True
# st()
mapnames = []
for i in [1, 201, 301, 401]:
mapname = 'FloorPlan' + str(i)
mapnames.append(mapname)
# random.shuffle(mapnames)
self.mapnames_train = mapnames
self.num_episodes = len(self.mapnames_train)
# get rest of the house in orders
a = np.arange(2, 30)
b = np.arange(202, 231)
c = np.arange(302, 331)
d = np.arange(402, 431)
abcd = np.hstack((a, b, c, d))
mapnames = []
for i in range(a.shape[0]):
mapname = 'FloorPlan' + str(a[i])
mapnames.append(mapname)
mapname = 'FloorPlan' + str(b[i])
mapnames.append(mapname)
mapname = 'FloorPlan' + str(c[i])
mapnames.append(mapname)
mapname = 'FloorPlan' + str(d[i])
mapnames.append(mapname)
self.mapnames_test = mapnames
self.ignore_classes = []
# classes to save
# self.include_classes = [
# 'ShowerDoor', 'Cabinet', 'CounterTop', 'Sink', 'Towel', 'HandTowel', 'TowelHolder', 'SoapBar',
# 'ToiletPaper', 'ToiletPaperHanger', 'HandTowelHolder', 'SoapBottle', 'GarbageCan', 'Candle', 'ScrubBrush',
# 'Plunger', 'SinkBasin', 'Cloth', 'SprayBottle', 'Toilet', 'Faucet', 'ShowerHead', 'Box', 'Bed', 'Book',
# 'DeskLamp', 'BasketBall', 'Pen', 'Pillow', 'Pencil', 'CellPhone', 'KeyChain', 'Painting', 'CreditCard',
# 'AlarmClock', 'CD', 'Laptop', 'Drawer', 'SideTable', 'Chair', 'Blinds', 'Desk', 'Curtains', 'Dresser',
# 'Watch', 'Television', 'WateringCan', 'Newspaper', 'FloorLamp', 'RemoteControl', 'HousePlant', 'Statue',
# 'Ottoman', 'ArmChair', 'Sofa', 'DogBed', 'BaseballBat', 'TennisRacket', 'VacuumCleaner', 'Mug', 'ShelvingUnit',
# 'Shelf', 'StoveBurner', 'Apple', 'Lettuce', 'Bottle', 'Egg', 'Microwave', 'CoffeeMachine', 'Fork', 'Fridge',
# 'WineBottle', 'Spatula', 'Bread', 'Tomato', 'Pan', 'Cup', 'Pot', 'SaltShaker', 'Potato', 'PepperShaker',
# 'ButterKnife', 'StoveKnob', 'Toaster', 'DishSponge', 'Spoon', 'Plate', 'Knife', 'DiningTable', 'Bowl',
# 'LaundryHamper', 'Vase', 'Stool', 'CoffeeTable', 'Poster', 'Bathtub', 'TissueBox', 'Footstool', 'BathtubBasin',
# 'ShowerCurtain', 'TVStand', 'Boots', 'RoomDecor', 'PaperTowelRoll', 'Ladle', 'Kettle', 'Safe', 'GarbageBag', 'TeddyBear',
# 'TableTopDecor', 'Dumbbell', 'Desktop', 'AluminumFoil', 'Window']
# These are all classes shared between aithor and coco
self.include_classes = [
'Sink',
'Toilet',
'Bed',
'Book',
'CellPhone',
'AlarmClock',
'Laptop',
'Chair',
'Television',
'RemoteControl',
'HousePlant',
'Ottoman',
'ArmChair',
'Sofa',
'BaseballBat',
'TennisRacket',
'Mug',
'Apple',
'Bottle',
'Microwave',
'Fork',
'Fridge',
'WineBottle',
'Cup',
'ButterKnife',
'Toaster',
'Spoon',
'Knife',
'DiningTable',
'Bowl',
'Vase',
'TeddyBear',
]
self.maskrcnn_to_ithor = {
81: 'Sink',
70: 'Toilet',
65: 'Bed',
84: 'Book',
77: 'CellPhone',
85: 'AlarmClock',
73: 'Laptop',
62: 'Chair',
72: 'Television',
75: 'RemoteControl',
64: 'HousePlant',
62: 'Ottoman',
62: 'ArmChair',
63: 'Sofa',
39: 'BaseballBat',
43: 'TennisRacket',
47: 'Mug',
53: 'Apple',
44: 'Bottle',
78: 'Microwave',
48: 'Fork',
82: 'Fridge',
44: 'WineBottle',
47: 'Cup',
49: 'ButterKnife',
80: 'Toaster',
50: 'Spoon',
49: 'Knife',
67: 'DiningTable',
51: 'Bowl',
86: 'Vase',
88: 'TeddyBear',
}
self.ithor_to_maskrcnn = {
'Sink': 81,
'Toilet': 70,
'Bed': 65,
'Book': 84,
'CellPhone': 77,
'AlarmClock': 85,
'Laptop': 73,
'Chair': 62,
'Television': 72,
'RemoteControl': 75,
'HousePlant': 64,
'Ottoman': 62,
'ArmChair': 62,
'Sofa': 63,
'BaseballBat': 39,
'TennisRacket': 43,
'Mug': 47,
'Apple': 53,
'Bottle': 44,
'Microwave': 78,
'Fork': 48,
'Fridge': 82,
'WineBottle': 44,
'Cup': 47,
'ButterKnife': 49,
'Toaster': 80,
'Spoon': 50,
'Knife': 49,
'DiningTable': 67,
'Bowl': 51,
'Vase': 86,
'TeddyBear': 88,
}
self.maskrcnn_to_catname = {
81: 'sink',
67: 'dining table',
65: 'bed',
84: 'book',
77: 'cell phone',
70: 'toilet',
85: 'clock',
73: 'laptop',
62: 'chair',
72: 'tv',
75: 'remote',
64: 'potted plant',
63: 'couch',
39: 'baseball bat',
43: 'tennis racket',
47: 'cup',
53: 'apple',
44: 'bottle',
78: 'microwave',
48: 'fork',
82: 'refrigerator',
46: 'wine glass',
49: 'knife',
79: 'oven',
80: 'toaster',
50: 'spoon',
67: 'dining table',
51: 'bowl',
86: 'vase',
88: 'teddy bear',
}
self.obj_conf_dict = {
'sink': [],
'dining table': [],
'bed': [],
'book': [],
'cell phone': [],
'clock': [],
'laptop': [],
'chair': [],
'tv': [],
'remote': [],
'potted plant': [],
'couch': [],
'baseball bat': [],
'tennis racket': [],
'cup': [],
'apple': [],
'bottle': [],
'microwave': [],
'fork': [],
'refrigerator': [],
'wine glass': [],
'knife': [],
'oven': [],
'toaster': [],
'spoon': [],
'dining table': [],
'bowl': [],
'vase': [],
'teddy bear': [],
}
self.data_store = {
'sink': {},
'dining table': {},
'bed': {},
'book': {},
'cell phone': {},
'clock': {},
'laptop': {},
'chair': {},
'tv': {},
'remote': {},
'potted plant': {},
'couch': {},
'baseball bat': {},
'tennis racket': {},
'cup': {},
'apple': {},
'bottle': {},
'microwave': {},
'fork': {},
'refrigerator': {},
'wine glass': {},
'knife': {},
'oven': {},
'toaster': {},
'spoon': {},
'dining table': {},
'bowl': {},
'vase': {},
'teddy bear': {},
}
self.data_store_features = []
self.feature_obj_ids = []
self.first_time = True
self.Softmax = nn.Softmax(dim=0)
self.action_space = {
0: "MoveLeft",
1: "MoveRight",
2: "MoveAhead",
3: "MoveBack",
4: "DoNothing"
}
self.num_actions = len(self.action_space)
cfg_det = get_cfg()
cfg_det.merge_from_file(
model_zoo.get_config_file(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml"))
cfg_det.MODEL.ROI_HEADS.SCORE_THRESH_TEST = 0.1 # set threshold for this model
cfg_det.MODEL.WEIGHTS = model_zoo.get_checkpoint_url(
"COCO-InstanceSegmentation/mask_rcnn_R_50_FPN_3x.yaml")
cfg_det.MODEL.DEVICE = 'cuda'
self.cfg_det = cfg_det
self.maskrcnn = DefaultPredictor(cfg_det)
self.normalize = transforms.Compose([
transforms.Resize(256, interpolation=PIL.Image.BILINEAR),
transforms.CenterCrop(256),
transforms.ToTensor(),
transforms.Normalize((0.485, 0.456, 0.406), (0.229, 0.224, 0.225)),
])
# Initialize vgg
vgg16 = torchvision.models.vgg16(pretrained=True).double().cuda()
vgg16.eval()
print(torch.nn.Sequential(*list(vgg16.features.children())))
self.vgg_feat_extractor = torch.nn.Sequential(
*list(vgg16.features.children())[:-2])
print(self.vgg_feat_extractor)
self.vgg_mean = torch.from_numpy(
np.array([0.485, 0.456, 0.406]).reshape(1, 3, 1, 1))
self.vgg_std = torch.from_numpy(
np.array([0.229, 0.224, 0.225]).reshape(1, 3, 1, 1))
self.conf_thresh_detect = 0.7 # for initially detecting a low confident object
self.conf_thresh_init = 0.8 # for after turning head toward object threshold
self.conf_thresh_end = 0.9 # if reach this then stop getting obs
self.BATCH_SIZE = 50 # frames (not episodes) - this is approximate - it could be higher
# self.percentile = 70
self.max_iters = 100000
self.max_frames = 10
self.val_interval = 10 #10
self.save_interval = 50
# self.BATCH_SIZE = 2
# self.percentile = 70
# self.max_iters = 100000
# self.max_frames = 2
# self.val_interval = 1
# self.save_interval = 1
self.small_classes = []
self.rot_interval = 5.0
self.radius_max = 3.5 #3 #1.75
self.radius_min = 1.0 #1.25
self.num_flat_views = 3
self.num_any_views = 7
self.num_views = 25
self.center_from_mask = False # get object centroid from maskrcnn (True) or gt (False)
self.obj_per_scene = 5
mod = 'test00'
# self.homepath = f'/home/nel/gsarch/aithor/data/test2'
self.homepath = '/home/sirdome/katefgroup/gsarch/ithor/data/' + mod
print(self.homepath)
if not os.path.exists(self.homepath):
os.mkdir(self.homepath)
else:
val = input("Delete homepath? [y/n]: ")
if val == 'y':
import shutil
shutil.rmtree(self.homepath)
os.mkdir(self.homepath)
else:
print("ENDING")
assert (False)
self.log_freq = 1
self.log_dir = self.homepath + '/..' + '/log_cem/' + mod
if not os.path.exists(self.log_dir):
os.mkdir(self.log_dir)
MAX_QUEUE = 10 # flushes when this amount waiting
self.writer = SummaryWriter(self.log_dir,
max_queue=MAX_QUEUE,
flush_secs=60)
self.W = 256
self.H = 256
self.fov = 90
self.utils = Utils(self.fov, self.W, self.H)
self.K = self.utils.get_habitat_pix_T_camX(self.fov)
self.camera_matrix = self.utils.get_camera_matrix(
self.W, self.H, self.fov)
self.controller = Controller(
scene='FloorPlan30', # will change
gridSize=0.25,
width=self.W,
height=self.H,
fieldOfView=self.fov,
renderObjectImage=True,
renderDepthImage=True,
)
self.init_network()
self.run_episodes()
def init_network(self):
input_shape = np.array([3, self.W, self.H])
self.localpnet = LocalPNET(input_shape=input_shape,
num_actions=self.num_actions).cuda()
self.loss = nn.CrossEntropyLoss()
self.optimizer = torch.optim.Adam(params=self.localpnet.parameters(),
lr=0.00001)
def batch_iteration(self, mapnames, BATCH_SIZE):
batch = {
"actions": [],
"obs_all": [],
"seg_ims": [],
"conf_end_change": [],
"conf_avg_change": [],
"conf_median_change": []
}
iter_idx = 0
total_loss = torch.tensor(0.0).cuda()
num_obs = 0
while True:
mapname = np.random.choice(mapnames)
# self.basepath = self.homepath + f"/{mapname}_{episode}"
# print("BASEPATH: ", self.basepath)
# # self.basepath = f"/hdd/ayushj/habitat_data/{mapname}_{episode}"
# if not os.path.exists(self.basepath):
# os.mkdir(self.basepath)
self.controller.reset(scene=mapname)
total_loss, obs, actions, seg_ims, confs = self.run(
"train", total_loss)
if obs is None:
print("NO EPISODE LOSS.. SKIPPING BATCH INSTANCE")
continue
num_obs += len(actions)
print("Total loss for train batch # ", iter_idx, " :", total_loss)
confs = np.array(confs)
conf_end_change = confs[-1] - confs[0]
conf_avg_change = np.mean(np.diff(confs))
conf_median_change = np.median(np.diff(confs))
batch["actions"].append(actions)
# These are only used for plotting
batch["obs_all"].append(obs)
batch["seg_ims"].append(seg_ims)
batch["conf_end_change"].append(conf_end_change)
batch["conf_avg_change"].append(conf_avg_change)
batch["conf_median_change"].append(conf_median_change)
iter_idx += 1
# if len(batch["obs_all"]) == BATCH_SIZE:
if num_obs >= BATCH_SIZE:
print("NUM OBS IN BATCH=", num_obs)
# batch["total_loss"] = total_loss
print("Total loss for iter: ", total_loss)
return total_loss, batch, num_obs
# iter_idx = 0
# total_loss = torch.tensor(0.0).cuda()
# batch = {"actions": [], "obs_all": [], "seg_ims": [], "conf_end_change": [], "conf_avg_change": []}
def run_episodes(self):
self.ep_idx = 0
# self.objects = []
for episode in range(len(self.mapnames_train)):
print("STARTING EPISODE ", episode)
mapname = self.mapnames_train[episode]
print("MAPNAME=", mapname)
self.controller.reset(scene=mapname)
# self.controller.start()
self.basepath = self.homepath + f"/{mapname}_{episode}"
print("BASEPATH: ", self.basepath)
# self.basepath = f"/hdd/ayushj/habitat_data/{mapname}_{episode}"
if not os.path.exists(self.basepath):
os.mkdir(self.basepath)
self.run(mode="train")
self.ep_idx += 1
self.ep_idx = 1
self.best_inner_prods = []
self.pred_ids = []
self.true_ids = []
self.pred_catnames = []
self.true_catnames = []
self.pred_catnames_all = []
self.true_catnames_all = []
self.conf_mats = []
# self.pred_catnames = []
for episode in range(len(self.mapnames_test)):
print("STARTING EPISODE ", episode)
mapname = self.mapnames_test[episode]
print("MAPNAME=", mapname)
self.controller.reset(scene=mapname)
# self.controller.start()
self.basepath = self.homepath + f"/{mapname}_{episode}"
print("BASEPATH: ", self.basepath)
# self.basepath = f"/hdd/ayushj/habitat_data/{mapname}_{episode}"
if not os.path.exists(self.basepath):
os.mkdir(self.basepath)
self.run(mode="test")
if self.ep_idx % 4 == 0:
self.best_inner_prods = np.array(self.best_inner_prods)
self.pred_ids = np.array(self.pred_ids)
self.true_ids = np.array(self.true_ids)
# for i in range(len(self.best_inner_prods)):s
correct_pred = self.best_inner_prods[self.pred_ids ==
self.true_ids]
incorrect_pred = self.best_inner_prods[
self.pred_ids != self.true_ids]
bins = 50
plt.figure(1)
plt.clf()
plt.hist([correct_pred, incorrect_pred],
alpha=0.5,
histtype='stepfilled',
label=['correct', 'incorrect'],
bins=bins)
plt.title(f'testhouse{self.ep_idx//4}')
plt.xlabel('inner product of nearest neighbor')
plt.ylabel('Counts')
plt.legend()
plt_name = self.homepath + f'/correct_incorrect_testhouse{self.ep_idx//4}.png'
plt.savefig(plt_name)
conf_mat = confusion_matrix(self.pred_catnames,
self.true_catnames,
labels=self.include_classes)
self.conf_mats.append(conf_mat)
plt.figure(1)
plt.clf()
df_cm = pd.DataFrame(conf_mat,
index=[i for i in self.include_classes],
columns=[i for i in self.include_classes])
plt.figure(figsize=(10, 7))
sn.heatmap(df_cm, annot=True)
plt_name = self.homepath + f'/confusion_matrix_testhouse{self.ep_idx//4}.png'
plt.savefig(plt_name)
# plt.show()
self.pred_catnames_all.extend(self.pred_catnames)
self.true_catnames_all.extend(self.true_catnames)
self.best_inner_prods = []
self.pred_ids = []
self.true_ids = []
self.true_catnames = []
self.pred_catnames = []
self.true_catnames = []
conf_mat = confusion_matrix(self.pred_catnames_all,
self.true_catnames_all,
labels=self.include_classes)
plt.figure(1)
plt.clf()
df_cm = pd.DataFrame(conf_mat,
index=[i for i in self.include_classes],
columns=[i for i in self.include_classes])
plt.figure(figsize=(10, 7))
sn.heatmap(df_cm, annot=True)
plt_name = self.homepath + f'/confusion_matrix_testhouses_all.png'
plt.savefig(plt_name)
self.ep_idx += 1
self.controller.stop()
time.sleep(1)
def run2(self):
event = self.controller.step('GetReachablePositions')
for obj in event.metadata['objects']:
if obj['objectType'] not in self.objects:
self.objects.append(obj['objectType'])
def get_detectron_conf_center_obj(self, im, obj_mask, frame=None):
im = Image.fromarray(im, mode="RGB")
im = cv2.cvtColor(np.asarray(im), cv2.COLOR_RGB2BGR)
outputs = self.maskrcnn(im)
pred_masks = outputs['instances'].pred_masks
pred_scores = outputs['instances'].scores
pred_classes = outputs['instances'].pred_classes
len_pad = 5
W2_low = self.W // 2 - len_pad
W2_high = self.W // 2 + len_pad
H2_low = self.H // 2 - len_pad
H2_high = self.H // 2 + len_pad
if False:
v = Visualizer(im[:, :, ::-1],
MetadataCatalog.get(self.cfg_det.DATASETS.TRAIN[0]),
scale=1.0)
out = v.draw_instance_predictions(outputs['instances'].to("cpu"))
seg_im = out.get_image()
plt.figure(1)
plt.clf()
plt.imshow(seg_im)
plt_name = self.homepath + f'/seg_all{frame}.png'
plt.savefig(plt_name)
seg_im[W2_low:W2_high, H2_low:H2_high, :] = 0.0
plt.figure(1)
plt.clf()
plt.imshow(seg_im)
plt_name = self.homepath + f'/seg_all_mask{frame}.png'
plt.savefig(plt_name)
ind_obj = None
# max_overlap = 0
sum_obj_mask = np.sum(obj_mask)
mask_sum_thresh = 7000
for idx in range(pred_masks.shape[0]):
pred_mask_cur = pred_masks[idx].detach().cpu().numpy()
pred_masks_center = pred_mask_cur[W2_low:W2_high, H2_low:H2_high]
sum_pred_mask_cur = np.sum(pred_mask_cur)
# print(torch.sum(pred_masks_center))
if np.sum(pred_masks_center) > 0:
if np.abs(sum_pred_mask_cur - sum_obj_mask) < mask_sum_thresh:
ind_obj = idx
mask_sum_thresh = np.abs(sum_pred_mask_cur - sum_obj_mask)
# max_overlap = torch.sum(pred_masks_center)
if ind_obj is None:
print("RETURNING NONE")
return None, None, None, None
v = Visualizer(im[:, :, ::-1],
MetadataCatalog.get(self.cfg_det.DATASETS.TRAIN[0]),
scale=1.0)
out = v.draw_instance_predictions(
outputs['instances'][ind_obj].to("cpu"))
seg_im = out.get_image()
if False:
plt.figure(1)
plt.clf()
plt.imshow(seg_im)
plt_name = self.homepath + f'/seg{frame}.png'
plt.savefig(plt_name)
# print("OBJ CLASS ID=", int(pred_classes[ind_obj].detach().cpu().numpy()))
# pred_boxes = outputs['instances'].pred_boxes.tensor
# pred_classes = outputs['instances'].pred_classes
# pred_scores = outputs['instances'].scores
obj_score = float(pred_scores[ind_obj].detach().cpu().numpy())
obj_pred_classes = int(pred_classes[ind_obj].detach().cpu().numpy())
obj_pred_mask = pred_masks[ind_obj].detach().cpu().numpy()
return obj_score, obj_pred_classes, obj_pred_mask, seg_im
def detect_object_centroid(self, im, event):
im = Image.fromarray(im, mode="RGB")
im = cv2.cvtColor(np.asarray(im), cv2.COLOR_RGB2BGR)
outputs = self.maskrcnn(im)
v = Visualizer(im[:, :, ::-1],
MetadataCatalog.get(self.cfg_det.DATASETS.TRAIN[0]),
scale=1.2)
out = v.draw_instance_predictions(outputs['instances'].to("cpu"))
seg_im = out.get_image()
if False:
plt.figure(1)
plt.clf()
plt.imshow(seg_im)
plt_name = self.homepath + '/seg_init.png'
plt.savefig(plt_name)
pred_masks = outputs['instances'].pred_masks
pred_boxes = outputs['instances'].pred_boxes.tensor
pred_classes = outputs['instances'].pred_classes
pred_scores = outputs['instances'].scores
obj_catids = []
obj_scores = []
obj_masks = []
for segs in range(len(pred_masks)):
if pred_scores[segs] <= self.conf_thresh_detect:
obj_catids.append(pred_classes[segs].item())
obj_scores.append(pred_scores[segs].item())
obj_masks.append(pred_masks[segs])
eulers_xyz_rad = np.radians(
np.array([
event.metadata['agent']['cameraHorizon'],
event.metadata['agent']['rotation']['y'], 0.0
]))
rx = eulers_xyz_rad[0]
ry = eulers_xyz_rad[1]
rz = eulers_xyz_rad[2]
rotation_ = self.utils.eul2rotm(-rx, -ry, rz)
translation_ = np.array(
list(event.metadata['agent']['position'].values())) + np.array(
[0.0, 0.675, 0.0])
# need to invert since z is positive here by convention
translation_[2] = -translation_[2]
T_world_cam = np.eye(4)
T_world_cam[0:3, 0:3] = rotation_
T_world_cam[0:3, 3] = translation_
if not obj_masks:
return None, None
elif self.center_from_mask:
# want an object not on the edges of the image
sum_interior = 0
while sum_interior == 0:
if len(obj_masks) == 0:
return None, None
random_int = np.random.randint(low=0, high=len(obj_masks))
obj_mask_focus = obj_masks.pop(random_int)
print("OBJECT ID INIT=", obj_catids[random_int])
sum_interior = torch.sum(obj_mask_focus[50:self.W - 50,
50:self.H - 50])
depth = event.depth_frame
xs, ys = np.meshgrid(np.linspace(-1 * 256 / 2., 1 * 256 / 2., 256),
np.linspace(1 * 256 / 2., -1 * 256 / 2., 256))
depth = depth.reshape(1, 256, 256)
xs = xs.reshape(1, 256, 256)
ys = ys.reshape(1, 256, 256)
xys = np.vstack(
(xs * depth, ys * depth, -depth, np.ones(depth.shape)))
xys = xys.reshape(4, -1)
xy_c0 = np.matmul(np.linalg.inv(self.K), xys)
xyz = xy_c0.T[:, :3].reshape(256, 256, 3)
xyz_obj_masked = xyz[obj_mask_focus]
xyz_obj_masked = np.matmul(
rotation_, xyz_obj_masked.T) + translation_.reshape(3, 1)
xyz_obj_mid = np.mean(xyz_obj_masked, axis=1)
xyz_obj_mid[2] = -xyz_obj_mid[2]
else:
# want an object not on the edges of the image
sum_interior = 0
while True:
if len(obj_masks) == 0:
return None, None
random_int = np.random.randint(low=0, high=len(obj_masks))
obj_mask_focus = obj_masks.pop(random_int)
# print("OBJECT ID INIT=", obj_catids[random_int])
sum_interior = torch.sum(obj_mask_focus[50:self.W - 50,
50:self.H - 50])
if sum_interior < 500:
continue # exclude too small objects
pixel_locs_obj = np.where(obj_mask_focus.cpu().numpy())
x_mid = np.round(np.median(pixel_locs_obj[1]) / self.W, 4)
y_mid = np.round(np.median(pixel_locs_obj[0]) / self.H, 4)
if False:
plt.figure(1)
plt.clf()
plt.imshow(obj_mask_focus)
plt.plot(np.median(pixel_locs_obj[1]),
np.median(pixel_locs_obj[0]), 'x')
plt_name = self.homepath + '/seg_mask.png'
plt.savefig(plt_name)
event = self.controller.step('TouchThenApplyForce',
x=x_mid,
y=y_mid,
handDistance=1000000.0,
direction=dict(x=0.0,
y=0.0,
z=0.0),
moveMagnitude=0.0)
obj_focus_id = event.metadata['actionReturn']['objectId']
xyz_obj_mid = None
for o in event.metadata['objects']:
if o['objectId'] == obj_focus_id:
if o['objectType'] not in self.include_classes_final:
continue
xyz_obj_mid = np.array(
list(o['axisAlignedBoundingBox']
['center'].values()))
if xyz_obj_mid is not None:
break
print("MIDPOINT=", xyz_obj_mid)
return xyz_obj_mid, obj_mask_focus
def run(self, mode=None, total_loss=None, summ_writer=None):
event = self.controller.step('GetReachablePositions')
if not event.metadata['reachablePositions']:
# Different versions this is empty/full
event = self.controller.step(action='MoveAhead')
self.nav_pts = event.metadata['reachablePositions']
self.nav_pts = np.array([list(d.values()) for d in self.nav_pts])
# objects = np.random.choice(event.metadata['objects'], self.obj_per_scene, replace=False)
objects = event.metadata['objects']
objects_inds = np.arange(len(event.metadata['objects']))
np.random.shuffle(objects_inds)
# objects = np.random.shuffle(event.metadata['objects'])
# for obj in event.metadata['objects']: #objects:
# print(obj['name'])
# objects = objects[0]
successes = 0
# meta_obj_idx = 0
num_obs = 0
# while successes < self.obj_per_scene and meta_obj_idx <= len(event.metadata['objects']) - 1:
for obj in objects:
# if meta_obj_idx > len(event.metadata['objects']) - 1:
# print("OUT OF OBJECT... RETURNING")
# return total_loss, None, None, None, None
# obj = objects[objects_inds[meta_obj_idx]]
# meta_obj_idx += 1
print("Center object is ", obj['objectType'])
st()
# if obj['name'] in ['Microwave_b200e0bc']:
# print(obj['name'])
# else:
# continue
# print(obj['name'])
if obj['objectType'] not in self.include_classes:
print("Continuing... Invalid Object")
continue
# Calculate distance to object center
obj_center = np.array(
list(obj['axisAlignedBoundingBox']['center'].values()))
obj_center = np.expand_dims(obj_center, axis=0)
distances = np.sqrt(np.sum((self.nav_pts - obj_center)**2, axis=1))
# Get points with r_min < dist < r_max
valid_pts = self.nav_pts[np.where(
(distances > self.radius_min) * (distances < self.radius_max))]
# Bin points based on angles [vertical_angle (10 deg/bin), horizontal_angle (10 deg/bin)]
valid_pts_shift = valid_pts - obj_center
dz = valid_pts_shift[:, 2]
dx = valid_pts_shift[:, 0]
dy = valid_pts_shift[:, 1]
# Get yaw for binning
valid_yaw = np.degrees(np.arctan2(dz, dx))
if mode == "train":
nbins = 10 #20
else:
nbins = 5
bins = np.linspace(-180, 180, nbins + 1)
bin_yaw = np.digitize(valid_yaw, bins)
num_valid_bins = np.unique(bin_yaw).size
if False:
import matplotlib.cm as cm
colors = iter(cm.rainbow(np.linspace(0, 1, nbins)))
plt.figure(2)
plt.clf()
print(np.unique(bin_yaw))
for bi in range(nbins):
cur_bi = np.where(bin_yaw == (bi + 1))
points = valid_pts[cur_bi]
x_sample = points[:, 0]
z_sample = points[:, 2]
plt.plot(z_sample, x_sample, 'o', color=next(colors))
plt.plot(self.nav_pts[:, 2],
self.nav_pts[:, 0],
'x',
color='red')
plt.plot(obj_center[:, 2],
obj_center[:, 0],
'x',
color='black')
plt_name = '/home/nel/gsarch/aithor/data/valid.png'
plt.savefig(plt_name)
if num_valid_bins == 0:
continue
if mode == "train":
spawns_per_bin = 3 #20
else:
spawns_per_bin = 1 #int(self.num_views / num_valid_bins) + 2
# print(f'spawns_per_bin: {spawns_per_bin}')
action = "do_nothing"
episodes = []
valid_pts_selected = []
camXs_T_camX0_4x4 = []
camX0_T_camXs_4x4 = []
origin_T_camXs = []
origin_T_camXs_t = []
cnt = 0
for b in range(nbins):
# get all angle indices in the current bin range
inds_bin_cur = np.where(
bin_yaw == (b + 1)) # bins start 1 so need +1
inds_bin_cur = list(inds_bin_cur[0])
if len(inds_bin_cur) == 0:
continue
for s in range(spawns_per_bin):
observations = {}
if len(inds_bin_cur) == 0:
continue
rand_ind = np.random.randint(0, len(inds_bin_cur))
s_ind = inds_bin_cur.pop(rand_ind)
pos_s = valid_pts[s_ind]
valid_pts_selected.append(pos_s)
# add height from center of agent to camera
pos_s[1] = pos_s[1] + 0.675
turn_yaw, turn_pitch = self.utils.get_rotation_to_obj(
obj_center, pos_s)
event = self.controller.step('TeleportFull',
x=pos_s[0],
y=pos_s[1],
z=pos_s[2],
rotation=dict(x=0.0,
y=int(turn_yaw),
z=0.0),
horizon=int(turn_pitch))
rgb = event.frame
object_id = obj['objectId']
instance_detections2D = event.instance_detections2D
if object_id not in instance_detections2D:
print("NOT in instance detections 2D.. continuing")
continue
obj_instance_detection2D = instance_detections2D[
object_id] # [start_x, start_y, end_x, end_y]
max_len = np.max(
np.array([
obj_instance_detection2D[2] -
obj_instance_detection2D[0],
obj_instance_detection2D[3] -
obj_instance_detection2D[1]
]))
pad_len = max_len // 8
if pad_len == 0:
print("pad len 0.. continuing")
continue
x_center = (obj_instance_detection2D[3] +
obj_instance_detection2D[1]) // 2
x_low = x_center - max_len - pad_len
if x_low < 0:
x_low = 0
x_high = x_center + max_len + pad_len #x_low + max_len + 2*pad_len
if x_high > self.W:
x_high = self.W
y_center = (obj_instance_detection2D[2] +
obj_instance_detection2D[0]) // 2
y_low = y_center - max_len - pad_len #-pad_len
if y_low < 0:
y_low = 0
y_high = y_center + max_len + pad_len #y_low + max_len + 2*pad_len
if y_high > self.H:
y_high = self.H
rgb_crop = rgb[x_low:x_high, y_low:y_high, :]
rgb_crop = Image.fromarray(rgb_crop)
normalize_cropped_rgb = self.normalize(rgb_crop).unsqueeze(
0).double().cuda()
obj_features = self.vgg_feat_extractor(
normalize_cropped_rgb).view((512, -1))
obj_features = obj_features.detach().cpu().numpy()
# pca = PCA(n_components=10)
# obj_features = pca.fit_transform(obj_features.T).flatten()
# obj_features = torch.from_numpy(obj_features).view(-1).cuda()
obj_features = obj_features.flatten()
if mode == "train":
if self.first_time:
self.first_time = False
self.data_store_features = obj_features
# self.data_store_features = self.data_store_features.cuda()
self.feature_obj_ids.append(
self.ithor_to_maskrcnn[obj['objectType']])
else:
# self.data_store_features = torch.vstack((self.data_store_features, obj_features))
self.data_store_features = np.vstack(
(self.data_store_features, obj_features))
self.feature_obj_ids.append(
self.ithor_to_maskrcnn[obj['objectType']])
elif mode == "test":
# obj_features = obj_features.unsqueeze(0)
# inner_prod = torch.abs(torch.mm(obj_features, self.data_store_features.T)).squeeze()
# inner_prod = inner_prod.detach().cpu().numpy()
# dist = np.squeeze(np.abs(np.matmul(obj_features, self.data_store_features.transpose())))
dist = np.linalg.norm(self.data_store_features -
obj_features,
axis=1)
k = 10
ind_knn = list(np.argsort(dist)[:k])
dist_knn = np.sort(dist)[:k]
dist_knn_norm = list(
self.Softmax(torch.from_numpy(-dist_knn)).numpy())
match_knn_id = [
self.feature_obj_ids[i] for i in ind_knn
]
# for i in range(1, len(match_knn_id)):
# add softmax values from the same class (probably a really complex way of doing this)
idx = 0
dist_knn_norm_add = []
match_knn_id_add = []
while True:
if not match_knn_id:
break
match_knn_cur = match_knn_id.pop(0)
dist_knn_norm_cur = dist_knn_norm.pop(0)
match_knn_id_add.append(match_knn_cur)
idxs_ = []
for i in range(len(match_knn_id)):
if match_knn_id[i] == match_knn_cur:
dist_knn_norm_cur += dist_knn_norm[i]
# match_knn_id_.pop(i)
else:
idxs_.append(i)
match_knn_id = [match_knn_id[idx] for idx in idxs_]
dist_knn_norm = [
dist_knn_norm[idx] for idx in idxs_
]
dist_knn_norm_add.append(dist_knn_norm_cur)
dist_knn_norm_add = np.array(dist_knn_norm_add)
dist_knn_argmax = np.argmax(dist_knn_norm_add)
match_nn_id = match_knn_id_add[
dist_knn_argmax] #self.feature_obj_ids[ind_nn]
match_nn_catname = self.maskrcnn_to_ithor[match_nn_id]
self.best_inner_prods.append(
dist_knn_norm_add[dist_knn_argmax])
self.pred_ids.append(match_nn_id)
# self.pred_catnames.append(match_nn_catname)
self.true_ids.append(
self.ithor_to_maskrcnn[obj['objectType']])
self.pred_catnames.append(match_nn_catname)
self.true_catnames.append(obj['objectType'])
print(match_nn_catname)
self.data_store_features = np.vstack(
(self.data_store_features, obj_features))
self.feature_obj_ids.append(
self.ithor_to_maskrcnn[obj['objectType']])
if False:
normalize_cropped_rgb = np.transpose(
normalize_cropped_rgb.squeeze(
0).detach().cpu().numpy(), (1, 2, 0))
plt.figure(1)
plt.clf()
plt.imshow(normalize_cropped_rgb)
# plt_name = self.homepath + '/seg_init.png'
plt.figure(2)
plt.clf()
plt.imshow(rgb)
plt.show()
plt.figure(3)
plt.clf()
plt.imshow(np.array(rgb_crop))
plt.show()
