<file_sep>from blox import AttrDict
from blox.torch.modules import DummyModule
from gcp.evaluation.evaluation_matching import DTWEvalBinding, BalancedPrunedDTWBinding, BalancedEvalBinding


class TreeDenseRec(DummyModule):
    def __init__(self, hp, *_, decoder, **__):
        super().__init__()
        self._hp = hp
        self.eval_binding = None
        self.decoder = decoder

    def get_sample_with_len(self, i_ex, length, outputs, inputs, pruning_scheme, name=None):
        """Perform evaluation matching, return dense sequence of specified length."""
        if self.eval_binding is None:
            self.eval_binding = self._get_eval_binding(pruning_scheme)
        return self.eval_binding(outputs, inputs, length, i_ex, name)
    
    def get_all_samples_with_len(self, length, outputs, inputs, pruning_scheme, name=None):
        """Perform evaluation matching, return dense sequence of specified length."""
        if self.eval_binding is None:
            self.eval_binding = self._get_eval_binding(pruning_scheme)
            
        if hasattr(self.eval_binding, 'get_all_samples'):
            return self.eval_binding.get_all_samples(outputs, inputs, length, name)
        else:
            return [self.eval_binding(outputs, inputs, length, i_ex, name) for i_ex in range(outputs.end_ind.shape[0])]
 
    def _get_eval_binding(self, pruning_scheme):
        if pruning_scheme == 'dtw':
            return DTWEvalBinding(self._hp)
        if pruning_scheme == 'pruned_dtw':
            assert self._hp.matching_type == 'balanced'
            return BalancedPrunedDTWBinding(self._hp)
        if pruning_scheme == 'basic':
            assert self._hp.matching_type == 'balanced'
            return BalancedEvalBinding(self._hp)
        else:
            raise ValueError("Eval pruning scheme {} not currently supported!".format(pruning_scheme))
        
    def forward(self, tree, inputs):
        decoded_seq = self.decoder.decode_seq(inputs, tree.bf.e_g_prime)
        tree.set_attr_bf(**decoded_seq)
        return AttrDict()
<file_sep>import os
from contextlib import contextmanager

import torch
import torch.nn as nn
from blox.torch.layers import LayerBuilderParams
from blox import AttrDict
from blox.torch.modules import Updater
from blox.torch.models import base as bloxm
from tensorflow.contrib.training import HParams


class BaseModel(bloxm.BaseModel):
    def __init__(self, logger):
        super().__init__()
        self._hp = None
        self._logger = logger

    @contextmanager
    def val_mode(self):
        """Sets validation parameters. To be used like: with model.val_mode(): ...<do something>..."""
        raise NotImplementedError("Need to implement val_mode context manager in subclass!")

    def step(self):
        self.call_children('step', Updater)

    def override_defaults(self, policyparams):
        for name, value in policyparams.items():
            print('overriding param {} to value {}'.format(name, value))
            if value == getattr(self._hp, name):
                raise ValueError("attribute is {} is identical to default value!!".format(name))
            self._hp.set_hparam(name, value)

    def _default_hparams(self):
        # Data Dimensions
        default_dict = AttrDict({
            'batch_size': -1,
            'max_seq_len': -1,
            'n_actions': -1,
            'state_dim': -1,
            'img_sz': 32,  # image resolution
            'input_nc': 3,  # number of input feature maps
            'n_conv_layers': None,  # Number of conv layers. Can be of format 'n-<int>' for any int for relative spec
        })
        
        # Network params
        default_dict.update({
            'use_convs': True,
            'use_batchnorm': True,  # TODO deprecate
            'normalization': 'batch',
            'predictor_normalization': 'group',
        })

        # Misc params
        default_dict.update({
            'filter_repeated_tail': False,      # whether to remove repeated states from the dataset
            'rep_tail': False,
            'dataset_class': None,
            'standardize': None,
            'split': None,
            'subsampler': None,
            'subsample_args': None,
            'checkpt_path': None,
        })
        
        # add new params to parent params
        parent_params = HParams()
        for k in default_dict.keys():
            parent_params.add_hparam(k, default_dict[k])
        return parent_params

    def postprocess_params(self):
        if not self._hp.use_convs:
            # self._hp.input_nc = self._hp.img_sz ** 2 * self._hp.input_nc
            self._hp.input_nc = self._hp.state_dim
        self._hp.add_hparam('builder', LayerBuilderParams(
            self._hp.use_convs, self._hp.use_batchnorm, self._hp.normalization, self._hp.predictor_normalization))
        self._hp.add_hparam('fc_builder', LayerBuilderParams(
            False, self._hp.use_batchnorm, self._hp.normalization, self._hp.predictor_normalization))

    def build_network(self):
        raise NotImplementedError("Need to implement this function in the subclass!")

    def forward(self, inputs):
        raise NotImplementedError("Need to implement this function in the subclass!")

    def loss(self, outputs, inputs):
        raise NotImplementedError("Need to implement this function in the subclass!")

    def log_outputs(self, outputs, inputs, losses, step, log_images, phase):
        # Log generally useful outputs
        self._log_losses(losses, step, log_images, phase)

        if phase == 'train':
            self.log_gradients(step, phase)
            
        for module in self.modules():
            if hasattr(module, '_log_outputs'):
                module._log_outputs(outputs, inputs, losses, step, log_images, phase, self._logger)

            if hasattr(module, 'log_outputs_stateful'):
                module.log_outputs_stateful(step, log_images, phase, self._logger)
            
    def _log_losses(self, losses, step, log_images, phase):
        for name, loss in losses.items():
            self._logger.log_scalar(loss.value, name + '_loss', step, phase)
            if 'breakdown' in loss and log_images:
                self._logger.log_graph(loss.breakdown, name + '_breakdown', step, phase)

    def _load_weights(self, weight_loading_info):
        """
        Loads weights of submodels from defined checkpoints + scopes.
        :param weight_loading_info: list of tuples: [(model_handle, scope, checkpoint_path)]
        """

        def get_filtered_weight_dict(checkpoint_path, scope):
            if os.path.isfile(checkpoint_path):
                checkpoint = torch.load(checkpoint_path, map_location=self._hp.device)
                filtered_state_dict = {}
                remove_key_length = len(scope) + 1      # need to remove scope from checkpoint key
                for key, item in checkpoint['state_dict'].items():
                    if key.startswith(scope):
                        filtered_state_dict[key[remove_key_length:]] = item
                if not filtered_state_dict:
                    raise ValueError("No variable with scope '{}' found in checkpoint '{}'!".format(scope, checkpoint_path))
                return filtered_state_dict
            else:
                raise ValueError("Cannot find checkpoint file '{}' for loading '{}'.".format(checkpoint_path, scope))

        print("")
        for loading_op in weight_loading_info:
            print(("=> loading '{}' from checkpoint '{}'".format(loading_op[1], loading_op[2])))
            filtered_weight_dict = get_filtered_weight_dict(checkpoint_path=loading_op[2],
                                                            scope=loading_op[1])
            loading_op[0].load_state_dict(filtered_weight_dict)
            print(("=> loaded '{}' from checkpoint '{}'".format(loading_op[1], loading_op[2])))
        print("")

    def log_gradients(self, step, phase):
        grad_norms = list([torch.norm(p.grad.data) for p in self.parameters() if p.grad is not None])
        if len(grad_norms) == 0:
            return
        grad_norms = torch.stack(grad_norms)

        self._logger.log_scalar(grad_norms.mean(), 'gradients/mean_norm', step, phase)
        self._logger.log_scalar(grad_norms.max(), 'gradients/max_norm', step, phase)
    
    def get_total_loss(self, inputs, losses):
        # compute total loss
        ## filtering is important when some losses are nan
        ## the unsqueeze is important when some of the weights or losses are 1-dim tensors.
        # TODO use the function from blox
        total_loss = torch.stack([loss[1].value[None] * loss[1].weight for loss in
                                  filter(lambda x: x[1].weight > 0, losses.items())]).sum()
        if torch.isnan(total_loss).any():
            import pdb; pdb.set_trace()
        return AttrDict(value=total_loss)
<file_sep>from math import floor

import numpy as np
import torch
from blox.tensor import ndim
from blox.torch.ops import batchwise_assign, batchwise_index
from blox.utils import timing
from torch import nn


def fast_gak(C, transition, begin_inds=0):
    """
    Computes the global alignment kernel (Cuturi'07, A kernel for time series based on global alignments).
    This version parallelizes the computation across the diagonal
    This version is able to process batched sequences of variable length by using begin_inds

    :param C: the cost matrix
    :return: the kernel value and the matrix of intermediate values (kernel values for smaller sequences)
    """
    """
    Computes Dynamic Time Warping (DTW) of two sequences.
    This version multiplies the cost instead of adding it. This can be used if the costs represent exponential values
    exp_dtw = exp(basic_dtw(log(C)))

    :param C: the cost matrix ... x n_x x n_y
    :param aggr: the function used for aggregating candidates.
    :param transition: if 'nohor', horizontal transitions are excluded, i.e. one x only matches one y
    (but not the other way around)
    :return: the minimum distance and the accumulated cost matrix
    """
    r, c = C.shape[-2:]
    D = torch.full_like(C, -np.inf)
    batchwise_assign(D[:, 0], begin_inds, batchwise_index(C[:, 0], begin_inds))
    
    assert transition == 'nohor'
    assert r >= c
    
    # impossible = (1 - torch.ones_like(D).tril()).byte()
    # impossible += (1 - torch.ones_like(D).triu(-diff)).byte()
    
    # Iterate over diagonals
    for i in range(1, r + c):
        ids = torch.arange(i + 1).flip(0)
        jds = torch.arange(i + 1)
        # exclude boundary violations
        ids = ids[max(0, i - r + 1):c]
        jds = jds[max(0, i - r + 1):c]
        
        # ToDo clean this up?
        ids = ids.flip(0)
        jds = jds.flip(0)
        
        skip = D[..., ids - 1, jds]
        step = D[..., ids - 1, jds - 1]
        
        # # Allow horizontal transitions where the y sequence ended
        # # jds[-1] is the column index of the lowest element on the diagonal
        # repeat = torch.full_like(skip, -np.inf)
        # allow_mask = jds[-1] > end_inds
        # repeat[allow_mask, -1] = D[allow_mask, ids[-1], jds[-1] - 1]
        
        # Recursion
        add = torch.logsumexp(ndim.stack([skip, step], -1), -1)
        new_cost = C[..., ids, jds] + add
        mask = D[..., ids, jds] != -np.inf
        new_cost[mask] = D[..., ids, jds][mask]  # If D was specified, leave as is (for begin_inds)
        
        # TODO change to simple array indexing once the backward pass for it is implemented
        mask = ids2mask(D, ids, jds)
        if hasattr(mask, 'bool'):
            mask = mask.bool()
        D = D.masked_scatter(mask, (new_cost).flatten())
    return D


def ids2mask(mat, ids, jds):
    mask = torch.zeros_like(mat, dtype=torch.uint8)
    mask[..., ids, jds] = 1
    return mask


def soft_dtw(C, end_inds=None):
    """
    Computes the expected edge frequencies. See https://www.overleaf.com/read/jksjyppbrdgn
    
    :param C: the cost matrix, n_x x n_y
    :param end_inds: the end indices for the sequences y. The sequences y will be assumed to have the length as per the
    end index and the remainer of the frames will not be matched
    :return: the matrix with expected edge frequencies
    """
    C = (-C).double()
    batch, r, c = C.shape
    if end_inds is None:
        end_inds = torch.full([batch], c - 1, dtype=torch.long)
    
    # mask = torch.zeros_like(C)
    # if end_inds is not None:
    #     for i, ind in enumerate(end_inds):
    #         mask[i, :-1, end_inds[i]:] = 1
    #     C[mask.byte()] = -np.inf

    # Compute forward-backward
    comb_C = torch.cat([C, ndim.flip(C, [-1, -2])], 0)
    # The backward begins with end indices, not (-1,-1)
    comb_begin_inds = torch.cat([torch.zeros_like(end_inds), c - end_inds - 1], 0)

    accum = fast_gak(comb_C, transition='nohor', begin_inds=comb_begin_inds)
    
    forward = accum[:batch]
    backward = ndim.flip(accum[batch:], [-1, -2])

    # Compute expected matrix
    z = batchwise_index(forward[:, -1], end_inds)[:, None, None]
    e = forward + backward - C
    e[C == -np.inf] = -np.inf
    w = (e - z).exp()
    
    if not equal(w.sum(2).max(), 1, eps=1e-2):
        print('warning: dtw is not stable with these cost values')
        import pdb; pdb.set_trace()

    return w.float()


def equal(n, m, eps):
    return m - eps < n < m + eps
<file_sep>import os

import matplotlib.pyplot as plt
import numpy as np
import torch
import torchvision
from gcp.prediction.utils.visualization import plot_gt_matching_overview, plot_val_tree, plot_pruned_seqs,\
    plot_balanced_tree, make_gif, plot_graph, param, fig2img, plot_inverse_model_actions
from tensorboardX import SummaryWriter

from blox.tensor.ops import broadcast_final, find_tensor
from gcp.prediction.utils.utils import get_pad_mask


class Logger:
    def __init__(self, log_dir, n_logged_samples=10, summary_writer=None, fps=4):
        self._log_dir = log_dir
        self._n_logged_samples = n_logged_samples
        self.fps = fps
        if summary_writer is not None:
            self._summ_writer = summary_writer
        else:
            self._summ_writer = SummaryWriter(log_dir, max_queue=1, flush_secs=1)

    def _loop_batch(self, fn, name, val, *argv, **kwargs):
        """Loops the logging function n times."""
        for log_idx in range(min(self._n_logged_samples, len(val))):
            name_i = os.path.join(name, "_%d" % log_idx)
            fn(name_i, val[log_idx], *argv, **kwargs)
     
    @staticmethod
    def _check_size(val, size):
        if isinstance(val, torch.Tensor) or isinstance(val, np.ndarray):
            assert len(val.shape) == size, "Size of tensor does not fit required size, {} vs {}".format(len(val.shape), size)
        elif isinstance(val, list):
            assert len(val[0].shape) == size-1, "Size of list element does not fit required size, {} vs {}".format(len(val[0].shape), size-1)
        else:
            raise NotImplementedError("Input type {} not supported for dimensionality check!".format(type(val)))
        # if (val[0].numel() > 1e9):
        #     print("Logging very large image with size {}px.".format(max(val[0].shape[1], val[0].shape[2])))
            #raise ValueError("This might be a bit too much")

    def log_scalar(self, scalar, name, step, phase):
        self._summ_writer.add_scalar('{}_{}'.format(name, phase), scalar, step)

    def log_scalars(self, scalar_dict, group_name, step, phase):
        """Will log all scalars in the same plot."""
        self._summ_writer.add_scalars('{}_{}'.format(group_name, phase), scalar_dict, step)

    def log_images(self, image, name, step, phase):
        self._check_size(image, 4)   # [N, C, H, W]
        if image[0].numel() > 3e9:
            print('skipping logging a giant image with {}px'.format(image[0].numel()))
            return
        
        self._loop_batch(
            self._summ_writer.add_image, '{}_{}'.format(name, phase), image, step)

    def log_video(self, video_frames, name, step, phase):
        assert len(video_frames.shape) == 4, "Need [T, C, H, W] input tensor for single video logging!"
        if not isinstance(video_frames, torch.Tensor): video_frames = torch.tensor(video_frames)
        #video_frames = torch.transpose(video_frames, 0, 1)  # tbX requires [C, T, H, W] <- not in tbX >= 1.6
        video_frames = video_frames.unsqueeze(0)    # add an extra dimension to get grid of size 1
        self._summ_writer.add_video('{}_{}'.format(name, phase), video_frames, step, fps=self.fps)

    def log_videos(self, video_frames, name, step, phase):
        assert len(video_frames.shape) == 5, "Need [N, T, C, H, W] input tensor for video logging!"
        video_frames = video_frames.unsqueeze(1)    # add an extra dimension after batch to get grid of size 1
        self._loop_batch(self._summ_writer.add_video, '{}_{}'.format(name, phase), video_frames, step, fps=self.fps)

    def log_image_grid(self, images, name, step, phase, nrow=8):
        assert len(images.shape) == 4, "Image grid logging requires input shape [batch, C, H, W]!"
        img_grid = torchvision.utils.make_grid(images, nrow=nrow)
        if img_grid.min() < 0:
            print("warning, image not rescaled!")
            img_grid = (img_grid + 1) / 2
        
        self._summ_writer.add_image('{}_{}'.format(name, phase), img_grid, step)

    def log_video_grid(self, video_frames, name, step, phase, fps=3):
        assert len(video_frames.shape) == 5, "Need [N, T, C, H, W] input tensor for video logging!"
        self._summ_writer.add_video('{}_{}'.format(name, phase), video_frames, step, fps=fps)

    def log_figures(self, figure, name, step, phase):
        """figure: matplotlib.pyplot figure handle"""
        assert figure.shape[0] > 0, "Figure logging requires input shape [batch x figures]!"
        self._loop_batch(self._summ_writer.add_figure, '{}_{}'.format(name, phase), figure, step)

    def log_figure(self, figure, name, step, phase):
        """figure: matplotlib.pyplot figure handle"""
        self._summ_writer.add_figure('{}_{}'.format(name, phase), figure, step)
        
    def log_graph(self, array, name, step, phase):
        """figure: matplotlib.pyplot figure handle"""
        im = plot_graph(array)
        self._summ_writer.add_image('{}_{}'.format(name, phase), im.transpose(2, 0, 1), step)

    def dump_scalars(self, log_path=None):
        log_path = os.path.join(self._log_dir, "scalar_data.json") if log_path is None else log_path
        self._summ_writer.export_scalars_to_json(log_path)


class HierarchyLogger(Logger):
    def __init__(self, log_dir, hp, max_seq_len, n_logged_samples=3, summary_writer=None, fps=4):
        Logger.__init__(self, log_dir, n_logged_samples, summary_writer, fps)
        self._gamma_width = 4       # width of the gamma visualization in px
        self._hp = hp
        self.max_seq_len = max_seq_len
        self.dummy_env = None

    def render(self, tensor):
        return tensor

    def log_hierarchy_image(self, model_output, inputs, name, step, phase):
        """Builds an image depicting the predicted hierarchy."""
        reference = find_tensor(inputs)
        batch_size, channels, im_height, _ = inputs.I_0_image.shape
        N = self.max_seq_len
        assert batch_size >= self._n_logged_samples
        depth = model_output.tree.depth
        
        if 'gamma' in model_output.tree.subgoals.keys():
            gamma_width = self._gamma_width
        else:
            gamma_width = 0
        
        level_height = im_height + gamma_width
        image = 0.7 * torch.ones((self._n_logged_samples, channels, depth * level_height, N * im_height),
                                 dtype=torch.float32, device=reference.device)
        for level in range(depth):
            # initialize gamma "strips" to 0
            image[:, :, level * level_height : level * level_height + gamma_width] = 0.0
    
        def tile_gamma(gamma):
            return gamma[:, None].repeat(1, im_height).view(-1)

        # TODO vectorize the batch computation
        for segment in model_output.tree.depth_first_iter():
            subgoal = segment.subgoal
            if not subgoal.keys(): break
            level = depth - segment.depth
            # visualize gamma on this depth level
            for batch_idx in range(self._n_logged_samples):
                if 'done_mask' in segment and (segment.done_mask is None or not segment.done_mask[batch_idx]):
                    time_idx = int(subgoal.ind[batch_idx])
                    if gamma_width != 0:
                        image[batch_idx, :, level * level_height:level * level_height + gamma_width] += \
                            tile_gamma(subgoal.gamma[batch_idx])

                    if 'images' in subgoal.keys():
                        image[batch_idx,
                              :,
                              level * level_height + gamma_width:(level+1) * level_height,
                              time_idx * im_height:(time_idx + 1) * im_height] = self.render(subgoal.images[batch_idx])
                
        def flatten_seq(seq):
            return seq[:self._n_logged_samples].transpose(1, 2).transpose(2, 3).\
                reshape(self._n_logged_samples, channels, im_height, N * im_height)
            
        def mask_extra(seq):
            return seq * broadcast_final(get_pad_mask(model_output.end_ind, N), seq)
        
        if 'traj_seq' in inputs.keys():
            input_image = flatten_seq(inputs.traj_seq_images)
        else:
            input_image = torch.zeros_like(image[:,:,:level_height])
            input_image[:,:,:,:im_height] = inputs.I_0_image[:self._n_logged_samples]
            input_image[:,:,:,-im_height:] = inputs.I_g_image[:self._n_logged_samples]
        image = torch.cat([input_image, image], dim=2)
            
        if 'dense_rec' in model_output.keys() and 'images' in model_output.dense_rec.keys()\
                and not 'p_n_hat' in model_output.tree.subgoals:     # don't log the dense rec here if it is pruned
            dense_images = mask_extra(self.render(model_output.dense_rec.images))
            image = torch.cat([image, flatten_seq(dense_images)], dim=2)

        if 'soft_matched_estimates' in model_output.keys():
            soft_estimates = mask_extra(self.render(model_output.soft_matched_estimates))
            image = torch.cat([image, flatten_seq(soft_estimates)], dim=2)

        image = (image + 1)/2   # rescale back to 0-1 range
        self.log_images(image, name, step, phase)
        return image

    def _log_plot_img(self, img, name, step, phase):
        if not isinstance(img, torch.Tensor): img = torch.tensor(img)
        img = img.float().permute(0, 3, 1, 2)
        self.log_images(img, name, step, phase)

    def log_gif(self, imgs, name, step, phase):
        if isinstance(imgs, list): imgs = np.concatenate(imgs)
        if not isinstance(imgs, torch.Tensor): imgs = torch.tensor(imgs)
        imgs = imgs.float().permute(0, 3, 1, 2)
        self.log_video(imgs, name + "_gif", step, phase)

    def log_gt_match_overview(self, model_output, inputs, name, step, phase):
        self._log_plot_img(plot_gt_matching_overview(model_output, inputs), name, step, phase)

    def log_attention_overview(self, model_output, inputs, name, step, phase):
        self._log_plot_img(plot_gt_matching_overview(model_output, inputs, plot_attr='gamma'), name, step, phase)

    def log_val_tree(self, model_output, inputs, name, step, phase):
        self._log_plot_img(plot_val_tree(model_output, inputs), name, step, phase)

    def log_balanced_tree(self, model_output, element, name, step, phase):
        with param(n_logged_samples=self._n_logged_samples):
            im = plot_balanced_tree(model_output, element).transpose((0, 2, 3, 1))
        self._log_plot_img(im, name, step, phase)  # show from first, last and predicted

    def log_pruned_pred(self, model_output, inputs, name, step, phase):
        im, seq = plot_pruned_seqs(model_output, inputs)
        self._log_plot_img(im, name, step, phase)
        self.log_gif(seq, name, step, phase)

    def log_dense_gif(self, model_output, inputs, name, step, phase):
        """Logs the dense reconstruction """
    
        rows = [model_output.dense_rec.images]
        if phase == 'train':
            rows = [inputs.traj_seq] + rows
            
        self.log_rows_gif(rows, name, step, phase)
        # self.log_loss_gif(model_output.dense_rec.images, inputs.traj_seq, name, step, phase)

    def log_rows_gif(self, rows, name, step, phase):
        """ Logs a gif with several rows
        
        :param rows: a list of tensors batch x time x channel x height x width
        :param name:
        :param step:
        :param phase:
        :return:
        """
        im = make_gif(rows)
        self.log_video(im, name, step, phase)

    def log_loss_gif(self, estimates, targets, name, step, phase):
        """Logs gifs showing a target and a ground truth sequence side by side,
         shows ground truth sequence at training time too."""
        
        estimate_imgs = self.render((estimates + 1) / 2)
        target_imgs = self.render((targets + 1) / 2)
        if phase == "train":
            # concat target images
            seq_shape = target_imgs.shape
            padding = torch.zeros(seq_shape[:3] + (2,) + seq_shape[4:], dtype=target_imgs.dtype,
                                  device=target_imgs.device)
            plot_imgs = torch.cat([target_imgs, padding, estimate_imgs], dim=3)
        else:
            plot_imgs = estimate_imgs
        plot_imgs = plot_imgs[:5]
        batch, time, channels, height, width = plot_imgs.shape
        plot_imgs = plot_imgs.permute(1, 2, 3, 0, 4).reshape(time, channels, height, width * batch)
        self.log_video(plot_imgs, name, step, phase)

    def log_single_topdown_traj(self, traj, name, step, phase):
        from gcp.planning.infra.envs.miniworld_env.multiroom3d.multiroom3d_env import Multiroom3dEnv
        if self.dummy_env is None:
            self.dummy_env = Multiroom3dEnv({'n_rooms': self._hp.n_rooms}, no_env=True)
        im = self.dummy_env.render_top_down(traj)
        self._log_plot_img(im[None], name, step, phase)

    def log_multiple_topdown_trajs(self, trajs, name, step, phase):
        from gcp.planning.infra.envs.miniworld_env.multiroom3d.multiroom3d_env import Multiroom3dEnv
        if self.dummy_env is None:
            self.dummy_env = Multiroom3dEnv({'n_rooms': self._hp.n_rooms}, no_env=True)
        im = self.dummy_env.render_top_down(trajs[0])
        for traj in trajs[1:]:
            im = self.dummy_env.render_top_down(traj, background=im)
        self._log_plot_img(im[None], name, step, phase)

    def log_dataset_specific_trajectory(
          self, model_output, inputs, name, step, phase, dataset, predictions=None, end_inds=None):
        """ Logs trajectory according to a protocol defined by the dataset class """
        # TODO any way to make logger class have access to the dataset directly?
        if not hasattr(dataset, 'render_trajectory'):
            return
    
        if predictions is None:
            predictions, end_inds = self.get_predictions_from_output(inputs, model_output)
    
        im = dataset.render_trajectory(model_output, inputs, predictions, end_inds, self._n_logged_samples)
        
        if 'image' in im:
            self._log_plot_img(im['image'], name, step, phase)
        if 'video' in im:
            self.log_video(im['video'], name, step, phase)
        if 'videos' in im:
            for i, video in enumerate(im['videos']):
                self.log_video(video, name + "_{}".format(i), step, phase)

    @staticmethod
    def get_predictions_from_output(inputs, model_output):
        if 'pruned_prediction' in model_output:
            predictions = model_output.pruned_prediction
            end_inds = model_output.end_ind
        elif 'images' in model_output.dense_rec:
            predictions = model_output.dense_rec.images
            end_inds = inputs.end_ind
        elif 'soft_matched_estimates' in model_output:
            predictions = model_output.soft_matched_estimates
            end_inds = inputs.end_ind
        else:
            predictions = model_output.tree.df.images
            end_inds = np.ones((predictions.shape[0],), dtype=int) * predictions.shape[1]   # use full length
        return predictions, end_inds

    def log_states_2d(self, model_output, inputs, name, step, phase):
        """Logs 2D plot of first 2 state dimensions."""
        predictions, end_inds = self.get_predictions_from_output(inputs, model_output)

        imgs = []
        for i, end_ind_inp, end_ind in zip(range(self._n_logged_samples), inputs.end_ind, end_inds):
            fig = plt.figure()
            gt = inputs.traj_seq_states[i, :end_ind_inp].data.cpu().numpy()
            plt.plot(gt[:, 0], gt[:, 1], 'g')
            pred = predictions[i][:end_ind].data.cpu().numpy()
            plt.plot(pred[:, 0], pred[:, 1], 'r')
            imgs.append(fig2img(fig))
        self._log_plot_img(np.stack(imgs), name, step, phase)

    def log_pred_actions(self, model_output, inputs, name, step, phase):
        self._log_plot_img(plot_inverse_model_actions(model_output, inputs), name, step, phase)


if __name__ == "__main__":
    logger = Logger(log_dir="./summaries")
    for step in range(10):
        print("Running step %d" % step)
        dummy_data = torch.rand([32, 10, 3, 64, 64])
        logger.log_scalar(dummy_data[0, 0, 0, 0, 0], name="scalar", step=step, phase="train")
        logger.log_scalars({
            'test1': dummy_data[0, 0, 0, 0, 0],
            'test2': dummy_data[0, 0, 0, 0, 1],
            'test3': dummy_data[0, 0, 0, 0, 2]
        }, group_name="scalar_group", step=step, phase="train")
        logger.log_images(dummy_data[:, 0], name="image", step=step, phase="train")
        logger.log_video(dummy_data, name="video", step=step, phase="train")
        logger.log_video_grid(dummy_data, name="video_grid", step=step, phase="train")
        fig = plt.figure()
        plt.plot(dummy_data.data.numpy()[:, 0, 0, 0, 0])
        logger.log_figures(np.asarray([fig for _ in range(10)]), name="figure", step=step, phase="train")
    logger.dump_scalars()
    print("Done!")


<file_sep>import glob
import imp
import os
import random

import h5py
import numpy as np
import torch.utils.data as data
from torchvision.transforms import Resize

from blox import AttrDict
from blox.basic_types import map_dict
from blox.torch.training import RepeatedDataLoader
from blox.vis import resize_video
from gcp.prediction import global_params


def set_seed(seed):
    random.seed(seed)
    np.random.seed(seed + 2)


class BaseVideoDataset(data.Dataset):
    def __init__(self, data_dir, mpar, data_conf, phase, shuffle=True, dataset_size=-1):
        """
        :param data_dir: path to data directory
        :param mpar: model parameters used to determine output resolution etc
        :param data_conf: dataset config
        :param phase: string indicating whether 'train'/'val'/'test'
        :param shuffle: whether to shuffle within batch, set to False for computing metrics
        :param dataset_size: (optional) if not full dataset should be used, specifies number of used sequences
        """

        self.phase = phase
        self.data_dir = data_dir
        self.spec = data_conf.dataset_spec
        self.data_conf = data_conf
        self.dataset_size = dataset_size
        
        self.shuffle = shuffle and phase == 'train'
        self.use_states = not mpar.use_convs
        self.img_sz = mpar.img_sz
        self.device = mpar.device
        
        if shuffle:
            self.n_worker = 4
        else:
            self.n_worker = 1
        if global_params.debug:
            self.n_worker = 0

        self.filenames = None
        self.states_mean = None
        self.states_std = None

    def process_data_dict(self, data_dict):
        if 'images' in data_dict:
            data_dict.traj_seq_images = data_dict.pop('images')
        if 'states' in data_dict:
            data_dict.traj_seq_states = data_dict.pop('states')
            if self.states_mean is not None:
                data_dict.traj_seq_states = self.standardize(data_dict.traj_seq_states)
                data_dict.traj_seq_states_mean = self.states_mean
                data_dict.traj_seq_states_std = self.states_std

        if 'traj_seq_images' in data_dict and len(data_dict.traj_seq_images.shape) > 1:    # some datasets don't have images
            data_dict.traj_seq_images = self.preprocess_images(data_dict.traj_seq_images)
        data_dict.traj_seq = data_dict.traj_seq_states if self.use_states else data_dict.traj_seq_images

        if 'start_ind' not in data_dict:
            data_dict.start_ind = 0
        if 'end_ind' not in data_dict:
            data_dict.end_ind = self.spec['max_seq_len'] - 1
        if 'pad_mask' not in data_dict:
            data_dict.pad_mask = np.ones(self.spec['max_seq_len'], dtype=np.float32)
            
        data_dict.I_0 = data_dict.traj_seq[0]
        data_dict.I_g = data_dict.traj_seq[data_dict.end_ind]
        if 'traj_seq_images' in data_dict:
            data_dict.I_0_image = data_dict.traj_seq_images[0]
            data_dict.I_g_image = data_dict.traj_seq_images[data_dict.end_ind]

    def get_data_loader(self, batch_size, n_repeat):
        print('len {} dataset {}'.format(self.phase, len(self)))
        assert self.device in ['cuda', 'cpu']  # Otherwise the logic below is wrong
        return RepeatedDataLoader(self, batch_size=batch_size, shuffle=self.shuffle, num_workers=self.n_worker,
                                  drop_last=True, n_repeat=n_repeat, pin_memory=self.device == 'cuda',
                                  worker_init_fn=lambda x: np.random.seed(np.random.randint(65536) + x))

    def preprocess_images(self, images):
        return images

    @staticmethod
    def visualize(*args, **kwargs):
        pass

    def standardize(self, states):
        return (states - self.states_mean)/(1e-6 + self.states_std)

    @staticmethod
    def get_dataset_spec(data_dir):
        return imp.load_source('dataset_spec', os.path.join(data_dir, 'dataset_spec.py')).dataset_spec


class VarLenVideoDataset(BaseVideoDataset):
    """Variable length video dataset"""
    def __init__(self, data_dir, mpar, data_conf, phase, shuffle=True, dataset_size=-1):
        """
        :param data_dir: path to data directory
        :param mpar: model parameters used to determine output resolution etc
        :param data_conf: dataset config
        :param phase: string indicating whether 'train'/'val'/'test'
        :param shuffle: whether to shuffle within batch, set to False for computing metrics
        :param dataset_size: (optional) if not full dataset should be used, specifies number of used sequences
        """
        super().__init__(data_dir, mpar, data_conf, phase, shuffle, dataset_size)

        print('loading files from', self.data_dir)
        self.filenames = self._get_filenames()
        self.traj_per_file = self.get_traj_per_file(self.filenames[0])

        self.randomize_length = mpar.randomize_length
        self.randomize_start = mpar.randomize_start
        self.transform = Resize([mpar.img_sz, mpar.img_sz])
        self.flatten_im = False

        if 'states_mean' in self.spec:
            self.states_mean = self.spec['states_mean']
            self.states_std = self.spec['states_std']

    def _get_filenames(self):
        raise NotImplementedError("Needs to be implemented in sub-class!")

    def get_traj_per_file(self, path):
        with h5py.File(path, 'r') as F:
            return F['traj_per_file'].value

    def __getitem__(self, index):
        file_index = index // self.traj_per_file
        path = self.filenames[file_index]

        try:
            with h5py.File(path, 'r') as F:
                ex_index = index % self.traj_per_file   # get the index
                key = 'traj{}'.format(ex_index)

                # Fetch data into a dict
                if key + '/images' in F.keys():
                    data_dict = AttrDict(images=(F[key + '/images'].value))
                else:
                    data_dict = AttrDict()
                for name in F[key].keys():
                    if name in ['states', 'actions', 'pad_mask']:
                        data_dict[name] = F[key + '/' + name].value.astype(np.float32)

                # Make length consistent
                end_ind = np.argmax(data_dict.pad_mask * np.arange(data_dict.pad_mask.shape[0], dtype=np.float32), 0)
                start_ind = np.random.randint(0, end_ind - 1) if self.randomize_start else 0
                start_ind, end_ind, data_dict = self.sample_max_len_video(data_dict, start_ind, end_ind)

                # Randomize length
                if self.randomize_length:
                    end_ind = self._randomize_length(start_ind, end_ind, data_dict)

                # Collect data into the format the model expects
                data_dict.end_ind = end_ind
                data_dict.start_ind = start_ind

                self.process_data_dict(data_dict)
        except:
            raise ValueError("Problem when loading file from {}".format(path))

        return data_dict
    
    def sample_max_len_video(self, data_dict, start_ind, end_ind):
        """ This function processes data tensors so as to have length equal to max_seq_len
        by sampling / padding if necessary """
        extra_length = (end_ind - start_ind + 1) - self.spec['max_seq_len']
        if self.phase == 'train':
            offset = max(0, int(np.random.rand() * (extra_length + 1))) + start_ind
        else:
            offset = 0 
        
        data_dict = map_dict(lambda tensor: self._maybe_pad(tensor, offset, self.spec['max_seq_len']), data_dict)
        if 'actions' in data_dict:
            data_dict.actions = data_dict.actions[:-1]
        end_ind = min(end_ind - offset, self.spec['max_seq_len'] - 1)

        return 0, end_ind, data_dict        # start index gets 0 by design
    
    def _randomize_length(self, start_ind, end_ind, data_dict):
        """ This function samples part of the input tensors so that the length of the result
        is uniform between 1 and max """
        
        length = 3 + int(np.random.rand() * (end_ind - 2))  # The length of the seq is from 2 to total length
        chop_length = int(np.random.rand() * (end_ind + 1 - length))  # from 0 to the reminder
        end_ind = length - 1
        pad_mask = np.logical_and((np.arange(self.spec['max_seq_len']) <= end_ind),
                                  (np.arange(self.spec['max_seq_len']) >= start_ind)).astype(np.float32)
    
        # Chop off the beginning of the arrays
        def pad(array):
            array = np.concatenate([array[chop_length:], np.repeat(array[-1:], chop_length, 0)], 0)
            array[end_ind + 1:] = 0
            return array
    
        for key in filter(lambda key: key != 'pad_mask', data_dict):
            data_dict[key] = pad(data_dict[key])
        data_dict.pad_mask = pad_mask
        
        return end_ind

    def preprocess_images(self, images):
        # Resize video
        if len(images.shape) == 5:
            images = images[:, 0]  # Number of cameras, used in RL environments
        assert images.dtype == np.uint8, 'image need to be uint8!'
        images = resize_video(images, (self.img_sz, self.img_sz))
        images = np.transpose(images, [0, 3, 1, 2])  # convert to channel-first
        images = images.astype(np.float32) / 255 * 2 - 1
        assert images.dtype == np.float32, 'image need to be float32!'
        if self.flatten_im:
            images = np.reshape(images, [images.shape[0], -1])
        return images

    def _split_with_percentage(self, frac, filenames):
        assert sum(frac.values()) <= 1.0    # fractions cannot sum up to more than 1
        assert self.phase in frac
        if self.phase == 'train':
            start, end = 0, frac['train']
        elif self.phase == 'val':
            start, end = frac['train'], frac['train'] + frac['val']
        else:
            start, end = frac['train'] + frac['val'], frac['train'] + frac['val'] + frac['test']
        start, end = int(len(filenames) * start), int(len(filenames) * end)
        return filenames[start:end]

    @staticmethod
    def _maybe_pad(val, offset, target_length):
        """Pads / crops sequence to desired length."""
        val = val[offset:]
        len = val.shape[0]
        if len > target_length:
            return val[:target_length]
        elif len < target_length:
            return np.concatenate((val, np.zeros([int(target_length - len)] + list(val.shape[1:]), dtype=val.dtype)))
        else:
            return val

    @staticmethod
    def _shuffle_with_seed(arr, seed=2):
        rng = random.Random()
        rng.seed(seed)
        rng.shuffle(arr)
        return arr

    def __len__(self):
        if self.dataset_size != -1:
            return self.dataset_size

        return len(self.filenames) * self.traj_per_file


class FolderSplitVarLenVideoDataset(VarLenVideoDataset):
    """Splits in train/val/test using given folder structure."""

    def _get_filenames(self):
        filenames = sorted(glob.glob(os.path.join(self.data_dir, 'hdf5', self.phase + '/*')))
        if not filenames:
            raise RuntimeError('No filenames found in {}'.format(self.data_dir))
        filenames = self._shuffle_with_seed(filenames)
        return filenames


class GlobalSplitVarLenVideoDataset(VarLenVideoDataset):
    """Splits in train/val/test using global percentages."""

    def _get_filenames(self):
        filenames = []
        for root, dirs, files in os.walk(self.data_dir):
            for file in files:
                if file.endswith(".h5") and not file == 'dataset_info.h5':
                    filenames.append(os.path.join(root, file))

        if not filenames:
            raise RuntimeError('No filenames found in {}'.format(self.data_dir))
        filenames = self._shuffle_with_seed(filenames)
        filenames = self._split_with_percentage(self.spec.split, filenames)
        return filenames


class MazeGlobalSplitVarLenVideoDataset(GlobalSplitVarLenVideoDataset):
    def process_data_dict(self, data_dict):
        if 'states' in data_dict:
            data_dict['states'] = data_dict['states'][..., :2]  # only use x,y position states
        return super().process_data_dict(data_dict)


class MazeTopRenderedGlobalSplitVarLenVideoDataset(MazeGlobalSplitVarLenVideoDataset):
    def __init__(self, data_dir, mpar, data_conf, phase, shuffle=True, dataset_size=-1):
        from gcp.planning.infra.envs.miniworld_env.multiroom3d.multiroom3d_env import Multiroom3dEnv
        super().__init__(data_dir, mpar, data_conf, phase, shuffle, dataset_size)
        assert 'n_rooms' in data_conf       # need to add this in config file!
        self._crop_window_px = data_conf.crop_window
        self._render_env = Multiroom3dEnv({'n_rooms': data_conf['n_rooms']}, no_env=True,
                                          crop_window=self._crop_window_px)

    def process_data_dict(self, data_dict):
        # replace images with topdown rendered images -> first render, then resize to scale
        if "images" in data_dict:
            assert "states" in data_dict and "end_ind" in data_dict
            rendered_imgs = np.zeros((data_dict.images.shape[0], 1, self._crop_window_px*2, self._crop_window_px*2, 3),
                                     dtype=data_dict.images.dtype)
            for t in range(data_dict.end_ind + 1):
                raw_img = self._render_env.render_pos_top_down(data_dict.states[t, :2],
                                                               data_dict.states[data_dict.end_ind, :2],)
                rendered_imgs[t, 0] = np.asarray(raw_img * 255, dtype=rendered_imgs.dtype)
            data_dict.images = rendered_imgs
        return super().process_data_dict(data_dict)
<file_sep>from blox import AttrDict
from blox.torch.ops import ten2ar
import numpy as np
import numbers
from gcp.datasets.data_loader import MazeTopRenderedGlobalSplitVarLenVideoDataset


class Nav9Rooms(MazeTopRenderedGlobalSplitVarLenVideoDataset):
    n_rooms = 9
    
    @classmethod
    def render_maze_trajectories(cls, states, end_inds, color, n_logged_samples=3, bckgrds=None):
        from gcp.planning.infra.envs.miniworld_env.multiroom3d.multiroom3d_env import Multiroom3dEnv
        dummy_env = Multiroom3dEnv({'n_rooms': cls.n_rooms}, no_env=True)
        if bckgrds is None:
            bckgrds = [None] * n_logged_samples
            
        if isinstance(color[0], numbers.Number):
            color = [color] * n_logged_samples

        imgs = []
        for i, end_ind in zip(range(n_logged_samples), end_inds):
            state_seq = ten2ar(states[i][:end_ind + 1])
            # if state_seq.shape[0] < 2:
            #     if bckgrds[i] is not None:
            #         imgs.append(bckgrds[i])
            #     continue
            imgs.append(dummy_env.render_top_down(state_seq, color=color[i], background=bckgrds[i]))
        return np.stack(imgs)

    @classmethod
    def render_trajectory(cls, outputs, inputs, predictions, end_inds, n_logged_samples=3):
        # render ground truth trajectories
        im = cls.render_maze_trajectories(inputs.traj_seq_states, inputs.end_ind, (0, 1.0, 0),  # green
                                                n_logged_samples=n_logged_samples)
    
        # render predicted trajectory on top
        color = np.asarray((1.0, 0, 0))  # red
    
        if 'tree' in outputs and 'match_dist' in outputs.tree.subgoals:
            # Color bottleneck frames
            bottleneck_frames = ten2ar(outputs.tree.bf.match_dist.argmax(2)[:, :7])
            end_inds_np = ten2ar(end_inds)
            colors = []
            for i in range(n_logged_samples):
                bottleneck_frames[i, bottleneck_frames[i] > end_inds_np[i]] = end_inds_np[i]
                color_seq = color[None].repeat(end_inds_np[i] + 1, 0)
                color_seq[bottleneck_frames[i]] = color_seq[bottleneck_frames[i]] * 0.5
                color_seq[bottleneck_frames[i] - 1] = color_seq[bottleneck_frames[i] - 1] * 0.5
                colors.append(color_seq)
            color = colors
    
        im = cls.render_maze_trajectories(predictions, end_inds, color, bckgrds=im)
        return {'image': im}


config = AttrDict(
    dataset_spec=AttrDict(
        max_seq_len=100,
        dataset_class=Nav9Rooms,
        split=AttrDict(train=0.994, val=0.006, test=0.00),
    ),
    n_rooms=9,
    crop_window=40,
)
<file_sep>import matplotlib
import numpy as np

matplotlib.use('agg')
from gcp.planning.infra.datasets.save_util.record_saver import HDF5SaverBase
from blox import AttrDict


def pad_traj_timesteps(traj, max_num_actions):
    """
    pad images and actions with zeros
    :param traj:
    :param max_num_actions:
    :return:
    """

    if 'images' in traj:
        im_shape = traj.images.shape
    ac_shape = traj.actions.shape

    if ac_shape[0] < max_num_actions:
        if 'images' in traj:
            zeros = np.zeros([max_num_actions - im_shape[0] + 1, im_shape[1], im_shape[2], im_shape[3], im_shape[4]], dtype=np.uint8)
            traj.images = np.concatenate([traj.images, zeros])

        if len(ac_shape) > 1:
            zeros = np.zeros([max_num_actions - ac_shape[0], ac_shape[1]])
        else:
            zeros = np.zeros([max_num_actions - ac_shape[0]])
        traj.actions = np.concatenate([traj.actions, zeros])

    if 'images' in traj:
        assert traj.images.shape[0] == max_num_actions + 1
    assert traj.actions.shape[0] == max_num_actions

    return traj


def get_pad_mask(action_len, max_num_actions):
    """
     create a 0/1 mask with 1 where there are images and 0 where there is padding
    :param action_len:  the number of actions in trajectory
    :param max_num_actions:  maximum number of actions allowed
    :return:
    """
    if action_len < max_num_actions:
        mask = np.concatenate([np.ones(action_len + 1), np.zeros(max_num_actions - action_len)])
    elif action_len == max_num_actions:
        mask = np.ones(max_num_actions + 1)
    else:
        raise ValueError

    assert mask.shape[0] == max_num_actions + 1

    return mask


class HDF5Saver(HDF5SaverBase):
    def __init__(self, save_dir, envparams, agentparams, traj_per_file,
                 offset=0, split=(0.90, 0.05, 0.05), split_train_val_test=True):

        self.do_not_save_images = agentparams.do_not_save_images
        if hasattr(agentparams, 'max_num_actions'):
            self.max_num_actions = envparams.max_num_actions
        else:
            self.max_num_actions = agentparams.T

        super().__init__(save_dir, traj_per_file, offset, split, split_train_val_test)

    def _save_manifests(self, agent_data, obs, policy_out):
        pass

    def make_traj(self, agent_data, obs, policy_out):
        traj = AttrDict()

        if not self.do_not_save_images:
            traj.images = obs['images']
        traj.states = obs['state']
        
        action_list = [action['actions'] for action in policy_out]
        traj.actions = np.stack(action_list, 0)
        
        traj.pad_mask = get_pad_mask(traj.actions.shape[0], self.max_num_actions)
        traj = pad_traj_timesteps(traj, self.max_num_actions)

        if 'robosuite_xml' in obs:
            traj.robosuite_xml = obs['robosuite_xml'][0]
        if 'robosuite_env_name' in obs:
            traj.robosuite_env_name = obs['robosuite_env_name'][0]
        if 'robosuite_full_state' in obs:
            traj.robosuite_full_state = obs['robosuite_full_state']

        # minimal state that contains all information to position entities in the env
        if 'regression_state' in obs:
            traj.regression_state = obs['regression_state']

        return traj

    def save_traj(self, itr, agent_data, obs, policy_out):
        traj = self.make_traj(agent_data, obs, policy_out)
        self._save_traj(traj)
<file_sep>from gcp.planning.infra.envs.base_env import BaseEnv


class BaseMiniworldEnv(BaseEnv):
    def __init__(self, hp):
        self._hp = hp
        self._ncam = 1
        self.num_objects = None

        self._goal = None
        self._goaldistances = []
        self._initial_shortest_dist, self._final_shortest_dist = None, None
        self._full_traj = []

    def reset(self):
        self._goaldistances = []
        self._initial_shortest_dist, self._final_shortest_dist = None, None
        self._full_traj = []

    def set_goal(self, goal):
        self._goal = goal

    def add_goal_dist(self, goal_dist):
        self._goaldistances.append(goal_dist)

    def valid_rollout(self):
        return True     # no invalid states in miniworld env

    def eval(self):
        assert self._initial_shortest_dist is not None and self._final_shortest_dist is not None    # need to be set by subclass before eval!

        stats = {}
        stats['improvement'] = self._initial_shortest_dist - self._final_shortest_dist
        stats['initial_dist'] = self._initial_shortest_dist
        stats['final_dist'] = self._final_shortest_dist
        stats['all_goal_distances'] = self._goaldistances
        stats['full_traj'] = self._full_traj
        stats['goal'] = self._goal
        # TODO add success rate
        return stats

    def get_reset_from_obs(self, obs_dict):
        reset_state = {}
        reset_state['state'] = obs_dict['state'][0]
        reset_state['qpos_full'] = obs_dict['qpos_full'][0]
        return reset_state

    @staticmethod
    def default_ncam():
        return 1

    @property
    def ncam(self):
        return self._ncam

    @property
    def cam_res(self):
        return self._frame_height, self._frame_width



<file_sep>import numpy as np
import torch

from blox import AttrDict
from gcp.planning.infra.policy.policy import Policy
from gcp.prediction.models.auxilliary_models.bc_mdl import TestTimeBCModel


class BehavioralCloningPolicy(Policy):
    """
    Behavioral Cloning Policy
    """
    def __init__(self, ag_params, policyparams, gpu_id, ngpu, conversion_fcns=None, n_rooms=None):
        super(BehavioralCloningPolicy, self).__init__()

        self._hp = self._default_hparams()
        self.override_defaults(policyparams)

        self.log_dir = ag_params.log_dir
        self._hp.params['batch_size'] = 1
        # self._hp.params['n_actions'] = self._hp.params.n_actions  # todo get this from env!
        self.policy = TestTimeBCModel(self._hp.params, None)
        self.policy.eval()
        self.hidden_var = None      # stays None for non-recurrent policy

    def reset(self):
        super().reset()
        self.hidden_var = None

    def _default_hparams(self):
        default_dict = {
            'params': {},
            'checkpt_path': None,
            'model': None,
            'logger': None,
        }

        parent_params = super()._default_hparams()
        parent_params.ncam = 1
        for k in default_dict.keys():
            parent_params.add_hparam(k, default_dict[k])
        return parent_params

    def act(self, t=None, i_tr=None, images=None, state=None, goal=None, goal_image=None):
        # Note: goal_image provides n (2) images starting from the last images of the trajectory
        self.t = t
        self.i_tr = i_tr
        self.goal_image = goal_image

        if self.policy.has_image_input:
            inputs = AttrDict(
                I_0=self._preprocess_input(images[t]),
                I_g=self._preprocess_input(goal_image[-1] if len(goal_image.shape) > 4 else goal_image),
                hidden_var=self.hidden_var
            )
        else:
            current = state[-1:, :2]
            goal = goal[-1:, :2] #goal_state = np.concatenate([state[-1:, -2:], state[-1:, 2:]], axis=-1)
            inputs = AttrDict(
                I_0=current,
                I_g=goal,
                hidden_var=self.hidden_var
            )

        actions, self.hidden_var = self.policy(inputs)

        output = AttrDict()
        output.actions = actions.data.cpu().numpy()[0]
        return output

    @staticmethod
    def _preprocess_input(input):
        assert len(input.shape) == 4    # can currently only handle inputs with 4 dims
        if input.max() > 1.0: 
            input = input / 255.
        if input.min() >= 0.0:
            input = 2*input - 1.0
        if input.shape[-1] == 3:
            input = input.transpose(0, 3, 1, 2)
        return input

    @property
    def default_action(self):
        return np.zeros(self.policy._hp.n_actions)

    def log_outputs_stateful(self, logger=None, global_step=None, phase=None, dump_dir=None, exec_seq=None, goal=None, index=None, env=None, goal_pos=None, traj=None, topdown_image=None):
        logger.log_video(np.transpose(exec_seq, [0, 3, 1, 2]), 'control/traj{}_'.format(index), global_step, phase)
        goal_img = np.transpose(goal, [2, 0, 1])[None]
        goal_img = torch.tensor(goal_img)
        logger.log_images(goal_img, 'control/traj{}_goal'.format(index), global_step, phase)

class BehavioralCloningPolicy_RegressionState(BehavioralCloningPolicy):
    def act(self, t=None, i_tr=None, images=None, regression_state=None, goal=None, goal_image=None):
        return super().act(t, i_tr, images, regression_state, goal, goal_image)

import cv2

def resize_image(im):
    return cv2.resize(im.squeeze(), (64, 64), interpolation=cv2.INTER_AREA)


<file_sep>import matplotlib;
from gcp.prediction.utils.utils import datetime_str, make_path, set_seeds, get_dataset_path, download_data

matplotlib.use('Agg')
import torch
import argparse
import os
from shutil import copy
import imp
import importlib
from tensorflow.contrib.training import HParams
from tensorboardX import SummaryWriter
from torch import autograd
from torch.optim import Adam, RMSprop, SGD
from functools import partial

from gcp.datasets.data_loader import FolderSplitVarLenVideoDataset
from gcp.prediction.training.checkpoint_handler import save_cmd, save_git, get_config_path
from blox.utils import dummy_context
from blox import AttrDict
from blox.torch.training import LossSpikeHook, NanGradHook, NoneGradHook, DataParallelWrapper, \
    get_clipped_optimizer
from gcp.evaluation.compute_metrics import Evaluator
from gcp.prediction.training.base_trainer import BaseTrainer
from gcp.prediction import global_params
from gcp.prediction.utils import visualization
from blox.torch.radam import RAdam



class GCPBuilder(BaseTrainer):
    """ This class constructs the GCP model, dataset and the optimizers """

    def __init__(self):
        self.batch_idx = 0
    
        ## Set up params
        cmd_args, model_conf, conf_path, data_conf = self.get_configs()
    
        ## Set up logging
        self.log_dir = log_dir = os.path.join(self._hp.exp_path, 'events')
        print('using log dir: ', log_dir)
        if not cmd_args.dont_save:
            # Log
            print('Writing to the experiment directory: {}'.format(self._hp.exp_path))
            if not os.path.exists(self._hp.exp_path):
                os.makedirs(self._hp.exp_path)
            save_cmd(self._hp.exp_path)
            save_git(self._hp.exp_path)
            # Copy config file
            copy(conf_path, os.path.join(self._hp.exp_path, "conf_" + datetime_str() + ".py"))
            writer = SummaryWriter(log_dir, max_queue=1, flush_secs=1)
            self.logger = self._hp.logger(log_dir, self._hp, max_seq_len=data_conf.dataset_spec.max_seq_len,
                                          summary_writer=writer)
        else:
            self.logger = None

        ## Set up CUDA
        self.use_cuda = torch.cuda.is_available() and not global_params.debug
        self.device = torch.device('cuda') if self.use_cuda else torch.device('cpu')
        if cmd_args.gpu != -1:
            os.environ["CUDA_VISIBLE_DEVICES"] = str(cmd_args.gpu)
        if cmd_args.deterministic:
            set_seeds()

        ## Set up model conf
        # copy over data specs, but cannot copy list into hparams
        model_conf.update({k: v for k, v in data_conf.dataset_spec.items() if not isinstance(v, list)})
        model_conf['device'] = self.device.type
        model_conf['batch_size'] = self._hp.batch_size
        if self.use_cuda:
            model_conf['batch_size'] = int(self._hp.batch_size / torch.cuda.device_count())

        ## Build model
        model = self._hp.model(model_conf, self.logger)
        if torch.cuda.device_count() > 1:
            print("\nUsing {} GPUs!\n".format(torch.cuda.device_count()))
            model = DataParallelWrapper(model)
        model.device = self.device
        self.model = model = model.to(self.device)
    
        ## Build data loading
        self.train_loader = self.get_dataset(model, data_conf, 'train', self._hp.epoch_cycles_train, -1)
        phase = 'test' if self.cmd_args.metric else 'val'
        self.val_loader = self.get_dataset(model, data_conf, phase, self._hp.epoch_cycles_train, cmd_args.val_data_size)
    
        ## Build optimizer
        self.optimizer = self.get_optimizer_class()(filter(lambda p: p.requires_grad, self.model.parameters()),
                                                    lr=self._hp.lr)
    
        ## Build evaluator
        if hasattr(self.model, 'dense_rec') and not self.cmd_args.skip_top100_val:
            self.evaluator = Evaluator(self.model, self.log_dir, self._hp, self.cmd_args.metric, self.logger)
        else:
            self.evaluator = None
    
        ## Set up training options: context, hooks
        self.training_context = autograd.detect_anomaly if cmd_args.detect_anomaly else dummy_context
        self.hooks = []
        self.hooks.append(LossSpikeHook('sg_img_mse_train'))
        self.hooks.append(NanGradHook(self))
        # SVG has none gradients for the tree part of the network
        if self.model._hp.dense_rec_type != 'svg': self.hooks.append(NoneGradHook(self))

    def get_dataset(self, model, data_conf, phase, n_repeat, dataset_size=-1):
        if self.cmd_args.feed_random_data:
            from gcp.datasets.data_generator import RandomVideoDataset
            dataset_class = RandomVideoDataset
        elif 'dataset_class' in data_conf.dataset_spec:
            dataset_class = data_conf.dataset_spec.dataset_class
            visualization.PARAMS.visualize = dataset_class.visualize
        elif hasattr(data_conf, 'dataset_conf') and 'dataset_class' in data_conf.dataset_conf:
            dataset_class = data_conf.dataset_conf.dataset_class
            visualization.PARAMS.visualize = dataset_class.visualize
        else:
            dataset_class = FolderSplitVarLenVideoDataset
    
        loader = dataset_class(get_dataset_path(self._hp.dataset_name), model._hp, data_conf,
                               phase=phase, shuffle=not self.cmd_args.metric, dataset_size=dataset_size). \
            get_data_loader(self._hp.batch_size, n_repeat)
    
        return loader

    def get_configs(self):
        # Cmd arguments
        self.cmd_args = cmd_args = self.get_cmd_args()
        exp_dir = os.environ['GCP_EXP_DIR']

        # Config file
        conf_path = get_config_path(cmd_args.path)
        print('loading from the config file {}'.format(conf_path))
        conf_module = imp.load_source('conf', conf_path)

        # Trainer config
        trainer_conf = conf_module.configuration
        self._hp = self._default_hparams()
        self.override_defaults(trainer_conf)  # override defaults with config file
        self._hp.set_hparam('exp_path', make_path(exp_dir, cmd_args.path, cmd_args.prefix, cmd_args.new_dir))

        # Model config
        model_conf = conf_module.model_config

        # Data config
        download_data(trainer_conf.dataset_name)
        data_conf = self.get_data_config(conf_module)
    
        return cmd_args, model_conf, conf_path, data_conf

    def get_data_config(self, conf_module):
        # get default data config
        
        path = os.path.join(get_dataset_path(conf_module.configuration['dataset_name']), 'dataset_spec.py')
        data_conf_file = imp.load_source('dataset_spec', path)
        data_conf = AttrDict()
        data_conf.dataset_spec = AttrDict(data_conf_file.dataset_spec)

        # update with custom params if available
        update_data_conf = {}
        if hasattr(conf_module, 'data_config'):
            update_data_conf = conf_module.data_config
        elif conf_module.configuration.dataset_name is not None:
            update_data_conf = importlib.import_module('gcp.datasets.configs.' + conf_module.configuration.dataset_name).config
            
        for key in update_data_conf:
            if key == "dataset_spec":
                data_conf.dataset_spec.update(update_data_conf.dataset_spec)
            else:
                data_conf[key] = update_data_conf[key]

        if not 'fps' in data_conf:
            data_conf.fps = 4
        return data_conf

    def get_optimizer_class(self):
        optim = self._hp.optimizer
        if optim == 'adam':
            get_optim = partial(get_clipped_optimizer, optimizer_type=Adam, betas=(self._hp.adam_beta, 0.999))
        elif optim == 'radam':
            get_optim = partial(get_clipped_optimizer, optimizer_type=RAdam, betas=(self._hp.adam_beta, 0.999))
        elif optim == 'rmsprop':
            get_optim = partial(get_clipped_optimizer, optimizer_type=RMSprop, momentum=self._hp.momentum)
        elif optim == 'sgd':
            get_optim = partial(get_clipped_optimizer, optimizer_type=SGD, momentum=self._hp.momentum)
        else:
            raise ValueError("Optimizer '{}' not supported!".format(optim))
        return partial(get_optim, gradient_clip=self._hp.gradient_clip)

    def get_cmd_args(self):
        parser = argparse.ArgumentParser()
        parser.add_argument("--path", help="path to the config file directory")
    
        # Folder settings
        parser.add_argument("--prefix", help="experiment prefix, if given creates subfolder in experiment directory")
        parser.add_argument('--new_dir', default=False, type=int, help='If True, concat datetime string to exp_dir.')
        parser.add_argument('--dont_save', default=False, type=int,
                            help="if True, nothing is saved to disk. Note: this doesn't work")  # TODO this doesn't work
    
        parser.add_argument("--visualize", default='', help="path to model file to visualize")  # TODO what uses this?
    
        # Running protocol
        parser.add_argument('--resume', default='latest', type=str, metavar='PATH',
                            help='path to latest checkpoint (default: none)')
        parser.add_argument('--train', default=True, type=int,
                            help='if False, will run one validation epoch')
        parser.add_argument('--test_control', default=True, type=int,
                            help="if False, control isn't run at validation time")
        parser.add_argument('--test_prediction', default=True, type=int,
                            help="if False, prediction isn't run at validation time")
        parser.add_argument('--skip_first_val', default=False, type=int,
                            help='if True, will skip the first validation epoch')
        parser.add_argument('--skip_top100_val', default=False, type=int,
                            help="if True, skips top of 100 eval")
        parser.add_argument('--metric', default=False, type=int,
                            help='if True, run test metrics')
        parser.add_argument('--val_sweep', default=False, type=int,
                            help='if True, runs validation on all existing model checkpoints')
        parser.add_argument('--dataset_val_sweep', default=False, type=int,
                            help='if True, runs validation on a given collection of datasets')
    
        # Misc
        parser.add_argument('--gpu', default=-1, type=int,
                            help='will set CUDA_VISIBLE_DEVICES to selected value')
        parser.add_argument('--strict_weight_loading', default=True, type=int,
                            help='if True, uses strict weight loading function')
        parser.add_argument('--deterministic', default=False, type=int,
                            help='if True, sets fixed seeds for torch and numpy')
        parser.add_argument('--imepoch', default=4, type=int,
                            help='number of image loggings per epoch')
        parser.add_argument('--val_data_size', default=-1, type=int,
                            help='number of sequences in the validation set. If -1, the full dataset is used')
        parser.add_argument('--log_images', default=True, type=int,
                            help='')
        parser.add_argument('--log_outputs_interval', default=10, type=int,
                            help='')
    
        # Debug
        parser.add_argument('--detect_anomaly', default=False, type=int,
                            help='if True, uses autograd.detect_anomaly()')
        parser.add_argument('--feed_random_data', default=False, type=int,
                            help='if True, we feed random data to the model to test its performance')
        parser.add_argument('--train_loop_pdb', default=False, type=int,
                            help='if True, opens a pdb into training loop')
        parser.add_argument('--debug', default=False, type=int,
                            help='if True, runs in debug mode')
        parser.add_argument('--verbose_timing', default=False, type=int,
                            help='if True, prints additional time measurements.')
        return parser.parse_args()

    def _default_hparams(self):
        # put new parameters in here:
        default_dict = {
            'model': None,
            'logger': None,
            'dataset_name': None,  # directory where dataset is in
            'batch_size': 64,
            'exp_path': None,  # Path to the folder with experiments
            'num_epochs': 200,
            'epoch_cycles_train': 1,
            'optimizer': 'radam',  # supported: 'radam', 'adam', 'rmsprop', 'sgd'
            'lr': None,
            'gradient_clip': None,
            'momentum': 0,  # momentum in RMSProp / SGD optimizer
            'adam_beta': 0.9,  # beta1 param in Adam
            'metric_pruning_scheme': 'dtw',
            'top_of_100_eval': True,
            'n_rooms': None,
        }
        # add new params to parent params
        parent_params = HParams()
        for k in default_dict.keys():
            parent_params.add_hparam(k, default_dict[k])
        return parent_params
<file_sep>from blox import AttrDict


def get_default_gcp_hyperparameters():
    
    # Params that actually should be elsewhere
    default_dict = AttrDict({
        'randomize_length': False,
        'randomize_start': False,
    })
    
    # Network size
    default_dict.update({
        'ngf': 4,  # number of feature maps in shallowest level
        'nz_enc': 32,  # number of dimensions in encoder-latent space
        'nz_vae': 32,  # number of dimensions in vae-latent space
        'nz_vae2': 256,  # number of dimensions in 2nd level vae-latent space (if used)
        'nz_mid': 32,  # number of dimensions for internal feature spaces
        'nz_mid_lstm': 32,
        'n_lstm_layers': 1,
        'n_processing_layers': 3,  # Number of layers in MLPs
        'conv_inf_enc_kernel_size': 3,  # kernel size of convolutional inference encoder
        'conv_inf_enc_layers': 1,  # number of layers in convolutional inference encoder
        'n_attention_heads': 1,  # number of attention heads (needs to divide nz_enc evenly)
        'n_attention_layers': 1,  # number of layers in attention network
        'nz_attn_key': 32,  # dimensionality of the attention key
        'init_mlp_layers': 3,  # number of layers in the LSTM initialization MLP (if used)
        'init_mlp_mid_sz': 32,  # size of hidden layers inside LSTM initialization MLP (if used)
        'n_conv_layers': None,  # Number of conv layers. Can be of format 'n-<int>' for any int for relative spec
    })
    
    # Network params
    default_dict.update(AttrDict(
        action_activation=None,
        device=None,
        context_every_step=True,
    ))
    
    # Loss weights
    default_dict.update({
        'kl_weight': 1.,
        'kl_weight_burn_in': None,
        'entropy_weight': .0,
        'length_pred_weight': 1.,
        'dense_img_rec_weight': 1.,
        'dense_action_rec_weight': 1.,
        'free_nats': 0,
    })
    
    # Architecture params
    default_dict.update({
        'use_skips': True,  # only works with conv encoder/decoder
        'skips_stride': 2,
        'add_weighted_pixel_copy': False,  # if True, adds pixel copying stream for decoder
        'pixel_shift_decoder': False,
        'skip_from_parents': False,  # If True, parents are added to the pixel copy/shift sources
        'seq_enc': 'none',  # Manner of sequence encoding. ['none', 'conv', 'lstm']
        'regress_actions': False,  # whether or not to regress actions
        'learn_attn_temp': True,  # if True, makes attention temperature a trainable variable
        'attention_temperature': 1.0,  # temperature param used in attention softmax
        'attach_inv_mdl': False,  # if True, attaches an inverse model to output that computes actions
        'attach_cost_mdl': False,  # if True, attaches a cost function MLP that estimates cost from pairs of states
        'run_cost_mdl': True,   # if False, does not run cost model (but might still build it
        'attach_state_regressor': False,    # if True, attaches network that regresses states from pre-decoding-latents
        'action_conditioned_pred': False,  # if True, conditions prediction on actions
        'learn_beta': True,  # if True, learns beta
        'initial_sigma': 1.0,  # if True, learns beta
        'separate_cnn_start_goal_encoder': False,   # if True, builds separate conv encoder for start/goal image
        'decoder_distribution': 'gaussian'  # [gaussian, categorical]
    })

    # RNN params
    default_dict.update({
        'use_conv_lstm': False,
    })
    
    # Variational inference parameters
    default_dict.update(AttrDict(
        prior_type='learned',  # type of prior to be used ['fixed', 'learned']
        var_inf='standard',  # type of variation inference ['standard', '2layer', 'deterministic']
    ))
    
    # RecPlan params
    default_dict.update({
        'hierarchy_levels': 3,  # number of levels in the subgoal tree
        
        'one_hot_attn_time_cond': False,  # if True, conditions attention on one-hot time step index
        'attentive_inference': False,  # if True, forces attention to single matching frame
        'non_goal_conditioned': False,  # if True, does not condition prediction on goal frame
        
        'tree_lstm': '',  # ['', 'sum' or 'linear']
        'lstm_init': 'zero',  # defines how treeLSTM is initialized, ['zero', 'mlp'], #, 'warmup']
        
        'matching_temp': 1.0,  # temperature used in TAP-style matching softmax
        'matching_temp_tenthlife': -1,
        'matching_temp_min': 1e-3,
        'matching_type': 'latent',  # evidence binding procedure
        # ['image', 'latent', 'fraction', 'balanced', 'tap']
        'leaves_bias': 0.0,
        'top_bias': 1.0,
        'n_top_bias_nodes': 1,
        'supervise_match_weight': 0.0,
        
        'regress_index': False,
        'regress_length': False,
        
        'inv_mdl_params': {},  # params for the inverse model, if attached
        'train_inv_mdl_full_seq': False,  # if True, omits sampling for inverse model and trains on full seq

        'cost_mdl_params': {},   # cost model parameters

        'act_cond_inference': False,  # if True, conditions inference on actions

        'train_on_action_seqs': False,  # if True, trains the predictive network on action sequences

        'learned_pruning_threshold': 0.5,   # confidence thresh for learned pruning network after which node gets pruned
        'untied_layers': False,
        'supervised_decoder': False,
        'states_inference': False,
    })
    
    # Outdated GCP params
    default_dict.update({
        'dense_rec_type': 'none',  # ['none', 'discrete', 'softmax', 'linear', 'node_prob', action_prob].
        'one_step_planner': 'discrete',  # ['discrete', 'continuous', 'sh_pred']. Always 'sh_pred' for HEDGE models
        'mask_inf_attention': False,  # if True, masks out inference attention outside the current subsegment
        'binding': 'frames',  # Matching loss form ['loss', 'frames', 'exp', 'lf']. Always 'loss'.
    })
    
    # Matching params
    default_dict.update(AttrDict(
        learn_matching_temp=True,  # If true, the matching temperature is learned
    ))
    
    # Logging params
    default_dict.update(AttrDict(
        dump_encodings='',  # Specifies the directory where to dump the encodings
        dump_encodings_inv_model='',  # Specifies the directory where to dump the encodings
        log_states_2d=False,  # if True, logs 2D plot of first two state dimensions
        log_cartgripper=False,  # if True, logs sawyer from states
        data_dir='',   # necessary for sawyer logging
    ))

    # Hyperparameters that shouldn't exist
    default_dict.update(AttrDict(
        log_d2b_3x3maze=0,  # Specifies the directory where to dump the encodings
    ))
    
    
    return default_dict
<file_sep>from setuptools import setup
from Cython.Build import cythonize
import numpy

setup(
    name='gcp',
    version='0.2dev',
    packages=['gcp', 'blox', 'gym-miniworld'],
    license='MIT License',
    ext_modules=cythonize(['gcp/evaluation/*.pyx']),
    include_dirs=[numpy.get_include(),],
)
<file_sep>import torch
import numpy as np

from blox import AttrDict


class GCPSimulator:
    """Implements simple simulator interface for GCP models."""
    def __init__(self, model, append_latent):
        self._model = model
        self._append_latent = append_latent  # if True, appends the latent to the state rollout
        self._logs = []

    def rollout(self, state, goal_state, samples, rollout_len, prune=False):
        """Performs one model rollout."""
        # prepare inputs
        batch_size = samples.shape[0]
        state, goal_state = state.repeat(batch_size, 0), goal_state.repeat(batch_size, 0)
        input_dict = AttrDict(I_0=torch.tensor(state, device=self._model.device, dtype=torch.float32),
                              I_g=torch.tensor(goal_state, device=self._model.device, dtype=torch.float32),
                              start_ind=torch.tensor(np.zeros((batch_size,)), device=self._model.device).long(),
                              end_ind=torch.tensor(np.ones((batch_size,)) * (rollout_len - 1),
                                                   device=self._model.device).long(),
                              z=torch.tensor(samples, device=self._model.device, dtype=torch.float32))
        input_dict = self._postprocess_inputs(input_dict)

        # perform rollout, collect outputs
        outputs = AttrDict()
        with self._model.val_mode():
            model_output = self._model(input_dict)
            end_ind = torch.max(model_output.end_ind, torch.ones_like(model_output.end_ind))
        # self._logs.append(model_output)

        if prune:
            outputs.predictions = self._list2np(model_output.pruned_prediction)
        else:
            outputs.predictions = self._list2np(self._get_state_rollouts(input_dict, model_output, end_ind))

        outputs.actions = self._list2np(self._cap_to_length(model_output.actions, end_ind))
        outputs.states = self._list2np(self._cap_to_length(model_output.regressed_state, end_ind))
        outputs.latents = self._list2np(self._cap_to_length(input_dict.model_enc_seq, end_ind))

        return outputs

    def _postprocess_inputs(self, input_dict):
        return input_dict

    def _get_state_rollouts(self, input_dict, model_output, end_ind):
        batch_size = model_output.end_ind.shape[0]
        state_plans = []
        for i in range(batch_size):
            seq_len = end_ind[i] + 1
            out = self._model.dense_rec.get_sample_with_len(i, seq_len, model_output, input_dict, 'basic')
            rollout = out[0].reshape(seq_len, -1)
            if self._append_latent:
                name = 'encodings' if model_output.dense_rec else 'e_g_prime'
                latent_rollout = self._model.dense_rec.get_sample_with_len(
                    i, seq_len, model_output, input_dict, 'basic', name=name)[0].reshape(seq_len, -1)
                rollout = torch.cat((rollout, latent_rollout), dim=-1)
            state_plans.append(rollout)
        return state_plans

    @staticmethod
    def _cap_to_length(vals, end_inds):
        assert vals.shape[0] == end_inds.shape[0]
        return [val[:end_ind + 1] for val, end_ind in zip(vals, end_inds)]

    @staticmethod
    def _list2np(list):
        return [elem.data.cpu().numpy() for elem in list]

    def dump_logs(self, dump_file='rollout_dump.pkl'):
        import pickle
        with open(dump_file, 'wb') as F:
            pickle.dump(self._logs, F)
        self._logs = []


class GCPImageSimulator(GCPSimulator):
    def _postprocess_inputs(self, input_dict):
        input_dict = super()._postprocess_inputs(input_dict)
        input_dict.z = input_dict.z[..., None, None]
        input_dict.I_0 = self._env2planner(input_dict.I_0)
        input_dict.I_g = self._env2planner(input_dict.I_g)
        return input_dict

    @staticmethod
    def _env2planner(img):
        """Converts images to the [-1...1] range of the hierarchical planner."""
        if img.max() > 1.0:
            img = img / 255.0
        if len(img.shape) == 5:
            img = img[0]
        if len(img.shape) == 4:
            img = img.permute(0, 3, 1, 2)
        return img * 2 - 1.0


class ActCondGCPImageSimulator(GCPImageSimulator):
    def _postprocess_inputs(self, input_dict):
        input_dict = super()._postprocess_inputs(input_dict)
        input_dict.actions = input_dict.pop('z')[..., 0, 0]     # input (action) samples are stored in z by default
        input_dict.pad_mask = torch.ones(input_dict.actions.shape[:2], device=input_dict.actions.device)
        return input_dict



<file_sep>import numpy as np
import copy
import os
import torch
import pickle as pkl
from collections import defaultdict
from blox import AttrDict
from blox.utils import ParamDict
from blox.basic_types import listdict2dictlist

from gcp.planning.cem.cost_fcn import EuclideanPathLength, LearnedCostEstimate
from gcp.planning.cem.sampler import FlatCEMSampler, HierarchicalTreeCEMSampler


class CEMPlanner:
    """Generic CEM planner."""
    def __init__(self, hp, simulator):
        self._hp = self._default_hparams().overwrite(hp)
        self._simulator = simulator
        self._cost_fcn = self._build_cost()
        self._sampler = self._build_sampler()
        self._logs = []

    def _default_hparams(self):
        default_dict = ParamDict(
            horizon=None,        # CEM optimization horizon (i.e. how many sequential actions get optimized)
            action_dim=None,     # dimensionality of the actions that are optimized
            n_iters=1,           # number of CEM iterations
            batch_size=64,       # number of rollouts per iteration
            max_rollout_bs=100,  # maximum batch size for rollout (splits 'batch_size' if too large)
            elite_frac=0.1,      # percentage of 'best' trajectories
        )
        # cost params
        default_dict.update(ParamDict(
            cost_fcn=EuclideanPathLength,
            dense_cost=False,
            final_step_cost_weight=1.0,
        ))
        # sampler params
        default_dict.update(ParamDict(
            sampler=FlatCEMSampler,
            sampler_clip_val=float("Inf"),
            initial_std=3e-1,
        ))
        # misc
        default_dict.update(ParamDict(
            verbose=False,                  # whether to visualize planning procedure (for debugging)
            dump_planning_data=False,       # whether to dump raw planning data
            use_delta_state_actions=False,  # if True, uses delta between inferred states as action plan
            use_inferred_actions=True,      # if True, uses model-inferred actions for action plan
            max_seq_len=None,               # used for model during rollout
        ))
        return default_dict

    def __call__(self, state, goal_state):
        logs = []
        self._sampler.init()
        for cem_iter in range(self._hp.n_iters):
            # sample actions
            samples = self._sampler.sample(self._hp.batch_size)

            # rollout simulator
            rollouts = self._rollout(state, goal_state, samples)

            best_rollouts, best_rollouts_states, best_scores, best_samples, elite_idxs = \
                self._get_best_rollouts(rollouts, goal_state, samples)

            # refit action distribution
            self._sampler.fit(best_samples, best_scores)

            # store all logs
            logs.append(AttrDict(
                elite_rollouts=copy.deepcopy(best_rollouts),
                elite_scores=best_scores,
                dists=self._sampler.get_dists(),
                goal_state=goal_state,
                elite_states=copy.deepcopy(best_rollouts_states),
            ))

        # perform final rollout with best actions
        final_rollouts = self._rollout(state, goal_state, best_samples)
        logs.append(AttrDict(
            elite_rollouts=copy.deepcopy(self._maybe_split_image(final_rollouts.predictions)),
            elite_scores=best_scores,
            dists=self._sampler.get_dists(),
            goal_state=goal_state,
            elite_states=copy.deepcopy(final_rollouts.states),
        ))

        # extract output action plan
        best_actions = self._get_action_plan(final_rollouts, best_samples)

        # save logs
        self._logs.append(logs)

        return final_rollouts.predictions[0], best_actions[0], final_rollouts.latents[0], best_scores[0]

    def log_verbose(self, logger, step, phase, i_tr, dump_dir):
        if self._hp.dump_planning_data:
            os.makedirs(os.path.join(dump_dir, "planning"), exist_ok=True)
            with open(os.path.join(dump_dir, "planning/traj{}_raw_data.pkl".format(i_tr)), "wb") as F:
                pkl.dump(self._logs, F)

        self._logs = []

    def _build_cost(self):
        return self._hp.cost_fcn(self._hp.dense_cost, self._hp.final_step_cost_weight)

    def _build_sampler(self):
        return self._hp.sampler(self._hp.sampler_clip_val,
                                self._hp.max_seq_len,
                                self._hp.action_dim,
                                self._hp.initial_std)

    def _rollout(self, state, goal, samples):
        output = defaultdict(list)
        for i in range(max(samples.shape[0] // self._hp.max_rollout_bs, 1)):
            sim_output = self._simulator.rollout(state, goal,
                                                 samples[i * self._hp.max_rollout_bs: (i + 1) * self._hp.max_rollout_bs],
                                                 self._hp.max_seq_len)
            output = self._join_dicts(sim_output, output)
        return AttrDict({key: self._cap_to_horizon(output[key]) for key in output})

    def _get_best_rollouts(self, rollouts, goal_state, samples):
        # compute rollout scores
        scores = self._cost_fcn(rollouts.predictions, goal_state)

        # get idxs of best rollouts
        full_elite_idxs = scores.argsort()
        elite_idxs = full_elite_idxs[:int(self._hp.batch_size * self._hp.elite_frac)]

        best_rollouts, best_rollouts_states = \
            [rollouts.predictions[idx] for idx in elite_idxs], [rollouts.states[idx] for idx in elite_idxs],
        best_scores, best_samples = scores[elite_idxs], samples[elite_idxs]
        return self._maybe_split_image(best_rollouts), best_rollouts_states, best_scores, best_samples, elite_idxs

    def _maybe_split_image(self, rollout):
        if hasattr(self._cost_fcn, "_split_state_rollout"):
            # separate latent and image in case that latent got attached to rollout
            return self._cost_fcn._split_state_rollout(rollout).image_rollout
        return rollout

    def _get_action_plan(self, final_rollouts, best_samples):
        if self._hp.use_delta_state_actions:
            return [b[1:] - b[:-1] for b in final_rollouts.states]
        elif self._hp.use_inferred_actions:
            return final_rollouts.actions
        else:
            return best_samples

    def _cap_to_horizon(self, input):
        if self._hp.horizon is not None:
            return [elem[:self._hp.horizon] for elem in input]
        else:
            return input

    @property
    def append_latent(self):
        return self._sampler.append_latent

    @staticmethod
    def _join_dicts(d1, d2):
        return AttrDict({key: d1[key] + d2[key] for key in d1})


class HierarchicalCEMPlanner(CEMPlanner):
    """CEM planner for hierarchical optimization."""
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        if self._hp.sampling_rates_per_layer is not None:
            assert self._hp.n_iters == len(self._hp.sampling_rates_per_layer) + 1

    def _default_hparams(self):
        default_dict = super()._default_hparams()
        # general params
        default_dict.update(ParamDict(
            horizon=None,       # for GCP we do not need to define horizon
        ))
        # cost params
        default_dict.update(ParamDict(
            cost_fcn=LearnedCostEstimate,
            cost_config={},     # cost function for subgoal optimization
            LL_cost_fcn=None,   # if None cost_fcn is used
        ))
        # sampler params
        default_dict.update(ParamDict(
            sampler=HierarchicalTreeCEMSampler,
            n_level_hierarchy=None,
            sampling_rates_per_layer=None,
            n_ll_samples=5,
        ))
        return default_dict

    def _build_cost(self):
        cost_fcn = self._hp.cost_fcn(self._hp.cost_config)
        self._ll_cost_fcn = cost_fcn if self._hp.LL_cost_fcn is None \
            else self._hp.LL_cost_fcn(self._hp.dense_cost, self._hp.final_step_cost_weight)
        return cost_fcn

    def _build_sampler(self):
        return self._hp.sampler(self._hp.sampler_clip_val,
                                self._hp.max_seq_len,
                                self._hp.action_dim,
                                self._hp.initial_std,
                                n_level_hierarchy=self._hp.n_level_hierarchy,
                                sampling_rates_per_layer=self._hp.sampling_rates_per_layer,
                                subgoal_cost_fcn=self._cost_fcn,
                                ll_cost_fcn=self._ll_cost_fcn,
                                n_ll_samples=self._hp.n_ll_samples)

    def _get_best_rollouts(self, rollouts, goal_state, samples):
        if not isinstance(self._sampler, HierarchicalTreeCEMSampler):
            # in case we use non-hierarchical optimization with tree-based model
            return super()._get_best_rollouts(rollouts, goal_state, samples)
        best_rollouts, best_scores = self._sampler.optimize(rollouts.predictions, goal_state)
        best_samples = self._sampler.sample(self._hp.batch_size)
        elite_idxs = np.arange(len(best_rollouts))      # dummy value
        return best_rollouts, rollouts.states, best_scores, best_samples, elite_idxs


class ImageCEMPlanner(CEMPlanner):
    def log_verbose(self, logger, step, phase, i_tr, dump_dir):
        if self._hp.verbose:
            for replan_idx, replan_log in enumerate(self._logs):
                for cem_iter_idx, iter_log in enumerate(replan_log):

                    # visualize all plans in order
                    plan_stack = []
                    for plan in iter_log.elite_rollouts:
                        time, c, h, w = plan.shape
                        plan = np.clip((plan+1) / 2, 0, 1.0)
                        if time < self._hp.horizon:
                            plan = np.concatenate((plan, np.ones((self._hp.horizon - time, c, h, w))))
                        plan_stack.append(plan)
                    plan_stack = np.array(plan_stack)
                    n_plans = plan_stack.shape[0]
                    log_img = torch.tensor(plan_stack.transpose(0, 2, 3, 1, 4)
                                           .reshape(n_plans, c, h, self._hp.horizon*w)
                                           .transpose(1, 0, 2, 3).reshape(c, h*n_plans, self._hp.horizon*w))
                    logger.log_images(log_img[None],
                            "elite_trajs_{}_test/plan_r{}_iter{}_overview".format(i_tr, replan_idx, cem_iter_idx),
                            step, phase)

                    if 'elite_states' in iter_log:
                        logger.log_single_topdown_traj(iter_log.elite_states[0],
                                                       "elite_trajs_{}_test/plan_r{}_iter{}_z_inferStateTraj".
                                                       format(i_tr, replan_idx, cem_iter_idx), step, phase)
                        logger.log_multiple_topdown_trajs(iter_log.elite_states,
                                                       "elite_trajs_{}_test/plan_r{}_iter{}_z_inferStateTrajDist".
                                                       format(i_tr, replan_idx, cem_iter_idx), step, phase)
                    if 'goal_state' in iter_log:
                        logger.log_images(torch.tensor(iter_log['goal_state'].transpose(0, 3, 1, 2)),
                                          "elite_trajs_{}_test/plan_r{}_iter{}_z_goal".
                                          format(i_tr, replan_idx, cem_iter_idx), step, phase)
        super().log_verbose(logger, step, phase, i_tr, dump_dir)

    def hack_add_state(self, state):
        self._logs[-1][-1].state = state.copy()


class HierarchicalImageCEMPlanner(HierarchicalCEMPlanner, ImageCEMPlanner):
    def log_verbose(self, *args, **kwargs):
        ImageCEMPlanner.log_verbose(self, *args, **kwargs)



<file_sep>import numpy as np


def _rand(n):
    return np.random.rand(n) - 0.5      # sample in range [-0.5, 0.5]


def _M(*elems):
    return np.array(elems)


class RoomSampler2d:
    """Samples 2d coordinates in a room environment with given number of rooms."""
    def __init__(self, rooms_per_side, sample_wide=False):
        """If sample_wide is True, the sampling fills the whole room all the way to the walls."""
        super().__init__()
        self._rooms_per_side = rooms_per_side

        self._agent_size = 0.02
        self._sampling_width = 1/3 - (not sample_wide) * 3 * self._agent_size      # equivalent to width of one room in mujoco
        self._room_offset = 1/3       # equivalent to middle point of rooms
        self._door_sampling_width = 1.5 * 0.0667 - 3 * self._agent_size

        self._hor_door_sampling_width = _M(2*self._agent_size, self._door_sampling_width)
        self._vert_door_sampling_width = _M(self._door_sampling_width, 2 * self._agent_size)

    def sample(self, room=None):
        if room is None:
            room = np.random.randint(self._rooms_per_side**2)
        room = self._ridx2coords(room)
        center = _M(*[self._room_offset / 2 + i * self._room_offset
                     - self._rooms_per_side / 2 * self._room_offset for i in room])  # centered around (0, 0)
        return _rand(2) * self._sampling_width + center

    def sample_door(self, room1, room2, sample_center=False):
        """Samples in the door way between two given rooms."""
        center = self.get_door_pos(room1, room2)
        if sample_center: return center
        room1, room2 = self._ridx2coords(room1), self._ridx2coords(room2)
        if room1[0] != room2[0] and room1[1] == room2[1]:
            # horizontal room connection
            return _rand(2) * self._hor_door_sampling_width + center
        elif room1[0] == room2[0] and room1[1] != room2[1]:
            # vertical room connection
            return _rand(2) * self._vert_door_sampling_width + center
        else:
            raise ValueError("Rooms don't have connection for door.")

    def get_door_pos(self, room1, room2):
        assert room1 < room2  # room1 needs to be on top or left of room2
        room1, room2 = self._ridx2coords(room1), self._ridx2coords(room2)
        assert np.abs(
            room1[0] - room2[0] + room1[1] - room2[1]) == 1  # difference between given rooms needs to be exactly 1
        center = _M(*[self._room_offset / 2 + (i + j) / 2 * self._room_offset
                      - self._rooms_per_side / 2 * self._room_offset for i, j in zip(room1, room2)])
        return center

    def get_door_path(self, room1, room2):
        """Returns path through the door between two rooms (i.e. three waypoints)."""
        lefttop = room1 < room2     # check if room 1 is on left/top of room2
        center = self.get_door_pos(min(room1, room2), max(room1, room2))
        room1, room2 = self._ridx2coords(room1), self._ridx2coords(room2)
        if room1[0] != room2[0] and room1[1] == room2[1]:
            # horizontal room connection
            offset = _M(3 * self._door_sampling_width, 0)
        elif room1[0] == room2[0] and room1[1] != room2[1]:
            # vertical room connection
            offset = _M(0, -3 * self._door_sampling_width)
        else:
            raise ValueError("Rooms don't have connection for door.")
        if lefttop:
            return [center - offset, center, center + offset]
        else:
            return [center + offset, center, center - offset]

    def _ridx2coords(self, room_idx):
        """Converts room index to coordinates based on grid size."""
        return (int(np.floor(room_idx / self._rooms_per_side)),
                int(self._rooms_per_side - 1 - room_idx % self._rooms_per_side))


if __name__ == "__main__":
    import matplotlib.pyplot as plt

    DOORS = [(0, 3), (3, 6), (1, 4), (4, 7), (5, 8), (3, 4), (1, 2), (7, 8)]
    sampler = RoomSampler2d(rooms_per_side=3, sample_wide=True)
    samples = np.asarray([[sampler.sample(r) for _ in range(100)] for r in range(36)]).transpose(2, 0, 1).reshape(2, -1)
    plt.scatter(samples[0], samples[1], c='black')

    samples = np.asarray([[sampler.sample_door(d[0], d[1]) for _ in range(10)] for d in DOORS]).transpose(2, 0, 1).reshape(2, -1)
    plt.scatter(samples[0], samples[1], c='red')
    plt.show()

<file_sep>import numpy as np

from blox import AttrDict
from blox.basic_types import listdict2dictlist


class HierarchicalTreeLatentOptimizer:
    """Optimizes latent distributions for GCP-tree layers recursively, one layer at a time.
       After N layers have been optimized hierarchically, all remaining layers are jointly optimized as 'segments'
       connecting between the planned subgoals."""
    def __init__(self, latent_dim, sampling_rates, depth, subgoal_cost_fcn, ll_cost_fcn, final_layer_samples):
        """
        :param latent_dim: dimensionality of optimized latent
        :param sampling_rates: per-layer sampling rates (except for last layer) as list
        :param depth: depth of GCP tree model who's latents are getting optimized
        :param subgoal_cost_fcn: cost function for estimating cost of sampled predictions
        :param ll_cost_fcn: cost function for estimating cost of dense prediction in last layer
        :param final_layer_samples: number of samples for optimizing last layer's dense trajectory predictions
        """
        self._latent_dim = latent_dim
        self._depth = depth
        self._subgoal_cost_fcn = subgoal_cost_fcn
        self._ll_cost_fcn = ll_cost_fcn
        self._is_optimized = False          # indicates whether this layer of the tree is already optimized
        self._opt_z = None                  # holds optimal subgoal latent once optimized
        self._latest_z_samples = None       # holds last-sampled z samples for subgoal optimization
        self._dummy_env = None              # used for visualization purposes
        if sampling_rates:
            # not yet at bottom-most layer
            self._n_samples = sampling_rates.pop(0)
            self._n_latents = 1
            self._children = [[type(self)(latent_dim, sampling_rates.copy(), depth - 1,
                                          self._subgoal_cost_fcn, self._ll_cost_fcn,
                                          final_layer_samples)
                               for _ in range(self._n_samples)] for _ in range(2)]
        else:
            # final layer, create remaining samples for non-hierarchical 'segment' optimization
            self._n_samples = final_layer_samples
            self._n_latents = 2**depth - 1
            self._children = None

        self.mean = np.zeros((self._n_latents, self._latent_dim))
        self.std = np.ones((self._n_latents, self._latent_dim))

    def sample(self, below_opt_layer=False):
        """Samples latents from all layers of the tree, returns concatenated result.
           Samples N latents for layer that's currently getting optimized, only 1 latent for all layers above and below.
        :param below_opt_layer: indicates whether layer is below the layer that is currently getting optimized."""
        # sample current layer's latents
        if self._is_optimized:      # if layer is already optimized --> sample optimized subgoal latent
            z = self._opt_z.copy()[None]
        else:
            # sample N latents if this is currently optimized layer,
            # if below optimized layer sample only single latent (since decoding won't be used for optimization)
            z = self._sample()[:1] if below_opt_layer else self._sample()
            self._latest_z_samples = z.copy()
        next_below_opt_layer = (not self._is_optimized and not below_opt_layer) \
                                    or below_opt_layer  # in the first case the current layer is getting optimized
        if self._children is not None:
            # sample children's latents and concatenate
            samples = []
            for child_left, child_right, z_i in zip(self._children[0], self._children[1], z):
                z_left, z_right = child_left.sample(next_below_opt_layer), child_right.sample(next_below_opt_layer)
                assert z_left.shape == z_right.shape    # latent tree needs to be balanced
                samples.append(np.concatenate([z_left, np.tile(z_i[0], (z_left.shape[0], 1, 1)), z_right], axis=1))
            z = np.concatenate(samples)

        return z

    def optimize(self, all_rollouts, goal):
        """Optimizes subgoal in all layers sequentially, optimizes full rollout in last layer."""
        if self._children is None:      # final layer --> optimize dense segment
            return self._optimize_segment(all_rollouts, goal)
        elif not self._is_optimized:    # non-final layer, not optimized --> optimize subgoal
            return self._optimize_subgoal(all_rollouts, goal)
        else:                           # non-final layer, already optimized --> recurse
            return self._recurse_optimization(all_rollouts, goal)

    def _optimize_segment(self, all_rollouts, goal):
        """Optimizes final-layer 'segment' between subgoals."""
        best_rollout, best_cost, best_idx = self._best_of_n_segments(all_rollouts, goal, self._ll_cost_fcn)
        self._opt_z = self._latest_z_samples[best_idx]
        self._is_optimized = True
        return best_rollout, best_cost

    def _optimize_subgoal(self, all_rollouts, goal):
        """Optimizes subgoal latent by minimizing pairwise cost estimate with both parents."""
        # first split off latents from concatenated rollouts
        rollouts = self._split_state_rollout(all_rollouts)

        # prepare start + goal + subgoal arrays
        starts, start_latents = np.stack([r[0] for r in rollouts.state_rollout]), \
                                np.stack([r[0] for r in rollouts.latent_rollout])
        subgoals = np.stack([r[int(np.floor(r.shape[0] / 2))] for r in rollouts.state_rollout])
        subgoal_latents = np.stack([r[int(np.floor(r.shape[0] / 2))] for r in rollouts.latent_rollout])
        goals = np.stack([self._split_state_rollout([goal[None]]).state_rollout[0][0] \
                              if goal.shape[-1] == all_rollouts[0].shape[-1] else goal
                          for _ in rollouts.state_rollout])
        goal_latents = np.stack([self._split_state_rollout([goal[None]]).latent_rollout[0][0] \
                                     if goal.shape[-1] == all_rollouts[0].shape[-1] else r[-1]
                                 for r in rollouts.latent_rollout])

        # compute pairwise cost
        to_cost, from_cost = self._subgoal_cost_fcn(start_latents, subgoal_latents), \
                             self._subgoal_cost_fcn(subgoal_latents, goal_latents)
        total_cost = to_cost + from_cost

        # find optimal subgoal
        opt_z_idx = np.argmin(total_cost)
        self._opt_z = self._latest_z_samples[opt_z_idx]

        # construct output rollout + output_cost
        outputs = [starts[opt_z_idx]]
        if (subgoals[opt_z_idx] != outputs[-1]).any():  # they can be the same if sequence is too short
            outputs.append(subgoals[opt_z_idx])
        if not goal.shape[-1] == all_rollouts[0].shape[-1]:  # only append very final goal once
            if goals[opt_z_idx].shape == outputs[-1].shape:
                outputs.append(goals[opt_z_idx])
            else:
                outputs.append(goals[opt_z_idx][0].transpose(2, 0, 1))  # for image-based CEM
        output_rollout = np.stack(outputs)
        output_cost = total_cost[opt_z_idx]

        # remove children for all but optimal latent, indicate layer is optimized for subsequent optimization passes
        self._children = [c[:1] for c in self._children]
        self._n_samples = 1
        self._is_optimized = True

        # (optional) log all options to output for debugging, here in first layer of 8-layer GCP only
        # if self._depth == 8:
        #     self._log_all_subgoal_plans(starts, subgoals, goals, to_cost, from_cost)
        return output_rollout, output_cost

    def _recurse_optimization(self, all_rollouts, goal):
        """Splits sequence around subgoal and optimizes both parts independently.
           Handles edge cases of sub-trajetories too short for further recursion."""
        per_latent_rollouts = np.array_split(all_rollouts, self._n_samples)
        best_costs, best_rollouts = [], []
        for child_left, child_right, rollouts in zip(self._children[0], self._children[1], per_latent_rollouts):
            rollouts = [r for r in rollouts]  # convert from array of arrays to list

            # filter too short rollouts that don't need hierarchical expansion, replace with dummy
            short_rollouts = []
            for r_idx, r in enumerate(rollouts):
                if r.shape[0] < 3:
                    short_rollouts.append(r)
                    rollouts[r_idx] = self._make_dummy_seq(r[0])

            # expand hierarchically
            subgoal_inds = [int(np.floor(r.shape[0] / 2)) for r in rollouts]
            subgoal = rollouts[0][subgoal_inds[0]]  # across batch dimension all of the subgoals are identical
            best_rollout_left, best_cost_left = \
                child_left.optimize([r[:si] for r, si in zip(rollouts, subgoal_inds)], subgoal)
            best_rollout_right, best_cost_right = \
                child_right.optimize([r[si:] for r, si in zip(rollouts, subgoal_inds)], goal)
            best_rollout = np.concatenate([best_rollout_left, best_rollout_right])
            best_cost = best_cost_left + best_cost_right

            # check whether too short trajectories are better, if so: replace results
            if short_rollouts:
                best_rollout_short, best_cost_short, _ = self._best_of_n_segments(short_rollouts, goal, self._ll_cost_fcn)
                if best_cost_short < best_cost or np.isnan(best_cost):
                    best_rollout, best_cost = best_rollout_short, best_cost_short

            # dump best results for this latent
            best_rollouts.append(best_rollout); best_costs.append(best_cost)

        best_cost_idx = np.argmin(np.array(best_costs))
        best_rollout, best_cost = best_rollouts[best_cost_idx], best_costs[best_cost_idx]
        return best_rollout, best_cost

    def _sample(self):
        return np.random.normal(loc=self.mean, scale=self.std, size=(self._n_samples, self._n_latents, self._latent_dim))

    def _best_of_n_segments(self, all_rollouts, goal, cost_fcn):
        """Applies dense cost function to segment samples, returns min-cost segment + min-cost + idx."""
        all_rollouts_opt, goal_opt = self._prep_segment_opt_inputs(all_rollouts, goal)
        cost = cost_fcn(all_rollouts_opt, goal_opt)
        best_cost_idx = np.argmin(cost)
        return self._split_state_rollout(all_rollouts).state_rollout[best_cost_idx], cost[best_cost_idx], best_cost_idx

    def _prep_segment_opt_inputs(self, all_rollouts, goal):
        """Splits off input to cost function from combined inputs (rollouts are concat of both state and latent)"""
        rollouts = self._split_state_rollout(all_rollouts).state_rollout
        state_goal = self._split_state_rollout([goal]).state_rollout[0] if goal.shape[-1] == all_rollouts[0].shape[-1] \
            else goal
        return rollouts, state_goal

    def _split_state_rollout(self, rollouts):
        """Splits off latents from states in joined rollouts."""
        return listdict2dictlist([AttrDict(state_rollout=r[..., :-self._subgoal_cost_fcn.input_dim],
                                           latent_rollout=r[..., -self._subgoal_cost_fcn.input_dim:]) for r in rollouts])

    @staticmethod
    def _make_dummy_seq(reference_array):
        return np.stack([np.ones_like(reference_array) * float("inf"),  # fill with dummy w/ max cost
                         np.zeros_like(reference_array),
                         np.ones_like(reference_array) * float("inf")])

    @property
    def fully_optimized(self):
        if self._children is not None:
            return not (not self._is_optimized or
                        np.any([not c.fully_optimized for c in self._children[0]]) or
                        np.any([not c.fully_optimized for c in self._children[1]]))
        else:
            return self._is_optimized

    def _log_all_subgoal_plans(self, starts, subgoals, goals, to_cost, from_cost):
        for s, sg, g, tc, fc in zip(starts, subgoals, goals, to_cost, from_cost):
           r = np.stack([s, sg, g])
           c = tc + fc
           self._log_subgoal_plan(r, c)

    def _log_subgoal_plan(self, rollout, cost):
        from gcp.planning.infra.envs.miniworld_env.multiroom3d.multiroom3d_env import Multiroom3dEnv
        import time
        import cv2, os
        if self._dummy_env is None:
            self._dummy_env = Multiroom3dEnv({'n_rooms': 25}, no_env=True)
        im = self._dummy_env.render_top_down(rollout)
        name = "subgoal_{}_{}.png".format(cost, time.time())
        cv2.imwrite(os.path.join("/parent/tmp", name), im*255.)


class ImageHierarchicalTreeLatentOptimizer(HierarchicalTreeLatentOptimizer):
    def _prep_segment_opt_inputs(self, all_rollouts, goal):
        rollouts = self._split_state_rollout(all_rollouts).latent_rollout
        if len(goal.shape) > 2:     # in case we dont have goal encoding use final state of rollout
            state_goal = [r[-1:] for r in rollouts]
        else:
            state_goal = [self._split_state_rollout([goal[None]]).latent_rollout[0] for _ in rollouts]
        return rollouts, state_goal

    def _split_state_rollout(self, rollouts):
        """Splits off latents from states in joined rollouts."""
        def reshape_to_image(flat):
            if len(flat.shape) != 2:
                import pdb; pdb.set_trace()
            assert len(flat.shape) == 2
            res = int(np.sqrt(flat.shape[1] / 3))   # assumes 3-channel image
            return flat.reshape(flat.shape[0], 3, res, res)
        return listdict2dictlist([AttrDict(state_rollout=reshape_to_image(r[..., :-self._subgoal_cost_fcn.input_dim]),
                                           latent_rollout=r[..., -self._subgoal_cost_fcn.input_dim:]) for r in rollouts])

    def _log_all_subgoal_plans(self, starts, subgoals, goals, to_cost, from_cost):
        import cv2, time, os
        img_stack = []
        for s, sg, g, tc, fc in zip(starts, subgoals, goals, to_cost, from_cost):
            if len(g.shape) == 4:
                g = g[0].transpose(2, 0, 1) * 2 - 1
            img_strip = (np.concatenate((s, sg, g), axis=2).transpose(1, 2, 0) + 1) / 2
            img_strip = cv2.resize(img_strip[..., ::-1], (0,0), fx=4, fy=4)
            img_strip = cv2.putText(img_strip, "{:.2f}".format(float(fc+tc)),
                                    (10,img_strip.shape[0]-10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (0,0,0), 2)
            img_strip = cv2.putText(img_strip, "{:.2f}".format(float(fc+tc)),
                                    (10,10), cv2.FONT_HERSHEY_SIMPLEX, 0.5, (1.0, 1.0, 1.0), 2)
            img_stack.append(img_strip)
        name = "subgoals_{}.png".format(time.time())
        cv2.imwrite(os.path.join("/parent/tmp", name), np.concatenate(img_stack, axis=0) * 255)
<file_sep>import numpy as np
import torch
import torch.nn as nn

import gcp.prediction.models.adaptive_binding.adaptive
import gcp.prediction.models.tree.frame_binding
from blox import AttrDict
from blox.basic_types import subdict
from blox.torch.dist import safe_entropy, ProbabilisticModel
from blox.torch.losses import PenaltyLoss
from blox.torch.subnetworks import GeneralizedPredictorModel, Predictor
from gcp.prediction.models.tree.tree_lstm import build_tree_lstm
from blox.torch.variational import setup_variational_inference
from gcp.prediction.models.adaptive_binding.attentive_inference import AttentiveInference
from gcp.prediction.utils.tree_utils import SubgoalTreeLayer
from gcp.prediction.models.tree.inference import Inference


class TreeModule(nn.Module, ProbabilisticModel):
    def __init__(self, hp, decoder):
        nn.Module.__init__(self)
        ProbabilisticModel.__init__(self)
        self._hp = hp
        self.decoder = decoder

        self.build_network()

    def build_network(self):
        hp = self._hp

        q, self.prior = setup_variational_inference(self._hp, self._hp.nz_enc, self._hp.nz_enc * 2)
        if self._hp.attentive_inference:
            self.inference = AttentiveInference(self._hp, q)
        else:
            self.inference = Inference(self._hp, q)

        # todo clean this up with subclassing?
        pred_inp_dim = hp.nz_enc * 2 + hp.nz_vae
        if self._hp.context_every_step:
            pred_inp_dim = pred_inp_dim + hp.nz_enc * 2
            
        if hp.tree_lstm:
            self.subgoal_pred, self.lstm_initializer = build_tree_lstm(hp, pred_inp_dim, hp.nz_enc)
        else:
            self.subgoal_pred = GeneralizedPredictorModel(hp, input_dim=pred_inp_dim, output_dims=[hp.nz_enc],
                                                          activations=[None])

        self.build_binding()
        
        if self._hp.regress_index:
            self.index_predictor = Predictor(
                self._hp, self._hp.nz_enc * 2, self._hp.max_seq_len, detached=False, spatial=False)

    def build_binding(self):
        # TODO this has to be specified with classes, not strings
        hp = self._hp
        
        if self._hp.matching_type == 'balanced':
            binding_class = gcp.prediction.models.tree.frame_binding.BalancedBinding
        elif 'dtw' in self._hp.matching_type:
            binding_class = gcp.prediction.models.adaptive_binding.adaptive.AdaptiveBinding
        else:
            raise NotImplementedError
        
        self.binding = binding_class(hp, self.decoder)

    def produce_subgoal(self, inputs, layerwise_inputs, start_ind, end_ind, left_parent, right_parent, depth=None):
        """
        Divides the subsequence by producing a subgoal inside it.
         This function represents one step of recursion of the model
        """
        subgoal = AttrDict()

        e_l = left_parent.e_g_prime
        e_r = right_parent.e_g_prime

        subgoal.p_z = self.prior(e_l, e_r)

        if 'z' in layerwise_inputs:
            z = layerwise_inputs.z
            if self._hp.prior_type == 'learned':    # reparametrize if learned prior is used
                z = subgoal.p_z.reparametrize(z)
        elif self._sample_prior:
            z = subgoal.p_z.sample()
        else:
            ## Inference
            if self._hp.attentive_inference:
                subgoal.update(self.inference(inputs, e_l, e_r, start_ind, end_ind))
            else:
                subgoal.match_timesteps = self.binding.comp_timestep(left_parent.match_timesteps,
                                                                     right_parent.match_timesteps)
                subgoal.update(self.inference(inputs, e_l, e_r, start_ind, end_ind, subgoal.match_timesteps.float()))
                
            z = subgoal.q_z.sample()

        ## Predict the next node
        pred_input = [e_l, e_r, z]
        if self._hp.context_every_step:
            mult = int(z.shape[0] / inputs.e_0.shape[0])
            pred_input += [inputs.e_0.repeat_interleave(mult, 0),
                           inputs.e_g.repeat_interleave(mult, 0)]
        
        if self._hp.tree_lstm:
            if left_parent.hidden_state is None and right_parent.hidden_state is None:
                left_parent.hidden_state, right_parent.hidden_state = self.lstm_initializer(e_l, e_r, z)
                
            subgoal.hidden_state, subgoal.e_g_prime = \
                self.subgoal_pred(left_parent.hidden_state, right_parent.hidden_state, *pred_input)
        else:
            subgoal.e_g_prime_preact = self.subgoal_pred(*pred_input)
            subgoal.e_g_prime = torch.tanh(subgoal.e_g_prime_preact)

        subgoal.ind = (start_ind + end_ind) / 2     # gets overwritten w/ argmax of matching at training time (in loss)
        return subgoal, left_parent, right_parent

    def loss(self, inputs, outputs):
        if outputs.tree.depth == 0:
            return {}

        losses = AttrDict()

        losses.update(self.get_node_loss(inputs, outputs))

        # Explaining loss
        losses.update(self.binding.loss(inputs, outputs))
        
        # entropy penalty
        losses.entropy = PenaltyLoss(weight=self._hp.entropy_weight)(outputs.entropy)

        return losses

    def compute_matching(self, inputs, outputs):
        """ Match the tree nodes to ground truth and compute relevant values """
        tree = outputs.tree
    
        # compute matching distributions
        if 'gt_match_dists' in outputs:
            gt_match_dists = outputs.gt_match_dists
        else:
            gt_match_dists = self.binding.get_w(inputs.pad_mask, inputs, outputs, log=True)
            
        tree.bf.match_dist = outputs.gt_match_dists = gt_match_dists

        # compute additional vals
        outputs.entropy = safe_entropy(outputs.gt_match_dists, dim=-1)
        # probability of the node existing
        tree.bf.p_n = outputs.p_n = torch.sum(outputs.gt_match_dists, dim=2).clamp(0, 1)

    def get_node_loss(self, inputs, outputs):
        """ Reconstruction and KL divergence loss """
        losses = AttrDict()
        tree = outputs.tree
        
        losses.update(self.binding.reconstruction_loss(inputs, outputs))
        losses.update(self.inference.loss(tree.bf.q_z, tree.bf.p_z))

        return losses

    @staticmethod
    def _log_outputs(outputs, inputs, losses, step, log_images, phase, logger):
        if log_images:
            # Log layerwise loss
            layerwise_keys = ['dense_img_rec', 'kl'] & losses.keys()
            for name, loss in subdict(losses, layerwise_keys).items():
                if len(loss.error_mat.shape) > 2:   # reduce to two dimensions
                    loss.error_mat = loss.error_mat.mean([i for i in range(len(loss.error_mat.shape))][2:])
                layerwise_loss = SubgoalTreeLayer.split_by_layer_bf(loss.error_mat, dim=1)
                layerwise_loss = torch.tensor([l[l != 0].mean() for l in layerwise_loss])
                logger.log_graph(layerwise_loss, '{}_{}'.format(name, 'loss_layerwise'), step, phase)
<file_sep>import torch

from blox import AttrDict, batch_apply, rmap
from blox.tensor.ops import batchwise_index
from blox.torch.losses import BCELogitsLoss
from blox.torch.modules import DummyModule
from blox.torch.ops import make_one_hot, broadcast_final
from blox.torch.subnetworks import Predictor


class BaseBinding(DummyModule):
    def __init__(self, hp, decoder=None):
        super().__init__()
        self._hp = hp
        self.decoder = decoder

        self.build_network()

    def get_init_inds(self, inputs):
        return inputs.start_ind.float()[:, None], inputs.end_ind.float()[:, None]

    def apply_tree(self, tree, inputs):
        # recursive_add_dim = make_recursive(lambda x: add_n_dims(x, n=1, dim=1))
        start_ind, end_ind = self.get_init_inds(inputs)
        tree.apply_fn(
            {}, fn=self, left_parents=AttrDict(timesteps=start_ind), right_parents=AttrDict(timesteps=end_ind))
    
    def get_matched_sequence(self, tree, key):
        latents = tree.bf[key]
        indices = tree.bf.match_dist.argmax(1)
        # Two-dimensional indexing
        matched_sequence = rmap(lambda x: batchwise_index(x, indices), latents)
    
        return matched_sequence


class BalancedBinding(BaseBinding):
    # TODO almost all of this class is no longer needed I think
    def build_network(self):
        self.existence_predictor = Predictor(self._hp, self._hp.nz_enc, 1, spatial=False)
    
    def __call__(self, inputs, subgoal, left_parent, right_parent):
        timesteps = self.comp_timestep(left_parent.timesteps, right_parent.timesteps)
        c_n_prime = make_one_hot(timesteps.long(), self._hp.max_seq_len)

        # TODO implement the alternative of not doing this. Then would need to renormalize
        c_n_prime[left_parent.timesteps.long() == timesteps, :] = 0
        c_n_prime[right_parent.timesteps.long() == timesteps, :] = 0
        
        return AttrDict(timesteps=timesteps.long(), c_n_prime=c_n_prime.float())

    @staticmethod
    def comp_timestep(t_l, t_r, *unused_args):
        return (t_l + t_r) / 2
    
    def get_w(self, pad_mask, inputs, outputs, log=False):
        """ Match to the middle frame between parents """
        self.apply_tree(outputs.tree, inputs)
        match_dists = outputs.tree.bf.c_n_prime
        return match_dists

    def get_init_inds(self, outputs):
        # TODO this actually mostly gets passed 'inputs'
        return torch.zeros_like(outputs.end_ind[:, None]) - 1, \
               outputs.end_ind[:, None] + 1

    def prune_sequence(self, inputs, outputs, key='images'):
        seq = getattr(outputs.tree.df, key)
        latent_seq = outputs.tree.df.e_g_prime
    
        existence = batch_apply(self.existence_predictor, latent_seq)[..., 0]
        outputs.existence_predictor = AttrDict(existence=existence)
    
        existing_frames = torch.sigmoid(existence) > 0.5
        existing_frames[:, 0] = 1
        pruned_seq = [seq[i][existing_frames[i]] for i in range(seq.shape[0])]
    
        return pruned_seq

    def loss(self, inputs, outputs):
        losses = AttrDict()
    
        if 'existence_predictor' in outputs:
            losses.existence_predictor = BCELogitsLoss()(
                outputs.existence_predictor.existence, outputs.tree.df.match_dist.sum(2).float())
        return losses

    def reconstruction_loss(self, inputs, outputs, weights=1):
        """ Balanced tree can have a simpler loss version which doesn't use cdist """
        tree = outputs.tree
    
        estimates = self.get_matched_sequence(tree, 'distr')
        outputs.soft_matched_estimates = self.get_matched_sequence(tree, 'images')
        targets = inputs.traj_seq

        weights = broadcast_final(weights * inputs.pad_mask, targets)
        losses = self.decoder.nll(estimates, targets, weights, log_error_arr=True)
    
        return losses
<file_sep>import torch.nn as nn
import torch
import torch.nn.functional as F

from blox import AttrDict
from blox.torch.ops import make_one_hot, apply_linear
from blox.torch.subnetworks import Predictor, MultiheadAttention
from gcp.prediction.models.tree.inference import Inference


class AttentiveInference(Inference):
    def __init__(self, hp, q):
        super().__init__(hp, q)
        self.attention = Attention(hp)
    
    def forward(self, inputs, e_l, e_r, start_ind, end_ind, timestep=None):
        assert timestep is None
        
        output = AttrDict()
        if self.deterministic:
            output.q_z = self.q(e_l)
            return output
        
        values = inputs.inf_enc_seq
        keys = inputs.inf_enc_key_seq
        
        # Get (initial) attention key
        query_input = [e_l, e_r]
        
        e_tilde, output.gamma = self.attention(values, keys, query_input, start_ind, end_ind, inputs)
        output.q_z = self.q(e_l, e_r, e_tilde)
        return output


class Attention(nn.Module):
    def __init__(self, hp):
        super().__init__()
        self._hp = hp
        time_cond_length = self._hp.max_seq_len if self._hp.one_hot_attn_time_cond else 1
        input_size = hp.nz_enc * 2
        self.query_net = Predictor(hp, input_size, hp.nz_attn_key)
        self.attention_layers = nn.ModuleList([MultiheadAttention(hp) for _ in range(hp.n_attention_layers)])
        self.predictor_layers = nn.ModuleList([Predictor(hp, hp.nz_enc, hp.nz_attn_key, num_layers=2)
                                               for _ in range(hp.n_attention_layers)])
        self.out = nn.Linear(hp.nz_enc, hp.nz_enc)

    def forward(self, values, keys, query_input, start_ind, end_ind, inputs):
        """
        Performs multi-layered, multi-headed attention.

        Note: the query can have a different batch size from the values/keys. In that case, the query is interpreted as
        multiple queries, i.e. the values are tiled to match the query tensor size.
        
        :param values: tensor batch x length x dim_v
        :param keys: tensor batch x length x dim_k
        :param query_input: input to the query network, batch2 x dim_k
        :param start_ind:
        :param end_ind:
        :param inputs:
        :param timestep: specify the timestep of the attention directly. tensor batch2 x 1
        :param attention_weights:
        :return:
        """

        query = self.query_net(*query_input)
        s_ind, e_ind = (torch.floor(start_ind), torch.ceil(end_ind)) if self._hp.mask_inf_attention \
                                                                     else (inputs.start_ind, inputs.end_ind)
        
        # Reshape values, keys, inputs if not enough dimensions
        mult = int(query.shape[0] / keys.shape[0])
        tile = lambda x: x[:, None][:, [0] * mult].reshape((-1,) + x.shape[1:])
        values = tile(values)
        keys = tile(keys)
        s_ind = tile(s_ind)
        e_ind = tile(e_ind)
        
        # Attend
        norm_shape_k = query.shape[1:]
        norm_shape_v = values.shape[2:]
        raw_attn_output, att_weights = None, None
        for attention, predictor in zip(self.attention_layers, self.predictor_layers):
            raw_attn_output, att_weights = attention(query, keys, values, s_ind, e_ind)
            x = F.layer_norm(raw_attn_output, norm_shape_v)
            query = F.layer_norm(predictor(x) + query, norm_shape_k)  # skip connections around attention and predictor

        return apply_linear(self.out, raw_attn_output, dim=1), att_weights     # output non-normalized output of final attention layer
<file_sep>from contextlib import contextmanager

import numpy as np
import torch

from blox import AttrDict
from blox.tensor.ops import broadcast_final, batch_apply, batchwise_index, remove_spatial
from blox.torch.encoder_decoder import Encoder
from blox.torch.losses import L2Loss
from blox.torch.subnetworks import Predictor
from gcp.prediction.models.auxilliary_models.base_model import BaseModel
from gcp.prediction.training.checkpoint_handler import CheckpointHandler


class InverseModel(BaseModel):
    def __init__(self, params, logger):
        super().__init__(logger)
        self._hp = self._default_hparams()
        self.override_defaults(params)  # override defaults with config file
        self.postprocess_params()
        assert self._hp.n_actions != -1     # make sure action dimensionality was overridden
        self.build_network()

        # load only the encoder params during training
        if self._hp.enc_params_checkpoint is not None:
            assert self._hp.build_encoder   # provided checkpoint but did not build encoder
            self._load_weights([
                (self.encoder, 'encoder', self._hp.enc_params_checkpoint),
            ])
        self.detach_enc = not self._hp.finetune_enc

    def _default_hparams(self):
        # put new parameters in here:
        default_dict = {
            'ngf': 4,  # number of feature maps in shallowest level
            'nz_enc': 128,  # number of dimensions in encoder-latent space
            'nz_mid': 128,  # number of hidden units in fully connected layer
            'n_processing_layers': 3,  # Number of layers in MLPs
            'temp_dist': 1,  # sample temporal distances between 1 and temp_dist, regress only first action
            'enc_params_checkpoint': None,   # specify pretrained encoder weights to load for training
            'take_first_tstep': False,   # take only first and second time step, no shuffling.
            'use_states': False,
            'aggregate_actions': False,           # when taking two images that are more than one step apart sum the actions along that
            'pred_states': False,
            'finetune_enc': False,
            'checkpt_path': None,
            'build_encoder': True,      # if False, does not build an encoder, assumes that inputs are encoded from model
            'add_lstm_state_enc': False,    # if True, expects lstm state as additional encoded input
            'log_topdown_maze': False,
            'train_full_seq': False,
            'train_im0_enc': True,  # If True, the first frame latent is passed in as `enc_traj_seq`
        }

        # loss weights
        default_dict.update({
            'action_rec_weight': 1.0,
            'state_rec_weight': 1.0,
        })

        # misc params
        default_dict.update({
            'use_skips': False,
            'dense_rec_type': None,
            'device': None,
            'randomize_length': False,
        })

        # add new params to parent params
        parent_params = super()._default_hparams()
        for k in default_dict.keys():
            parent_params.add_hparam(k, default_dict[k])
        return parent_params
    
    def build_network(self, build_encoder=True):
        if self._hp.build_encoder:
            self.encoder = Encoder(self._hp)
        input_sz = self._hp.nz_enc * 3 if self._hp.add_lstm_state_enc else self._hp.nz_enc * 2
        self.action_pred = Predictor(self._hp, input_sz, self._hp.n_actions)

    def sample_offsets(self, end_ind):
        """
         # sample temporal distances between 1 and temp_dist, regress only first action
        :return:  None, call by reference
        """
        bs = end_ind.shape[0]
        if self._hp.take_first_tstep:
            t0 = torch.zeros(bs, device=self._hp.device).long()
            t1 = torch.ones_like(t0)
        else:
            t0 = np.zeros(bs)
            for b in range(bs):
                assert end_ind[b].cpu().numpy() >= self._hp.temp_dist
                t0[b] = np.random.randint(0, end_ind[b].cpu().numpy() - self._hp.temp_dist + 1, 1)
            delta_t = np.random.randint(1, self._hp.temp_dist + 1, bs)
            t1 = t0 + delta_t
            t0 = torch.tensor(t0, device=self._hp.device, dtype=torch.long)
            t1 = torch.tensor(t1, device=self._hp.device, dtype=torch.long)
        return t0, t1

    def index_input(self, input, t, aggregate=False, t1=None):
        if aggregate:
            assert t1 is not None       # need end time step for aggregation
            selected = torch.zeros_like(input[:, 0])
            for b in range(input.shape[0]):
                selected[b] = torch.sum(input[b, t[b]:t1[b]], dim=0)
        else:
            selected = batchwise_index(input, t)
        return selected

    def full_seq_forward(self, inputs):
        if 'model_enc_seq' in inputs:
            enc_seq_1 = inputs.model_enc_seq[:, 1:]
            if self._hp.train_im0_enc and 'enc_traj_seq' in inputs:
                enc_seq_0 = inputs.enc_traj_seq.reshape(inputs.enc_traj_seq.shape[:2] + (self._hp.nz_enc,))[:, :-1]
                enc_seq_0 = enc_seq_0[:, :enc_seq_1.shape[1]]
            else:
                enc_seq_0 = inputs.model_enc_seq[:, :-1]
        else:
            enc_seq = batch_apply(self.encoder, inputs.traj_seq)
            enc_seq_0, enc_seq_1 = enc_seq[:, :-1], enc_seq[:, 1:]

        if self.detach_enc:
            enc_seq_0 = enc_seq_0.detach()
            enc_seq_1 = enc_seq_1.detach()

        # TODO quite sure the concatenation is automatic
        actions_pred = batch_apply(self.action_pred, torch.cat([enc_seq_0, enc_seq_1], dim=2))
       
        output = AttrDict()
        output.actions = actions_pred  #remove_spatial(actions_pred)
        if 'actions' in inputs:
            output.action_targets = inputs.actions
            output.pad_mask = inputs.pad_mask
        return output

    def forward(self, inputs, full_seq=None):
        """
        forward pass at training time
        :arg full_seq: if True, outputs actions for the full sequence, expects input encodings
        """
        if full_seq is None:
            full_seq = self._hp.train_full_seq
        
        if full_seq:
            return self.full_seq_forward(inputs)

        t0, t1 = self.sample_offsets(inputs.norep_end_ind if 'norep_end_ind' in inputs else inputs.end_ind)
        im0 = self.index_input(inputs.traj_seq, t0)
        im1 = self.index_input(inputs.traj_seq, t1)
        if 'model_enc_seq' in inputs:
            if self._hp.train_im0_enc and 'enc_traj_seq' in inputs:
                enc_im0 = self.index_input(inputs.enc_traj_seq, t0).reshape(inputs.enc_traj_seq.shape[:1] + (self._hp.nz_enc,))
            else:
                enc_im0 = self.index_input(inputs.model_enc_seq, t0)
            enc_im1 = self.index_input(inputs.model_enc_seq, t1)
        else:
            assert self._hp.build_encoder       # need encoder if no encoded latents are given
            enc_im0 = self.encoder.forward(im0)[0]
            enc_im1 = self.encoder.forward(im1)[0]

        if self.detach_enc:
            enc_im0 = enc_im0.detach()
            enc_im1 = enc_im1.detach()

        selected_actions = self.index_input(inputs.actions, t0, aggregate=self._hp.aggregate_actions, t1=t1)
        selected_states = self.index_input(inputs.traj_seq_states, t0)

        if self._hp.pred_states:
            actions_pred, states_pred = torch.split(self.action_pred(enc_im0, enc_im1), 2, 1)
        else:
            actions_pred = self.action_pred(enc_im0, enc_im1)

        output = AttrDict()
        output.actions = remove_spatial(actions_pred)
        output.action_targets = selected_actions
        output.state_targets = selected_states
        output.img_t0, output.img_t1 = im0, im1
        
        return output

    def loss(self, inputs, outputs, add_total=True):
        losses = AttrDict()

        # subgoal reconstruction loss
        n_action_output = outputs.actions.shape[1]
        loss_weights = broadcast_final(outputs.pad_mask[:, :n_action_output], inputs.actions) if 'pad_mask' in outputs else 1
        losses.action_reconst = L2Loss(self._hp.action_rec_weight)(outputs.actions, outputs.action_targets[:, :n_action_output], weights=loss_weights)
        if self._hp.pred_states:
            losses.state_reconst = L2Loss(self._hp.state_rec_weight)(outputs.states, outputs.state_targets)

        return losses

    def log_outputs(self, outputs, inputs, losses, step, log_images, phase):
        super()._log_losses(losses, step, log_images, phase)
        
        if 'actions' not in outputs:
            # TODO figure out why this happens
            return
        
        if log_images and len(inputs.traj_seq.shape) == 5:
            self._logger.log_pred_actions(outputs, inputs, 'pred_actions', step, phase)
        if self._hp.pred_states:
            self._logger.log_pred_states(outputs, inputs, 'pred_states', step, phase)

        if log_images:
            dataset = self._hp.dataset_class
            if len(outputs.actions.shape) == 3:
                actions = outputs.actions
            else:
                # Training, need to get the action sequence
                actions = self(inputs, full_seq=True).actions
    
            cum_action_traj = torch.cat((inputs.traj_seq_states[:, :1], actions), dim=1).cumsum(1)
            self._logger.log_dataset_specific_trajectory(outputs, inputs, "action_traj_topdown", step, phase, dataset,
                                                         predictions=cum_action_traj, end_inds=inputs.end_ind)
            
            cum_action_traj = torch.cat((inputs.traj_seq_states[:, :1], inputs.actions), dim=1).cumsum(1)
            self._logger.log_dataset_specific_trajectory(outputs, inputs, "action_traj_gt_topdown", step, phase, dataset,
                                                         predictions=cum_action_traj, end_inds=inputs.end_ind)

    def run_single(self, enc_latent_img0, model_latent_img1):
        """Runs inverse model on first input encoded by encoded and second input produced by model."""
        assert self._hp.train_im0_enc   # inv model needs to be trained from
        return remove_spatial(self.action_pred(enc_latent_img0, model_latent_img1))

    @contextmanager
    def val_mode(self, *args, **kwargs):
        yield


class TestTimeInverseModel(InverseModel):
    def __init__(self, params, logger):
        super().__init__(params, logger)
        if torch.cuda.is_available():
            self.cuda()
            self.device = torch.device('cuda')
        else:
            self.device = torch.device('cpu')

        assert self._hp.checkpt_path is not None
        weights_file = CheckpointHandler.get_resume_ckpt_file('latest', self._hp.checkpt_path)
        success = CheckpointHandler.load_weights(weights_file, self)
        if not success: raise ValueError("Could not load checkpoint from {}!".format(weights_file))

    def forward(self, inputs):
        for k in inputs:
            if not isinstance(inputs[k], torch.Tensor):
                inputs[k] = torch.Tensor(inputs[k])
            if not inputs[k].device == self.device:
                inputs[k] = inputs[k].to(self.device) 

        enc_im0 = self.encoder.forward(inputs['img_t0'])[0]
        enc_im1 = self.encoder.forward(inputs['img_t1'])[0]
        return remove_spatial(self.action_pred(enc_im0, enc_im1))


class FromStatesInverseModel(InverseModel):
    def __init__(self, params, logger):
        super().__init__(params, logger)

    def get_timesteps(self, inputs):
        """
         # sample temporal distances between 1 and temp_dist, regress only first action
        :return:  None, call by reference
        """

        t0 = np.zeros(self._hp.batch_size)
        for b in range(self._hp.batch_size):
            t0[b] = np.random.randint(0, abs(inputs.end_ind[b].cpu().numpy() - self._hp.temp_dist), 1)
        delta_t = np.random.randint(1, self._hp.temp_dist + 1, self._hp.batch_size)
        t1 = t0 + delta_t
        t0 = torch.tensor(t0, device=inputs.traj_seq_states.device, dtype=torch.long)
        t1 = torch.tensor(t1, device=inputs.traj_seq_states.device, dtype=torch.long)
        inputs.state_t0 = batchwise_index(inputs.traj_seq_states, t0)
        inputs.state_t1 = batchwise_index(inputs.traj_seq_states, t1)
        inputs.selected_action = batchwise_index(inputs.actions, t0)

    def build_network(self, build_encoder=True):
        self.action_pred = Predictor(self._hp, self._hp.state_dim*2, self._hp.n_actions, 3)

    def forward(self, inputs):
        self.get_timesteps(inputs)
        actions_pred = self.action_pred(inputs.state_t0[:,:,None, None], inputs.state_t1[:,:,None, None])
        output = AttrDict()
        output.actions = torch.squeeze(actions_pred)
        return output

    def log_outputs(self, outputs, inputs, losses, step, log_images, phase):
        super()._log_losses(losses, step, phase)


class EarlyFusionInverseModel(InverseModel):
    def __init__(self, params, logger):
        super().__init__(params, logger)

    def build_network(self, build_encoder=True):
        self._hp.input_nc = 6
        self.encoder = Encoder(self._hp)

        if self._hp.pred_states:
            outdim = self._hp.n_actions + self._hp.state_dim
        else:
            outdim = self._hp.n_actions
        self.action_pred = Predictor(self._hp, self._hp.nz_enc, outdim, 3)

    def forward(self, inputs):
        self.get_timesteps(inputs)
        enc = self.encoder.forward(torch.cat([inputs.img_t0, inputs.img_t1], dim=1))[0]
        output = AttrDict()
        out = self.action_pred(enc)
        if self._hp.pred_states:
            output.actions, output.states = torch.split(torch.squeeze(out), [2,2], 1)
        else:
            output.actions = torch.squeeze(out)
        return output



<file_sep>import cv2
import copy
import gym
import numpy as np
from blox import AttrDict
from gcp.planning.infra.envs.miniworld_env.base_miniworld_env import BaseMiniworldEnv
from gcp.planning.infra.envs.miniworld_env.utils.multiroom2d_layout import define_layout, draw_layout_overview, \
                                                                                        default_texture_dir
from gcp.planning.infra.envs.miniworld_env.utils.sampling_fcns import RoomSampler2d
import numbers


def fcn_apply(fcn, arg):
    return lambda: fcn(arg())


class Multiroom3dEnv(BaseMiniworldEnv):
    def __init__(self, hp, reset_state=None, no_env=False, crop_window=None):
        self._hp = self._default_hparams()
        for name, value in hp.items():
            print('setting param {} to value {}'.format(name, value))
            self._hp.set_hparam(name, value)
        super().__init__(self._hp)

        self._texture_dir = default_texture_dir()
        self._rooms_per_side = int(np.sqrt(self._hp.n_rooms))
        self._layout = define_layout(self._rooms_per_side, self._texture_dir)
        self._topdown_render_scale = 256       # table_size * scale = px size of output render img
        self._static_img_topdown = draw_layout_overview(self._rooms_per_side,
                                                        self._topdown_render_scale,
                                                        texture_dir=self._texture_dir)
        self._crop_window = crop_window
        if crop_window is not None:
            # top-down rendering will get cropped -> pad static background
            padded_bg = np.zeros((self._static_img_topdown.shape[0] + 2 * crop_window,
                                  self._static_img_topdown.shape[1] + 2 * crop_window, 3), dtype=self._static_img_topdown.dtype)
            padded_bg[crop_window:-crop_window, crop_window:-crop_window] = self._static_img_topdown
            self._static_img_topdown = padded_bg

        self._adim, self._sdim = 2, 3
        if not no_env:
            import gym_miniworld    # keep! important for env registration
            self.env = gym.make("MiniWorld-Multiroom3d-v0",
                            obs_height=self._hp.obs_height,
                            obs_width=self._hp.obs_width,
                            rooms_per_side=self._rooms_per_side,
                            doors=self._layout.doors,
                            heading_smoothing=self._hp.heading_smoothing,
                            layout_params=AttrDict(room_size=self._layout.room_size,
                                                   door_size=self._layout.door_size,
                                                   textures=self._layout.textures))

        # Define the sample_*_state method by looking up the function with the corresponding name
        self.state_sampler = RoomSampler2d(self._rooms_per_side)
        self.current_pos = None
        self.goal_pos = None
        self.prm_policy = None  # used to compute shortest distance between pos and goal

    def _default_hparams(self):
        default_dict = {
            'obs_height': 300,
            'obs_width': 400,
            'goal_pos': None,
            'init_pos': None,
            'n_rooms': 9,
            'heading_smoothing': 0.2,   # how much of new angle is taken into average
        }
        parent_params = super()._default_hparams()
        for k in default_dict.keys():
            parent_params.add_hparam(k, default_dict[k])
        return parent_params

    def reset(self, reset_state):
        super().reset()

        if reset_state is None:
            start_pos = self.env.mj2mw(self.state_sampler.sample(self._hp.init_pos))
            start_angle = 2 * np.pi * np.random.rand()
            goal_pos = self.env.mj2mw(self.state_sampler.sample(self._hp.goal_pos))
        else:
            start_pos = reset_state[:2]
            start_angle = reset_state[2]
            goal_pos = reset_state[-2:]

        reset_state = AttrDict(start_pos=start_pos,
                               start_angle=start_angle,
                               goal=goal_pos)

        img_obs = self.env.reset(reset_state)
        self.goal_pos = goal_pos
        qpos_full = np.concatenate((start_pos, np.array([start_angle])))

        obs = AttrDict(images=np.expand_dims(img_obs, axis=0),      # add camera dimension
                       qpos_full=qpos_full,
                       goal=goal_pos,
                       env_done=False,
                       state=np.concatenate((qpos_full, goal_pos)),
                       topdown_image=self.render_pos_top_down(qpos_full, self.goal_pos)
                        )
        self._post_step(start_pos)
        self._initial_shortest_dist = self.comp_shortest_dist(start_pos, goal_pos)
        return obs, reset_state

    def get_reset_from_obs(self, obs_dict):
        return obs_dict['state'][0]

    def get_goal_from_obs(self, obs_dict):
        self.goal = obs_dict['goal'][-1]
        return self.goal

    def step(self, action):
        img_obs, reward, done, agent_pos = self.env.step(action)
        obs = AttrDict(images=np.expand_dims(img_obs, axis=0),  # add camera dimension
                       qpos_full=agent_pos,
                       goal=self.goal_pos,
                       env_done=done,
                       state=np.concatenate((agent_pos, self.goal_pos)),
                       topdown_image=self.render_pos_top_down(agent_pos, self.goal_pos)
                       )
        self._post_step(agent_pos)
        return obs

    def _post_step(self, agent_pos):
        self.current_pos = agent_pos
        self.add_goal_dist(self.comp_shortest_dist(self.current_pos[:2], self.goal_pos))
        self._full_traj.append(agent_pos)

    def eval(self):
        self._final_shortest_dist = self.comp_shortest_dist(self.current_pos[:2], self.goal_pos)
        return super().eval()

    def comp_shortest_dist(self, p1, p2):
        """Uses PRM to get the shortest distance between two points within the maze."""
        if self.prm_policy is None:
            from gcp.planning.infra.policy.prm_policy.prm_policy import PrmPolicy
            self.prm_policy = PrmPolicy(None, AttrDict(n_samples_per_room=200), None, None, **self.env_policy_params())
        dist, _ = self.prm_policy.compute_shortest_path(p1, p2)
        return dist

    def env_policy_params(self):
        def transform_plan(state_plan, action_plan):
            state_plan = self.env.mj2mw(state_plan)
            action_plan = state_plan[:, 1:] - state_plan[:, :-1]
            return state_plan, action_plan

        conversion_fcns = AttrDict(transform_plan=transform_plan,
                                   env2prm=self.env.mw2mj,
                                   prm2env=self.env.mj2mw)
        return {'conversion_fcns': conversion_fcns, 'n_rooms': self._hp.n_rooms}

    def render_top_down(self, traj, background=None, goal=None, line_thickness=4, color=(1.0, 0, 0), mark_pts=False):
        """Renders a state trajectory in a top-down view."""
        if isinstance(color[0], numbers.Number):
            color = [color] * (traj.shape[0] - 1)
         
        img = self._static_img_topdown.copy() if background is None else background.copy()
        traj = traj.copy()  # very important!!!
        if goal is not None:
            goal = goal.copy()
            if traj.shape[1] == 5 or traj.shape[1] == 2: goal = goal[:2]; goal[1] *= -1
            if traj.max() > 1.0 or traj.min() < -1.0: goal = goal / 27.0
            goal = goal + 0.5 * self._layout.table_size
        if traj.shape[1] == 5 or traj.shape[1] == 2: traj = traj[:, :2]; traj[:, 1] *= -1
        if traj.max() > 1.0 or traj.min() < -1.0: traj = traj / 27.0    # scale from miniworld env to [-1...1]
        traj = traj + 0.5 * self._layout.table_size
        for i in range(traj.shape[0] - 1):
            cv2.line(img, (int(traj[i, 0] * self._topdown_render_scale),
                           img.shape[0] - int(traj[i, 1] * self._topdown_render_scale)),
                          (int(traj[i+1, 0] * self._topdown_render_scale),
                           img.shape[0] - int(traj[i+1, 1] * self._topdown_render_scale)),
                           color[i], line_thickness)
            if mark_pts and i > 0 and i < (traj.shape[0] - 2):
                cv2.line(img, (int(traj[i, 0] * self._topdown_render_scale),
                               img.shape[0] - int(traj[i, 1] * self._topdown_render_scale)),
                              (int(traj[i, 0] * self._topdown_render_scale),
                               img.shape[0] - int(traj[i, 1] * self._topdown_render_scale)),
                         (1.0, 0, 0), int(3*line_thickness))

        # print start+end position
        img = self.render_pos_top_down(traj[0], traj[-1], background=img, mirror_scale=False)
        if goal is not None:
            img = self.render_pos_top_down(traj[0], goal, background=img, mirror_scale=False, large_goal=True)
        return img

    def render_pos_top_down(self,
                            current_pose,
                            goal_pos,
                            background=None,
                            mirror_scale=True,
                            large_goal=False):
        """Renders a state trajectory in a top-down view."""
        img = self._static_img_topdown.copy() if background is None else background.copy()

        def convert_sim2topdown(pos, img_shape):
            pos = pos.copy()  # very important !!!!!
            if mirror_scale:
                pos[1] *= -1
                pos = pos / 27.0    # scale from miniworld env to [-1...1]
                pos = pos + 0.5 * self._layout.table_size
            return (int(pos[0] * self._topdown_render_scale), img_shape[0] - int(pos[1] * self._topdown_render_scale))

        curr_pos = convert_sim2topdown(current_pose, img.shape)
        goal_pos = convert_sim2topdown(goal_pos, img.shape)

        if self._crop_window is not None:
            # we need to correct for the too large size of img.shape above, therefore -2*self._crop_window
            curr_pos = (curr_pos[0] + self._crop_window, curr_pos[1] + self._crop_window - 2*self._crop_window)
            goal_pos = (goal_pos[0] + self._crop_window, goal_pos[1] + self._crop_window - 2*self._crop_window)

        cv2.line(img, curr_pos, curr_pos, (0.0, 0, 1.0), 10)
        cv2.line(img, goal_pos, goal_pos, (0.0, 1.0, 0), 10 if not large_goal else 20)

        if self._crop_window is not None:
            # catch rounding errors
            curr_pos = (max(self._crop_window, curr_pos[0]), max(self._crop_window, curr_pos[1]))
            lower, upper = np.asarray(curr_pos) - self._crop_window, np.asarray(curr_pos) + self._crop_window
            img = img[lower[1]:upper[1], lower[0]:upper[0]]

        return img

    @property
    def adim(self):
        return self._adim

    @property
    def sdim(self):
        return self._sdim


class TopdownMultiroom3dEnv(Multiroom3dEnv):
    """Image observations are rendered topdown in a window around the agent."""
    def __init__(self, hp, reset_state=None, no_env=False, crop_window=None):
        assert "crop_window" in hp      # need to specify the crop window for topdown rendering
        temp_hp = copy.deepcopy(hp)
        crop_window = temp_hp.pop("crop_window")
        super().__init__(temp_hp, reset_state, no_env, crop_window=crop_window)

    def reset(self, reset_state):
        obs, reset_state = super().reset(reset_state)
        obs.images = np.asarray(255*obs.topdown_image.copy(), dtype=np.uint8)[None]
        return obs, reset_state

    def step(self, action):
        obs = super().step(action)
        obs.images = np.asarray(255*obs.topdown_image.copy(), dtype=np.uint8)[None]
        return obs

<file_sep>import pickle as pkl
import numpy as np
import os

from robosuite.utils.mjcf_utils import postprocess_model_xml
from gcp.planning.infra.agent.benchmarking_agent import BenchmarkAgent


class BenchmarkAgentLoadHDF5(BenchmarkAgent):
    def __init__(self, hyperparams, start_goal_list=None):
        super.__init__(hyperparams, start_goal_list)

    def _load_raw_data(self, itr):
        """
        doing the reverse of save_raw_data
        :param itr:
        :return:
        """
        traj = self._start_goal_confs + 'itr.hdf5'
        self._demo_images = demo_images.astype(np.float32)/255.
        self._goal_image = self._demo_images[-1]

        with open('{}/obs_dict.pkl'.format(traj), 'rb') as file:
            obs_dict.update(pkl.load(file))


        self._goal = self.env.get_goal_from_obs(obs_dict)
        reset_state = self.get_reset_state(obs_dict)

        if os.path.exists(traj  + '/robosuite.xml'):
            with open(traj  + '/robosuite.xml', "r") as model_f:
                model_xml = model_f.read()

            xml = postprocess_model_xml(model_xml)
            reset_state['robosuite_xml'] = xml

        return reset_state

    def get_reset_state(self, obs_dict):
        return self.env.get_reset_from_obs(obs_dict)


<file_sep>import numpy as np

from blox import AttrDict
from gcp.planning.tree_optimizer import HierarchicalTreeLatentOptimizer, ImageHierarchicalTreeLatentOptimizer


class CEMSampler:
    """Defines interface for sampler used in CEM optimization loop."""
    def __init__(self, clip_val, n_steps, action_dim, initial_std):
        self._clip_val = clip_val
        self._n_steps = n_steps
        self._action_dim = action_dim
        self._initial_std = initial_std
        self.init()

    def init(self):
        """Initialize the sampling distributions."""
        raise NotImplementedError

    def sample(self, n_samples):
        """Sample n_samples from the sampling distributions."""
        raise NotImplementedError

    def fit(self, data, scores):
        """Refits distributions to data."""
        raise NotImplementedError

    def get_dists(self):
        """Returns a representation of the current sampling distributions."""
        raise NotImplementedError


class FlatCEMSampler(CEMSampler):
    """Samples flat arrays from Gaussian distributions."""
    def init(self):
        """Initialize the sampling distributions."""
        self.mean = np.zeros((self._n_steps, self._action_dim))
        self.std = self._initial_std * np.ones((self._n_steps, self._action_dim))

    def sample(self, n_samples):
        raw_actions = np.random.normal(loc=self.mean, scale=self.std, size=(n_samples, self._n_steps, self._action_dim))
        return np.clip(raw_actions, -self._clip_val, self._clip_val)

    def fit(self, data, scores):
        self.mean = np.mean(data, axis=0)
        self.std = np.std(data, axis=0)

    def get_dists(self):
        return AttrDict(mean=self.mean, std=self.std)


class PDDMSampler(FlatCEMSampler):
    """Samples correlated noise, uses path integral formulation to fit it."""
    BETA = 0.5      # noise correlation factor
    GAMMA = 1.0     # reward weighting factor

    def sample(self, n_samples):
        noise = np.random.normal(loc=np.zeros_like(self.mean), scale=self.std,
                                 size=(n_samples, self._n_steps, self._action_dim))
        correlated_noise, n_i = [], np.zeros((n_samples, self._action_dim))
        for i in range(noise.shape[1]):
            u_i = noise[:, i]
            n_i = self.BETA * u_i + (1 - self.BETA) * n_i
            correlated_noise.append(n_i)
        correlated_noise = np.stack(correlated_noise, axis=1)
        return np.clip(correlated_noise + self.mean[None], -self._clip_val, self._clip_val)

    def fit(self, actions, scores):
        """Assumes that scores are better the lower (ie cost function output)."""
        self.mean = np.sum(actions * np.exp(-self.GAMMA * scores)[:, None, None], axis=0) \
                    / np.sum(np.exp(-self.GAMMA * scores))


class SimpleTreeCEMSampler(FlatCEMSampler):
    """CEM sampler for tree-GCPs that optimizes all levels at once with same number of samples
       (ie like flat CEM sampler)."""
    def __init__(self, *args, n_level_hierarchy, **kwargs):
        self._n_layer_hierarchy = n_level_hierarchy
        super().__init__(*args)
        self._n_steps = 2**n_level_hierarchy - 1


class HierarchicalTreeCEMSampler(SimpleTreeCEMSampler):
    """Tree-GCP CEM sampler that optimizes the layers of the hierarchy sequentially, starting from the top."""
    def __init__(self, *args, sampling_rates_per_layer, subgoal_cost_fcn, ll_cost_fcn, n_ll_samples, **kwargs):
        self._sampling_rates_per_layer = sampling_rates_per_layer
        self._subgoal_cost_fcn = subgoal_cost_fcn
        self._ll_cost_fcn = ll_cost_fcn
        self._n_ll_samples = n_ll_samples
        super().__init__(*args, **kwargs)
        assert self._n_layer_hierarchy >= len(sampling_rates_per_layer)     # not enough layers in tree

    def init(self):
        self._optimizer = HierarchicalTreeLatentOptimizer(self._action_dim,
                                                          self._sampling_rates_per_layer.copy(),
                                                          self._n_layer_hierarchy,
                                                          self._subgoal_cost_fcn,
                                                          self._ll_cost_fcn,
                                                          self._n_ll_samples)

    def sample(self, n_samples):
        raw_actions = self._optimizer.sample()
        return np.clip(raw_actions, -self._clip_val, self._clip_val)

    def optimize(self, rollouts, goal):
        best_rollout, best_cost = self._optimizer.optimize(rollouts, goal)
        if (best_rollout[-1] != goal).any():    # this can happen if too few frames on right tree side
            best_rollout = np.concatenate((best_rollout, goal[None]))
        return [best_rollout], best_cost

    def fit(*args, **kwargs):
        """Does not currently support refitting distributions."""
        pass

    def get_dists(self):
        return AttrDict(mean=0., std=1.)    # dummy values

    @property
    def append_latent(self):
        return True  # we need latent rollouts to compute subgoal costs

    @property
    def fully_optimized(self):
        return self._optimizer.fully_optimized


class ImageHierarchicalTreeCEMSampler(HierarchicalTreeCEMSampler):
    """Hierarchical GCP-tree CEM sampler for image prediction GCPs."""
    def init(self):
        self._optimizer = ImageHierarchicalTreeLatentOptimizer(self._action_dim,
                                                               self._sampling_rates_per_layer.copy(),
                                                               self._n_layer_hierarchy,
                                                               self._subgoal_cost_fcn,
                                                               self._ll_cost_fcn,
                                                               self._n_ll_samples)

    def optimize(self, rollouts, goal):
        best_rollout, best_cost = self._optimizer.optimize(rollouts, goal)
        if (best_rollout[-1] != goal[0].transpose(2, 0, 1)).any():    # can happen if too few frames on right tree side
            best_rollout = np.concatenate((best_rollout, goal.transpose(0, 3, 1, 2)))
        if not hasattr(best_cost, "__len__"):
            best_cost = [best_cost]         # need to return array-shaped cost, no scalar
        return [best_rollout], best_cost
<file_sep>import glob
import itertools
import os

import cv2
import numpy as np
from blox import AttrDict
from gcp.planning.infra.agent.utils.hdf5_saver import HDF5SaverBase
from tqdm import tqdm


def process_frame(frame_in, frame_output_size=(64, 64)):
    """Standardizes input frame width and height, and removes dummy channels.
    """
    
    # NB: OpenCV takes new size as (X, Y), not (Y, X)!!!
    frame_out = cv2.resize(
        frame_in.astype(np.float32),
        (frame_output_size[1], frame_output_size[0])).astype(
        frame_in.dtype)
    
    if frame_out.shape[-1] == 1:
        # Chop off redundant dimensions
        frame_out = frame_out[..., 0]
    elif frame_out.shape[-1] == 3:
        # Convert OpenCV's BGR to RGB
        frame_out = frame_out[..., ::-1]
    else:
        raise ValueError("Unexpected frame shape!")
    
    return frame_out

def read_video(video_path, n_downsample=8):
    """Reads a video from a file and returns as an array."""
    assert os.path.isfile(video_path)
    cap = cv2.VideoCapture(video_path)
    
    all_frames = []
    i = 0
    while cap.isOpened():
        ret, frame = cap.read()
        if ret:
            if i % n_downsample == 0:
                all_frames.append(process_frame(frame))
        else:
            cap.release()
        i += 1
    
    return np.asarray(all_frames)


class H36Maker(HDF5SaverBase):
    def __init__(self, data_dir, folders, save_dir, traj_per_file, offset=0, split=(0.90, 0.05, 0.05), n_downsample=8):
        super().__init__(save_dir, traj_per_file, offset, split)
        
        data_folders = list([data_dir + '/' + folder + '/Videos' for folder in folders])
        listlist_names = list([glob.glob(folder + '/*') for folder in data_folders])
        self.filenames = list(itertools.chain(*listlist_names))
        self.max_seq_len = 1500
        self.n_downsample = n_downsample
    
    def get_traj(self, name):
        vid = read_video(name, self.n_downsample)
        print(vid.shape)
        # pad_mask = np.concatenate([np.ones(vid.shape[0]), np.zeros(self.max_seq_len - vid.shape[0])])
        pad_mask = np.ones(vid.shape[0])
        
        return AttrDict(images=vid, pad_mask=pad_mask)
    
    def make_phase(self, filenames, phase):
        traj_list = []
        
        for name in tqdm(filenames):
            traj = self.get_traj(name)
            traj_list.append(traj)

            if len(traj_list) == self.traj_per_file:
                self.save_hdf5(traj_list, phase)
                traj_list = []


if __name__ == '__main__':
    folders_list = [['S1', 'S5', 'S6', 'S7'], ['S8'], ['S9', 'S11']]
    phase_list = ['train', 'val', 'test']
    
    for i in range(3):
        maker = H36Maker('/parent/nfs/kun1/users/oleg/h36m', folders_list[i],
                         '/workspace/recplan_data/h36m_long', 10, split=(1, 0, 0), n_downsample=1)
        
        maker.make_phase(maker.filenames, phase_list[i])
<file_sep>from contextlib import contextmanager

import numpy as np
import torch

from blox import AttrDict
from blox.tensor.ops import batchwise_index
from blox.torch.losses import L2Loss
from blox.torch.subnetworks import Predictor
from gcp.prediction.models.auxilliary_models.base_model import BaseModel
from gcp.prediction.training.checkpoint_handler import CheckpointHandler


class CostModel(BaseModel):
    def __init__(self, params, logger):
        super().__init__(logger)
        self._hp = self._default_hparams()
        params.update({'use_convs': False})
        self.override_defaults(params)  # override defaults with config file
        self.postprocess_params()
        self.build_network()

        if self._hp.cost_fcn is not None:
            self._gt_cost_fcn = self._hp.cost_fcn(True)

    def _default_hparams(self):
        # put new parameters in here:
        default_dict = {
            'nz_enc': 128,  # number of dimensions in encoder-latent space
            'nz_mid': 128,  # number of hidden units in fully connected layer
            'n_processing_layers': 3,  # Number of layers in MLPs
            'checkpt_path': None,
            'load_epoch': None,  # checkpoint epoch which should be loaded, if 'none' loads latest
            'cost_fcn': None,   # holds ground truth cost function
            'use_path_dist_cost': False,   # if True, uses fast path distance cost computation, ignores cost_fcn
        }

        # misc params
        default_dict.update({
            'use_skips': False,
            'dense_rec_type': None,
            'device': None,
            'randomize_length': False,
        })

        # add new params to parent params
        parent_params = super()._default_hparams()
        for k in default_dict.keys():
            parent_params.add_hparam(k, default_dict[k])
        return parent_params
    
    def build_network(self, build_encoder=True):
        self.cost_pred = Predictor(self._hp, self._hp.nz_enc * 2, 1, detached=True)

    def forward(self, inputs):
        """Forward pass at training time."""
        # randomly sample start and end state, compute GT cost
        start, end, gt_cost = self._fast_path_dist_cost(inputs) if self._hp.use_path_dist_cost \
                                else self._general_cost(inputs)

        # compute cost estimate
        cost_pred = self.cost_pred(torch.cat([start, end], dim=-1))

        output = AttrDict(
            cost=cost_pred,
            cost_target=gt_cost,
        )
        return output

    def loss(self, inputs, outputs, add_total=True):
        return AttrDict(
            cost_estimation=L2Loss(1.0)(outputs.cost, outputs.cost_target)
        )

    def log_outputs(self, outputs, inputs, losses, step, log_images, phase):
        super()._log_losses(losses, step, log_images, phase)

    @contextmanager
    def val_mode(self):
        yield

    def _fast_path_dist_cost(self, inputs):
        """Vectorized computation of path distance cost."""
        # sample start goal indices
        batch_size = inputs.end_ind.shape[0]
        start_idx = torch.rand((batch_size,), device=inputs.end_ind.device) * (inputs.end_ind.float() - 1)
        end_idx = torch.rand((batch_size,), device=inputs.end_ind.device) * (inputs.end_ind.float() - (start_idx + 1)) + (start_idx + 1)
        start_idx, end_idx = start_idx.long(), end_idx.long()

        # get start goal latents
        start = batchwise_index(inputs.model_enc_seq, start_idx).detach()
        end = batchwise_index(inputs.model_enc_seq, end_idx).detach()

        # compute path distance cost
        cum_diff = torch.cumsum(torch.norm(inputs.traj_seq[:, 1:] - inputs.traj_seq[:, :-1], dim=-1), dim=1)
        cum_diff = torch.cat((torch.zeros((batch_size, 1), dtype=cum_diff.dtype, device=cum_diff.device),
                              cum_diff), dim=1)    # prepend 0
        gt_cost = batchwise_index(cum_diff, end_idx) - batchwise_index(cum_diff, start_idx)

        return start, end, gt_cost[:, None].detach()

    def _general_cost(self, inputs):
        """Computes cost with generic cost function, not vectorized."""
        batch_size = inputs.end_ind.shape[0]

        start, end, gt_cost = [], [], []
        for b in range(batch_size):
            start_idx = np.random.randint(0, inputs.end_ind[b].cpu().numpy(), 1)[0]
            end_idx = np.random.randint(start_idx + 1, inputs.end_ind[b].cpu().numpy() + 1, 1)[0]
            start.append(inputs.model_enc_seq[b, start_idx])
            end.append(inputs.model_enc_seq[b, end_idx])
            gt_cost.append(self._gt_cost_fcn([inputs.traj_seq[b, start_idx:end_idx+1].data.cpu().numpy()],
                                             inputs.traj_seq[b, end_idx].data.cpu().numpy()))
        start, end = torch.stack(start).detach(), torch.stack(end).detach()  # no gradients in encoder
        gt_cost = torch.tensor(np.stack(gt_cost), device=start.device, requires_grad=False).float()
        return start, end, gt_cost

    @property
    def input_dim(self):
        return self._hp.nz_enc


class TestTimeCostModel(CostModel):
    def __init__(self, params, logger):
        super().__init__(params, logger)
        if torch.cuda.is_available():
            self.cuda()
            self.device = torch.device('cuda')
        else:
            self.device = torch.device('cpu')

        assert self._hp.checkpt_path is not None
        load_epoch = self._hp.load_epoch if self._hp.load_epoch is not None else 'latest'
        weights_file = CheckpointHandler.get_resume_ckpt_file(load_epoch, self._hp.checkpt_path)
        success = CheckpointHandler.load_weights(weights_file, self, submodule_name='cost_mdl')
        if not success: raise ValueError("Could not load checkpoint from {}!".format(weights_file))

    def forward(self, inputs):
        for k in inputs:
            if not isinstance(inputs[k], torch.Tensor):
                inputs[k] = torch.Tensor(inputs[k])
            if not inputs[k].device == self.device:
                inputs[k] = inputs[k].to(self.device) 

        return self.cost_pred(inputs['enc1'], inputs['enc2'])

<file_sep>from itertools import chain

import torch
import torch.nn as nn

from blox.tensor.ops import concat_inputs
from blox.torch.subnetworks import HiddenStatePredictorModel
from blox.torch.recurrent_modules import ZeroLSTMCellInitializer, MLPLSTMCellInitializer


class SumTreeHiddenStatePredictorModel(HiddenStatePredictorModel):
    """ A HiddenStatePredictor for tree morphologies. Averages parents' hidden states """

    def forward(self, hidden1, hidden2, *inputs):
        hidden_state = hidden1 + hidden2
        return super().forward(hidden_state, *inputs)


class LinTreeHiddenStatePredictorModel(HiddenStatePredictorModel):
    """ A HiddenStatePredictor for tree morphologies. Averages parents' hidden states """
    def build_network(self):
        super().build_network()
        self.projection = nn.Linear(self.get_state_dim() * 2, self.get_state_dim())

    def forward(self, hidden1, hidden2, *inputs):
        hidden_state = self.projection(concat_inputs(hidden1, hidden2))
        return super().forward(hidden_state, *inputs)


class SplitLinTreeHiddenStatePredictorModel(HiddenStatePredictorModel):
    """ A HiddenStatePredictor for tree morphologies. Averages parents' hidden states """
    def build_network(self):
        super().build_network()
        split_state_size = int(self.get_state_dim() / (self._hp.n_lstm_layers * 2))
        
        if self._hp.use_conv_lstm:
            projection = lambda: nn.Conv2d(split_state_size * 2, split_state_size, kernel_size=3, padding=1)
        else:
            projection = lambda: nn.Linear(split_state_size * 2, split_state_size)
        
        self.projections = torch.nn.ModuleList([projection() for _ in range(self._hp.n_lstm_layers*2)])

    def forward(self, hidden1, hidden2, *inputs):
        chunked_hidden1 = list(chain(*[torch.chunk(h, 2, 1) for h in torch.chunk(hidden1, self._hp.n_lstm_layers, 1)]))
        chunked_hidden2 = list(chain(*[torch.chunk(h, 2, 1) for h in torch.chunk(hidden2, self._hp.n_lstm_layers, 1)]))
        chunked_projected = [projection(concat_inputs(h1, h2))
                             for projection, h1, h2 in zip(self.projections, chunked_hidden1, chunked_hidden2)]
        hidden_state = torch.cat(chunked_projected, dim=1)
        return super().forward(hidden_state, *inputs)


def build_tree_lstm(hp, input_dim, output_dim):
    if hp.tree_lstm == 'sum':
        cls = SumTreeHiddenStatePredictorModel
    elif hp.tree_lstm == 'linear':
        cls = LinTreeHiddenStatePredictorModel
    elif hp.tree_lstm == 'split_linear':
        cls = SplitLinTreeHiddenStatePredictorModel
    else:
        raise ValueError("don't know this TreeLSTM type")

    subgoal_pred = cls(hp, input_dim, output_dim)
    lstm_initializer = get_lstm_initializer(hp, subgoal_pred)
    
    return subgoal_pred, lstm_initializer


def get_lstm_initializer(hp, cell):
    if hp.lstm_init == 'zero':
        return ZeroLSTMCellInitializer(hp, cell)
    elif hp.lstm_init == 'mlp':
        return MLPLSTMCellInitializer(hp, cell, 2 * hp.nz_enc + hp.nz_vae)
    else:
        raise ValueError('dont know lstm init type {}!'.format(hp.lstm_init))
<file_sep>import os
from contextlib import contextmanager

import torch
import torch.nn as nn
from torch.nn.utils.rnn import pad_sequence as pad_sequence

from blox import batch_apply, AttrDict, optional, rmap
from blox.basic_types import subdict
from blox.tensor.core import find_tensor
from blox.tensor.ops import remove_spatial
from blox.torch.dist import ProbabilisticModel
from blox.torch.encoder_decoder import Encoder, DecoderModule
from blox.torch.layers import BaseProcessingNet
from blox.torch.losses import L2Loss
from blox.torch.modules import Identity, LinearUpdater
from blox.torch.subnetworks import ConvSeqEncodingModule, RecurrentSeqEncodingModule, BidirectionalSeqEncodingModule, \
    Predictor
from gcp.prediction.models.auxilliary_models.misc import AttnKeyEncodingModule, LengthPredictorModule, \
    ActionConditioningWrapper
from gcp.prediction import global_params
from gcp.prediction.hyperparameters import get_default_gcp_hyperparameters
from gcp.prediction.models.auxilliary_models.base_model import BaseModel
from gcp.prediction.models.auxilliary_models.cost_mdl import CostModel
from gcp.prediction.models.auxilliary_models.inverse_mdl import InverseModel
from gcp.prediction.utils import visualization


class BaseGCPModel(BaseModel):
    def __init__(self, params, logger):
        super().__init__(logger)
        self._hp = self._default_hparams()
        self.override_defaults(params)  # override defaults with config file
        self.postprocess_params()
        visualization.PARAMS.hp = self._hp
        assert self._hp.batch_size != -1   # make sure that batch size was overridden
        if self._hp.regress_actions or self._hp.action_conditioned_pred:
            assert self._hp.n_actions != -1     # make sure action dimensionality was overridden

        self.build_network()
        self._use_pred_length = False
        self._inv_mdl_full_seq = False

    @contextmanager
    def val_mode(self, pred_length=True):
        """Sets validation parameters. To be used like: with model.val_mode(): ...<do something>..."""
        self.call_children('switch_to_prior', ProbabilisticModel)
        self._use_pred_length = pred_length
        self._inv_mdl_full_seq = True
        yield
        self.call_children('switch_to_inference', ProbabilisticModel)
        self._use_pred_length = False
        self._inv_mdl_full_seq = False

    def postprocess_params(self):
        super().postprocess_params()
        if self._hp.action_activation is None:
            pass
        elif self._hp.action_activation == 'sigmoid':
            self._hp.action_activation = torch.sigmoid
        elif self._hp.action_activation == 'tanh':
            self._hp.action_activation = torch.tanh
        else:
            raise ValueError('Action activation {} not supported!'.format(self._hp.action_activation))
        
        global_params.hp = self._hp
    
    def _default_hparams(self):
        parent_params = super()._default_hparams()
        default_dict = get_default_gcp_hyperparameters()
        
        for k in default_dict.keys():
            parent_params.add_hparam(k, default_dict[k])
        return parent_params

    def build_network(self, build_encoder=True):
        if build_encoder:
            self.encoder = Encoder(self._hp)
        self.decoder = DecoderModule(self._hp,  # infer actions in decoder if not using SH-Pred model
                               regress_actions=self._hp.regress_actions and self._hp.one_step_planner is not 'sh_pred')

        self.build_inference()

        if self._hp.regress_length:
            self.length_pred = LengthPredictorModule(self._hp)

        if self._hp.attach_inv_mdl:
            self.inv_mdl = InverseModel(self._hp.inv_mdl_params, self._logger)

        if self._hp.attach_cost_mdl:
            self.cost_mdl = CostModel(self._hp.cost_mdl_params, self._logger)

        if self._hp.attach_state_regressor:
            self.state_regressor = BaseProcessingNet(self._hp.nz_enc, self._hp.nz_mid, self._hp.state_dim,
                                                     self._hp.n_processing_layers, self._hp.fc_builder)

        if self._hp.separate_cnn_start_goal_encoder:
            from blox.torch.layers import LayerBuilderParams
            from tensorflow.contrib.training import HParams
            with_conv_hp = HParams()
            for k in self._hp.values().keys():
                with_conv_hp.add_hparam(k, self._hp.values()[k])
            #with_conv_hp = self._hp 
            with_conv_hp.set_hparam('use_convs', True)
            with_conv_hp.set_hparam('input_nc', 3)
            with_conv_hp.set_hparam('builder', LayerBuilderParams(
                            True, self._hp.use_batchnorm, self._hp.normalization, self._hp.predictor_normalization))
            self.start_goal_enc = Encoder(with_conv_hp)

    def build_inference(self):
        if self._hp.states_inference:
            self.inf_encoder = Predictor(self._hp, self._hp.nz_enc + 2, self._hp.nz_enc)
        elif self._hp.act_cond_inference:
            self.inf_encoder = ActionConditioningWrapper(self._hp, self.build_temporal_inf_encoder())
        else:
            self.inf_encoder = self.build_temporal_inf_encoder()
    
        self.inf_key_encoder = nn.Sequential(self.build_temporal_inf_encoder(),
                                             AttnKeyEncodingModule(self._hp, add_time=False))
    
        if self._hp.kl_weight_burn_in is not None:
            target_kl_weight = self._hp.kl_weight
            delattr(self._hp, 'kl_weight')
            del self._hp._hparam_types['kl_weight']
            self._hp.add_hparam('kl_weight', nn.Parameter(torch.zeros(1)))
            self._hp.kl_weight.requires_grad_(False)
            self.kl_weight_updater = LinearUpdater(
                self._hp.kl_weight, self._hp.kl_weight_burn_in, target_kl_weight, name='kl_weight')

    def build_temporal_inf_encoder(self):
        if self._hp.seq_enc == 'none':
            return Identity()
        elif self._hp.seq_enc == 'conv':
            return ConvSeqEncodingModule(self._hp)
        elif self._hp.seq_enc == 'lstm':
            return RecurrentSeqEncodingModule(self._hp)
        elif self._hp.seq_enc == 'bi-lstm':
            return BidirectionalSeqEncodingModule(self._hp)
            
    def forward(self, inputs, phase='train'):
        """
        forward pass at training time
        :param
            images shape = batch x time x height x width x channel
            pad mask shape = batch x time, 1 indicates actual image 0 is padded
        :return:
        """
        outputs = AttrDict()
        inputs.reference_tensor = find_tensor(inputs)
        self.optional_preprocessing(inputs)
    
        # Get features
        self.run_encoder(inputs)
        # Optionally predict sequence length
        end_ind = self.get_end_ind(inputs, outputs)
        # Generate sequences
        outputs.update(self.predict_sequence(inputs, outputs, inputs.start_ind, end_ind, phase))
        # Predict other optional values
        outputs.update(self.run_auxilliary_models(inputs, outputs, phase))
    
        return outputs

    def optional_preprocessing(self, inputs):
        if self._hp.non_goal_conditioned:
            # Set goal to zeros
            if 'traj_seq' in inputs:
                inputs.traj_seq[torch.arange(inputs.traj_seq.shape[0]), inputs.end_ind] = 0.0
                inputs.traj_seq_images[torch.arange(inputs.traj_seq.shape[0]), inputs.end_ind] = 0.0
            inputs.I_g = torch.zeros_like(inputs.I_g)
            if "I_g_image" in inputs:
                inputs.I_g_image = torch.zeros_like(inputs.I_g_image)
            if inputs.I_0.shape[-1] == 5: # special hack for maze
                inputs.I_0[..., -2:] = 0.0
                if "traj_seq" in inputs:
                    inputs.traj_seq[..., -2:] = 0.0

        if self._hp.train_on_action_seqs:
            # swap in actions if we want to train action sequence decoder
            inputs.traj_seq = torch.cat([inputs.actions, torch.zeros_like(inputs.actions[:, :1])], dim=1)

        if 'start_ind' not in inputs:
            inputs.start_ind = torch.zeros(self._hp.batch_size, dtype=torch.long, device=inputs.reference_tensor.device)

    def run_encoder(self, inputs):
        if 'traj_seq' in inputs:
            # Encode sequence
            if not 'enc_traj_seq' in inputs:
                inputs.enc_traj_seq, inputs.skips = batch_apply(self.encoder, inputs.traj_seq)
                if 'skips' in inputs:
                    inputs.skips = rmap(lambda s: s[:, 0], inputs.skips)  # only use start image activations

            # Encode sequence for inference and attention
            if self._hp.act_cond_inference:
                inputs.inf_enc_seq = self.inf_encoder(inputs.enc_traj_seq, inputs.actions)
            elif self._hp.states_inference:
                inputs.inf_enc_seq = batch_apply(self.inf_encoder,
                                                 inputs.enc_traj_seq, inputs.traj_seq_states[..., None, None])
            else:
                inputs.inf_enc_seq = self.inf_encoder(inputs.enc_traj_seq)
            inputs.inf_enc_key_seq = self.inf_key_encoder(inputs.enc_traj_seq)
            
        # Encode I_0, I_g
        if self._hp.separate_cnn_start_goal_encoder:
            e_0, inputs.skips = self.start_goal_enc(inputs.I_0_image)
            inputs.e_0 = remove_spatial(e_0)
            inputs.e_g = remove_spatial(self.start_goal_enc(inputs.I_g_image)[0])
        else:
            inputs.e_0, inputs.skips = self.encoder(inputs.I_0)
            inputs.e_g = self.encoder(inputs.I_g)[0]
            
        # Encode actions
        if self._hp.action_conditioned_pred:
            inputs.enc_action_seq = batch_apply(self.action_encoder, inputs.actions)
            
    def get_end_ind(self, inputs, outputs):
        # TODO clean this up. outputs.end_ind is not currently used anywhere
        end_ind = inputs.end_ind if 'end_ind' in inputs else None
        
        if self._hp.regress_length:
            # predict total sequence length
            outputs.update(self.length_pred(inputs.e_0, inputs.e_g))
            if self._use_pred_length and (self._hp.length_pred_weight > 0 or end_ind is None):
                end_ind = torch.clamp(torch.argmax(outputs.seq_len_pred.sample().long(), dim=1), min=2)   # min pred seq len needs to be >= 3 for planning
                if self._hp.action_conditioned_pred or self._hp.non_goal_conditioned:
                    # don't use predicted length when action conditioned
                    end_ind = torch.ones_like(end_ind) * (self._hp.max_seq_len - 1)

        outputs.end_ind = end_ind
        return end_ind
    
    def predict_sequence(self, inputs, outputs, start_ind, end_ind, phase):
        raise NotImplementedError("Needs to be implemented in the child class")

    def run_auxilliary_models(self, inputs, outputs, phase):
        aux_outputs = AttrDict()
        
        if self.prune_sequences:
            if phase == 'train':
                inputs.model_enc_seq = self.get_matched_pruned_seqs(inputs, outputs)
            else:
                inputs.model_enc_seq = self.get_predicted_pruned_seqs(inputs, outputs)
            inputs.model_enc_seq = pad_sequence(inputs.model_enc_seq, batch_first=True)
            if len(inputs.model_enc_seq.shape) == 5:
                inputs.model_enc_seq = inputs.model_enc_seq[..., 0, 0]
                
            if inputs.model_enc_seq.shape[1] == 0:
                # This happens usually in the beginning of training where the pruning is inaccurate
                return aux_outputs
        
            if self._hp.attach_inv_mdl and phase == 'train':
                aux_outputs.update(self.inv_mdl(inputs, full_seq=self._inv_mdl_full_seq or self._hp.train_inv_mdl_full_seq))
            if self._hp.attach_state_regressor:
                regressor_inputs = inputs.model_enc_seq
                if not self._hp.supervised_decoder:
                    regressor_inputs = regressor_inputs.detach()
                    aux_outputs.regressed_state = batch_apply(self.state_regressor, regressor_inputs)
            if self._hp.attach_cost_mdl and self._hp.run_cost_mdl and phase == 'train':
                # There is an issue here since SVG doesn't output a latent for the first imagge
                # Beyong conceptual problems, this breaks if end_ind = 199
                aux_outputs.update(self.cost_mdl(inputs))
                
        return aux_outputs

    def loss(self, inputs, outputs, log_error_arr=False):
        losses = AttrDict()
        
        # Length prediction loss
        if self._hp.regress_length:
            losses.update(self.length_pred.loss(inputs, outputs))
        
        # Dense Reconstruction loss
        losses.update(self.dense_rec.loss(inputs, outputs.dense_rec, log_error_arr))

        # Inverse Model loss
        if self._hp.attach_inv_mdl:
            losses.update(self.inv_mdl.loss(inputs, outputs, add_total=False))

        # Cost model loss
        if self._hp.attach_cost_mdl and self._hp.run_cost_mdl:
            losses.update(self.cost_mdl.loss(inputs, outputs))

        # State regressor cost
        if self._hp.attach_state_regressor:
            reg_len = outputs.regressed_state.shape[1]
            losses.state_regression = L2Loss(1.0)(outputs.regressed_state, inputs.traj_seq_states[:, :reg_len],
                                                  weights=inputs.pad_mask[:, :reg_len][:, :, None])

        # Negative Log-likelihood (upper bound)
        if 'dense_img_rec' in losses and 'kl' in losses:
            losses.nll = AttrDict(value=losses.dense_img_rec.value + losses.kl.value, weight=0.0)

        return losses
    
    def get_total_loss(self, inputs, losses):
        # compute total loss
        ## filtering is important when some losses are nan
        ## the unsqueeze is important when some of the weights or losses are 1-dim tensors.
        # TODO use the function from blox
        total_loss = torch.stack([loss[1].value[None] * loss[1].weight for loss in
                                  filter(lambda x: x[1].weight > 0, losses.items())]).sum()
        total_loss = total_loss / torch.prod(torch.tensor(inputs.traj_seq.shape[1:]))
        if torch.isnan(total_loss).any():
            import pdb; pdb.set_trace()
        return AttrDict(value=total_loss)

    def log_outputs(self, outputs, inputs, losses, step, log_images, phase):
        super().log_outputs(outputs, inputs, losses, step, log_images, phase)
        if self._hp.attach_inv_mdl:
            self.inv_mdl.log_outputs(outputs, inputs, losses, step, log_images, phase)
        
        if log_images:
            dataset = self._hp.dataset_class
            if 'regressed_state' in outputs:
                self._logger.log_dataset_specific_trajectory(outputs, inputs, "regressed_state_topdown", step, phase,
                                                             dataset, predictions=outputs.regressed_state,
                                                             end_inds=inputs.end_ind)
    
            if 'regressed_state' in outputs and self._hp.attach_inv_mdl and 'actions' in outputs:
                if len(outputs.actions.shape) == 3:
                    actions = outputs.actions
                else:
                    # Training, need to get the action sequence
                    actions = self.inv_mdl(inputs, full_seq=True).actions
                    
                cum_action_traj = torch.cat((outputs.regressed_state[:, :1], actions), dim=1).cumsum(1)
                self._logger.log_dataset_specific_trajectory(outputs, inputs, "action_traj_topdown", step, phase, dataset,
                                                             predictions=cum_action_traj, end_inds=inputs.end_ind)
    
            if not self._hp.use_convs:
                self._logger.log_dataset_specific_trajectory(outputs, inputs, "prediction_topdown", step, phase, dataset)
                
                # TODO remove this
                if self._hp.log_states_2d:
                    self._logger.log_states_2d(outputs, inputs, "prediction_states_2d", step, phase)
                    
                if self._hp.train_on_action_seqs:
                    action_seq = outputs.dense_rec.images
                    cum_action_seq = torch.cumsum(action_seq, dim=1)
                    self._logger.log_dataset_specific_trajectory(outputs, inputs, "cum_action_prediction_topdown", step,
                                                                 phase, dataset, predictions=cum_action_seq,
                                                                 end_inds=inputs.end_ind)

        if self._hp.dump_encodings:
            os.makedirs(self._logger._log_dir + '/stored_data/', exist_ok=True)
            torch.save(subdict(inputs, ['enc_traj_seq', 'traj_seq', 'traj_seq_states', 'actions']),
                       self._logger._log_dir + '/stored_data/encodings_{}'.format(step))

        if self._hp.dump_encodings_inv_model:
            os.makedirs(self._logger._log_dir + '/stored_data_inv_model/', exist_ok=True)
            torch.save(subdict(inputs, ['model_enc_seq', 'traj_seq_states', 'actions']),
                       self._logger._log_dir + '/stored_data_inv_model/encodings_{}.th'.format(step))

    @property
    def prune_sequences(self):
        return self._hp.attach_inv_mdl or (self._hp.attach_cost_mdl and self._hp.run_cost_mdl) \
               or self._hp.attach_state_regressor

    def get_predicted_pruned_seqs(self, inputs, outputs):
        raise NotImplementedError

    def get_matched_pruned_seqs(self, inputs, outputs):
        name = 'encodings' if outputs.dense_rec else 'e_g_prime'  # for SVG vs tree
        if 'dtw' in self._hp.matching_type:
            # use precomputed matching dists for pruning
            matched_latents = self.tree_module.binding.get_matched_sequence(outputs.tree, 'e_g_prime')
            
            batch, time = inputs.traj_seq.shape[:2]
            model_enc_seq = [matched_latents[i_ex, :inputs.end_ind[i_ex] + 1] for i_ex in range(batch)]
            model_enc_seq = model_enc_seq
        else:
            # batched collection for SVG and balanced tree
            model_enc_seq = self.dense_rec.get_all_samples_with_len(
                inputs.end_ind, outputs, inputs, 'basic', name=name)[0]
        return model_enc_seq

<file_sep>import numpy as np
import torch
import torch.nn as nn

from blox import AttrDict
from blox.tensor.ops import batchwise_index
from blox.torch.losses import L2Loss, PenaltyLoss
from blox.torch.ops import batch_cdist, like, list2ten
from gcp.prediction.utils.visualization import draw_frame
from gcp.prediction import global_params


class LossAveragingCriterion(nn.Module):
    def __init__(self, hp):
        self._loss = L2Loss
        self._hp = hp
        super().__init__()
    
    def loss(self, outputs, targets, weights, pad_mask, weight, log_sigma):
        predictions = outputs.tree.bf.images
        gt_match_dists = outputs.gt_match_dists
        
        # Compute likelihood
        loss_val = batch_cdist(predictions, targets, reduction='sum')
        
        log_sigmas = log_sigma - WeightsHacker.hack_weights(torch.ones_like(loss_val)).log()
        n = np.prod(predictions.shape[2:])
        loss_val = 0.5 * loss_val * torch.pow(torch.exp(-log_sigmas), 2) + n * (log_sigmas + 0.5 * np.log(2 * np.pi))
        
        # Weigh by matching probability
        match_weights = gt_match_dists
        match_weights = match_weights * pad_mask[:, None]  # Note, this is now unnecessary since both tree models handle it already
        loss_val = loss_val * match_weights * weights
        
        losses = AttrDict()
        losses.dense_img_rec = PenaltyLoss(weight, breakdown=2)(loss_val, log_error_arr=True, reduction=[-1, -2])
        
        # if self._hp.top_bias > 0.0:
        #     losses.n_top_bias_nodes = PenaltyLoss(
        #         self._hp.supervise_match_weight)(1 - WeightsHacker.get_n_top_bias_nodes(targets, weights))
        
        return losses

    def get_soft_estimates(self, gt_match_dists, vals):
        """ This function is only used to produce visualization now. Move it. """
    
        # soft_matched_estimates = torch.sum(add_n_dims(gt_match_dists, len(vals.shape)-2) * vals[:, :, None], dim=1)
        def soft_average(values):
            """ Averages per-node values to compute per-frame values """
        
            return torch.einsum('int, in...->it...', gt_match_dists, values).detach()
    
        soft_matched_estimates = soft_average(vals)
    
        # Mark top nodes
        if self._hp.use_convs:
            color = torch.zeros(vals.shape[:2]).to(vals.device)
            color[:, :3] = 0.5
            color_t = soft_average(color)
            soft_matched_estimates = draw_frame(soft_matched_estimates, color_t)
    
        return soft_matched_estimates


class WeightsHacker():
    @staticmethod
    def hack_weights(weights):
        if abs(global_params.hp.leaves_bias) > 0.0:
            w_1 = weights[:, :-global_params.hp.max_seq_len]
            w_2 = weights[:, -global_params.hp.max_seq_len:] * (1 - global_params.hp.leaves_bias)
            weights = torch.cat([w_1, w_2], 1)
        
        if global_params.hp.top_bias != 1.0:
            w_1 = weights[:, :global_params.hp.n_top_bias_nodes] * global_params.hp.top_bias
            w_2 = weights[:, global_params.hp.n_top_bias_nodes:]
            weights = torch.cat([w_1, w_2], 1)
        
        return weights
    
    @staticmethod
    def hack_weights_df(weights):
        # TODO implement bf2df for indices and use here
        
        if global_params.hp.top_bias != 1.0:
            n_top_bias_layers = np.int(np.log2(global_params.hp.n_top_bias_nodes + 1))
            depth = np.int(np.log2(weights.shape[1] + 1))
            m = torch.ones(weights.shape[:2], device=weights.device)
            for l in range(n_top_bias_layers):
                m[:, 2 ** (depth - l - 1) - 1:: 2 ** (depth - l)] = global_params.hp.top_bias
            weights = weights * m[:, :, None]
        return weights
    
    @staticmethod
    def get_n_top_bias_nodes(targets, weights):
        """ Return the probability that the downweighted nodes match the noisy frame"""
        inds = WeightsHacker.get_index(targets)
        
        noise_frames = batchwise_index(weights, inds, 2)
        n = noise_frames.mean(0)[:global_params.hp.n_top_bias_nodes].sum() / global_params.hp.top_bias
        
        return n
    
    @staticmethod
    def can_get_index():
        return 'dataset_conf' in global_params.data_conf and 'dataset_class' in global_params.data_conf.dataset_conf \
               and global_params.data_conf.dataset_conf.dataset_class == PointMassDataset
    
    @staticmethod
    def can_get_d2b(inputs):
        if 'actions' in inputs and inputs.actions.shape[2] == 3 and \
              torch.equal(torch.unique(inputs.actions[:, :, 2]), like(list2ten, inputs.actions)([-1, 1])):
            # Looks like sawyer
            return True
        else:
            return False
    
    @staticmethod
    def distance2bottleneck(inputs, outputs):
        dists = []
        for i in range(inputs.actions.shape[0]):
            gripper = inputs.actions[i, :, -1]
            picks = (gripper[1:] == gripper[:-1] + 2).nonzero()[:, 0]
            places = (gripper[1:] == gripper[:-1] - 2).nonzero()[:, 0]
            bottlenecks = torch.cat([picks, places], -1)
            
            top_inds = outputs.tree.bf.match_dist[i, :3].argmax(-1)
            
            # top_inds = torch.from_numpy(np.asarray([20, 40, 60])).long().cuda()
            
            def closest_point_distance(a, b):
                mat = torch.abs(a[:, None] - b[None, :])
                ind = mat.argmin(-1)
                return torch.abs(a - b[ind])
            
            dist = closest_point_distance(top_inds, bottlenecks)
            dists.append(dist)
        
        return torch.mean(torch.stack(dists).float(), 0)
<file_sep>import datetime
import os
from functools import partial

import dload
import numpy as np
import torch
from blox import AttrDict
from blox.tensor.ops import batchwise_index


def select_e_0_e_g(seq, start_ind, end_ind):
    e_0 = batchwise_index(seq, start_ind)
    e_g = batchwise_index(seq, end_ind)
    return e_0, e_g


def get_end_ind(pad_mask):
    """
    :param pad_mask: torch tensor with 1 where there is an actual image and zeros where there's padding
    pad_mask has shape batch_size x max_seq_len
    :return:
    """
    max_seq_len = pad_mask.shape[1]
    end_ind = torch.argmax(pad_mask * torch.arange(max_seq_len, dtype=torch.float, device=pad_mask.device), 1)

    return end_ind


def get_pad_mask(end_ind, max_seq_len):
    """
    :param pad_mask: torch tensor with 1 where there is an actual image and zeros where there's padding
    pad_mask has shape batch_size x max_seq_len
    :return:
    """
    use_torch = isinstance(end_ind, torch.Tensor)
    
    if use_torch:
        arange_fn = partial(torch.arange, dtype=torch.long, device=end_ind.device)
    else:
        arange_fn = np.arange
    
    pad_mask = arange_fn(max_seq_len) <= end_ind[:, None]
    
    if use_torch:
        pad_mask = pad_mask.float()
    else:
        pad_mask = pad_mask.astype(np.float)
    
    return pad_mask


def datetime_str():
    return datetime.datetime.now().strftime("_%Y-%m-%d_%H-%M-%S")


def make_path(exp_dir, conf_path, prefix, make_new_dir):
    # extract the subfolder structure from config path
    path = conf_path.split('experiments/', 1)[1]
    if make_new_dir:
        prefix += datetime_str()
    base_path = os.path.join(exp_dir, path)
    return os.path.join(base_path, prefix) if prefix else base_path


def set_seeds(seed=0):
    """Sets all seeds and disables non-determinism in cuDNN backend."""
    torch.manual_seed(seed)
    torch.backends.cudnn.deterministic = True
    torch.backends.cudnn.benchmark = False
    np.random.seed(seed)


def get_dataset_path(dataset_name):
    """Returns path to dataset."""
    return os.path.join(os.environ["GCP_DATA_DIR"], dataset_name)


def download_data(dataset_name):
    """Downloads data if not yet existent."""
    DATA_URLs = AttrDict(
        nav_9rooms='https://www.seas.upenn.edu/~oleh/datasets/gcp/nav_9rooms.zip',
        nav_25rooms='https://www.seas.upenn.edu/~oleh/datasets/gcp/nav_25rooms.zip',
        sawyer='https://www.seas.upenn.edu/~oleh/datasets/gcp/sawyer.zip',
        h36m='https://www.seas.upenn.edu/~oleh/datasets/gcp/h36m.zip',
    )
    if dataset_name not in DATA_URLs:
        raise ValueError("Dataset identifier {} is not known!".format(dataset_name))
    if not os.path.exists(get_dataset_path(dataset_name)):
        print("Downloading dataset from {} to {}.".format(DATA_URLs[dataset_name], os.environ["GCP_DATA_DIR"]))
        print("This may take a few minutes...")
        dload.save_unzip(DATA_URLs[dataset_name], os.environ["GCP_DATA_DIR"], delete_after=True)
        print("...Done!")
<file_sep>from blox import AttrDict
from experiments.prediction.base_configs import base_tree as base_conf

configuration = AttrDict(base_conf.configuration)

model_config = AttrDict(base_conf.model_config)
model_config.update({
    'matching_type': 'dtw_image',
    'learn_matching_temp': False,
    'attentive_inference': True,
})
<file_sep>import torch
import torch.distributions
from gcp.prediction.models.tree.untied_layers_tree import UntiedLayersTree
from gcp.prediction.models.tree.tree_module import TreeModule
from gcp.prediction.models.tree.tree_dense_rec import TreeDenseRec

from blox import AttrDict, rmap
from blox.basic_types import subdict
from blox.tensor.ops import add_n_dims
from gcp.prediction.models.base_gcp import BaseGCPModel
from gcp.prediction.utils.tree_utils import SubgoalTreeLayer


class TreeModel(BaseGCPModel):
    def build_network(self, build_encoder=True):
        super().build_network(build_encoder)
        
        cls = TreeModule
        if self._hp.untied_layers:
            cls = UntiedLayersTree
        self.tree_module = cls(self._hp, self.decoder)

        self.dense_rec = TreeDenseRec(
            hp=self._hp, input_size=self._hp.nz_enc, output_size=self._hp.nz_enc, decoder=self.decoder)
        
    def _create_initial_nodes(self, inputs):
        start_node, end_node = AttrDict(e_g_prime=inputs.e_0, images=inputs.I_0), \
                               AttrDict(e_g_prime=inputs.e_g, images=inputs.I_g)
        if not self._hp.attentive_inference:
            start_match_timestep, end_match_timestep = self.tree_module.binding.get_init_inds(inputs)
            start_node.update(AttrDict(match_timesteps=start_match_timestep))
            end_node.update(AttrDict(match_timesteps=end_match_timestep))
        if self._hp.tree_lstm:
            start_node.hidden_state, end_node.hidden_state = None, None
        return start_node, end_node

    def filter_layerwise_inputs(self, inputs):
        layerwise_input_keys = ['z']  # these inputs are assumed to be depth-first inputs per node in dim 1
        layerwise_inputs = subdict(inputs, layerwise_input_keys, strict=False)
        return layerwise_inputs

    def predict_sequence(self, inputs, outputs, start_ind, end_ind, phase):
        layerwise_inputs = self.filter_layerwise_inputs(inputs)
        start_node, end_node = self._create_initial_nodes(inputs)

        outputs.tree = root = SubgoalTreeLayer()
        tree_inputs = [layerwise_inputs, start_ind, end_ind, start_node, end_node]
        tree_inputs = rmap(lambda x: add_n_dims(x, n=1, dim=1), tree_inputs)
        tree_inputs = [inputs] + tree_inputs

        root.produce_tree(*tree_inputs, self.tree_module, self._hp.hierarchy_levels)
        outputs.dense_rec = self.dense_rec(root, inputs)

        if 'traj_seq' in inputs and phase == 'train':
            # compute matching between nodes & frames of input sequence, needed for loss and inv mdl etc.
            self.tree_module.compute_matching(inputs, outputs)

        # TODO the binding has to move to this class
        outputs.pruned_prediction = self.tree_module.binding.prune_sequence(inputs, outputs)

        # add pruned reconstruction if necessary
        if not outputs.dense_rec and self._hp.matching_type == 'balanced':
            # TODO this has to be unified with the balanced tree case
            outputs.pruned_prediction = self.dense_rec.get_all_samples_with_len(
                outputs.end_ind, outputs, inputs, pruning_scheme='basic')[0]

        return outputs

    def get_predicted_pruned_seqs(self, inputs, outputs):
        return self.tree_module.binding.prune_sequence(inputs, outputs, 'e_g_prime')

    def loss(self, inputs, outputs, log_error_arr=False):
        losses = super().loss(inputs, outputs, log_error_arr)

        losses.update(self.tree_module.loss(inputs, outputs))

        return losses

    def log_outputs(self, outputs, inputs, losses, step, log_images, phase):
        super().log_outputs(outputs, inputs, losses, step, log_images, phase)

        if outputs.tree.subgoals is None:
            outputs.tree.subgoals = AttrDict()

        if log_images:
            dataset = self._hp.dataset_class
            if self._hp.use_convs:
                self._logger.log_hierarchy_image(outputs, inputs, "hierarchical_splits", step, phase)
                self._logger.log_rows_gif([outputs.pruned_prediction], 'pruned_seq_gif', step, phase)
                
                if 'match_dist' in outputs.tree.subgoals:
                    # Any model that has matching
                    self._logger.log_gt_match_overview(outputs, inputs, "match_overview", step, phase)
    
                    if 'soft_matched_estimates' in outputs:
                        self._logger.log_loss_gif(outputs.soft_matched_estimates, inputs.traj_seq,
                                                  'gt_target_gif', step, phase)

                if phase == 'val' and 'images' in outputs.tree.subgoals:
                    self._logger.log_val_tree(outputs, inputs, "output_tree", step, phase)
                if 'pixel_copy_mask' in outputs.tree.subgoals:
                    self._logger.log_balanced_tree(outputs, "pixel_copy_mask", "pixel_copy_masks", step, phase)
                if 'gamma' in outputs.tree.subgoals and outputs.tree.subgoals.gamma is not None:
                    self._logger.log_attention_overview(outputs, inputs, "attention_masks", step, phase)
                if outputs.dense_rec:
                    self._logger.log_pruned_pred(outputs, inputs, "pruned_pred", step, phase)
                    if outputs.tree.pruned is not None:
                        self._logger.log_pruned_tree(outputs, "pruned_tree", step, phase)

            log_prior_images = False
            if log_prior_images:
                # Run the model N times
                with torch.no_grad(), self.val_mode():
                    rows = list([self(inputs).pruned_prediction for i in range(4)])
                    self._logger.log_rows_gif(rows, "prior_samples", step, phase)
                    
                    for i in range(4):
                        if 'regressed_state' in outputs:
                            out = self(inputs, 'test')
                            self._logger.log_dataset_specific_trajectory(outputs, inputs,
                                                                         "prior_regressed_state_topdown_" + str(i),
                                                                         step, phase, dataset,
                                                                         predictions=out.regressed_state,
                                                                         end_inds=inputs.end_ind)
<file_sep>import matplotlib; matplotlib.use('Agg')
import argparse
import copy
import glob
import importlib.machinery
import importlib.util
import os
import random
from multiprocessing import Pool

import numpy as np
from blox import AttrDict
from gcp.planning.infra.sim.benchmarks import run_trajectories


def bench_worker(conf, iex=-1, ngpu=1):
    print('started process with PID:', os.getpid())
    random.seed(None)
    np.random.seed(None)
    print('start ind', conf['start_index'])
    print('end ind', conf['end_index'])
    run_trajectories(conf, iex, gpu_id=conf['gpu_id'], ngpu=ngpu)


def check_and_pop(dict, key):
    if dict.pop(key, None) is not None:
        print('popping key: {}'.format(key))


def postprocess_hyperparams(hyperparams, args):
    if args.data_save_postfix:
        hyperparams['data_save_dir'] = os.path.join(hyperparams['data_save_dir'], args.data_save_postfix)
    return hyperparams


class ControlManager:
    def __init__(self, args_in=None, hyperparams=None):
        parser = argparse.ArgumentParser(description='run parallel data collection')
        parser.add_argument('experiment', type=str, help='experiment name')
        parser.add_argument('--nworkers', type=int, help='use multiple threads or not', default=1)
        parser.add_argument('--gpu_id', type=int, help='the starting gpu_id', default=0)
        parser.add_argument('--ngpu', type=int, help='the number of gpus to use', default=1)
        parser.add_argument('--gpu', type=int, help='the gpu to use', default=-1)
        parser.add_argument('--nsplit', type=int, help='number of splits', default=-1)
        parser.add_argument('--isplit', type=int, help='split id', default=-1)
        parser.add_argument('--iex', type=int, help='if different from -1 use only do example', default=-1)
        parser.add_argument('--data_save_postfix', type=str, help='appends to the data_save_dir path', default='')
        parser.add_argument('--nstart_goal_pairs', type=int, help='max number of start goal pairs', default=None)
        parser.add_argument('--resume_from', type=int, help='from which traj idx to continue generating', default=None)

        args = parser.parse_args(args_in)

        print("Resume from")
        print(args.resume_from)

        if args.gpu != -1:
            os.environ["CUDA_VISIBLE_DEVICES"] = str(args.gpu)

        if hyperparams is None:
            hyperparams_file = args.experiment
            loader = importlib.machinery.SourceFileLoader('mod_hyper', hyperparams_file)
            spec = importlib.util.spec_from_loader(loader.name, loader)
            mod = importlib.util.module_from_spec(spec)
            loader.exec_module(mod)
            hyperparams = AttrDict(mod.config)

        self.args = args
        self.hyperparams = postprocess_hyperparams(hyperparams, args)
        
    def run(self, logging_conf=None):
        args = self.args
        hyperparams = self.hyperparams
    
        gpu_id = args.gpu_id
    
        n_worker = args.nworkers
        if args.nworkers == 1:
            parallel = False
        else:
            parallel = True
        print('parallel ', bool(parallel))

        if args.nsplit != -1:
            assert args.isplit >= 0 and args.isplit < args.nsplit, "isplit should be in [0, nsplit-1]"
           
            n_persplit = max((hyperparams['end_index'] + 1 - hyperparams['start_index']) / args.nsplit, 1)
            hyperparams['end_index'] = int((args.isplit + 1) * n_persplit + hyperparams['start_index'] - 1)
            hyperparams['start_index'] = int(args.isplit * n_persplit + hyperparams['start_index'])

        n_traj = hyperparams['end_index'] - hyperparams['start_index'] + 1
        traj_per_worker = int(n_traj // np.float32(n_worker))
        offset = int(args.resume_from // np.float32(n_worker)) if args.resume_from is not None else 0
        start_idx = [hyperparams['start_index'] + offset + traj_per_worker * i for i in range(n_worker)]
        end_idx = [hyperparams['start_index'] + traj_per_worker * (i+1)-1 for i in range(n_worker)]

        if 'gen_xml' in hyperparams['agent']:
            try:
                os.system("rm {}".format('/'.join(str.split(hyperparams['agent']['filename'], '/')[:-1]) + '/auto_gen/*'))
            except: pass

        self.set_paths(hyperparams, args)
        record_queue, record_saver_proc, counter = None, None, None

        if args.iex != -1:
            hyperparams['agent']['iex'] = args.iex
    
        conflist = []
        for i in range(n_worker):
            modconf = copy.deepcopy(hyperparams)
            modconf['start_index'] = start_idx[i]
            modconf['end_index'] = end_idx[i]
            modconf['ntraj'] = n_traj
            modconf['gpu_id'] = i + gpu_id
            if logging_conf is not None:
                modconf['logging_conf'] = logging_conf
            conflist.append(modconf)
        if parallel:
            self.start_parallel(conflist, n_worker)
        else:
            bench_worker(conflist[0], args.iex, args.ngpu)
    
        if args.save_thread:
            record_queue.put(None)           # send flag to background thread that it can end saving after it's done
            record_saver_proc.join()         # joins thread and continues execution

    def set_paths(self, hyperparams, args):
        subpath = str.partition(hyperparams['current_dir'], 'experiments')[-1]

        if 'data_save_dir' not in hyperparams:
            data_save_dir = os.environ['GCP_DATA_DIR'] + '/' + subpath
            hyperparams['data_save_dir'] = data_save_dir
        print('setting data_save_dir to', hyperparams['data_save_dir'])
        if 'log_dir' not in hyperparams:
            log_dir = os.environ['GCP_EXP_DIR'] + '/' + subpath
            if args.data_save_postfix:
                log_dir = os.path.join(log_dir, args.data_save_postfix)
            hyperparams['log_dir'] = log_dir
        print('setting log_dir to', hyperparams['log_dir'])
        result_dir = hyperparams['data_save_dir'] + '/log'
        if not os.path.exists(result_dir):
            os.makedirs(result_dir)
        for file in glob.glob(result_dir + '/*.pkl'):
            os.remove(file)
        hyperparams['result_dir'] = result_dir

    def start_parallel(self, conflist, n_worker):
        # mp.set_start_method('spawn')  # this is important for parallelism with xvfb
        p = Pool(n_worker)
        p.map(bench_worker, conflist)


if __name__ == '__main__':
    ControlManager().run()
<file_sep>import numpy as np
import torch

from blox import AttrDict
from blox.torch.ops import cdist
from blox.utils import PriorityQueue
from gcp.prediction.models.tree.frame_binding import BalancedBinding

try:
    print("\nUsing fast C-version of DTW!\n")
    import gcp.evaluation.cutils as cutils
    from gcp.evaluation.dtw_utils import c_dtw as dtw
except:
    print("\nC-version of DTW not compiled! Falling back to slower Numpy version!\n")
    from gcp.evaluation.dtw_utils import basic_dtw as dtw


def torch2np(tensor):
    return tensor.detach().cpu().numpy()


class BaseEvalBinding:
    def __init__(self, hp):
        self._hp = hp
    
    def __call__(self, *args, **kwargs):
        raise NotImplementedError("This function needs to be implemented by the subclass!")
    
    def _init_tree(self, tree, inputs):
        bsize = tree.end_ind.shape[0]
        for node in tree:
            node.selected = np.zeros(bsize, dtype=bool)
    
    def _collect_sequence(self, tree, inputs, i_ex):
        return torch.stack([node.subgoal.images[i_ex] for node in tree.depth_first_iter()])
    
    @staticmethod
    def _check_output(output_seq, length):
        if output_seq.shape[0] != length:
            print("Expected output length {}, got sequence of length {}. Repeating last frame".format(length,
                                                                                                      output_seq.shape[
                                                                                                          0]))
            output_seq = torch.cat((output_seq, output_seq[-1:].expand(length - output_seq.shape[0], -1, -1, -1)))
            return output_seq
        return output_seq


class GreedyExistEvalBinding(BaseEvalBinding):
    def __call__(self, tree, inputs, length, i_ex):
        """Perform greedy search over the tree prioritizing with respect to existence probability."""
        if tree.selected is None:
            self._init_tree(tree)
        root = tree
        
        p_queue = PriorityQueue()
        p_queue.push(root, torch2np(root.subgoal.p_n_hat[i_ex]))
        
        for i in range(length):
            node = p_queue.pop()
            node.selected[i_ex] = True
            
            s0 = node.s0
            s1 = node.s1
            if s0.subgoal is not None:
                p_queue.push(s0, torch2np(s0.subgoal.p_n_hat[i_ex]))
            if s1.subgoal is not None:
                p_queue.push(s1, torch2np(s1.subgoal.p_n_hat[i_ex]))
        
        gen_images = self._collect_sequence(tree, inputs, i_ex)
        gen_images = self._check_output(gen_images, length)
        return gen_images, None


class GreedyL2EvalBinding(BaseEvalBinding):
    def _init_tree(self, tree, inputs):
        super()._init_tree(tree, inputs)
        tree.min_l2_match(np.zeros_like(inputs.end_ind.cpu().numpy()), inputs.end_ind.cpu().numpy() + 1,
                                  inputs.traj_seq,
                                  np.asarray(np.ones_like(inputs.end_ind.cpu().numpy()), dtype=np.uint8))
    
    @staticmethod
    def _get_best_filler_seq(partial_gt_seq, frame_1, frame_2):
        frames = torch.stack((frame_1, frame_2), dim=0)
        gt_seg_len = partial_gt_seq.shape[0]
        l2_distances = torch.nn.MSELoss(reduction='none')(partial_gt_seq[None], frames[:, None]) \
            .view(2, gt_seg_len, -1).mean(-1)
        frame_choices = torch.argmin(l2_distances, dim=0)
        frames = frames[frame_choices]
        return [t[0] for t in torch.split(frames, 1)]
    
    @staticmethod
    def _collect_sequence(tree, inputs, i_ex):
        sel_list = []
        
        def maybe_fill(prev_matched_step, prev_matched_img, matched_step, matched_img):
            diff = matched_step - prev_matched_step
            if diff > 1:  # add missing nodes
                sel_list.extend(
                    GreedyL2EvalBinding._get_best_filler_seq(inputs.traj_seq[i_ex, prev_matched_step + 1:matched_step],
                                                             prev_matched_img, matched_img))
        
        prev_matched_step, prev_matched_img = -1, None
        for node in tree:  # iterate through the tree
            if node.selected[i_ex]:
                matched_step, matched_img = node.match_eval_idx[i_ex], node.subgoal.images[i_ex]
                prev_matched_img = matched_img if prev_matched_img is None else prev_matched_img  # fill with first predicted image
                maybe_fill(prev_matched_step, prev_matched_img, matched_step, matched_img)
                sel_list.append(matched_img)
                prev_matched_step, prev_matched_img = matched_step, matched_img
        maybe_fill(prev_matched_step, prev_matched_img, inputs.end_ind[i_ex] + 1,
                   prev_matched_img)  # fill with last predicted image
        return torch.stack(sel_list, dim=0)
    
    def __call__(self, tree, inputs, length, i_ex):
        """Perform greedy minimal-L2-matchings from the root of the tree."""
        if tree.selected is None:
            self._init_tree(tree, inputs)
        gen_images = self._collect_sequence(tree, inputs, i_ex)
        gen_images = self._check_output(gen_images, length)
        return gen_images, None


class DTWEvalBinding(BaseEvalBinding):
    def __call__(self, outputs, inputs, length, i_ex, targets=None, estimates=None):
        """ Match """
        if estimates is None:
            estimates = self._collect_sequence(outputs.tree, inputs, i_ex)
        
        if targets is None:
            targets = inputs.traj_seq[i_ex, :inputs.end_ind[i_ex] + 1]

        return self.get_single_matches(targets, estimates)
    
    @staticmethod
    def get_single_matches(targets, estimates):
        # Get dtw
        matrix = cdist(estimates, targets, reduction='mean').data.cpu().numpy()
        # matrix = ((estimates[:, None] - targets[None]) ** 2).mean(dim=[2, 3, 4]).data.cpu().numpy()
        # norm = lambda x,targets: torch.nn.MSELoss()(x,targets).data.cpu().numpy()
        d, cost_matrix, path = dtw(matrix)
        
        # Get best matches for gt frames
        match_matrix = np.zeros_like(cost_matrix)
        match_matrix[:, :] = np.inf
        match_matrix[path[0], path[1]] = cost_matrix[path[0], path[1]]
        inds = np.argmin(match_matrix, axis=0)
        gen_images = estimates[inds]
        matching_output = AttrDict(targets=targets, estimates=estimates, matching_path=path, gen_images=gen_images)
        return gen_images, matching_output
    
    @staticmethod
    def vis_matching(matching_output):
        """
        Visualizes the DTW matching path between GT and predicted sequence
        :param matching_output: Dict that gets returned in 'get_single_matches'
                (targets: n_t, channels, res, res; dtype: torch.Tensor)
                (estimates: n_n, channels, res, res; dtype: torch.Tensor)
                ([2x(max(n_t, n_n))]; dtype: ndarray)
                (gen_images: n_t, channels, res, res; dtype: torch.Tensor)
        """
        n_t, channels, res, _ = matching_output.targets.shape
        n_n = matching_output.estimates.shape[0]
        img = -torch.ones((channels, res * (n_t + 1), res * (n_n + 2)), dtype=matching_output.targets.dtype)
        img[:, res:, :res] = matching_output.targets.transpose(0, 1).contiguous().view(channels, n_t * res, res)
        img[:, res:, res:2 * res] = matching_output.gen_images.transpose(0, 1).contiguous().view(channels, n_t * res,
                                                                                                 res)
        img[:, :res, 2 * res:] = matching_output.estimates.permute(1, 2, 0, 3).contiguous().view(channels, res,
                                                                                                 n_n * res)
        for pn, pt in zip(*matching_output.matching_path):
            img[:, (pt + 1) * res: (pt + 2) * res, (pn + 2) * res: (pn + 3) * res] = 1.0
        return img.permute(1, 2, 0).data.cpu().numpy()


class BalancedEvalBinding(BaseEvalBinding):
    def __init__(self, hp):
        self.binding = BalancedBinding(hp)
        super().__init__(hp)

    def __call__(self, outputs, inputs, length, i_ex, name=None, targets=None, estimates=None):
        start_ind, end_ind = self.binding.get_init_inds(outputs)
        if i_ex == 0:
            # Only for the first element
            outputs.tree.compute_matching_dists({}, matching_fcn=self.binding,
                                                left_parents=AttrDict(timesteps=start_ind),
                                                right_parents=AttrDict(timesteps=end_ind))

        name = 'images' if name is None else name
        estimates = torch.stack([node.subgoal[name][i_ex] for node in outputs.tree.depth_first_iter()])
        leave = torch.stack([node.subgoal.c_n_prime[i_ex] for node in outputs.tree.depth_first_iter()]).bool().any(1)
        return estimates[leave], None
    
    def get_all_samples(self, outputs, inputs, length, name=None, targets=None, estimates=None):
        start_ind, end_ind = self.binding.get_init_inds(outputs)
        
        # Only for the first element
        outputs.tree.compute_matching_dists({}, matching_fcn=self.binding,
                                            left_parents=AttrDict(timesteps=start_ind),
                                            right_parents=AttrDict(timesteps=end_ind))

        name = 'images' if name is None else name
        estimates = torch.stack([node.subgoal[name] for node in outputs.tree.depth_first_iter()])
        leave = torch.stack([node.subgoal.c_n_prime for node in outputs.tree.depth_first_iter()]).bool().any(-1)

        pruned_seqs = [estimates[:, i][leave[:, i]] for i in range(leave.shape[1])]
        
        return pruned_seqs, None
    
    
class BalancedPrunedDTWBinding():
    def __init__(self, hp):
        self.pruning_binding = BalancedEvalBinding(hp)
        self.dtw_binding = DTWEvalBinding(hp)

    def __call__(self, outputs, inputs, length, i_ex, targets=None, estimates=None):
        
        pruned_sequence = self.pruning_binding(outputs, inputs, length, i_ex)
        warped_sequence = self.dtw_binding(outputs, inputs, length, i_ex, estimates=pruned_sequence[0])
        
        return warped_sequence
    
    @staticmethod
    def vis_matching(matching_output):
        return DTWEvalBinding.vis_matching(matching_output)
        
<file_sep>import random
import torch.utils.data as data
import numpy as np
from PIL import Image
import glob
import h5py
import random
import torch
import matplotlib.pyplot as plt
from PIL import Image

def get_traj_per_file(file):
    with h5py.File(file, 'r') as F:
        return len(list(F.keys()))

class Video_Dataset(data.Dataset):
    def __init__(self,
                 data_dir,
                 phase
                 ):

        self.data_dir = data_dir + '/' + phase
        self.train_val_split = 0.95

        self.filenames = sorted(glob.glob(self.data_dir + '/*.h5'))
        if not self.filenames:
            raise RuntimeError('No filenames found')
        random.seed(1)
        random.shuffle(self.filenames)

        self.traj_per_file = get_traj_per_file(self.filenames[0])

    def __getitem__(self, index):
        file_index = index//self.traj_per_file
        path = self.filenames[file_index]

        in_file_ind = index % self.traj_per_file

        with h5py.File(path, 'r') as F:
            images = np.asarray(F['traj{}/images'.format(in_file_ind)])
            actions = np.asarray(F['traj{}/actions'.format(in_file_ind)])
        return images, actions

    def __len__(self):
        return len(self.filenames)*self.traj_per_file


def make_data_loader(data_dir, phase):
    """
    :param data_dir:
    :param phase: either train, val or test
    :return: dataset iterator
    """
    d = Video_Dataset(data_dir, phase)
    return torch.utils.data.DataLoader(d, batch_size=3, shuffle=True,num_workers=1, pin_memory=True)

if __name__ == '__main__':

    data_dir = '/mnt/sda1/recursive_planning_data/sim/cartgripper/multi_block_2d'
    d = Video_Dataset(data_dir, phase='train')
    loader = torch.utils.data.DataLoader(d, batch_size=3, shuffle=True,num_workers=1, pin_memory=True)

    deltat = []
    for i_batch, sample_batched in enumerate(loader):

        images, actions = sample_batched

        print(actions)
        for t in range(images.shape[1]):
            plt.imshow(images[0,t,0])
            plt.show()
        # images = images.numpy().transpose((0, 1, 3, 4, 2))
        # file = '/'.join(str.split(config['agent']['data_save_dir'], '/')[:-1]) + '/example'
        # comp_single_video(file, images)

<file_sep># Long-Horizon Visual Planning with Goal-Conditioned Hierarchical Predictors
#### [[Project Website]](https://orybkin.github.io/video-gcp/)

[<NAME>](https://kpertsch.github.io/)<sup>*1</sup>, [<NAME>](https://www.seas.upenn.edu/~oleh/)<sup>*2</sup>, 
[<NAME>](https://febert.github.io/)<sup>3</sup>, [<NAME>](http://people.eecs.berkeley.edu/~cbfinn/)<sup>4</sup>,
[<NAME>](https://www.seas.upenn.edu/~dineshj/)<sup>2</sup>, [<NAME>](https://people.eecs.berkeley.edu/~svlevine/)<sup>3</sup><br/>
(&#42; equal contribution)

<sup>1</sup>University of Southern California 
<sup>2</sup>University of Pennsylvania 
<sup>3</sup>UC Berkeley 
<sup>4</sup>Stanford 

<a href="https://orybkin.github.io/video-gcp/">
<p align="center">
<img src="https://orybkin.github.io/video-gcp/resources/teaser.gif" width="800">
</p>
</img></a>

This is the official Pytorch implementation for our paper **Long-Horizon Visual Planning with Goal-Conditioned Hierarchical Predictors**.


## Installation

To install the module, run the following commands:

```
git clone --recursive <EMAIL>:orybkin/video-gcp.git
cd video-gcp

virtualenv -p $(which python3) ./venv
source ./venv/bin/activate
pip3 install -r requirements.txt
python3 setup.py develop
```

## Model Training

To train a tree-structured prediction model start by setting paths for data and experiment logs:
```
export GCP_DATA_DIR=./data
export GCP_EXP_DIR=./experiment_logs
``` 
Note that the data will be automatically downloaded upon first running the model training code.  

To start training the GCP model on the 25-room navigation dataset, run: 
```
cd gcp
python3 prediction/train.py --path=../experiments/prediction/25room/gcp_tree/ --skip_first_val=1
```
The training logs get written into `GCP_EXP_DIR` and can be displayed by opening a Tensorboard in this folder:
`tensorboard --logdir=$GCP_EXP_DIR`.

## Planning Experiments
For running planning experiments we need to set up a virtual display for rendering (skip this step if you are not running on a headless machine).
```
Xvfb -screen 0 320x240x24 &
export DISPLAY=:0
```

Next, we need to install the `gym-miniworld` submodule:
```
cd gym-miniworld
python3 setup.py develop
```

To run a planning evaluation run with the previously trained model, execute:
```
cd gcp
python3 planning/run.py ../experiments/control/25room/gcp_tree/mod_hyper.py
```

To compute the control performance metrics (success rate and trajectory cost), run:
```
python evaluation/compute_control_perf.py --path=${GCP_EXP_DIR}/control/25room/gcp_tree/scores_0to99.pkl --n_rooms=25
```


## Modifying the Code
### Running on a new dataset
If you want to run our model on a new dataset, the easiest is to subclass [```FolderSplitVarLenVideoDataset```](gcp/datasets/data_loader.py)
 and overwrite the ```__getitem__()``` method to load your data. Then just point to your custom 
data loader in the config file. Our code expects the data loader's ```__getitem__()``` function to output a dictionary with
the following structure:
```
dict({
    'traj_seq': (time, channels, heigth, width) # zero-padded sequence of max_seq_len
    'pad_mask': (time,)                         # one for all time steps that are zero-padded, 0 otherwise
    'I_0': (channels, height, width)            # start image
    'I_g': (channels, height, width)            # goal image
    'end_ind': int                              # index of goal image in the demo sequence
    'states': (time, state_dim)                 # (optional) state sequence corresponding to image sequence, for training state regressor
    'actions': (time, action_dim)               # (optional) action sequence for training inverse model
})
```
Further your dataset folder is expected to have the following structure:
```
< dataset_path >
       |-- dataset_spec.py      # specifies state/action dimensions + sequence length
       |-- train/               # training trajectory files
       |-- val/                 # validation trajectory files
       |-- test/                # test trajectory files
```
You can find an example ```dataset_spec.py``` file  [here](experiments/data_gen/misc/example_dataset_spec.py). 
By default the data loader will search for all ```.h5``` files in the ```train/val/test``` subfolders to compile the list of filenames. 
If you want to load from a different file format or have a different folder structure you need to adjust the ```get_filenames()``` function accordingly.

### Modifying the predictive model
The proposed goal-conditioned predictors are implemented in the [```prediction```](gcp/prediction) directory.
Below we describe the functionality implemented in the different files to allow for quick adjustments.

|File        | Description         |
|:------------- |:-------------|
| [```train.py```](gcp/prediction/train.py) | Tranining script. |
| [```sequential.py```](gcp/prediction/sequential.py) | A naive sequential goal-conditioned predictor (GCP). This works well on some data, but fails to scale to long horizons. |
| [```blox/torch/models/vrnn.py```](blox/torch/models/vrnn.py) | The variational RNN, which is the basis of the sequential GCP. |
| [```tree/tree.py```](gcp/prediction/tree/tree.py) | The proposed hierarchical GCP which scales to long horizons. |
| [```tree/tree_module.py```](gcp/prediction/tree/tree_module.py) | The basic module for hierarchical prediction: predicting a state between two other states. |
| [```tree/tree_lstm.py```](gcp/prediction/tree/tree_lstm.py) | Tree LSTM implementation. |
| [```blox/torch/encoder_decoder.py```](blox/torch/encoder_decoder.py) | The encoder and decoder networks and various decoding distributions. |
| [```adaptive_binding```](gcp/prediction/adaptive_binding) | The hierarchical model wih adaptive binding. |

A significant amount of utils developed by us for this and other projects are in the [blox](https://github.com/orybkin/blox-nn) repo,
including tools for quick fully connected, convolutional, and recurrent architectures, operations on nested tensor dictionaries,
tools for probabilistic deep learning, and tools to interface numpy and torch code.

### Modifying the planning algorithm
The code that implements the image-based planning and control approach described in the paper is in the [```planning```](gcp/planning) directory. 
Below we describe the functionality implemented in the different files to allow for quick adjustments.

|File        | Description         |
|:------------- |:-------------|
| [```run.py```](gcp/planning/run.py) | Launch script for planning & control experiments. |
| [```tree_optimizer.py```](gcp/planning/tree_optimizer.py) | Core hierarchical planning algorithm for GCP-tree latent optimization. |
| [```planner_policy.py```](gcp/planning/planner_policy.py) | Policy wrapper that calls planner and follows plan using inverse model. |
| [```cem/cem_planner.py```](gcp/planning/cem/cem_planner.py) | CEM planners for flat and hierarchical prediction models. |
| [```cem/cem_simulator.py```](gcp/planning/cem/cem_simulator.py) | Wrapper around predictive model for interface with CEM planner. |
| [```cem/sampler.py```](gcp/planning/cem/sampler.py) | Samples latents / actions for CEM planning, calls latent optimizer. |
| [```cem/cost_fcn.py```](gcp/planning/cem/cost_fcn.py) | Cost functions for CEM planning. |

Most of the infrastructure for the planning experiments (start / goal loading, logging, interaction between environment and agent)
directly builds on the [Visual Foresight codebase](https://github.com/SudeepDasari/visual_foresight) and can be found in the 
[```infra```](gcp/planning/infra) directory.

### Generating navigation data with new layout
We provide an example script for using a PRM planner to generate a new navigation dataset in an environment with 16 rooms.
The layout can be further adjusted [here](gcp/planning/infra/envs/miniworld_env/utils/multiroom2d_layout.py). 
The ```nworkers``` argument allows for parallelized data generation.
```
cd recursive_planning
python3 planning/run.py ../experiments/data_gen/nav_16rooms/mod_hyper.py --nworkers=4
```

## Downloading the datasets
The training code automatically downloads the required datasets if they are not already in the expected folders. 
However, if you want to download the datasets independently, you can find the zip files here:

|Dataset        | Link         | Size |
|:------------- |:-------------|:-----|
| 9-room navigation | [https://www.seas.upenn.edu/~oleh/datasets/gcp/nav_9rooms.zip](https://www.seas.upenn.edu/~oleh/datasets/gcp/nav_9rooms.zip) | 140MB |
| 25-room navigation |[https://www.seas.upenn.edu/~oleh/datasets/gcp/nav_25rooms.zip](https://www.seas.upenn.edu/~oleh/datasets/gcp/nav_25rooms.zip)| 395MB|
| Sawyer |[https://www.seas.upenn.edu/~oleh/datasets/gcp/sawyer.zip](https://www.seas.upenn.edu/~oleh/datasets/gcp/sawyer.zip)|395MB|
| Human3.6M 500-frame | [https://www.seas.upenn.edu/~oleh/datasets/gcp/h36m.zip](https://www.seas.upenn.edu/~oleh/datasets/gcp/h36m.zip) | 14GB|


## Bibtex
If you find this work useful in your research, please consider citing:
```
@article{pertsch2020gcp,
    title={Long-Horizon Visual Planning with Goal-Conditioned Hierarchical Predictors},
    author={<NAME> and <NAME> and <NAME> 
    and <NAME> and <NAME> and <NAME>},
    year={2020},
    journal={arXiv preprint arXiv:2006.13205},
}
```

## Acknowledgements
Parts of the planning code are based on the [Visual Foresight codebase](https://github.com/SudeepDasari/visual_foresight).

In parallel with this method, we developed the [blox](https://github.com/orybkin/blox-nn) python module, which contains
many useful utilities and architectures for pytorch. Check it out and let us know if it was useful!
<file_sep>import numpy as np
import torch

from blox import AttrDict
from blox.basic_types import listdict2dictlist
from gcp.prediction.models.auxilliary_models.cost_mdl import TestTimeCostModel


class CostFcn:
    """Base class to define CEM cost functions."""
    def __init__(self, dense_cost, final_step_weight=1.0, *unused_args):
        self._dense_cost = dense_cost
        self._final_step_weight = final_step_weight

    def __call__(self, cem_outputs, goal):
        cost_per_step = self._compute(cem_outputs, goal)
        for i in range(len(cost_per_step)):
            cost_per_step[i][-1] *= self._final_step_weight
        if self._dense_cost:
            return np.array([np.sum(c) for c in cost_per_step])
        else:
            return np.array([c[-1] for c in cost_per_step])

    def _compute(self, cem_outputs, goal):
        raise NotImplementedError


class ImageCost:
    """Provides method to split off image and latent sequence from input sequence."""
    def _split_state_rollout(self, rollouts):
        """Splits off latents from states in joined rollouts."""
        def reshape_to_image(flat):
            if len(flat.shape) != 2:
                import pdb; pdb.set_trace()
            assert len(flat.shape) == 2
            res = int(np.sqrt(flat.shape[1] / 3))   # assumes 3-channel image
            return flat.reshape(flat.shape[0], 3, res, res)
        return listdict2dictlist([AttrDict(image_rollout=reshape_to_image(r[..., :-self.input_dim]),
                                           latent_rollout=r[..., -self.input_dim:]) for r in rollouts])


class EuclideanDistance(CostFcn):
    """Euclidean distance between vals and goal."""
    def _compute(self, cem_outputs, goal):
        euclid_dists = [np.linalg.norm(cem_output - goal[None], axis=-1) for cem_output in cem_outputs]
        return euclid_dists


class EuclideanPathLength(CostFcn):
    """Euclidean length of the whole path to the goal."""
    def _compute(self, cem_outputs, goal):
        assert self._dense_cost     # need dense cost for path length computation
        return [np.linalg.norm(np.concatenate([cem_output[1:], goal[None]]) - cem_output, axis=-1)
                        for cem_output in cem_outputs]


class StepPathLength(CostFcn):
    """Cost is equivalent to number of steps in path."""
    def _compute(self, cem_outputs, goal):
        path_lengths = [cem_output.shape[0] for cem_output in cem_outputs]
        return [np.concatenate((np.zeros(cem_output.shape[0]-1), np.array([path_length]))) 
                        for cem_output, path_length in zip(cem_outputs, path_lengths)]


class L2ImageCost(CostFcn, ImageCost):
    """Cost is equivalent to L2 distance in image space."""
    LATENT_SIZE = 128        # TODO: make this configurable

    def _compute(self, cem_outputs, goal_raw):
        image_sequences = self._split_state_rollout(cem_outputs).image_rollout
        goal = goal_raw.transpose(0, 3, 1, 2) * 2 - 1.0
        return [np.sqrt(np.sum((seq - goal)**2, axis=(1, 2, 3))) for seq in image_sequences]

    @property
    def input_dim(self):
        return self.LATENT_SIZE


class LearnedCostEstimate:
    """Uses learned network to estimate cost between to latent states."""
    def __init__(self, config):
        self.net = TestTimeCostModel(params=config, logger=None)

    def __call__(self, start_enc, goal_enc):
        if isinstance(start_enc, np.ndarray):
            # compute cost for single start goal pair
            return self.net(AttrDict(enc1=start_enc, enc2=goal_enc)).data.cpu().numpy()
        elif isinstance(start_enc, list):
            # compute summed cost for sequence
            costs = []
            for seq, goal in zip(start_enc, goal_enc):
                seq_input = torch.cat((torch.tensor(seq).to(self.net.device), torch.tensor(goal).to(self.net.device)))
                cost_per_step = self.net(AttrDict(enc1=seq_input[:-1], enc2=seq_input[1:]))
                costs.append(cost_per_step.sum().data.cpu().numpy())
            return np.array(costs)
        else:
            raise ValueError("Dimensionality of input to learned cost function not supported!")

    @property
    def input_dim(self):
        return self.net.input_dim


class ImageLearnedCostEstimate(LearnedCostEstimate, ImageCost):
    pass


class ImageWrappedLearnedCostFcn(LearnedCostEstimate, ImageCost):
    """Shallow wrapper around LearnedCostEstimate that unpacks image input."""
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)

    def __call__(self, start_enc, goal_enc):
        start_enc = self._split_state_rollout(start_enc).latent_rollout
        goal_enc = [start_enc[-1] for _ in range(len(start_enc))]   # HACK that only works for goal-cond prediction!
        return super().__call__(start_enc, goal_enc)
<file_sep>from contextlib import contextmanager

import torch

from blox import AttrDict
from blox.tensor.ops import broadcast_final, batch_apply, remove_spatial
from blox.torch.encoder_decoder import Encoder
from blox.torch.losses import L2Loss
from blox.torch.subnetworks import Predictor
from gcp.prediction.models.auxilliary_models.misc import RecurrentPolicyModule
from gcp.prediction.models.auxilliary_models.base_model import BaseModel
from gcp.prediction.training.checkpoint_handler import CheckpointHandler


class BehavioralCloningModel(BaseModel):
    def __init__(self, params, logger):
        super().__init__(logger)
        self._hp = self._default_hparams()
        self.override_defaults(params)  # override defaults with config file
        self.postprocess_params()
        assert self._hp.n_actions != -1     # make sure action dimensionality was overridden
        self.build_network()

        # load only the encoder params during training
        if self._hp.enc_params_checkpoint is not None:
            self._load_weights([
                (self.encoder, 'encoder', self._hp.enc_params_checkpoint),
            ])
        self.detach_enc = self._hp.enc_params_checkpoint is not None and not self._hp.finetune_enc

    def _default_hparams(self):
        # put new parameters in here:
        default_dict = {
            'ngf': 4,  # number of feature maps in shallowest level
            'nz_enc': 32,  # number of dimensions in encoder-latent space
            'nz_mid': 32,       # number of hidden units in fully connected layer
            'nz_mid_lstm': 32,
            'n_lstm_layers': 1,
            'n_processing_layers': 3,  # Number of layers in MLPs
            'reactive': True,   # if False, adds recurrent cell to policy
            'enc_params_checkpoint': None,   # specify pretrained encoder weights to load for training
            'finetune_enc': False,
            'checkpt_path': None,
            'train_first_action_only': False,     # if True, only trains on initial action of sequence
            'n_conv_layers': None,  # Number of conv layers. Can be of format 'n-<int>' for any int for relative spec'n_conv_layers': None,  # Number of conv layers. Can be of format 'n-<int>' for any int for relative spec
        }

        # misc params
        default_dict.update({
            'use_skips': False,
            'dense_rec_type': None,
            'device': None,
            'randomize_length': False,
            'randomize_start': False,
        })

        # add new params to parent params
        parent_params = super()._default_hparams()
        for k in default_dict.keys():
            parent_params.add_hparam(k, default_dict[k])
        return parent_params
    
    def build_network(self, build_encoder=True):
        self.encoder = Encoder(self._hp)
        if not self._hp.reactive:
            self.policy = RecurrentPolicyModule(self._hp, 2*self._hp.nz_enc, self._hp.n_actions)
        else:
            self.policy = Predictor(self._hp, 2*self._hp.nz_enc, self._hp.n_actions)

    def forward(self, inputs, phase='train'):
        """
        forward pass at training time
        """
        if not 'enc_traj_seq' in inputs:
            enc_traj_seq, _ = self.encoder(inputs.traj_seq[:, 0]) if self._hp.train_first_action_only \
                                    else batch_apply(self.encoder, inputs.traj_seq)
            if self._hp.train_first_action_only: enc_traj_seq = enc_traj_seq[:, None]
            enc_traj_seq = enc_traj_seq.detach() if self.detach_enc else enc_traj_seq

        enc_goal, _ = self.encoder(inputs.I_g)
        n_dim = len(enc_goal.shape)
        fused_enc = torch.cat((enc_traj_seq, enc_goal[:, None].repeat(1, enc_traj_seq.shape[1], *([1]*(n_dim-1)))), dim=2)
        #fused_enc = torch.cat((enc_traj_seq, enc_goal[:, None].repeat(1, enc_traj_seq.shape[1], 1, 1, 1)), dim=2)

        if self._hp.reactive:
            actions_pred = batch_apply(self.policy, fused_enc)
        else:
            policy_output = self.policy(fused_enc) 
            actions_pred = policy_output

        # remove last time step to match ground truth if training on full sequence
        actions_pred = actions_pred[:, :-1]	if not self._hp.train_first_action_only else actions_pred

        output = AttrDict()
        output.actions = remove_spatial(actions_pred) if len(actions_pred.shape) > 3 else actions_pred
        return output

    def loss(self, inputs, outputs):
        losses = AttrDict()

        # action prediction loss
        n_actions = outputs.actions.shape[1]
        losses.action_reconst = L2Loss(1.0)(outputs.actions, inputs.actions[:, :n_actions],
                                            weights=broadcast_final(inputs.pad_mask[:, :n_actions], inputs.actions))

        # compute total loss
        #total_loss = torch.stack([loss[1].value * loss[1].weight for loss in losses.items()]).sum()
        #losses.total = AttrDict(value=total_loss)
        # losses.total = total_loss*torch.tensor(np.nan)   # for checking if backprop works
        return losses

    def get_total_loss(self, inputs, losses):
        total_loss = torch.stack([loss[1].value * loss[1].weight for loss in losses.items()]).sum()
        return AttrDict(value=total_loss)

    def log_outputs(self, outputs, inputs, losses, step, log_images, phase):
        super().log_outputs(outputs, inputs, losses, step, log_images, phase)
        if log_images and self._hp.use_convs:
            self._logger.log_pred_actions(outputs, inputs, 'pred_actions', step, phase)

    @contextmanager
    def val_mode(self, *args, **kwargs):
        yield

    @property
    def has_image_input(self):
        return self._hp.use_convs


class TestTimeBCModel(BehavioralCloningModel):
    def __init__(self, params, logger):
        super().__init__(params, logger) 
        if torch.cuda.is_available(): 
            self.cuda()
            self.device = torch.device('cuda')
        else:
            self.device = torch.device('cpu')
        if not self._hp.reactive: self.init_hidden_var = self.init_hidden_var.to(self.device)

        assert self._hp.checkpt_path is not None
        weights_file = CheckpointHandler.get_resume_ckpt_file('latest', self._hp.checkpt_path)
        success = CheckpointHandler.load_weights(weights_file, self)
        if not success: raise ValueError("Could not load checkpoint from {}!".format(weights_file))

    def forward(self, inputs):
        for k in inputs:
            if inputs[k] is None:
                continue
            if not isinstance(inputs[k], torch.Tensor):
                inputs[k] = torch.Tensor(inputs[k])
            if not inputs[k].device == self.device:
                inputs[k] = inputs[k].to(self.device)

        enc, _ = self.encoder(inputs['I_0'])
        enc_goal, _ = self.encoder(inputs['I_g'])
        fused_enc = torch.cat((enc, enc_goal), dim=1)
        if self._hp.reactive:
            action_pred = self.policy(fused_enc)
            hidden_var = None
        else:
            hidden_var = self.init_hidden_var if inputs.hidden_var is None else inputs.hidden_var
            policy_output = self.policy(fused_enc[:, None], hidden_var)
            action_pred, hidden_var = policy_output.output, policy_output.hidden_state[:, 0]
        if self._hp.use_convs:
            return remove_spatial(action_pred if len(action_pred.shape)==4 else action_pred[:, 0]), hidden_var
        else:
            return action_pred, hidden_var

<file_sep>""" This file defines an agent for the MuJoCo simulator environment. """
import copy
import os

import cv2
import matplotlib.pyplot as plt
import numpy as np
from gcp.planning.infra.policy.policy import get_policy_args
from gcp.planning.infra.utils.im_utils import resize_store, npy_to_gif
from tensorflow.contrib.training import HParams


class Image_Exception(Exception):
    def __init__(self):
        pass


class Environment_Exception(Exception):
    def __init__(self):
        pass


class GeneralAgent(object):
    """
    All communication between the algorithms and MuJoCo is done through
    this class.
    """
    N_MAX_TRIALS = 100

    def __init__(self, hyperparams, start_goal_list=None):

        self._hp = self._default_hparams()
        self.override_defaults(hyperparams)

        self.T = self._hp.T
        self._start_goal_list = start_goal_list
        self._goal = None
        self._goal_seq = None
        self._goal_image = None
        self._demo_images = None
        self._reset_state = None
        self._setup_world(0)

    def override_defaults(self, config):
        """
        :param config:  override default valus with config dict
        :return:
        """
        for name, value in config.items():
            print('overriding param {} to value {}'.format(name, value))
            if value == getattr(self._hp, name):
                raise ValueError("attribute is {} is identical to default value!!".format(name))
            if name in self._hp and self._hp.get(name) is None:   # don't do a type check for None default values
                setattr(self._hp, name, value)
            else: self._hp.set_hparam(name, value)

    def _default_hparams(self):
        default_dict = {
            'T': None,
            'adim': None,
            'sdim': None,
            'ncam': 1,
            'rejection_sample': False,   # repeatedly attemp to collect a trajectory if error occurs
            'type': None,
            'env': None,
            'image_height': 48,
            'image_width': 64,
            'nchannels': 3,
            'data_save_dir': '',
            'log_dir': '',
            'make_final_gif': True,   # whether to make final gif
            'make_final_gif_freq': 1,   # final gif, frequency
            'make_final_gif_pointoverlay': False,
            'gen_xml': (True, 1),  # whether to generate xml, and how often
            'start_goal_confs': None,
            'show_progress': False,
            'state_resets': False,   # reset the simluator state zeroing velocities according to policies replan frequency
            'do_not_save_images': False  # dataset savers will not save images if True
        }
        # add new params to parent params
        parent_params = HParams()
        for k in default_dict.keys():
            parent_params.add_hparam(k, default_dict[k])
        return parent_params

    def _setup_world(self, itr):
        """
        Helper method for handling setup of the MuJoCo world.
        Args:
            filename: Path to XML file containing the world information.
        """
        env_type, env_params = self._hp.env
        if self._start_goal_list is not None:
            env_params['init_pos'] = self._start_goal_list[itr, 0]
            env_params['goal_pos'] = self._start_goal_list[itr, 1]
        self.env = env_type(env_params, self._reset_state)

        self._hp.adim = self.adim = self.env.adim
        self._hp.sdim = self.sdim = self.env.sdim
        self._hp.ncam = self.ncam = self.env.ncam
        self.num_objects = self.env.num_objects

    def sample(self, policy, i_traj):
        """
        Runs a trial and constructs a new sample containing information
        about the trial.
        """
        self.i_traj = i_traj
        if self._hp.gen_xml[0]:
            if i_traj % self._hp.gen_xml[1] == 0 and i_traj > 0:
                self._setup_world(i_traj)

        traj_ok, obs_dict, policy_outs, agent_data = False, None, None, None
        i_trial = 0
        imax = self.N_MAX_TRIALS

        while not traj_ok and i_trial < imax:
            i_trial += 1
            try:
                agent_data, obs_dict, policy_outs = self.rollout(policy, i_trial, i_traj)
                traj_ok = agent_data['traj_ok']
            except Image_Exception:
                traj_ok = False
            if not traj_ok:
                print('traj_ok: ', traj_ok)

        print('needed {} trials'.format(i_trial))

        if self._hp.make_final_gif or self._hp.make_final_gif_pointoverlay:
            if i_traj % self._hp.make_final_gif_freq == 0:
                self.save_gif(i_traj, self._hp.make_final_gif_pointoverlay)
                # self.plot_endeff_traj(obs_dict)

        self._reset_state = None # avoid reusing the same reset state

        return agent_data, obs_dict, policy_outs

    def _post_process_obs(self, env_obs, agent_data, initial_obs=False, stage=None):
        """
        Handles conversion from the environment observations, to agent observation
        space. Observations are accumulated over time, and images are resized to match
        the given image_heightximage_width dimensions.

        Original images from cam index 0 are added to buffer for saving gifs (if needed)

        Data accumlated over time is cached into an observation dict and returned. Data specific to each
        time-step is returned in agent_data

        :param env_obs: observations dictionary returned from the environment
        :param initial_obs: Whether or not this is the first observation in rollout
        :return: obs: dictionary of observations up until (and including) current timestep
        """
        agent_img_height = self._hp.image_height
        agent_img_width = self._hp.image_width

        if stage is not None:
            env_obs['stage'] = stage

        if initial_obs:
            T = self._hp.T + 1
            self._agent_cache = {}

            for k in env_obs:
                if k == 'images':
                    if 'obj_image_locations' in env_obs:
                        self.traj_points = []
                    self._agent_cache['images'] = np.zeros((T, self._hp.ncam, agent_img_height, agent_img_width, self._hp.nchannels), dtype=np.uint8)
                elif isinstance(env_obs[k], np.ndarray):
                    obs_shape = [T] + list(env_obs[k].shape)
                    self._agent_cache[k] = np.zeros(tuple(obs_shape), dtype=env_obs[k].dtype)
                else:
                    self._agent_cache[k] = []
            self._cache_cntr = 0

        t = self._cache_cntr
        self._cache_cntr += 1

        point_target_width = agent_img_width

        obs = {}

        if self._hp.show_progress:
            plt.imshow(env_obs['images'][0])
            path = self._hp.log_dir + '/verbose/traj{}/'.format(self.i_traj)
            if not os.path.exists(path):
                os.makedirs(path)
            plt.savefig(path + '/im{}.png'.format(t))

        for k in env_obs:
            if k == 'images':
                resize_store(t, self._agent_cache['images'], env_obs['images'])
                self.gif_images_traj.append(self._agent_cache['images'][t,0])  # only take first camera
            elif k == 'obj_image_locations':
                self.traj_points.append(copy.deepcopy(env_obs['obj_image_locations'][0]))  # only take first camera
                env_obs['obj_image_locations'] = np.round((env_obs['obj_image_locations'] *
                                                           point_target_width / env_obs['images'].shape[2])).astype(
                    np.int64)
                self._agent_cache['obj_image_locations'][t] = env_obs['obj_image_locations']
            elif isinstance(env_obs[k], np.ndarray):
                self._agent_cache[k][t] = env_obs[k]
            else:
                self._agent_cache[k].append(env_obs[k])
            obs[k] = self._agent_cache[k][:self._cache_cntr]

            if k == 'topdown_image':
                self.topdown_images.append((self._agent_cache['topdown_image'][t]*255).astype(np.uint8))  # only take first camera

        if 'obj_image_locations' in env_obs:
            agent_data['desig_pix'] = env_obs['obj_image_locations']
        if self._goal_image is not None:
            agent_data['goal_image'] = self._goal_image

        if self._goal is not None:
            agent_data['goal'] = self._goal
        if self._demo_images is not None:
            agent_data['demo_images'] = self._demo_images
        if self._reset_state is not None:
            agent_data['reset_state'] = self._reset_state
            obs['reset_state'] = self._reset_state
        return obs

    def _required_rollout_metadata(self, agent_data, traj_ok, t, i_tr):
        """
        Adds meta_data such as whether the goal was reached and the total number of time steps
        into the agent dictionary that is MANDATORY for later parts of pipeline
        :param agent_data: Agent data dictionary
        :param traj_ok: Whether or not rollout succeeded
        :return: None
        """
        agent_data['term_t'] = t - 1
        if self.env.has_goal():
            agent_data['goal_reached'] = self.env.goal_reached()
        agent_data['traj_ok'] = traj_ok

    def rollout(self, policy, i_trial, i_traj):
        """
        Rolls out policy for T timesteps
        :param policy: Class extending abstract policy class. Must have act method (see arg passing details)
        :param i_trial: Rollout attempt index (increment each time trajectory fails rollout)
        :return: - agent_data: Dictionary of extra statistics/data collected by agent during rollout
                 - obs: dictionary of environment's observations. Each key maps to that values time-history
                 - policy_ouputs: list of policy's outputs at each timestep.
                 Note: tfrecord saving assumes all keys in agent_data/obs/policy_outputs point to np arrays or primitive int/float
        """
        self._init()

        agent_data, policy_outputs = {}, []

        # Take the sample.
        t = 0
        done = self._hp.T <= 0
        initial_env_obs, self._reset_state = self.env.reset(self._reset_state)
        obs = self._post_process_obs(initial_env_obs, agent_data, True, stage=0)
        policy.reset()

        while not done:
            """
            Every time step send observations to policy, acts in environment, and records observations

            Policy arguments are created by
                - populating a kwarg dict using get_policy_arg
                - calling policy.act with given dictionary

            Policy returns an object (pi_t) where pi_t['actions'] is an action that can be fed to environment
            Environment steps given action and returns an observation
            """

            if self._hp.state_resets:   # reset the simulator to state so that mujoco-based cem-planning can strat from there.
                if t % policy.replan_interval == 0 and t != 0:
                    print('_____')
                    print('gen_ag: performing state reset ')
                    self.env.qpos_reset(obs['qpos_full'][t], obs['qvel_full'][t])

                    new_obs = self.env._get_obs()
                    print('qpos of t ', new_obs['qpos'])
                    print('qvel of t', new_obs['qvel'])
                    print('_____')

            pi_t = policy.act(**get_policy_args(policy, obs, t, i_traj, agent_data))
            policy_outputs.append(pi_t)

            if 'done' in pi_t:
                done = pi_t['done']
            try:
                obs = self._post_process_obs(self.env.step(pi_t['actions']), agent_data)
                # obs = self._post_process_obs(self.env.step(copy.deepcopy(pi_t['actions']), stage=stage), agent_data, stage=pi_t['policy_index'])
            except Environment_Exception as e:
                print(e)
                return {'traj_ok': False}, None, None


            if (self._hp.T - 1) == t or obs['env_done'][-1]:   # environements can include the tag 'env_done' in the observations to signal that time is over
                done = True
            t += 1
            print('t', t)

        traj_ok = self.env.valid_rollout()
        if self._hp.rejection_sample:
            if self._hp.rejection_sample > i_trial:
                assert self.env.has_goal(), 'Rejection sampling enabled but env has no goal'
                traj_ok = self.env.goal_reached()
            print('goal_reached', self.env.goal_reached())

        self._required_rollout_metadata(agent_data, traj_ok, t, i_trial)
        obs.update(self.env.add_extra_obs_info())

        return agent_data, obs, policy_outputs

    def save_gif(self, i_traj, overlay=False):
        if self.traj_points is not None and overlay:
            colors = [tuple([np.random.randint(0, 256) for _ in range(3)]) for __ in range(self.num_objects)]
            for pnts, img in zip(self.traj_points, self.gif_images_traj):
                for i in range(self.num_objects):
                    center = tuple([int(np.round(pnts[i, j])) for j in (1, 0)])
                    cv2.circle(img, center, 4, colors[i], -1)

        file_path = self._hp.log_dir
        # plt.switch_backend('tkagg')
        # plt.imshow(self.gif_images_traj[0])
        # plt.show()
        npy_to_gif(self.gif_images_traj, file_path + '/verbose/traj{}/video'.format(i_traj)) # todo make extra folders for each run?

        if False: #len(self.topdown_images) > 0:
            npy_to_gif(self.topdown_images, file_path + '/verbose/traj{}/topdownvideo'.format(i_traj))

    def plot_endeff_traj(self, obs_dict):
        endeff_pos = obs_dict['regression_state'][:,:3]
        xpos = endeff_pos[:,0]
        zpos = endeff_pos[:,2]
        plt.switch_backend('TkAgg')
        plt.plot(xpos, zpos)
        plt.show()

    def _init(self):
        """
        Set the world to a given model
        """
        self.gif_images_traj, self.topdown_images, self.traj_points = [], [], None


<file_sep>import torch.nn as nn

from blox import AttrDict
from blox.tensor.ops import batchwise_index
from blox.torch.losses import KLDivLoss2
from blox.torch.variational import FixedPrior


class Inference(nn.Module):
    def __init__(self, hp, q):
        super().__init__()
        self._hp = hp
        self.q = q
        self.deterministic = isinstance(self.q, FixedPrior)
    
    def forward(self, inputs, e_l, e_r, start_ind, end_ind, timestep):
        assert timestep is not None
        output = AttrDict(gamma=None)
        
        if self.deterministic:
            output.q_z = self.q(e_l)
            return output
        
        values = inputs.inf_enc_seq
        keys = inputs.inf_enc_key_seq
        
        mult = int(timestep.shape[0] / keys.shape[0])
        if mult > 1:
            timestep = timestep.reshape(-1, mult)
            result = batchwise_index(values, timestep.long())
            e_tilde = result.reshape([-1] + list(result.shape[2:]))
        else:
            e_tilde = batchwise_index(values, timestep[:, 0].long())
        
        output.q_z = self.q(e_l, e_r, e_tilde)
        return output
    
    def loss(self, q_z, p_z, weights=1):
        if q_z.mu.numel() == 0:
            return {}
        
        return AttrDict(kl=KLDivLoss2(self._hp.kl_weight, breakdown=1, free_nats_per_dim=self._hp.free_nats)(
            q_z, p_z, weights=weights, log_error_arr=True))
<file_sep>"""based on: https://github.com/google-research/google-research/blob/master/frechet_video_distance/example.py
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse

import numpy as np
import tensorflow as tf

from blox import AttrDict
from gcp.evaluation.fvd import frechet_video_distance as fvd

# Number of videos must be divisible by 16.
NUMBER_OF_VIDEOS = 16


def get_trainer_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--gt", help="path to ground truth sequence .npy file")
    parser.add_argument("--pred", help="path to predicted sequence .npy file")
    return parser.parse_args()


def main(argv):
    args = get_trainer_args()

    def load_videos(path):
        print("Loading trajectories from {}".format(path))
        if not path.endswith('.npy'): raise ValueError("Can only read in .npy files!")
        seqs = (np.load(path).transpose(0, 1, 3, 4, 2) + 1) / 2
        assert len(seqs.shape) == 5     # need [batch, T, W, H, C] input data
        assert seqs.shape[-1] == 3      # assume 3-channeled seq with channel in last dim
        if seqs.max() <= 1: seqs *= 255     # assume [0...255] range
        return seqs

    gt_seqs = load_videos(args.gt)
    pred_seqs = load_videos(args.pred)

    assert gt_seqs.shape == pred_seqs.shape

    batch, time, h, w, c = gt_seqs.shape
    n_batches = int(np.floor(batch / NUMBER_OF_VIDEOS))     # needs to be dividable by NUMBER_OF_VIDEOS
    print("Evaluating batch of {} sequences of shape {}...".format(NUMBER_OF_VIDEOS, (time, h, w, c)))

    # Step 1: collect all embeddings (needs to run in loop bc network can only handle batch_size 16)
    config = tf.ConfigProto()
    config.gpu_options.allow_growth = True

    end_ind = np.argmax(np.all(np.abs(gt_seqs) < 1e-6, axis=(2, 3, 4)), axis=-1)  # check for black images
    end_ind[end_ind == 0] = time
    embeddings = AttrDict()
    for key, seq in [['gt', gt_seqs], ['pred', pred_seqs]]:
        stored_embeds = []
        for i, s in enumerate(seq):
            length = end_ind[i]
            if length < 10: continue
            with tf.Graph().as_default():
                # construct embedding graph
                seq_ph = tf.placeholder(dtype=tf.float32, shape=(1, length, h, w, c))
                embed = fvd.create_id3_embedding(fvd.preprocess(seq_ph, (224, 224)))
                with tf.Session(config=config) as sess:
                    sess.run(tf.global_variables_initializer())
                    sess.run(tf.tables_initializer())
                    feed_dict = {seq_ph: s[:length][None]}
                    #feed_dict = {seq_ph: np.concatenate((s[:length][None], np.zeros([NUMBER_OF_VIDEOS-1, length] + s.shape[1:])))}
                    print("{} - Seq {} - Length: {}".format(key, i, length))
                    e = sess.run(embed, feed_dict=feed_dict)
                stored_embeds.append(e[0])
        embeddings[key] = np.stack(stored_embeds)
    print("Generated embeddings!")

    # Step 2: evaluate the FVD
    with tf.Graph().as_default():
        gt_embed_ph = tf.placeholder(dtype=tf.float32, shape=embeddings.gt.shape)
        pred_embed_ph = tf.placeholder(dtype=tf.float32, shape=embeddings.pred.shape)
        result = fvd.calculate_fvd(gt_embed_ph, pred_embed_ph)
        with tf.Session(config=config) as sess:
            sess.run(tf.global_variables_initializer())
            sess.run(tf.tables_initializer())
            feed_dict = {gt_embed_ph: embeddings.gt, pred_embed_ph: embeddings.pred}
            print("FVD is: %.2f." % sess.run(result, feed_dict=feed_dict))


if __name__ == "__main__":
    tf.app.run(main)
<file_sep>import argparse
import numpy as np
import torch

import models
from blox import batch_apply
from blox.basic_types import map_dict

def get_trainer_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--gt", help="path to ground truth sequence .npy file")
    parser.add_argument("--pred", help="path to predicted sequence .npy file")
    parser.add_argument("--batch_size", default=8, type=int, help="batch size for network forward pass")
    parser.add_argument('--use_gpu', action='store_true', help='turn on flag to use GPU')
    return parser.parse_args()


def main():
    args = get_trainer_args()

    def load_videos(path):
        print("Loading trajectories from {}".format(path))
        if not path.endswith('.npy'): raise ValueError("Can only read in .npy files!")
        seqs = np.load(path)
        assert len(seqs.shape) == 5     # need [batch, T, C, H, W] input data
        assert seqs.shape[2] == 3      # assume 3-channeled seq with channel in last dim
        seqs = torch.Tensor(seqs)
        if args.use_gpu: seqs = seqs.cuda()
        return seqs     # range [-1, 1]

    gt_seqs = load_videos(args.gt)
    pred_seqs = load_videos(args.pred)
    print('shape: ', gt_seqs.shape)

    assert gt_seqs.shape == pred_seqs.shape

    n_seqs, time, c, h, w = gt_seqs.shape
    n_batches = int(np.floor(n_seqs / args.batch_size))

    # import pdb; pdb.set_trace()
    # get sequence mask (for sequences with variable length
    mask = 1 - torch.all(torch.all(torch.all((gt_seqs + 1.0).abs() < 1e-6, dim=-1), dim=-1), dim=-1)  # check for black images
    mask2 = 1 - torch.all(torch.all(torch.all((gt_seqs).abs() < 1e-6, dim=-1), dim=-1), dim=-1)  # check for gray images
    mask = mask * mask2

    # Initializing the model
    model = models.PerceptualLoss(model='net-lin', net='alex', use_gpu=args.use_gpu)

    # run forward pass to compute LPIPS distances
    distances = []
    for b in range(n_batches):
        x, y = gt_seqs[b*args.batch_size : (b+1)*args.batch_size], pred_seqs[b*args.batch_size : (b+1)*args.batch_size]
        lpips_dist = batch_apply(model, x, y)
        distances.append(lpips_dist)
    distances = torch.cat(distances)
    mean_distance = distances[mask].mean()

    print("LPIPS distance: {}".format(mean_distance))


if __name__ == "__main__":
    main()
<file_sep>import os

from blox import AttrDict

current_dir = os.path.dirname(os.path.realpath(__file__))
from experiments.prediction.base_configs import gcp_sequential as base_conf

configuration = AttrDict(base_conf.configuration)

model_config = base_conf.model_config
model_config.update({
    'action_conditioned_pred': True,
    'non_goal_conditioned': True,
    'nz_vae': 0,
    'var_inf': 'deterministic',
})
<file_sep>import os.path

from blox import AttrDict
from gcp.planning.infra.agent.benchmarking_agent import BenchmarkAgent
from gcp.planning.infra.envs.miniworld_env.multiroom3d.multiroom3d_env import TopdownMultiroom3dEnv
from gcp.planning.planner_policy import ImageCEMPolicy
from gcp.planning.cem.sampler import ImageHierarchicalTreeCEMSampler
from gcp.planning.cem.cem_planner import HierarchicalImageCEMPlanner
from gcp.planning.cem.cost_fcn import ImageLearnedCostEstimate

from experiments.prediction.base_configs import gcp_tree as base_conf

BASE_DIR = '/'.join(str.split(__file__, '/')[:-1])
current_dir = os.path.dirname(os.path.realpath(__file__))

env_params = {
    'init_pos': None,
    'goal_pos': None,
    'n_rooms': 25,
    'heading_smoothing': 0.1,
    'crop_window': 40,
}

agent = AttrDict(
    type=BenchmarkAgent,
    env=(TopdownMultiroom3dEnv, env_params),
    T=400,
    image_height=32,
    image_width=32,
    start_goal_confs=os.environ['GCP_DATA_DIR'] + '/nav_25rooms/start_goal_configs/raw',
)

h_config = AttrDict(base_conf.model_config)
h_config.update({
    'state_dim': 2,
    'ngf': 16,
    'max_seq_len': 200,
    'hierarchy_levels': 8,
    'nz_mid_lstm': 512,
    'n_lstm_layers': 3,
    'nz_mid': 128,
    'nz_enc': 128,
    'nz_vae': 256,
    'regress_length': True,
    'attach_state_regressor': True,
    'attach_inv_mdl': True,
    'inv_mdl_params': AttrDict(
        n_actions=2,
        use_convs=False,
        build_encoder=False,
    ),
    'untied_layers': True,
    'decoder_distribution': 'discrete_logistic_mixture',
})
h_config.pop("add_weighted_pixel_copy")

cem_params = AttrDict(
    prune_final=True,
    horizon=200,
    action_dim=256,
    verbose=True,
    n_iters=3,
    batch_size=10,
    n_level_hierarchy=8,
    sampler=ImageHierarchicalTreeCEMSampler,
    sampling_rates_per_layer=[10, 10],
    cost_fcn=ImageLearnedCostEstimate,
    cost_config=AttrDict(
        checkpt_path=os.environ['GCP_EXP_DIR'] + '/prediction/25room/gcp_tree/weights'
    ),
)

policy = AttrDict(
    type=ImageCEMPolicy,
    params=h_config,
    checkpt_path=cem_params.cost_config.checkpt_path,
    cem_planner=HierarchicalImageCEMPlanner,
    cem_params=cem_params,
    replan_interval=agent.T+2,
    #load_epoch=10,
    closed_loop_execution=True
)

config = AttrDict(
    current_dir=current_dir,
    start_index=0,
    end_index=99,
    agent=agent,
    policy=policy,
    save_format=['raw'],
    data_save_dir=os.environ['GCP_EXP_DIR'] + '/control/nav_25rooms/gcp_tree',
    split_train_val_test=False,
    traj_per_file=1,
)
<file_sep>import numpy as np
import torch
import torch.nn as nn

from blox import batch_apply, AttrDict
from blox.torch.dist import normalize
from blox.torch.losses import BCELogitsLoss, PenaltyLoss
from blox.torch.modules import ExponentialDecayUpdater
from blox.torch.ops import batch_cdist, cdist
from blox.torch.subnetworks import Predictor
from gcp.prediction.models.adaptive_binding.probabilistic_dtw import soft_dtw
from gcp.prediction.models.tree.frame_binding import BaseBinding
from gcp.prediction.utils.tree_utils import depthfirst2breadthfirst
from gcp.prediction.models.adaptive_binding.binding_loss import LossAveragingCriterion, WeightsHacker


class AdaptiveBinding(BaseBinding):
    def build_network(self):
        self.temp = nn.Parameter(self._hp.matching_temp * torch.ones(1))
        if not self._hp.learn_matching_temp:
            self.temp.requires_grad_(False)

        if self._hp.matching_temp_tenthlife != -1:
            assert not self._hp.learn_matching_temp
            self.matching_temp_updater = ExponentialDecayUpdater(
                self.temp, self._hp.matching_temp_tenthlife, min_limit=self._hp.matching_temp_min)
            
        self.distance_predictor = Predictor(self._hp, self._hp.nz_enc * 2, 1, spatial=False)

        self.criterion = LossAveragingCriterion(self._hp)
    
    def get_w(self, pad_mask, inputs, model_output, log=False):
        """ Matches according to the dynamic programming-based posterior. """
        
        # TODO add a prior over w - this is currently trained as an AE.
        # (oleg) it seems that a uniform prior wouldn't change the computation
        # A prior on specific edges would change it somewhat similarly to weighting the cost (but not exactly with p)
    
        # Get cost matrix
        tree = model_output.tree

        if self._hp.matching_type == 'dtw_image':
            imgs = tree.df.images
            cost_matrix = batch_cdist(imgs, inputs.traj_seq, reduction='mean')
        elif self._hp.matching_type == 'dtw_latent':
            img_latents = tree.df.e_g_prime
            cost_matrix = batch_cdist(img_latents, inputs.enc_traj_seq, reduction='mean')
        
        # TODO remove the detachment to propagate the gradients!
        cost_matrix = WeightsHacker.hack_weights_df(cost_matrix)
        w_matrix = soft_dtw(cost_matrix.detach() / self.temp, inputs.end_ind)
        
        # TODO write this up
        # (oleg) There is some magic going on here. To define a likelihood, we define a mixture model for each frame
        # that consists of the nodes and the respective weights. We normalize the weights for it to be a distribution.
        # Then, we invoke Jensen's!
        # Since we expect all elements in the mixture to be either x or have zero probability, the bound is tight.
        w_matrix = normalize(w_matrix, 1)
        
        return depthfirst2breadthfirst(w_matrix)
    
    def prune_sequence(self, inputs, outputs, key='images'):
        seq = getattr(outputs.tree.df, key)
        latent_seq = outputs.tree.df.e_g_prime
        
        distances = batch_apply(self.distance_predictor,
                                latent_seq[:, :-1].contiguous(), latent_seq[:, 1:].contiguous())[..., 0]
        outputs.distance_predictor = AttrDict(distances=distances)

        # distance_predictor outputs true if the two frames are too close
        close_frames = torch.sigmoid(distances) > self._hp.learned_pruning_threshold
        # Add a placeholder for the first frame
        close_frames = torch.cat([torch.zeros_like(close_frames[:, [0]]), close_frames], 1)
        
        pruned_seq = [seq[i][~close_frames[i]] for i in range(seq.shape[0])]
        
        return pruned_seq
    
    def loss(self, inputs, outputs):
        losses = super().loss(inputs, outputs)

        if self._hp.top_bias != 1.0 and WeightsHacker.can_get_index():
            losses.n_top_bias_nodes = PenaltyLoss(self._hp.supervise_match_weight) \
                (1 - WeightsHacker.get_n_top_bias_nodes(inputs.traj_seq, outputs.tree.bf.match_dist))

        if WeightsHacker.can_get_d2b(inputs):
            dist = WeightsHacker.distance2bottleneck(inputs, outputs)
            losses.distance_to_bottleneck_1 = PenaltyLoss(0)(dist[0])
            losses.distance_to_bottleneck_2 = PenaltyLoss(0)(dist[1])
            losses.distance_to_bottleneck_3 = PenaltyLoss(0)(dist[2])

        if self._hp.log_d2b_3x3maze:
            top_nodes = outputs.tree.bf.images[:, :self._hp.n_top_bias_nodes].reshape(-1, 2)
    
            def get_bottleneck_states():
                if inputs.traj_seq_states.max() > 13.5:
                    scale = 45  # 5x5 maze
                else:
                    scale = 27  # 3x3 maze
                start = -0.5 * scale
                end = 0.5 * scale
        
                doors_x = torch.linspace(start, end, self._hp.log_d2b_3x3maze + 1).to(self._hp.device)[1:-1]
                doors_y = torch.linspace(start, end, self._hp.log_d2b_3x3maze * 2 + 1).to(self._hp.device)[1:-1:2]
                n_x = doors_x.shape[0]
                n_y = doors_y.shape[0]
                doors_x = doors_x.repeat(n_y)
                doors_y = doors_y.repeat_interleave(n_x)
                doors = torch.stack([doors_x, doors_y], 1)
                # And the other way around
                return torch.cat([doors, doors.flip(1)], 0)
    
            doors = get_bottleneck_states()
            dist = cdist(top_nodes, doors)
            avg_dist = dist.min(-1).values.mean()
    
            losses.distance_to_doors = PenaltyLoss(self._hp.supervise_match_weight)(avg_dist)
            
        if 'distance_predictor' in outputs:
            df_match_dists = outputs.tree.df.match_dist
            best_matching = df_match_dists.argmax(-1)
            targets = best_matching[:, 1:] == best_matching[:, :-1]  # 1 if frames are too close, i.e. if best matching gt is the same
            losses.distance_predictor = BCELogitsLoss()(outputs.distance_predictor.distances, targets.float())
        
        return losses

    def reconstruction_loss(self, inputs, outputs, weights):
        losses = AttrDict()
    
        outputs.soft_matched_estimates = self.criterion.get_soft_estimates(outputs.gt_match_dists,
                                                                           outputs.tree.bf.images)
        losses.update(self.criterion.loss(
            outputs, inputs.traj_seq, weights, inputs.pad_mask, self._hp.dense_img_rec_weight, self.decoder.log_sigma))
    
        return losses
<file_sep>import os
from contextlib import contextmanager

import numpy as np
import torch
from skimage.io import imsave

from blox import AttrDict
from blox.basic_types import dict_concat
from blox.tensor.ops import batchwise_index
from blox.torch.evaluation import ssim, psnr, mse
from blox.utils import timed
from gcp.prediction.utils.visualization import plot_pruned_tree, make_gif


def make_image_strips(input, gen_seq, phase, outdir, ind):
    """
    :param input:
    :param gen_seq:   t, channel, r, c
    :param ind: batch index to make images for
    :param phase:
    :param outdir:
    :return:
    """

    gen_seq = gen_seq.detach().cpu().numpy()
    gen_seq = np.split(gen_seq, gen_seq.shape[0], axis=0)
    gen_seq = [l.squeeze() for l in gen_seq]
    gen_seq = [np.transpose(item,(1,2,0)) for item in gen_seq]

    input = input.detach().cpu().numpy()
    input = np.split(input, input.shape[0], axis=0)
    input = [l.squeeze() for l in input]
    input = [np.transpose(item,(1,2,0)) for item in input]

    input = np.concatenate(input, axis=1)
    gen_seq = np.concatenate(gen_seq, axis=1)

    input = (input + 1)/2
    gen_seq = (gen_seq + 1)/2

    out = np.concatenate([input, gen_seq], axis=0)

    if not os.path.exists(outdir):
        os.makedirs(outdir)
    imsave(outdir + '/outfile{}_{}.png'.format(ind, phase), out)


class Evaluator:
    """Performs evaluation of metrics etc."""
    N_PLOTTED_ELEMENTS = 5
    LOWER_IS_BETTER_METRICS = ['mse']
    HIGHER_IS_BETTER_METRICS = ['psnr', 'ssim']

    def __init__(self, model, logdir, hp, log_to_file, tb_logger, top_comp_metric='mse'):
        self._logdir = logdir + '/metrics'
        self._logger = FileEvalLogger(self._logdir) if log_to_file else TBEvalLogger(logdir, tb_logger)
        self._hp = hp
        self._pruning_scheme = hp.metric_pruning_scheme
        self._dense_rec_module = model.dense_rec
        self.use_images = model._hp.use_convs
        self._top_of_100 = hp.top_of_100_eval
        self._top_of = 100
        self._top_comp_metric = top_comp_metric
        if not os.path.exists(self._logdir): os.makedirs(self._logdir)
        self.evaluation_buffer = None
        self.full_evaluation = None
        self.dummy_env = None

    def reset(self):
        self.evaluation_buffer = None
        self.full_evaluation = None

    def _erase_eval_buffer(self):
        def get_init_array(val):
            n_eval_samples = self._top_of if self._top_of_100 else 1
            return val * np.ones((self._hp.batch_size, n_eval_samples))
        self.evaluation_buffer = AttrDict(ssim=get_init_array(0.),
                                          psnr=get_init_array(0.),
                                          mse=get_init_array(np.inf),
                                          gen_images=np.empty(self._hp.batch_size, dtype=np.object),
                                          rand_seqs=np.empty(self._hp.batch_size, dtype=np.object))
        for b in range(self._hp.batch_size):
            self.evaluation_buffer.rand_seqs[b] = []
        if not self.use_images:
            self.evaluation_buffer.pop('ssim')
            self.evaluation_buffer.pop('psnr')

    def eval_single(self, inputs, outputs, sample_n=0):
        input_images = inputs.traj_seq
        bsize = input_images.shape[0]
        store_states = "traj_seq_states" in inputs and (inputs.traj_seq_states.shape[-1] == 2 or
                                                        inputs.traj_seq_states.shape[-1] == 5)
        # TODO paralellize DTW
        for b in range(bsize):
            input_seq = input_images[b, :inputs.end_ind[b]+1]
            input_len = input_seq.shape[0]
            gen_seq, matching_output = self._dense_rec_module.get_sample_with_len(b, input_len, outputs, inputs, self._pruning_scheme)
            input_seq, gen_seq = input_seq[1:-1], gen_seq[1:-1]     # crop first and last frame for eval (conditioning frames)
            state_seq = inputs.traj_seq_states[b, :input_len] if store_states else None

            full_gen_seq, gen_seq = self.compute_metrics(b, gen_seq, input_seq, outputs, sample_n)

            if self._is_better(self.evaluation_buffer[self._top_comp_metric][b, sample_n],
                               self.evaluation_buffer[self._top_comp_metric][b]):
                # log visualization results for the best sample only, replace if better
                self.evaluation_buffer.gen_images[b] = AttrDict(gt_seq=input_images.cpu().numpy()[b],
                                                                gen_images=gen_seq,
                                                                full_gen_seq=full_gen_seq,
                                                                matching_outputs=matching_output,
                                                                state_seq=state_seq)

            if sample_n < self.N_PLOTTED_ELEMENTS:
                pred_len = outputs.end_ind[b].data.cpu().numpy() + 1 if 'end_ind' in outputs else input_len
                pred_len_seq, _ = self._dense_rec_module.get_sample_with_len(b, pred_len, outputs, inputs,
                                                                             self._pruning_scheme)
                self.evaluation_buffer.rand_seqs[b].append(pred_len_seq.data.cpu().numpy())

    def compute_metrics(self, b, gen_seq, input_seq, outputs, sample_n):
        input_seq = input_seq.detach().cpu().numpy()
        gen_seq = gen_seq.detach().cpu().numpy()
        full_gen_seq = torch.stack([n.subgoal.images[b] for n in outputs.tree.depth_first_iter()]) \
            .detach().cpu().numpy() if 'tree' in outputs \
                                       and outputs.tree.subgoals is not None else gen_seq
        self.evaluation_buffer.mse[b, sample_n] = mse(gen_seq, input_seq)
        if 'psnr' in self.evaluation_buffer:
            self.evaluation_buffer.psnr[b, sample_n] = psnr(gen_seq, input_seq)
        if 'ssim' in self.evaluation_buffer:
            self.evaluation_buffer.ssim[b, sample_n] = ssim(gen_seq, input_seq)
        return full_gen_seq, gen_seq

    @timed("Eval time for batch: ")
    def eval(self, inputs, outputs, model):
        self._erase_eval_buffer()
        if self._top_of_100:
            for n in range(self._top_of):
                outputs = model(inputs)
                self.eval_single(inputs, outputs, sample_n=n)
        else:
            self.eval_single(inputs, outputs)
        self._flush_eval_buffer()

    def _flush_eval_buffer(self):
        if self.full_evaluation is None:
            self.full_evaluation = self.evaluation_buffer
        else:
            dict_concat(self.full_evaluation, self.evaluation_buffer)

    def dump_results(self, it):
        self.dump_metrics(it)
        if self.use_images:
            self.dump_seqs(it)
            if 'matching_outputs' in self.full_evaluation.gen_images[0] \
                    and self.full_evaluation.gen_images[0].matching_outputs is not None:
                self.dump_matching_vis(it)
        self.reset()

    def dump_trees(self, it):
        no_pruning = lambda x, b: False     # show full tree, not pruning anything
        img_dict = self.full_evaluation.gen_images[0]
        plot_matched = img_dict.outputs.tree.match_eval_idx is not None
        assert Evaluator.N_PLOTTED_ELEMENTS <= len(img_dict.gen_images)   # can currently only max plot as many trees as in batch

        def make_padded_seq_img(tensor, target_width, prepend=0):
            assert len(tensor.shape) == 4     # assume [n_frames, channels, res, res]
            n_frames, channels, res, _ = tensor.shape
            seq_im = np.transpose(tensor, (1, 2, 0, 3)).reshape(channels, res, n_frames * res)
            concats = [np.zeros((channels, res, prepend * res), dtype=np.float32)] if prepend > 0 else []
            concats.extend([seq_im, np.zeros((channels, res, target_width - seq_im.shape[2] - prepend * res), dtype=np.float32)])
            seq_im = np.concatenate(concats, axis=-1)
            return seq_im

        with self._logger.log_to('trees', it, 'image'):
            tree_imgs = plot_pruned_tree(img_dict.outputs.tree, no_pruning, plot_matched).detach().cpu().numpy()
            for i in range(Evaluator.N_PLOTTED_ELEMENTS):
                im = tree_imgs[i]
                if plot_matched:
                    gt_seq_im = make_padded_seq_img(img_dict.gt_seq[i], im.shape[-1])
                    pred_seq_im = make_padded_seq_img(img_dict.gen_images[i], im.shape[-1], prepend=1)  # prepend for cropped first frame
                    im = np.concatenate((gt_seq_im, im, pred_seq_im), axis=1)
                im = np.transpose(im, [1, 2, 0])
                self._logger.log(im)

    def dump_metrics(self, it):
        with self._logger.log_to('results', it, 'metric'):
            best_idxs = self._get_best_idxs(self.full_evaluation[self._top_comp_metric])
            print_st = []
            for metric in sorted(self.full_evaluation):
                vals = self.full_evaluation[metric]
                if metric in ['psnr', 'ssim', 'mse']:
                    if metric not in self.evaluation_buffer: continue
                    best_vals = batchwise_index(vals, best_idxs)
                    print_st.extend([best_vals.mean(), best_vals.std(), vals.std(axis=1).mean()])
                    self._logger.log(metric, vals if self._top_of_100 else None, best_vals)
            print(*print_st, sep=',')

    def dump_seqs(self, it):
        """Dumps all predicted sequences and all ground truth sequences in separate .npy files"""
        DUMP_KEYS = ['gt_seq', 'gen_images', 'full_gen_seq']
        batch = len(self.full_evaluation.gen_images)
        _, c, h, w = self.full_evaluation.gen_images[0].gt_seq.shape
        stacked_seqs = AttrDict()

        for key in DUMP_KEYS:
            if key == 'full_gen_seq':
                time = max([i[key].shape[0] for i in self.full_evaluation.gen_images])
            else:
                time = self.full_evaluation.gen_images[0]['gt_seq'].shape[0] - 1
            stacked_seqs[key] = np.zeros((batch, time, c, h, w), dtype=self.full_evaluation.gen_images[0][key].dtype)
        for b, seqs in enumerate(self.full_evaluation.gen_images):
            stacked_seqs['gt_seq'][b] = seqs['gt_seq'][1:]  # skip the first (conditioning frame)
            stacked_seqs['gen_images'][b, :seqs['gen_images'].shape[0]] = seqs['gen_images']
            stacked_seqs['full_gen_seq'][b, :seqs['full_gen_seq'].shape[0]] = seqs['full_gen_seq']
        for b, seqs in enumerate(self.full_evaluation.rand_seqs[:self.N_PLOTTED_ELEMENTS]):
            key = 'seq_samples_{}'.format(b)
            time = self.full_evaluation.gen_images[0]['gt_seq'].shape[0] - 1
            stacked_seqs[key] = np.zeros((self.N_PLOTTED_ELEMENTS, time, c, h, w), dtype=self.full_evaluation.rand_seqs[0][0].dtype)
            for i, seq_i in enumerate(seqs):
                stacked_seqs[key][i, :seq_i.shape[0]] = seq_i[:time]
        for key in DUMP_KEYS:
            with self._logger.log_to(key, it, 'array'):
                self._logger.log(stacked_seqs[key])
        self.dump_gifs(stacked_seqs, it)

        if self._hp.n_rooms is not None and self.full_evaluation.gen_images[0].state_seq is not None:
            self.dump_traj_overview(it)

    def dump_matching_vis(self, it):
        """Dumps some visualization of the matching procedure."""
        with self._logger.log_to('matchings', it, 'image'):
            try:
                for i in range(min(Evaluator.N_PLOTTED_ELEMENTS, self.full_evaluation.gen_images.shape[0])):
                    im = self._dense_rec_module.eval_binding.vis_matching(self.full_evaluation.gen_images[i].matching_outputs)
                    self._logger.log(im)
            except AttributeError:
                print("Binding does not provide matching visualization")
                pass

    def dump_gifs(self, seqs, it):
        """Dumps gif visualizations of pruned and full sequences."""
        with self._logger.log_to('pruned_seq', it, 'gif'):
            im = make_gif([torch.Tensor(seqs.gt_seq), (torch.Tensor(seqs.gen_images))])
            self._logger.log(im)
        with self._logger.log_to('full_gen_seq', it, 'gif'):
            im = make_gif([torch.Tensor(seqs.full_gen_seq)])
            self._logger.log(im)
        for key in seqs:
            if 'seq_samples' in key:
                with self._logger.log_to(key, it, 'gif'):
                    im = make_gif([torch.Tensor(seqs[key])])
                    self._logger.log(im)

    def dump_traj_overview(self, it):
        """Dumps top-down overview of trajectories in Multiroom datasets."""
        from gcp.planning.infra import Multiroom3dEnv
        if self.dummy_env is None:
            self.dummy_env = Multiroom3dEnv({'n_rooms': self._hp.n_rooms}, no_env=True)
        with self._logger.log_to('trajectories', it, 'image'):
            for b in range(min(Evaluator.N_PLOTTED_ELEMENTS, self.full_evaluation.gen_images.shape[0])):
                im = self.dummy_env.render_top_down(self.full_evaluation.gen_images[b].state_seq.data.cpu().numpy())
                self._logger.log(im * 2 - 1)

    def _is_better(self, val, other):
        """Comparison function for different metrics.
           returns True if val is "better" than any of the values in the array other
        """
        if self._top_comp_metric in self.LOWER_IS_BETTER_METRICS:
            return np.all(val <= other)
        elif self._top_comp_metric in self.HIGHER_IS_BETTER_METRICS:
            return np.all(val >= other)
        else:
            raise ValueError("Currently only support comparison on the following metrics: {}. Got {}."
                             .format(self.LOWER_IS_BETTER_METRICS + self.HIGHER_IS_BETTER_METRICS, self._top_comp_metric))

    def _get_best_idxs(self, vals):
        assert len(vals.shape) == 2     # assumes batch in first dimension, N samples in second dim
        if self._top_comp_metric in self.LOWER_IS_BETTER_METRICS:
            return np.argmin(vals, axis=1)
        else:
            return np.argmax(vals, axis=1)


class EvalLogger:
    def __init__(self, log_dir):
        self._log_dir = log_dir
        self.log_target = None
        self.log_type = None
        self.log_tag = None
        self.log_counter = None

    @contextmanager
    def log_to(self, tag, it, type):
        """Sets logging context (e.g. what file to log to)."""
        raise NotImplementedError

    def log(self, *vals):
        """Implements logging within the 'log_to' context."""
        assert self.log_target is not None      # cannot log without 'log_to' context
        if self.log_type == 'metric':
            self._log_metric(*vals)
        elif self.log_type == 'image':
            self._log_img(*vals)
        elif self.log_type == 'array':
            self._log_array(*vals)
        elif self.log_type == 'gif':
            self._log_gif(*vals)
        self.log_counter += 1

    def _log_metric(self, name, vals, best_vals):
        raise NotImplementedError

    def _log_img(self, img):
        raise NotImplementedError

    def _log_array(self, array):
        np.save(os.path.join(self.log_target, "{}_{}.npy".format(self.log_tag, self.log_counter)), array)

    def _log_gif(self, gif):
        pass

    def _make_dump_dir(self, tag, it):
        dump_dir = os.path.join(self._log_dir, '{}/it_{}'.format(tag, it))
        if not os.path.exists(dump_dir): os.makedirs(dump_dir)
        return dump_dir


class FileEvalLogger(EvalLogger):
    """Logs evaluation results on disk."""
    @contextmanager
    def log_to(self, tag, it, type):
        """Creates logging file."""
        self.log_type, self.log_tag, self.log_counter = type, tag, 0
        if type == 'metric':
            self.log_target = open(os.path.join(self._log_dir, '{}_{}.txt'.format(tag, it)), 'w')
        elif type == 'image' or type == 'array':
            self.log_target = self._make_dump_dir(tag, it)
        elif type == 'gif':
            self.log_target = 'no log'
        else:
            raise ValueError("Type {} is not supported for logging in eval!".format(type))
        yield
        if type == 'metric':
            self.log_target.close()
        self.log_target, self.log_type, self.log_tag, self.log_counter = None, None, None, None

    def _log_metric(self, name, vals, best_vals):
        str = 'mean {} {}, standard error of the mean (SEM) {}'.format(name, best_vals.mean(), best_vals.std())
        str += ', mean std of 100 samples {}\n'.format(vals.std(axis=1).mean()) if vals is not None else '\n'
        self.log_target.write(str)
        print(str)

    def _log_img(self, img):
        #assert -1.0 <= img.min() and img.max() <= 1.0   # expect image to be in range [-1...1]
        imsave(os.path.join(self.log_target, "{}_{}.png".format(self.log_tag, self.log_counter)), (img + 1) / 2)


class TBEvalLogger(EvalLogger):
    """Logs evaluation results to Tensorboard."""
    def __init__(self, log_dir, tb_logger):
        super().__init__(log_dir)
        self._tb_logger = tb_logger
        self.log_step = None

    @contextmanager
    def log_to(self, tag, it, type):
        self.log_type, self.log_tag, self.log_counter, self.log_step = type, tag, 0, it
        if type == 'array':
            self.log_target = self._make_dump_dir(tag, it)
        else:
            self.log_target = 'TB'
        yield
        self.log_target, self.log_type, self.log_tag, self.log_counter, self.log_step = None, None, None, None, None

    def _log_metric(self, name, vals, best_vals):
        self._tb_logger.log_scalar(best_vals.mean(), self.group_tag + '/metric/{}/top100_mean'.format(name), self.log_step, '')
        self._tb_logger.log_scalar(best_vals.std(), self.group_tag + '/verbose/{}/top100_std'.format(name), self.log_step, '')
        if vals is not None:
            self._tb_logger.log_scalar(vals.mean(), self.group_tag + '/verbose/{}/all100_mean'.format(name), self.log_step, '')
            self._tb_logger.log_scalar(vals.std(axis=1).mean(), self.group_tag + '/verbose/{}/all100_std'.format(name), self.log_step, '')

    def _log_img(self, img):
        #assert -1.0 <= img.min() and img.max() <= 1.0  # expect image to be in range [-1...1]
        if not isinstance(img, torch.Tensor): img = torch.tensor(img)
        img = (img.permute(2, 0, 1) + 1) / 2
        self._tb_logger.log_images(img[None], self.group_tag + '/{}'.format(self.log_counter), self.log_step, '')

    def _log_gif(self, gif):
        self._tb_logger.log_video(gif, self.group_tag + '/{}'.format(self.log_counter), self.log_step, '')

    @property
    def group_tag(self):
        assert self.log_tag is not None     # need to set logging context first
        return 'eval/{}'.format(self.log_tag)
<file_sep>from blox import AttrDict
from gcp.prediction.models.tree.tree import TreeModel
from gcp.prediction.utils.logger import HierarchyLogger

configuration = {
    'model': TreeModel,
    'logger': HierarchyLogger,
}
configuration = AttrDict(configuration)

model_config = {
    'one_step_planner': 'sh_pred',
    'hierarchy_levels': 7,
    'binding': 'loss',
    'seq_enc': 'conv',
    'tree_lstm': 'split_linear',
    'lstm_init': 'mlp',
    'add_weighted_pixel_copy': True,
    'dense_rec_type': 'node_prob',
}

<file_sep>import matplotlib; matplotlib.use('Agg')
import pickle as pkl
import numpy as np
import argparse

from gcp.planning.infra.envs.miniworld_env.utils.multiroom2d_layout import define_layout
from gcp.planning.infra.policy.prm_policy.prm_policy import plan_room_seq


def n_room_path(start, end, layout):
    return len(plan_room_seq(start, end, layout.doors))


def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--path", help="path to the config file directory")
    parser.add_argument('--n_rooms', default=9, type=int,
                        help='number of rooms in navigation layout')
    parser.add_argument('--n_tasks', default=100, type=int,
                        help='number of task instances in eval set')
    parser.add_argument('--max_seq_len', default=200, type=int,
                        help='maximum length of sequence (for cost computation')
    return parser.parse_args()


def main():
    args = parse_args()
    FILE = args.path

    rooms_per_side = int(np.sqrt(args.n_rooms))
    layout = define_layout(rooms_per_side, None)

    with open(FILE, 'rb') as pickle_file:
        data = pkl.load(pickle_file)

    paths = data['full_traj']

    success, rooms_to_goal, rooms_traversed = 0, [], []
    penalized_length = 0.

    for i in range(args.n_tasks):
        # extract start / final / goal position and room
        goal_pos = data['reset_state'][i]['goal'][-2:] / 27
        final_pos = paths[i][-1][:2] / 27
        start_pos = paths[i][0][:2] / 27
        goal_pos[1] *= -1; final_pos[1] *= -1; start_pos[1] *= -1
        goal_room = layout.coords2ridx(goal_pos[0], goal_pos[1])
        final_room = layout.coords2ridx(final_pos[0], final_pos[1])
        start_room = layout.coords2ridx(start_pos[0], start_pos[1])

        # compute success
        if final_room == goal_room:
            success += 1

        # compute length
        path = np.stack([p[:2] for p in paths[i]])
        path_len = np.sum(np.linalg.norm(path[1:] - path[:-1], axis=-1))
        penalized_length += path_len if final_room == goal_room else args.max_seq_len

        # compute number of rooms to goal / traversed
        rooms_to_goal += [n_room_path(final_room, goal_room, layout)]
        rooms_traversed += [n_room_path(start_room, final_room, layout)]

    print("Success: \t{}".format(success / args.n_tasks))
    print("Cost: \t{:.2f}".format(penalized_length / args.n_tasks))

    print("")
    print("Room2Goal: \t{}\t{}".format(np.mean(rooms_to_goal), np.std(rooms_to_goal)))
    print("RTravers: \t{}\t{}".format(np.mean(rooms_traversed), np.std(rooms_traversed)))


if __name__ == "__main__":
    main()
<file_sep>import os
import time
import warnings
import math
import numpy as np
from copy import deepcopy
from tensorboardX import SummaryWriter

import torch
from torch import autograd

from blox import AttrDict
from blox.basic_types import map_dict
from blox.utils import AverageMeter
from blox.utils import RecursiveAverageMeter
from gcp.prediction.training.checkpoint_handler import CheckpointHandler, NoCheckpointsException
from gcp.prediction.training.gcp_builder import GCPBuilder
from gcp.evaluation.compute_metrics import Evaluator

warnings.simplefilter('once')


class ModelTrainer(GCPBuilder):
    """ This class defines the training loop of the GCP model"""

    def run(self):
        """ Runs training """
        args = self.cmd_args
        model_conf = self.model._hp
        data_conf = self.train_loader.dataset.data_conf
    
        # TODO clean up resuming
        self.global_step = 0
        start_epoch = 0
        if args.resume or ('checkpt_path' in model_conf and model_conf.checkpt_path is not None):
            ckpt_path = model_conf.checkpt_path if 'checkpt_path' in model_conf else None
            start_epoch = self.resume(args.resume, ckpt_path)
    
        if args.val_sweep:
            epochs = CheckpointHandler.get_epochs(os.path.join(self._hp.exp_path, 'weights'))
            for epoch in list(sorted(epochs))[::4]:
                self.resume(epoch)
                self.val()
            return
    
        if args.dataset_val_sweep:
            self.run_dataset_val_sweep(args, data_conf, model_conf)
            return
    
        ## Train
        if args.train:
            self.train(start_epoch)
        else:
            self.val()

    def resume(self, ckpt, path=None):
        path = os.path.join(self._hp.exp_path, 'weights') if path is None else os.path.join(path, 'weights')
        assert ckpt is not None  # need to specify resume epoch for loading checkpoint
        try:
            weights_file = CheckpointHandler.get_resume_ckpt_file(ckpt, path)
        except NoCheckpointsException:
            return 0
        self.global_step, start_epoch, _ = \
            CheckpointHandler.load_weights(weights_file, self.model,
                                           load_step_and_opt=True, optimizer=self.optimizer,
                                           dataset_length=len(self.train_loader) * self._hp.batch_size,
                                           strict=self.cmd_args.strict_weight_loading)
        self.model.to(self.model.device)
        return start_epoch

    def run_dataset_val_sweep(self, args, data_conf, model_conf):
        assert 'sweep_specs' in data_conf.dataset_spec and data_conf.dataset_spec.sweep_specs  # need to define sweep_specs
        for sweep_spec in data_conf.dataset_spec.sweep_specs:
            print("\nStart eval of dataset {}...".format(sweep_spec.name))
            dc, mc = deepcopy(data_conf), deepcopy(model_conf)
            dc.dataset_spec.dataset_class = sweep_spec.dataset_class
            dc.dataset_spec.split = sweep_spec.split
            dc.dataset_spec.max_seq_len = sweep_spec.max_seq_len
            mc.update({i: dc.dataset_spec[i] for i in dc.dataset_spec if not isinstance(dc.dataset_spec[i], list)})
            if "dense_rec_type" not in mc or mc["dense_rec_type"] is not "svg":
                mc["hierarchy_levels"] = int(np.ceil(math.log2(sweep_spec.max_seq_len)))
            log_dir = self.log_dir + "_" + sweep_spec.name
            writer = SummaryWriter(log_dir)
        
            def rebuild_phase(logger, model, phase, n_repeat=1, dataset_size=-1):
                logger = logger(log_dir, self._hp, max_seq_len=sweep_spec.max_seq_len, summary_writer=writer)
                model = model(mc, logger).to(self.device)
                model.device = self.device
                loader = self.get_dataset(args, model, dc, phase, n_repeat, dataset_size)
                return logger, model, loader
        
            self.logger, self.model, self.val_loader = \
                rebuild_phase(self._hp.logger, self._hp.model, "val", dataset_size=args.val_data_size)
            self.evaluator = Evaluator(self.model, log_dir, self._hp, self.cmd_args.metric, self.logger)
            if args.resume:
                self.resume(args.resume)
            else:
                self.resume("latest")
            self.val()
        print("...Done!")

    def train(self, start_epoch):
        if not self.cmd_args.skip_first_val:
            self.val()
            
        for epoch in range(start_epoch, self._hp.num_epochs):
            self.train_epoch(epoch)
        
            if not self.cmd_args.dont_save:
                self.save_checkpoint(epoch)
            self.val(not (epoch - start_epoch) % 3)
            
    def save_checkpoint(self, epoch):
        state = {
            'epoch': epoch,
            'global_step': self.global_step,
            'state_dict': self.model.state_dict(),
            'optimizer': self.optimizer.state_dict(),
        }
        folder = os.path.join(self._hp.exp_path, 'weights')
        os.makedirs(folder, exist_ok=True)
        torch.save(state, os.path.join(folder, CheckpointHandler.get_ckpt_name(epoch)))
    
    @property
    def log_images_now(self):
        return self.global_step % self.log_images_interval == 0 and self.cmd_args.log_images
    
    @property
    def log_outputs_now(self):
        return self.global_step % self.cmd_args.log_outputs_interval == 0 or self.global_step % self.log_images_interval == 0

    def try_move_to_dev(self, data):
        try:
            return data.to(self.device)
        except:
            # print('warning: could not move {} to gpu'.format(type(data)))
            return data

    def train_epoch(self, epoch):
        self.model.train()
        epoch_len = len(self.train_loader)
        end = time.time()
        batch_time = AverageMeter()
        upto_log_time = AverageMeter()
        data_load_time = AverageMeter()
        forward_backward_time = AverageMeter()
        self.log_images_interval = int(epoch_len / self.cmd_args.imepoch)
        
        print('starting epoch ', epoch)
        
        for self.batch_idx, sample_batched in enumerate(self.train_loader):
            data_load_time.update(time.time() - end)
            inputs = AttrDict(map_dict(self.try_move_to_dev, sample_batched))
            with self.training_context():
                self.optimizer.zero_grad()
                start_fw_bw = time.time()
                output = self.model(inputs)
                losses = self.model.loss(inputs, output)
                losses.total = self.model.get_total_loss(inputs, losses)
                losses.total.value.backward()
                self.call_hooks(inputs, output, losses, epoch)
                self.optimizer.step()
                self.model.step()
                forward_backward_time.update(time.time() - start_fw_bw)
            
            if self.cmd_args.train_loop_pdb:
                import pdb; pdb.set_trace()
            
            upto_log_time.update(time.time() - end)
            if self.log_outputs_now and not self.cmd_args.dont_save:
                self.model.log_outputs(output, inputs, losses, self.global_step,
                                       log_images=self.log_images_now, phase='train')
            batch_time.update(time.time() - end)
            end = time.time()
            
            if self.log_outputs_now:
                print('GPU {}: {}'.format(os.environ["CUDA_VISIBLE_DEVICES"] if self.use_cuda else 'none', self._hp.exp_path))
                print(('itr: {} Train Epoch: {} [{}/{} ({:.0f}%)]\tLoss: {:.6f}'.format(
                    self.global_step, epoch, self.batch_idx, len(self.train_loader),
                    100. * self.batch_idx / len(self.train_loader), losses.total.value.item())))
                
                print('avg time for loading: {:.2f}s, logs: {:.2f}s, compute: {:.2f}s, total: {:.2f}s'
                      .format(data_load_time.avg,
                              batch_time.avg - upto_log_time.avg,
                              upto_log_time.avg - data_load_time.avg,
                              batch_time.avg))
                togo_train_time = batch_time.avg * (self._hp.num_epochs - epoch) * epoch_len / 3600.
                print('ETA: {:.2f}h'.format(togo_train_time))
                if self.cmd_args.verbose_timing: print("avg FW/BW time: {:.3f}s/batch".format(forward_backward_time.avg))
            
            del output, losses
            self.global_step = self.global_step + 1

    def val(self, test_control=True):
        print('Running Testing')
        if self.cmd_args.test_prediction:
            start = time.time()
            losses_meter = RecursiveAverageMeter()
            infer_time = AverageMeter()

            # self.model.eval()
            with autograd.no_grad():
                for batch_idx, sample_batched in enumerate(self.val_loader):
                    inputs = AttrDict(map_dict(self.try_move_to_dev, sample_batched))
                    with self.model.val_mode(pred_length=False):
                        infer_start = time.time()
                        output = self.model(inputs, 'test')
                        infer_time.update(time.time() - infer_start)
                        if self.evaluator is not None:    # force eval on all batches for reduced noise
                            self.evaluator.eval(inputs, output, self.model)
                    # run train model to get NLL on validation data
                    output_train_mdl = self.model(inputs)
                    losses = self.model.loss(inputs, output_train_mdl)
                    losses.total = self.model.get_total_loss(inputs, losses)
                    losses_meter.update(losses)
                    del losses
                    del output_train_mdl

                    # if batch_idx == 0:
                    #     break
                
                if not self.cmd_args.dont_save:
                    if self.evaluator is not None:
                        self.evaluator.dump_results(self.global_step)
                    if self.cmd_args.metric:
                        print("Finished Evaluation! Exiting...")
                        exit(0)
    
                    self.model.log_outputs(
                        output, inputs, losses_meter.avg, self.global_step, log_images=self.cmd_args.log_images, phase='val')
                    print(('\nTest set: Average loss: {:.4f} in {:.2f}s\n'
                           .format(losses_meter.avg.total.value.item(), time.time() - start)))
                    if self.cmd_args.verbose_timing: print("avg Inference time: {:.3f}s/batch".format(infer_time.avg))
            del output

        
if __name__ == '__main__':
    trainer = ModelTrainer()
    trainer.run()
<file_sep>from blox import AttrDict
from gcp.datasets.configs.nav_9rooms import Nav9Rooms


class Nav25Rooms(Nav9Rooms):
    n_rooms = 25


config = AttrDict(
    dataset_spec=AttrDict(
        max_seq_len=200,
        dataset_class=Nav25Rooms,
        split=AttrDict(train=0.994, val=0.006, test=0.00),
    ),
    n_rooms=25,
    crop_window=40,
)

<file_sep>from contextlib import contextmanager

import cv2
import matplotlib.cm as cm
import matplotlib.pyplot as plt
import numpy as np
import torch
from PIL import Image
import numbers
from torch.nn.utils.rnn import pad_sequence as pad_sequence

from blox.tensor.ops import batchwise_index
from blox.tensor import ndim

PLOT_BINARY_DISTS = True    # if True, plots greyscale tiles instead of line plots for matching functions

class Params():
    """ Singleton holding visualization params """
    n_logged_samples = 3
    
PARAMS = Params()


@contextmanager
def param(**kwargs):
    """ A context manager that sets global params to specified values for the context duration """
    # Update values
    old_kwargs = {}
    for name, value in kwargs.items():
        old_kwargs[name] = getattr(Params, name)
        setattr(Params, name, value)
        
    yield
    
    # Reset old values
    for name, value in old_kwargs.items():
        setattr(Params, name, value)
    

def fig2img(fig):
    """Converts a given figure handle to a 3-channel numpy image array."""
    fig.canvas.draw()
    w, h = fig.canvas.get_width_height()
    buf = np.fromstring(fig.canvas.tostring_argb(), dtype=np.uint8)
    buf.shape = (w, h, 4)
    buf = np.roll(buf, 3, axis=2)
    w, h, d = buf.shape
    return np.array(Image.frombytes("RGBA", (w, h), buf.tostring()), dtype=np.float32)[:, :, :3] / 255.


def plot_greyscale_dist(dist, h, w):
    n_tiles = dist.shape[0]
    tile_width = int(w / n_tiles)
    tiled_im = np.repeat(np.repeat(np.repeat(dist, tile_width, axis=0)[None], h, axis=0)[..., None], 3, axis=-1)
    return tiled_im


def plot_dists(dists, h=400, w=400, dpi=10, linewidth=1.0):
    if PLOT_BINARY_DISTS and len(dists) == 1:
        return plot_greyscale_dist(dists[0], h, w)
    COLORS = ['red', 'blue', 'green']
    assert len(dists) <= 3      # only have 3 different colors for now, add more if necessary!
    fig = plt.figure(figsize=(w / dpi, h / dpi), dpi=dpi)
    for dist, color in zip(dists, COLORS[:len(dists)]):
        plt.plot(dist, color=color, linewidth=linewidth)
    plt.ylim(0, 1)
    plt.xlim(0, dist.shape[0]-1)
    plt.tight_layout()
    fig_img = fig2img(fig)
    plt.close(fig)
    return fig_img


def plot_graph(array, h=400, w=400, dpi=10, linewidth=3.0):
    fig = plt.figure(figsize=(w / dpi, h / dpi), dpi=dpi)
    if isinstance(array, torch.Tensor):
        array = array.cpu().numpy()
    plt.xlim(0, array.shape[0] - 1)
    plt.xticks(fontsize=100)
    plt.yticks(fontsize=100)
    plt.plot(array)
    plt.grid()
    plt.tight_layout()
    fig_img = fig2img(fig)
    plt.close(fig)
    return fig_img


def tensor2np(tensor, n_logged_samples=None):
    if tensor is None: return tensor
    if n_logged_samples is None: return tensor.data.cpu().numpy()
    return tensor[:n_logged_samples].data.cpu().numpy()


def visualize(tensor):
    """ Visualizes the state. Returns blank image as default, otherwise calls the dataset method"""
    n_logged_samples = PARAMS.n_logged_samples
    array = np.ones(tensor.shape[:-1] + (3, PARAMS.hp.img_sz, PARAMS.hp.img_sz), dtype=np.float32)[:n_logged_samples]
    
    if hasattr(PARAMS, 'visualize'):
        array = PARAMS.visualize(tensor, array, PARAMS.hp) # 'visualize' is set in train_planner_mode
        
    return array


def imgtensor2np(tensor, n_logged_samples=None, gt_images=False):
    if tensor is None: return tensor
    
    if not PARAMS.hp.use_convs and not gt_images:
        with param(n_logged_samples=n_logged_samples):
            return visualize(tensor[:n_logged_samples])
        
    return (tensor2np(tensor, n_logged_samples) + 1)/2


def np2imgtensor(array, device, n_logged_samples=None):
    if array is None: return array
    if n_logged_samples is not None: array = array[:n_logged_samples]
    return torch.tensor(array * 2 - 1)


def action2img(action, res, channels):
    action_scale = 50   # how much action arrows should get magnified
    assert action.size == 2   # can only plot 2-dimensional actions
    img = np.zeros((res, res, channels), dtype=np.float32).copy()
    start_pt = res/2 * np.ones((2,))
    end_pt = start_pt + action * action_scale * (res/2 - 1) * np.array([1, -1])     # swaps last dimension
    np2pt = lambda x: tuple(np.asarray(x, int))
    img = cv2.arrowedLine(img, np2pt(start_pt), np2pt(end_pt), (255, 255, 255), 1, cv2.LINE_AA, tipLength=0.2)
    return img * 255.0


def batch_action2img(actions, res, channels):
    batch, seq_len, _ = actions.shape
    im = np.empty((batch, seq_len, res, res, channels), dtype=np.float32)
    for b in range(batch):
        for s in range(seq_len):
            im[b, s] = action2img(actions[b, s], res, channels)
    return im


class channel_last:
    """ A function decorator that transposes the input and output image if necessary
     
     The first input to the function has to be the image, and the function must output an image back
     """
    
    def __init__(self):
        """ Decorator parameters """
        pass
    
    def __call__(self, func):
        """ Wrapping """
        
        def wrapper(img, *args, **kwargs):
            transpose = self.is_channel_first(img.shape[-3:])
            
            if transpose:
                sh_offset = len(img.shape) - 3
                offset_order = lambda x: np.concatenate([np.arange(sh_offset), np.array(x) + sh_offset])
                img = ndim.permute(img, offset_order([1, 2, 0]))
                
            result = func(img, *args, **kwargs)
            
            if transpose:
                result = ndim.permute(result, offset_order([2, 0, 1]))
            return result
        
        return wrapper
    
    @staticmethod
    def is_channel_first(shape):
        # The allowed channel dimesions are [1, 3]
    
        # If the last dim is not an allowed channel dimension
        if shape[2] not in [1, 3]:
            assert shape[0] in [1, 3]
            return True
    
        # If the first dim is not an allowed channel dimension
        if shape[0] not in [1, 3]:
            return False
    
        # If the last dim is image size but the first is not
        if hasattr(PARAMS, 'hp') and shape[2] == PARAMS.hp.img_sz and shape[0] != PARAMS.hp.img_sz:
            return True


@ndim.torched
@channel_last()
def draw_frame(img, prob):
    """
    
    :param img: array dims x width x height x colors
    :param prob: array dims
    :return:
    """
    if img.shape[2] == 1: return img
    if isinstance(prob, numbers.Number):
        prob = np.full_like(img[..., 0, 0, 0], prob)
    img = ndim.copy(img)
    
    cmap = cm.get_cmap('RdYlGn')
    rgb = cmap(prob[..., None, None])[..., :3]
    img[..., :, :2, :], img[..., :, -2:, :] = rgb, rgb
    img[..., :2, :, :], img[..., -2:, :, :] = rgb, rgb
    return img


def framed_action2img(action, prob, res, channel):
    """Draws frame around action image indicating the probability of usage."""
    img = action2img(action, res, channel)
    img = draw_frame(img, prob)
    return img


def sort_actions_depth_first(model_output, attrs, n_logged_samples):
    """Sorts actions in depth-first ordering. attrs is a list with attribute names for the left/right
       distributions that should be sorted."""
    assert len(attrs) == 2  # need one left and one right action attribute
    tree = model_output.tree
    n_sg = (2 ** tree.depth) - 1
    dist_shape = [n_logged_samples, 2 * n_sg + 1]
    dist_shape = dist_shape + list(tree.subgoals[attrs[0]].shape[2:]) if len(tree.subgoals[attrs[0]].shape) > 2 else dist_shape
    match_dists = torch.zeros(dist_shape, device=tree.subgoals[attrs[0]].device)
    for i, segment in enumerate(tree.depth_first_iter()):
        match_dists[:, 2 * (i + 1) - 2 ** (segment.depth - 1) - 1] = segment.subgoal[attrs[0]][:n_logged_samples]
        match_dists[:, 2 * (i + 1) + 2 ** (segment.depth - 1) - 1] = segment.subgoal[attrs[1]][:n_logged_samples]
    return match_dists


def plot_balanced_tree(model_output, elem="images"):
    """Plots all subgoals of the tree in a balanced format."""
    tree = model_output.tree
    max_depth = tree.depth
    batch = tree.subgoals[elem].shape[0]
    res = PARAMS.hp.img_sz
    n_logged_samples = PARAMS.n_logged_samples
    n_logged_samples = batch if n_logged_samples is None else n_logged_samples
    
    n_sg = (2 ** max_depth) - 1
    im_height = max_depth * res
    im_width = n_sg * res
    im = 0.7 * np.ones((n_logged_samples, 3, im_height, im_width))
    
    if 'gt_match_dists' in model_output:
        usage_prob = tree.df.match_dist.sum(2)
    elif 'existence_predictor' in model_output:
        usage_prob = model_output.existence_predictor.existence
    else:
        usage_prob = None
    usage_prob = tensor2np(usage_prob, n_logged_samples)
    
    for n, node in enumerate(tree.depth_first_iter()):
        level = max_depth - node.depth
        # Handle 5-dimensional masks
        im_batch = imgtensor2np(node.subgoal[elem][:, [0, 1, -1]], n_logged_samples)
        if usage_prob is not None:
            imgs = draw_frame(im_batch, usage_prob[:, n])
            im_batch = np.stack(imgs, 0)
        im[:, :, level*res : (level+1)*res, n*res : (n+1)*res] = im_batch
    return im


def plot_balanced_tree_with_actions(model_output, inputs, n_logged_samples, get_prob_fcn=None):
    tree = model_output.tree
    batch, channels, res, _ = tree.subgoal.images.shape
    _, action_dim = tree.subgoal.a_l.shape
    max_depth = tree.depth
    n_sg = (2 ** max_depth) - 1

    im_height = (max_depth * 2 + 1) * res  # plot all actions (*2) and start/end frame again (+1)
    im_width = (n_sg + 2) * res
    im = np.asarray(0.7 * np.ones((n_logged_samples, im_height, im_width, 3)), dtype=np.float32)

    # insert start and goal frame
    if inputs is not None:
        im[:, :res, :res] = imgtensor2np(inputs.traj_seq[:n_logged_samples, 0], n_logged_samples).transpose(0, 2, 3, 1)
        im[:, :res, -res:] = imgtensor2np(batchwise_index(inputs.traj_seq[:n_logged_samples], model_output.end_ind[:n_logged_samples]),
                                          n_logged_samples).transpose(0, 2, 3, 1)

    if 'norm_gt_action_match_dists' in model_output:
        action_usage_prob = np.max(tensor2np(model_output.norm_gt_action_match_dists, n_logged_samples), axis=2)

    step = 1
    for i, segment in enumerate(tree):
        level = 2 * (max_depth - segment.depth + 1)
        dx = 2 ** (segment.depth - 2)
        im[:, level * res : (level + 1) * res, step * res: (step + 1) * res] = \
            imgtensor2np(segment.subgoal.images[:n_logged_samples], n_logged_samples).transpose(0, 2, 3, 1)
        a_l, a_r = tensor2np(segment.subgoal.a_l, n_logged_samples), tensor2np(segment.subgoal.a_r, n_logged_samples)
        if get_prob_fcn is not None:
            usage_prob_l, usage_prob_r = get_prob_fcn(segment)
        else:
            usage_prob_l, usage_prob_r = action_usage_prob[:, 2*i], action_usage_prob[:, 2*i+1]
        for b in range(n_logged_samples):
            im[b, (level-1) * res : level * res, int((step-dx) * res): int((step - dx + 1) * res)] = \
                framed_action2img(a_l[b], usage_prob_l[b], res, channels)
            im[b, (level - 1) * res: level * res, int((step + dx) * res): int((step + dx + 1) * res)] = \
                framed_action2img(a_r[b], usage_prob_r[b], res, channels)
        step += 1
    return im


def plot_pruned_tree(tree, check_pruned_fcn=lambda x, b: x.pruned[b], plot_matched=False):
    """Plots subgoal tree, but only non-pruned nodes.
       'check_pruned_fcn' allows flexible definition of what 'pruned' means
       'plot_matched': if True, plot nodes at positions where they were matched
    """
    max_depth = tree.depth
    batch, channels, res, _ = tree.subgoal.images.shape
    n_sg = (2 ** max_depth) - 1
    im_height = max_depth * res
    im_width = n_sg * res
    im = 1.0 * torch.ones((batch, channels, im_height, im_width))
    step = 0
    for segment in tree:
        level = max_depth - segment.depth
        for b in range(batch):
            if check_pruned_fcn(segment, b): continue      # only plot non-pruned elements of the tree
            pos = segment.match_eval_idx[b] if plot_matched else step
            im[b, :, level * res: (level + 1) * res, pos * res: (pos + 1) * res] = (segment.subgoal.images[b]+1 / 2)
        step += 1
    return im


def plot_val_tree(model_output, inputs, n_logged_samples=3):
    tree = model_output.tree
    batch, _, channels, res, _ = tree.subgoals.images.shape
    max_depth = tree.depth
    n_sg = (2 ** max_depth) - 1

    dpi = 10
    fig_height, fig_width = 2 * res, n_sg * res

    im_height, im_width = max_depth*res + fig_height, 2*res + fig_width
    im = np.asarray(0.7 * np.ones((n_logged_samples, im_height, im_width, 3)), dtype=np.float32)

    # plot existence probabilities
    if 'p_n_hat' in tree.subgoals:
        p_n_hat = tensor2np(tree.df.p_n_hat, n_logged_samples)
        for i in range(n_logged_samples):
            im[i, :res, res:-res] = plot_dists([p_n_hat[i]], res, fig_width, dpi)
    if 'p_a_l_hat' in tree.subgoals:
        p_a_hat = tensor2np(sort_actions_depth_first(model_output, ['p_a_l_hat', 'p_a_r_hat'], n_logged_samples))
        for i in range(n_logged_samples):
            im[i, res:2*res, int(3*res/4):int(-3*res/4)] = plot_dists([p_a_hat[i]], res, fig_width + int(res/2), dpi)
        im = np.concatenate((im[:, :fig_height],
                             plot_balanced_tree_with_actions(model_output, inputs, n_logged_samples,
                                                             get_prob_fcn=lambda s: (tensor2np(s.subgoal.p_a_l_hat),
                                                                                     tensor2np(s.subgoal.p_a_r_hat)))), axis=1)
    else:
        with param(n_logged_samples=n_logged_samples):
            im[:, fig_height:, res:-res] = plot_balanced_tree(model_output).transpose((0, 2, 3, 1))

    # insert start and goal frame
    if inputs is not None:
        im[:, :res, :res] = imgtensor2np(inputs.traj_seq[:n_logged_samples, 0], n_logged_samples).transpose(0, 2, 3, 1)
        im[:, :res, -res:] = imgtensor2np(batchwise_index(inputs.traj_seq[:n_logged_samples], model_output.end_ind[:n_logged_samples]),
                                          n_logged_samples).transpose(0, 2, 3, 1)

    return im


def plot_gt_matching_overview(model_output, inputs, plot_attr='match_dist'):
    """Plots overview of which predicted frames contributed to which subgoals."""
    if len(inputs.traj_seq_images.shape) > 2:
        assert inputs.traj_seq_images.shape[3] == inputs.traj_seq_images.shape[4]     # code can only handle square images atm
    batch, n_gt, channels, res, _ = inputs.traj_seq_images.shape
    n_logged_samples = PARAMS.n_logged_samples
    assert batch >= n_logged_samples
    tree = model_output.tree
    max_depth = tree.depth
    n_sg = (2**max_depth) - 1

    im_height = (n_gt+max_depth) * res
    im_width = (n_sg + 2) * res
    im = np.asarray(0.7 * np.ones((n_logged_samples, im_height, im_width, 3)), dtype=np.float32)

    # insert ground truth images on the left, soft estimates on the right, top to bottom
    # insert raw subgoal predictions tree at the bottom, left to right in depth-first order
    get_strip = lambda x, gt=False: imgtensor2np(x, n_logged_samples, gt).transpose(0, 1, 3, 4, 2)\
        .reshape(n_logged_samples, res * n_gt, res, channels)
    
    im[:, :res*n_gt, :res] = get_strip(inputs.traj_seq_images, gt=True)
    if 'soft_matched_estimates' in model_output:
        im[:, :res * n_gt, res:2*res] = get_strip(model_output.soft_matched_estimates)
    
    with param(n_logged_samples=n_logged_samples):
        im[:, -max_depth*res:, 2*res:] = plot_balanced_tree(model_output).transpose((0, 2, 3, 1))

    fig_height, fig_width = res, n_sg * res
    dpi = 10
    match_dists = tensor2np(tree.get_attr_df(plot_attr), n_logged_samples)
    for i in range(n_gt):
        for b in range(n_logged_samples):
            match_plot = plot_dists([match_dists[b, :, i]], fig_height, fig_width, dpi, linewidth=3.0)
            im[b, i*res : (i+1)*res, 2*res:] = match_plot

    return im


def plot_pruned_seqs(model_output, inputs, n_logged_samples=3, max_seq_len=None):
    """Plots the pruned output sequences of the SH-Pred model."""
    assert "images" in model_output.dense_rec       # need pruned predicted images of SH-Pred model
    if inputs is not None:
        batch, n_gt_imgs, channels, res, _ = inputs.traj_seq.shape
    else:
        batch = len(model_output.dense_rec.images)
        assert batch == 1      # con currently only handle batch size 1
        n_gt_imgs, channels, res, _ = model_output.dense_rec.images[0].shape
    MAX_SEQ_LEN = int(n_gt_imgs * 1.5) if not max_seq_len else max_seq_len

    im_height = 2 * res
    im_width = (MAX_SEQ_LEN+1) * res
    im = np.asarray(0.7 * np.ones((n_logged_samples, im_height, im_width, 3)), dtype=np.float32)
    pred_imgs = list(map(imgtensor2np, model_output.dense_rec.images[:n_logged_samples]))
    max_len = min(n_gt_imgs, MAX_SEQ_LEN)
    for b in range(n_logged_samples):
        if pred_imgs[b] is None: continue
        seq_len = min(pred_imgs[b].shape[0], MAX_SEQ_LEN)
        max_len = max(max_len, seq_len)
        im[b, -res:, res:(seq_len+1)*res] = pred_imgs[b][:seq_len].transpose(2, 0, 3, 1).reshape(res, seq_len*res, channels)
    if inputs is not None:
        im[:, :res, :(n_gt_imgs*res)] = imgtensor2np(inputs.traj_seq, n_logged_samples).transpose(0, 3, 1, 4, 2)\
                                            .reshape(n_logged_samples, res, n_gt_imgs*res, channels)
    if "actions" in model_output.dense_rec \
            and model_output.dense_rec.actions is not None \
            and (True in [a is not None for a in model_output.dense_rec.actions]) \
            and (inputs is None or inputs.actions.shape[-1] == 2):
        ac_im = np.asarray(0.7 * np.ones((n_logged_samples, res, im_width, 3)), dtype=np.float32)
        pred_ac = list(map(tensor2np, model_output.dense_rec.actions[:n_logged_samples]))
        for b in range(n_logged_samples):
            if pred_ac[b] is None: continue
            seq_len = min(pred_ac[b].shape[0], MAX_SEQ_LEN)
            ac_im[b, :, :seq_len*res] = batch_action2img(pred_ac[b][None, :seq_len], res, channels).transpose(0, 2, 1, 3, 4)\
                .reshape(res, seq_len*res, channels)
        im = np.concatenate((im, ac_im), axis=1)

    # prepare GIF version
    gif_imgs = np.swapaxes(im.reshape(n_logged_samples, im.shape[1], MAX_SEQ_LEN+1, res, channels), 0, 2)[:max_len+1] \
                    .reshape(max_len+1, im.shape[1], res * n_logged_samples, channels)

    return im, gif_imgs


def unstack(arr):
        arr = np.split(arr, arr.shape[0], 0)
        arr = [a.squeeze() for a in arr]
        return arr


def plot_inverse_model_actions(model_output, inputs, n_logged_samples=5):
    #assert inputs.actions.shape[-1] == 2  # code can only handle 2-dim actions
    batch, n_gt_imgs, channels, res, _ = inputs.traj_seq.shape

    def make_row(arr):
        """stack images in a row along batch dimension"""
        return np.concatenate(unstack(arr), 1)

    if len(model_output.action_targets.shape) <= 2:
        action_targets, actions = model_output.action_targets, model_output.actions
    else:
        action_targets, actions = model_output.action_targets[:, 0], model_output.actions[:, 0]
        model_output.img_t0, model_output.img_t1 = inputs.traj_seq[:, 0], inputs.traj_seq[:, 1]
    if action_targets.shape[-1] > 2:
        actions, action_targets = actions[..., :2], action_targets[..., :2]

    input_action_imgs = batch_action2img(tensor2np(action_targets[:, None], n_logged_samples), res, channels)
    pred_action_imgs = batch_action2img(tensor2np(actions[:, None], n_logged_samples), res, channels)

    image_rows = []
    image_rows.append(np.transpose(tensor2np(model_output.img_t0, n_logged_samples), [0, 2, 3, 1]))
    image_rows.append(np.transpose(tensor2np(model_output.img_t1, n_logged_samples), [0, 2, 3, 1]))
    image_rows.append(input_action_imgs.squeeze())
    image_rows.append(pred_action_imgs.squeeze())

    image_rows = [make_row(item) for item in image_rows]

    im = (np.concatenate(image_rows, 0)[None] + 1.0)/2

    return im


def make_gif(seqs, n_seqs_logged=5):
    """Fuse sequences in list + bring in gif format. Uses the imgtensor2np function"""
    seqs = [pad_sequence(seq, batch_first=True) for seq in seqs]
    seqs = [imgtensor2np(s, n_logged_samples=n_seqs_logged) for s in seqs]
    stacked_seqs = seqs[0]
    
    if len(seqs) > 1:
        padding = np.zeros_like(stacked_seqs)[:, :, :, :2]
        padded_seqs = list(np.concatenate([padding, seq], axis=3) for seq in seqs[1:])
        stacked_seqs = np.concatenate([stacked_seqs] + padded_seqs, axis=3)
    
    batch, time, channels, height, width = stacked_seqs.shape
    return stacked_seqs.transpose(1, 2, 3, 0, 4).reshape(time, channels, height, width * batch)


PREV_OBJS = None
def eval_mem_objs():
    """A simple helper function to evaluate the number of objects currently in memory (CPU or GPU) and print the
       difference to the objects in memory when previously calling this function."""
    import gc
    gc.collect()
    param_size, tensor_size = 0, 0
    objs = dict()
    for obj in gc.get_objects():
        try:
            if torch.is_tensor(obj) or (hasattr(obj, 'data') and torch.is_tensor(obj.data)):
                # print(type(obj), obj.size())
                if isinstance(obj, torch.nn.parameter.Parameter):
                    param_size = param_size + 1
                else:
                    tensor_size = tensor_size + 1
                    key = tuple(obj.size())
                    if key in objs:
                        objs[key] = objs[key] + 1
                    else:
                        objs[key] = 1
        except:
            pass
    print("#Params: {}".format(param_size))
    print("#Tensors: {}".format(tensor_size))

    global PREV_OBJS
    if PREV_OBJS is not None:
        diff = dict()
        for key in objs:
            if key in PREV_OBJS:
                d = objs[key] - PREV_OBJS[key]
                if d != 0:
                    diff[key] = d
            else:
                diff[key] = objs[key]

        import pprint
        pprint.pprint(diff)

    PREV_OBJS = objs

<file_sep>import torch
from torch.distributions import OneHotCategorical
import torch.nn as nn

from blox import batch_apply, AttrDict
from blox.tensor.ops import remove_spatial, broadcast_final
from blox.torch.losses import CELogitsLoss
from blox.torch.recurrent_modules import BaseProcessingLSTM
from blox.torch.subnetworks import SeqEncodingModule, Predictor


class AttnKeyEncodingModule(SeqEncodingModule):
    def build_network(self, input_size, hp):
        self.net = Predictor(hp, input_size, hp.nz_attn_key, num_layers=1)

    def forward(self, seq):
        return batch_apply(self.net, seq.contiguous())


class RecurrentPolicyModule(SeqEncodingModule):
    def __init__(self, hp, input_size, output_size, add_time=True):
        super().__init__(hp, False)
        self.hp = hp
        self.output_size = output_size
        self.net = BaseProcessingLSTM(hp, input_size, output_size)

    def build_network(self, input_size, hp):
        pass

    def forward(self, seq):
        sh = list(seq.shape)
        seq = seq.view(sh[:2] + [-1])
        proc_seq = self.run_net(seq)
        proc_seq = proc_seq.view(sh[:2] + [self.output_size] + sh[3:])
        return proc_seq


class LengthPredictorModule(nn.Module):
    """Predicts the length of a segment given start and goal image encoding of that segment."""
    def __init__(self, hp):
        super().__init__()
        self._hp = hp
        self.p = Predictor(hp, hp.nz_enc * 2, hp.max_seq_len)

    def forward(self, e0, eg):
        """Returns the logits of a OneHotCategorical distribution."""
        output = AttrDict()
        output.seq_len_logits = remove_spatial(self.p(e0, eg))
        output.seq_len_pred = OneHotCategorical(logits=output.seq_len_logits)
        
        return output
    
    def loss(self, inputs, model_output):
        losses = AttrDict()
        losses.len_pred = CELogitsLoss(self._hp.length_pred_weight)(model_output.seq_len_logits, inputs.end_ind)
        return losses


class ActionConditioningWrapper(nn.Module):
    def __init__(self, hp, net):
        super().__init__()
        self.net = net
        self.ac_net = Predictor(hp, hp.nz_enc + hp.n_actions, hp.nz_enc)

    def forward(self, input, actions):
        net_outputs = self.net(input)
        padded_actions = torch.nn.functional.pad(actions, (0, 0, 0, net_outputs.shape[1] - actions.shape[1], 0, 0))
        # TODO quite sure the concatenation is automatic
        net_outputs = batch_apply(self.ac_net, torch.cat([net_outputs, broadcast_final(padded_actions, input)], dim=2))
        return net_outputs
<file_sep>import torch
import numpy as np

from blox import AttrDict
from gcp.planning.infra.policy.policy import Policy
from gcp.prediction.models.tree.tree import TreeModel
from gcp.prediction.training.checkpoint_handler import CheckpointHandler
from gcp.planning.cem.cem_planner import ImageCEMPlanner, CEMPlanner
from gcp.planning.cem.cem_simulator import GCPSimulator, \
    GCPImageSimulator, ActCondGCPImageSimulator


class PlannerPolicy(Policy):
    """Policy that uses predictive planning algorithm to devise plan, and then follows it."""
    def __init__(self, ag_params, policyparams, gpu_id=None, ngpu=None, conversion_fcns=None, n_rooms=None):
        """
        :param ag_params: Agent parameters for infrastructure
        :param policyparams: Parameters for the policy, including model parameters
        :param gpu_id: unused arg (to comply with infrastructure definition)
        :param ngpu: unused arg (to comply with infrastructure definition)
        :param conversion_fcns: unused arg (to comply with infrastructure definition)
        :param n_rooms: unused arg (to comply with infrastructure definition)
        """
        super(PlannerPolicy, self).__init__()

        self._hp = self._default_hparams()
        self.override_defaults(policyparams)
        self.verbose = self._hp.verbose
        self.log_dir = ag_params.log_dir
        self._hp.params['batch_size'] = 1
        self.max_seq_len = ag_params.T
        if 'max_seq_len' not in self._hp.params:
            self._hp.params['max_seq_len'] = ag_params.T

        # create planner predictive model
        model = policyparams['model_cls'] if 'model_cls' in policyparams else TreeModel
        self.planner = model(self._hp.params, None)
        assert self.planner._hp.img_sz == ag_params.image_width

        # move planner model to device
        self.use_cuda = torch.cuda.is_available()
        self.device = torch.device('cuda') if self.use_cuda else torch.device('cpu')
        self.planner.to(self.device)
        self.planner.device = torch.device('cuda')
        self.planner._hp.device = self.planner.device

        # load weights for predictive model
        load_epoch = 'latest' if self._hp.load_epoch is None else self._hp.load_epoch
        weights_file = CheckpointHandler.get_resume_ckpt_file(load_epoch, self._hp.checkpt_path)
        CheckpointHandler.load_weights(weights_file, self.planner, strict=False)
        self.planner.eval()

        self.current_exec_step = None
        self.image_plan = None
        self.action_plan = None
        self.planner_outputs = []
        self.num_replans = 0

    def reset(self):
        super().reset()
        self.current_exec_step = None
        self.action_plan = None
        self.image_plan = None
        self.num_replans = 0
        self.planner_outputs = []
        self.img_t0_history = []
        self.img_t1_history = []

    def _default_hparams(self):
        default_dict = {
            'params': {},               # parameters for predictive model
            'model_cls': None,  # class for predictive model
            'checkpt_path': None,       # checkpoint path for predictive model
            'load_epoch': None,         # epoch that weigths should be loaded from
            'logger': None,
            'verbose': False,           # whether verbose planning outputs are logged
            'max_dump_rollouts': 5,     # max number of rollouts to dump
            'replan_interval': 1,       # interval at which replanning is triggered
            'num_max_replans': 10,      # maximum number of replannings per episode
        }

        parent_params = super()._default_hparams()
        parent_params.ncam = 1
        for k in default_dict.keys():
            parent_params.add_hparam(k, default_dict[k])
        return parent_params

    def act(self, t=None, i_tr=None, images=None, goal_image=None):
        """
        Triggers planning if no plan is made yet / last plan is completely executed. Then executes current plan.
        :param t: current time step in task execution
        :param i_tr: index of currently executed task
        :param images: images of so-far executed trajectory
        :param goal_image: goal-image that should be planned towards
        """
        self.t = t
        self.i_tr = i_tr
        self.goal_image = goal_image
        self.log_dir_verb = self.log_dir + '/verbose/traj{}'.format(self.i_tr)
        output = AttrDict()

        if self.image_plan is None \
              or self.image_plan.shape[0] - 1 <= self.current_exec_step \
              or (t % self._hp.replan_interval == 0 and self.num_replans < self._hp.num_max_replans):
            self._plan(images[t], goal_image, t)
            self.num_replans += 1

        output.actions = self.get_action(images[t])
        self.current_exec_step = self.current_exec_step + 1
        return output

    def get_action(self, current_image):
        assert self.action_plan is not None     # need to attach inverse model to planner to get actions!
        action = self.action_plan[self.current_exec_step]
        return action

    def _plan(self, image, goal_image, step):
        """Runs planning algorithm to obtain image and action plans."""
        raise NotImplementedError

    def log_outputs_stateful(self, logger=None, global_step=None, phase=None, dump_dir=None, exec_seq=None, goal=None,
                             index=None, env=None, goal_pos=None, traj=None, topdown_image=None):
        """Logs planner outputs for visualization."""
        raise NotImplementedError


class CEMPolicy(PlannerPolicy):
    """Implements a simple CEM planning policy."""

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self._hp.cem_params.update({'max_seq_len': self._hp.params['max_seq_len']})
        self._cem_simulator = self.simulator_type(self.planner, append_latent=True)
        self._cem_planner = self._hp.cem_planner(self._hp.cem_params, self._cem_simulator)

    def _default_hparams(self):
        default_dict = {
            'cem_planner': None,
            'cem_params': {},
        }
        parent_params = super()._default_hparams()
        parent_params.ncam = 1
        for k in default_dict.keys():
            parent_params.add_hparam(k, default_dict[k])
        return parent_params

    def act(self, t=None, i_tr=None, state=None, goal=None, regression_state=None, images=None, run_super=False):
        """Plans a state trajectory with CEM, output actions are delta-states."""
        if run_super:
            return super().act(t, i_tr, state, goal)
        self._images = images[:, 0]
        self._states = state
        return super().act(t, i_tr, state, goal)

    def _plan(self, state, goal, step):
        """Runs CEM with planner model to generate state/action plan."""
        # run CEM to get state plan
        input_goal = goal[-1] if len(goal.shape) > 1 else goal
        self.image_plan, action_plan, _, self.plan_cost = self._cem_planner(state, input_goal)
        self.current_exec_step = 0  # reset internal execution counter used to index plan

        # compute action plan as difference between states
        self.action_plan = self.image_plan[1:] - self.image_plan[:-1]

    def log_outputs_stateful(self, logger=None, global_step=None, phase=None, dump_dir=None, exec_seq=None, goal=None,
                             index=None, env=None, goal_pos=None, traj=None, topdown_image=None):
        # self._cem_simulator.dump_logs()
        self._cem_planner.log_verbose(logger, global_step, phase, self.i_tr, dump_dir)

        # log executed image sequence
        executed_traj = self._images.astype(np.float32) / 255
        logger.log_video(executed_traj.transpose(0, 3, 1, 2), "elite_trajs_{}_test/execution".format(self.i_tr),
                         global_step, phase)

    def get_action(self, current_image):
        assert self.action_plan is not None  # need to attach inverse model to planner to get actions!
        if self.action_plan.size < 1:
            return 0.05 * np.random.rand(2, )
        action = self.action_plan[self.current_exec_step]
        return action

    @property
    def simulator_type(self):
        return GCPSimulator

    @property
    def planner_type(self):
        return CEMPlanner


class ImageCEMPolicy(CEMPolicy):
    """CEM planning policy for image-based tasks. Uses inverse model to follow plan"""

    def _default_hparams(self):
        default_dict = {
            'closed_loop_execution': False,  # if True, will execute state plan in closed loop
            'act_cond': False,  # if action-conditioned simulator should be used
        }
        parent_params = super()._default_hparams()
        for k in default_dict.keys():
            parent_params.add_hparam(k, default_dict[k])
        return parent_params

    def act(self, t=None, i_tr=None, state=None, images=None, goal_image=None):
        self._images = images[:, 0]
        self._states = state
        return super().act(t, i_tr, images, goal_image, run_super=True)

    def _plan(self, state, goal, step):
        """Planner directly outputs action plan via inverse model."""
        self.image_plan, self.action_plan, self.latent_plan, self.plan_cost = self._cem_planner(state, goal)
        self.current_exec_step = 0
        self._cem_planner.hack_add_state(self._states[-1].copy())

    def get_action(self, current_image):
        """Executes plan, optional closed-loop by re-inferring actions with the inverse model."""
        if self._hp.closed_loop_execution:
            return self._infer_action(current_image, self.latent_plan[self.current_exec_step + 1])
        else:
            return super().get_action(current_image)

    def _infer_action(self, current_img, target_latent):
        """Uses inverse model to infer closed loop execution action."""
        img = torch.tensor(current_img, device=self.device, dtype=torch.float32)
        enc_img0 = self.planner.encoder(self._cem_simulator._env2planner(img))[0][:, :, 0, 0]
        return self.planner.inv_mdl.run_single(
            enc_img0, torch.tensor(target_latent[None], device=self.device))[0].data.cpu().numpy()

    @property
    def simulator_type(self):
        return GCPImageSimulator if not self._hp.act_cond else ActCondGCPImageSimulator

    @property
    def planner_type(self):
        return ImageCEMPlanner

<file_sep>import os
import shutil
import pickle as pkl
import cv2

import copy

class RawSaver():
    def __init__(self, save_dir, ngroup=1000):
        self.save_dir = save_dir
        self.ngroup = ngroup

    def save_traj(self, itr, agent_data=None, obs_dict=None, policy_outputs=None):

        igrp = itr // self.ngroup
        group_folder = self.save_dir + '/raw/traj_group{}'.format(igrp)
        if not os.path.exists(group_folder):
            os.makedirs(group_folder)

        traj_folder = group_folder + '/traj{}'.format(itr)
        if os.path.exists(traj_folder):
            print('trajectory folder {} already exists, deleting the folder'.format(traj_folder))
            shutil.rmtree(traj_folder)

        os.makedirs(traj_folder)
        print('writing: ', traj_folder)

        if 'robosuite_xml' in obs_dict:
            save_robosuite_xml(traj_folder + '/robosuite.xml', obs_dict['robosuite_xml'][-1])

        if 'images' in obs_dict:
            images = obs_dict['images'].copy()
            T, n_cams = images.shape[:2]
            for i in range(n_cams):
                os.mkdir(traj_folder + '/images{}'.format(i))
            for t in range(T):
                for i in range(n_cams):
                    cv2.imwrite('{}/images{}/im_{}.png'.format(traj_folder, i, t), images[t, i, :, :, ::-1])

        if agent_data is not None:
            with open('{}/agent_data.pkl'.format(traj_folder), 'wb') as file:
                pkl.dump(agent_data, file)
        if obs_dict is not None:
            with open('{}/obs_dict.pkl'.format(traj_folder), 'wb') as file:
                obs_dict_cpy = copy.deepcopy(obs_dict)
                if 'topdown_image' in obs_dict_cpy:
                    obs_dict_cpy.pop('topdown_image') # don't save topdown image, takes too much memory!
                pkl.dump(obs_dict_cpy, file)
        if policy_outputs is not None:
            with open('{}/policy_out.pkl'.format(traj_folder), 'wb') as file:
                pkl.dump(policy_outputs, file)




import os
import xml.dom.minidom
import xml.etree.ElementTree as ET
import io

def save_robosuite_xml(fname, xml_str, pretty=False):
    with open(fname, "w") as f:
        if pretty:
            # TODO: get a better pretty print library
            parsed_xml = xml.dom.minidom.parseString(xml_str)
            xml_str = parsed_xml.toprettyxml(newl="")
        f.write(xml_str)<file_sep>import queue

import numpy as np
from blox import AttrDict
from gcp.planning.infra.envs.miniworld_env.utils.multiroom2d_layout import define_layout
from gcp.planning.infra.envs.miniworld_env.utils.sampling_fcns import RoomSampler2d
from gcp.planning.infra.policy.policy import Policy
from gcp.planning.infra.policy.prm_policy.prm import PRM_planning
from scipy import interpolate


class PrmPolicy(Policy):
    """
    Cross Entropy Method Stochastic Optimizer
    """
    VAR_SAMPLING_RATES = [30, 300]

    def __init__(self, ag_params, policyparams, gpu_id, ngpu, conversion_fcns=None, n_rooms=None):
        super(PrmPolicy, self).__init__()

        self._hp = self._default_hparams()
        policyparams['n_rooms'] = n_rooms
        self.override_defaults(policyparams)

        self._rooms_per_side = int(np.sqrt(self._hp.n_rooms))
        self.layout = define_layout(self._rooms_per_side)
        self.state_sampler = RoomSampler2d(self._rooms_per_side, sample_wide=self.layout.non_symmetric)
        self.plan_params = AttrDict(n_knn=self._hp.n_knn,
                                    max_edge_len=self._hp.max_edge_len,
                                    cost_fcn=lambda d: d ** self._hp.cost_power)

        self.current_action = None
        self.state_plan = None
        self.action_plan = None
        self.convert = conversion_fcns      # AttrDict containing env2prm, transform_plan
        self._room_plan = None

    def reset(self):
        self.current_action = None
        self.state_plan = None
        self.action_plan = None
        self._room_plan = None

    def _default_hparams(self):
        default_dict = {
            'n_samples_per_room': 50,  # number of sample_points in first try, then gets increased
            'n_samples_per_door': 3,   # number of samples per door
            'n_knn': 10,  # number of edge from one sampled point
            'max_edge_len': 0.1,  # Maximum edge length (in layout units)
            'replan_eps': 0.05,     # distance btw planned and executed state that triggers replan, in % of table size
            'max_planning_retries': 2,   # maximum number of replans before inverting the last action
            'cost_power': 2,    # power on the distance for cost function
            'bottleneck_sampling': True,    # sample explicitly in bottlenecks to ease planning
            'use_var_sampling': False,  # if True, uses variable PRM sampling rates for different rooms
            'subsample_factor': 1.0,    # how much to subsample the plan in state space
            'max_traj_length': None,     # maximum length of planned trajectory
            'smooth_trajectory': False,     # if True, uses spline interpolation to smooth trajectory
            'sample_door_center': False,     # if True, samples door samples in center position of door
            'use_scripted_path': False,     # if True, uses scripted waypoints to construct path
            'straight_through_door': False,     # if True, crosses through door in a straight line
            'n_rooms': None,   # number of rooms in the layout
            'use_fallback_plan': True,  # if True executes fallback plan if planning fails
            'logger': None,     # dummy variable
        }

        parent_params = super()._default_hparams()
        parent_params.ncam = 1
        for k in default_dict.keys():
            parent_params.add_hparam(k, default_dict[k])
        return parent_params

    def act(self, t=None, i_tr=None, qpos_full=None, goal=None):
        self.i_tr = i_tr
        output = AttrDict()

        if self.action_plan is None or \
                self._check_deviate(qpos_full[t, :2],
                                    self.state_plan[:, min(self.current_action, self.state_plan.shape[1]-1)]):
            self._plan(qpos_full[t], goal[t], t)
            self.current_action = 0

        done = False
        if self.current_action < self.action_plan.shape[1]:
            output.actions = self.action_plan[:, self.current_action]
        else:   # if required number of steps > planned steps
            done = True
            output.actions = np.zeros(2)
        self.current_action = self.current_action + 1
        output.done = done
        return output

    def _sample_uniform(self):
        px, py = [], []
        for _ in range(self._hp.n_samples_per_room * self._hp.n_rooms):
            p = self.state_sampler.sample()
            px.append(p[0]); py.append(p[1])
        return px, py

    def _sample_per_room(self, room_path):
        px, py = [], []
        room_path = range(self._hp.n_rooms) if room_path is None else room_path
        for room in room_path:
            n_samples = int(np.random.choice(PrmPolicy.VAR_SAMPLING_RATES, 1)) if self._hp.use_var_sampling \
                            else self._hp.n_samples_per_room
            for _ in range(n_samples):
                p = self.state_sampler.sample(room)
                px.append(p[0]); py.append(p[1])
        return px, py

    def _sample_per_door(self, room_path=None):
        doors = self.layout if room_path is None else \
            [(min(room_path[i], room_path[i+1]), max(room_path[i], room_path[i+1])) for i in range(len(room_path) - 1)]
        if not doors: return [], []
        samples = np.asarray([[self.state_sampler.sample_door(d[0], d[1], self._hp.sample_door_center)
                               for _ in range(self._hp.n_samples_per_door)]
                              for d in doors]).transpose(2, 0, 1).reshape(2, -1)
        return samples[0], samples[1]

    def _sample_points(self, room_path=None):
        px, py = self._sample_per_room(room_path)
        if self._hp.bottleneck_sampling:
            dx, dy = self._sample_per_door(room_path)
            px.extend(dx); py.extend(dy)
        return [px, py]

    def _check_deviate(self, pos, target_pos):
        print(np.linalg.norm(pos - target_pos))
        return np.linalg.norm(pos - target_pos) > self._hp.replan_eps

    def _plan(self, agent_pos, goal_pos, t):
        ## UNCOMMENT for random exploration policcy
        #from gcp.infra.policy.cem.utils.sampler import PDDMSampler
        #sampler = PDDMSampler(clip_val=float("Inf"), n_steps=self._hp.max_traj_length - t, action_dim=2, initial_std=2)
        #self.action_plan = sampler.sample(1)[0].transpose(1, 0)  #(np.random.rand(2, self._hp.max_traj_length - t) - 0.5) * 2 * 2#* 6e-1
        #self.state_plan = agent_pos[:2][:, None].repeat(self.action_plan.shape[1], axis=1) + np.cumsum(self.action_plan, axis=1)
        #return self.action_plan, True

        if self.convert is not None:
            pos = self.convert.env2prm(agent_pos[:2])
            goal_pos = self.convert.env2prm(goal_pos)
        else:
            pos = agent_pos[:2]

        length, path = self.compute_shortest_path(pos, goal_pos, transform_pose=False)
        if self._hp.use_scripted_path:
            planned_x, planned_y = [p[0] for p in path], [p[1] for p in path]
            success = True
        else:
            sx, sy = pos[0], pos[1]
            gx, gy = goal_pos[0], goal_pos[1]
            ox, oy = self.layout.ox, self.layout.oy
            if self._room_plan is None:
                room_path = self.plan_room_seq(self.layout.coords2ridx(*pos),
                                               self.layout.coords2ridx(*goal_pos), self.layout.doors)
                print("Planned room sequence with {} rooms!".format(len(room_path)))
                self._room_plan = room_path
            else:
                room_path = self._room_plan
                print("Reused existing room plan!")

            for _ in range(self._hp.max_planning_retries):
                pts = self._sample_points(room_path)
                planned_x, planned_y, success = PRM_planning(sx, sy, gx, gy, ox, oy, self.layout.robot_size, self.plan_params,
                                                             self._hp.n_samples_per_room * self._hp.n_rooms, pts)
                if success: break     # when planning is successful

        if not success:
            if self._hp.use_fallback_plan:
                print("Did not find a plan in {} tries!".format(self._hp.max_planning_retries))
                self._fallback_plan()
            return None, False

        n_steps = min(int(length * 20), self._hp.max_traj_length - t) # min(int(self._hp.subsample_factor * len(planned_x)), self._hp.max_traj_length)
        try:
            if self._hp.max_traj_length is not None: n_steps = min(n_steps, self._hp.max_traj_length - t)
            tck, u = interpolate.splprep([planned_x, planned_y], s=0.0)
            x_i, y_i = interpolate.splev(np.linspace(0, 1, n_steps), tck)
            # x_i, y_i = planned_x, planned_y
            self.state_plan = np.stack((x_i, y_i))
        except TypeError:
            print("Could not interpolate!")     # this happens if duplicate values in plan
            self._fallback_plan()
            return None, False
            #self.state_plan = np.array([planned_x, planned_y])
        self.action_plan = self.state_plan[:, 1:] - self.state_plan[:, :-1]

        raw_plan = self.state_plan.copy()
        if self.convert is not None:
            self.state_plan, self.action_plan = self.convert.transform_plan(self.state_plan, self.action_plan)
        return raw_plan, True

    def _fallback_plan(self):
        if self.action_plan is not None:
            self.action_plan = -2 * self.action_plan[:, self.current_action-1:]    # TODO: come up with a better fallback solution!)
        else:
            self.action_plan = self.state_plan = 0.02 * np.random.rand(2, 1)

    def compute_shortest_path(self, p1, p2, transform_pose=True, straight_through_door=False):
        if self.convert is not None and transform_pose:
            p1, p2 = self.convert.env2prm(p1), self.convert.env2prm(p2)
        if (np.stack((p1,p2)) < -0.5).any() or (np.stack((p1,p2)) > 0.5).any():
            return 10., []      # coordinates invalid
        room_path = plan_room_seq(self.layout.coords2ridx(p1[0], p1[1]),
                                  self.layout.coords2ridx(p2[0], p2[1]), self.layout.doors)
        waypoints = [p1]
        for n in range(len(room_path)-1):
            # input rooms must be in ascending order
            if straight_through_door:
                waypoints.extend(self.state_sampler.get_door_path(room_path[n], room_path[n + 1]))
            else:
                waypoints.append(self.state_sampler.get_door_pos(min(room_path[n], room_path[n+1]),
                                                                 max(room_path[n], room_path[n+1])))
        waypoints.append(p2)
        waypoints = np.array(waypoints)
        length = np.linalg.norm(waypoints[1:] - waypoints[:-1], axis=1).sum()
        return length, waypoints

    def plan_room_seq(self, *args, **kwargs):
        if self.layout.multimodal:
            return plan_room_seq_multimodal(*args, **kwargs)
        else:
            return plan_room_seq(*args, **kwargs)

    def avg_step_length(self, px, py):
        return np.mean(np.sqrt((np.array(px[1:]) - np.array(px[:-1]))**2 + (np.array(py[1:]) - np.array(py[:-1]))**2))


def plan_room_seq(start, goal, doors):
    """Implements a breadth-first room search to find the sequence of rooms that reaches the goal."""
    frontier = queue.Queue()
    visited = []

    def expand(node):
        if node.room == goal: return node
        visited.append(node.room)
        neighbors = []
        for d in doors:
            if d[0] == node.room and d[1] not in visited:
                neighbors.append(d[1])
            elif d[1] == node.room and d[0] not in visited:
                neighbors.append(d[0])
        [frontier.put(AttrDict(room=neighbor, parent=node)) for neighbor in neighbors]
        return expand(frontier.get())

    linked_path = expand(AttrDict(room=start, parent=None))
    room_path = []

    def collect(node):
        room_path.append(node.room)
        if node.parent is None: return
        collect(node.parent)

    collect(linked_path)
    return room_path[::-1]


def plan_room_seq_multimodal(start, goal, doors):
    """Finds all paths between start and goal that visit each room at most once. Returns one of them at random."""
    frontier = queue.Queue()
    goal_nodes = []

    def collect_path(start_node):
        room_path = []
        def collect(node):
            room_path.append(node.room)
            if node.parent is None: return
            collect(node.parent)
        collect(start_node)
        return room_path

    def expand(node):
        if node.room == goal:
            goal_nodes.append(node)
        else:
            neighbors = []
            for d in doors:
                if d[0] == node.room and d[1] not in collect_path(node):
                    neighbors.append(d[1])
                elif d[1] == node.room and d[0] not in collect_path(node):
                    neighbors.append(d[0])
            [frontier.put(AttrDict(room=neighbor, parent=node)) for neighbor in neighbors]
        if frontier.empty(): return
        expand(frontier.get())

    # collect list of all possible paths, is sorted by length (short to long)
    expand(AttrDict(room=start, parent=None))

    # sample one of the possible paths at random
    return collect_path(np.random.choice(goal_nodes))[::-1]


if __name__ == "__main__":
    layout = define_layout(3)
    room_seq = plan_room_seq_multimodal(0, 8, layout.doors)
    print(room_seq)



<file_sep>from blox import AttrDict
from experiments.prediction.base_configs import base_tree as base_conf

configuration = AttrDict(base_conf.configuration)
configuration.metric_pruning_scheme = 'pruned_dtw'

model_config = AttrDict(base_conf.model_config)
model_config.update({
    'matching_type': 'balanced',
})
<file_sep>import os

from blox import AttrDict
from gcp.planning.cem.cost_fcn import EuclideanPathLength

current_dir = os.path.dirname(os.path.realpath(__file__))
from experiments.prediction.base_configs import gcp_sequential as base_conf

configuration = AttrDict(base_conf.configuration)
configuration.update({
    'dataset_name': 'nav_9rooms',
    'batch_size': 16,
    'lr': 2e-4,
    'epoch_cycles_train': 2,
    'n_rooms': 9,
    'metric_pruning_scheme': 'basic',
})

model_config = AttrDict(base_conf.model_config)
model_config.update({
    'ngf': 16,
    'nz_mid_lstm': 512,
    'n_lstm_layers': 3,
    'nz_mid': 128,
    'nz_enc': 128,
    'nz_vae': 256,
    'regress_length': True,
    'attach_state_regressor': True,
    'attach_cost_mdl': True,
    'cost_mdl_params': AttrDict(
        cost_fcn=EuclideanPathLength,
    ),
    'attach_inv_mdl': True,
    'inv_mdl_params': AttrDict(
        n_actions=2,
        use_convs=False,
        build_encoder=False,
    ),
    'decoder_distribution': 'discrete_logistic_mixture',
})
model_config.pop("add_weighted_pixel_copy")<file_sep>import numpy as np
import torch
from scipy.spatial.distance import cdist
from math import isinf
try:
    import gcp.evaluation.cutils as cutils
except:
    pass


def dtw_dist(x, y, dist=None, warp=1, w=np.inf, s=1.0):
    """
    Computes Dynamic Time Warping (DTW) of two sequences.

    :param array x: N1*M array
    :param array y: N2*M array
    :param func dist: distance used as cost measure
    :param int warp: how many shifts are computed.
    :param int w: window size limiting the maximal distance between indices of matched entries |i,j|.
    :param float s: weight applied on off-diagonal moves of the path. As s gets larger, the warping path is increasingly biased towards the diagonal
    Returns the minimum distance, the cost matrix, the accumulated cost matrix, and the wrap path.
    """
    r, c = len(x), len(y)
    D1 = np.zeros((r, c))
    for i in range(r):
        for j in range(c):
            # if (isinf(w) or (max(0, i - w) <= j <= min(c, i + w))):
            D1[i, j] = dist(x[i], y[j])

    return dtw(D1, warp, w, s)


def dtw(inp_D0, warp=1, w=np.inf, s=1.0):
    """
    Computes Dynamic Time Warping (DTW) of two sequences.

    :param int warp: how many shifts are computed.
    :param int w: window size limiting the maximal distance between indices of matched entries |i,j|.
    :param float s: weight applied on off-diagonal moves of the path. As s gets larger, the warping path is increasingly biased towards the diagonal
    Returns the minimum distance, the cost matrix, the accumulated cost matrix, and the wrap path.
    """
    r, c = inp_D0.shape
    assert w >= abs(r - c)
    assert s > 0
    if not isinf(w):
        D0 = np.full((r + 1, c + 1), np.inf)
        for i in range(1, r + 1):
            D0[i, max(1, i - w):min(c + 1, i + w + 1)] = 0
        D0[0, 0] = 0
    else:
        D0 = np.zeros((r + 1, c + 1))
        D0[0, 1:] = np.inf
        D0[1:, 0] = np.inf

    D1 = D0[1:, 1:]  # view
    D0[1:, 1:] = inp_D0  # TODO to support w, this needs to be modified to mask the assignment.
    C = D1.copy()
    jrange = range(c)
    for i in range(r):
        if not isinf(w):
            jrange = range(max(0, i - w), min(c, i + w + 1))
        for j in jrange:
            min_list = [D0[i, j]]
            for k in range(1, warp + 1):
                i_k = min(i + k, r)
                j_k = min(j + k, c)
                min_list += [D0[i_k, j] * s, D0[i, j_k] * s]
            D1[i, j] += min(min_list)
    if r == 1:
        path = np.zeros(c), range(c)
    elif c == 1:
        path = range(r), np.zeros(r)
    else:
        path = _traceback(D0)
    return D1[-1, -1] / sum(D1.shape), C, D1, path

def basic_dtw(C):
    """
    Computes Dynamic Time Warping (DTW) of two sequences.

    Returns the minimum distance, the cost matrix, the accumulated cost matrix, and the wrap path.

    :param C: the cost matrix
    :return:
    """
    r, c = C.shape

    D = np.zeros((r + 1, c + 1))
    D[0, 1:] = np.inf
    D[1:, 0] = np.inf
    D[1:, 1:] = C
    for i in range(r):
        for j in range(c):
            candidates = [D[i, j], D[i + 1, j], D[i, j + 1]]
            D[i + 1, j + 1] += min(candidates)
    path = _traceback(D)
    return D[-1, -1] / (r + c), D[1:, 1:], path

def c_dtw(C):
    """
        Computes Dynamic Time Warping (DTW) of two sequences efficiently in C.

        Returns the minimum distance, the cost matrix, the accumulated cost matrix, and the wrap path.

        :param C: the cost matrix
        :return:
        """
    r, c = C.shape

    D = np.zeros((r + 1, c + 1))
    D[0, 1:] = np.inf
    D[1:, 0] = np.inf
    D[1:, 1:] = C
    cutils.min_cumsum(D)
    path = _traceback(D)
    return D[-1, -1] / (r + c), D[1:, 1:], path


def batched_dtw(C, end_ind):
    b, r, c = C.shape
    D = np.zeros((b, r + 1, c + 1))
    D[:, 0, 1:] = np.inf
    D[:, 1:, 0] = np.inf
    D[:, 1:, 1:] = C
    for i in range(r):
       for j in range(c):
           candidates = [D[:, i, j], D[:, i + 1, j], D[:, i, j + 1]]
           D[:, i + 1, j + 1] += np.min(np.stack(candidates), axis=0)
    paths, path_lengths = _batched_traceback(D, end_ind)
    return D[np.arange(b), -1, end_ind+1] / (r + end_ind+1), D[:, 1:, 1:], paths, path_lengths


def torch_dtw(C, end_ind):
    b, r, c = C.shape
    D = torch.zeros((b, r + 1, c + 1))
    D[:, 0, 1:] = torch.Tensor([float("Inf")])
    D[:, 1:, 0] = torch.Tensor([float("Inf")])
    D[:, 1:, 1:] = C

    for i in range(r):
        for j in range(c):
            candidates = [D[:, i, j], D[:, i + 1, j], D[:, i, j + 1]]
            D[:, i + 1, j + 1].add_(torch.min(torch.stack(candidates), dim=0).values)

    paths, path_lengths = _torched_traceback(D, end_ind)
    return D[torch.arange(b), -1, (end_ind.float()+1).long()] / (r + end_ind.float()+1), D[:, 1:, 1:], paths, path_lengths


def accelerated_dtw(x, y, dist=None, inp_D0=None, warp=1):
    """
    Computes Dynamic Time Warping (DTW) of two sequences in a faster way.
    Instead of iterating through each element and calculating each distance,
    this uses the cdist function from scipy (https://docs.scipy.org/doc/scipy/reference/generated/scipy.spatial.distance.cdist.html)

    :param array x: N1*M array
    :param array y: N2*M array
    :param string or func dist: distance parameter for cdist. When string is given, cdist uses optimized functions for the distance metrics.
    If a string is passed, the distance function can be 'braycurtis', 'canberra', 'chebyshev', 'cityblock', 'correlation', 'cosine', 'dice', 'euclidean', 'hamming', 'jaccard', 'kulsinski', 'mahalanobis', 'matching', 'minkowski', 'rogerstanimoto', 'russellrao', 'seuclidean', 'sokalmichener', 'sokalsneath', 'sqeuclidean', 'wminkowski', 'yule'.
    :param int warp: how many shifts are computed.
    Returns the minimum distance, the cost matrix, the accumulated cost matrix, and the wrap path.
    """
    assert len(x)
    assert len(y)
    if np.ndim(x) == 1:
        x = x.reshape(-1, 1)
    if np.ndim(y) == 1:
        y = y.reshape(-1, 1)
    r, c = len(x), len(y)
    D0 = np.zeros((r + 1, c + 1))
    D0[0, 1:] = np.inf
    D0[1:, 0] = np.inf
    D1 = D0[1:, 1:]
    if inp_D0 is not None:
        D0[1:, 1:] = inp_D0
    else:
        D0[1:, 1:] = cdist(x, y, dist)
    C = D1.copy()
    for i in range(r):
        for j in range(c):
            min_list = [D0[i, j]]
            for k in range(1, warp + 1):
                min_list += [D0[min(i + k, r), j],
                             D0[i, min(j + k, c)]]
            D1[i, j] += min(min_list)
    if len(x) == 1:
        path = np.zeros(len(y)), range(len(y))
    elif len(y) == 1:
        path = range(len(x)), np.zeros(len(x))
    else:
        path = _traceback(D0)
    return D1[-1, -1] / sum(D1.shape), C, D1, path


def _traceback(D):
    # TODO I suspect this doesn't work with fancy stuff (w, s, warp)
    i, j = np.array(D.shape) - 2
    p, q = [i], [j]
    while (i > 0) or (j > 0):
        tb = np.argmin((D[i, j], D[i, j + 1], D[i + 1, j]))
        if tb == 0:
            i -= 1
            j -= 1
        elif tb == 1:
            i -= 1
        elif tb == 2:
            j -= 1
        else:
            raise ValueError

        p.insert(0, i)
        q.insert(0, j)
    return np.array(p), np.array(q)


def _batched_traceback(D, end_ind):
    b, r, c = D.shape
    i, j = np.asarray(np.ones((b,)) * (r - 2), dtype=int), end_ind
    p, q = [i.copy()], [j.copy()]
    path_lengths = np.zeros_like(i)
    cnt = 0
    while (i > 0).any() or (j > 0).any():
        cnt += 1
        path_lengths[(i == 0) & (j == 0) & (path_lengths == 0)] = cnt
        tb = np.argmin(np.stack((D[np.arange(b), i, j], D[np.arange(b), i, j + 1], D[np.arange(b), i + 1, j])), axis=0)
        i[(tb == 0) & (i > 0)] -= 1
        j[(tb == 0) & (j > 0)] -= 1
        i[(tb == 1) & (i > 0)] -= 1
        j[(tb == 2) & (j > 0)] -= 1
        p.insert(0, i.copy())
        q.insert(0, j.copy())
    return (np.array(p), np.array(q)), path_lengths


def _torched_traceback(D, end_ind):
    b, r, c = D.shape
    i, j = (torch.ones((b,)) * (r - 2)).long(), end_ind
    p, q = [i.clone()], [j.clone()]
    path_lengths = torch.zeros_like(i)
    cnt = 0
    while (i > 0).any() or (j > 0).any():
        cnt += 1
        path_lengths[(i == 0) & (j == 0) & (path_lengths == 0)] = cnt
        tb = torch.argmin(
            torch.stack((D[torch.arange(b), i, j], D[torch.arange(b), i, j + 1], D[torch.arange(b), i + 1, j])), dim=0)
        i[(tb == 0) & (i > 0)] -= 1
        j[(tb == 0) & (j > 0)] -= 1
        i[(tb == 1) & (i > 0)] -= 1
        j[(tb == 2) & (j > 0)] -= 1
        p.insert(0, i.clone())
        q.insert(0, j.clone())
    return (torch.stack(p), torch.stack(q)), path_lengths


if __name__ == "__main__":
    b, r, c = 8, 1024, 1000
    min_length = int(c - 1)

    EPS = 1e-5
    import numpy as np
    import time
    np.random.seed(40)
    DD = np.random.rand(b, r, c)
    end_ind = min_length + np.asarray(np.random.rand(b) * (c - min_length - 1), dtype=int)
    
    dd, dd2, pp, pp2, t1, t2 = [], [], [], [], 0.0, 0.0
    for D, i in zip(DD, end_ind):
        s = time.time()
        d, cost_matrix, acc_cost_matrix, path = dtw(D[:, :i+1])
        t1 += time.time() - s
        dd.append(d); pp.append(path)
        s = time.time()
        d2, acc_cost_matrix_2, path_2 = c_dtw(D[:, :i+1])
        t2 += time.time() - s
        dd2.append(d2); pp2.append(path_2)
    print("DTW: {}".format(t1))
    print("C DTW: {}".format(t2))
    
    def check(cond, name):
        print("{}: PASS".format(name)) if cond else print("{}: FAIL".format(name))

    check(not np.any((np.array(dd) - dd2) > EPS), "Distance")
    check(not np.any(np.concatenate([(np.array(pp[i][0]) - np.array(pp2[i][0])) > EPS for i in range(b)])) and \
          not np.any(np.concatenate([(np.array(pp[i][1]) - np.array(pp2[i][1])) > EPS for i in range(b)])), "Paths")

<file_sep>import glob
import os
import random
from random import Random

import cv2
import numpy as np
from blox import AttrDict

# don't change these values, they need to correspond to the multiroom2d.xml file!
ROBOT_SIZE = 0.02
ROOM_SIZE = 1/3
DOOR_SIZE = 1.5 * 0.0667

MAZE_SEED = 42
MULTIMODAL = True
NON_SYMMETRIC = False #True


def define_layout_raw(rooms_per_side, _add_horizontal_line=None, _add_vertical_line=None):

    if _add_vertical_line is None:
        coord_offset = 0.5 * rooms_per_side * ROOM_SIZE     # center around (0,0)

        def _add_horizontal_line(x_range, y):
            ox = np.linspace(x_range[0], x_range[1], int((x_range[1] - x_range[0] + 1) * 100))
            oy = y * np.ones_like(ox)
            return np.stack([ox, oy], axis=0) - coord_offset

        def _add_vertical_line(y_range, x):
            oy = np.linspace(y_range[0], y_range[1], int((y_range[1] - y_range[0] + 1) * 100))
            ox = x * np.ones_like(oy)
            return np.stack([ox, oy], axis=0) - coord_offset

    # add outer boundaries
    table_size = ROOM_SIZE * rooms_per_side
    o = _add_horizontal_line([0, table_size], 0)
    o = np.concatenate((o, _add_horizontal_line([0, table_size], table_size)), axis=1)
    o = np.concatenate((o, _add_vertical_line([0, table_size], 0)), axis=1)
    o = np.concatenate((o, _add_vertical_line([0, table_size], table_size)), axis=1)

    # add wall segments
    rng = Random()
    rng.seed(MAZE_SEED)
    for wall_add_fcn in [_add_horizontal_line, _add_vertical_line]:
        for r in range(rooms_per_side):
            o = np.concatenate((o, wall_add_fcn([0, 1 * ROOM_SIZE/2 - DOOR_SIZE/2], (r+1) * ROOM_SIZE)), axis=1)
            for seg_idx in range(rooms_per_side - 1):
                if NON_SYMMETRIC and rng.random() < 0.1: continue
                o = np.concatenate((o, wall_add_fcn(
                    [(2*seg_idx+1) * ROOM_SIZE/2 + DOOR_SIZE/2, (2*(seg_idx+1)+1) * ROOM_SIZE/2 - DOOR_SIZE/2],
                    (r+1) * ROOM_SIZE)), axis=1)
            o = np.concatenate((o, wall_add_fcn([(rooms_per_side-0.5)*ROOM_SIZE + DOOR_SIZE/2, rooms_per_side*ROOM_SIZE],
                                                         (r + 1) * ROOM_SIZE)), axis=1)

    # generate maze and add doors
    doors = gen_doors_multimodal(rooms_per_side) if MULTIMODAL else generate_maze(rooms_per_side)
    for rx in range(rooms_per_side):
        for ry in range(rooms_per_side):
            if rx + 1 < rooms_per_side and \
                    (((rx, ry), (rx+1, ry)) not in doors and ((rx+1, ry), (rx, ry)) not in doors):
                door_center = ROOM_SIZE/2 + ry * ROOM_SIZE
                o = np.concatenate((o, _add_vertical_line([door_center - DOOR_SIZE/2, door_center + DOOR_SIZE/2],
                                   (rx + 1) * ROOM_SIZE)), axis=1)
            if ry + 1 < rooms_per_side and \
                    (((rx, ry), (rx, ry+1)) not in doors and ((rx, ry+1), (rx, ry)) not in doors):
                door_center = ROOM_SIZE/2 + rx * ROOM_SIZE
                o = np.concatenate((o, _add_horizontal_line([door_center - DOOR_SIZE/2, door_center + DOOR_SIZE/2],
                                   (ry + 1) * ROOM_SIZE)), axis=1)

    def coords2ridx(x, y):
        return x * rooms_per_side + (rooms_per_side-1) - y

    # translate to idx and make sure that smaller room idx comes first
    doors = [sorted((coords2ridx(d[0][0], d[0][1]), coords2ridx(d[1][0], d[1][1]))) for d in doors]

    return o, ROBOT_SIZE, table_size, doors


def generate_maze(rooms_per_side):
    """Returns a set of doors that, when open, generate a maze without shortcuts.
       Algorithm from here: https://github.com/maximecb/gym-miniworld"""
    doors = []
    rng = Random()
    rng.seed(MAZE_SEED)

    visited = []
    neighbors = [(0, 1), (1, 0), (-1, 0), (0, -1)]

    def visit(x, y):
        visited.append((x, y))
        rng.shuffle(neighbors)
        for dx, dy in neighbors.copy():
            nx, ny = x + dx, y + dy
            if nx < 0 or nx >= rooms_per_side or ny < 0 or ny >= rooms_per_side:
                continue    # not a valid neighbor
            if (nx, ny) in visited:
                continue    # neighbor already in visited states
            doors.append(((x, y), (nx, ny)))      # open door to neighbor
            visit(nx, ny)

    visit(0, 0)     # generate maze starting from room 0
    return doors


def gen_doors_multimodal(rooms_per_side):
    """Generates open layout with many doors that allows for multimodal trajectories."""
    # generate list of all doors
    doors = []
    neighbors = [(0, 1), (1, 0), (-1, 0), (0, -1)]

    def add_doors(x, y):
        for dx, dy in neighbors.copy():
            nx, ny = x + dx, y + dy
            if nx < 0 or nx >= rooms_per_side or ny < 0 or ny >= rooms_per_side:
                continue    # not a valid neighbor
            if ((x, y), (nx, ny)) in doors or ((nx, ny), (x, y)) in doors:
                continue
            doors.append(((x, y), (nx, ny)))
            add_doors(nx, ny)

    add_doors(0, 0)

    def maybe_remove(r1, r2):
        if (r1, r2) in doors:
            doors.remove((r1, r2))
        elif (r2, r1) in doors:
            doors.remove((r2, r1))

    # remove a few doors (i.e. add walls)
    if rooms_per_side == 3:     # add two walls around the middle room
        maybe_remove((1, 1), (1, 2))
        maybe_remove((1, 1), (1, 0))
    elif rooms_per_side == 4:
        maybe_remove((0, 3), (1, 3))
        maybe_remove((1, 0), (2, 0))
        maybe_remove((2, 1), (3, 1))
        maybe_remove((2, 2), (3, 2))
        maybe_remove((2, 3), (3, 3))
        maybe_remove((1, 3), (1, 4))
    elif rooms_per_side == 5:
        maybe_remove((0, 3), (1, 3))
        maybe_remove((1, 0), (2, 0))
        maybe_remove((2, 1), (3, 1))
        maybe_remove((2, 2), (3, 2))
        maybe_remove((2, 3), (3, 3))
        maybe_remove((1, 3), (1, 4))
        maybe_remove((1, 1), (1, 2))
        maybe_remove((2, 1), (2, 2))
        maybe_remove((3, 1), (3, 2))
        maybe_remove((4, 2), (4, 3))
    else:
        raise NotImplementedError
    return doors


def define_layout(rooms_per_side, texture_dir=None):
    if texture_dir is None:
        texture_dir = default_texture_dir()
    o, robot_size, table_size, doors = define_layout_raw(rooms_per_side)
    ox, oy = list(o[0]), list(o[1])

    def coords2ridx(x, y):
        """Translates float x, y coords into room index."""
        xy_room = [np.floor((c + table_size/2) / ROOM_SIZE) for c in [x, y]]
        ridx = xy_room[0] * rooms_per_side + (rooms_per_side-1) - xy_room[1]
        return int(ridx) if ridx.size == 1 else np.asarray(ridx, dtype=int)

    textures = load_textures(texture_dir) if texture_dir is not None else None

    return AttrDict(ox=ox, 
                    oy=oy, 
                    robot_size=robot_size,
                    table_size=table_size,
                    room_size=ROOM_SIZE,
                    door_size=DOOR_SIZE,
                    doors=doors,
                    coords2ridx=coords2ridx,
                    textures=textures,
                    texture_dir=texture_dir,
                    multimodal=MULTIMODAL,
                    non_symmetric=NON_SYMMETRIC,)


def default_texture_dir():
    texture_dir = None
    for dir in ['nav_9rooms', 'nav_25rooms']:
        path = os.path.join(os.environ["GCP_DATA_DIR"], dir, "textures")
        if os.path.exists(path):
            texture_dir = path
    assert texture_dir is not None   # need to download either 'nav_9room' or 'nav_25room' dataset to get textures
    return texture_dir


def load_textures(texture_dir):
    """Loads all textures from asset folder"""
    texture_files = glob.glob(os.path.join(texture_dir, "*.png"))
    texture_files = [os.path.basename(p) for p in texture_files]
    texture_files.sort()
    rng = random.Random()   # shuffle texture files
    rng.seed(42)
    rng.shuffle(texture_files)
    texture_files.remove("asphalt_1.png")   # this one is used for the floor
    # there was a bug in initial data collection, this is a hack to synchronize with the generated data
    # TODO remove the hardcoded textures when collecting new data
    HARDCODED_TEXTURE_FILES = ['floor_tiles_white.png', 'lg_style_01_4tile_d_result.png', 'lg_style_01_wall_blue_1.png',
                               'wood_1.png', 'lg_style_04_wall_cerise_d_result.png',
                               'lg_style_05_floor_blue_bright_d_result.png', 'cardboard_4.png',
                               'lg_style_03_wall_light_m_result.png',
                               'lg_style_02_wall_dblue_d_result.png',
                               'lg_style_02_wall_purple_d_result.png', 'cinder_blocks_1.png', 'wood_2.png',
                               'ceiling_tiles_1.png',   # to avoid aliasing
                               'lg_style_03_wall_purple_d_result.png', 'airduct_grate_1.png',
                               'lg_style_03_wall_orange_1.png', 'grass_2.png',
                               'lg_style_01_wall_light_m_result.png',
                               'lg_style_04_wall_purple_d_result.png',
                               'lg_style_03_floor_light1_m_result.png',
                               'lg_style_05_wall_red_d_result.png', 'slime_1.png',
                               'lg_style_05_wall_yellow_d_result.png', 'floor_tiles_bw_1.png',
                               'lg_style_02_floor_orange_d_result.png', 'lg_style_05_wall_yellow_bright_d_result.png',
                               'concrete_1.png', 'lg_style_03_wall_gray_d_result.png',
                               'lg_style_04_wall_red_d_result.png',   # to avoid aliasing
                               'lg_style_04_floor_orange_bright_d_result.png',
                               'lg_style_01_floor_orange_bright_d_result.png', 'stucco_1.png',
                               'lg_style_04_wall_green_bright_d_result.png', 'door_steel_brown.png',
                               'lg_style_03_floor_blue_bright_d_result.png', 'lava_1.png',
                               'lg_style_05_floor_light1_m_result.png',
                               'lg_style_01_wall_red_bright_1.png',
                               'lg_style_01_wall_green_1.png', 'lg_style_01_wall_yellow_1.png',
                               'lg_style_01_wall_red_1.png', 'lg_style_02_wall_yellow_d_result.png', 'door_doom_1.png',
                               'wood_planks_1.png', 'lg_style_03_floor_blue_d_result.png',
                               'lg_style_04_floor_blue_d_result.png', 'lg_style_03_floor_orange_d_result.png',
                               'lg_style_04_wall_red_bright_d_result.png', 'lg_style_02_floor_blue_bright_d_result.png',
                               'door_garage_white.png', 'lg_style_04_floor_blue_bright_d_result.png',
                               'lg_style_01_floor_blue_d_result.png',
                               'lg_style_02_floor_light_m_result.png',
                               'marble_2.png', 'lg_style_04_floor_cyan_d_result.png',
                               'lg_style_05_floor_blue_d_result.png', 'lg_style_01_wall_cerise_1.png',
                               'lg_style_02_wall_yellow_bright_d_result.png',
                               'lg_style_01_floor_blue_bright_d_result.png', 'lg_style_04_wall_green_d_result.png',
                               'drywall_1.png', 'lg_style_01_floor_blue_team_d_result.png', 'door_steel_red.png',
                               'lg_style_01_floor_light_m_result.png',
                               'lg_style_03_wall_cyan_1.png', 'marble_1.png',
                               'picket_fence_1.png', 'door_steel_grey.png', 'water_1.png',
                               'lg_style_02_floor_green_d_result.png', 'lg_style_01_floor_orange_d_result.png',
                               'lg_style_01_wall_green_bright_1.png', 'lg_style_03_floor_green_bright_d_result.png',
                               'lg_style_04_floor_orange_d_result.png', 'door_garage_red.png', 'brick_wall_1.png',
                               'lg_style_03_wall_gray_bright_d_result.png', 'lg_style_03_wall_blue_d_result.png',
                               'rock_1.png', 'lg_style_05_wall_red_bright_d_result.png', 'grass_1.png',
                               'lg_style_03_floor_green_d_result.png', 'lg_style_02_floor_green_bright_d_result.png',
                               'lg_style_05_floor_orange_d_result.png', 'door_doom_2.png',
                               'lg_style_02_wall_blue_d_result.png', 'lg_style_04_floor_dorange_d_result.png',
                               'lg_style_03_floor_purple_d_result.png', 'lg_style_05_floor_orange_bright_d_result.png',
                               'lg_style_01_floor_red_team_d_result.png', 'metal_grill_1.png',
                               'lg_style_02_floor_blue_d_result.png', 'cardboard_3.png',
                               'lg_style_01_ceiling_d_result.png', 'lg_style_01_wall_purple_1.png',
                               'lg_style_03_wall_orange_bright_d_result.png',
                               'lg_style_02_wall_blue_bright_d_result.png', 'cardboard_1.png',
                               'ceiling_tile_noborder_1.png', 'lg_style_02_wall_lgreen_d_result.png',
                               'lg_style_03_floor_red_d_result.png']

    return HARDCODED_TEXTURE_FILES


def draw_layout_overview(rooms_per_side, render_scale, texture_dir, add_textures=True):
    textures = load_textures(texture_dir)
    layout = define_layout(rooms_per_side, texture_dir)

    # draw texture background
    n_textures = len(textures)
    res = int(layout.table_size * render_scale)
    room_size = int(res / rooms_per_side)
    img = np.ones((res, res, 3))
    if add_textures:
        for x in range(rooms_per_side):
            for y in range(rooms_per_side):
                texture = cv2.imread(os.path.join(texture_dir,
                                                  textures[(x * rooms_per_side + y) % n_textures]))
                texture = cv2.resize(texture, (room_size, room_size))[:, :, ::-1] / 255.
                img[int(y * room_size) : int((y+1) * room_size),
                    int(x * room_size) : int((x+1) * room_size)] = texture

    def _add_horizontal_line(x_range, y):
        cv2.line(img, (int(x_range[0] * render_scale), res - int(y * render_scale - 1)),
                      (int(x_range[1] * render_scale), res - int(y * render_scale - 1)), (0, 0, 0), 3)
        return [[None]]

    def _add_vertical_line(y_range, x):
        cv2.line(img, (int(x * render_scale), res - int(y_range[0] * render_scale - 1)),
                      (int(x * render_scale), res - int(y_range[1] * render_scale - 1)), (0, 0, 0), 3)
        return [[None]]

    define_layout_raw(rooms_per_side, _add_horizontal_line, _add_vertical_line)
    return img


if __name__ == '__main__':
    import matplotlib.pyplot as plt
    img = draw_layout_overview(rooms_per_side=5, render_scale=256, texture_dir="../../../../../assets/textures")
    # plt.imshow(img)
    # plt.show()
    plt.imsave("test.png", img)

    # import matplotlib.pyplot as plt
    # l = define_layout(10)
    # print(l.doors)
    # plt.scatter(l.ox, l.oy, c='black')
    # plt.axis('equal')
    # plt.show()
<file_sep>from gcp.prediction.models.sequential import SequentialModel
from gcp.prediction.utils.logger import HierarchyLogger

configuration = {
    'model': SequentialModel,
    'logger': HierarchyLogger,
}

model_config = {
    'one_step_planner': 'continuous',
    'dense_rec_type': 'svg',
    'hierarchy_levels': 0,
    'add_weighted_pixel_copy': True,
}<file_sep>dataset_spec = {
	'max_seq_len': 80,      # maximum sequence in dataset is 80 frames
	'n_actions': 2,         # actions in dataset are 2-dimensional
	'state_dim': 2,         # states in dataset are 2-dimensional
}<file_sep>Cython==0.28.1
decorator==4.3.0
dload==0.6
funcsigs==1.0.2
gym==0.14.0
h5py==2.7.1
imageio==2.1.2
matplotlib==2.2.2
moviepy==0.2.3.2
numpy==1.18.4
opencv-python==3.4.2.17
Pillow==6.0.0
pyglet==1.2.4
scikit-image==0.15.0
scipy==1.3.0
tensorflow==1.14.0
tensorboardX==2.0
torch==1.3.0
torchvision==0.2.0
tqdm==4.11.2
urllib3==1.22
numpy-stl==2.11.2
pyquaternion==0.9.5
<file_sep>import torch
import torch.nn as nn

from blox import AttrDict
from blox.tensor.ops import broadcast_final
from blox.torch.layers import BaseProcessingNet
from blox.torch.losses import KLDivLoss2
from blox.torch.models.vrnn import VRNNCell
from gcp.evaluation.evaluation_matching import DTWEvalBinding
from gcp.prediction.models.base_gcp import BaseGCPModel


class SequentialRecModule(nn.Module):

    def __init__(self, hp, input_size, output_size, decoder):
        # TODO make test time version
        assert input_size == output_size
        super().__init__()
        self._hp = hp
        self.decoder = decoder
        
        context_size = 0
        if hp.context_every_step:
            context_size += hp.nz_enc * 2
        if hp.action_conditioned_pred:
            context_size += hp.nz_enc
            
        self.lstm = VRNNCell(hp, input_size, context_size, hp.nz_enc * 2).make_lstm()
        self.eval_binding = DTWEvalBinding(hp)
        if hp.skip_from_parents:
            raise NotImplementedError("SVG doesn't support skipping from parents")
    
    def forward(self, root, inputs):
        lstm_inputs = AttrDict()
        initial_inputs = AttrDict(x=inputs.e_0)
        context = torch.cat([inputs.e_0, inputs.e_g], dim=1)
        static_inputs = AttrDict()

        if 'enc_traj_seq' in inputs:
            lstm_inputs.x_prime = inputs.enc_traj_seq[:, 1:]
        if 'z' in inputs:
            lstm_inputs.z = inputs.z
        if self._hp.context_every_step:
            static_inputs.context = context
        if self._hp.action_conditioned_pred:
            assert 'enc_action_seq' in inputs   # need to feed actions for action conditioned predictor
            lstm_inputs.update(more_context=inputs.enc_action_seq)
        
        self.lstm.cell.init_state(initial_inputs.x, context, lstm_inputs.get('more_context', None))
        # Note: the last image is also produced. The actions are defined as going to the image
        outputs = self.lstm(inputs=lstm_inputs,
                            initial_inputs=initial_inputs,
                            static_inputs=static_inputs,
                            length=self._hp.max_seq_len - 1)
        outputs.encodings = outputs.pop('x')
        outputs.update(self.decoder.decode_seq(inputs, outputs.encodings))
        outputs.images = torch.cat([inputs.I_0[:, None], outputs.images], dim=1)
        return outputs
    
    def loss(self, inputs, outputs, log_error_arr=False):
        losses = self.decoder.loss(inputs, outputs, extra_action=False, log_error_arr=log_error_arr)
        
        # TODO don't place loss on the final image
        weights = broadcast_final(inputs.pad_mask[:, 1:], outputs.p_z.mu)
        losses.kl = KLDivLoss2(self._hp.kl_weight, breakdown=1, free_nats_per_dim=self._hp.free_nats)\
            (outputs.q_z, outputs.p_z, weights=weights, log_error_arr=log_error_arr)
        
        return losses
    
    def get_sample_with_len(self, i_ex, len, outputs, inputs, pruning_scheme, name=None):
        """
        :param i_ex:  example index
        :param input_seq:
        :param outputs:
        :return gen_images
        """
        
        if pruning_scheme == 'dtw':
            # Cut the first image off for DTW - it is the GT image
            targets = inputs.traj_seq[i_ex, 1:inputs.end_ind[i_ex] + 1]
            estimates = outputs.dense_rec.images[i_ex, 1:inputs.end_ind[i_ex] + 1]
            images, matching_output = self.eval_binding(None, None, None, None, targets=targets, estimates=estimates)
            # TODO clean up
            # Add the first image back (eval cuts it off)..
            return torch.cat([outputs.dense_rec.images[i_ex, [1]], images], dim=0), matching_output
        elif pruning_scheme == 'basic':
            if name is None:
                return outputs.dense_rec.images[i_ex, :len], None
            elif name == 'encodings':
                # TODO fix this. This is necessary because the Hierarchical model outputs the first latent too.
                # This concatenates the first encoder latent to compensate
                return torch.cat((inputs.e_0[i_ex][None], outputs.dense_rec[name][i_ex]), 0)[:len], None
            else:
                return outputs.dense_rec[name][i_ex, :len], None

    def get_all_samples_with_len(self, end_idxs, outputs, inputs, pruning_scheme, name=None):
        return [self.get_sample_with_len(b, end_idxs[b] + 1, outputs, inputs, pruning_scheme, name=name)[0]
                for b in range(end_idxs.shape[0])], None
            

class SequentialModel(BaseGCPModel):
    def build_network(self, build_encoder=True):
        super().build_network(build_encoder)
        
        self.dense_rec = SequentialRecModule(
            hp=self._hp, input_size=self._hp.nz_enc, output_size=self._hp.nz_enc, decoder=self.decoder)
        
        if self._hp.action_conditioned_pred:
            self.action_encoder = BaseProcessingNet(self._hp.n_actions, self._hp.nz_mid, self._hp.nz_enc,
                                                    self._hp.n_processing_layers, self._hp.fc_builder)
            
    def predict_sequence(self, inputs, outputs, start_ind, end_ind, phase):
        outputs = AttrDict(dense_rec=self.dense_rec(None, inputs))
        return outputs

    def log_outputs(self, outputs, inputs, losses, step, log_images, phase):
        super().log_outputs(outputs, inputs, losses, step, log_images, phase)

        if log_images:
            if outputs.dense_rec and self._hp.use_convs:
                self._logger.log_dense_gif(outputs, inputs, "dense_rec", step, phase)
                
                log_prior_images = False
                if log_prior_images:
                    # Run the model N times
                    with torch.no_grad(), self.val_mode():
                        rows = list([self(inputs).dense_rec.images for i in range(4)])
                    self._logger.log_rows_gif(rows, "prior_samples", step, phase)

    def get_predicted_pruned_seqs(self, inputs, outputs):
        return [seq[:end_ind+1] for seq, end_ind in zip(outputs.dense_rec.encodings, outputs.end_ind)]
<file_sep>from .general_agent import GeneralAgent
import pickle as pkl
import numpy as np
import cv2
import PIL
from PIL import Image
import os
import glob


class BenchmarkAgent(GeneralAgent):
    N_MAX_TRIALS = 1    # only allow one trial per trajectory when benchmarking
    def __init__(self, hyperparams, start_goal_list=None):
        self._start_goal_confs = hyperparams.get('start_goal_confs', None)
        self.ncam = hyperparams['env'][1].get('ncam', hyperparams['env'][0].default_ncam()) # check if experiment has ncam set, otherwise get env default

        GeneralAgent.__init__(self, hyperparams, start_goal_list)
        self._is_robot_bench = 'robot_name' in self._hp.env[1]
        if not self._is_robot_bench:
            self._hp.gen_xml = (True, 1)        # this was = 1 but that did not work?!

    def _setup_world(self, itr):
        old_ncam = self.ncam
        GeneralAgent._setup_world(self, itr)
        if self._start_goal_confs is not None:
            self._reset_state = self._load_raw_data(itr)
        assert old_ncam == self.ncam, """Environment has {} cameras but benchmark has {}. 
                                            Feed correct ncam in agent_params""".format(self.ncam, old_ncam)

    def _required_rollout_metadata(self, agent_data, traj_ok, t, i_itr):
        GeneralAgent._required_rollout_metadata(self, agent_data, traj_ok, t, i_itr)

        if self._start_goal_confs is not None:
            agent_data.update(self.env.eval())

    def _init(self):
        return GeneralAgent._init(self)

    def _load_raw_data(self, itr):
        """
        doing the reverse of save_raw_data
        :param itr:
        :return:
        """
        if 'robot_name' in self._hp.env[1]:   # robot experiments don't have a reset state
            return None

        if 'iex' in self._hp:
            itr = self._hp.iex

        ngroup = 1000
        igrp = itr // ngroup
        group_folder = '{}/traj_group{}'.format(self._start_goal_confs, igrp)
        traj_folder = group_folder + '/traj{}'.format(itr)

        print('reading from: ', traj_folder)
        num_files = len(glob.glob("{}/images0/*.png".format(traj_folder)))
        assert num_files > 0, " no files found!"

        obs_dict = {}
        demo_images = np.zeros([num_files, self.ncam, self._hp.image_height, self._hp.image_width, 3])
        for t in [0, num_files-1]: #range(num_files):
            for c in range(self.ncam):
                image_file = '{}/images{}/im_{}.png'.format(traj_folder, c, t)
                if not os.path.isfile(image_file):
                    raise ValueError("Can't find goal image: {}".format(image_file))
                img = cv2.imread(image_file)[..., ::-1]
                if img.shape[0] != self._hp.image_height or img.shape[1] != self._hp.image_width:
                    img = Image.fromarray(img)
                    img = img.resize((self._hp.image_height, self._hp.image_width), PIL.Image.BILINEAR)
                    img = np.asarray(img, dtype=np.uint8)
                demo_images[t, c] = img
        self._demo_images = demo_images.astype(np.float32)/255.

        self._goal_image = self._demo_images[-1]

        with open('{}/obs_dict.pkl'.format(traj_folder), 'rb') as file:
            obs_dict.update(pkl.load(file))

        self._goal = self.env.get_goal_from_obs(obs_dict)
        reset_state = self.get_reset_state(obs_dict)

        if os.path.exists(traj_folder + '/robosuite.xml'):
            with open(traj_folder + '/robosuite.xml', "r") as model_f:
                model_xml = model_f.read()

            from robosuite.utils.mjcf_utils import postprocess_model_xml
            xml = postprocess_model_xml(model_xml)
            reset_state['robosuite_xml'] = xml

        return reset_state

    def get_reset_state(self, obs_dict):
        return self.env.get_reset_from_obs(obs_dict)


<file_sep>import numpy as np
import torch
from blox.torch.ops import slice_tensor, reduce_dim
from blox.torch import porch
from blox.basic_types import map_dict, listdict2dictlist
from blox.tensor.ops import batch_apply, make_recursive_list, rmap


class SubgoalTreeLayer:
    def __init__(self, parent=None):
        if isinstance(parent, map):
            raise ValueError("This happens when splitting across multiple GPUs, will get caught above")
        self.child_layer = None
        self.subgoals = None
        self.depth = None
        self.pruned = None
        # self.parent_layer = parent
        self.selected = None
        self.match_eval_idx = None

    def produce_tree(self, inputs, layerwise_inputs, start_inds, end_inds, left_parents, right_parents, producer, depth):
        """no done mask checks, assumes start_ind never None.
            all input tensors are of shape [batch, num_parent_nodes, ...]
        """
        self.depth = depth
        if depth == 0:
            return

        # slice out inputs for this layer

        layer_inputs = rmap(lambda x: depthfirst2layers(reduce_dim(x, dim=1))[-depth].contiguous(), layerwise_inputs)

        out = batch_apply(lambda x: producer.produce_subgoal(inputs, *x, depth=depth),
                          [layer_inputs, start_inds.float(), end_inds.float(), left_parents, right_parents])
        self.subgoals, left_parents, right_parents = out

        self.child_layer = SubgoalTreeLayer(self)
        self.child_layer.produce_tree(inputs,
                                      layerwise_inputs,
                                      rec_interleave([start_inds.float(), self.subgoals.ind.clone()]),
                                      rec_interleave([self.subgoals.ind.clone(), end_inds.float()]),
                                      rec_interleave([left_parents, self.subgoals]),
                                      rec_interleave([self.subgoals, right_parents]),
                                      producer, depth - 1)

    def compute_matching_dists(self, inputs, matching_fcn, left_parents, right_parents):
        """Computes the distribution of matches of subgoals to ground truth frames."""
        self.apply_fn(inputs, matching_fcn, left_parents, right_parents)
        
    def apply_fn(self, inputs, fn, left_parents, right_parents):
        """ Recursively applies fn to the tree.
        
        :param inputs:
        :param fn: a function that takes in (inputs, subgoal, left_parent, right_parent) and outputs a dict
        :param left_parents:
        :param right_parents:
        :return:
        """
        
        if self.depth == 0:
            return
        assert self.subgoals is not None      # need subgoal info to match to ground truth sequence

        self.subgoals.update(batch_apply(fn, inputs, self.subgoals, left_parents, right_parents, unshape_inputs=True))
        self.child_layer.apply_fn(rec_interleave([inputs, inputs]),
                                  fn,
                                  rec_interleave([left_parents, self.subgoals]),
                                  rec_interleave([self.subgoals, right_parents]))

    def __iter__(self):
        """Layer-wise iterator."""
        if self.subgoals is None:
            return
        yield self
        if self.child_layer is not None:
            for l in self.child_layer:
                yield l

    def depth_first_iter(self, current_node=0):
        """Depth-first subgoal iterator."""
        if self.subgoals is None or self.child_layer is None:
            return
        for n in self.child_layer.depth_first_iter(2*current_node):
            yield n
        self.subgoal = rmap(lambda x: x[:, current_node], self.subgoals)
        yield self
        for n in self.child_layer.depth_first_iter(2*current_node+1):
            yield n

    def get_attr_df(self, attr):
        # TODO make this faster
        return torch.stack([node.subgoal[attr] for node in self.depth_first_iter()], 1)
    
    def set_attr_df(self, **kwargs):
        # TODO check
        for name, value in kwargs:
            split = self.split_by_layer_df(value, 1)
            for chunk, node in zip(split, self):
                node[name] = chunk
        
    def get_attr_bf(self, attr):
        return porch.cat([node.subgoals[attr] for node in self], 1)

    def set_attr_bf(self, **kwargs):
        start = 0
        for i, node in enumerate(self):
            node.subgoals.update(rmap(lambda x: x[:,start:start+2**i].contiguous(), kwargs))
            start += 2**i

    def get_leaf_nodes(self):
        if self.depth == 0:
            raise ValueError("Depth 0 tree does not have leaf nodes!")
        elif self.depth == 1:
            return self.subgoals
        else:
            return self.child_layer.get_leaf_nodes()

    @staticmethod
    def cat(*argv):
        tree = SubgoalTreeLayer()
        for attr, val in argv[0].__dict__.items():
            if val is None or np.isscalar(val):
                tree.__dict__[attr] = val
            elif attr == 'subgoals':
                tree.__dict__[attr] = map_dict(concat, listdict2dictlist([d.subgoals for d in argv]))
            elif attr == 'child_layer':
                tree.__dict__[attr] = SubgoalTreeLayer.cat(*[d.child_layer for d in argv])
            else:
                raise ValueError("Cannot handle data type {} during tree concatenation!".format(type(val)))
        return tree

    @staticmethod
    def reduce(*argv):
        """Called inside result gathering for multi-GPU processing"""
        return SubgoalTreeLayer.cat(*argv)

    @staticmethod
    def split_by_layer_df(vals, dim):
        return depthfirst2layers(vals, dim)
        # """Splits depth-first vals into N lists along dimension dim, each containing vals for the corresp. layer."""
        # depth = int(np.log2(vals.shape[dim]) + 1)
        # output = [[] for _ in range(depth)]     # one list per layer
        #
        # def get_elem(l_idx, r_idx, d):
        #     if l_idx == r_idx - 1: return
        #     idx = int((r_idx - l_idx) / 2) + l_idx
        #     output[d].append(vals[:, idx])
        #     get_elem(l_idx, idx, d + 1)
        #     get_elem(idx, r_idx, d + 1)
        #
        # get_elem(-1, vals.shape[dim], 0)
        # return output

    @staticmethod
    def split_by_layer_bf(vals, dim):
        """Splits breadth-first vals into N arrays along dimension dim, each containing vals for the corresp. layer."""
        depth = int(np.log2(vals.shape[dim]) + 1)
        output = []     # one list per layer
        current_idx = 0
        for d in range(depth):
            output.append(vals[:, current_idx : current_idx + int(2**d)])
            current_idx += int(2**d)
        return output

    @property
    def bf(self):
        return AccessWrapper(self, 'bf')
    
    @property
    def df(self):
        return AccessWrapper(self, 'df')

    @property
    def size(self):
        return int(2**self.depth - 1)


class AccessWrapper():
    def __init__(self, obj, type):
        super().__setattr__('tree', obj)
        super().__setattr__('type', type)
        
    def __getattr__(self, item):
        if self.type == 'bf':
            return self.tree.get_attr_bf(item)
        elif self.type == 'df':
            return self.tree.get_attr_df(item)

    def __setattr__(self, key, value):
        if self.type == 'bf':
            return self.tree.set_attr_bf(**{key: value})
        elif self.type == 'df':
            return self.tree.set_attr_df(**{key: value})

    def __getitem__(self, item):
        return getattr(self, item)
        
    def __setitem__(self, key, value):
        return setattr(self, key, value)


def interleave(t1, t2):
    if t1 is None or t2 is None: return None
    assert t1.shape == t2.shape     # can only interleave vectors of equal shape
    return torch.stack((t1, t2), dim=2).view(t1.shape[0], 2*t1.shape[1], *t1.shape[2:])


rec_interleave = make_recursive_list(interleave)


def concat(*argv):
    if argv[0][0] is None: return None
    device = argv[0][0].device
    return torch.cat([v.to(device) for v in argv[0]], dim=0)


def depthfirst2breadthfirst(tensor, dim=1):
    """ Converts a sequence represented depth first to breadth first """
    return torch.cat(depthfirst2layers(tensor, dim), dim)


def depthfirst2layers(tensor, dim=1):
    """ Converts a sequence represented depth first to a list of layers """
    len = tensor.shape[dim]
    depth = np.int(np.log2(len + 1))
    
    slices = []
    for i in range(depth):
        slices.append(slice_tensor(tensor, 0, 2, dim))
        tensor = slice_tensor(tensor, 1, 2, dim)
    
    return list(reversed(slices))


def ind_df2bf(df_indices, depth):
    """ Transforms indices for a depth-first array such that the same elements can be retrieved from the corresponding
    breadth-first array """
    df_indices = (df_indices + 1).bool()  # starting from 1
    bf_indices = torch.zeros_like(df_indices)
    for i in range(depth):
        mask = (df_indices % (2**i) == 0) & (df_indices % (2**(i+1)) > 0)  # if in layer i from the bottom
        bf_indices[mask] = df_indices[mask] // (2**(i+1)) + (2**(depth - i - 1) - 1)  # order in layer + layer position
        
    return bf_indices


def ind_bf2df(bf_indices, depth):
    """ Transforms indices for a breadth-first array such that the same elements can be retrieved from the corresponding
    depth-first array """
    bf_indices = (bf_indices + 1).bool()  # starting from 1
    df_indices = torch.zeros_like(bf_indices)
    for i in range(depth):
        mask = (bf_indices >= 2 ** i) & (bf_indices < 2 ** (i + 1))  # if in layer i from the top
        ib = depth - i - 1  # layer from the bottom
        # order in layer * layer position
        df_indices[mask] = (bf_indices[mask] - 2**i) * (2**(ib+1)) + (2**ib) - 1
    
    return df_indices
<file_sep>import os.path

from blox import AttrDict
from gcp.planning.infra.agent.general_agent import GeneralAgent
from gcp.planning.infra.envs.miniworld_env.multiroom3d.multiroom3d_env import TopdownMultiroom3dEnv
from gcp.planning.infra.policy.prm_policy.prm_policy import PrmPolicy

BASE_DIR = '/'.join(str.split(__file__, '/')[:-1])
current_dir = os.path.dirname(os.path.realpath(__file__))

env_params = {
    'init_pos': None,
    'goal_pos': None,
    'n_rooms': 16,
    'heading_smoothing': 0.1,
    'crop_window': 40,
}

agent = AttrDict(
    type=GeneralAgent,
    env=(TopdownMultiroom3dEnv, env_params),
    T=100,
    make_final_gif=False,   # whether to make final gif
    #make_final_gif_freq=100,   # final gif, frequency
    image_height=128,
    image_width=128,
)

policy = AttrDict(
    type=PrmPolicy,
    max_traj_length=agent.T,
)

config = AttrDict(
    current_dir=current_dir,
    start_index=0,
    end_index=999,
    agent=agent,
    policy=policy,
    save_format=['hdf5'],
    data_save_dir=os.environ['GCP_DATA_DIR'] + '/nav_16rooms',
    split_train_val_test=False,
    traj_per_file=1,
)
<file_sep>import os
import os.path
import sys

from tensorflow.contrib.training import HParams

from gcp.planning.infra.agent.utils.hdf5_saver import HDF5Saver
from gcp.planning.infra.agent.utils.raw_saver import RawSaver
from gcp.prediction.utils.logger import HierarchyLogger

sys.path.append('/'.join(str.split(__file__, '/')[:-2]))
import numpy as np


class Sim(object):
    """ Main class to run algorithms and experiments. """

    def __init__(self, config, gpu_id=0, ngpu=1):
        self._start_goal_list = config.pop('start_goal_list') if 'start_goal_list' in config else None
        self._hp = self._default_hparams()
        self.override_defaults(config)
        self._hp.agent['log_dir'] = self._hp.log_dir
        self._hp.n_rooms = self._hp.agent['env'][1]['n_rooms'] if 'n_rooms' in self._hp.agent['env'][1] else None
        self.agent = self._hp.agent['type'](self._hp.agent, self._start_goal_list)
        self.agentparams = self._hp.agent

        self._record_queue = self._hp.record_saver
        self._counter = self._hp.counter

        if self._hp.logging_conf is None:
            self.logger = HierarchyLogger(self._hp.log_dir + '/verbose', self._hp, self._hp.agent['T'])
            self._hp.logging_conf = {'logger':self.logger, 'global_step':-1, 'phase':'test'}

        self._hp.policy['logger'] = self.logger
        self.policy = self._hp.policy['type'](self.agent._hp, self._hp.policy, gpu_id, ngpu,
                                              **self.agent.env.env_policy_params())

        self.trajectory_list = []
        self.im_score_list = []
        try:
            os.remove(self._hp.agent['image_dir'])
        except:
            pass

        self.savers = []
        if 'hdf5' in self._hp.save_format:
            self.savers.append(HDF5Saver(self._hp.data_save_dir, self.agent.env._hp, self.agent._hp,
                                     traj_per_file=self._hp.traj_per_file, offset=self._hp.start_index,
                                         split_train_val_test=self._hp.split_train_val_test))
        if 'raw' in self._hp.save_format:
            self.savers.append(RawSaver(self._hp.data_save_dir))

        self.logging_conf = self._hp.logging_conf

    def override_defaults(self, config):
        """
        :param config:  override default valus with config dict
        :return:
        """
        for name, value in config.items():
            print('overriding param {} to value {}'.format(name, value))
            if value == getattr(self._hp, name):
                raise ValueError("attribute {} is identical to default value!!".format(name))
            if name in self._hp and self._hp.get(name) is None:   # don't do a type check for None default values
                setattr(self._hp, name, value)
            else: self._hp.set_hparam(name, value)

    def _default_hparams(self):
        default_dict = {
            'save_format': ['hdf5', 'raw'],
            'save_data': True,
            'agent': {},
            'policy': {},
            'start_index': -1,
            'end_index': -1,
            'ntraj': -1,
            'gpu_id': -1,
            'current_dir': '',
            'record_saver': None,
            'counter': None,
            'traj_per_file': 10,
            'data_save_dir': '',
            'log_dir': '',
            'result_dir': '',
            'split_train_val_test': True,
            'logging_conf': None,   # only needed for training loop
        }
        # add new params to parent params
        parent_params = HParams()
        for k in default_dict.keys():
            parent_params.add_hparam(k, default_dict[k])
        return parent_params

    def run(self):
        if self._counter is None:
            for i in range(self._hp.start_index, self._hp.end_index+1):
                self.take_sample(i)
        else:
            itr = self._counter.ret_increment()
            while itr < self._hp.ntraj:
                print('taking sample {} of {}'.format(itr, self._hp.ntraj))
                self.take_sample(itr)
                itr = self._counter.ret_increment()

    def take_sample(self, index):
        """
        :param index:  run a single trajectory with index
        :return:
        """
        self.policy.reset()
        agent_data, obs_dict, policy_out = self.agent.sample(self.policy, index)
        if self._hp.save_data:
            self.save_data(index, agent_data, obs_dict, policy_out)
        if self.logging_conf is not None and 'goal_image' in agent_data and 'images' in obs_dict:
            goal = agent_data['goal_image'][-1, 0] if len(agent_data['goal_image'].shape)>4 else agent_data['goal_image'][-1]
            if 'goal_pos' in obs_dict:
                goal_pos = obs_dict['goal'][-1, :] if isinstance(obs_dict['goal'], np.ndarray) else obs_dict['goal']
            else:
                goal_pos = None
            topdown_image = obs_dict['topdown_image'] if 'topdown_image' in obs_dict else None
            self.policy.log_outputs_stateful(**self.logging_conf, dump_dir=self._hp.log_dir,
                                             exec_seq=obs_dict['images'][:, 0], goal=goal, goal_pos=goal_pos,
                                             index=index, topdown_image=topdown_image, env=self.agent.env)     # [:, 0] for cam0
        return agent_data

    def save_data(self, itr, agent_data, obs_dict, policy_outputs):
        if self._record_queue is not None:  # if using a queue to save data
            self._record_queue.put((agent_data, obs_dict, policy_outputs))
        else:
            for saver in self.savers: # if directly saving data
                saver.save_traj(itr, agent_data, obs_dict, policy_outputs)
<file_sep>import torch.nn as nn

from gcp.prediction.models.tree.tree_module import TreeModule

""" This module contains the hierarchical model in which every layer is predicted by a separately trained network. """


class UntiedLayersTree(nn.Module):
    def __init__(self, hp, decoder):
        super().__init__()
        self._hp = hp
        self.tree_modules = nn.ModuleList([TreeModule(hp, decoder) for i in range(self._hp.hierarchy_levels)])

    def produce_subgoal(self, *args, depth, **kwargs):
        return self.tree_modules[self._hp.hierarchy_levels - depth].produce_subgoal(*args, **kwargs)
    
    def __getattr__(self, item):
        if item in self._modules.keys():
            return super().__getattr__(item)
            
        return getattr(self.tree_modules[0], item)
<file_sep># import ipdb
import glob

import numpy as np
import gcp.planning.infra.datasets.save_util.configs.TAP_3obj_push as config
import tensorflow as tf
from blox import AttrDict
from blox.basic_types import str2int
from gcp.planning.infra.agent.utils.hdf5_saver import HDF5SaverBase
from tqdm import tqdm

phase = 'train'

count = 0
H = config.precrop_frame_ht
W = config.precrop_frame_wd
C = 3


class TAPMaker(HDF5SaverBase):
    def __init__(self, save_dir, offset=0, split=(0.90, 0.05, 0.05)):
        super().__init__(save_dir, traj_per_file=1, offset=offset, split=split)
        
        self.filenames = sorted(glob.glob(config.tfrecord_dir + phase + '/*.tfrecord?'))
        self.max_seq_len = 80
    
    def get_traj(self, string_record):
        example = tf.train.Example()
        example.ParseFromString(string_record)
        images = []
        
        numbered_keys = filter(lambda x: str2int(x.split('/')[0]) is not None, example.features.feature.keys())
        image_keys = filter(lambda x: 'image_view0/encoded' in x, numbered_keys)
        indices = np.array(list([str2int(x.split('/')[0]) for x in image_keys]))
        length = np.max(indices) + 1
        
        for i in range(length):
            key = '{}/image_view0/encoded'.format(i)
            val = np.frombuffer(example.features.feature[key].bytes_list.value[0], dtype=np.uint8)
            val = val.reshape(H, W, C)
            images.append(val)
        
        pad_mask = np.ones(len(images))
        images = np.array(images)
        
        return AttrDict(images=images, pad_mask=pad_mask)
    
    def make_phase(self, filenames, phase):
        for fn in filenames:
            record_iterator = tf.python_io.tf_record_iterator(path=fn)
            for i, string_record in enumerate(tqdm(record_iterator)):
                traj = self.get_traj(string_record)
                self.save_hdf5([traj], phase)
                
    # def save_hdf5(self, traj_list, phase):
    #     traj = traj_list[0]
    #     with h5py.File(config.h5record_dir + 'hdf5/' + phase + '/traj_{0:06d}'.format(count) + '.h5', 'w') as F:
    #         F['traj0/pad_mask'] = traj.pad_mask
    #         F['traj0/images'] = traj.images
    #         F['traj_per_file'] = 1


if __name__ == '__main__':
    maker = TAPMaker(config.h5record_dir)
    
    maker.make_dataset()