storage.py:
"""
Based on https://github.com/ikostrikov/pytorch-a2c-ppo-acktr

Used for on-policy rollout storages.
"""

import argparse
from typing import Any, Dict, Optional, Tuple

import gym
import numpy as np
import torch
from metarl.environments import utils as utl
from metarl.environments.utils import EnvLatents
from metarl.utils.torch import DeviceConfig
from torch.utils.data.sampler import BatchSampler, SubsetRandomSampler


def _flatten_helper(T, N, _tensor):
    return _tensor.reshape(T * N, *_tensor.size()[2:])


class OnlineStorage(object):
    def __init__(
        self,
        args: argparse.Namespace,
        model: torch.nn.Module,
        num_steps: int,
        num_processes: int,
        state_dim: int,  # TODO: Is this right?
        belief_dim: int,
        task_dim: int,
        action_space: gym.Space,
        hidden_size: int,
        latent_dim: int,
        normalize_rewards: bool,
        use_gae: bool,
        gamma: float,
        tau: float,
        use_proper_time_limits=True,
        use_popart: bool = False,
        add_exploration_bonus: bool = False,
        intrinsic_reward: Any = None,
        use_informedness: bool = False,
    ):
        """
        Stores data collected during a rollout.

        Args:
            args (argparse.Namespace): Arguments.
            num_steps (int): Number of steps to store.
            num_processes (int): Number of parallel processes.
            state_dim (int): Dimensionality of state space.
            belief_dim (int): Dimensionality of belief space.
            task_dim (int): Dimensionality of task space.
            action_space (gym.Space): Action space.
            hidden_size (int): Size of hidden RNN states.
            latent_dim (int): Dimensionality of latent space (of VAE).
            normalize_rewards (bool): Whether to normalize rewards or not.
            use_gae (bool): Whether to use Generalised Advantage Estimation.
            gamma: (float): The discount factor.
            tau: (float): The GAE parameter.
            use_proper_time_limits: (bool): Whether to use proper time limits.
            use_popart (bool): Whether to use pop-art for reward normalisation.
            add_exploration_bonus (bool): Whether to add exploration bonus to
                the reward.
            intrinsic_reward (Any): Intrinsic reward object.
                        use_gae: (bool): Whether to use Generalised Advantage Estimation.

        > What is the format of these buffers?

        prev_state (N+1) :    [  s_0* |  s_1  |  ...  |  s_{N-1}  |  s_N  ]
        beliefs    (N+1) :    [  b_0^ |  b_1  |  ...  |  b_{N-1}  |  b_N  ]
        rewards    (N)   :    [  r_0  |  r_1  |  ...  |  r_{N-1}  ]
        actions    (N)   :    [  a_0  |  a_1  |  ...  |  a_{N-1}  ]
        masks      (N+1) :    [   0 ^ |   0   |  ...  |     0     |  0/1  ]
        next_state (N)   :    [  s_1  |  s_2  |  ...  |    s_N    ]
        latents    (N+1) :    [  l_0  |  l_1  |  ...  |  l_{N-1}  |  l_N  ]
            * NOTE: s_0 is added in MetaLearner>kickstart_training()
            ^ NOTE: b_0 is never filled in, nor is mask_0

        > On 'insert' we do:
        -                       (step)      (step + 1)
        prev_state     ...  |     ...      |   state    |  ...
        beliefs        ...  |     ...      |   belief   |  ...
        actions        ...  |    action    |    ...     |  ...
        rewards        ...  |    reward    |    ...     |  ...
        masks          ...  |     ...      |    mask    |  ...
        next_state     ...  |    state*    |    ...     |  ...
        latents        ...  |     ...      |  latent^   |  ...
            * NOTE: metalearner adds this afterwards; not added in insert
            ^ NOTE: latent values are appended differently, but this is how the indices work out
        """

        self.args = args
        # Support tuple state dimensions # TODO: update docstring accordingly
        if isinstance(state_dim, int):
            # We set int state_dim
            self.state_dim = (state_dim,)
        elif isinstance(state_dim, tuple) or isinstance(state_dim, list):
            # We set tuple state_dim
            self.state_dim = tuple(state_dim)
        else:
            raise ValueError(
                "state_dim must be int, tuple or list, got {}".format(type(state_dim))
            )
        self.belief_dim = belief_dim
        self.task_dim = task_dim
        self.model = model

        self.num_steps = (
            num_steps  # how many steps to do per update (= size of online buffer)
        )
        self.num_processes = num_processes  # number of parallel processes
        self.step = 0  # keep track of current environment step

        # normalisation of the rewards
        self.normalize_rewards = normalize_rewards
        self.use_popart = use_popart  # TODO: Not implemented (look at DCD code)

        # computing returns
        self.use_gae = use_gae
        self.gamma = gamma
        self.tau = tau
        self.use_proper_time_limits = use_proper_time_limits

        # inputs to the policy
        # this will include s_0 when state was reset (hence num_steps+1)
        self.prev_state = torch.zeros(num_steps + 1, num_processes, *self.state_dim)
        if self.args.policy.pass_latent:
            # latent variables (of VAE)
            self.latent_dim = latent_dim
            self.latent_samples = []
            self.latent_mean = []
            self.latent_logvar = []
            # hidden states of RNN (necessary if we want to re-compute embeddings)
            self.hidden_size = hidden_size
            self.hidden_states = torch.zeros(num_steps + 1, num_processes, hidden_size)
        else:
            self.latent_mean = None
            self.latent_logvar = None
            self.latent_samples = None

        # next_state will include s_N when state was reset, skipping s_0
        # (only used if we need to re-compute embeddings after backpropagating RL loss through encoder)
        self.next_state = torch.zeros(num_steps, num_processes, *self.state_dim)

        if self.args.policy.pass_belief:
            self.beliefs = torch.zeros(num_steps + 1, num_processes, belief_dim)
        else:
            self.beliefs = None

        self.tasks = torch.zeros(num_steps + 1, num_processes, task_dim)

        # rewards and end of trials
        self.rewards_raw = torch.zeros(num_steps, num_processes, 1)
        self.rewards_normalized = torch.zeros(num_steps, num_processes, 1)
        self.done = torch.zeros(num_steps + 1, num_processes, 1)
        self.masks = torch.ones(num_steps + 1, num_processes, 1)
        # masks that indicate whether it's a true terminal state (false) or time limit end state (true)
        self.bad_masks = torch.ones(num_steps + 1, num_processes, 1)

        # Keep track of cliffhanger timesteps (from DCD)
        self.cliffhanger_masks = torch.ones(num_steps + 1, num_processes, 1)

        self.level_seeds = torch.zeros(num_steps, num_processes, 1, dtype=torch.int)

        # actions
        if action_space.__class__.__name__ == "Discrete":
            action_shape = 1
        else:
            action_shape = action_space.shape[0]
        self.actions = torch.zeros(num_steps, num_processes, action_shape)
        if action_space.__class__.__name__ == "Discrete":
            self.actions = self.actions.long()
        self.action_log_probs = None
        self.action_log_dist = None

        # values and returns
        self.value_preds = torch.zeros(num_steps + 1, num_processes, 1)
        self.returns = torch.zeros(num_steps + 1, num_processes, 1)
        self.denorm_value_preds = None

        # exploration bonus
        self.add_reward_bonus = add_exploration_bonus
        if self.add_reward_bonus:
            self.intrinsic_reward = intrinsic_reward

        # informedness
        self.use_informedness = use_informedness
        if use_informedness:
            self.informedness = []
            self.z_truths = []
        else:
            self.informedness = None
            self.z_truths = None

        self.images = []

        self.to_device()

    @property
    def rewards(self):
        if self.normalize_rewards:
            return self.rewards_normalized
        else:
            return self.rewards_raw

    def to_device(self):
        """
        Send all fields to GPU device.
        """
        if self.args.policy.pass_state:
            self.prev_state = self.prev_state.to(DeviceConfig.DEVICE)
        if self.args.policy.pass_latent:
            self.latent_samples = [
                t.to(DeviceConfig.DEVICE) for t in self.latent_samples
            ]
            self.latent_mean = [t.to(DeviceConfig.DEVICE) for t in self.latent_mean]
            self.latent_logvar = [t.to(DeviceConfig.DEVICE) for t in self.latent_logvar]
            self.hidden_states = self.hidden_states.to(DeviceConfig.DEVICE)
        self.next_state = self.next_state.to(DeviceConfig.DEVICE)
        if self.args.policy.pass_belief:
            self.beliefs = self.beliefs.to(DeviceConfig.DEVICE)
        if self.args.policy.pass_task or self.use_informedness:
            self.tasks = self.tasks.to(DeviceConfig.DEVICE)
        self.rewards_raw = self.rewards_raw.to(DeviceConfig.DEVICE)
        self.rewards_normalized = self.rewards_normalized.to(DeviceConfig.DEVICE)
        self.done = self.done.to(DeviceConfig.DEVICE)
        self.masks = self.masks.to(DeviceConfig.DEVICE)
        self.bad_masks = self.bad_masks.to(DeviceConfig.DEVICE)
        self.cliffhanger_masks = self.cliffhanger_masks.to(DeviceConfig.DEVICE)
        self.value_preds = self.value_preds.to(DeviceConfig.DEVICE)
        self.returns = self.returns.to(DeviceConfig.DEVICE)
        self.actions = self.actions.to(DeviceConfig.DEVICE)
        if self.level_seeds is not None:
            self.level_seeds = self.level_seeds.to(DeviceConfig.DEVICE)
        if self.use_informedness:
            self.informedness = [t.to(DeviceConfig.DEVICE) for t in self.informedness]
            self.z_truths = [t.to(DeviceConfig.DEVICE) for t in self.z_truths]

    def insert_initial_latents(
        self, env_latents: EnvLatents, informedness: Optional[torch.Tensor] = None
    ):
        """
        Insert the initial latent states into the buffer before rollout starts.
        """
        self.latent_samples.append(env_latents.z_samples.clone())
        self.latent_mean.append(env_latents.z_means.clone())
        self.latent_logvar.append(env_latents.z_logvars.clone())
        self.hidden_states[0].copy_(env_latents.hs)
        if self.use_informedness:
            # DEBUG - Do I need to clone?
            self.informedness.append(informedness.clone())
            self.z_truths.append(env_latents.z_truths.clone())  # DEBUG - "" ""

    def add_initial_state(self, state):
        if isinstance(state, dict):
            for k in state.keys():
                self.prev_state[k][0].copy_(state[k])
        else:
            self.prev_state[0].copy_(state)

    def add_next_state_before_reset(self, step_idx, state):
        """We assume state is already cloned!"""
        if isinstance(state, dict):
            for k in state.keys():
                self.next_state[k][step_idx] = state[k]
        else:
            self.next_state[step_idx] = state

    def insert(
        self,
        state: torch.Tensor,
        belief: torch.Tensor,
        task: torch.Tensor,
        actions: torch.Tensor,
        rewards_raw: torch.Tensor,
        rewards_normalized: torch.Tensor,
        value_preds: torch.Tensor,
        masks: torch.Tensor,
        bad_masks: torch.Tensor,
        done: torch.Tensor,
        hidden_states: torch.Tensor = None,
        latent_sample: torch.Tensor = None,
        latent_mean: torch.Tensor = None,
        latent_logvar: torch.Tensor = None,
        level_seeds: torch.Tensor = None,
        cliffhanger_masks: torch.Tensor = None,
        z_truths: torch.Tensor = None,
        informedness: torch.Tensor = None,
        images: np.array = None,
    ):
        """
        Insert a transition into the buffer.

        Args:
            state (torch.Tensor)
            belief (torch.Tensor)
            task (torch.Tensor)
            actions (torch.Tensor)
            rewards_raw (torch.Tensor)
            rewards_normalized (torch.Tensor)
            value_preds (torch.Tensor)
            masks (torch.Tensor): masks that indicate whether it's a true
                terminal state (false) or time limit end state (true)
            bad_masks (torch.Tensor): masks that indicate whether it's a time
                limit end state because of failure
            done (torch.Tensor):
            hidden_states (torch.Tensor):
            latent_sample (torch.Tensor)
            latent_mean (torch.Tensor)
            latent_logvar (torch.Tensor)
            level_seeds (torch.Tensor)
            cliffhanger_masks (torch.Tensor)
            z_truths (torch.Tensor)
            informedness (torch.Tensor)
            images (torch.Tensor)
        """
        self.prev_state[self.step + 1].copy_(state)
        if self.args.policy.pass_belief:
            self.beliefs[self.step + 1].copy_(belief)
        if self.args.policy.pass_task or self.use_informedness:
            self.tasks[self.step + 1].copy_(task)
        if self.args.policy.pass_latent:
            self.latent_samples.append(latent_sample.detach().clone())
            self.latent_mean.append(latent_mean.detach().clone())
            self.latent_logvar.append(latent_logvar.detach().clone())
            self.hidden_states[self.step + 1].copy_(hidden_states.detach())
        self.actions[self.step] = actions.detach().clone()
        self.rewards_raw[self.step].copy_(rewards_raw)
        self.rewards_normalized[self.step].copy_(rewards_normalized)
        if isinstance(value_preds, list):
            self.value_preds[self.step].copy_(value_preds[0].detach())
        else:
            self.value_preds[self.step].copy_(value_preds.detach())
        self.masks[self.step + 1].copy_(masks)
        self.bad_masks[self.step + 1].copy_(bad_masks)
        self.done[self.step + 1].copy_(done)

        if cliffhanger_masks is not None:
            self.cliffhanger_masks[self.step + 1].copy_(cliffhanger_masks)

        if level_seeds is not None:
            self.level_seeds[self.step].copy_(level_seeds)

        if self.use_informedness:
            self.informedness.append(informedness.detach().clone())
            self.z_truths.append(z_truths.clone())  # We want to preserve the graph!
            # But we need to copy to prevent double backpropagation

        if images is not None:
            self.images.append(images)

        self.step = (self.step + 1) % self.num_steps

    def after_update(self):
        """
        Copy the last state of the buffer to the first state.
        """
        self.prev_state[0].copy_(self.prev_state[-1])
        if self.args.policy.pass_belief:
            self.beliefs[0].copy_(self.beliefs[-1])
        if self.args.policy.pass_task or self.use_informedness:
            self.tasks[0].copy_(self.tasks[-1])
        if self.args.policy.pass_latent:
            self.latent_samples = []
            self.latent_mean = []
            self.latent_logvar = []
            self.hidden_states[0].copy_(self.hidden_states[-1])
        self.done[0].copy_(self.done[-1])
        self.masks[0].copy_(self.masks[-1])
        self.bad_masks[0].copy_(self.bad_masks[-1])
        self.cliffhanger_masks[0].copy_(self.cliffhanger_masks[-1])
        self.action_log_probs = None
        self.action_log_dist = None
        if self.use_informedness:
            self.informedness = []
            self.z_truths = []
        self.images = []

    def compute_returns(self, next_value: torch.Tensor, vae=None):
        """
        Compute the returns for each step in the buffer.

        Args:
            next_value: (torch.Tensor): The value of the next state.
            vae: (VAE): The VAE used for latent space.
        """

        if not self.add_reward_bonus:
            rewards = (
                self.rewards_normalized.clone()
                if self.normalize_rewards
                else self.rewards_raw.clone()
            )
        else:
            # compute reward bonus (we do it in here because then we can batch
            # the computation)
            with torch.no_grad():
                if self.args.policy.pass_latent:
                    # compute the rew bonus for s^+_{t+1}, i.e., skipping first
                    # observation and prior
                    belief = torch.cat(
                        (
                            torch.stack(self.latent_mean[1:]),
                            torch.stack(self.latent_logvar[1:]),
                        ),
                        dim=-1,
                    )
                    rew_bonus = self.intrinsic_reward.reward(
                        state=self.next_state,
                        belief=belief,
                        done=self.done,
                        update_normalisation=True,
                        vae=vae,
                        latent_mean=self.latent_mean,
                        latent_logvar=self.latent_logvar,
                        batch_prev_obs=self.prev_state[:-1],
                        batch_next_obs=self.prev_state[1:],
                        batch_actions=self.actions,
                        batch_rewards=self.rewards_raw,
                        batch_tasks=self.tasks[1:],
                    )
                else:
                    rew_bonus = self.intrinsic_reward.reward(
                        state=self.next_state,
                        belief=self.beliefs[1:],
                        done=self.done[1:],
                        update_normalisation=True,
                    )

            if self.normalize_rewards:
                rewards = self.rewards_normalized.clone() + rew_bonus.clone()

            else:
                rewards = self.rewards_raw.clone() + rew_bonus.clone()

        self._compute_returns(
            next_value=next_value,
            rewards=rewards,
            value_preds=self.value_preds,
            returns=self.returns,
            gamma=self.gamma,
            tau=self.tau,
            use_gae=self.use_gae,
            use_proper_time_limits=self.use_proper_time_limits,
        )

    def _compute_returns(
        self,
        next_value,
        rewards,
        value_preds,
        returns,
        gamma,
        tau,
        use_gae,
        use_proper_time_limits,
    ):
        if self.use_popart:
            self.denorm_value_preds = self.model.popart.denormalize(
                value_preds
            )  # denormalize all value predictions
            value_preds = self.denorm_value_preds

        if use_proper_time_limits:
            if use_gae:
                value_preds[-1] = next_value
                gae = 0
                for step in reversed(range(rewards.size(0))):
                    delta = (
                        rewards[step]
                        + gamma * value_preds[step + 1] * self.masks[step + 1]
                        - value_preds[step]
                    )
                    gae = delta + gamma * tau * self.masks[step + 1] * gae
                    gae = gae * self.bad_masks[step + 1]
                    returns[step] = gae + value_preds[step]
            else:
                returns[-1] = next_value
                for step in reversed(range(rewards.size(0))):
                    returns[step] = (
                        returns[step + 1] * gamma * self.masks[step + 1] + rewards[step]
                    ) * self.bad_masks[step + 1] + (
                        1 - self.bad_masks[step + 1]
                    ) * value_preds[step]
        else:
            if use_gae:
                value_preds[-1] = next_value
                gae = 0
                for step in reversed(range(rewards.size(0))):
                    delta = (
                        rewards[step]
                        + gamma * value_preds[step + 1] * self.masks[step + 1]
                        - value_preds[step]
                    )
                    gae = delta + gamma * tau * self.masks[step + 1] * gae
                    returns[step] = gae + value_preds[step]
            else:
                returns[-1] = next_value
                for step in reversed(range(rewards.size(0))):
                    returns[step] = (
                        returns[step + 1] * gamma * self.masks[step + 1] + rewards[step]
                    )

    def num_transitions(self):
        """
        Get the total number of transitions in the buffer.

        Returns:
            (int) The total number of transitions in the buffer.
        """
        return len(self.prev_state) * self.num_processes

    def before_update(
        self, policy: torch.nn.Module, informedness=None, env_latents=None
    ):
        """
        Compute the action log probabilities before the update step.

        Args:
            policy: (torch.nn.Module): The policy network.
        """
        z_samples = (
            torch.stack(self.latent_samples[:-1])  # [horizon, envs, z_dim]
            if self.latent_samples is not None
            else None
        )
        z_means = (
            torch.stack(self.latent_mean[:-1])  # [horizon, envs, z_dim]
            if self.latent_mean is not None
            else None
        )
        z_logvars = (
            torch.stack(self.latent_logvar[:-1])  # [horizon, envs, z_dim]
            if self.latent_logvar is not None
            else None
        )

        z_truths = None
        if env_latents.z_truths is not None:
            assert isinstance(env_latents.z_truths, torch.Tensor)
            # Extend same z_truths values across horizon dimension
            z_truths = (
                env_latents.z_truths.clone()
                .unsqueeze(0)
                .expand(z_samples.shape[0], -1, -1)
            )  # [envs, z_dim] -> [horizon, envs, z_dim]

        # We don't need `hs`---we're just getting processed zs for policy
        traj_env_latents = EnvLatents(
            z_samples=z_samples.reshape(-1, z_samples.size(-1)),
            z_means=z_means,
            z_logvars=z_logvars,
            hs=None,
        )
        traj_env_latents.set_z_truths(z_truths)

        # Get latents for each `horizon` step, for each `env`
        latent = utl.get_zs_for_policy(
            self.args, env_latents=traj_env_latents, informedness=informedness
        )  # [horizon, envs, z_dim*2] (usually)

        if informedness is not None and latent.dim() == 3:
            # informedness is currently shape [envs]
            # latent_dim is shape [horizon, envs, z_dim*2]
            # when we pass indices, we should reshape informedness to [horizon, envs]
            indices = informedness.unsqueeze(0).expand(latent.size(0), -1)
            # format so we can use .view() (make contiguous)
            indices = indices.contiguous()
        else:
            indices = informedness

        # Evaluate action log probs across the horizon, given the latents
        # we just computed
        _, action_log_probs, action_log_dist, _ = policy.evaluate_actions(
            self.prev_state[:-1],
            latent,
            self.beliefs[:-1] if self.beliefs is not None else None,
            self.tasks[:-1] if self.tasks is not None else None,
            self.actions,
            indices=indices,
        )
        self.action_log_probs = action_log_probs.detach()
        self.action_log_dist = action_log_dist.detach()

    def feed_forward_generator(
        self,
        advantages: Optional[torch.Tensor],
        num_mini_batch: Optional[int] = None,
        mini_batch_size: Optional[int] = None,
    ):
        """
        Generator that yields training data for PPO. The advantages are computed
        over num_steps steps and stored in the rollouts object.

        Args:
            advantages: (torch.Tensor)
            num_mini_batch: (int) number of batches for ppo
            mini_batch_size: (int) number of samples for each ppo batch
        """
        num_steps, num_processes = self.rewards_raw.size()[0:2]
        batch_size = num_processes * num_steps

        if mini_batch_size is None:
            assert batch_size >= num_mini_batch, (
                "PPO requires the number of processes ({}) "
                "* number of steps ({}) = {} "
                "to be greater than or equal to the number of PPO mini batches ({})."
                "".format(
                    num_processes, num_steps, num_processes * num_steps, num_mini_batch
                )
            )
            mini_batch_size = batch_size // num_mini_batch

        sampler = BatchSampler(
            SubsetRandomSampler(range(batch_size)), mini_batch_size, drop_last=True
        )
        for indices in sampler:
            # Initialize each batch component to None
            state_batch = latent_sample_batch = latent_mean_batch = (
                latent_logvar_batch
            ) = belief_batch = task_batch = adv_targ = informedness_batch = (
                z_truths_batch
            ) = None
            # Get relevant batch components from the rollouts object
            if self.args.policy.pass_state:
                state_batch = self.prev_state[:-1].reshape(
                    -1, *self.prev_state.size()[2:]
                )[indices]
            if self.args.policy.pass_latent:
                latent_sample_batch = torch.cat(self.latent_samples[:-1])[indices]
                latent_mean_batch = torch.cat(self.latent_mean[:-1])[indices]
                latent_logvar_batch = torch.cat(self.latent_logvar[:-1])[indices]
            if self.args.policy.pass_belief:
                belief_batch = self.beliefs[:-1].reshape(-1, *self.beliefs.size()[2:])[
                    indices
                ]
            if self.args.policy.pass_task or self.use_informedness:  # DEBUG
                task_batch = self.tasks[:-1].reshape(-1, *self.tasks.size()[2:])[
                    indices
                ]
            actions_batch = self.actions.reshape(-1, self.actions.size(-1))[indices]
            value_preds_batch = self.value_preds[:-1].reshape(-1, 1)[indices]
            return_batch = self.returns[:-1].reshape(-1, 1)[indices]
            old_action_log_probs_batch = self.action_log_probs.reshape(-1, 1)[indices]
            if advantages is not None:
                adv_targ = advantages.reshape(-1, 1)[indices]
            if self.use_informedness:
                informedness_batch = torch.cat(self.informedness[:-1])[indices]
                z_truths_batch = torch.cat(self.z_truths[:-1])[indices]
            # Return the batch
            yield (
                state_batch,
                belief_batch,
                task_batch,
                actions_batch,
                latent_sample_batch,
                latent_mean_batch,
                latent_logvar_batch,
                value_preds_batch,
                return_batch,
                old_action_log_probs_batch,
                adv_targ,
                informedness_batch,
                z_truths_batch,
            )

    def recurrent_generator(
        self,
        advantages: torch.Tensor,
        num_mini_batch: int = None,
        mini_batch_size: int = None,
    ):
        del advantages, num_mini_batch, mini_batch_size
        raise NotImplementedError


