<reponame>FineArtz/ILSwiss
from base64 import encode
from collections import OrderedDict
from copy import deepcopy
import itertools
from typing import Type

import numpy as np
import torch
from torch import dropout, nn as nn
import torch.nn.functional as F
import torch.optim as optim
from rlkit.core import logger
from rlkit.data_management.env_replay_buffer import EnvReplayBuffer
from rlkit.data_management.replay_buffer import ReplayBuffer

from rlkit.torch.algorithms.mbpo.encoder import Encoder, IdentityEncoder
import rlkit.torch.utils.pytorch_util as ptu
from rlkit.torch.utils.normalizer import preprocess_obs
from rlkit.core.trainer import Trainer
from rlkit.core.eval_util import create_stats_ordered_dict


class SoftActorCritic(Trainer):
    """
    version that:
        - uses encoder and CURL loss
        - uses reparameterization trick
        - has two Q functions
        - has auto-tuned alpha
    """

    def __init__(
        self,
        encoder: Encoder,
        policy: nn.Module,
        qf1: nn.Module,
        qf2: nn.Module,

        reward_scale: float = 1.0,
        discount: float = 0.99,
        policy_lr: float = 1e-3,
        qf_lr: float = 1e-3,
        alpha_lr: float = 1e-4,
        encoder_lr: float = 1e-3,
        soft_target_tau: float = 0.01,
        enc_soft_target_tau: float = 0.05,
        target_entropy: float = None,

        alpha: float = 0.1,
        train_alpha: bool = True,
        policy_mean_reg_weight: float = 1e-3,
        policy_std_reg_weight: float = 1e-3,
        optimizer_class: Type[optim.Optimizer] = optim.Adam,
        beta_1: float = 0.9,
        alpha_beta: float = 0.5,

        encoder_latent_lambda: float = 1e-6,
        ac_update_freq: int = 2,
        encoder_update_freq: int = 1,
        **kwargs
    ):
        self.encoder = encoder
        self.enc_no_target = kwargs.get('enc_no_target', False) # not use target encoder
        self.enc_model: nn.Module = kwargs.get('enc_model', None) # use model to regularize encoder
        if self.enc_model: 
            self.enc_model_reg: float = kwargs.get('enc_model_reg', 0.1)
            self.enc_model_lr: float = kwargs.get('enc_model_lr', 1e-3)
            self.enc_target_reg: float = kwargs.get('enc_target_reg', 0.001)
            self.enc_model_freq: int = kwargs.get('enc_model_freq', 1)
            self.train_model_after: int = kwargs.get('train_model_after', 0)
            print('Enc Model Freq =', self.enc_model_freq)
            print('Enc Model After =', self.train_model_after)
            self.enc_model_optimizer = optimizer_class(self.enc_model.parameters(), lr=self.enc_model_lr)
            
        self.latent_q: bool = kwargs.get('latent_q', True) # only use q in latent space
        self.distill_qf: nn.Module = kwargs.get('distill_qf', None) # distill q from original space to latent space
        self.double_q: bool = kwargs.get('double_q', False) # both q in original space and latent space
        # if self.double_q:
        #     assert not self.latent_q, self.latent_q
        #     assert not self.distill_qf, self.distill_qf
        #     self.lat_qf1: nn.Module = kwargs['lat_qf1']
        #     self.lat_qf2: nn.Module = kwargs['lat_qf2']
        #     self.target_lat_qf1: nn.Module = self.lat_qf1.copy()
        #     self.target_lat_qf2: nn.Module = self.lat_qf2.copy()
        #     self.lat_qf_optimizer = optimizer_class(
        #         list(self.lat_qf1.parameters())
        #         + list(self.lat_qf2.parameters()),
        #         lr=qf_lr,
        #         betas=(beta_1, 0.999),
        #     )
        #     # latent q update mode
        #     # 0: lat_q use origin_q as target for data in origin space, 
        #     # and target_lat_q as target for generated data in latent space
        #     self.lat_update_mode = kwargs.get('lat_update_mode', 0)
        #     self.lat_update_iter = kwargs.get('lat_update_iter', 1)
        #     self.lat_q_enc_reg = kwargs.get('lat_q_enc_reg', 0.0)

        self.policy = policy
        self.qf1 = qf1
        self.qf2 = qf2
        self.reward_scale = reward_scale
        self.discount = discount
        self.soft_target_tau = soft_target_tau
        self.enc_soft_target_tau = enc_soft_target_tau
        self.policy_mean_reg_weight = policy_mean_reg_weight
        self.policy_std_reg_weight = policy_std_reg_weight
        self.encoder_latent_lambda = encoder_latent_lambda

        self.ac_update_freq = ac_update_freq
        self.encoder_update_freq = encoder_update_freq

        self.train_alpha = train_alpha
        self.log_alpha = torch.tensor(np.log(alpha), requires_grad=train_alpha)
        self.alpha = self.log_alpha.detach().exp()
        assert "env" in kwargs.keys(), "env info should be taken into SAC alpha"
        self.target_entropy = target_entropy
        if target_entropy is None:
            self.target_entropy = -np.prod(kwargs["env"].action_space.shape)

        self.target_qf1 = qf1.copy()
        self.target_qf2 = qf2.copy()
        self.target_encoder = encoder.copy()

        self.eval_statistics = None

        policy_params = [
            {'params': self.policy.parameters(), 'lr': policy_lr},
        ]
        self.policy_optimizer = optimizer_class(
            policy_params,
            betas=(beta_1, 0.999),
        )
        qf_params = list(self.qf1.parameters()) + list(self.qf2.parameters())
        self.q_update_enc = kwargs.get('q_update_enc', False)
        if self.q_update_enc:
            print('Q function updates encoder.')
            qf_params += list(self.encoder.parameters())
        self.qf_optimizer = optimizer_class(
            qf_params,
            lr=qf_lr,
            betas=(beta_1, 0.999),
        )
        self.alpha_optimizer = optimizer_class(
            [self.log_alpha], lr=alpha_lr, betas=(alpha_beta, 0.999)
        )
        # if self.distill_qf:
        #     assert not self.latent_q, "If distill Q, can not use latent Q."
        #     print('Distill latent Q.')
        #     self.dis_qf1: nn.Module = self.distill_qf.copy()
        #     self.dis_qf2: nn.Module = self.distill_qf.copy()
        #     self.dis_qf_optimizer =  optimizer_class(
        #         list(self.dis_qf1.parameters())
        #         + list(self.dis_qf2.parameters()),
        #         lr=qf_lr,
        #         betas=(beta_1, 0.999),
        #     )
        #     self.dis_iters = kwargs.get('dis_iters', 1000)
        #     self.reset_dis_qf = kwargs.get('reset_dis_qf', False)

        self.total_train_step = 0

    def encode(self, x: torch.Tensor, detach: bool = False, target: bool = False, **kwargs) -> torch.Tensor:
        if target:
            with torch.no_grad():
                z_out = self.target_encoder(x, **kwargs)
        else:
            z_out = self.encoder(x, **kwargs)
        if detach:
            z_out = z_out.detach()
        return z_out

    def train_ac(self, batch):
        # batch["env"] is raw data sampled from real environment
        # batch["model"] is encoded data generated by the model

        if "env" in batch:
            env_batch = batch["env"]
            env_rewards = self.reward_scale * env_batch["rewards"]
            env_terminals = env_batch["terminals"]
            env_obs = self.encoder(env_batch["obserations"])
            env_actions = batch["actions"]
            env_next_obs = self.encoder(batch["next_observations"])
            env_target_next_obs = self.target_encoder(batch["next_observations"])
            batch = batch["model"]
        rewards = self.reward_scale * batch["rewards"]
        terminals = batch["terminals"]
        obs = batch["observations"]
        actions = batch["actions"]
        next_obs = batch["next_observations"]

        if "env" in batch:
            rewards = torch.cat([rewards, env_rewards], dim=0)
            terminals = torch.cat([terminals, env_terminals], dim=0)
            obs = torch.cat([obs, env_obs], dim=0)
            actions = torch.cat([actions, env_actions], dim=0)
            next_obs = torch.cat([next_obs, env_next_obs], dim=0)
            target_next_obs = torch.cat([next_obs, env_target_next_obs], dim=0)
        else:
            enc_obs = self.encode(obs)
            enc_next_obs = self.encode(next_obs)
            target_cur_obs = self.encode(obs, target=True)
            target_next_obs = self.encode(next_obs, target=True)
            if self.latent_q:
                obs = self.encode(obs)
                next_obs = self.encode(next_obs, target=True)

        """
        Model Regularization
        """
        if self.enc_model:
            if self.enc_no_target:
                enc_next_obs_ = enc_next_obs.detach()
            else:
                enc_next_obs_ = target_next_obs
            p_model_pred = self.enc_model(enc_obs, actions)
            p_model_next_s = p_model_pred[:, :-1]
            p_model_r = p_model_pred[:, -1]
            p_model_loss = torch.mean((p_model_next_s - enc_next_obs_) ** 2) \
                         + torch.mean((p_model_r - rewards) ** 2)
            if self.latent_q and not self.enc_no_target:
                enc_target_loss = torch.mean((enc_obs - target_cur_obs) ** 2)
            else:
                enc_target_loss = 0.0

        """
        QF Loss
        """
        self.qf_optimizer.zero_grad()
        q1_pred = self.qf1(obs, actions)
        q2_pred = self.qf2(obs, actions)

        # Make sure policy accounts for squashing functions like tanh correctly!
        next_policy_outputs = self.policy(
            next_obs.detach(), use_feature=self.latent_q, return_log_prob=True
        )
        # in this part, we only need new_actions and log_pi with no grad
        (
            next_new_actions,
            next_policy_mean,
            next_policy_log_std,
            next_log_pi,
        ) = next_policy_outputs[:4]
        with torch.no_grad():
            target_qf1_values = self.target_qf1(
                next_obs, next_new_actions
            )  # do not need grad || it's the shared part of two calculation
            target_qf2_values = self.target_qf2(
                next_obs, next_new_actions
            )  # do not need grad || it's the shared part of two calculation
            min_target_value = torch.min(target_qf1_values, target_qf2_values)
            q_target = rewards + (1.0 - terminals) * self.discount * (
                min_target_value - self.alpha * next_log_pi
            )  ## original implementation has detach

        qf1_loss = 0.5 * torch.mean((q1_pred - q_target) ** 2)
        qf2_loss = 0.5 * torch.mean((q2_pred - q_target) ** 2)
        qf_loss = qf1_loss + qf2_loss
        if self.q_update_enc and self.enc_model:
            qf_loss += self.enc_model_reg * p_model_loss \
                + self.enc_target_reg * enc_target_loss

        qf_loss.backward(retain_graph=True)
        self.qf_optimizer.step()

        """
        Double Space Q
        """
        # if self.double_q:
        #     enc_obs_ = enc_obs.detach()
        #     enc_next_obs_ = enc_next_obs.detach()

        #     for _ in range(self.lat_update_iter):
        #         lat_q1_pred = self.lat_qf1(enc_obs_, actions)
        #         lat_q2_pred = self.lat_qf2(enc_obs_, actions)
        #         enc_next_policy_outputs= self.policy(enc_next_obs_, use_feature=True, return_log_prob=True)
        #         (
        #             enc_next_new_actions,
        #             enc_next_policy_mean,
        #             enc_next_policy_log_std,
        #             enc_next_log_pi,
        #         ) = enc_next_policy_outputs[:4]

        #         with torch.no_grad():
        #             if self.lat_update_mode == 0:
        #                 target_lat_qf1_values = self.qf1(target_next_obs, enc_next_new_actions)
        #                 target_lat_qf2_values = self.qf2(target_next_obs, enc_next_new_actions)
        #             elif self.lat_update_mode == 1:
        #                 target_lat_qf1_values = self.target_qf1(target_next_obs, enc_next_new_actions)
        #                 target_lat_qf2_values = self.target_qf2(target_next_obs, enc_next_new_actions)
        #             min_lat_target_values = torch.min(target_lat_qf1_values, target_lat_qf2_values)
        #             lat_q_target = rewards + (1.0 - terminals) * self.discount * (min_lat_target_values - self.alpha * enc_next_log_pi)

        #         lat_qf1_loss = 0.5 * torch.mean((lat_q1_pred - lat_q_target) ** 2)
        #         lat_qf2_loss = 0.5 * torch.mean((lat_q2_pred - lat_q_target) ** 2)
        #         lat_qf_loss = lat_qf1_loss + lat_qf2_loss

        #         self.lat_qf_optimizer.zero_grad()
        #         lat_qf_loss.backward(retain_graph=True)
        #         self.lat_qf_optimizer.step()


        # """
        # Distill latent Q
        # """
        # if self.distill_qf:
        #     # reset latent Q
        #     if self.reset_dis_qf:
        #         self.dis_qf1.reset_parameters()
        #         self.dis_qf2.reset_parameters()
        #     enc_obs_ = enc_obs.detach()
        #     for _ in range(self.dis_iters):
        #         qf1_pred = self.qf1(obs, actions).detach()
        #         qf2_pred = self.qf2(obs, actions).detach()
        #         dis_qf1_pred = self.dis_qf1(enc_obs_, actions)
        #         dis_qf2_pred = self.dis_qf2(enc_obs_, actions)
        #         dis_loss = 0.5 * torch.mean((dis_qf1_pred - qf1_pred) ** 2) \
        #                  + 0.5 * torch.mean((dis_qf2_pred - qf2_pred) ** 2)
        #         self.dis_qf_optimizer.zero_grad()
        #         dis_loss.backward()
        #         self.dis_qf_optimizer.step()

        """
        Policy Loss
        """
        if self.q_update_enc:
            policy_outputs = self.policy(
                obs.detach(), use_feature=self.latent_q, return_log_prob=True
            )  # policy update the encoder
        else:
            policy_outputs = self.policy(
                obs, use_feature=self.latent_q, return_log_prob=True
            )  # policy update the encoder
        
        new_actions, policy_mean, policy_log_std, log_pi = policy_outputs[:4]
        # if self.distill_qf:
        #     q1_new_acts = self.dis_qf1(enc_obs, new_actions)
        #     q2_new_acts = self.dis_qf2(enc_obs, new_actions)
        # elif self.double_q:
        #     q1_new_acts = self.lat_qf1(enc_obs.detach(), new_actions)
        #     q2_new_acts = self.lat_qf2(enc_obs.detach(), new_actions)
        # else:
        q1_new_acts = self.qf1(obs.detach(), new_actions)
        q2_new_acts = self.qf2(obs.detach(), new_actions)
        q_new_actions = torch.min(q1_new_acts, q2_new_acts)

        # model regularization

        self.policy_optimizer.zero_grad()
        policy_loss = torch.mean(self.alpha * log_pi - q_new_actions)  ##
        mean_reg_loss = self.policy_mean_reg_weight * (policy_mean ** 2).mean()
        std_reg_loss = self.policy_std_reg_weight * (policy_log_std ** 2).mean()
        policy_reg_loss = mean_reg_loss + std_reg_loss
        policy_loss = policy_loss + policy_reg_loss

        if not self.q_update_enc and self.enc_model:
            policy_loss += self.enc_model_reg * p_model_loss \
                + self.enc_target_reg * enc_target_loss

        # if self.enc_model:
        #     # if self.total_train_step % self.enc_model_freq == 0:
        #     #     model_pred = self.enc_model(target_cur_obs, actions)
        #     #     model_next_s = model_pred[:, :-1]
        #     #     model_r = model_pred[:, -1]
        #     #     model_loss = torch.mean((model_next_s - enc_next_obs_) ** 2) + torch.mean((model_r - rewards) ** 2)
        #     #     self.enc_model_optimizer.zero_grad()
        #     #     model_loss.backward()
        #     #     self.enc_model_optimizer.step()

        #     p_model_pred = self.enc_model(enc_obs, actions)
        #     p_model_next_s = p_model_pred[:, :-1]
        #     p_model_r = p_model_pred[:, -1]
        #     p_model_loss = torch.mean((p_model_next_s - enc_next_obs_) ** 2) \
        #                  + torch.mean((p_model_r - rewards) ** 2)
        #     policy_loss = policy_loss + self.enc_model_reg * p_model_loss
            
        #         policy_loss = policy_loss + self.enc_target_reg * enc_target_loss
        #     # if self.double_q:
        #     #     enc_lat_q1_pred = self.lat_qf1(enc_obs, actions)
        #     #     enc_lat_q2_pred = self.lat_qf2(enc_obs, actions)
        #     #     target_enc_lat_q1_pred = self.lat_qf1(target_cur_obs, actions)
        #     #     target_enc_lat_q2_pred = self.lat_qf2(target_cur_obs, actions)
        #     #     lat_q_enc_loss = 0.5 * torch.mean((enc_lat_q1_pred - target_enc_lat_q1_pred) ** 2) \
        #     #                    + 0.5 * torch.mean((enc_lat_q2_pred - target_enc_lat_q2_pred) ** 2)
        #     #     policy_loss = policy_loss + self.lat_q_enc_reg * lat_q_enc_loss
        policy_loss.backward()
        self.policy_optimizer.step()

        """
        Update alpha
        """
        if self.train_alpha:
            log_prob = log_pi.detach() + self.target_entropy
            alpha_loss = -(self.log_alpha * log_prob).mean()
            self.alpha_optimizer.zero_grad()
            alpha_loss.backward()
            self.alpha_optimizer.step()
            self.alpha = self.log_alpha.detach().exp()

        """
        Update networks
        """

        self._update_target_network()

        """
        Save some statistics for eval
        """
        if self.eval_statistics is None:
            """
            Eval should set this to None.
            This way, these statistics are only computed for one batch.
            """
            self.eval_statistics = OrderedDict()
        self.eval_statistics["QF1 Loss"] = np.mean(ptu.get_numpy(qf1_loss))
        self.eval_statistics["QF2 Loss"] = np.mean(ptu.get_numpy(qf2_loss))
        if self.train_alpha:
            self.eval_statistics["Alpha Loss"] = np.mean(ptu.get_numpy(alpha_loss))
        self.eval_statistics["Policy Loss"] = np.mean(ptu.get_numpy(policy_loss))
        if self.enc_model:
            self.eval_statistics["Enc Model Loss"] = np.mean(ptu.get_numpy(p_model_loss))
        # if self.double_q:
        #     self.eval_statistics["Latent QF1 Loss"] = np.mean(ptu.get_numpy(lat_qf1_loss))
        #     self.eval_statistics["Latent QF2 Loss"] = np.mean(ptu.get_numpy(lat_qf2_loss))
        #     self.eval_statistics["Latent Q Encoder Reg Loss"] = np.mean(ptu.get_numpy(lat_q_enc_loss))
        #     self.eval_statistics.update(
        #         create_stats_ordered_dict(
        #             "Latent Q1 Predictions",
        #             ptu.get_numpy(lat_q1_pred),
        #         )
        #     )
        #     self.eval_statistics.update(
        #         create_stats_ordered_dict(
        #             "Latent Q2 Predictions",
        #             ptu.get_numpy(lat_q2_pred),
        #         )
        #     )
        self.eval_statistics.update(
            create_stats_ordered_dict(
                "Q1 Predictions",
                ptu.get_numpy(q1_pred),
            )
        )
        self.eval_statistics.update(
            create_stats_ordered_dict(
                "Q2 Predictions",
                ptu.get_numpy(q2_pred),
            )
        )
        # if self.distill_qf:
        #     self.eval_statistics.update(
        #         create_stats_ordered_dict(
        #             "Latent Q1 Predictions",
        #             ptu.get_numpy(dis_qf1_pred),
        #         )
        #     )
        #     self.eval_statistics.update(
        #         create_stats_ordered_dict(
        #             "Latent Q2 Predictions",
        #             ptu.get_numpy(dis_qf2_pred),
        #         )
        #     )
        #     self.eval_statistics['Distillation Loss'] = np.mean(ptu.get_numpy(dis_loss))
        self.eval_statistics.update(
            create_stats_ordered_dict(
                "Alpha",
                [ptu.get_numpy(self.alpha)],
            )
        )
        self.eval_statistics.update(
            create_stats_ordered_dict(
                "Log Pis",
                ptu.get_numpy(log_pi),
            )
        )
        self.eval_statistics.update(
            create_stats_ordered_dict(
                "Policy mu",
                ptu.get_numpy(policy_mean),
            )
        )
        self.eval_statistics.update(
            create_stats_ordered_dict(
                "Policy log std",
                ptu.get_numpy(policy_log_std),
            )
        )

    def train_model(
        self, 
        replay_buffer: EnvReplayBuffer, 
        batch_size: int = 256, 
        holdout_ratio: float = 0.2,
        max_epochs_since_update: int = 5,
        **kwargs
    ) -> None:
        batch = replay_buffer.get_all()
        rewards: np.array = self.reward_scale * batch["rewards"]
        obs: np.array = batch["observations"]
        actions: np.array = batch["actions"]
        next_obs: np.array = batch["next_observations"]
        
        self._epochs_since_update = 0
        self._snapshots = dict(epoch=-1, loss=1e12, params=self.enc_model.state_dict())

        num_holdout = int(obs.shape[0] * holdout_ratio)
        perm = np.random.permutation(obs.shape[0])
        ho_ids = perm[:num_holdout]
        ids = perm[num_holdout:]

        ho_obs = ptu.from_numpy(obs[ho_ids])
        ho_actions = ptu.from_numpy(actions[ho_ids])
        ho_rewards = ptu.from_numpy(rewards[ho_ids])
        ho_next_obs = ptu.from_numpy(next_obs[ho_ids])

        obs: np.array = obs[ids]
        actions: np.array = actions[ids]
        rewards: np.array = rewards[ids]
        next_obs: np.array = next_obs[ids]

        break_train = False
        for epoch in itertools.count():
            train_idx = np.random.permutation(obs.shape[0])
            for st in range(0, obs.shape[0], batch_size):
                idx = train_idx[st: min(st + batch_size, len(train_idx))]
                # if len(idx == 1):
                #     # discard the last instance to avoid error in BatchNorm
                #     break 
                tr_obs = ptu.from_numpy(obs[idx])
                tr_actions = ptu.from_numpy(actions[idx])
                tr_rewards = ptu.from_numpy(rewards[idx])
                tr_next_obs = ptu.from_numpy(next_obs[idx])

                with torch.no_grad():
                    enc_tr_obs = self.encode(tr_obs, detach=True)
                    enc_tr_next_obs = self.encode(tr_next_obs, target=True)

                pred = self.enc_model(enc_tr_obs, tr_actions)
                pred_next_obs = pred[:, :-1]
                pred_rewards = torch.unsqueeze(pred[:, -1], dim=1)
                obs_loss = F.mse_loss(pred_next_obs, enc_tr_next_obs) 
                r_loss = F.mse_loss(pred_rewards, tr_rewards)
                loss = obs_loss + r_loss
                # print(f'epoch {epoch} | st {st} | obs loss {obs_loss:.3f} | r loss {r_loss:.3f}')

                self.enc_model_optimizer.zero_grad()
                loss.backward()
                self.enc_model_optimizer.step()

            with torch.no_grad():
                enc_ho_obs = self.encode(ho_obs, detach=True)
                enc_ho_next_obs = self.encode(ho_next_obs, target=True)
                ho_pred = self.enc_model(enc_ho_obs, ho_actions)
                pred_ho_next_obs = ho_pred[:, :-1]
                pred_ho_rewards = torch.unsqueeze(ho_pred[:, -1], dim=1)
                # print('pred: ', pred_ho_next_obs)
                # print('tgt: ', enc_ho_next_obs)
                # print(torch.max(pred_ho_next_obs))
                # print(torch.max(enc_ho_next_obs, dim=0)[0])
                # print(torch.min(enc_ho_next_obs, dim=0)[0])
                # print(torch.std(enc_ho_next_obs, dim=0))
                ho_obs_loss = F.mse_loss(pred_ho_next_obs, enc_ho_next_obs) 
                ho_r_loss = F.mse_loss(pred_ho_rewards, ho_rewards)
                # print('ho obs loss: ', ho_obs_loss)
                # print('ho r loss: ', ho_r_loss)
                ho_loss = ho_obs_loss + ho_r_loss
                ho_loss = ho_loss.detach().cpu().numpy()
                # print(f'epoch {epoch} | obs loss {ho_obs_loss:.3f} | r loss {ho_r_loss:.3f}')

            updated = False
            cur_loss = ho_loss
            best_loss = self._snapshots['loss']
            # print(f'epoch {epoch} | cur loss {cur_loss:.3f} | best_loss {best_loss:.3f}')
            if (best_loss - cur_loss) / best_loss > 0.01:
                self._snapshots['epoch'] = epoch
                self._snapshots['loss'] = cur_loss
                self._snapshots['params'] = self.enc_model.state_dict()
                updated = True
            if updated:
                self._epochs_since_update = 0
            else:
                self._epochs_since_update += 1
            if self._epochs_since_update > max_epochs_since_update:
                break_train = True
            if break_train:
                break
        
        logger.log(f'Train Model | epoch: {epoch} | best loss: {self._snapshots["loss"]:.3f}')
        self.enc_model.load_state_dict(self._snapshots['params'])


    def train_step(self, batch: torch.Tensor):
        if self.total_train_step % self.ac_update_freq == 0:
            self.train_ac(batch)
        # if self.enc_model and self.total_train_step % self.enc_model_freq == 0:
        #     self.train_model(replay_buffer, **kwargs)
        self.total_train_step += 1

    @property
    def networks(self):
        nets = [
            self.policy,
            self.qf1,
            self.qf2,
            self.target_qf1,
            self.target_qf2,
            self.encoder,
            self.target_encoder,
        ]
        if self.distill_qf:
            nets.extend([self.distill_qf, self.dis_qf1, self.dis_qf2])
        if self.enc_model:
            nets.append(self.enc_model)
        if self.double_q:
            nets.extend([
                self.lat_qf1,
                self.lat_qf2,
                self.target_lat_qf1,
                self.target_lat_qf2
            ])
        return nets

    def _update_target_network(self):
        ptu.soft_update_from_to(self.qf1, self.target_qf1, self.soft_target_tau)
        ptu.soft_update_from_to(self.qf2, self.target_qf2, self.soft_target_tau)
        if self.double_q:
            ptu.soft_update_from_to(self.lat_qf1, self.target_lat_qf1, self.soft_target_tau)
            ptu.soft_update_from_to(self.lat_qf2, self.target_lat_qf2, self.soft_target_tau)
        ptu.soft_update_from_to(self.encoder, self.target_encoder, self.enc_soft_target_tau)

    def get_snapshot(self):
        return dict(
            qf1=self.qf1,
            qf2=self.qf2,
            policy=self.policy,
            target_qf1=self.target_qf1,
            target_qf2=self.target_qf2,
            log_alpha=self.log_alpha,
            policy_optimizer=self.policy_optimizer,
            qf_optimizer=self.qf_optimizer,
            alpha_optimizer=self.alpha_optimizer,
        )

    def load_snapshot(self, snapshot):
        self.qf1 = snapshot["qf1"]
        self.qf2 = snapshot["qf2"]
        self.policy = snapshot["policy"]
        self.target_qf1 = snapshot["target_qf1"]
        self.target_qf2 = snapshot["target_qf2"]
        self.log_alpha = snapshot["log_alpha"]
        self.policy_optimizer = snapshot["policy_optimizer"]
        self.qf_optimizer = snapshot["qf_optimizer"]
        self.alpha_optimizer = snapshot["alpha_optimizer"]

    def get_eval_statistics(self):
        return self.eval_statistics

    def end_epoch(self):
        self.eval_statistics = None

    def to(self, device):
        self.log_alpha = self.log_alpha.to(device)
        for net in self.networks:
            net.to(device)
