text: import torch as th
  import torch.nn as nn
  from typing import List, Tuple, Union
  from torchvision import transforms
  
  from autoencoding_rl.latent_extractors.Autoencoder import Autoencoder
  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.latent_extractors.Autoencoder import Autoencoder
  
  class VAE(Autoencoder):
      def __init__(self,  latent_space_size : int,
                          image_width : int,
                          image_height : int,
                          image_channels : int,
                          kld_loss_weight : float = 0.0001):
          super().__init__()
          self._latent_space_size = latent_space_size
          self._image_width = image_width
          self._image_height = image_height
          self._image_channels = image_channels
          self._kld_loss_weight = kld_loss_weight
          self._resizeToInput = transforms.Resize((self._image_height,self._image_width))
  
          self._buildModels()
  
      def _buildModels(self):
          self._encoder = VAE_encoder(latent_space_size = self._latent_space_size,
                                      image_channels_num = self._image_channels,
                                      net_input_width=self._image_width,
                                      net_input_height=self._image_height)
          self._decoder = VAE_decoder(latent_space_size=self._latent_space_size,
                                      output_channels_num=self._image_channels,
                                      net_output_width=self._image_width,
                                      net_output_height=self._image_height)
  
      def reset(self):        
          self._buildModels()
          
      def extract_latent_distribution(self, x : th.Tensor) -> Tuple[th.Tensor, th.Tensor]:
          return self._encoder(x)
  
      def encode(self, x : th.Tensor) -> th.Tensor:
          return self.extract_latent_distribution(x)[0]
      
      def decode(self, z : th.Tensor) -> th.Tensor:
          return self._decoder(z)
  
      def forward(self, image_batch: th.Tensor) -> Tuple[th.Tensor,th.Tensor,th.Tensor]:
          mu, log_var = self.extract_latent_distribution(image_batch)
          z = self._encoder.sample(mu, log_var)
          reconstruction = self.decode(z)
          return  (reconstruction, mu, log_var)
  
      def loss_function(self,     input_image_batch : th.Tensor,
                                  mu_batch : th.Tensor,
                                  log_var_batch : th.Tensor,
                                  reconstructed_image_batch : th.Tensor) -> th.Tensor:
  
          # 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)
          kld_loss = -th.mean(0.5 * th.sum(1 + log_var_batch - mu_batch ** 2 - log_var_batch.exp(), dim = 1), dim = 0)
  
          # 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.
          reconstruction_loss = th.nn.functional.mse_loss(input_image_batch, reconstructed_image_batch)
  
          loss = reconstruction_loss + self._kld_loss_weight*kld_loss
  
          # print(f"{reconstruction_loss}+{self._kld_loss_weight}*{kld_loss}")
  
          return loss
  
      def compute_loss(self, input_image_batch : th.Tensor):
          reconstructedImgsBatch, mu_batch, logvar_batch = self(input_image_batch)
          return self.loss_function(input_image_batch, mu_batch, logvar_batch, reconstructedImgsBatch)
  
  
  
  
  
  
  
  
  
  
      @property
      def encoding_size(self):
          return self._latent_space_size
  
      @property
      def input_width(self):
          return self._input_width
  
      @property
      def input_height(self):
          return self._input_height
  
      @property
      def input_channels(self):
          return self._input_channels
  
  
      
      def preprocess_observations(self, observation_batch : th.Tensor):
          resize_batch = self._resizeToInput(observation_batch)
          # Input should be in the [0,1] range, as this is what torchvision.transforms.ToTensor does
          normalized = resize_batch*2 - 1
          return normalized
  
      def postprocess_observations(self, observation_batch : th.Tensor):
          return (observation_batch+1)/2