<reponame>c-rizz/autoencoding_rl

from typing import Tuple

import torch as th
import torch.nn as nn
from torchvision import transforms
from typing import Dict, Union, List, Tuple

import lr_gym.utils.dbg.ggLog as ggLog
import lr_gym

from autoencoding_rl.latent_extractors.vae.VAE_encoder import VAE_encoder
from autoencoding_rl.latent_extractors.vae.VAE_decoder import VAE_decoder
from autoencoding_rl.utils import Transition
from autoencoding_rl.utils import ObsConverter
from autoencoding_rl.latent_extractors.Autoencoder import Autoencoder

class MultiObsDynVAE(Autoencoder):
    
    def __init__(self,  observation_shape : Union[Dict[str,List[int]],Tuple[int,int,int]],
                        encoding_size: int,
                        action_size: int,
                        dynamics_nn_arch: Tuple[int, int],
                        rewardLossWeight : float = 1.0,
                        kdlLossWeight : float = 0.0001,
                        vecLossWeight : float = 0.005,
                        torchDevice : str = "cuda",
                        encoder_ensemble_size = 1,
                        dynamics_ensemble_size = 1):

        super().__init__()

        self._obs_converter = ObsConverter(observation_shape = observation_shape)
        img_shape_chw = self._obs_converter.imageSizeCHW()
        self._vec_part_size = self._obs_converter.vectorPartSize()
        self._torchDevice = torchDevice

        self._img_channels = img_shape_chw[0]
        self._img_height = img_shape_chw[1]
        self._img_width = img_shape_chw[2]
        self._img_encoding_size = encoding_size
        self._action_size = action_size
        self._dynamics_nn_arch = dynamics_nn_arch
        self._rewardWeight = rewardLossWeight
        self._kld_loss_weight = kdlLossWeight
        self._vec_lossWeight = vecLossWeight
        self._encoder_ensemble_size = encoder_ensemble_size
        self._dynamics_ensemble_size = dynamics_ensemble_size

        self._resizeToInput = transforms.Resize((self._img_height,self._img_width))

        self._buildModels()

    def _buildModels(self):
        self._encoder = VAE_encoder( latent_space_size = self._img_encoding_size,
                                        image_channels_num = self._img_channels,
                                        net_input_width = self._img_width,
                                        net_input_height = self._img_height,
                                        ensemble_size = self._encoder_ensemble_size)
        dyn_nets = []
        for _ in range(self._dynamics_ensemble_size):
            layersizes = []
            layersizes.append(self.encoding_size+self._action_size)
            layersizes += list(self._dynamics_nn_arch)
            layersizes.append(self.encoding_size+1)
            layers = []
            for i in range(len(layersizes)-1):
                layers.append(th.nn.Linear(layersizes[i],layersizes[i+1]))
                layers.append(th.nn.ReLU())
            dyn_nets.append(th.nn.Sequential(*layers))
        self._dynamics_nets = nn.ModuleList(dyn_nets)


        self._decoder = VAE_decoder(latent_space_size = self.encoding_size,
                                    output_channels_num = self._img_channels,
                                    net_output_width = self._img_width,
                                    net_output_height = self._img_height)

    def reset(self):        
        self._buildModels()

    def extract_img_latent_distribution(self, observation_batch : Dict[str,th.Tensor]):
        return self._encoder(observation_batch)

    def encode(self, observation_batch : Dict[str,th.Tensor]):
        img_batch = observation_batch["image"]
        vec_batch = observation_batch["vector"]
        enc_img_batch = self.extract_img_latent_distribution(img_batch)[0] #Use directly the mean, don't do sampling, sampling is only done for training
        return th.cat([enc_img_batch,vec_batch], dim = 1)

    def decode(self, predicted_state_batch : th.Tensor):
        #return self._decoder(dynamic_encoding_batch) #Gives a (batch_size, observations_channels_num, observation_height, observation_width) output
        predicted_vec_batch = predicted_state_batch[:,self._img_encoding_size : self._img_encoding_size + self._vec_part_size]
        predicted_img_batch = self._decoder(predicted_state_batch) #Gives a (batch_size, observations_channels_num, observation_height, observation_width) output

        predicted_observation_batch = { "image" : predicted_img_batch,
                                        "vector" : predicted_vec_batch}
        return predicted_observation_batch


    def predict_dynamics(self, state_batch : th.Tensor, action_batch : th.Tensor):
        assert state_batch.size()[0] == action_batch.size()[0], \
               f"state batch and action batch have different sizes, respectively {state_batch.size()[0]} and {action_batch.size()[0]}"
        assert state_batch.size()[1] == self.encoding_size, \
               f"States have wrong size, should be {self.encoding_size}, but it's {state_batch.size()[1]}"
        assert action_batch.size()[1] == self._action_size, \
               f"Actions have wrong size, should be {self._action_size} but it's {action_batch.size()[1]}"

        #Concatenate states and actions
        state_action_batch = th.cat((state_batch, action_batch), 1) #Gives a (batch_size, dyn_encoding_size+action_size) output
        nextstate_reward_batches = [dnet(state_action_batch) for dnet in self._dynamics_nets] #Gives a list of (batch_size, dyn_encoding_size+1) outputs
        nextstate_reward_batch = th.mean(th.stack(nextstate_reward_batches), dim=0) # Average out the ensemble outputs
        nextstate_batch, reward_batch = th.split(nextstate_reward_batch, [self.encoding_size, 1], 1)
        #nextstate_batch now has size (batch_size, dyn_encoding_size)
        #reward_batch now has size (batch_size, 1) (still 2-dimensional)
        return nextstate_batch, reward_batch

    def forward(self, transition_batch : Transition):
        img_batch = transition_batch.observation["image"]
        vec_batch = transition_batch.observation["vector"]
        action_batch = transition_batch.action
        assert action_batch.size()[0] == img_batch.size()[0], \
                f"Observation batch and action batch should have the same length. Action batch size = {action_batch.size()[0]}, image batch size = {img_batch.size()[0]}. Action tensor size = {action_batch.size()[0]}. Image tensor size = {img_batch.size()[0]}"
        assert action_batch.size()[0] == vec_batch.size()[0], \
                f"Vector batch and action batch should have the same length. Action batch size = {action_batch.size()[0]}, vector batch size = {vec_batch.size()[0]}. Action tensor size = {action_batch.size()}. Vector tensor size = {vec_batch.size()}"
        assert img_batch.size() == (img_batch.size()[0], self._img_channels, self._img_height, self._img_width), \
                f"Image size should be (Any, {self._img_channels}, {self._img_height}, {self._img_width}), instead it is  {img_batch.size()}"
        assert action_batch.size()[1] == self._action_size, \
                f"Each action should have size {self._action_size}, not {action_batch.size()[1]}. Tensor has size {action_batch.size()}"
        assert vec_batch.size()[1] == self._vec_part_size, \
                f"Each vector observation should have size {self._vec_part_size}, not {vec_batch.size()[1]}. Tensor has size {vec_batch.size()}"

        #Compute 'dynamic' features encoding
        if not th.all(th.isfinite(img_batch)):
            nan_detected = th.any(th.isnan(img_batch))
            raise RuntimeError(f"Invalid value detected in img_batch! nan_detected = {nan_detected}")
        if not th.all(th.isfinite(vec_batch)):
            nan_detected = th.any(th.isnan(vec_batch))
            raise RuntimeError(f"Invalid value detected in vec_batch! nan_detected = {nan_detected}")
        if not th.all(th.isfinite(action_batch)):
            nan_detected = th.any(th.isnan(action_batch))
            raise RuntimeError(f"Invalid value detected in action_batch! nan_detected = {nan_detected}")

        mu, log_var = self.extract_img_latent_distribution(img_batch)
        enc_img_batch = self._encoder.sample(mu, log_var) #Gives a (batch_size, dyn_encoding_size) output
        state_batch = th.cat([enc_img_batch,vec_batch], dim = 1)
        
        if not th.all(th.isfinite(state_batch)):
            raise RuntimeError(f"Invalid value detected in state_batch! state_batch = {state_batch}")
        
        predicted_state_batch, predicted_reward_batch = self.predict_dynamics(state_batch, action_batch)
        #predicted_state_batch now has size (batch_size, dyn_encoding_size)
        #predicted_reward_batch now has size (batch_size, 1) (still 2-dimensional)

        #Will now use 'static' features vectors and predicted states to predict the observation
        predicted_observation_batch = self.decode(predicted_state_batch) #Gives a (batch_size, observations_channels_num, observation_height, observation_width) output
        
        
        assert [th.all(th.isfinite(o)) for o in predicted_observation_batch.values()]
        assert th.all(th.isfinite(predicted_reward_batch))
        assert th.all(th.isfinite(mu))
        assert th.all(th.isfinite(log_var))
        
        return (predicted_observation_batch, predicted_reward_batch, mu, log_var)


    def loss_functions(self,    input_observation_batch : Dict[str,th.Tensor],
                                next_observation_batch : Dict[str,th.Tensor],
                                reward_batch : th.Tensor,
                                mu_batch : th.Tensor,
                                log_var_batch : th.Tensor,
                                predicted_observation_batch : Dict[str,th.Tensor],
                                predicted_reward_batch) -> Tuple[th.Tensor,th.Tensor,th.Tensor]:
        # ggLog.info("computing dyn vae loss")
        # ggLog.info(f"next_observation_batch.size() = {next_observation_batch.size()}")
        # ggLog.info(f"predicted_observation_batch.size() = {predicted_observation_batch.size()}")
        # ggLog.info(f"predicted_reward_batch.size() = {predicted_reward_batch.size()}")
        # ggLog.info(f"reward_batch.size() = {reward_batch.size()}")
        # x_i = input_image
        # (mu, logvar) describe the latent distribution conditioned by x_i

        next_img_batch = next_observation_batch["image"]
        predicted_img_batch =  predicted_observation_batch["image"]
        next_vec_batch = next_observation_batch["vector"]
        predicted_vec_batch =  predicted_observation_batch["vector"]
        assert th.all(th.isfinite(next_img_batch))
        assert th.all(th.isfinite(predicted_img_batch))
        assert th.all(th.isfinite(next_vec_batch))
        assert th.all(th.isfinite(predicted_vec_batch))

        # Compute KL-divergence loss
        # D_kl(q_theta(z|x_i) || p(z))
        # with:  q_theta(z|x_i) = N(mu(x_i), e^logvar(x_i))
        #        p(z) = N(0,I)
        # using mean instead of sum for the inner term to avoid dependance on the latent vector size
        kld_loss = -th.mean(0.5 * th.mean(1 + log_var_batch - mu_batch ** 2 - log_var_batch.exp(), dim = 1), dim = 0)
        assert th.all(th.isfinite(kld_loss))

        # Second part of the loss should theoretically maxmise:
        #    E_{z~q_theta}[log(p_phi(x_i|z))]
        # Which with q_theta(z|x_i) = N(mu, e^logvar) can be reparametrized:
        #    1/L*∑_i∈[0,L] [log(p_phi(x_i | mu_theta + var_theta * eps)] with eps~N(0,I)
        # i.e. it's a maximum log-likelihood term, in which phi and theta must be optimized
        # At this point there's some considerations:
        #  - we can just set L = 1 if batch sizes are big enough (~100)
        #  - we don't actually output a probability distribution with our decoder, we directly output the image
        # Because of the second consideration we actually just use the MSE loss for this term.
        reward_loss = self._rewardWeight*th.nn.MSELoss()(predicted_reward_batch, reward_batch)
        if next_vec_batch.size()[1] != 0:
            vec_loss = self._vec_lossWeight*th.nn.MSELoss()(next_vec_batch, predicted_vec_batch)
        else:
            vec_loss = th.tensor(0)
        img_loss = th.nn.MSELoss()(next_img_batch, predicted_img_batch)
        reconstruction_loss = img_loss + reward_loss + vec_loss
        assert th.all(th.isfinite(reward_loss))
        assert th.all(th.isfinite(vec_loss))
        assert th.all(th.isfinite(img_loss))
        assert th.all(th.isfinite(reconstruction_loss))

        # print(f"{reconstruction_loss}+{self._kld_loss_weight}*{kld_loss}")

        loss = reconstruction_loss + self._kld_loss_weight*kld_loss

        # ggLog.info("computed dyn vae loss")
        return loss, reconstruction_loss, kld_loss, reward_loss


    def loss_function(self,     input_observation_batch : Dict[str,th.Tensor],
                                next_observation_batch : Dict[str,th.Tensor],
                                reward_batch : th.Tensor,
                                mu_batch : th.Tensor,
                                log_var_batch : th.Tensor,
                                predicted_observation_batch : Dict[str,th.Tensor],
                                predicted_reward_batch) -> th.Tensor:
        loss, reconstruction_loss, kld_loss, reward_loss = self.loss_functions(input_observation_batch,
                                                                    next_observation_batch,
                                                                    reward_batch,
                                                                    mu_batch,
                                                                    log_var_batch,
                                                                    predicted_observation_batch,
                                                                    predicted_reward_batch)

        return loss

    def compute_loss(self,  transition_batch : Transition):
        predicted_observation_batch, predicted_reward_batch, mu_batch, log_var_batch = self(transition_batch)
        return self.loss_function(  transition_batch.observation,
                                    transition_batch.next_observation,
                                    transition_batch.reward,
                                    mu_batch,
                                    log_var_batch,
                                    predicted_observation_batch,
                                    predicted_reward_batch)

    def compute_losses(self,  transition_batch : Transition):
        predicted_observation_batch, predicted_reward_batch, mu_batch, log_var_batch = self(transition_batch)
        return  self.loss_functions(    transition_batch.observation,
                                        transition_batch.next_observation,
                                        transition_batch.reward,
                                        mu_batch,
                                        log_var_batch,
                                        predicted_observation_batch,
                                        predicted_reward_batch)








    @property
    def encoding_size(self):
        return self._img_encoding_size + self._vec_part_size

    @property
    def input_height(self):
        return self._img_height

    @property
    def input_width(self):
        return self._img_width

    @property
    def input_channels(self):
        return self._img_channels

    @property 
    def action_size(self):
        return self._action_size



    def preprocess_observations(self, observation_batch : Dict[str, th.Tensor]):
        #print("observation_batch.dtype = ",self._obs_converter.getImgPart(observation_batch).dtype)
        resized_img_batch = self._resizeToInput(self._obs_converter.getImgPart(observation_batch).to(self._torchDevice))
        #print("resized_img_batch = ",resized_img_batch.dtype)
        # Input should be in the [0,1] range, as this is what torchvision.transforms.ToTensor does
        # We move it to [-1,1]
        normalized_img_batch = resized_img_batch*2 - 1
        #print("normalized_img_batch = ",normalized_img_batch.dtype)
        return { "image" : normalized_img_batch,
                 "vector" : self._obs_converter.getVectorPart(observation_batch).to(self._torchDevice)}   

    def postprocess_observations(self, observation_batch : Dict[str,th.Tensor]):
        unnormalized_img_batch = (observation_batch["image"] + 1)/2
        return self._obs_converter.buildDictObs(vectorPart_batch = observation_batch["vector"],
                                                imgPart_batch = unnormalized_img_batch)

