<gh_stars>0
import torch as th
import torchvision
# import time
import cv2
import numpy as np
import datetime
import time
from typing import Dict, OrderedDict, Any, Optional, Union, List, Tuple

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

from autoencoding_rl.utils import tensorToCvImageRgb, Transition, ImageComparisonWindow
from autoencoding_rl.latent_extractors.autoencoder.SimpleAutoencoder import SimpleAutoencoder
from autoencoding_rl.latent_extractors.Autoencoder import Autoencoder
from autoencoding_rl.latent_extractors.LatentExtractor import LatentExtractor

class SimpleLatentExtractor(LatentExtractor):
    def __init__(self,
                 autoencoder : Autoencoder,
                 debug_out_folder : Optional[str] = None,
                 torchDevice : str = "cuda",
                 learning_rate : float = 0.001,
                 dbg_info_keys : List[str] = ["train_iterations_done","training_samples_used","time_from_start","loss","skippedMinibatches"]):
        super().__init__(debug_out_folder, dbg_info_keys)
        
        self._autoencoder = autoencoder
        self._torchDevice = torchDevice
        self._learning_rate = learning_rate

        self.reset()

    def torchDevice(self):
        return self._torchDevice

    def displayDebugInfo(self, transition_batch_example : Transition, id : str = None) -> None:
        if id is None:
            id = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
        if self._comparisonWindow is not None:
            with th.no_grad():
                preprocessed_obs_batch = self._autoencoder.preprocess_observations(transition_batch_example.observation.to(self._torchDevice))
                encodedDecoded_obs_batch = self._autoencoder(preprocessed_obs_batch)
                if isinstance(encodedDecoded_obs_batch,tuple):
                    encodedDecoded_obs_batch = encodedDecoded_obs_batch[0]

                encodedDecoded_obs_batch = self._autoencoder.postprocess_observations(encodedDecoded_obs_batch)
                preprocessed_obs_batch = self._autoencoder.postprocess_observations(preprocessed_obs_batch)
            
            self._comparisonWindow.displayExample(  predicted = encodedDecoded_obs_batch,
                                                    truth = preprocessed_obs_batch,
                                                    input = preprocessed_obs_batch,
                                                    id = id)

    def train_model(self, transition_batch : Transition) -> float:
        """Perform one training iteration on the provided batch

        Parameters
        ----------
        transition_batch : Transition
            A batch of transitions (Actually a transition of batches)
        """
        obs_batch = transition_batch.observation.to(self._torchDevice)
        lossVal : float = 0
        self._autoencoder.train() # put model in train mode
        self._optimizer.zero_grad()
        with th.set_grad_enabled(True):
            #print("observation_batch.type()="+str(observation_batch.type()))
            preprocessed_obss = self._autoencoder.preprocess_observations(obs_batch)
            loss = self._autoencoder.compute_loss(preprocessed_obss)
            lossVal = loss.item()
            loss.backward()
            th.nn.utils.clip_grad_value_(self._autoencoder.parameters(), 10)
            self._optimizer.step()
        
        self._info["train_iterations_done"] += 1
        self._info["training_samples_used"] += obs_batch.size()[0]
        self._info["time_from_start"] = time.monotonic()-self._initTime
        self._info["loss"] = lossVal
        self._writeToLog()
        return lossVal


    def evaluate(self, transition_batch : Transition) -> th.Tensor:
        """Evaluate model loss on the provided batch

        Parameters
        ----------
        transition_batch : Transition
            A batch of transitions (Actually a transition of batches)
        """

        with th.set_grad_enabled(False):
            obs_batch = transition_batch.observation.to(self._torchDevice)
            self._autoencoder.eval() # put model in eval mode
            preprocessed_obss = self._autoencoder.preprocess_observations(obs_batch)
            loss = self._autoencoder.compute_loss(preprocessed_obss)

            self._autoencoder.train() # put model in train mode
        return loss

    def extractFeatures(self, observation_batch : Dict[Union[str, int], th.Tensor]):
        
        # t0 = time.monotonic()
        with th.no_grad():
            preprocessed_obs_batch = self._autoencoder.preprocess_observations(observation_batch)
            # t1 = time.monotonic()
            #print(f"resize took {t1-t0}s")
            feature_vec_batch = self._autoencoder.encode(preprocessed_obs_batch)
            assert th.all(th.isfinite(feature_vec_batch))
            # t2 = time.monotonic()
            #print(f"encode took {t2-t1}s")
        # t4 = time.monotonic()
        return feature_vec_batch

    def preAndPostProcess(self, observation_batch : Dict[Union[str, int], th.Tensor]):
        with th.no_grad():
            pre = self._autoencoder.preprocess_observations(observation_batch)
            post = self._autoencoder.postprocess_observations(pre)
        return post

    def decode(self, encoded_obs_batch : th.Tensor, action_batch : th.Tensor = None):
        with th.set_grad_enabled(False):
            decoded_obs_batch = self._autoencoder.postprocess_observations(self._autoencoder.decode(encoded_obs_batch))
        return decoded_obs_batch

    @property
    def encoding_size(self):
        return self._autoencoder.encoding_size

    def reset(self):
        self._autoencoder.reset()
        self._autoencoder.to(self._torchDevice)
        self._optimizer = th.optim.Adam(self._autoencoder.parameters(), lr = self._learning_rate)
import torch as th
import torchvision
# import time
import cv2
import numpy as np
import datetime
import time
from typing import Dict, OrderedDict, Any, Optional, Union, List, Tuple

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