###############################################################################


class DictOnlineStorage(OnlineStorage):
    def __init__(
        self,
        args: argparse.Namespace,
        model: torch.nn.Module,
        num_steps: int,
        num_processes: int,
        state_dim: Dict[str, Tuple[int, ...]],  # TODO: Is this right?
        belief_dim: int,
        task_dim: int,
        action_space: gym.spaces.Dict,
        hidden_size: int,
        latent_dim: int,
        normalize_rewards: bool,
        use_popart: bool = False,
        add_exploration_bonus: bool = False,
        intrinsic_reward: Any = None,
    ):
        raise NotImplementedError("Deprecated")
        # TODO: Implement OnlineStorage as special case of DictOnlineStorage,
        # so we only have to make modifications to the Dict one. I tried other
        # way around, but too much redundancy in storage (i.e. only obs are
        # different, rest is the same).
        """
        Stores data collected during a rollout (with Dict observation space)

        Args:
            args (argparse.Namespace): parsed command line arguments
            num_steps (int): number of steps to collect per rollout
            num_processes (int): number of parallel processes
            state_dim (Dict[str, Tuple[int, ...]]): dimensionality of the state space
            belief_dim (int): dimensionality of the belief space
            task_dim (int): dimensionality of the task space
            action_space (gym.spaces.Dict): action space of the environment
            hidden_size (int): Size of hidden RNN states.
            latent_dim (int): Dimensionality of latent space (of VAE).
            normalize_rewards (bool): Whether to normalize rewards or not.
            use_popart (bool): Whether to use pop-art for reward normalisation.
            add_exploration_bonus (bool): Whether to add exploration bonus to
                the reward.
            intrinsic_reward (Any): Intrinsic reward object.
        """

        self.args = args
        self.state_dim = state_dim
        self.belief_dim = belief_dim
        self.task_dim = task_dim
        self.model = model

        self.num_steps = (
            num_steps  # how many steps to do per update (= size of online buffer)
        )
        self.num_processes = num_processes  # number of parallel processes
        self.step = 0  # keep track of current environment step

        # normalisation of the rewards
        self.normalize_rewards = normalize_rewards
        self.use_popart = use_popart

        # inputs to the policy
        # this will include s_0 when state was reset (hence num_steps+1)
        self.prev_state = {
            k: torch.zeros(num_steps + 1, num_processes, *v)
            for k, v in state_dim.items()
        }
        if self.args.policy.pass_latent:
            # latent variables (of VAE)
            self.latent_dim = latent_dim
            self.latent_samples = []
            self.latent_mean = []
            self.latent_logvar = []
            # hidden states of RNN (necessary if we want to re-compute embeddings)
            self.hidden_size = hidden_size
            self.hidden_states = torch.zeros(num_steps + 1, num_processes, hidden_size)
        else:
            self.latent_mean = None
            self.latent_logvar = None
            self.latent_samples = None
        # next_state will include s_N when state was reset, skipping s_0
        # (only used if we need to re-compute embeddings after backpropagating RL loss through encoder)
        self.next_state = {
            k: torch.zeros(num_steps, num_processes, *v) for k, v in state_dim.items()
        }
        if self.args.policy.pass_belief:
            self.beliefs = torch.zeros(num_steps + 1, num_processes, belief_dim)
        else:
            self.beliefs = None
        if self.args.policy.pass_task:
            self.tasks = torch.zeros(num_steps + 1, num_processes, task_dim)
        else:
            self.tasks = None

        # rewards and end of trials
        self.rewards_raw = torch.zeros(num_steps, num_processes, 1)
        self.rewards_normalized = torch.zeros(num_steps, num_processes, 1)
        self.done = torch.zeros(num_steps + 1, num_processes, 1)
        self.masks = torch.ones(num_steps + 1, num_processes, 1)
        # masks that indicate whether it's a true terminal state (false) or time limit end state (true)
        self.bad_masks = torch.ones(num_steps + 1, num_processes, 1)

        # Keep track of cliffhanger timesteps (from DCD)
        self.cliffhanger_masks = torch.ones(num_steps + 1, num_processes, 1)

        self.level_seeds = torch.zeros(num_steps, num_processes, 1, dtype=torch.int)

        # actions
        if action_space.__class__.__name__ == "Discrete":
            action_shape = 1
        else:
            action_shape = action_space.shape[0]
        self.actions = torch.zeros(num_steps, num_processes, action_shape)
        if action_space.__class__.__name__ == "Discrete":
            self.actions = self.actions.long()
        self.action_log_probs = None
        self.action_log_dist = None

        # values and returns
        self.value_preds = torch.zeros(num_steps + 1, num_processes, 1)
        self.returns = torch.zeros(num_steps + 1, num_processes, 1)

        # exploration bonus
        self.add_reward_bonus = add_exploration_bonus
        if self.add_reward_bonus:
            self.intrinsic_reward = intrinsic_reward

        self.to_device()

    def to_device(self):
        """
        Send all tensors to GPU device.
        """
        if self.args.policy.pass_state:
            self.prev_state = {
                k: v.to(DeviceConfig.DEVICE) for k, v in self.prev_state.items()
            }
        if self.args.policy.pass_latent:
            self.latent_samples = [
                t.to(DeviceConfig.DEVICE) for t in self.latent_samples
            ]
            self.latent_mean = [t.to(DeviceConfig.DEVICE) for t in self.latent_mean]
            self.latent_logvar = [t.to(DeviceConfig.DEVICE) for t in self.latent_logvar]
            self.hidden_states = self.hidden_states.to(DeviceConfig.DEVICE)
        self.next_state = {
            k: v.to(DeviceConfig.DEVICE) for k, v in self.next_state.items()
        }
        if self.args.policy.pass_belief:
            self.beliefs = self.beliefs.to(DeviceConfig.DEVICE)
        if self.args.policy.pass_task or self.use_informedness:  # DEBUG
            self.tasks = self.tasks.to(DeviceConfig.DEVICE)
        self.rewards_raw = self.rewards_raw.to(DeviceConfig.DEVICE)
        self.rewards_normalized = self.rewards_normalized.to(DeviceConfig.DEVICE)
        self.done = self.done.to(DeviceConfig.DEVICE)
        self.masks = self.masks.to(DeviceConfig.DEVICE)
        self.bad_masks = self.bad_masks.to(DeviceConfig.DEVICE)
        self.cliffhanger_masks = self.cliffhanger_masks.to(DeviceConfig.DEVICE)
        self.value_preds = self.value_preds.to(DeviceConfig.DEVICE)
        self.returns = self.returns.to(DeviceConfig.DEVICE)
        self.actions = self.actions.to(DeviceConfig.DEVICE)

        if self.level_seeds is not None:
            self.level_seeds = self.level_seeds.to(DeviceConfig.DEVICE)

    def insert(
        self,
        state: Dict[str, torch.Tensor],
        belief: torch.Tensor,
        task: torch.Tensor,
        actions: torch.Tensor,
        rewards_raw: torch.Tensor,
        rewards_normalized: torch.Tensor,
        value_preds: torch.Tensor,
        masks: torch.Tensor,
        bad_masks: torch.Tensor,
        done: torch.Tensor,
        hidden_states: torch.Tensor = None,
        latent_sample: torch.Tensor = None,
        latent_mean: torch.Tensor = None,
        latent_logvar: torch.Tensor = None,
        level_seeds: torch.Tensor = None,
        cliffhanger_masks: torch.Tensor = None,
    ):
        """
        Insert a transition into the buffer.

        Args:
            state: (Dict[str, torch.Tensor]) containing all the environment state representations
            belief: (torch.Tensor)
            task: (torch.Tensor)
            actions: (torch.Tensor)
            rewards_raw: (torch.Tensor)
            rewards_normalized: (torch.Tensor)
            value_preds: (torch.Tensor)
            masks: (torch.Tensor)
            bad_masks: (torch.Tensor)
            done: (torch.Tensor)
            hidden_states: (torch.Tensor)
            latent_sample: (torch.Tensor)
            latent_mean: (torch.Tensor)
            latent_logvar: (torch.Tensor)
            level_seeds: (torch.Tensor)
        """
        for key in state.keys():
            self.prev_state[key][self.step + 1].copy_(state[key])
        if self.args.policy.pass_belief:
            self.beliefs[self.step + 1].copy_(belief)
        if self.args.policy.pass_task:
            self.tasks[self.step + 1].copy_(task)
        if self.args.policy.pass_latent:
            self.latent_samples.append(latent_sample.detach().clone())
            self.latent_mean.append(latent_mean.detach().clone())
            self.latent_logvar.append(latent_logvar.detach().clone())
            self.hidden_states[self.step + 1].copy_(hidden_states.detach())
        self.actions[self.step] = actions.detach().clone()
        self.rewards_raw[self.step].copy_(rewards_raw)
        self.rewards_normalized[self.step].copy_(rewards_normalized)
        if isinstance(value_preds, list):
            self.value_preds[self.step].copy_(value_preds[0].detach())
        else:
            self.value_preds[self.step].copy_(value_preds.detach())
        self.masks[self.step + 1].copy_(masks)
        self.bad_masks[self.step + 1].copy_(bad_masks)
        self.done[self.step + 1].copy_(done)

        if cliffhanger_masks is not None:
            self.cliffhanger_masks[self.step + 1].copy_(cliffhanger_masks)

        if level_seeds is not None:
            self.level_seeds[self.step].copy_(level_seeds)

        self.step = (self.step + 1) % self.num_steps

    def after_update(self):
        """
        Copy the last state of the buffer to the first state.
        """
        for k in self.prev_state.keys():
            self.prev_state[k][0].copy_(self.prev_state[k][-1])
        if self.args.policy.pass_belief:
            self.beliefs[0].copy_(self.beliefs[-1])
        if self.args.policy.pass_task:
            self.tasks[0].copy_(self.tasks[-1])
        if self.args.policy.pass_latent:
            self.latent_samples = []
            self.latent_mean = []
            self.latent_logvar = []
            self.hidden_states[0].copy_(self.hidden_states[-1])
        self.done[0].copy_(self.done[-1])
        self.masks[0].copy_(self.masks[-1])
        self.bad_masks[0].copy_(self.bad_masks[-1])
        self.cliffhanger_masks[0].copy_(self.cliffhanger_masks[-1])
        self.action_log_probs = None
        self.action_log_dist = None

    def num_transitions(self):
        """
        Get the total number of transitions in the buffer.

        Returns:
            (int) The total number of transitions in the buffer.
        """
        k = self.prev_state.keys()[0]  # All keys should have same num
        return len(self.prev_state[k]) * self.num_processes

    def before_update(self, policy: torch.nn.Module):
        """
        Compute the action log probabilities before the update step.

        Args:
            policy: (torch.nn.Module): The policy network.
        """
        latent = utl.get_zs_for_policy(
            self.args,
            latent_sample=torch.stack(self.latent_samples[:-1])
            if self.latent_samples is not None
            else None,
            latent_mean=torch.stack(self.latent_mean[:-1])
            if self.latent_mean is not None
            else None,
            latent_logvar=torch.stack(self.latent_logvar[:-1])
            if self.latent_mean is not None
            else None,
        )

        # Construct
        prev_state = {k: self.prev_state[k][:-1] for k in self.prev_state.keys()}

        _, action_log_probs, action_log_dist, _ = policy.evaluate_actions(
            prev_state,
            latent,
            self.beliefs[:-1] if self.beliefs is not None else None,
            self.tasks[:-1] if self.tasks is not None else None,
            self.actions,
        )
        self.action_log_probs = action_log_probs.detach()
        self.action_log_dist = action_log_dist.detach()

    def feed_forward_generator(
        self,
        advantages: torch.Tensor,
        num_mini_batch: int = None,
        mini_batch_size: int = None,
    ):
        """
        Generator that yields training data for PPO. The advantages are computed
        over num_steps steps and stored in the rollouts object.

        Args:
            advantages: (torch.Tensor)
            num_mini_batch: (int) number of batches for ppo
            mini_batch_size: (int) number of samples for each ppo batch
        """
        num_steps, num_processes = self.rewards_raw.size()[0:2]
        batch_size = num_processes * num_steps

        # if (mini_batch_size is None):
        if mini_batch_size is None:
            assert batch_size >= num_mini_batch, (
                "PPO requires the number of processes ({}) "
                "* number of steps ({}) = {} "
                "to be greater than or equal to the number of PPO mini batches ({})."
                "".format(
                    num_processes, num_steps, num_processes * num_steps, num_mini_batch
                )
            )
            mini_batch_size = batch_size // num_mini_batch

        sampler = BatchSampler(
            SubsetRandomSampler(range(batch_size)), mini_batch_size, drop_last=True
        )

        for indices in sampler:
            if self.args.policy.pass_state:
                state_batch = dict()
                for k, v in self.prev_state.items():
                    state_batch[k] = v[:-1].reshape(-1, *self.prev_state[k].size()[2:])[
                        indices
                    ]
            else:
                state_batch = None
            if self.args.policy.pass_latent:
                latent_sample_batch = torch.cat(self.latent_samples[:-1])[indices]
                latent_mean_batch = torch.cat(self.latent_mean[:-1])[indices]
                latent_logvar_batch = torch.cat(self.latent_logvar[:-1])[indices]
            else:
                latent_sample_batch = latent_mean_batch = latent_logvar_batch = None
            if self.args.policy.pass_belief:
                belief_batch = self.beliefs[:-1].reshape(-1, *self.beliefs.size()[2:])[
                    indices
                ]
            else:
                belief_batch = None
            if self.args.policy.pass_task:
                task_batch = self.tasks[:-1].reshape(-1, *self.tasks.size()[2:])[
                    indices
                ]
            else:
                task_batch = None

            actions_batch = self.actions.reshape(-1, self.actions.size(-1))[indices]

            value_preds_batch = self.value_preds[:-1].reshape(-1, 1)[indices]
            return_batch = self.returns[:-1].reshape(-1, 1)[indices]

            old_action_log_probs_batch = self.action_log_probs.reshape(-1, 1)[indices]
            if advantages is None:
                adv_targ = None
            else:
                adv_targ = advantages.reshape(-1, 1)[indices]

            yield (
                state_batch,
                belief_batch,
                task_batch,
                actions_batch,
                latent_sample_batch,
                latent_mean_batch,
                latent_logvar_batch,
                value_preds_batch,
                return_batch,
                old_action_log_probs_batch,
                adv_targ,
            )

    def recurrent_generator(
        self,
        advantages: torch.Tensor,
        num_mini_batch: int = None,
        mini_batch_size: int = None,
    ):
        del advantages, num_mini_batch, mini_batch_size
        raise NotImplementedError