from typing import Tuple

import torch as th
import torch.nn as nn
from torchvision import transforms
from typing import Dict, Union, List, Tuple

import lr_gym.utils.dbg.ggLog as ggLog
import lr_gym

from autoencoding_rl.latent_extractors.vae.VAE_encoder import VAE_encoder
from autoencoding_rl.latent_extractors.vae.VAE_decoder import VAE_decoder
from autoencoding_rl.utils import Transition
from autoencoding_rl.utils import ObsConverter
from autoencoding_rl.latent_extractors.Autoencoder import Autoencoder

class MultiObsDynVAE(Autoencoder):
    
    def __init__(self,  observation_shape : Union[Dict[str,List[int]],Tuple[int,int,int]],
                        encoding_size: int,
                        action_size: int,
                        dynamics_nn_arch: Tuple[int, int],
                        rewardLossWeight : float = 1.0,
                        kdlLossWeight : float = 0.0001,
                        vecLossWeight : float = 0.005,
                        torchDevice : str = "cuda",
                        encoder_ensemble_size = 1,
                        dynamics_ensemble_size = 1):

        super().__init__()

        self._obs_converter = ObsConverter(observation_shape = observation_shape)
        img_shape_chw = self._obs_converter.imageSizeCHW()
        self._vec_part_size = self._obs_converter.vectorPartSize()
        self._torchDevice = torchDevice

        self._img_channels = img_shape_chw[0]
        self._img_height = img_shape_chw[1]
        self._img_width = img_shape_chw[2]
        self._img_encoding_size = encoding_size
        self._action_size = action_size
        self._dynamics_nn_arch = dynamics_nn_arch
        self._rewardWeight = rewardLossWeight
        self._kld_loss_weight = kdlLossWeight
        self._vec_lossWeight = vecLossWeight
        self._encoder_ensemble_size = encoder_ensemble_size
        self._dynamics_ensemble_size = dynamics_ensemble_size

        self._resizeToInput = transforms.Resize((self._img_height,self._img_width))

        self._buildModels()

    def _buildModels(self):
        self._encoder = VAE_encoder( latent_space_size = self._img_encoding_size,
                                        image_channels_num = self._img_channels,
                                        net_input_width = self._img_width,
                                        net_input_height = self._img_height,
                                        ensemble_size = self._encoder_ensemble_size)
        dyn_nets = []
        for _ in range(self._dynamics_ensemble_size):
            layersizes = []
            layersizes.append(self.encoding_size+self._action_size)
            layersizes += list(self._dynamics_nn_arch)
            layersizes.append(self.encoding_size+1)
            layers = []
            for i in range(len(layersizes)-1):
                layers.append(th.nn.Linear(layersizes[i],layersizes[i+1]))
                layers.append(th.nn.ReLU())
            dyn_nets.append(th.nn.Sequential(*layers))
        self._dynamics_nets = nn.ModuleList(dyn_nets)


        self._decoder = VAE_decoder(latent_space_size = self.encoding_size,
                                    output_channels_num = self._img_channels,
                                    net_output_width = self._img_width,
                                    net_output_height = self._img_height)

    def reset(self):        
        self._buildModels()

    def extract_img_latent_distribution(self, observation_batch : Dict[str,th.Tensor]):
        return self._encoder(observation_batch)

    def encode(self, observation_batch : Dict[str,th.Tensor]):
        img_batch = observation_batch["image"]
        vec_batch = observation_batch["vector"]
        enc_img_batch = self.extract_img_latent_distribution(img_batch)[0] #Use directly the mean, don't do sampling, sampling is only done for training
        return th.cat([enc_img_batch,vec_batch], dim = 1)

    def decode(self, predicted_state_batch : th.Tensor):
        #return self._decoder(dynamic_encoding_batch) #Gives a (batch_size, observations_channels_num, observation_height, observation_width) output
        predicted_vec_batch = predicted_state_batch[:,self._img_encoding_size : self._img_encoding_size + self._vec_part_size]
        predicted_img_batch = self._decoder(predicted_state_batch) #Gives a (batch_size, observations_channels_num, observation_height, observation_width) output

        predicted_observation_batch = { "image" : predicted_img_batch,
                                        "vector" : predicted_vec_batch}
        return predicted_observation_batch


    def predict_dynamics(self, state_batch : th.Tensor, action_batch : th.Tensor):
        assert state_batch.size()[0] == action_batch.size()[0], \
               f"state batch and action batch have different sizes, respectively {state_batch.size()[0]} and {action_batch.size()[0]}"
        assert state_batch.size()[1] == self.encoding_size, \
               f"States have wrong size, should be {self.encoding_size}, but it's {state_batch.size()[1]}"
        assert action_batch.size()[1] == self._action_size, \
               f"Actions have wrong size, should be {self._action_size} but it's {action_batch.size()[1]}"

        #Concatenate states and actions
        state_action_batch = th.cat((state_batch, action_batch), 1) #Gives a (batch_size, dyn_encoding_size+action_size) output
        nextstate_reward_batches = [dnet(state_action_batch) for dnet in self._dynamics_nets] #Gives a list of (batch_size, dyn_encoding_size+1) outputs
        nextstate_reward_batch = th.mean(th.stack(nextstate_reward_batches), dim=0) # Average out the ensemble outputs
        nextstate_batch, reward_batch = th.split(nextstate_reward_batch, [self.encoding_size, 1], 1)
        #nextstate_batch now has size (batch_size, dyn_encoding_size)
        #reward_batch now has size (batch_size, 1) (still 2-dimensional)
        return nextstate_batch, reward_batch

    def forward(self, transition_batch : Transition):
        img_batch = transition_batch.observation["image"]
        vec_batch = transition_batch.observation["vector"]
        action_batch = transition_batch.action
        assert action_batch.size()[0] == img_batch.size()[0], \
                f"Observation batch and action batch should have the same length. Action batch size = {action_batch.size()[0]}, image batch size = {img_batch.size()[0]}. Action tensor size = {action_batch.size()[0]}. Image tensor size = {img_batch.size()[0]}"
        assert action_batch.size()[0] == vec_batch.size()[0], \
                f"Vector batch and action batch should have the same length. Action batch size = {action_batch.size()[0]}, vector batch size = {vec_batch.size()[0]}. Action tensor size = {action_batch.size()}. Vector tensor size = {vec_batch.size()}"
        assert img_batch.size() == (img_batch.size()[0], self._img_channels, self._img_height, self._img_width), \
                f"Image size should be (Any, {self._img_channels}, {self._img_height}, {self._img_width}), instead it is  {img_batch.size()}"
        assert action_batch.size()[1] == self._action_size, \
                f"Each action should have size {self._action_size}, not {action_batch.size()[1]}. Tensor has size {action_batch.size()}"
        assert vec_batch.size()[1] == self._vec_part_size, \
                f"Each vector observation should have size {self._vec_part_size}, not {vec_batch.size()[1]}. Tensor has size {vec_batch.size()}"

        #Compute 'dynamic' features encoding
        if not th.all(th.isfinite(img_batch)):
            nan_detected = th.any(th.isnan(img_batch))
            raise RuntimeError(f"Invalid value detected in img_batch! nan_detected = {nan_detected}")
        if not th.all(th.isfinite(vec_batch)):
            nan_detected = th.any(th.isnan(vec_batch))
            raise RuntimeError(f"Invalid value detected in vec_batch! nan_detected = {nan_detected}")
        if not th.all(th.isfinite(action_batch)):
            nan_detected = th.any(th.isnan(action_batch))
            raise RuntimeError(f"Invalid value detected in action_batch! nan_detected = {nan_detected}")

        mu, log_var = self.extract_img_latent_distribution(img_batch)
        enc_img_batch = self._encoder.sample(mu, log_var) #Gives a (batch_size, dyn_encoding_size) output
        state_batch = th.cat([enc_img_batch,vec_batch], dim = 1)
        
        if not th.all(th.isfinite(state_batch)):
            raise RuntimeError(f"Invalid value detected in state_batch! state_batch = {state_batch}")
        
        predicted_state_batch, predicted_reward_batch = self.predict_dynamics(state_batch, action_batch)
        #predicted_state_batch now has size (batch_size, dyn_encoding_size)
        #predicted_reward_batch now has size (batch_size, 1) (still 2-dimensional)

        #Will now use 'static' features vectors and predicted states to predict the observation
        predicted_observation_batch = self.decode(predicted_state_batch) #Gives a (batch_size, observations_channels_num, observation_height, observation_width) output
        
        
        assert [th.all(th.isfinite(o)) for o in predicted_observation_batch.values()]
        assert th.all(th.isfinite(predicted_reward_batch))
        assert th.all(th.isfinite(mu))
        assert th.all(th.isfinite(log_var))
        
        return (predicted_observation_batch, predicted_reward_batch, mu, log_var)


    def loss_functions(self,    input_observation_batch : Dict[str,th.Tensor],
                                next_observation_batch : Dict[str,th.Tensor],
                                reward_batch : th.Tensor,
                                mu_batch : th.Tensor,
                                log_var_batch : th.Tensor,
                                predicted_observation_batch : Dict[str,th.Tensor],
                                predicted_reward_batch) -> Tuple[th.Tensor,th.Tensor,th.Tensor]:
        # ggLog.info("computing dyn vae loss")
        # ggLog.info(f"next_observation_batch.size() = {next_observation_batch.size()}")
        # ggLog.info(f"predicted_observation_batch.size() = {predicted_observation_batch.size()}")
        # ggLog.info(f"predicted_reward_batch.size() = {predicted_reward_batch.size()}")
        # ggLog.info(f"reward_batch.size() = {reward_batch.size()}")
        # x_i = input_image
        # (mu, logvar) describe the latent distribution conditioned by x_i

        next_img_batch = next_observation_batch["image"]
        predicted_img_batch =  predicted_observation_batch["image"]
        next_vec_batch = next_observation_batch["vector"]
        predicted_vec_batch =  predicted_observation_batch["vector"]
        assert th.all(th.isfinite(next_img_batch))
        assert th.all(th.isfinite(predicted_img_batch))
        assert th.all(th.isfinite(next_vec_batch))
        assert th.all(th.isfinite(predicted_vec_batch))

        # Compute KL-divergence loss
        # D_kl(q_theta(z|x_i) || p(z))
        # with:  q_theta(z|x_i) = N(mu(x_i), e^logvar(x_i))
        #        p(z) = N(0,I)
        # using mean instead of sum for the inner term to avoid dependance on the latent vector size
        kld_loss = -th.mean(0.5 * th.mean(1 + log_var_batch - mu_batch ** 2 - log_var_batch.exp(), dim = 1), dim = 0)
        assert th.all(th.isfinite(kld_loss))

        # Second part of the loss should theoretically maxmise:
        #    E_{z~q_theta}[log(p_phi(x_i|z))]
        # Which with q_theta(z|x_i) = N(mu, e^logvar) can be reparametrized:
        #    1/L*∑_i∈[0,L] [log(p_phi(x_i | mu_theta + var_theta * eps)] with eps~N(0,I)
        # i.e. it's a maximum log-likelihood term, in which phi and theta must be optimized
        # At this point there's some considerations:
        #  - we can just set L = 1 if batch sizes are big enough (~100)
        #  - we don't actually output a probability distribution with our decoder, we directly output the image
        # Because of the second consideration we actually just use the MSE loss for this term.
        reward_loss = self._rewardWeight*th.nn.MSELoss()(predicted_reward_batch, reward_batch)
        if next_vec_batch.size()[1] != 0:
            vec_loss = self._vec_lossWeight*th.nn.MSELoss()(next_vec_batch, predicted_vec_batch)
        else:
            vec_loss = th.tensor(0)
        img_loss = th.nn.MSELoss()(next_img_batch, predicted_img_batch)
        reconstruction_loss = img_loss + reward_loss + vec_loss
        assert th.all(th.isfinite(reward_loss))
        assert th.all(th.isfinite(vec_loss))
        assert th.all(th.isfinite(img_loss))
        assert th.all(th.isfinite(reconstruction_loss))

        # print(f"{reconstruction_loss}+{self._kld_loss_weight}*{kld_loss}")

        loss = reconstruction_loss + self._kld_loss_weight*kld_loss

        # ggLog.info("computed dyn vae loss")
        return loss, reconstruction_loss, kld_loss, reward_loss


    def loss_function(self,     input_observation_batch : Dict[str,th.Tensor],
                                next_observation_batch : Dict[str,th.Tensor],
                                reward_batch : th.Tensor,
                                mu_batch : th.Tensor,
                                log_var_batch : th.Tensor,
                                predicted_observation_batch : Dict[str,th.Tensor],
                                predicted_reward_batch) -> th.Tensor:
        loss, reconstruction_loss, kld_loss, reward_loss = self.loss_functions(input_observation_batch,
                                                                    next_observation_batch,
                                                                    reward_batch,
                                                                    mu_batch,
                                                                    log_var_batch,
                                                                    predicted_observation_batch,
                                                                    predicted_reward_batch)

        return loss

    def compute_loss(self,  transition_batch : Transition):
        predicted_observation_batch, predicted_reward_batch, mu_batch, log_var_batch = self(transition_batch)
        return self.loss_function(  transition_batch.observation,
                                    transition_batch.next_observation,
                                    transition_batch.reward,
                                    mu_batch,
                                    log_var_batch,
                                    predicted_observation_batch,
                                    predicted_reward_batch)

    def compute_losses(self,  transition_batch : Transition):
        predicted_observation_batch, predicted_reward_batch, mu_batch, log_var_batch = self(transition_batch)
        return  self.loss_functions(    transition_batch.observation,
                                        transition_batch.next_observation,
                                        transition_batch.reward,
                                        mu_batch,
                                        log_var_batch,
                                        predicted_observation_batch,
                                        predicted_reward_batch)








    @property
    def encoding_size(self):
        return self._img_encoding_size + self._vec_part_size

    @property
    def input_height(self):
        return self._img_height

    @property
    def input_width(self):
        return self._img_width

    @property
    def input_channels(self):
        return self._img_channels

    @property 
    def action_size(self):
        return self._action_size



    def preprocess_observations(self, observation_batch : Dict[str, th.Tensor]):
        #print("observation_batch.dtype = ",self._obs_converter.getImgPart(observation_batch).dtype)
        resized_img_batch = self._resizeToInput(self._obs_converter.getImgPart(observation_batch).to(self._torchDevice))
        #print("resized_img_batch = ",resized_img_batch.dtype)
        # Input should be in the [0,1] range, as this is what torchvision.transforms.ToTensor does
        # We move it to [-1,1]
        normalized_img_batch = resized_img_batch*2 - 1
        #print("normalized_img_batch = ",normalized_img_batch.dtype)
        return { "image" : normalized_img_batch,
                 "vector" : self._obs_converter.getVectorPart(observation_batch).to(self._torchDevice)}   

    def postprocess_observations(self, observation_batch : Dict[str,th.Tensor]):
        unnormalized_img_batch = (observation_batch["image"] + 1)/2
        return self._obs_converter.buildDictObs(vectorPart_batch = observation_batch["vector"],
                                                imgPart_batch = unnormalized_img_batch)