from autoencoding_rl.utils import tensorToCvImageRgb, Transition, ImageComparisonWindow
from autoencoding_rl.latent_extractors.autoencoder.SimpleAutoencoder import SimpleAutoencoder
from autoencoding_rl.latent_extractors.Autoencoder import Autoencoder
from autoencoding_rl.latent_extractors.LatentExtractor import LatentExtractor

class SimpleLatentExtractor(LatentExtractor):
    def __init__(self,
                 autoencoder : Autoencoder,
                 debug_out_folder : Optional[str] = None,
                 torchDevice : str = "cuda",
                 learning_rate : float = 0.001,
                 dbg_info_keys : List[str] = ["train_iterations_done","training_samples_used","time_from_start","loss","skippedMinibatches"]):
        super().__init__(debug_out_folder, dbg_info_keys)
        
        self._autoencoder = autoencoder
        self._torchDevice = torchDevice
        self._learning_rate = learning_rate

        self.reset()

    def torchDevice(self):
        return self._torchDevice

    def displayDebugInfo(self, transition_batch_example : Transition, id : str = None) -> None:
        if id is None:
            id = datetime.datetime.now().strftime('%Y%m%d-%H%M%S')
        if self._comparisonWindow is not None:
            with th.no_grad():
                preprocessed_obs_batch = self._autoencoder.preprocess_observations(transition_batch_example.observation.to(self._torchDevice))
                encodedDecoded_obs_batch = self._autoencoder(preprocessed_obs_batch)
                if isinstance(encodedDecoded_obs_batch,tuple):
                    encodedDecoded_obs_batch = encodedDecoded_obs_batch[0]

                encodedDecoded_obs_batch = self._autoencoder.postprocess_observations(encodedDecoded_obs_batch)
                preprocessed_obs_batch = self._autoencoder.postprocess_observations(preprocessed_obs_batch)
            
            self._comparisonWindow.displayExample(  predicted = encodedDecoded_obs_batch,
                                                    truth = preprocessed_obs_batch,
                                                    input = preprocessed_obs_batch,
                                                    id = id)

    def train_model(self, transition_batch : Transition) -> float:
        """Perform one training iteration on the provided batch

        Parameters
        ----------
        transition_batch : Transition
            A batch of transitions (Actually a transition of batches)
        """
        obs_batch = transition_batch.observation.to(self._torchDevice)
        lossVal : float = 0
        self._autoencoder.train() # put model in train mode
        self._optimizer.zero_grad()
        with th.set_grad_enabled(True):
            #print("observation_batch.type()="+str(observation_batch.type()))
            preprocessed_obss = self._autoencoder.preprocess_observations(obs_batch)
            loss = self._autoencoder.compute_loss(preprocessed_obss)
            lossVal = loss.item()
            loss.backward()
            th.nn.utils.clip_grad_value_(self._autoencoder.parameters(), 10)
            self._optimizer.step()
        
        self._info["train_iterations_done"] += 1
        self._info["training_samples_used"] += obs_batch.size()[0]
        self._info["time_from_start"] = time.monotonic()-self._initTime
        self._info["loss"] = lossVal
        self._writeToLog()
        return lossVal


    def evaluate(self, transition_batch : Transition) -> th.Tensor:
        """Evaluate model loss on the provided batch

        Parameters
        ----------
        transition_batch : Transition
            A batch of transitions (Actually a transition of batches)
        """

        with th.set_grad_enabled(False):
            obs_batch = transition_batch.observation.to(self._torchDevice)
            self._autoencoder.eval() # put model in eval mode
            preprocessed_obss = self._autoencoder.preprocess_observations(obs_batch)
            loss = self._autoencoder.compute_loss(preprocessed_obss)

            self._autoencoder.train() # put model in train mode
        return loss

    def extractFeatures(self, observation_batch : Dict[Union[str, int], th.Tensor]):
        
        # t0 = time.monotonic()
        with th.no_grad():
            preprocessed_obs_batch = self._autoencoder.preprocess_observations(observation_batch)
            # t1 = time.monotonic()
            #print(f"resize took {t1-t0}s")
            feature_vec_batch = self._autoencoder.encode(preprocessed_obs_batch)
            assert th.all(th.isfinite(feature_vec_batch))
            # t2 = time.monotonic()
            #print(f"encode took {t2-t1}s")
        # t4 = time.monotonic()
        return feature_vec_batch

    def preAndPostProcess(self, observation_batch : Dict[Union[str, int], th.Tensor]):
        with th.no_grad():
            pre = self._autoencoder.preprocess_observations(observation_batch)
            post = self._autoencoder.postprocess_observations(pre)
        return post

    def decode(self, encoded_obs_batch : th.Tensor, action_batch : th.Tensor = None):
        with th.set_grad_enabled(False):
            decoded_obs_batch = self._autoencoder.postprocess_observations(self._autoencoder.decode(encoded_obs_batch))
        return decoded_obs_batch

    @property
    def encoding_size(self):
        return self._autoencoder.encoding_size

    def reset(self):
        self._autoencoder.reset()
        self._autoencoder.to(self._torchDevice)
        self._optimizer = th.optim.Adam(self._autoencoder.parameters(), lr = self._learning_rate)
# import time Perform one training iteration on the provided batch          Parameters         ----------         transition_batch : Transition             A batch of transitions (Actually a transition of batches) # put model in train mode #print("observation_batch.type()="+str(observation_batch.type())) Evaluate model loss on the provided batch          Parameters         ----------         transition_batch : Transition             A batch of transitions (Actually a transition of batches) # put model in eval mode # put model in train mode # t0 = time.monotonic() # t1 = time.monotonic() #print(f"resize took {t1-t0}s") # t2 = time.monotonic() #print(f"encode took {t2-t1}s") # t4 = time.monotonic()