main.py:
import atexit
import cProfile
import os
os.environ["MUJOCO_GL"] = 'glfw'
import pstats
import subprocess
import sys
import warnings
from collections.abc import Mapping
from io import StringIO

import hydra
import matplotlib
import numpy as np
import torch
from loguru import logger as logr
from numba import NumbaDeprecationWarning
from omegaconf import DictConfig, OmegaConf

import metarl.environments.register  # Register environments  # noqa: E402, F401
from metarl.environments.make_envs import make_vec_envs
from metarl.environments.utils import get_vec_env_kwargs
from metarl.metalearner import MetaLearner
from metarl.utils.config import omegaconf_to_dict, print_dict
from metarl.utils.constructors import get_plr_args_dict
from metarl.utils.torch import (
    DeviceConfig,  # Import it early so we can set the device
    select_device,
)

# Suppress specific warnings
os.environ["TF_CPP_MIN_LOG_LEVEL"] = "2"  # Disables annoying TF warnings
os.environ["PYTHONWARNINGS"] = "ignore::DeprecationWarning"
os.environ["PYTHONWARNINGS"] = "ignore"
warnings.simplefilter(action="ignore", category=DeprecationWarning)
warnings.filterwarnings("ignore", category=NumbaDeprecationWarning)
warnings.simplefilter(action="ignore", category=FutureWarning)
matplotlib.use("Agg")  # For headless plotting

