def generate_semantic_layers(env:habitat.RLEnv,mapper:Mapper, M:int, mapper_config:CN, global_seen_map_np:torch.Tensor):
grid_size = mapper_config.TASK.GT_EGO_MAP.MAP_SCALE
start_y = env.habitat_env.sim.get_agent_state().position[1]
# init the semantic layers
semantic_layers = []
scene = env.habitat_env.sim.semantic_scene
obj_pos = {}
for name in label_idx.keys():
obj_pos[name] = {obj.id : [obj.aabb.center,obj.aabb.sizes]
for obj in scene.objects
if name == obj.category.name()}
for objects in tqdm.tqdm(obj_pos,desc='semantic map', position=3):
if objects == 'chair':
obj_mapping = np.zeros((M,M))
gc.collect()
for object in obj_pos[objects].items():
pos = pos_real_to_map(object[1][0], env.habitat_env.current_episode)
obj_height = object[1][0][1]
# Check if obj is part of the current floor. the centre of the floor is 1 meter higher than the start position and we look for object inside a 2 meters bound.
if np.abs(start_y - obj_height + 1) > 1:
continue
edges_w = max(int(object[1][1][0] / grid_size), 1)
edges_l = max(int(object[1][1][2] / grid_size), 1)
trans_obj = asnumpy(map_ego_to_global(mapper, edges_w, edges_l,torch.Tensor(pos).unsqueeze(0),M).squeeze())
obj_mapping = (obj_mapping + trans_obj) > 0
reduced_obj_layer = obj_mapping
semantic_layers.append(reduced_obj_layer)
return np.array(semantic_layers)
def _load_transformed_wall_maps(self, scene_map_info, episode):
seen_maps = []
wall_maps = []
start_position = episode.start_position # (X, Y, Z)
start_rotation = quaternion_xyzw_to_wxyz(episode.start_rotation)
start_heading = compute_heading_from_quaternion(start_rotation)
for floor_data in scene_map_info:
seen_map = np.load(floor_data["seen_map_path"])
wall_map = np.load(floor_data["wall_map_path"])
# ===== Transform the maps relative to the episode start pose =====
map_view_position = floor_data["world_position"]
map_view_heading = floor_data["world_heading"]
# Originally, Z is downward and X is rightward.
# Convert it to X upward and Y rightward
x_map, y_map = -map_view_position[2], map_view_position[0]
theta_map = map_view_heading
x_start, y_start = -start_position[2], start_position[0]
theta_start = start_heading
# Compute relative coordinates
r_rel = math.sqrt((x_start - x_map) ** 2 + (y_start - y_map) ** 2)
phi_rel = math.atan2(y_start - y_map, x_start - x_map) - theta_map
x_rel = r_rel * math.cos(phi_rel) / self.config.MAP_SCALE
y_rel = r_rel * math.sin(phi_rel) / self.config.MAP_SCALE
theta_rel = theta_start - theta_map
# Convert these to image coordinates with X being rightward and Y
# being downward
x_img_rel = y_rel
y_img_rel = -x_rel
theta_img_rel = theta_rel
x_trans = torch.Tensor([[x_img_rel, y_img_rel, theta_img_rel]])
# Perform the transformations
p_seen_map = rearrange(torch.Tensor(seen_map), "h w c -> () c h w")
p_wall_map = rearrange(torch.Tensor(wall_map), "h w c -> () c h w")
p_seen_map_trans = spatial_transform_map(p_seen_map, x_trans)
p_wall_map_trans = spatial_transform_map(p_wall_map, x_trans)
seen_map_trans = asnumpy(p_seen_map_trans)
seen_map_trans = rearrange(seen_map_trans, "() c h w -> h w c")
wall_map_trans = asnumpy(p_wall_map_trans)
wall_map_trans = rearrange(wall_map_trans, "() c h w -> h w c")
seen_maps.append(seen_map_trans)
wall_maps.append(wall_map_trans)
return seen_maps, wall_maps
def test_chainer_layer():
chainer_is_present = any(
'chainer' in backend.framework_name for backend in collect_test_backends(symbolic=False, layers=True)
)
if chainer_is_present:
# checked that gluon present
import chainer
import chainer.links as L
import chainer.functions as F
from einops.layers.chainer import Rearrange, Reduce, EinMix
from einops import asnumpy
import numpy as np
def create_model():
return chainer.Sequential(
L.Convolution2D(3, 6, ksize=(5, 5)),
Reduce('b c (h h2) (w w2) -> b c h w', 'max', h2=2, w2=2),
L.Convolution2D(6, 16, ksize=(5, 5)),
Reduce('b c (h h2) (w w2) -> b c h w', 'max', h2=2, w2=2),
Rearrange('b c h w -> b (c h w)'),
L.Linear(16 * 5 * 5, 120),
L.Linear(120, 84),
F.relu,
EinMix('b c1 -> (b c2)', weight_shape='c1 c2', bias_shape='c2', c1=84, c2=84),
EinMix('(b c2) -> b c3', weight_shape='c2 c3', bias_shape='c3', c2=84, c3=84),
L.Linear(84, 10),
)
model1 = create_model()
model2 = create_model()
x = np.random.normal(size=[10, 3, 32, 32]).astype('float32')
x = chainer.Variable(x)
assert not numpy.allclose(asnumpy(model1(x)), asnumpy(model2(x)))
with tempfile.TemporaryDirectory() as dir:
filename = f'{dir}/file.npz'
chainer.serializers.save_npz(filename, model1)
chainer.serializers.load_npz(filename, model2)
assert numpy.allclose(asnumpy(model1(x)), asnumpy(model2(x)))
def test_gluon_layer():
gluon_is_present = any(
'mxnet' in backend.framework_name for backend in collect_test_backends(symbolic=False, layers=True)
)
if gluon_is_present:
# checked that gluon present
import mxnet
from mxnet.gluon.nn import HybridSequential, Dense, Conv2D, LeakyReLU
from einops.layers.gluon import Rearrange, Reduce, EinMix
from einops import asnumpy
def create_model():
model = HybridSequential()
layers = [
Conv2D(6, kernel_size=5),
Reduce('b c (h h2) (w w2) -> b c h w', 'max', h2=2, w2=2),
Conv2D(16, kernel_size=5),
Reduce('b c (h h2) (w w2) -> b c h w', 'max', h2=2, w2=2),
Rearrange('b c h w -> b (c h w)'),
Dense(120),
LeakyReLU(alpha=0.0),
Dense(84),
LeakyReLU(alpha=0.0),
Dense(10),
]
for layer in layers:
model.add(layer)
model.initialize(mxnet.init.Xavier(), ctx=mxnet.cpu())
return model
model1 = create_model()
model2 = create_model()
x = mxnet.ndarray.random_normal(shape=[10, 3, 32, 32])
assert not numpy.allclose(asnumpy(model1(x)), asnumpy(model2(x)))
with tempfile.NamedTemporaryFile(mode='r+b') as fp:
model1.save_parameters(fp.name)
model2.load_parameters(fp.name)
assert numpy.allclose(asnumpy(model1(x)), asnumpy(model2(x)))
# testing with symbolic (NB with fixed dimensions!)
input = mxnet.sym.Variable('data', shape=x.shape)
json = model1(input).tojson()
model3 = mxnet.gluon.SymbolBlock(outputs=mxnet.sym.load_json(json), inputs=input)
model4 = mxnet.gluon.SymbolBlock(outputs=mxnet.sym.load_json(json), inputs=input)
model3.initialize(ctx=mxnet.cpu())
model3(x)
with tempfile.NamedTemporaryFile(mode='r+b') as fp:
model3.save_parameters(fp.name)
model4.load_parameters(fp.name)
assert numpy.allclose(asnumpy(model3(x)), asnumpy(model4(x)))
try:
model1.hybridize(static_alloc=True, static_shape=True)
model1(x)
except:
# hybridization is not supported
pass
def unprocess_image(img):
"""Undo imagenet normalization to a batch of images."""
# img - (bs, C, H, W)
mean = [0.485, 0.456, 0.406]
std = [0.229, 0.224, 0.225]
img_unproc = np.copy(asnumpy(img))
img_unproc[:, 0] = img_unproc[:, 0] * std[0] + mean[0]
img_unproc[:, 1] = img_unproc[:, 1] * std[1] + mean[1]
img_unproc[:, 2] = img_unproc[:, 2] * std[2] + mean[2]
img_unproc = np.clip(img_unproc, 0.0, 1.0) * 255.0
img_unproc = img_unproc.astype(np.uint8)
img_unproc = rearrange(img_unproc, "b c h w -> b h w c")
return img_unproc
def forward(self, all_voting_maps_):
"""Each visually similar viewpoint votes for the location of the
target viewpoint in the top-down view. These voting maps are then
summed and the final pose is estimated by taking the location with the
highest overall vote.
Inputs:
all_voting_maps_ - (k, bs, 1, mh, mw)
Outputs:
pose - (bs, num_poses) tensor
mask_maxmin - (bs, ) tensor
"""
bs = all_voting_maps_.shape[1]
# ======================== Aggregate votes ============================
all_voting_maps = all_voting_maps_.sum(dim=0) # (bs, 1, mh, mw)
# Add small non-zero votes to all locations
all_voting_maps = all_voting_maps + 1e-9
voting_sum = reduce(all_voting_maps, "b c h w -> b () () ()", "sum")
voting_map = all_voting_maps / voting_sum # (bs, 1, mh, mw)
# =================== Compute argmax locations ========================
device = voting_map.device
mh, mw = voting_map.shape[2], voting_map.shape[3]
midx, midy = mw // 2, mh // 2
voting_map_flat = rearrange(asnumpy(voting_map[:, 0]), "b h w -> b (h w)")
predy, predx = np.unravel_index(np.argmax(voting_map_flat, axis=1), (mh, mw))
predy = (torch.Tensor(predy).to(device).float() - midy) * self.map_scale
predx = (torch.Tensor(predx).to(device).float() - midx) * self.map_scale
predr = torch.norm(torch.stack([predx, predy], dim=1), dim=1)
predt = torch.atan2(predy, predx)
# Handle edge-case where no votes were cast. When no votes are cast,
# the pose prediction is set to zeros by default.
diff_maxmin = np.max(voting_map_flat, axis=1) - np.min(
voting_map_flat, axis=1
) # (bs, )
mask_maxmin = (torch.Tensor(diff_maxmin) == 0).to(device)
predr[mask_maxmin] = 0
predt[mask_maxmin] = 0
# Converts predicted pose to required format.
pose = process_pose(torch.stack([predr, predt], dim=1))
return pose, mask_maxmin
def get_eval_map(env:ExpRLEnv, trainer:SemAntExpTrainer, M:int,depth_projection_net:DepthProjectionNet, step:int, objectName:str):
"""Given the environment and the configuration, compute the global
top-down wall and seen area maps by sampling maps for individual locations
along a uniform grid in the environment, and registering them.
step (m): the length between two measure points
"""
# Initialize a global map for the episode
scene_name = env.habitat_env.current_episode.scene_id.split('/')[-1].split('.')[0]
dirPath = './data/debug/data/' + objectName + "/" + scene_name + '_' + env.habitat_env.current_episode.episode_id + '/' + "record/"
safe_mkdir(dirPath)
mapper = trainer.mapper
config = trainer.config.TASK_CONFIG
device = trainer.device
scene = env.habitat_env.sim.semantic_scene
obj_pos = {}
obj_pos = {obj.id : [obj.aabb.center,obj.aabb.sizes]
for obj in scene.objects
if objectName == obj.category.name()}
with open(dirPath + 'obj_pos.pkl', 'wb') as f: pickle.dump(obj_pos, f)
objectIndexes = [int(key.split('_')[-1]) for key in obj_pos.keys()]
global_wall_map = torch.zeros(1, 2, M, M).to(device)
global_seen_map = torch.zeros(1, 2, M, M).to(device)
global_object_map = torch.zeros(1, 2, M, M).to(device)
global_se_map = torch.zeros(1, 2, M, M).to(device)
global_se_filtered_map = torch.zeros(1, 2, M, M).to(device)
global_se_seen_map = torch.zeros(1, 2, M, M).to(device)
global_se_map_mit = torch.zeros(1, 2, M, M).to(device)
global_se_map_gt = torch.zeros(1, 2, M, M).to(device)
grid_size = config.TASK.GT_EGO_MAP.MAP_SCALE
coordinate_max = maps.COORDINATE_MAX
coordinate_min = maps.COORDINATE_MIN
resolution = (coordinate_max - coordinate_min) / grid_size
grid_resolution = (int(resolution), int(resolution))
top_down_map = maps.get_topdown_map(
env.habitat_env.sim, grid_resolution, 20000, draw_border=False,
)
map_w, map_h = top_down_map.shape
intervals = (max(int(step / grid_size), 1), max(int(step / grid_size), 1))
x_vals = np.arange(0, map_w, intervals[0], dtype=int)
y_vals = np.arange(0, map_h, intervals[1], dtype=int)
coors = np.stack(np.meshgrid(x_vals, y_vals), axis=2) # (H, W, 2)
coors = coors.reshape(-1, 2) # (H*W, 2)
map_vals = top_down_map[coors[:, 0], coors[:, 1]]
valid_coors = coors[map_vals > 0]
real_x_vals = coordinate_max - valid_coors[:, 0] * grid_size
real_z_vals = coordinate_min + valid_coors[:, 1] * grid_size
start_y = env.habitat_env.sim.get_agent_state().position[1]
index = 0
records = []
totalOffsetX = np.random.uniform(-1.5, 1.5)
totalOffsetY = np.random.uniform(-1.5, 1.5)
for j in tqdm.tqdm(range(real_x_vals.shape[0]),desc='occupacy map', position=2):
for theta in np.arange(-np.pi, np.pi, np.pi ):
index +=1
randomAngle = np.random.uniform(-np.pi,np.pi)
randomX = np.random.uniform(-0.5, 0.5)
randomY = np.random.uniform(-0.5, 0.5)
position = [
real_x_vals[j].item() + randomX + totalOffsetX,
start_y.item(),
real_z_vals[j].item() + randomY + totalOffsetY,
]
rotation = [
0.0,
np.sin((theta + randomAngle) / 2).item(),
0.0,
np.cos((theta + randomAngle) / 2).item(),
]
sim_obs = env.habitat_env.sim.get_observations_at(
position, rotation, keep_agent_at_new_pose=True,
)
episode = env.habitat_env.current_episode
obs = obs, _, _, _ = env.step(action={'action':1})
batch = batch_obs([obs], device=trainer.device)
semantic = batch["semantic"]
mask = torch.zeros_like(semantic,dtype=bool)
for objectIndex in objectIndexes:
mask ^= (semantic == objectIndex)
mask = rearrange(mask, "b h w -> b h w ()")
batch["object_mask_gt"] = mask
ego_map_gt_b = depth_projection_net(
rearrange(batch["depth"], "b h w c -> b c h w")
)
object_depth = batch["depth"]*(mask > 0) + (mask == 0)*100
ego_map_gt_object = depth_projection_net(rearrange(object_depth, "b h w c -> b c h w"))
batch["ego_map_gt_object"] = ego_map_gt_object
#begin use MIT Seg
img_data = pil_to_tensor(rearrange(batch["rgb"], "b h w c -> b c h w")[0] / 255)
singleton_batch = {'img_data': img_data[None]}
output_size = img_data.shape[1:]
with torch.no_grad():
scores = segmentation_module(singleton_batch, segSize=output_size)
_, pred = torch.max(scores, dim=1)
# visualize_result(pred)
#end use MIT Seg
batch["mit"] = pred
#mask for chair
if objectName == "chair":
mask = (pred == 19)^(pred == 30)^(pred == 75)
elif objectName == "bed":
mask = (pred == 7)
elif objectName == "table":
mask = (pred == 15)
else:
mask = pred > 0
mask = rearrange(mask, "b h w -> b h w ()")
batch["object_mask_mit"] = mask
object_depth2 = batch["depth"]*(mask > 0) + (mask == 0)*100
ego_map_gt_object2 = depth_projection_net(rearrange(object_depth2, "b h w c -> b c h w"))
batch["ego_map_mit_object"] = ego_map_gt_object2
ego_map_gt_anti = transpose_image(batch["ego_map_gt_anticipated"]).to(trainer.device)
pu_inputs_t = {
"rgb": process_image(batch["rgb"], trainer.mapper.img_mean_t, trainer.mapper.img_std_t).to(trainer.device),
"depth": transpose_image(batch["depth"]).to(trainer.device),
"ego_map_gt": transpose_image(
rearrange(ego_map_gt_b, "b c h w -> b h w c")
).to(trainer.device),
"ego_map_gt_anticipated": ego_map_gt_anti,
}
pu_inputs = trainer.mapper._safe_cat(pu_inputs_t, pu_inputs_t)
torch.save(batch, dirPath + 'pu_inputs_t_'+str(index)+'.pt')
pu_output = trainer.mapper.projection_unit(pu_inputs)
se_map = asnumpy(pu_output['sem_estimate'])[0,:,:,:]
ego_map_gt_b_np = asnumpy(ego_map_gt_b[0,...])
se_seen = ego_map_gt_b_np*se_map
dilation_mask = np.ones((100, 100))
current_mask = cv2.dilate(
se_seen.astype(np.float32), dilation_mask, iterations=1,
).astype(np.float32)
se_filtered_map = se_map
se_seen_map = torch.Tensor(se_seen).to(device)
se_seen_map = rearrange(se_seen_map, "c h w -> () c h w")
se_seen_map = (se_seen_map[:,0])[None,:]
se_map = torch.Tensor(se_map).to(device)
se_map = rearrange(se_map, "c h w -> () c h w")
se_filtered_map = torch.Tensor(se_filtered_map).to(device)
se_filtered_map = rearrange(se_filtered_map, "c h w -> () c h w")
torch.save(se_filtered_map, dirPath + 'se_filtered_map_'+str(index)+'.pt')
ego_map_gt = torch.Tensor(obs["ego_map_gt"]).to(device)
ego_map_gt = rearrange(ego_map_gt, "h w c -> () c h w")
ego_wall_map_gt = torch.Tensor(obs["ego_wall_map_gt"]).to(device)
ego_wall_map_gt = rearrange(ego_wall_map_gt, "h w c -> () c h w")
pose_gt = torch.Tensor(obs["pose_gt"]).unsqueeze(0).to(device)
global_se_seen_map = mapper.ext_register_map(
global_se_seen_map, se_seen_map, pose_gt
)
global_se_map_gt = mapper.ext_register_map(
global_se_map_gt, ego_map_gt_anti, pose_gt
)
global_seen_map = mapper.ext_register_map(
global_seen_map, ego_map_gt_b, pose_gt
)
global_object_map = mapper.ext_register_map(
global_object_map, ego_map_gt_object, pose_gt
)
global_wall_map = mapper.ext_register_map(
global_wall_map, ego_wall_map_gt, pose_gt
)
global_se_map = mapper.ext_register_map(
global_se_map, se_map, pose_gt
)
global_se_filtered_map = mapper.ext_register_map(
global_se_filtered_map, se_filtered_map, pose_gt
)
global_se_map_mit = mapper.ext_register_map(
global_se_map_mit, ego_map_gt_object2, pose_gt
)
se_filtered_map = None
batch = None
se_map = None
ego_map_gt = None
ego_wall_map_gt = None
obs = None
pu_inputs_t = None
pu_output=None
_ = None
pu_inputs=None
gc.collect()
torch.cuda.empty_cache()
mask = dilate_tensor(global_seen_map*(1 - global_wall_map > 0)*global_se_map,(51, 51))
global_se_filtered_map = global_se_filtered_map*mask
global_se_seen_map = global_se_seen_map*mask
mask = dilate_tensor(global_seen_map*(1 - global_wall_map > 0)*global_object_map,(51, 51))
global_se_map_gt = global_se_map_gt*mask
global_wall_map_np = asnumpy(rearrange(global_wall_map, "b c h w -> b h w c")[0])
global_seen_map_np = asnumpy(rearrange(global_seen_map, "b c h w -> b h w c")[0])
global_se_map_np = asnumpy(rearrange(global_se_map, "b c h w -> b h w c")[0])
global_se_global_se_filtered_map_np= asnumpy(rearrange(global_se_filtered_map, "b c h w -> b h w c")[0])
global_se_map_mit_np= asnumpy(rearrange(global_se_map_mit, "b c h w -> b h w c")[0])
global_se_map_gt_np= asnumpy(rearrange(global_se_map_gt, "b c h w -> b h w c")[0])
global_se_seen_map_np = asnumpy(rearrange(global_se_seen_map, "b c h w -> b h w c")[0])
global_object_map_np = asnumpy(rearrange(global_object_map, "b c h w -> b h w c")[0])
return global_seen_map_np, global_wall_map_np, global_se_map_np,global_se_global_se_filtered_map_np, global_se_map_mit_np,global_se_map_gt_np,global_se_seen_map_np,global_object_map_np
safe_mkdir(mapPath)
save_episode_info(env, mapPath + "epInfo")
for key in mapsDic.keys():
np.save(mapPath + key, mapsDic[key])
dataPath = './data/debug/data/' + objectName + "/" + scene_name + '_' + env.habitat_env.current_episode.episode_id + '/' + "record/"
imgPath = './data/debug/img/' + objectName + "/" + scene_name + '_' + env.habitat_env.current_episode.episode_id +"/" + "ep" + str(i) + '/'
file = open(dataPath + "obj_pos.pkl",'rb')
obj_pos = pickle.load(file)
filelist = os.listdir(dataPath)
for id in range(len(filelist) // 2 ):
idx = id+1
se_filtered_map = torch.load(dataPath + 'se_filtered_map_'+str(idx)+'.pt')
pu_inputs_t = torch.load(dataPath + 'pu_inputs_t_'+str(idx)+'.pt')
rbg_numpy = asnumpy(pu_inputs_t["rgb"][0])/255
semantic_numpy = asnumpy(pu_inputs_t["semantic"][0]) % 40
indexes = np.unique(semantic_numpy)
if pu_inputs_t["ego_map_gt_anticipated"][0,:,:,0].sum() > 0:
safe_mkdir(imgPath + str(idx))
visualize_result(pu_inputs_t['mit'],imgPath, colors)
plt.imsave(imgPath + str(idx)+'/predict_'+ str(idx) + '.png', asnumpy(se_filtered_map[0,0]))
plt.imsave(imgPath + str(idx)+'/ego_map_gt_anticipated_'+ str(idx) + '.png', asnumpy(pu_inputs_t["ego_map_gt_anticipated"][0,:,:,0]))
plt.imsave(imgPath + str(idx)+'/rbg_image_'+ str(idx) + '.png', rbg_numpy)
plt.imsave(imgPath + str(idx)+'/ego_map_gt_'+ str(idx) + '.png', asnumpy(pu_inputs_t["ego_map_gt"][0,:,:,0]))
plt.imsave(imgPath + str(idx)+'/depth_image_'+ str(idx) + '.png', asnumpy(np.squeeze(pu_inputs_t["depth"][0,:,:,0])))
plt.imsave(imgPath + str(idx)+'/object_mask_'+ str(idx) + '.png', asnumpy(np.squeeze(pu_inputs_t["object_mask_gt"][0,:,:,0])))
plt.imsave(imgPath + str(idx)+'/ego_map_ob_gt_'+ str(idx) + '.png', asnumpy(np.squeeze(pu_inputs_t["ego_map_gt_object"][0,0,:,:])))
plt.imsave(imgPath + str(idx)+'/ego_map_ob_mit_'+ str(idx) + '.png', asnumpy(np.squeeze(pu_inputs_t["ego_map_mit_object"][0,0,:,:])))
safe_mkdir('./data/debug/data/' + objectName + '/' + scene_name + '_' + env.habitat_env.current_episode.episode_id + '/iouscene/')
def _compute_plans_and_local_goals(
self, global_map, agent_map_xy, SAMPLE_LOCAL_GOAL_FLAGS,
):
"""
global_map - (bs, 2, V, V) tensor
agent_map_xy - (bs, 2) agent's current position on the map
"""
s = self.map_scale
goal_map_xy = self.states["curr_global_goals"]
plans_xy = self._compute_plans(
global_map,
agent_map_xy,
goal_map_xy,
SAMPLE_LOCAL_GOAL_FLAGS,
cache_map=True,
crop_map_flag=self.config.crop_map_for_planning,
)
# ========================= Sample a local goal =======================
# Sample a local goal and measure the shortest path distance to that
# goal according to the current map.
# Update step counts for sample_random_explored calls
self.states["sample_random_explored_timer"] += 1
# Pick a local goal to reach with the local policy
for i in range(self.nplanners):
if SAMPLE_LOCAL_GOAL_FLAGS[i] != 1:
continue
path_x, path_y = plans_xy[i]
# If planning failed, sample random local goal.
if path_x is None:
# Note: This is an expensive call, especially when the map is messy
# and planning keeps failing. Call sample_random_explored only after
# so many steps elapsed since the last call.
if self.config.recovery_heuristic == "random_explored_towards_goal":
if self.states["sample_random_explored_timer"][i].item() > 10:
goal_x, goal_y = self._sample_random_explored_towards_goal(
global_map[i],
asnumpy(agent_map_xy[i]).tolist(),
asnumpy(goal_map_xy[i]).tolist(),
s,
)
# Reset count
self.states["sample_random_explored_timer"][i] = 0
else:
goal_x, goal_y = self._sample_random_towards_goal(
global_map[i],
asnumpy(agent_map_xy[i]).tolist(),
asnumpy(goal_map_xy[i]).tolist(),
s,
)
goal_x, goal_y = asnumpy(agent_map_xy[i]).tolist()
elif self.config.recovery_heuristic == "random_explored":
if self.states["sample_random_explored_timer"][i].item() > 10:
goal_x, goal_y = self._sample_random_explored(
global_map[i], asnumpy(agent_map_xy[i]).tolist(), s
)
# Reset count
self.states["sample_random_explored_timer"][i] = 0
else:
goal_x, goal_y = self._sample_random_near_agent(
global_map[i], asnumpy(agent_map_xy[i]).tolist(), s
)
# goal_x, goal_y = asnumpy(agent_map_xy[i]).tolist()
else:
raise ValueError
# When planning fails, default to euclidean distance
curr_x, curr_y = agent_map_xy[i].tolist()
splength = (
math.sqrt((goal_x - curr_x) ** 2 + (goal_y - curr_y) ** 2) * s
)
else:
dl = min(int(self.planning_step_mts / s), len(path_x) - 1)
# The path is in reverse order
goal_x, goal_y = path_x[-dl], path_y[-dl]
sp_xy = np.array([path_x[-dl:], path_y[-dl:]]).T # (dl, 2)
splength = (
np.linalg.norm(sp_xy[:-1] - sp_xy[1:], axis=1).sum().item() * s
)
# Ensure goals are within map bounds
goal_x = np.clip(goal_x, 0, global_map[i].shape[-1] - 1).item()
goal_y = np.clip(goal_y, 0, global_map[i].shape[-2] - 1).item()
# Set the local goals as well as the corresponding path length
# measures
self.states["curr_local_goals"][i, 0] = goal_x
self.states["curr_local_goals"][i, 1] = goal_y
self.states["local_shortest_path_length"][i] = splength
self.states["local_path_length"][i] = 0.0
def _compute_gt_local_action(
self, global_map, agent_world_xyt, goal_world_xyt, M, s
):
"""
Estimate the shortest-path action from agent position to goal position.
"""
agent_map_xy = convert_world2map(agent_world_xyt, (M, M), s)
goal_map_xy = convert_world2map(goal_world_xyt, (M, M), s)
# ============ Crop a region covering agent_map_xy and goal_map_xy ============
abs_delta_xy = torch.abs(agent_map_xy - goal_map_xy)
S = max(int(torch.max(abs_delta_xy).item()), 80)
old_center_xy = (agent_map_xy + goal_map_xy) // 2 # (bs, 2)
cropped_global_map = crop_map(global_map, old_center_xy, int(S))
global_map_np = self._process_maps(cropped_global_map) # (bs, M, M)
new_center_xy = torch.zeros_like(old_center_xy)
new_center_xy[:, 0] = global_map_np.shape[2] // 2
new_center_xy[:, 1] = global_map_np.shape[1] // 2
# ================ Transform points to new coordinate system ==================
new_agent_map_xy = agent_map_xy + new_center_xy - old_center_xy
new_goal_map_xy = goal_map_xy + new_center_xy - old_center_xy
map_xy_np = asnumpy(new_agent_map_xy).astype(np.int32) # (bs, 2)
goal_xy_np = asnumpy(new_goal_map_xy).astype(np.int32)
# Ensure that points do not go outside map limits
map_W = global_map_np.shape[2]
map_H = global_map_np.shape[1]
map_xy_np[:, 0] = np.clip(map_xy_np[:, 0], 0, map_W - 1)
map_xy_np[:, 1] = np.clip(map_xy_np[:, 1], 0, map_H - 1)
goal_xy_np[:, 0] = np.clip(goal_xy_np[:, 0], 0, map_W - 1)
goal_xy_np[:, 1] = np.clip(goal_xy_np[:, 1], 0, map_H - 1)
sample_flag = [1.0 for _ in range(self.nplanners)]
# Compute plan
plans_xy = self.planner.plan(global_map_np, map_xy_np, goal_xy_np, sample_flag)
# ===================== Sample an action to a nearby point ====================
# 0 is forward, 1 is left, 2 is right
gt_actions = []
forward_step = self.config.LOCAL_POLICY.AGENT_DYNAMICS.forward_step
turn_angle = math.radians(self.config.LOCAL_POLICY.AGENT_DYNAMICS.turn_angle)
for i in range(self.nplanners):
path_x, path_y = plans_xy[i]
# If planning failed, sample random action
if path_x is None:
gt_actions.append(random.randint(0, 2))
continue
# Plan to navigate to a point 0.5 meters away
dl = min(int(0.5 / s), len(path_x) - 1)
# The path is in reverse order
goal_x, goal_y = path_x[-dl], path_y[-dl]
agent_x, agent_y = map_xy_np[i].tolist()
# Decide action
agent_heading = agent_world_xyt[i, 2].item() - math.pi / 2
reqd_heading = math.atan2(goal_y - agent_y, goal_x - agent_x)
diff_angle = reqd_heading - agent_heading
diff_angle = math.atan2(math.sin(diff_angle), math.cos(diff_angle))
# Move forward if facing the correct direction
if abs(diff_angle) < 1.5 * turn_angle:
gt_actions.append(0)
# Turn left if the goal is to the left
elif diff_angle < 0:
gt_actions.append(1)
# Turn right otherwise
else:
gt_actions.append(2)
return torch.Tensor(gt_actions).unsqueeze(1).long()
def act(
self,
observations,
prev_observations,
prev_state_estimates,
ep_time,
masks,
deterministic=False,
):
# ============================ Set useful variables ===========================
ep_step = ep_time[0].item()
M = prev_state_estimates["map_states"].shape[2]
s = self.map_scale
device = observations['rgb'].device
assert M % 2 == 1, "The code is tested only for odd map sizes!"
# =================== Update states from current observation ==================
# Update map, pose and visitation map
mapper_inputs = self._create_mapper_inputs(
observations, prev_observations, prev_state_estimates
)
mapper_outputs = self.mapper(mapper_inputs)
global_map = mapper_outputs["mt"]
global_pose = mapper_outputs["xt_hat"]
if self.config.MAPPER.debug_log:
print(f"Global Pose: {global_pose}")
map_xy = convert_world2map(global_pose[:, :2], (M, M), s)
map_xy = torch.clamp(map_xy, 0, M - 1)
visited_states = self._update_state_visitation(
prev_state_estimates["visited_states"], map_xy
) # (bs, 1, M, M)
# Update local ANM state variables
curr_map_position = map_xy
if ep_step > 0:
# Compute state updates
prev_dist2localgoal = self.states["curr_dist2localgoal"]
curr_dist2localgoal = self._compute_dist2localgoal(
global_map, map_xy, self.states["curr_local_goals"],
)
prev_map_position = self.states["curr_map_position"]
prev_step_size = (
torch.norm(curr_map_position - prev_map_position, dim=1).unsqueeze(1)
* s
)
# Update the state variables
self.states["prev_dist2localgoal"] = prev_dist2localgoal
self.states["curr_dist2localgoal"] = curr_dist2localgoal
self.states["prev_map_position"] = prev_map_position
self.states["local_path_length"] += prev_step_size
self.states["curr_map_position"] = curr_map_position
# Initialize collision and visited maps at t=0
if ep_step == 0:
self.states["collision_map"] = torch.zeros(self.nplanners, M, M).to(device)
self.states["visited_map"] = torch.zeros(self.nplanners, M, M).to(device)
self.states["col_width"] = torch.ones(self.nplanners)
# Monitors number of steps elapsed since last call to sample random explored
self.states["sample_random_explored_timer"] = torch.zeros(self.nplanners)
if ep_step > 0:
# Update collision maps
forward_step = self.config.LOCAL_POLICY.AGENT_DYNAMICS.forward_step
for i in range(self.nplanners):
prev_action_i = observations["prev_actions"][i, 0].item()
# If not forward action, skip
if prev_action_i != 0:
continue
x1, y1 = asnumpy(self.states["prev_map_position"][i]).tolist()
x2, y2 = asnumpy(self.states["curr_map_position"][i]).tolist()
t2 = global_pose[i, 2].item() - math.pi / 2
if abs(x1 - x2) < 1 and abs(y1 - y2) < 1:
self.states["col_width"][i] += 2
self.states["col_width"][i] = min(self.states["col_width"][i], 9)
else:
self.states["col_width"][i] = 1
dist_trav_i = math.sqrt((x1 - x2) ** 2 + (y1 - y2) ** 2) * s
# Add an obstacle infront of the agent if a collision happens
if dist_trav_i < 0.7 * forward_step: # Collision
length = 2
width = int(self.states["col_width"][i].item())
buf = 3
cmH, cmW = self.states["collision_map"][i].shape
for j in range(length):
for k in range(width):
wx = (
x2
+ ((j + buf) * math.cos(t2))
+ ((k - width / 2) * math.sin(t2))
)
wy = (
y2
+ ((j + buf) * math.sin(t2))
- ((k - width / 2) * math.cos(t2))
)
wx, wy = int(wx), int(wy)
if wx < 0 or wx >= cmW or wy < 0 or wy >= cmH:
continue
self.states["collision_map"][i, wy, wx] = 1
# Update visitation maps
for i in range(self.nplanners):
mx, my = asnumpy(self.states["curr_map_position"][i]).tolist()
mx, my = int(mx), int(my)
self.states["visited_map"][i, my - 2: my + 3, mx - 2: mx + 3] = 1
# ===================== Compute rewards for previous action ===================
local_rewards = self._compute_local_rewards(ep_step, s).to(device)
# ====================== Global policy action selection =======================
SAMPLE_GLOBAL_GOAL_FLAG = ep_step % self.goal_interval == 0
# Sample global goal if needed
if SAMPLE_GLOBAL_GOAL_FLAG:
global_policy_inputs = self._create_global_policy_inputs(
global_map, visited_states, self.states["curr_map_position"]
)
# diff_policy baseline add
if hasattr(self.config.GLOBAL_POLICY, "difference_baseline_explorer"):
if self.config.GLOBAL_POLICY.difference_baseline_explorer:
update_dict = {"diff_scores": prev_state_estimates.pop("diff_scores")}
global_policy_inputs.update(update_dict)
(
global_value,
global_action,
global_action_log_probs,
_,
) = self.global_policy.act(global_policy_inputs, None, None, None)
# Convert action to location (row-major format)
G = self.global_policy.G
global_action_map_x = torch.fmod(
global_action.squeeze(1), G
).float() # (bs, )
global_action_map_y = (global_action.squeeze(1).float() / G).float() # (bs, )
# Convert to MxM map coordinates
global_action_map_x = global_action_map_x * M / G
global_action_map_y = global_action_map_y * M / G
if self.config.fixed_global_goal:
starting_point = (M - 1) // 2
global_action_map_x[0] = starting_point + (self.config.fixed_delta_x / s)
global_action_map_y[0] = starting_point - (self.config.fixed_delta_y / s)
global_action_map_xy = torch.stack(
[global_action_map_x, global_action_map_y], dim=1
)
# Set the goal (bs, 2) --- (x, y) in map image coordinates
self.states["curr_global_goals"] = global_action_map_xy
# Update the current map to prev_map_states in order to facilitate
# future reward computation.
self.states["prev_map_states"] = global_map.detach()
else:
global_policy_inputs = None
global_value = None
global_action = None
global_action_log_probs = None
# ======================= Local policy action selection =======================
# Initialize states at t=0
if ep_step == 0:
self.states["curr_local_goals"] = torch.zeros(self.nplanners, 2).to(device)
self.states["local_shortest_path_length"] = torch.zeros(
self.nplanners, 1
).to(device)
self.states["local_path_length"] = torch.zeros(self.nplanners, 1).to(device)
# Should local goals be sampled now?
if SAMPLE_GLOBAL_GOAL_FLAG:
if self.config.MAPPER.debug_log:
print('NEW LOCAL GOAL SAMPLED!')
# Condition 1: A new global goal was sampled
SAMPLE_LOCAL_GOAL_FLAGS = [1 for _ in range(self.nplanners)]
else:
if self.config.MAPPER.debug_log:
print('REACHING LOCAL GOAL:')
print('Distance to Goal: {}'.format(self.states["curr_dist2localgoal"]))
print('Goal Position: {}'.format(self.states["curr_local_goals"]))
print('Agent Position: {}'.format(map_xy))
# Condition 2 (a): The previous local goal was reached
prev_goal_reached = (
self.states["curr_dist2localgoal"] < self.goal_success_radius
)
# Condition 2 (b): The previous local goal is occupied.
goals = self.states["curr_local_goals"].long().to(device)
prev_gcells = global_map[
torch.arange(0, goals.shape[0]).long(), :, goals[:, 1], goals[:, 0]
]
prev_goal_occupied = (prev_gcells[:, 0] > self.config.thresh_obstacle) & (
prev_gcells[:, 1] > self.config.thresh_explored
)
SAMPLE_LOCAL_GOAL_FLAGS = asnumpy(
(prev_goal_reached | prev_goal_occupied).float()
).tolist()
# Execute planner and compute local goals
self._compute_plans_and_local_goals(
global_map, self.states["curr_map_position"], SAMPLE_LOCAL_GOAL_FLAGS
)
# Update state variables to account for new local goals
self.states["curr_dist2localgoal"] = self._compute_dist2localgoal(
global_map,
self.states["curr_map_position"],
self.states["curr_local_goals"],
)
# Sample action with local policy
local_masks = 1 - torch.Tensor(SAMPLE_LOCAL_GOAL_FLAGS).to(device).unsqueeze(1)
recurrent_hidden_states = prev_state_estimates["recurrent_hidden_states"]
relative_goals = self._compute_relative_local_goals(global_pose, M, s)
local_policy_inputs = {
"rgb_at_t": observations["rgb"],
"goal_at_t": relative_goals,
"t": ep_time,
}
outputs = self.local_policy.act(
local_policy_inputs,
recurrent_hidden_states,
None,
local_masks,
deterministic=deterministic,
)
(
local_value,
local_action,
local_action_log_probs,
recurrent_hidden_states,
) = outputs
# If imitation learning is used, also sample the ground-truth action to take
if "gt_global_map" in observations.keys():
gt_global_map = observations["gt_global_map"] # (bs, 2, M, M)
gt_global_pose = observations["pose_gt"] # (bs, 3)
relative_goals_aug = torch.cat(
[relative_goals, torch.zeros_like(relative_goals[:, 0:1])], dim=1
)
gt_goals = add_pose(gt_global_pose, relative_goals_aug) # (bs, 3)
gt_actions = self._compute_gt_local_action(
gt_global_map, gt_global_pose, gt_goals, M, s
) # (bs, 1)
gt_actions = gt_actions.to(device)
else:
gt_actions = torch.zeros_like(local_action)
# ============================== Create output dicts ==========================
state_estimates = {
"recurrent_hidden_states": recurrent_hidden_states,
"map_states": mapper_outputs["mt"],
"visited_states": visited_states,
"pose_estimates": global_pose,
}
local_policy_outputs = {
"values": local_value,
"actions": local_action,
"action_log_probs": local_action_log_probs,
"local_masks": local_masks,
"gt_actions": gt_actions,
}
global_policy_outputs = {
"values": global_value,
"actions": global_action,
"action_log_probs": global_action_log_probs,
}
rewards = {
"local_rewards": local_rewards,
}
mapper_outputs.update({'curr_map_position': curr_map_position})
return (
mapper_inputs,
local_policy_inputs,
global_policy_inputs,
mapper_outputs,
local_policy_outputs,
global_policy_outputs,
state_estimates,
rewards,
)
def _compute_plans(
self,
global_map,
agent_map_xy,
goal_map_xy,
sample_goal_flags,
cache_map=False,
crop_map_flag=True,
):
"""
global_map - (bs, 2, V, V) tensor
agent_map_xy - (bs, 2) agent's current position on the map
goal_map_xy - (bs, 2) goal's current position on the map
sample_goal_flags - list of zeros and ones should a new goal be sampled?
"""
# ==================== Process the map to get planner map =====================
s = self.map_scale
# Processed map has zeros for free-space and ones for obstacles
global_map_proc = self._process_maps(
global_map, goal_map_xy
) # (bs, M, M) tensor
# =================== Crop a local region around agent, goal ==================
if crop_map_flag:
# Determine crop size
abs_diff_xy = torch.abs(agent_map_xy - goal_map_xy)
S = int(torch.max(abs_diff_xy).item())
# Add a small buffer space around the agent location
buffer_size = int(3.0 / s) # 3 meters buffer space in total
S += buffer_size
old_center_xy = (agent_map_xy + goal_map_xy) / 2
# Crop a SxS region centered around old_center_xy
# Note: The out-of-bound regions will be zero-padded by default. In this case,
# since zeros correspond to free-space, that is not a problem.
cropped_global_map = crop_map(
global_map_proc.unsqueeze(1), old_center_xy, S
).squeeze(
1
) # (bs, S, S)
# Add zero padding to ensure the plans don't fail due to cropping
pad_size = 5
cropped_global_map = F.pad(
cropped_global_map, (pad_size, pad_size, pad_size, pad_size)
)
S += pad_size * 2
# Transform to new coordinates
new_center_xy = torch.ones_like(old_center_xy) * (S / 2)
new_agent_map_xy = agent_map_xy + (new_center_xy - old_center_xy)
new_goal_map_xy = goal_map_xy + (new_center_xy - old_center_xy)
else:
cropped_global_map = global_map_proc # (bs, M, M)
old_center_xy = (agent_map_xy + goal_map_xy) / 2
new_center_xy = old_center_xy
new_agent_map_xy = agent_map_xy
new_goal_map_xy = goal_map_xy
S = cropped_global_map.shape[1]
if cache_map:
self._cropped_global_map = cropped_global_map
# Clip points to ensure they are within map limits
new_agent_map_xy = torch.clamp(new_agent_map_xy, 0, S - 1)
new_goal_map_xy = torch.clamp(new_goal_map_xy, 0, S - 1)
# Convert to numpy
agent_map_xy_np = asnumpy(new_agent_map_xy).astype(np.int32)
goal_map_xy_np = asnumpy(new_goal_map_xy).astype(np.int32)
global_map_np = asnumpy(cropped_global_map)
# =================== Plan path from agent to goal positions ==================
plans = self.planner.plan(
global_map_np, agent_map_xy_np, goal_map_xy_np, sample_goal_flags
) # List of tuple of lists
# Convert plans back to original coordinates
final_plans = []
for i in range(len(plans)):
plan_x, plan_y = plans[i]
# Planning failure
if plan_x is None:
final_plans.append((plan_x, plan_y))
continue
offset_x = int((old_center_xy[i, 0] - new_center_xy[i, 0]).item())
offset_y = int((old_center_xy[i, 1] - new_center_xy[i, 1]).item())
final_plan_x, final_plan_y = [], []
for px, py in zip(plan_x, plan_y):
final_plan_x.append(px + offset_x)
final_plan_y.append(py + offset_y)
final_plans.append((final_plan_x, final_plan_y))
return final_plans
def get_observation(self, *args: Any, observations, episode, **kwargs: Any):
sim_depth = asnumpy(observations["depth"])
ego_wall_map_gt = self._get_depth_projection(sim_depth)
return ego_wall_map_gt
assert caps_norms.max() <= 1.001, 'capsules outputs should be bound by unit norm'
# correct capsule is enforced is not penalized if norm > m_plus,
# while incorrect ones are not penalized if norm < m_minus
loss_sig = torch.clamp(m_plus - caps_norms, 0) ** 2
loss_bkg = torch.clamp(caps_norms - m_minus, 0) ** 2
loss = target_one_hot * loss_sig + loss_lambda * (1.0 - target_one_hot) * loss_bkg
return reduce(loss, 'b cls -> b', 'sum')
for epoch in range(100):
for i, (images, labels) in enumerate(train_loader):
digit_capsules = encoder(images.to(device))
labels_one_hot = torch.nn.functional.one_hot(labels, num_classes=10).float().to(device)
loss = margin_loss(digit_capsules, labels_one_hot).mean()
reconstructed = decoder(digit_capsules * rearrange(labels_one_hot, 'b caps -> b caps 1'))
reconstruction_loss_mse = (images.to(device) - reconstructed).pow(2).mean()
loss += reconstruction_loss_mse * 10. # pick a weight
loss.backward()
optimizer.step()
optimizer.zero_grad()
accuracies = []
for images, labels in test_loader:
digit_capsules = encoder(images.to(device)).cpu()
# predicted capsule is capsule with largest norm
predicted_labels = digit_capsules.norm(dim=2).argmax(dim=1)
accuracies += asnumpy(predicted_labels == labels).tolist()
print(f'epoch {epoch} accuracy: {np.mean(accuracies)}')
def get_episode_map(env:habitat.RLEnv, mapper:Mapper, M:int, config:CN, device:torch.device):
"""Given the environment and the configuration, compute the global
top-down wall and seen area maps by sampling maps for individual locations
along a uniform grid in the environment, and registering them.
"""
# Initialize a global map for the episode
global_wall_map = torch.zeros(1, 2, M, M).to(device)
global_seen_map = torch.zeros(1, 2, M, M).to(device)
grid_size = config.TASK.GT_EGO_MAP.MAP_SCALE
coordinate_max = maps.COORDINATE_MAX
coordinate_min = maps.COORDINATE_MIN
resolution = (coordinate_max - coordinate_min) / grid_size
grid_resolution = (int(resolution), int(resolution))
top_down_map = maps.get_topdown_map(
env.habitat_env.sim, grid_resolution, 20000, draw_border=False,
)
map_w, map_h = top_down_map.shape
intervals = (max(int(0.5 / grid_size), 1), max(int(0.5 / grid_size), 1))
x_vals = np.arange(0, map_w, intervals[0], dtype=int)
y_vals = np.arange(0, map_h, intervals[1], dtype=int)
coors = np.stack(np.meshgrid(x_vals, y_vals), axis=2) # (H, W, 2)
coors = coors.reshape(-1, 2) # (H*W, 2)
map_vals = top_down_map[coors[:, 0], coors[:, 1]]
valid_coors = coors[map_vals > 0]
real_x_vals = coordinate_max - valid_coors[:, 0] * grid_size
real_z_vals = coordinate_min + valid_coors[:, 1] * grid_size
start_y = env.habitat_env.sim.get_agent_state().position[1]
for j in tqdm.tqdm(range(real_x_vals.shape[0]),desc='occupacy map', position=2):
for theta in np.arange(-np.pi, np.pi, np.pi / 3.0):
position = [
real_x_vals[j].item(),
start_y.item(),
real_z_vals[j].item(),
]
rotation = [
0.0,
np.sin(theta / 2).item(),
0.0,
np.cos(theta / 2).item(),
]
sim_obs = env.habitat_env.sim.get_observations_at(
position, rotation, keep_agent_at_new_pose=True,
)
episode = env.habitat_env.current_episode
obs = env.habitat_env.task.sensor_suite.get_observations(
observations=sim_obs, episode=episode, task=env.habitat_env.task
)
ego_map_gt = torch.Tensor(obs["ego_map_gt"]).to(device)
ego_map_gt = rearrange(ego_map_gt, "h w c -> () c h w")
ego_wall_map_gt = torch.Tensor(obs["ego_wall_map_gt"]).to(device)
ego_wall_map_gt = rearrange(ego_wall_map_gt, "h w c -> () c h w")
pose_gt = torch.Tensor(obs["pose_gt"]).unsqueeze(0).to(device)
global_seen_map = mapper.ext_register_map(
global_seen_map, ego_map_gt, pose_gt
)
global_wall_map = mapper.ext_register_map(
global_wall_map, ego_wall_map_gt, pose_gt
)
global_wall_map_np = asnumpy(rearrange(global_wall_map, "b c h w -> b h w c")[0])
global_seen_map_np = asnumpy(rearrange(global_seen_map, "b c h w -> b h w c")[0])
return global_seen_map_np, global_wall_map_np
def _eval(self):
start_time = time.time()
if self.config.MANUAL_COMMANDS:
init_time = None
manual_step_start_time = None
total_manual_time = 0.0
checkpoint_index = int(
(re.findall('\d+', self.config.EVAL_CKPT_PATH_DIR))[-1])
ckpt_dict = torch.load(self.config.EVAL_CKPT_PATH_DIR,
map_location="cpu")
print(
f'Number of steps of the ckpt: {ckpt_dict["extra_state"]["step"]}')
config = self._setup_config(ckpt_dict)
ppo_cfg = config.RL.PPO
ans_cfg = config.RL.ANS
self.mapper_rollouts = None
self._setup_actor_critic_agent(ppo_cfg, ans_cfg)
self.mapper_agent.load_state_dict(ckpt_dict["mapper_state_dict"])
if self.local_agent is not None:
self.local_agent.load_state_dict(ckpt_dict["local_state_dict"])
self.local_actor_critic = self.local_agent.actor_critic
else:
self.local_actor_critic = self.ans_net.local_policy
self.global_agent.load_state_dict(ckpt_dict["global_state_dict"])
self.mapper = self.mapper_agent.mapper
self.global_actor_critic = self.global_agent.actor_critic
# Set models to evaluation
self.mapper.eval()
self.local_actor_critic.eval()
self.global_actor_critic.eval()
M = ans_cfg.overall_map_size
V = ans_cfg.MAPPER.map_size
s = ans_cfg.MAPPER.map_scale
imH, imW = ans_cfg.image_scale_hw
num_steps = self.config.T_EXP
prev_action = torch.zeros(1, 1, device=self.device, dtype=torch.long)
masks = torch.zeros(1, 1, device=self.device)
try:
self.sim = make_sim('PyRobot-v1',
config=self.config.TASK_CONFIG.PYROBOT)
except (KeyboardInterrupt, SystemExit):
sys.exit()
pose = defaultdict()
self.sim._robot.camera.set_tilt(math.radians(self.config.CAMERA_TILT),
wait=True)
print(
f"\nStarting Camera State: {self.sim.get_agent_state()['camera']}")
print(f"Starting Agent State: {self.sim.get_agent_state()['base']}")
obs = [self.sim.reset()]
if self.config.SAVE_OBS_IMGS:
cv2.imwrite(f'obs/depth_dirty_s.jpg', obs[0]['depth'] * 255.0)
obs[0]['depth'][..., 0] = self._correct_depth(obs, -1)
if self.config.SAVE_OBS_IMGS:
cv2.imwrite(f'obs/rgb_s.jpg', obs[0]['rgb'][:, :, ::-1])
cv2.imwrite(f'depth_s.jpg', obs[0]['depth'] * 255.0)
starting_agent_state = self.sim.get_agent_state()
locobot2relative = CoordProjection(starting_agent_state['base'])
pose['base'] = locobot2relative(starting_agent_state['base'])
print(f"Starting Agent Pose: {pose['base']}\n")
batch = self._prepare_batch(obs, -1, device=self.device)
if ans_cfg.MAPPER.use_sensor_positioning:
batch['pose'] = pose['base'].to(self.device)
batch['pose'][0][1:] = -batch['pose'][0][1:]
prev_batch = batch
num_envs = self.config.NUM_PROCESSES
agent_poses_over_time = []
for i in range(num_envs):
agent_poses_over_time.append(
torch.tensor([(M - 1) / 2, (M - 1) / 2, 0]))
state_estimates = {
"pose_estimates":
torch.zeros(num_envs, 3).to(self.device),
"map_states":
torch.zeros(num_envs, 2, M, M).to(self.device),
"recurrent_hidden_states":
torch.zeros(1, num_envs,
ans_cfg.LOCAL_POLICY.hidden_size).to(self.device),
"visited_states":
torch.zeros(num_envs, 1, M, M).to(self.device),
}
ground_truth_states = {
"visible_occupancy": torch.zeros(num_envs, 2, M,
M).to(self.device),
"pose": torch.zeros(num_envs, 3).to(self.device),
"environment_layout": torch.zeros(num_envs, 2, M,
M).to(self.device)
}
# Reset ANS states
self.ans_net.reset()
# Frames for video creation
rgb_frames = []
if len(self.config.VIDEO_OPTION) > 0:
os.makedirs(self.config.VIDEO_DIR, exist_ok=True)
step_start_time = time.time()
for i in range(num_steps):
print(
f"\n\n---------------------------------------------------<<< STEP {i} >>>---------------------------------------------------"
)
ep_time = torch.zeros(num_envs, 1, device=self.device).fill_(i)
(
mapper_inputs,
local_policy_inputs,
global_policy_inputs,
mapper_outputs,
local_policy_outputs,
global_policy_outputs,
state_estimates,
intrinsic_rewards,
) = self.ans_net.act(
batch,
prev_batch,
state_estimates,
ep_time,
masks,
deterministic=True,
)
if self.config.SAVE_MAP_IMGS:
cv2.imwrite(
f'maps/test_map_{i - 1}.jpg',
self._round_map(state_estimates['map_states']) * 255)
action = local_policy_outputs["actions"][0][0]
distance2ggoal = torch.norm(
mapper_outputs['curr_map_position'] -
self.ans_net.states["curr_global_goals"],
dim=1) * s
print(f"Distance to Global Goal: {distance2ggoal}")
reached_flag = distance2ggoal < ans_cfg.goal_success_radius
if self.config.MANUAL_COMMANDS:
if init_time is None:
init_time = time.time() - start_time
total_manual_time = total_manual_time + init_time
if manual_step_start_time is not None:
manual_step_time = time.time() - manual_step_start_time
total_manual_time = total_manual_time + manual_step_time
action = torch.tensor(
int(input('Waiting input to start new action: ')))
manual_step_start_time = time.time()
if action.item() == 3:
reached_flag = True
prev_action.copy_(action)
if not reached_flag and action.item() != 3:
print(f'Doing Env Step [{self.ACT_2_NAME[action.item()]}]...')
action_command = self.ACT_2_COMMAND[action.item()]
obs = self._do_action(action_command)
if self.config.SAVE_OBS_IMGS:
cv2.imwrite(f'obs/depth_dirty_{i}.jpg',
obs[0]['depth'] * 255.0)
# Correcting invalid depth pixels
obs[0]['depth'][..., 0] = self._correct_depth(obs, i)
if self.config.SAVE_OBS_IMGS:
cv2.imwrite(f'obs/rgb_{i}.jpg', obs[0]['rgb'][:, :, ::-1])
cv2.imwrite(f'obs/depth_{i}.jpg', obs[0]['depth'] * 255.0)
agent_state = self.sim.get_agent_state()
prev_batch = batch
batch = self._prepare_batch(obs, i, device=self.device)
pose = defaultdict()
pose['base'] = locobot2relative(agent_state['base'])
if ans_cfg.MAPPER.use_sensor_positioning:
batch['pose'] = pose['base'].to(self.device)
batch['pose'][0][1:] = -batch['pose'][0][1:]
map_coords = convert_world2map(
batch['pose'], (M, M), ans_cfg.OCCUPANCY_ANTICIPATOR.
EGO_PROJECTION.map_scale).squeeze()
map_coords = torch.cat(
(map_coords, batch['pose'][0][-1].reshape(1)))
if self.config.COORD_DEBUG:
print('COORDINATES CHECK')
print(
f'Starting Agent State: {starting_agent_state["base"]}'
)
print(f'Current Agent State: {agent_state["base"]}')
print(
f'Current Sim Agent State: {self.sim.get_agent_state()["base"]}'
)
print(f'Current Global Coords: {batch["pose"]}')
print(f'Current Map Coords: {map_coords}')
agent_poses_over_time.append(map_coords)
step_time = time.time() - step_start_time
print(f"\nStep Time: {step_time}")
step_start_time = time.time()
# Create new frame of the video
if (len(self.config.VIDEO_OPTION) > 0):
frame = observations_to_image(
obs[0],
observation_size=300,
collision_flag=self.config.DRAW_COLLISIONS)
# Add ego_map_gt to frame
ego_map_gt_i = asnumpy(batch["ego_map_gt"][0]) # (2, H, W)
ego_map_gt_i = convert_gt2channel_to_gtrgb(ego_map_gt_i)
ego_map_gt_i = cv2.resize(ego_map_gt_i, (300, 300))
# frame = np.concatenate([frame], axis=1)
# Generate ANS specific visualizations
environment_layout = asnumpy(
ground_truth_states["environment_layout"][0]) # (2, H, W)
visible_occupancy = mapper_outputs["gt_mt"][0].cpu().numpy(
) # (2, H, W)
anticipated_occupancy = mapper_outputs["hat_mt"][0].cpu(
).numpy() # (2, H, W)
H = frame.shape[0]
visible_occupancy_vis = generate_topdown_allocentric_map(
environment_layout,
visible_occupancy,
agent_poses_over_time,
thresh_explored=ans_cfg.thresh_explored,
thresh_obstacle=ans_cfg.thresh_obstacle,
zoom=False)
visible_occupancy_vis = cv2.resize(visible_occupancy_vis,
(H, H))
anticipated_occupancy_vis = generate_topdown_allocentric_map(
environment_layout,
anticipated_occupancy,
agent_poses_over_time,
thresh_explored=ans_cfg.thresh_explored,
thresh_obstacle=ans_cfg.thresh_obstacle,
zoom=False)
anticipated_occupancy_vis = cv2.resize(
anticipated_occupancy_vis, (H, H))
anticipated_action_map = generate_topdown_allocentric_map(
environment_layout,
anticipated_occupancy,
agent_poses_over_time,
zoom=False,
thresh_explored=ans_cfg.thresh_explored,
thresh_obstacle=ans_cfg.thresh_obstacle,
)
global_goals = self.ans_net.states["curr_global_goals"]
local_goals = self.ans_net.states["curr_local_goals"]
if global_goals is not None:
cX = int(global_goals[0, 0].item())
cY = int(global_goals[0, 1].item())
anticipated_action_map = cv2.circle(
anticipated_action_map,
(cX, cY),
10,
(255, 0, 0),
-1,
)
if local_goals is not None:
cX = int(local_goals[0, 0].item())
cY = int(local_goals[0, 1].item())
anticipated_action_map = cv2.circle(
anticipated_action_map,
(cX, cY),
10,
(0, 255, 255),
-1,
)
anticipated_action_map = cv2.resize(anticipated_action_map,
(H, H))
maps_vis = np.concatenate(
[
visible_occupancy_vis, anticipated_occupancy_vis,
anticipated_action_map, ego_map_gt_i
],
axis=1,
)
if self.config.RL.ANS.overall_map_size == 2001 or self.config.RL.ANS.overall_map_size == 961:
if frame.shape[1] < maps_vis.shape[1]:
diff = maps_vis.shape[1] - frame.shape[1]
npad = ((0, 0), (diff // 2, diff // 2), (0, 0))
frame = np.pad(frame,
pad_width=npad,
mode='constant',
constant_values=0)
elif frame.shape[1] > maps_vis.shape[1]:
diff = frame.shape[1] - maps_vis.shape[1]
npad = ((0, 0), (diff // 2, diff // 2), (0, 0))
maps_vis = np.pad(maps_vis,
pad_width=npad,
mode='constant',
constant_values=0)
frame = np.concatenate([frame, maps_vis], axis=0)
rgb_frames.append(frame)
if self.config.SAVE_VIDEO_IMGS:
try:
os.mkdir("fig1")
except:
pass
print("Saved imgs for Fig. 1!")
cv2.imwrite(f'fig1/rgb_{step_start_time}.jpg',
obs[0]['rgb'][:, :, ::-1])
cv2.imwrite(f'fig1/depth_{step_start_time}.jpg',
obs[0]['depth'] * 255.0)
cv2.imwrite(f'fig1/aap_{step_start_time}.jpg',
anticipated_action_map[..., ::-1])
cv2.imwrite(f'fig1/em_{step_start_time}.jpg',
ego_map_gt_i[..., ::-1])
if self.config.DEBUG_VIDEO_FRAME:
cv2.imwrite('last_frame.jpg', frame)
if reached_flag:
for f in range(20):
rgb_frames.append(frame)
# Video creation
video_dict = {"t": start_time}
if (i + 1) % 10 == 0 or reached_flag:
generate_video(
video_option=self.config.VIDEO_OPTION,
video_dir=self.config.VIDEO_DIR,
images=rgb_frames,
episode_id=0,
checkpoint_idx=checkpoint_index,
metrics=video_dict,
tb_writer=TensorboardWriter('tb/locobot'),
)
if reached_flag:
if self.config.MANUAL_COMMANDS:
manual_step_time = time.time() - manual_step_start_time
total_manual_time = total_manual_time + manual_step_time
print(f"Manual elapsed time: {total_manual_time}")
print(f"Number of steps: {i + 1}")
print(f"Elapsed time: {time.time() - start_time}")
print(f"Final Distance to Goal: {distance2ggoal}")
if "bump" in obs[0]:
print(f"Collision: {obs[0]['bump']}")
print("Exiting...")
break
return
