def calculate_local_path_costs()

in habitat_baselines/slambased/path_planners.py [0:0]


    def calculate_local_path_costs(self, non_obstacle_cost_map=None):
        coords = self.coords
        h = coords.size(2)
        w = coords.size(3)
        obstacles_pd = F.pad(self.obstacles, (1, 1, 1, 1), "replicate")
        if non_obstacle_cost_map is None:
            learned_bias = torch.ones_like(self.obstacles).to(
                obstacles_pd.device
            )
        else:
            learned_bias = non_obstacle_cost_map.to(obstacles_pd.device)
        left_diff_sq = (
            self.gx_to_left(
                F.pad(coords[:, 1:2, :, :], (1, 1, 0, 0), "replicate")
            )
            ** 2
        )
        right_diff_sq = (
            self.gx_to_right(
                F.pad(coords[:, 1:2, :, :], (1, 1, 0, 0), "replicate")
            )
            ** 2
        )
        up_diff_sq = (
            self.gy_to_up(
                F.pad(coords[:, 0:1, :, :], (0, 0, 1, 1), "replicate")
            )
            ** 2
        )
        down_diff_sq = (
            self.gy_to_down(
                F.pad(coords[:, 0:1, :, :], (0, 0, 1, 1), "replicate")
            )
            ** 2
        )
        out = torch.cat(
            [
                # Order in from up to down, from left to right
                # hopefully same as in PyTorch
                torch.sqrt(left_diff_sq + up_diff_sq + self.eps)
                + self.ob_cost
                * torch.max(
                    obstacles_pd[:, :, 0:h, 0:w],
                    obstacles_pd[:, :, 1 : h + 1, 1 : w + 1],
                ),
                torch.sqrt(left_diff_sq + self.eps)
                + self.ob_cost
                * torch.max(
                    obstacles_pd[:, :, 0:h, 1 : w + 1],
                    obstacles_pd[:, :, 1 : h + 1, 1 : w + 1],
                ),
                torch.sqrt(left_diff_sq + down_diff_sq + self.eps)
                + self.ob_cost
                * torch.max(
                    obstacles_pd[:, :, 2 : h + 2, 0:w],
                    obstacles_pd[:, :, 1 : h + 1, 1 : w + 1],
                ),
                torch.sqrt(up_diff_sq + self.eps)
                + self.ob_cost
                * torch.max(
                    obstacles_pd[:, :, 0:h, 1 : w + 1],
                    obstacles_pd[:, :, 1 : h + 1, 1 : w + 1],
                ),
                0 * right_diff_sq
                + self.ob_cost
                * obstacles_pd[:, :, 1 : h + 1, 1 : w + 1],  # current center
                torch.sqrt(down_diff_sq + self.eps)
                + self.ob_cost
                * torch.max(
                    obstacles_pd[:, :, 2 : h + 2, 1 : w + 1],
                    obstacles_pd[:, :, 1 : h + 1, 1 : w + 1],
                ),
                torch.sqrt(right_diff_sq + up_diff_sq + self.eps)
                + self.ob_cost
                * torch.max(
                    obstacles_pd[:, :, 0:h, 2 : w + 2],
                    obstacles_pd[:, :, 1 : h + 1, 1 : w + 1],
                ),
                torch.sqrt(right_diff_sq + self.eps)
                + self.ob_cost
                * torch.max(
                    obstacles_pd[:, :, 1 : h + 1, 2 : w + 2],
                    obstacles_pd[:, :, 1 : h + 1, 1 : w + 1],
                ),
                torch.sqrt(right_diff_sq + down_diff_sq + self.eps)
                + self.ob_cost
                * torch.max(
                    obstacles_pd[:, :, 2 : h + 2, 2 : w + 2],
                    obstacles_pd[:, :, 1 : h + 1, 1 : w + 1],
                ),
            ],
            dim=1,
        )
        return out + torch.clamp(
            learned_bias.expand_as(out), min=0, max=self.ob_cost
        )