#Use directly the mean, don't do sampling, sampling is only done for training #return self._decoder(dynamic_encoding_batch) #Gives a (batch_size, observations_channels_num, observation_height, observation_width) output #Gives a (batch_size, observations_channels_num, observation_height, observation_width) output #Concatenate states and actions #Gives a (batch_size, dyn_encoding_size+action_size) output #Gives a list of (batch_size, dyn_encoding_size+1) outputs # Average out the ensemble outputs #nextstate_batch now has size (batch_size, dyn_encoding_size) #reward_batch now has size (batch_size, 1) (still 2-dimensional) #Compute 'dynamic' features encoding #Gives a (batch_size, dyn_encoding_size) output #predicted_state_batch now has size (batch_size, dyn_encoding_size) #predicted_reward_batch now has size (batch_size, 1) (still 2-dimensional) #Will now use 'static' features vectors and predicted states to predict the observation #Gives a (batch_size, observations_channels_num, observation_height, observation_width) output # ggLog.info("computing dyn vae loss") # ggLog.info(f"next_observation_batch.size() = {next_observation_batch.size()}") # ggLog.info(f"predicted_observation_batch.size() = {predicted_observation_batch.size()}") # ggLog.info(f"predicted_reward_batch.size() = {predicted_reward_batch.size()}") # ggLog.info(f"reward_batch.size() = {reward_batch.size()}") # x_i = input_image # (mu, logvar) describe the latent distribution conditioned by x_i # Compute KL-divergence loss # D_kl(q_theta(z|x_i) || p(z)) # with:  q_theta(z|x_i) = N(mu(x_i), e^logvar(x_i)) #        p(z) = N(0,I) # using mean instead of sum for the inner term to avoid dependance on the latent vector size # Second part of the loss should theoretically maxmise: #    E_{z~q_theta}[log(p_phi(x_i|z))] # Which with q_theta(z|x_i) = N(mu, e^logvar) can be reparametrized: #    1/L*∑_i∈[0,L] [log(p_phi(x_i | mu_theta + var_theta * eps)] with eps~N(0,I) # i.e. it's a maximum log-likelihood term, in which phi and theta must be optimized # At this point there's some considerations: #  - we can just set L = 1 if batch sizes are big enough (~100) #  - we don't actually output a probability distribution with our decoder, we directly output the image # Because of the second consideration we actually just use the MSE loss for this term. # print(f"{reconstruction_loss}+{self._kld_loss_weight}*{kld_loss}") # ggLog.info("computed dyn vae loss") #print("observation_batch.dtype = ",self._obs_converter.getImgPart(observation_batch).dtype) #print("resized_img_batch = ",resized_img_batch.dtype) # Input should be in the [0,1] range, as this is what torchvision.transforms.ToTensor does # We move it to [-1,1] #print("normalized_img_batch = ",normalized_img_batch.dtype)