text: # Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
  #
  # This source code is licensed under the MIT license found in the
  # LICENSE file in the root directory of this source tree.
  import os
  from typing import Optional, Tuple, cast
  
  import gym
  import hydra.utils
  import numpy as np
  import omegaconf
  import torch
  
  import mbrl.constants
  import mbrl.models
  import mbrl.planning
  import mbrl.third_party.pytorch_sac as pytorch_sac
  import mbrl.types
  import mbrl.util
  import mbrl.util.common
  import mbrl.util.math
  from mbrl.planning.sac_wrapper import SACAgent
  
  MBPO_LOG_FORMAT = mbrl.constants.EVAL_LOG_FORMAT + [
      ("epoch", "E", "int"),
      ("rollout_length", "RL", "int"),
  ]
  
  def rollout_model_and_populate_sac_buffer(
      model_env: mbrl.models.ModelEnv,
      replay_buffer: mbrl.util.ReplayBuffer,
      agent: SACAgent,
      sac_buffer: pytorch_sac.ReplayBuffer,
      sac_samples_action: bool,
      rollout_horizon: int,
      batch_size: int,
  ):
  
      batch = replay_buffer.sample(batch_size)
      initial_obs, *_ = cast(mbrl.types.TransitionBatch, batch).astuple()
      model_state = model_env.reset(
          initial_obs_batch=cast(np.ndarray, initial_obs),
          return_as_np=True,
      )
      accum_dones = np.zeros(initial_obs.shape[0], dtype=bool)
      obs = initial_obs
      for i in range(rollout_horizon):
          action = agent.act(obs, sample=sac_samples_action, batched=True)
          pred_next_obs, pred_rewards, pred_dones, model_state = model_env.step(
              action, model_state, sample=True
          )
          sac_buffer.add_batch(
              obs[~accum_dones],
              action[~accum_dones],
              pred_rewards[~accum_dones],
              pred_next_obs[~accum_dones],
              pred_dones[~accum_dones],
              pred_dones[~accum_dones],
          )
          obs = pred_next_obs
          accum_dones |= pred_dones.squeeze()
  
  
  def double_masked_rollout_model_and_populate_sac_buffer(
      model_env: mbrl.models.ModelEnv,
      replay_buffer: mbrl.util.ReplayBuffer,
      carry_agent: SACAgent,
      carry_sac_buffer: pytorch_sac.ReplayBuffer,
      sup_sac_buffer: pytorch_sac.ReplayBuffer,
      sac_samples_action: bool,
      rollout_horizon: int,
      batch_size: int,
      penalty: float = 1e-3,
      bonus: float = 1e-3,
  ):
  
      batch = replay_buffer.sample(batch_size)
      initial_obs, *_ = cast(mbrl.types.TransitionBatch, batch).astuple()
      model_state = model_env.reset(
          initial_obs_batch=cast(np.ndarray, initial_obs),
          return_as_np=True,
      )
      accum_dones = np.zeros(initial_obs.shape[0], dtype=bool)
      
      accum_ovr_score = np.zeros(initial_obs.shape[0])
      rollout_width = batch_size
      obs = initial_obs
  
      previous_cache = None
  
      for i in range(rollout_horizon):
  
          carry_mask_rate = (rollout_horizon - i) / (2 * (rollout_horizon + 1)) if rollout_horizon>1 else 0.5
          # carry_mask_rate = 0.5
  
          action = carry_agent.act(obs, sample=sac_samples_action, batched=True)
          pred_next_obs, pred_rewards, pred_dones, model_state, ovr_score = model_env.step_w_ovr(
              action, model_state, sample=True
          )
  
          # rank by accumu ovr score, but bonus/penalt by single-step ovr
          ovr_score = ovr_score.cpu().numpy()
          accum_ovr_score += ovr_score
          ovr_ranked_ids = np.argsort(accum_ovr_score)
          ovr_score = np.reshape(ovr_score, [batch_size, 1])
          
          carry_pred_rewards = pred_rewards - ovr_score * penalty
          carry_masked_ids = ovr_ranked_ids[:int(carry_mask_rate*rollout_width)]
          carry_mask = np.zeros(initial_obs.shape[0], dtype=bool)
          carry_mask[carry_masked_ids] = True
  
  
          carry_sac_buffer.add_batch(
              obs[~accum_dones & carry_mask],
              action[~accum_dones & carry_mask],
              carry_pred_rewards[~accum_dones & carry_mask],
              pred_next_obs[~accum_dones & carry_mask],
              pred_dones[~accum_dones & carry_mask],
              pred_dones[~accum_dones & carry_mask],
          )
  
          current_cache = (
              obs,
              action,
              pred_rewards,
              pred_next_obs,
              pred_dones,
          )
  
          if previous_cache != None:
              p_obs, p_action, p_pred_rewards, p_pred_next_obs, p_pred_dones = previous_cache
              p_sup_pred_rewards = p_pred_rewards + ovr_score * bonus
              sup_sac_buffer.add_batch(
                  p_obs[~old_accum_dones],
                  p_action[~old_accum_dones],
                  p_sup_pred_rewards[~old_accum_dones],
                  p_pred_next_obs[~old_accum_dones],
                  p_pred_dones[~old_accum_dones],
                  p_pred_dones[~old_accum_dones],
              )
              print(f'RS {i} | add {(~old_accum_dones).sum()} rollouts to sup pool')
  
          obs = pred_next_obs
          previous_cache = current_cache
          print(f'RS {i} | CMS {carry_mask_rate} | RW {rollout_width} | add {(~accum_dones & carry_mask).sum()} rollouts to carry pool |/ mean accum ovr {accum_ovr_score[ovr_ranked_ids[int(carry_mask_rate*rollout_width)]]}') 
          old_accum_dones = accum_dones
          accum_dones |= pred_dones.squeeze()
          rollout_width = (~accum_dones).sum()
      
      action = carry_agent.act(obs, sample=sac_samples_action, batched=True)
      pred_next_obs, pred_rewards, pred_dones, model_state, ovr_score = model_env.step_w_ovr(
          action, model_state, sample=True
      )
  
      ovr_score = ovr_score.cpu().numpy()
      # ovr_ranked_ids = np.argsort(ovr_score)
      ovr_score = np.reshape(ovr_score, [batch_size, 1])
  
      p_obs, p_action, p_pred_rewards, p_pred_next_obs, p_pred_dones = previous_cache
      p_sup_pred_rewards = p_pred_rewards + ovr_score * bonus
      sup_sac_buffer.add_batch(
          p_obs[~old_accum_dones],
          p_action[~old_accum_dones],
          p_sup_pred_rewards[~old_accum_dones],
          p_pred_next_obs[~old_accum_dones],
          p_pred_dones[~old_accum_dones],
          p_pred_dones[~old_accum_dones],
      )
      print(f'RS {i} | add {(~old_accum_dones).sum()} rollouts to sup pool')
  
  
  
  
          
  
  
  def evaluate(
      env: gym.Env,
      agent: pytorch_sac.Agent,
      num_episodes: int,
      video_recorder: pytorch_sac.VideoRecorder,
  ) -> float:
      avg_episode_reward = 0
      for episode in range(num_episodes):
          obs = env.reset()
          video_recorder.init(enabled=(episode == 0))
          done = False
          episode_reward = 0
          while not done:
              action = agent.act(obs)
              obs, reward, done, _ = env.step(action)
              video_recorder.record(env)
              episode_reward += reward
          avg_episode_reward += episode_reward
      return avg_episode_reward / num_episodes
  
  
  def maybe_replace_sac_buffer(
      sac_buffer: Optional[pytorch_sac.ReplayBuffer],
      new_capacity: int,
      obs_shape: Tuple[int],
      act_shape: Tuple[int],
      device: torch.device,
  ) -> pytorch_sac.ReplayBuffer:
      if sac_buffer is None or new_capacity != sac_buffer.capacity:
          new_buffer = pytorch_sac.ReplayBuffer(
              obs_shape, act_shape, new_capacity, device
          )
          if sac_buffer is None:
              return new_buffer
          n = len(sac_buffer)
          new_buffer.add_batch(
              sac_buffer.obses[:n],
              sac_buffer.actions[:n],
              sac_buffer.rewards[:n],
              sac_buffer.next_obses[:n],
              np.logical_not(sac_buffer.not_dones[:n]),
              np.logical_not(sac_buffer.not_dones_no_max[:n]),
          )
          return new_buffer
      return sac_buffer
  
  
  def train(
      env: gym.Env,
      test_env: gym.Env,
      termination_fn: mbrl.types.TermFnType,
      cfg: omegaconf.DictConfig,
      silent: bool = False,
      work_dir: Optional[str] = None,
  ) -> np.float32:
      # ------------------- Initialization -------------------
      debug_mode = cfg.get("debug_mode", False)
  
      obs_shape = env.observation_space.shape
      act_shape = env.action_space.shape
  
      mbrl.planning.complete_agent_cfg(env, cfg.algorithm.agent)
      carry_agent = hydra.utils.instantiate(cfg.algorithm.agent)
      sup_agent = hydra.utils.instantiate(cfg.algorithm.agent)
  
      work_dir = work_dir or os.getcwd()
      # enable_back_compatible to use pytorch_sac agent
      logger = mbrl.util.Logger(work_dir, enable_back_compatible=True)
      logger.register_group(
          mbrl.constants.RESULTS_LOG_NAME,
          MBPO_LOG_FORMAT,
          color="green",
          dump_frequency=1,
      )
      save_video = cfg.get("save_video", False)
      video_recorder = pytorch_sac.VideoRecorder(work_dir if save_video else None)
  
      rng = np.random.default_rng(seed=cfg.seed)
      torch_generator = torch.Generator(device=cfg.device)
      if cfg.seed is not None:
          torch_generator.manual_seed(cfg.seed)
  
      # -------------- Create initial overrides. dataset --------------
      dynamics_model = mbrl.util.common.create_one_dim_tr_model(cfg, obs_shape, act_shape)
      use_double_dtype = cfg.algorithm.get("normalize_double_precision", False)
      dtype = np.double if use_double_dtype else np.float32
      replay_buffer = mbrl.util.common.create_replay_buffer(
          cfg,
          obs_shape,
          act_shape,
          rng=rng,
          obs_type=dtype,
          action_type=dtype,
          reward_type=dtype,
      )
      random_explore = cfg.algorithm.random_initial_explore
      mbrl.util.common.rollout_agent_trajectories(
          env,
          cfg.algorithm.initial_exploration_steps,
          mbrl.planning.RandomAgent(env) if random_explore else sup_agent,
          {} if random_explore else {"sample": True, "batched": False},
          replay_buffer=replay_buffer,
      )
  
      # ---------------------------------------------------------
      # --------------------- Training Loop ---------------------
      rollout_batch_size = (
          cfg.overrides.effective_model_rollouts_per_step * cfg.algorithm.freq_train_model
      )
      init_rollout_batch_size = rollout_batch_size
  
      trains_per_epoch = int(
          np.ceil(cfg.overrides.epoch_length / cfg.overrides.freq_train_model)
      )
      updates_made = 0
      sup_updates_made =0
      env_steps = 0
      model_env = mbrl.models.ModelEnv(
          env, dynamics_model, termination_fn, None, generator=torch_generator
      )
      model_trainer = mbrl.models.ModelTrainer(
          dynamics_model,
          optim_lr=cfg.overrides.model_lr,
          weight_decay=cfg.overrides.model_wd,
          logger=None if silent else logger,
      )
      best_eval_reward = -np.inf
      epoch = 0
      carry_sac_buffer = None
      sup_sac_buffer = None
      num_sac_updates_per_step = cfg.overrides.num_sac_updates_per_step
  
      print(f'[INFO] current step = {env_steps} | rollout batch size = {rollout_batch_size} | sac updates {num_sac_updates_per_step} per step')
  
      
  
      # --- it takes two: --- 
  
      print(f'[INFO] training stage 1 of {cfg.overrides.s1_num_steps} steps begin...')
  
      while env_steps < cfg.overrides.s1_num_steps :
          rollout_length = int(
              mbrl.util.math.truncated_linear(
                  *(cfg.overrides.rollout_schedule + [epoch + 1])
              )
          )
          
          carry_sac_buffer_capacity = rollout_batch_size * trains_per_epoch
          carry_sac_buffer_capacity *= cfg.overrides.num_epochs_to_retain_sac_buffer
  
  
          sup_sac_buffer_capacity = rollout_length * rollout_batch_size * trains_per_epoch
          sup_sac_buffer_capacity *= cfg.overrides.num_epochs_to_retain_sac_buffer
  
          
          carry_sac_buffer = maybe_replace_sac_buffer(
              carry_sac_buffer,
              carry_sac_buffer_capacity,
              obs_shape,
              act_shape,
              torch.device(cfg.device),
          )
          sup_sac_buffer = maybe_replace_sac_buffer(
              sup_sac_buffer,
              sup_sac_buffer_capacity,
              obs_shape,
              act_shape,
              torch.device(cfg.device),
          )
          obs, done = None, False
  
          for steps_epoch in range(cfg.overrides.epoch_length):
              if steps_epoch == 0 or done:
                  obs, done = env.reset(), False
              # --- Doing env step and adding to model dataset ---
              next_obs, reward, done, _ = mbrl.util.common.step_env_and_add_to_buffer(
                  env, obs, sup_agent, {}, replay_buffer
              )
  
              # --------------- Model Training -----------------
              if (env_steps + 1) % cfg.overrides.freq_train_model == 0:
                  mbrl.util.common.train_model_and_save_model_and_data(
                      dynamics_model,
                      model_trainer,
                      cfg.overrides,
                      replay_buffer,
                      work_dir=work_dir,
                  )
  
                  # --------- Rollout new model and store imagined trajectories --------
                  # Batch all rollouts for the next freq_train_model steps together
                  double_masked_rollout_model_and_populate_sac_buffer(
                      model_env,
                      replay_buffer,
                      carry_agent,
                      carry_sac_buffer,
                      sup_sac_buffer,
                      cfg.algorithm.sac_samples_action,
                      rollout_length,
                      rollout_batch_size,
                  )
  
                  if debug_mode:
                      print(
                          f"Epoch: {epoch}. "
                          f"carry SAC buffer size: {len(carry_sac_buffer)}. "
                          f"sup SAC buffer size: {len(sup_sac_buffer)}. "
                          f"Rollout length: {rollout_length}. "
                          f"Steps: {env_steps}"
                      )
  
              # --------------- Agent Training -----------------
              for _ in range(num_sac_updates_per_step):
                  if (env_steps + 1) % cfg.overrides.sac_updates_every_steps != 0 or len(
                      sup_sac_buffer
                  ) < rollout_batch_size:
                      break  # only update every once in a while
                  sup_agent.update(sup_sac_buffer, logger, sup_updates_made)
                  sup_updates_made += 1
                  # if not silent and updates_made % cfg.log_frequency_agent == 0:
                  #     logger.dump(updates_made, save=True)
              
              for _ in range(num_sac_updates_per_step):
                  if (env_steps + 1) % (2 * cfg.overrides.sac_updates_every_steps) != 0 or len(
                      carry_sac_buffer
                  ) < rollout_batch_size:
                      break  # only update every once in a while
                  carry_agent.update(carry_sac_buffer, logger, updates_made)
                  updates_made += 1
                  if not silent and updates_made % cfg.log_frequency_agent == 0:
                      logger.dump(updates_made, save=True)
  
  
              # ------ Epoch ended (evaluate and save model) ------
              if (env_steps + 1) % cfg.overrides.epoch_length == 0:
                  avg_reward = evaluate(
                      test_env, carry_agent, cfg.algorithm.num_eval_episodes, video_recorder
                  )
                  logger.log_data(
                      mbrl.constants.RESULTS_LOG_NAME,
                      {
                          "epoch": epoch,
                          "env_step": env_steps,
                          "episode_reward": avg_reward,
                          "rollout_length": rollout_length,
                      },
                  )
                  if avg_reward > best_eval_reward:
                      video_recorder.save(f"{epoch}.mp4")
                      best_eval_reward = avg_reward
                      torch.save(
                          carry_agent.critic.state_dict(), os.path.join(work_dir, "critic.pth")
                      )
                      torch.save(
                          sup_agent.actor.state_dict(), os.path.join(work_dir, "actor.pth")
                      )
                  epoch += 1
  
              env_steps += 1
              obs = next_obs
  
      # --- It Takes One...
      print(f'[INFO] training stage 2 of {cfg.overrides.num_steps - cfg.overrides.s1_num_steps} steps begin... | current step = {env_steps}')
  
      sac_buffer = carry_sac_buffer
      while env_steps < cfg.overrides.num_steps:
          # adjust sac updates frequency and rollout batch size
          if (env_steps + 1 - cfg.overrides.s1_num_steps) == cfg.overrides.s1_num_steps:
              rollout_batch_size += init_rollout_batch_size // 2
              num_sac_updates_per_step += cfg.num_sac_updates_per_step // 2
              print(f'[INFO] current step = {env_steps} | rollout batch size = {rollout_batch_size} | sac updates {num_sac_updates_per_step} per step')
  
          if (env_steps + 1 - cfg.overrides.s1_num_steps) == 3*cfg.overrides.s1_num_steps:
              rollout_batch_size += init_rollout_batch_size // 2
              num_sac_updates_per_step += cfg.num_sac_updates_per_step // 2
              print(f'[INFO] current step = {env_steps} | rollout batch size = {rollout_batch_size} | sac updates {num_sac_updates_per_step} per step')
  
          rollout_length = int(
              mbrl.util.math.truncated_linear(
                  *(cfg.overrides.rollout_schedule + [epoch + 1])
              )
          )
          sac_buffer_capacity = rollout_length * rollout_batch_size * trains_per_epoch
          sac_buffer_capacity *= cfg.overrides.num_epochs_to_retain_sac_buffer
          
          sac_buffer = maybe_replace_sac_buffer(
              sac_buffer,
              sac_buffer_capacity,
              obs_shape,
              act_shape,
              torch.device(cfg.device),
          )
          obs, done = None, False
          for steps_epoch in range(cfg.overrides.epoch_length):
              if steps_epoch == 0 or done:
                  obs, done = env.reset(), False
              # --- Doing env step and adding to model dataset ---
              next_obs, reward, done, _ = mbrl.util.common.step_env_and_add_to_buffer(
                  env, obs, carry_agent, {}, replay_buffer
              )
  
              # --------------- Model Training -----------------
              if (env_steps + 1) % cfg.overrides.freq_train_model == 0:
                  mbrl.util.common.train_model_and_save_model_and_data(
                      dynamics_model,
                      model_trainer,
                      cfg.overrides,
                      replay_buffer,
                      work_dir=work_dir,
                  )
  
                  # --------- Rollout new model and store imagined trajectories --------
                  # Batch all rollouts for the next freq_train_model steps together
                  rollout_model_and_populate_sac_buffer(
                      model_env,
                      replay_buffer,
                      carry_agent,
                      sac_buffer,
                      cfg.algorithm.sac_samples_action,
                      rollout_length,
                      rollout_batch_size,
                  )
  
                  if debug_mode:
                      print(
                          f"Epoch: {epoch}. "
                          f"SAC buffer size: {len(sac_buffer)}. "
                          f"Rollout length: {rollout_length}. "
                          f"Steps: {env_steps}"
                      )
  
              # --------------- Agent Training -----------------
              for _ in range(num_sac_updates_per_step):
                  if (env_steps + 1) % cfg.overrides.sac_updates_every_steps != 0 or len(
                      sac_buffer
                  ) < rollout_batch_size:
                      break  # only update every once in a while
                  carry_agent.update(sac_buffer, logger, updates_made)
                  updates_made += 1
                  if not silent and updates_made % cfg.log_frequency_agent == 0:
                      logger.dump(updates_made, save=True)
  
              # ------ Epoch ended (evaluate and save model) ------
              if (env_steps + 1) % cfg.overrides.epoch_length == 0:
                  avg_reward = evaluate(
                      test_env, carry_agent, cfg.algorithm.num_eval_episodes, video_recorder
                  )
                  logger.log_data(
                      mbrl.constants.RESULTS_LOG_NAME,
                      {
                          "epoch": epoch,
                          "env_step": env_steps,
                          "episode_reward": avg_reward,
                          "rollout_length": rollout_length,
                      },
                  )
                  if avg_reward > best_eval_reward:
                      video_recorder.save(f"{epoch}.mp4")
                      best_eval_reward = avg_reward
                      torch.save(
                          carry_agent.critic.state_dict(), os.path.join(work_dir, "critic.pth")
                      )
                      torch.save(
                          carry_agent.actor.state_dict(), os.path.join(work_dir, "actor.pth")
                      )
                  epoch += 1
  
              env_steps += 1
      return np.float32(best_eval_reward)
  