# Start Xvfb
display_number = "1"
lock_file = f"/tmp/.X{display_number}-lock"


def start_xvfb(display):
    return subprocess.Popen(["Xvfb", display, "-screen", "0", "1024x768x24"])


def cleanup():
    if "xvfb_process" in locals():
        xvfb_process.terminate()


if not os.path.exists(lock_file):
    xvfb_process = start_xvfb(f":{display_number}")
atexit.register(cleanup)
os.environ["DISPLAY"] = f":{display_number}"

logr.info("Starting experiment!")


def recursive_merge(dict1, dict2):
    """Recursively merge two dictionaries."""
    merged = dict1.copy()
    for key, value in dict2.items():
        if isinstance(value, Mapping) and key in merged:
            merged[key] = recursive_merge(merged.get(key, {}), value)
        else:
            merged[key] = value
    return merged


def apply_domain_overrides(domain_name, arbitrary_cfg):
    # Fetch overrides specific to the db_type if available
    domain_specific = arbitrary_cfg.domain_overrides.get(domain_name, {})
    # Create a new dict that starts with the base model config
    final_params = {k: v for k, v in arbitrary_cfg.items() if k != "domain_overrides"}
    # Apply overrides
    final_params = recursive_merge(final_params, domain_specific)
    return final_params


@hydra.main(version_base=None, config_path="cfg", config_name="main")
def main(cfg: DictConfig) -> None:
    OmegaConf.set_struct(cfg, False)  # Allows for dynamic attribute assignment

    # Remove the old handler. Else, the old one will continue to run along with
    # the new one specified below
    logr.remove()

    # Specify logger behavior
    logger_format = (
        "<level>{level: <2}</level> <cyan>{name}</cyan>:<cyan>{function}</cyan>:<cyan>{line}</cyan>"
        "- <level>{message}</level>"
    )
    if cfg.log_level == "debug" or cfg.debug:
        cfg.log_level = "debug"
        logr.add(sys.stdout, level="DEBUG", format=logger_format)
    elif cfg.log_level == "info":
        logr.add(sys.stdout, level="INFO", format=logger_format)
    else:
        raise ValueError(f"Invalid log level: {cfg.log_level}")

    dict = omegaconf_to_dict(cfg)
    print_dict(dict)
    args = cfg

    # For dubugging
    if args.debug:
        if args.use_qd:
            args.qd.warm_start_updates = 20
            args.qd.warm_start_no_sim_objective = True
        args.policy.num_processes = 3
        args.domain.num_frames = 10_000

    # Set evaluation environment
    if args.ued:
        args.domain.env_name = args.domain.ued_env_id
    else:
        args.domain.env_name = args.domain.reg_env_id
    args.domain.eval_env_name = args.domain.reg_env_id

    print("\nUED: ", args.ued)
    print("Train env name: ", args.domain.env_name)
    print("Eval  env name: ", args.domain.eval_env_name)

    # Designed to avoid issues with domains overwriting default algo parameters:
    if args.plr and args.plr.domain_randomization:
        args.plr.level_replay_strategy = "random"

    if args.qd:
        # If only warm-start for QD, set num-processes to 3 to preserve memory
        if args.use_qd and args.qd.warm_start_only:
            args.policy.num_processes = 3
        # If no initial warm starts, set boolean for convenience
        if args.qd.init_warm_start_updates == 0:
            args.qd.use_two_stage_ws = False
        else:
            args.qd.use_two_stage_ws = True
        if args.qd.dual_warm_start_updates is not None:
            args.qd.init_warm_start_updates = args.qd.dual_warm_start_updates
            args.qd.warm_start_updates = args.qd.dual_warm_start_updates

    # check if we're adding an exploration bonus
    if args.explore:
        args.add_exploration_bonus = (
            args.exploration_bonus_hyperstate
            or args.exploration_bonus_state
            or args.exploration_bonus_belief
            or args.exploration_bonus_vae_error
        )
    else:
        args.add_exploration_bonus = False

    # warning for deterministic execution
    if args.deterministic_execution:
        print("Envoking deterministic code execution.")
        if torch.backends.cudnn.enabled:
            warnings.warn("Running with deterministic CUDNN.")
        if args.policy.num_processes > 1:
            raise RuntimeError(
                "If you want fully deterministic, run with num_processes=1."
                "Warning: This will slow things down and might break A2C if "
                "policy_num_steps < env._max_episode_steps."
            )

    os.environ["TF_TENSORRT_DISABLED"] = "1"

    # Select device
    # GPU selection
    if "CUDA_VISIBLE_DEVICES" not in os.environ:  # If not defined we use GPU by default
        os.environ["CUDA_VISIBLE_DEVICES"] = "0"
    device = os.environ["CUDA_VISIBLE_DEVICES"]
    if len(device) == 0:
        # -1 is CPU
        select_device(-1)
        print("\nDevice: CPU")
    else:
        # Non-negative integer is the index of GPU
        select_device(0)
        print("\nDevice: GPU", device)

    # if we're normalising the actions, we have to make sure that the env
    # expects actions within [-1, 1]
    if args.policy.norm_actions_pre_sampling or args.policy.norm_actions_post_sampling:
        # NOTE: We don't actually use these environments; just for assertions
        envs, _ = make_vec_envs(
            env_name=args.domain.env_name,
            seed=0,
            num_processes=args.policy.num_processes,
            gamma=args.policy.gamma,
            device="cpu",
            episodes_per_trial=args.domain.episodes_per_trial,
            normalize_rew=args.policy.norm_rew,
            ret_rms=None,
            plr=args.use_plr,
            tasks=None,
        )
        assert np.unique(envs.action_space.low) == [-1]
        assert np.unique(envs.action_space.high) == [1]
        envs.close()
        del envs

    # clean up arguments
    if args.vae.disable_metalearner or args.vae.disable_decoder:
        args.vae.decode_reward = False
        args.vae.decode_state = False
        args.vae.decode_task = False

    if hasattr(args, "decode_only_past") and args.vae.decode_only_past:
        args.vae.split_batches_by_elbo = True

    # Init profiler
    if args.profile:
        args.domain.num_frames = 96_000  # 50_000
        pr = cProfile.Profile()
        pr.enable()

    # Weights and Biases
    if args.use_wandb:
        assert args.wandb_label is not None
        args.wb_exp_index = args.wandb_label.split("_")[0]

    # We only use one seed for the main training loop
    assert isinstance(args.seed, int) or isinstance(args.seed[0], int)
    args.seed = args.seed if isinstance(args.seed, int) else args.seed[0]
    print("Training with seed: ", args.seed)

    # Prepare environment arguments
    args.action_space = None  # In case it wasn't defined (?)
    args.vec_env_kwargs = get_vec_env_kwargs(args)

    # Prepare some things before QD
    archive_dims = None
    try:
        if args.qd:
            # Account for two-stage warm-start
            if args.qd.use_two_stage_ws:
                init_archive_dims = args.qd.init_archive_dims
                reg_archive_dims = args.qd.archive_dims
                init_total_cells = np.prod(init_archive_dims)
                reg_total_cells = np.prod(reg_archive_dims)
                # We initialize PLR with maximum total cells between the two archives
                total_cells = max(init_total_cells, reg_total_cells)
                # Set current archive dims to initial archive dims
                archive_dims = init_archive_dims
            else:
                # Set PLR buffer size normally
                archive_dims = args.qd.archive_dims  # Default
                total_cells = np.prod(archive_dims)

        if args.plr:
            if args.plr.buffer_size_from_qd:
                print(
                    "Setting PLR seed buffer size based on QD archive: ",
                    total_cells,
                    "+ 10_000",
                )
                args.plr.seed_buffer_size = int(total_cells) + 10_000
            elif args.use_qd and args.use_plr:
                print(
                    "Adding 1000 to PLR seed buffer size to account for temporary QD excess during eval: ",
                    total_cells,
                    "+ 1000",
                )
                args.plr.seed_buffer_size = args.plr.seed_buffer_size + 10_000
            else:
                print("Not setting PLR seed buffer size based on QD archive.")
            print("args.plr.seed_buffer_size:", args.plr.seed_buffer_size)
        # Adding 10_000 to account for when archive gets too full---we need
        # this extra space because otherwise, when we find new solutions
        # and our archive/buffer is full, when we try to insert, the
        # fully plr buffer will prevent us from being able to insert these
        # solutions. This is a temporary fix!
    except Exception as e:
        # TODO: We should remove this try: except if possible
        print("Exception with QD preparation:", e)
        print("Continuing. Assuming unnecessary for this experiment.")
        archive_dims = np.prod(args.qd.archive_dims)

    # Initialize environments here for memory efficiency
    if args.use_plr:
        if args.use_qd and args.qd.no_sim_objective:
            args.plr.level_replay_strategy = "random"  # Ignore objectives
        plr_num_actors = (
            args.qd.batch_size * 5 if args.use_qd else args.policy.num_processes
        )
        plr_level_sampler_args = get_plr_args_dict(args, plr_num_actors)
    else:
        plr_level_sampler_args = None
    # Initialize environments and plr level sampler (if applicable)
    envs, plr_components = make_vec_envs(
        env_name=args.domain.env_name,
        seed=args.seed,
        num_processes=args.policy.num_processes,
        gamma=args.policy.gamma,
        device=DeviceConfig.DEVICE,
        episodes_per_trial=args.domain.episodes_per_trial,
        normalize_rew=args.policy.norm_rew,
        ret_rms=None,
        tasks=None,
        qd=args.use_qd,
        qd_tasks=None,
        plr=args.use_plr,
        plr_level_sampler_args=plr_level_sampler_args,
        plr_level_sampler=None,  # We initialize inside the function for now
        plr_env_generator=args.plr.env_generator if args.plr else None,
        dense_rewards=args.dense_rewards,
        normalize_rew_by_tt=(args.informed is not None)
        and (args.informed.normalize_rew_by_tt),
        **args.vec_env_kwargs,
    )

    # Meta-learn
    learner_fn = (
        None if args.vae.disable_metalearner else MetaLearner
    )  # TODO: reimplement base RL learner
    learner = learner_fn(args, envs, plr_components, archive_dims=archive_dims)
    learner.train()

    # Profiling output
    if args.profile:
        s = StringIO()
        pr.disable()
        sortby = "cumtime"
        ps = pstats.Stats(pr, stream=s).sort_stats(sortby)
        ps.print_stats(250)
        print(s.getvalue())


if __name__ == "__main__":
    main()


vae.py:
import os
import pickle
import tempfile
import warnings
from typing import Optional

import gym
import gymnasium
import numpy as np
import torch
import torch.nn as nn
import wandb
from loguru import logger
from metarl.components.decoder import RewardDecoder, StateTransitionDecoder, TaskDecoder
from metarl.components.encoder import RNNEncoder
from metarl.components.vae.storage import DictRolloutStorageVAE, RolloutStorageVAE
from metarl.environments.utils import get_num_tasks, get_task_dim
from metarl.utils.torch import DeviceConfig
from torch.distributions.normal import Normal
from torch.nn import functional as F
from torch.optim import Adam


