text: #!/usr/bin/env python3
  
  # Copyright (c) Facebook, Inc. and its affiliates.
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  
  import glob
  import os
  from collections import defaultdict
  from typing import Dict, List, Optional
  
  import numpy as np
  import torch
  import torch.nn as nn
  
  from habitat.utils.visualizations.utils import images_to_video
  from habitat_baselines.common.tensorboard_utils import TensorboardWriter
  
  
  class Flatten(nn.Module):
      def forward(self, x):
          return x.view(x.size(0), -1)
  
  
  class CustomFixedCategorical(torch.distributions.Categorical):
      def sample(self, sample_shape=torch.Size()):
          return super().sample(sample_shape).unsqueeze(-1)
  
      def log_probs(self, actions):
          return (
              super()
              .log_prob(actions.squeeze(-1))
              .view(actions.size(0), -1)
              .sum(-1)
              .unsqueeze(-1)
          )
  
      def mode(self):
          return self.probs.argmax(dim=-1, keepdim=True)
  
  
  class CategoricalNet(nn.Module):
      def __init__(self, num_inputs, num_outputs):
          super().__init__()
  
          self.linear = nn.Linear(num_inputs, num_outputs)
  
          nn.init.orthogonal_(self.linear.weight, gain=0.01)
          nn.init.constant_(self.linear.bias, 0)
  
      def forward(self, x):
          x = self.linear(x)
          return CustomFixedCategorical(logits=x)
  
  
  def linear_decay(epoch: int, total_num_updates: int) -> float:
      r"""Returns a multiplicative factor for linear value decay
  
      Args:
          epoch: current epoch number
          total_num_updates: total number of epochs
  
      Returns:
          multiplicative factor that decreases param value linearly
      """
      return 1 - (epoch / float(total_num_updates))
  
  
  def _to_tensor(v):
      if torch.is_tensor(v):
          return v
      elif isinstance(v, np.ndarray):
          return torch.from_numpy(v)
      else:
          return torch.tensor(v, dtype=torch.float)
  
  
  def batch_obs(
      observations: List[Dict], device: Optional[torch.device] = None
  ) -> Dict[str, torch.Tensor]:
      r"""Transpose a batch of observation dicts to a dict of batched
      observations.
  
      Args:
          observations:  list of dicts of observations.
          device: The torch.device to put the resulting tensors on.
              Will not move the tensors if None
  
      Returns:
          transposed dict of lists of observations.
      """
      batch = defaultdict(list)
  
      for obs in observations:
          for sensor in obs:
              batch[sensor].append(_to_tensor(obs[sensor]))
  
      for sensor in batch:
          batch[sensor] = torch.stack(batch[sensor], dim=0).to(
              device=device, dtype=torch.float
          )
  
      return batch
  
  
  def poll_checkpoint_folder(
      checkpoint_folder: str, previous_ckpt_ind: int
  ) -> Optional[str]:
      r""" Return (previous_ckpt_ind + 1)th checkpoint in checkpoint folder
      (sorted by time of last modification).
  
      Args:
          checkpoint_folder: directory to look for checkpoints.
          previous_ckpt_ind: index of checkpoint last returned.
  
      Returns:
          return checkpoint path if (previous_ckpt_ind + 1)th checkpoint is found
          else return None.
      """
      assert os.path.isdir(checkpoint_folder), "invalid checkpoint folder path"
      models_paths = list(
          filter(os.path.isfile, glob.glob(checkpoint_folder + "/*"))
      )
      models_paths.sort(key=os.path.getmtime)
      ind = previous_ckpt_ind + 1
      if ind < len(models_paths):
          return models_paths[ind]
      return None
  
  
  def generate_video(
      video_option: List[str],
      video_dir: Optional[str],
      images: List[np.ndarray],
      episode_id: int,
      checkpoint_idx: int,
      metric_name: str,
      metric_value: float,
      tb_writer: TensorboardWriter,
      fps: int = 10,
  ) -> None:
      r"""Generate video according to specified information.
  
      Args:
          video_option: string list of "tensorboard" or "disk" or both.
          video_dir: path to target video directory.
          images: list of images to be converted to video.
          episode_id: episode id for video naming.
          checkpoint_idx: checkpoint index for video naming.
          metric_name: name of the performance metric, e.g. "spl".
          metric_value: value of metric.
          tb_writer: tensorboard writer object for uploading video.
          fps: fps for generated video.
      Returns:
          None
      """
      if len(images) < 1:
          return
  
      video_name = f"episode{episode_id}_ckpt{checkpoint_idx}_{metric_name}{metric_value:.2f}"
      if "disk" in video_option:
          assert video_dir is not None
          images_to_video(images, video_dir, video_name)
      if "tensorboard" in video_option:
          tb_writer.add_video_from_np_images(
              f"episode{episode_id}", checkpoint_idx, images, fps=fps
          )
  