class VaribadVAE:
    """
    VAE of VariBAD:
    - has an encoder and decoder
    - can compute the ELBO loss
    - can update the VAE (encoder+decoder)
    """

    def __init__(
        self,
        args,
        logger,
        get_iter_idx,
        state_feature_extractor,
        state_dtype,
        use_informedness=False,
    ):
        self.args = args
        self.logger = logger
        self.get_iter_idx = get_iter_idx
        self.state_feature_extractor = state_feature_extractor
        self.state_dtype = state_dtype
        self.task_dim = (
            get_task_dim(self.args)
            if self.args.vae.decode_task or use_informedness
            else None
        )
        self.num_tasks = get_num_tasks(self.args) if self.args.vae.decode_task else None
        self.use_informedness = use_informedness

        # initialize the encoder
        self.encoder = self.initialize_encoder()

        # initialize the decoders (returns None for unused decoders)
        self.state_decoder, self.reward_decoder, self.task_decoder = (
            self.initialize_decoder()
        )

        # initialize rollout storage for the VAE update
        # (this differs from the data that the on-policy RL algorithm uses)
        if isinstance(self.args.state_dim, dict):
            RSVAE = DictRolloutStorageVAE
        else:
            RSVAE = RolloutStorageVAE

        self.rollout_storage = RSVAE(
            num_processes=self.args.policy.num_processes,
            max_trajectory_len=self.args.max_trajectory_len,
            zero_pad=True,
            max_num_rollouts=self.args.vae.buffer_size,
            state_dim=self.args.state_dim,
            action_dim=self.args.action_dim,
            vae_buffer_add_thresh=self.args.vae.buffer_add_thresh,
            task_dim=self.task_dim,
            latent_dim=self.args.vae.latent_dim,
            state_dtype=self.state_dtype,
            logger=self.logger,
        )

        # initialize optimizer for the encoder and decoders
        decoder_params = []
        if not self.args.vae.disable_decoder:
            if self.args.vae.decode_reward:
                decoder_params.extend(self.reward_decoder.parameters())
            if self.args.vae.decode_state:
                decoder_params.extend(self.state_decoder.parameters())
            if self.args.vae.decode_task:
                decoder_params.extend(self.task_decoder.parameters())

        # Partial beliefs
        self.encode_z_to_h = nn.Sequential(
            nn.Linear(
                self.args.vae.latent_dim * 2, self.args.vae.encoder_gru_hidden_size
            ),
            nn.ReLU(),
            nn.Linear(
                self.args.vae.encoder_gru_hidden_size,
                self.args.vae.encoder_gru_hidden_size,
            ),
            nn.ReLU(),
        ).to(DeviceConfig.DEVICE)

        self.optimizer_vae = Adam(
            [
                *self.encoder.parameters(),
                *decoder_params,
                *self.encode_z_to_h.parameters(),
            ],
            lr=self.args.vae.lr,
        )

        # Create environment so we can get task information
        if 'Ant' in self.args.domain.env_name:
            self.dummy_env = gymnasium.make(self.args.domain.env_name)
        else:
            kwargs = self.args.domain.kwargs
            kwargs = {key: self.args.domain[key] for key in kwargs}
            self.dummy_env = gym.make(self.args.domain.env_name, **kwargs)
        # TODO: What if we need other environment parameters to ensure the
        #   genotype etc. are right? We should pass in the actual environments
        #   so that we can get the right task information. This only works
        #   for gridworld at the moment.
        if self.args.vae.decode_task or self.args.vae.multihead_for_reward:
            # assert (
            #     "gridworld" in self.args.domain.domain
            # ), "The dummy_env for task decoding only works for GridWorld at the moment."
            self.task_to_id = self.dummy_env.task_to_id_fn

        # Alignment coeff (we have a schedule for this, updated in metalearner)
        if self.use_informedness and self.args.informed.alignment_loss:
            self.alignment_coeff = self.args.informed.alignment_init_coeff
        else:
            self.alignment_coeff = None

        self.update_count = 0

        self.attributes = [
            self.encoder,
            self.state_decoder,
            self.reward_decoder,
            self.task_decoder,
            self.encode_z_to_h,
        ]
        self.attributes_names = [
            "encoder",
            "state_decoder",
            "reward_decoder",
            "task_decoder",
            "encode_z_to_h",
        ]

    ###########################################################################
    #                            INITIALIZATION                               #
    ###########################################################################

    def get_encoder_params(self):
        return self.encoder.parameters()

    def initialize_encoder(self):
        """Initializes and returns an RNN encoder"""
        encoder = RNNEncoder(
            args=self.args,
            layers_before_gru=self.args.vae.encoder_layers_before_gru,
            hidden_size=self.args.vae.encoder_gru_hidden_size,
            layers_after_gru=self.args.vae.encoder_layers_after_gru,
            latent_dim=self.args.vae.latent_dim,
            action_dim=self.args.action_dim,
            action_embed_dim=self.args.vae.action_embedding_size,
            state_dim=self.args.state_dim,
            state_embed_dim=self.args.vae.state_embedding_size,
            reward_size=1,
            reward_embed_size=self.args.vae.reward_embedding_size,
            state_feature_extractor=self.state_feature_extractor,
            state_is_image=self.args.domain.state_is_image,
        ).to(DeviceConfig.DEVICE)
        return encoder

    def initialize_decoder(self):
        """
        Initializes and returns the (state/reward/task) decoder as specified
        in self.args
        """

        if self.args.vae.disable_decoder:
            return None, None, None

        latent_dim = self.args.vae.latent_dim
        # if we don't sample embeddings for the decoder, we feed in
        # mean & variance
        if self.args.vae.disable_stochasticity_in_latent:
            latent_dim *= 2

        # initialize state decoder for VAE
        if self.args.vae.decode_state:
            state_decoder = StateTransitionDecoder(
                args=self.args,
                layers=self.args.vae.state_decoder_layers,
                latent_dim=latent_dim,
                action_dim=self.args.action_dim,
                action_embed_dim=self.args.vae.action_embedding_size,
                state_dim=self.args.state_dim,
                state_embed_dim=self.args.vae.state_embedding_size,
                pred_type=self.args.vae.state_pred_type,
                state_feature_extractor=self.state_feature_extractor,
            ).to(DeviceConfig.DEVICE)
        else:
            state_decoder = None

        # initialize reward decoder for VAE
        if self.args.vae.decode_reward:
            reward_decoder = RewardDecoder(
                args=self.args,
                layers=self.args.vae.reward_decoder_layers,
                latent_dim=latent_dim,
                state_dim=self.args.state_dim,
                state_embed_dim=self.args.vae.state_embedding_size,
                action_dim=self.args.action_dim,
                action_embed_dim=self.args.vae.action_embedding_size,
                num_states=self.args.num_states,
                multi_head=self.args.vae.multihead_for_reward,
                pred_type=self.args.vae.rew_pred_type,
                input_prev_state=self.args.vae.input_prev_state,
                input_action=self.args.vae.input_action,
                state_feature_extractor=self.state_feature_extractor,
            ).to(DeviceConfig.DEVICE)
        else:
            reward_decoder = None

        # initialize task decoder for VAE
        if self.args.vae.decode_task:
            assert self.task_dim != 0
            task_decoder = TaskDecoder(
                latent_dim=latent_dim,
                layers=self.args.vae.task_decoder_layers,
                task_dim=self.task_dim,
                num_tasks=self.num_tasks,
                pred_type=self.args.vae.task_pred_type,
            ).to(DeviceConfig.DEVICE)
        else:
            task_decoder = None

        return state_decoder, reward_decoder, task_decoder

    ###########################################################################
    #                        MAIN VAE LOSS COMPUTATION                        #
    ###########################################################################

    def compute_loss(
        self,
        latent_mean: torch.Tensor,
        latent_logvar: torch.Tensor,
        prior_means: torch.Tensor,
        prior_logvars: torch.Tensor,
        vae_prev_obs: torch.Tensor,
        vae_next_obs: torch.Tensor,
        vae_actions: torch.Tensor,
        vae_rewards: torch.Tensor,
        vae_tasks: torch.Tensor,
        trajectory_lens: np.ndarray,
        keep_batch_dim=False,
    ):
        """
        Computes the VAE loss for the given data.
        Batches everything together and therefore needs all trajectories to be
        of the same length.
        (Important because we need to separate ELBOs and decoding terms so
        can't collapse those dimensions)
        """
        num_unique_trajectory_lens = len(np.unique(trajectory_lens))

        # Ensure that if we have different trajectory lengths, we are
        # subsampling both ELBO terms and decoding terms # DEBUG: why?
        assert (num_unique_trajectory_lens == 1) or (
            self.args.vae.subsample_elbos and self.args.vae.subsample_decodes
        )
        # Ensure that we are not decoding only the past (because loss
        # computation requires decoding the entire trajectory)
        assert not self.args.vae.decode_only_past

        # Cut down the batch to the longest trajectory length to preserve the
        # structure. NOTE: this means we will possibly waste some computation
        # on zero-padded trajectories that are shorter than max_traj_len.
        H = np.max(trajectory_lens)  # max_traj_len (H for horizon)

        (latent_mean, latent_logvar) = (latent_mean[: H + 1], latent_logvar[: H + 1])
        if isinstance(vae_prev_obs, dict):
            for key in vae_prev_obs.keys():
                vae_prev_obs[key] = vae_prev_obs[key][:H]
                vae_next_obs[key] = vae_next_obs[key][:H]
        else:
            (vae_prev_obs, vae_next_obs) = (vae_prev_obs[:H], vae_next_obs[:H])
        (vae_actions, vae_rewards) = (vae_actions[:H], vae_rewards[:H])

        # take one sample for each ELBO term
        if not self.args.vae.disable_stochasticity_in_latent:
            latent_samples = self.encoder._sample_gaussian(latent_mean, latent_logvar)
        else:
            latent_samples = torch.cat((latent_mean, latent_logvar), dim=-1)

        E = latent_samples.shape[0]  # number of ELBO terms
        D = vae_actions.shape[0]  # number of decoding terms
        B = latent_samples.shape[1]  # number of trajectories

        # Subsample elbo terms: [E B dim] -> [_E B dim]
        if self.args.vae.subsample_elbos is not None:
            _E = self.args.vae.subsample_elbos
            # randomly choose which elbo's to subsample
            if num_unique_trajectory_lens == 1:
                # Trajectories are all of the same length, so we can subsample
                # elbos uniformly across the batch
                elbo_indices = torch.tensor(_E * B, dtype=torch.long).random_(
                    0, E
                )  # select diff elbos for each task
            else:
                # Trajectories are different lengths; we subsample elbo
                # indices separately up to their maximum possible encoding
                # length; only allow duplicates if the sample size would be
                # larger than the number of samples
                elbo_indices = np.concatenate(
                    [
                        np.random.choice(range(0, t + 1), _E, replace=_E > (t + 1))
                        for t in trajectory_lens
                    ]
                )
                if _E > H:
                    warnings.warn(
                        "_E > H so there will be duplicates. To avoid this use "
                        "--split-batches-by-elbo or --split-batches-by-task."
                    )
            # Create selection mask
            task_indices = torch.arange(B).repeat(_E)
            latent_samples = latent_samples[elbo_indices, task_indices, :].reshape(
                (_E, B, -1)
            )
            # Update number of elbos
            E = _E
        else:
            elbo_indices = None
            E = E  # just to make it explicit

        # Expand the s/r/a decoder inputs (to match size of latents)
        # New shape should be [E, H, B, a/s/r_dim]
        if isinstance(vae_prev_obs, dict):
            dec_prev_obs = dict()
            dec_next_obs = dict()
            for key in vae_prev_obs.keys():
                _expanded_shape = (E, *vae_prev_obs[key].shape)
                dec_prev_obs[key] = (
                    vae_prev_obs[key].unsqueeze(0).expand(_expanded_shape)
                )
                dec_next_obs[key] = (
                    vae_next_obs[key].unsqueeze(0).expand(_expanded_shape)
                )
        else:
            _expanded_shape = (E, *vae_prev_obs.shape)
            dec_prev_obs = vae_prev_obs.unsqueeze(0).expand(_expanded_shape)
            dec_next_obs = vae_next_obs.unsqueeze(0).expand(_expanded_shape)
        _expanded_shape = (E, *vae_actions.shape)
        dec_actions = vae_actions.unsqueeze(0).expand(_expanded_shape)
        _expanded_shape = (E, *vae_rewards.shape)
        dec_rewards = vae_rewards.unsqueeze(0).expand(_expanded_shape)

        # Subsample reconstruction terms [E D B dim] -> [E _D B dim]
        if self.args.vae.subsample_decodes is not None:
            _D = self.args.vae.subsample_decodes
            # NOTE(vb_authors): This current code will produce duplicates
            I0 = torch.arange(E).repeat(_D * B)  # indices for dim 0
            if num_unique_trajectory_lens == 1:
                I1 = torch.LongTensor(E * _D * B).random_(0, D)
            else:
                I1 = np.concatenate(  # indices for dim 1
                    [
                        np.random.choice(range(0, t), E * _D, replace=True)
                        for t in trajectory_lens
                    ]
                )
            I2 = torch.arange(B).repeat(E * _D)  # indices for dim 2

            _new_shape = (E, _D, B, -1)
            # Dictionary observations
            if isinstance(dec_prev_obs, dict):
                dec_prev_obs = {
                    key: dec_prev_obs[key][I0, I1, I2, :].reshape(_new_shape)
                    for key in dec_prev_obs.keys()
                }
                dec_next_obs = {
                    key: dec_next_obs[key][I0, I1, I2, :].reshape(_new_shape)
                    for key in dec_next_obs.keys()
                }
            # Vectorized observations
            else:
                dec_prev_obs = dec_prev_obs[I0, I1, I2, :].reshape(_new_shape)
                dec_next_obs = dec_next_obs[I0, I1, I2, :].reshape(_new_shape)
            # Actions and rewards
            dec_actions = dec_actions[I0, I1, I2, :].reshape(_new_shape)
            dec_rewards = dec_rewards[I0, I1, I2, :].reshape(_new_shape)
            D = _D

        # Expand the latent (to match the number of s/r/a decoder inputs)
        # New shape will be [B, E, H, latent_dim]
        dec_embedding = (
            latent_samples.unsqueeze(0)
            .expand((D, *latent_samples.shape))
            .transpose(1, 0)
        )

        # Reconstruction loss computation; for each timestep encoded,
        # decode everything and sum it up
        (
            rew_rec_loss,
            state_rec_loss,
            task_rec_loss,
        ) = 0, 0, 0  # by default
        # Shape for each computed loss is [E, D, B]
        if self.args.vae.decode_reward:  # Rewards
            rew_rec_loss = self.compute_rew_rec_loss(
                dec_embedding, dec_prev_obs, dec_next_obs, dec_actions, dec_rewards
            )
            # (E) Aggregate across individual ELBO terms
            _agg_fn = torch.mean if self.args.vae.avg_elbo_terms else torch.sum
            rew_rec_loss = _agg_fn(rew_rec_loss, dim=0)
            # (D) Aggregate across individual reconstruction (decoding) terms
            _agg_fn = (
                torch.mean if self.args.vae.avg_reconstruction_terms else torch.sum
            )
            rew_rec_loss = _agg_fn(rew_rec_loss, dim=0)
            # (B) Aggregate across tasks (in batch)
            rew_rec_loss_mean = rew_rec_loss.mean()
        else:
            rew_rec_loss_mean = 0

        if self.args.vae.decode_state:  # States
            state_rec_loss = self.compute_state_rec_loss(
                dec_embedding, dec_prev_obs, dec_next_obs, dec_actions
            )
            # (E) Aggregate across individual ELBO terms
            _agg_fn = torch.mean if self.args.vae.avg_elbo_terms else torch.sum
            state_rec_loss = _agg_fn(state_rec_loss, dim=0)
            # (D) Aggregate across individual reconstruction (decoding) terms
            _agg_fn = (
                torch.mean if self.args.vae.avg_reconstruction_terms else torch.sum
            )
            state_rec_loss = _agg_fn(state_rec_loss, dim=0)
            # (B) Aggregate across tasks (in batch)
            state_rec_loss_mean = state_rec_loss.mean()
        else:
            state_rec_loss_mean = 0

        if self.args.vae.decode_task:  # Tasks
            task_rec_loss = self.compute_task_rec_loss(latent_samples, vae_tasks)
            # (E) Aggregate across individual ELBO terms
            _agg_fn = torch.mean if self.args.vae.avg_elbo_terms else torch.sum
            task_rec_loss = _agg_fn(task_rec_loss, dim=0)
            # (D) Aggregate across individual reconstruction (decoding) terms
            # DEBUG - why always sum here? Is it possible they forgot to add avg_reconstruction_terms for task_rec_loss?
            task_rec_loss = task_rec_loss.sum(dim=0)
            # (B) Aggregate across tasks (in batch)
            task_rec_loss_mean = task_rec_loss.mean()
        else:
            task_rec_loss_mean = 0

        # Compute KL loss
        if not self.args.vae.disable_kl_term:
            # compute the KL term for each ELBO term of the current trajectory
            # shape: [num_elbo_terms] x [num_trajectories]
            kl_loss = self.compute_kl_loss(
                latent_mean, latent_logvar, elbo_indices, prior_means, prior_logvars
            )
            # (E) Aggregate across individual ELBO terms
            _agg_fn = torch.mean if self.args.vae.avg_elbo_terms else torch.sum
            kl_loss = _agg_fn(kl_loss, dim=0)
            # (D) Aggregate across individual reconstruction (decoding) terms
            kl_loss = kl_loss.sum(dim=0)
            # (B) Aggregate across tasks (in batch)
            kl_loss_mean = kl_loss.mean()
        else:
            kl_loss = 0
            kl_loss_mean = 0

        if keep_batch_dim:
            return (rew_rec_loss, state_rec_loss, task_rec_loss, kl_loss)
        else:
            return (
                rew_rec_loss_mean,
                state_rec_loss_mean,
                task_rec_loss_mean,
                kl_loss_mean,
            )

    def compute_alignment_loss(
        self,
        latent_mean: torch.Tensor,
        latent_logvar: torch.Tensor,
        z_truths: torch.Tensor,
        keep_batch_dim: bool = False,
        mask: Optional[bool] = None,
        post_gru_args=None,
    ):
        """
        Compute alignment loss between tt-0 and tt-2 latents.

        Args:
            latent_mean:        [H+1, B, latent_dim]
            latent_logvar:      [H+1, B, latent_dim]
            z_truths:                [B, latent_dim]
            keep_batch_dim: bool
            mask: bool

        Returns mean alignment loss across horizon and batch; shape=[1] (unless
        keep_batch_dim=True; shape=[B]).
        """
        if mask is None:
            mask = 1.0
        # Compute IMPORT alignment loss here
        alignment_loss = 0.0

        detach_mode = self.args.informed.alignment_detach_mode

        # If we're freezing GRU but updating post-GRU layers, we need to recompute
        # the latent means and logvars after the GRU, with input detached
        if detach_mode == "extra":
            assert post_gru_args is not None
            # Detach so we don't backprop through the GRU
            post_gru_args[0] = post_gru_args[0].detach()
            post_gru_args[1] = post_gru_args[1].detach()
            # Get new latent_means that have grads
            _, latent_mean, latent_logvar, _ = self.encoder.post_gru_processing(*post_gru_args)

        # Repeat z_truths; structure latent means and logvars
        if self.args.informed.alignment_last_timestep_only:
            _zts = z_truths.clone().unsqueeze(0)  # [1, B, latent_dim]
            _lms = latent_mean[-1].unsqueeze(0)  # [1, B, latent_dim]
            _llvs = latent_logvar[-1].unsqueeze(0)  # [1, B, latent_dim]
        else:
            _zts = z_truths.unsqueeze(0).repeat(  # [H+1, B, latent_dim]
                latent_mean.shape[0], 1, 1
            )
            _lms = latent_mean.clone()  # [H+1, B, latent_dim]
            _llvs = latent_logvar.clone()  # [H+1, B, latent_dim]

        # Create mask across horizon dimension
        if self.args.informed.alignment_last_timestep_only:
            horizon_weights = torch.ones(_lms.shape[0]).to(DeviceConfig.DEVICE)  # [1]
        else:
            if self.args.informed.alignment_weighting_type == "uniform":
                # Use uniform weighting for alignment loss
                horizon_weights = torch.ones(_lms.shape[0]).to(
                    DeviceConfig.DEVICE
                )  # [H+1]
            elif self.args.informed.alignment_weighting_type == "exponential":
                # Use exponential weighting for alignment loss
                # (later in trajectory means less uncertainty; should align more)
                H = _lms.shape[0] - 1
                coeff = self.args.informed.alignment_weighting_exp_coeff
                # Create tensor like [H, H-1, ..., 1, 0] of shape [H+1],
                # and then use coefficient to create exponential weights
                horizon_weights = torch.tensor(
                    [coeff ** ((H + 1) - (i + 1)) for i in range((H + 1))]
                ).to(DeviceConfig.DEVICE)
                # Prevent underflow by clamping to 1e-8
                horizon_weights = torch.clamp(horizon_weights, min=1e-8)  # [H+1]
                # Normalize weights
                horizon_weights /= horizon_weights.sum()
                horizon_weights *= H + 1
            elif self.args.informed.alignment_weighting_type == "rewards":
                # Use reward reconstruction loss to weight alignment loss
                # (lower loss means less uncertainty; should align more)
                raise NotImplementedError
            else:
                raise ValueError(
                    "Unknown alignment weighting type: "
                    f"{self.args.informed.alignment_weighting_type}"
                )
            # So we can apply weighting to unreduced loss of shape [-1, B, latent_dim]
            horizon_weights = horizon_weights.unsqueeze(-1).unsqueeze(-1)  # [-1, 1, 1]

        # Define loss functions
        """Distributional alignment loss."""

        def _dist_al_loss(lms, llvs, zts):
            dist = torch.distributions.Normal(loc=lms, scale=torch.exp(llvs / 2))
            zts_mean = zts[..., : self.args.vae.latent_dim]  # first half is mean
            return -dist.log_prob(zts_mean)

        """Deterministic alignment loss."""

        def _det_al_loss(lms, llvs, zts):
            ls = torch.concatenate((lms, llvs), dim=-1)  # no semantic diff for det.
            return F.mse_loss(ls, zts, reduction="none")

        # Choose loss function based on latent type
        # TODO: we should be able to use det-type loss function on distr. latents
        # (i.e. only use mean; can implement)
        loss_fn = (
            _dist_al_loss if self.args.informed.distributional_latents else _det_al_loss
        )

        # Detach latents if necessary
        if detach_mode == "horizon_latent":
            alignment_loss = loss_fn(_lms.detach(), _llvs.detach(), _zts)
        elif detach_mode == "task_latent" or detach_mode == "extra":
            alignment_loss = loss_fn(_lms, _llvs, _zts.detach())
        elif detach_mode == "none":
            alignment_loss = loss_fn(_lms, _llvs, _zts)
        else:
            raise NotImplementedError(f"Unknown detach mode: {detach_mode}")

        # Apply horizon mask (axis 0) to unreduced shape: [-1, B, latent_dim]
        horizon_weights = horizon_weights.to(alignment_loss.device)
        alignment_loss *= horizon_weights

        # Aggregate over horizon dimension (axis 0)
        _agg_fn = torch.mean if self.args.vae.avg_elbo_terms else torch.sum
        # NOTE: we currently always mean across the latent dimension
        alignment_loss = _agg_fn(alignment_loss, dim=0).mean(dim=-1)  # [B]

        # Apply `mask`` to remaining batch dimension to ignore informed
        alignment_loss *= mask  # [B]
        alignment_loss_mean = alignment_loss.mean()  # [1]

        if keep_batch_dim:
            return alignment_loss
        else:
            return alignment_loss_mean

    def compute_vae_loss(
        self,
        update=False,
        pretrain_index=None,
        task_encoder=None,
        policy_optimizer=None,
        keep_batch_dim=False,
    ):
        """Returns the VAE loss"""
        if not self.rollout_storage.ready_for_update():
            if keep_batch_dim:
                return 0, None
            return 0

        # We don't perform alignment if (1) we don't use informedness at all,
        # or (2) we use informedness but don't use alignment loss
        no_alignment = (
            self.args.informed is None or not self.args.informed.alignment_loss
        )

        if (
            self.args.vae.disable_decoder
            and self.args.vae.disable_kl_term
            and no_alignment
        ):
            if keep_batch_dim:
                return 0, None
            return 0

        B = self.args.vae.batch_num_trajs
        # Get a mini-batch; each will be shape [H, B, -1]
        (
            vae_prev_obs,
            vae_next_obs,
            vae_actions,
            vae_rewards,
            vae_tasks,
            trajectory_lens,
            prior_deltas,
            prior_sigmas,
            informedness,
        ) = self.rollout_storage.get_batch(
            batchsize=B,
            sample_informedness=(self.args.informed.sample_vae_batch_informedness 
                                 if self.args.informed is not None else False),
            informed_perc=(self.args.informed.vae_informed_batch_perc 
                           if self.args.informed is not None else 0.0),
        )

        # Compute the RNN hidden state to act as the prior
        hidden_state = None
        if self.use_informedness:
            # create mask where informedness is 1 or 2
            if self.args.informed.vae_naive_percentage > 0:
                informedness_mask = (informedness > 0).squeeze(-1)  # only partial
                # Create mask of random bits, sampled from bernoulli with p
                p = self.args.informed.vae_naive_percentage
                to_zero_out = torch.bernoulli(
                    p * torch.ones_like(informedness)
                ).squeeze(-1)
                # Turn off bits of informedness_mask with random_mask
                informedness_mask = (informedness_mask * (1.0 - to_zero_out)).bool()
            else:
                informedness_mask = (informedness > 0).squeeze(-1)  # partial or full
            # initialize prior beliefs with zeros
            prior_beliefs = (
                torch.zeros((B, self.args.vae.latent_dim * 2))
                .float()
                .to(DeviceConfig.DEVICE)
            )
            # First, use tasks to compute the z_truths
            z_truths = task_encoder(vae_tasks)  # [B dim] NOTE: we want these gradients!
            # Then, use the deltas to get the perturbed z's
            z_means_pert = z_truths[:, : self.args.vae.latent_dim] + prior_deltas
            # Get logvars from sigmas
            z_logvars = torch.log(prior_sigmas**2 + 1e-20)
            # Clamp to lowest logvar we use
            z_logvars = torch.clamp(z_logvars, min=self.args.informed.min_logvar)
            # Finally, construct the prior beliefs with means and sigmas
            prior_beliefs[informedness_mask] = torch.cat(
                (z_means_pert, z_logvars), dim=-1
            )[informedness_mask]
            # And then encode them to get the hidden state
            hidden_state = self.encode_z_to_h(prior_beliefs)

        # Encode; shape will be [H+1, B, latent_dim] --- includes the prior!
        _, latent_mean, latent_logvar, _post_gru_args = self.encoder(
            actions=vae_actions,
            states=vae_next_obs,
            rewards=vae_rewards,
            hidden_state=hidden_state,
            return_prior=True,
            detach_every=self.args.vae.tbptt_stepsize
                if hasattr(self.args, "tbptt_stepsize") else None,
            return_post_gru_args=True  # for alignment loss (maybe)
        )

        # Create KL priors to use for KL loss (each are [1, B, latent_dim])
        prior_means = torch.zeros(1, *latent_mean.shape[1:]).to(DeviceConfig.DEVICE)
        prior_logvars = torch.zeros(1, *latent_logvar.shape[1:]).to(DeviceConfig.DEVICE)
        if self.use_informedness:
            _informed_means = z_means_pert  # [B, latent_sim]
            _informed_logvars = z_logvars  # [B, latent_sim]
            prior_means[:, informedness_mask] = _informed_means[informedness_mask]
            prior_logvars[:, informedness_mask] = _informed_logvars[informedness_mask]
        prior_means.detach()  # we detach because these are used as anchors for KL
        prior_logvars.detach()  # " "

        _cl_args = [latent_mean, latent_logvar, prior_means, prior_logvars, 
                    vae_prev_obs, vae_next_obs, vae_actions, vae_rewards,
                    vae_tasks, trajectory_lens]
        _cl_kwargs = {"keep_batch_dim": keep_batch_dim}

        def weight_vae_losses(_L_r_rec, _L_s_rec, _L_m_rec, _L_kl):
            return (
                self.args.vae.rew_loss_coeff * _L_r_rec +
                self.args.vae.state_loss_coeff * _L_s_rec +
                self.args.vae.task_loss_coeff * _L_m_rec +
                self.args.vae.kl_weight * _L_kl
            )

        # Compute VAE decoding loss using HORIZON ENCODER
        losses = self.compute_loss(*_cl_args, **_cl_kwargs)
        (L_r_rec, L_s_rec, L_m_rec, L_kl) = losses  # collect individual losses
        loss = weight_vae_losses(L_r_rec, L_s_rec, L_m_rec, L_kl)
        
        # Compute VAE decoding loss using TASK ENCODER
        if self.use_informedness and self.args.informed.vae_decode_from_m:
            _cl_args[0] = z_truths[:, : self.args.vae.latent_dim].unsqueeze(0).repeat(
                latent_mean.shape[0], 1, 1)
            # Use min_logvar for task latents
            _cl_args[1] = torch.full_like(_cl_args[0], self.args.informed.min_logvar)
            _losses_from_m = self.compute_loss(*_cl_args, **_cl_kwargs)
            (Lm_r_rec, Lm_s_rec, Lm_m_rec, Lm_kl) = _losses_from_m
            Lm_kl = 0  # we don't want KL loss from task latents, just reconstruction
            loss_from_m = weight_vae_losses(Lm_r_rec, Lm_s_rec, Lm_m_rec, Lm_kl)
            # Tack on this new loss to the original loss
            loss = loss + self.args.informed.vae_decode_from_m_coeff * loss_from_m
        else:
            Lm_r_rec = Lm_s_rec = Lm_m_rec = Lm_kl = None
            loss_from_m = None

        # Compute alignment loss
        alignment_loss = 0
        if self.use_informedness and self.args.informed.alignment_loss:
            # DEBUG - are we computing alignment loss two times?---one in ppo and one in VAE?
            # (the informed latents are meaningless, since they are not used!)
            if self.args.informed.alignment_include_informed_trajs:
                keep_mask = (
                    torch.ones_like(informedness)
                    .float()
                    .squeeze()
                    .to(DeviceConfig.DEVICE)
                )
            else:
                keep_mask = (informedness != 2).float().squeeze()
            alignment_loss = self.compute_alignment_loss(
                latent_mean,
                latent_logvar,
                z_truths,
                keep_batch_dim=keep_batch_dim,
                mask=keep_mask,
                post_gru_args=_post_gru_args,
            )
            loss += self.alignment_coeff * alignment_loss

        # Make sure we can compute gradients
        assert self._can_compute_gradients(losses)

        # Overall loss
        elbo_loss = loss  # [B]
        if self.use_informedness and not self.args.informed.vae_train_on_informed:
            # Mask out informed trajectories if applicable!
            non_informed_mask = (informedness != 2).float().squeeze()
            elbo_loss = elbo_loss * non_informed_mask
        elbo_loss_mean = elbo_loss.mean()  # [1]

        # Update
        if update:
            policy_optimizer.zero_grad()  # Zero policy optimizer gradients
            self.optimizer_vae.zero_grad()  # Zero gradients
            elbo_loss_mean.backward()  # Compute gradients
            self._clip_gradients()  # Clip gradients
            self.optimizer_vae.step()  # Update parameters
            if self.use_informedness and self.args.informed.alignment_loss:
                # step through policy as well---this is where the task encoder is
                policy_optimizer.step()

            self.update_count += 1

        self.log(
            elbo_loss_mean,
            L_r_rec,
            L_s_rec,
            L_m_rec,
            L_kl,
            pretrain_index,
            alignment_loss=alignment_loss,
        )
        if keep_batch_dim:
            return elbo_loss, informedness
        else:
            return elbo_loss_mean

    ###########################################################################
    #                         INDIVIDUAL LOSSES                               #
    ###########################################################################

    def compute_state_rec_loss(
        self, latent, prev_obs, next_obs, action, return_predictions=False
    ):
        """
        Compute state reconstruction loss. (No reduction of loss along batch
        dimension is done here; sum/avg has to be done outside)

        TODO(rsc): Explain in which cases this function is used.
        """

        state_pred = self.state_decoder(latent, prev_obs, action)

        if self.args.vae.state_pred_type == "deterministic":
            loss_state = (state_pred - next_obs).pow(2).mean(dim=-1)
        elif self.args.vae.state_pred_type == "gaussian":  # TODO: untested!
            state_pred_mean = state_pred[:, : state_pred.shape[1] // 2]
            state_pred_std = torch.exp(0.5 * state_pred[:, state_pred.shape[1] // 2 :])
            m = Normal(state_pred_mean, state_pred_std)
            loss_state = -m.log_prob(next_obs).mean(dim=-1)
        else:
            raise NotImplementedError

        if return_predictions:
            return loss_state, state_pred
        else:
            return loss_state

    def compute_rew_rec_loss(
        self,
        latent: torch.Tensor,
        prev_obs: torch.Tensor,
        next_obs: torch.Tensor,
        action: torch.Tensor,
        reward: torch.Tensor,
        return_predictions: bool = False,
    ):
        """
        Compute reward reconstruction loss. (No reduction of loss along batch
        dimension is done here; sum/avg has to be done outside)

        E = number of ELBO terms
        H = horizon length; number of steps in trajectory
        B = batch size; number of sampled trajectories

        Often T = E+1 (i.e. we have one ELBO term per timestep in the
        trajectory, and one for the prior)

        Args:
            latent:     [E, H, B, latent_dim]
            prev_obs:   [E, H, B, obs_dim]
            next_obs:   [E, H, B, obs_dim]
            action:     [E, H, B, action_dim]
            reward:     [E, H, B, reward_dim]
            return_predictions: bool

        TODO(robby): Explain the difference between multi-head reward prediction
        and single-head prediction in the docstring (currently in the comments).
        """

        if self.args.vae.multihead_for_reward:
            # Use multiple heads per reward pred (i.e. per state).
            # NOTE: This means that we need to be able to enumerate all states
            # and num_states should therefore be small. Otherwise we will have
            # too many heads, and we are better off sending in the state as
            # input instead of creating a separate head for each.

            # Use reward decoder to predict reward from encoded latent
            rew_pred = self.reward_decoder(latent, None)
            if self.args.vae.rew_pred_type == "categorical":
                # Categorical reward
                rew_pred = F.softmax(rew_pred, dim=-1)
            elif self.args.vae.rew_pred_type == "bernoulli":
                # Bernoulli reward
                rew_pred = torch.sigmoid(rew_pred)

            # Use next observations to get the ID of the task we will have
            # completed by landing in that state (?)
            state_indices = self.task_to_id(next_obs).to(DeviceConfig.DEVICE)

            # Ensure state indices and reward predictions have matching
            # dimensions; if state_indices has one fewer, we unsqueeze (i.e.
            # adding dimension at the end)
            if state_indices.dim() + 1 == rew_pred.dim():
                state_indices = state_indices.unsqueeze(-1)
            elif state_indices.dim() == rew_pred.dim():
                pass
            else:
                # They should either already match, or we should match by
                # unsqueezing once.
                raise ValueError
            # NOTE: state_indices is now shape:
            #                       [task_batch_size, num_steps, 1]

            # Gather values along final axis (because dim=-1), using task/state
            # IDs as indicies.
            # NOTE: rew_pred is currently of shape:
            #                       [task_batch_size, num_steps*4, num_states]
            # We use state_indices (which index the states) to choose which
            # reward prediction to use for each item in the batch
            rew_pred = rew_pred.gather(dim=-1, index=state_indices)
            # New shape of rew_pred:
            #                       [task_batch_size, num_steps*4, 1]

            # Depending on the reward prediction type, compute the loss between
            # the reward predictions and the targets
            rew_target = (reward == 1).float()
            if self.args.vae.rew_pred_type == "deterministic":  # TODO: untested!
                loss_rew = (rew_pred - reward).pow(2).mean(dim=-1)
            elif self.args.vae.rew_pred_type in ["categorical", "bernoulli"]:
                loss_rew = F.binary_cross_entropy(
                    rew_pred, rew_target, reduction="none"
                ).mean(dim=-1)
            else:
                raise NotImplementedError
        else:
            # Use one head per reward pred
            # NOTE: This is better when we have too many states to create a new
            # head for each, and we're better of sending in the state as
            # input.
            rew_pred = self.reward_decoder(latent, next_obs, prev_obs, action.float())
            if self.args.vae.rew_pred_type == "bernoulli":  # TODO: untested!
                rew_pred = torch.sigmoid(rew_pred)
                rew_target = (reward == 1).float()  # TODO: necessary?
                loss_rew = F.binary_cross_entropy(
                    rew_pred, rew_target, reduction="none"
                ).mean(dim=-1)
            elif self.args.vae.rew_pred_type == "deterministic":
                loss_rew = (rew_pred - reward).pow(2).mean(dim=-1)
            else:
                raise NotImplementedError

        if return_predictions:
            return loss_rew, rew_pred
        else:
            return loss_rew

    def compute_task_rec_loss(self, latent, task, return_predictions=False):
        """
        Compute task reconstruction loss. (No reduction of loss along batch
        dimension is done here; sum/avg has to be done outside)

        TODO(rsc): Explain when this function is used--it's not used for Varibad!
        Or even the belief_oracle... so who uses it?
        """
        # raise NotImplementedError
        task_pred = self.task_decoder(latent)
        if self.args.vae.task_pred_type == "task_id":
            task_target = self.task_to_id(task).to(DeviceConfig.DEVICE)
            # expand along first axis (number of ELBO terms)
            task_target = task_target.expand(task_pred.shape[:-1]).reshape(-1)
            loss_task = F.cross_entropy(
                task_pred.view(-1, task_pred.shape[-1]),
                task_target,
                reduction="none",
            ).view(task_pred.shape[:-1])
        elif self.args.vae.task_pred_type == "task_description":
            loss_task = (task_pred - task).pow(2).mean(dim=-1)
        else:
            raise NotImplementedError

        if return_predictions:
            return loss_task, task_pred
        else:
            return loss_task

    def compute_kl_loss(
        self,
        latent_mean,
        latent_logvar,
        elbo_indices,
        prior_means=None,
        prior_logvars=None,
    ):
        # -- KL divergence
        if self.args.vae.kl_to_gauss_prior:
            kl_divergences = -0.5 * (
                1 + latent_logvar - latent_mean.pow(2) - latent_logvar.exp()
            ).sum(dim=-1)
        else:
            gauss_dim = latent_mean.shape[-1]
            # add the gaussian prior

            # Assume uninformed
            if prior_means is None or prior_logvars is None:
                prior_means = (
                    torch.zeros_like(latent_mean[0])
                    .unsqueeze(0)
                    .to(DeviceConfig.DEVICE)
                )
                prior_logvars = (
                    torch.zeros_like(latent_logvar[0])
                    .unsqueeze(0)
                    .to(DeviceConfig.DEVICE)
                )

            all_means = torch.cat(
                (
                    prior_means,
                    latent_mean,
                )
            )
            all_logvars = torch.cat(
                (
                    prior_logvars,
                    latent_logvar,
                )
            )
            # https://arxiv.org/pdf/1811.09975.pdf
            # KL(N(mu,E)||N(m,S)) = 0.5 * (log(|S|/|E|) - K + tr(S^-1 E) +
            #                       (m-mu)^T S^-1 (m-mu)))
            mu = all_means[1:]
            m = all_means[:-1]
            logE = all_logvars[1:]
            logS = all_logvars[:-1]
            kl_divergences = 0.5 * (
                torch.sum(logS, dim=-1)
                - torch.sum(logE, dim=-1)
                - gauss_dim
                + torch.sum(1 / torch.exp(logS) * torch.exp(logE), dim=-1)
                + ((m - mu) / torch.exp(logS) * (m - mu)).sum(dim=-1)
            )

        # returns, for each ELBO_t term, one KL (so H+1 kl's)
        if elbo_indices is not None:
            batchsize = kl_divergences.shape[-1]
            task_indices = torch.arange(batchsize).repeat(self.args.vae.subsample_elbos)
            kl_divergences = kl_divergences[elbo_indices, task_indices].reshape(
                (self.args.vae.subsample_elbos, batchsize)
            )

        return kl_divergences

    ###########################################################################
    #                          HELPER FUNCTIONS                               #
    ###########################################################################

    def _can_compute_gradients(self, losses):
        (
            rew_reconstruction_loss,
            state_reconstruction_loss,
            task_reconstruction_loss,
            kl_loss,
        ) = losses
        if not self.args.vae.disable_kl_term:
            assert kl_loss.requires_grad
        if self.args.vae.decode_reward:
            assert rew_reconstruction_loss.requires_grad
        if self.args.vae.decode_state:
            assert state_reconstruction_loss.requires_grad
        if self.args.vae.decode_task:
            assert task_reconstruction_loss.requires_grad
        return True

    def _clip_gradients(self):
        # clip gradients
        if self.args.vae.encoder_max_grad_norm is not None:
            nn.utils.clip_grad_norm_(
                self.encoder.parameters(), self.args.encoder_max_grad_norm
            )
        if self.args.vae.decoder_max_grad_norm is not None:
            if self.args.vae.decode_reward:
                nn.utils.clip_grad_norm_(
                    self.reward_decoder.parameters(), self.args.decoder_max_grad_norm
                )
            if self.args.vae.decode_state:
                nn.utils.clip_grad_norm_(
                    self.state_decoder.parameters(), self.args.decoder_max_grad_norm
                )
            if self.args.vae.decode_task:
                nn.utils.clip_grad_norm_(
                    self.task_decoder.parameters(), self.args.decoder_max_grad_norm
                )

    ###########################################################################
    #                              LOGGING                                    #
    ###########################################################################

    def log(
        self,
        elbo_loss,
        rew_reconstruction_loss,
        state_reconstruction_loss,
        task_reconstruction_loss,
        kl_loss,
        pretrain_index=None,
        alignment_loss=None,
        Lm_r_rec=None,
        Lm_s_rec=None,
        Lm_m_rec=None,
        Lm_kl=None,
    ):
        if pretrain_index is None:
            curr_iter_idx = self.get_iter_idx()
        else:
            curr_iter_idx = (
                -self.args.vae.pretrain_len * self.args.vae.num_updates_per_pretrain
                + pretrain_index
            )

        if curr_iter_idx % self.args.domain.log_interval == 0:
            if self.args.vae.decode_reward:
                self.logger.add(
                    "vae_losses/reward_reconstr_err",
                    rew_reconstruction_loss.mean(),
                    curr_iter_idx,
                )
            if self.args.vae.decode_state:
                self.logger.add(
                    "vae_losses/state_reconstr_err",
                    state_reconstruction_loss.mean(),
                    curr_iter_idx,
                )
            if self.args.vae.decode_task:
                self.logger.add(
                    "vae_losses/task_reconstr_err",
                    task_reconstruction_loss.mean(),
                    curr_iter_idx,
                )

            if not self.args.vae.disable_kl_term:
                self.logger.add("vae_losses/kl", kl_loss.mean(), curr_iter_idx)

            if self.use_informedness and self.args.informed.alignment_loss:
                self.logger.add(
                    "vae_losses/alignment_loss_value",
                    alignment_loss.mean(),
                    curr_iter_idx,
                )

                self.logger.add(
                    "vae_losses/alignment_loss_coeff",
                    self.alignment_coeff,
                    curr_iter_idx,
                )

            self.logger.add("vae_losses/sum", elbo_loss, curr_iter_idx)

sparse_ant_goal.py:

import numpy as np
from gymnasium.envs.mujoco.ant_v4 import AntEnv


class AntGoalEnv(AntEnv):
    def __init__(self):
        super().__init(use_contact_forces=True)
        self.set_task(self.sample_tasks(1)[0])
        self.task_dim = 2

    def step(self, action):
        obs, reward, terminated, truncated, info = super().step(action)
        xposafter = self.get_body_com("torso").copy()

        goal_reward = -np.sum(np.abs(xposafter[:2] - self.goal_pos))

        ctrl_cost = 0.1 * np.square(action).sum()
        contact_cost = 0.5 * 1e-3 * np.sum(
            np.square(np.clip(self.data.cfrc_ext, -1, 1))
        )
        survive_reward = 0.0
        reward = goal_reward - ctrl_cost - contact_cost + survive_reward

        info.update({
            "goal_forward": goal_reward,
            "reward_ctrl": -ctrl_cost,
            "reward_contact": -contact_cost,
            "reward_survive": survive_reward,
            "task": self.get_task(),
        })

        return obs, reward, terminated, truncated, info

    def sample_tasks(self, num_tasks):
        a = np.random.uniform(0, 2 * np.pi, size=num_tasks)
        r = 3 * np.sqrt(np.random.uniform(0, 1, size=num_tasks))
        return np.stack((r * np.cos(a), r * np.sin(a)), axis=-1)

    def set_task(self, task):
        self.goal_pos = task

    def get_task(self):
        return np.array(self.goal_pos)

    def _get_obs(self):
        return np.concatenate([
            self.get_body_com("torso"),
            self.data.qpos.flat,
            self.data.qvel.flat,
            np.clip(self.data.cfrc_ext, -1, 1).flat,
        ])

class AntGoalSparseEnv(AntEnv):
    def __init__(self, max_episode_steps=200, dense_rewards=False, seed=None):
        super().__init__()
        self.task_dim = 2
        self.goal_radius = 0.75
        self._max_episode_steps = max_episode_steps  # Maximum number of steps per episode
        self.dense_rewards = dense_rewards  # Option to use dense rewards
        self.np_random = None  # Random generator, will be set in seed method
        self.seed(seed)  # Initialize the seed for reproducibility
        self.goal_pos = self.sample_tasks(1)[0]  # Set an initial goal position

    def seed(self, seed=None):
        """Set the seed for randomness."""
        self.np_random = np.random.RandomState(seed)
        return [seed]

    def step(self, action):
        obs, reward, terminated, truncated, info = super().step(action)
        xposafter = self.get_body_com("torso").copy()

        dist = np.linalg.norm(self.goal_pos - xposafter[:2], ord=2)
        success = dist <= self.goal_radius

        if success:
            goal_reward = 1.0 - dist if self.dense_rewards else 1.0
        else:
            goal_reward = -self.goal_radius

        # Calculate costs
        contact_cost = 0.5 * 1e-3 * np.sum(np.square(np.clip(self.data.cfrc_ext, -1, 1)))
        ctrl_cost = 0.01 * np.square(action).sum()
        survive_reward = 0.0

        # Calculate final reward
        reward = goal_reward - ctrl_cost - contact_cost + survive_reward

        # Update info dictionary with task details
        info.update({
            'goal_forward': goal_reward,
            'reward_ctrl': -ctrl_cost,
            'reward_contact': -contact_cost,
            'reward_survive': survive_reward,
            'task': self.get_task(),
            'success': success,
        })

        return obs, reward, terminated, truncated, info

    def sample_tasks(self, num_tasks):
        """Sample new goal positions for the environment."""
        angles = self.np_random.uniform(0, 2 * np.pi, size=num_tasks)
        radii = 3 * np.sqrt(self.np_random.uniform(0, 1, size=num_tasks))
        return np.stack((radii * np.cos(angles), radii * np.sin(angles)), axis=-1)

    def set_task(self, task):
        """Set the current goal task."""
        self.goal_pos = task

    def get_task(self):
        """Get the current goal position."""
        return np.array(self.goal_pos)

    def reset(self, *, seed=None, options=None, return_info=False):
        """Reset the environment and optionally set a new task."""
        # Reset the random seed if provided
        super().reset(seed=seed)
        
        # Sample a new task if not provided in options
        if options and "task" in options:
            task = options["task"]
        else:
            task = self.sample_tasks(1)[0]

        # Set the new task (goal)
        self.set_task(task)

        # Call Mujoco's reset model to reset the environment
        obs = self._get_obs()

        if return_info:
            return obs, {"task": self.get_task()}
        return obs

    def _get_obs(self):
        """Get the current observation."""
        return np.concatenate([
            self.get_body_com("torso"),
            self.data.qpos.flat,
            self.data.qvel.flat,
            np.clip(self.data.cfrc_ext, -1, 1).flat
        ])

    def reset_task(self, task=None):
        """Reset the task with an optional new task."""
        if task is None:
            task = self.sample_tasks(1)[0]
        self.set_task(task)

Error in storage.py:
Exception has occurred: RuntimeError       (note: full exception trace is shown but execution is paused at: _run_module_as_main)
The size of tensor a (28) must match the size of tensor b (117) at non-singleton dimension 1
  File "/home/ayan/home/partial-beliefs/metarl/components/policy/storage.py", line 262, in add_initial_state
    self.prev_state[0].copy_(state)
  File "/home/ayan/home/partial-beliefs/metarl/metalearner.py", line 374, in kickstart_training
    self.policy_storage.add_initial_state(self.env_salients.states)
  File "/home/ayan/home/partial-beliefs/metarl/metalearner.py", line 388, in train
    skip_training = self.kickstart_training()
                    ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayan/home/partial-beliefs/metarl/main.py", line 319, in main
    learner.train()
  File "/home/ayan/home/partial-beliefs/env/lib/python3.11/site-packages/hydra/core/utils.py", line 186, in run_job
    ret.return_value = task_function(task_cfg)
                       ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/ayan/home/partial-beliefs/env/lib/python3.11/site-packages/hydra/core/utils.py", line 260, in return_value
    raise self._return_value
  File "/home/ayan/home/partial-beliefs/env/lib/python3.11/site-packages/hydra/_internal/hydra.py", line 132, in run
    _ = ret.return_value
        ^^^^^^^^^^^^^^^^
  File "/home/ayan/home/partial-beliefs/env/lib/python3.11/site-packages/hydra/_internal/utils.py", line 458, in <lambda>
    lambda: hydra.run(
            ^^^^^^^^^^
  File "/home/ayan/home/partial-beliefs/env/lib/python3.11/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/home/ayan/home/partial-beliefs/env/lib/python3.11/site-packages/hydra/_internal/utils.py", line 223, in run_and_report
    raise ex
  File "/home/ayan/home/partial-beliefs/env/lib/python3.11/site-packages/hydra/_internal/utils.py", line 457, in _run_app
    run_and_report(
  File "/home/ayan/home/partial-beliefs/env/lib/python3.11/site-packages/hydra/_internal/utils.py", line 394, in _run_hydra
    _run_app(
  File "/home/ayan/home/partial-beliefs/env/lib/python3.11/site-packages/hydra/main.py", line 94, in decorated_main
    _run_hydra(
  File "/home/ayan/home/partial-beliefs/metarl/main.py", line 332, in <module>
    main()
  File "/home/ayan/home/partial-beliefs/env/lib/python3.11/runpy.py", line 88, in _run_code
    exec(code, run_globals)
  File "/home/ayan/home/partial-beliefs/env/lib/python3.11/runpy.py", line 198, in _run_module_as_main (Current frame)
    return _run_code(code, main_globals, None,
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: The size of tensor a (28) must match the size of tensor b (117) at non-singleton dimension 1