<gh_stars>10-100
import os
import functools
import tools
import tensorflow as tf
import numpy as np
import pickle
import pathlib
import off_policy
from dreamer import Dreamer, setup_dreamer, create_envs, count_steps, make_dataset, parse_dreamer_args


class GCDreamer(Dreamer):
  def __init__(self, config, logger, dataset):
    if config.offpolicy_opt:
      self._off_policy_handler = off_policy.GCOffPolicyOpt(config)
    super().__init__(config, logger, dataset)
    self._should_expl_ep = tools.EveryNCalls(config.expl_every_ep)
    self.skill_to_use = tf.zeros([0], dtype=tf.float16)

  def get_one_time_skill(self):
    skill = self.skill_to_use
    self.skill_to_use = tf.zeros([0], dtype=tf.float16)
    return skill

  def _policy(self, obs, state, training, reset):
    # Choose the goal for the next episode
    # TODO this would be much more elegant if it was implemented in the training loop (can simulate a single episode)
    if state is None:
      if training and self._config.training_goals == 'batch':
        # TODO this modifies the underlying dict, is that fine?
        obs['image_goal'], obs['goal'] = self.sample_replay_goal(obs)
      obs['skill'] = self.get_one_time_skill()
      # The world model adds the observation to the state in this case

    if reset.any() and state is not None:
      # Replace the goal in the agent state at new episode
      # The actor always takes goal from the state, not the observation
      if training and self._config.training_goals == 'batch':
        state[0]['image_goal'], state[0]['goal'] = self.sample_replay_goal(obs)
      else:
        state[0]['image_goal'] = tf.cast(obs['image_goal'], self._float) / 255.0 - 0.5
      state[0]['skill'] = self.get_one_time_skill()

      # Toggle exploration
      self._should_expl_ep()

    # TODO double check everything
    #return super()._policy(obs, state, training, reset, should_expl=self._should_expl_ep.value)
    #if not training:
    return super()._policy(obs, state, training, reset, should_expl=self._should_expl_ep.value)

  def sample_replay_goal(self, obs):
    """ Sample goals from replay buffer """
    assert self._config.gc_input != 'state'
    random_batch = next(self._dataset)
    random_batch = self._wm.preprocess(random_batch)
    
    images = random_batch['image']
    states = random_batch['state']
    if self._config.labelled_env_multiplexing:
      assert obs['env_idx'].shape[0] == 1
      env_ids = random_batch['env_idx'][:, 0]
      if tf.reduce_any(env_ids == obs['env_idx']):
        ids = np.nonzero(env_ids == obs['env_idx'])[0]
        images = tf.gather(images, ids)
        states = tf.gather(states, ids)
    
    random_goals = tf.reshape(images, (-1,) + tuple(images.shape[2:]))
    random_goal_states = tf.reshape(states, (-1,) + tuple(states.shape[2:]))
    # random_goals = tf.random.shuffle(random_goals)
    return random_goals[:obs['image_goal'].shape[0]], random_goal_states[:obs['image_goal'].shape[0]]


def process_eps_data(eps_data):
  keys = eps_data[0].keys()
  new_data = {}
  for key in keys:
    new_data[key] = np.array([eps_data[i][key] for i in range(len(eps_data))]).squeeze()
  return new_data

def main(logdir, config):
  logdir, logger = setup_dreamer(config, logdir)
  eval_envs, eval_eps, train_envs, train_eps, acts = create_envs(config, logger)

  prefill = max(0, config.prefill - count_steps(config.traindir))
  print(f'Prefill dataset ({prefill} steps).')
  random_agent = lambda o, d, s: ([acts.sample() for _ in d], s)
  tools.simulate(random_agent, train_envs, prefill)
  if count_steps(config.evaldir) == 0:
    tools.simulate(random_agent, eval_envs, episodes=1)
  logger.step = config.action_repeat * count_steps(config.traindir)

  print('Simulate agent.')
  train_dataset = make_dataset(train_eps, config)
  eval_dataset = iter(make_dataset(eval_eps, config))
  agent = GCDreamer(config, logger, train_dataset)
  if (logdir / 'variables.pkl').exists():
    agent.load(logdir / 'variables.pkl')
    agent._should_pretrain._once = False

  pathlib.Path(logdir / "distance_func_logs_trained_model").mkdir(parents=True, exist_ok = True)

  state = None
  assert len(eval_envs) == 1
  while agent._step.numpy().item() < config.steps:
    logger.write()
    print('Start gc evaluation.')
    executions = []
    goals = []
    #rews_across_goals = []
    num_goals = min(100, len(eval_envs[0].get_goals()))
    all_eps_data = []
    num_eval_eps = 1
    for ep_idx in range(num_eval_eps):
      ep_data_across_goals = []
      for idx in range(num_goals):
        eval_envs[0].set_goal_idx(idx)
        eval_policy = functools.partial(agent, training=False)
        sim_out = tools.simulate(eval_policy, eval_envs, episodes=1)
        obs, eps_data = sim_out[4], sim_out[6]

        ep_data_across_goals.append(process_eps_data(eps_data))
        video = eval_envs[0]._convert([t['image'] for t in eval_envs[0]._episode])
        executions.append(video[None])
        goals.append(obs[0]['image_goal'][None])

      all_eps_data.append(ep_data_across_goals)

    if ep_idx == 0:
      executions = np.concatenate(executions, 0)
      goals = np.stack(goals, 0)
      goals = np.repeat(goals, executions.shape[1], 1)
      gc_video = np.concatenate([goals, executions], -3)
      agent._logger.video(f'eval_gc_policy', gc_video)
      logger.write()

    with pathlib.Path(logdir / ("distance_func_logs_trained_model/step_"+str(logger.step)+".pkl") ).open('wb') as f:
      pickle.dump(all_eps_data, f)

    if config.sync_s3:
      os.system('aws s3 sync '+str(logdir)+ ' s3://goalexp2021/research_code/goalexp_data/'+str(logdir))

    if not config.training:
        continue
    print('Start training.')
    state = tools.simulate(agent, train_envs, config.eval_every, state=state)
    agent.save(logdir / 'variables.pkl')
  for env in train_envs + eval_envs:
    try:
      env.close()
    except Exception:
      pass


if __name__ == '__main__':
  args, remaining = parse_dreamer_args()
  main(args.logdir, remaining)

import os
import functools
import tools
import tensorflow as tf
import numpy as np
import pickle
import pathlib
import off_policy
from dreamer import Dreamer, setup_dreamer, create_envs, count_steps, make_dataset, parse_dreamer_args


class GCDreamer(Dreamer):
  def __init__(self, config, logger, dataset):
    if config.offpolicy_opt:
      self._off_policy_handler = off_policy.GCOffPolicyOpt(config)
    super().__init__(config, logger, dataset)
    self._should_expl_ep = tools.EveryNCalls(config.expl_every_ep)
    self.skill_to_use = tf.zeros([0], dtype=tf.float16)

  def get_one_time_skill(self):
    skill = self.skill_to_use
    self.skill_to_use = tf.zeros([0], dtype=tf.float16)
    return skill

  def _policy(self, obs, state, training, reset):
    # Choose the goal for the next episode
    # TODO this would be much more elegant if it was implemented in the training loop (can simulate a single episode)
    if state is None:
      if training and self._config.training_goals == 'batch':
        # TODO this modifies the underlying dict, is that fine?
        obs['image_goal'], obs['goal'] = self.sample_replay_goal(obs)
      obs['skill'] = self.get_one_time_skill()
      # The world model adds the observation to the state in this case

    if reset.any() and state is not None:
      # Replace the goal in the agent state at new episode
      # The actor always takes goal from the state, not the observation
      if training and self._config.training_goals == 'batch':
        state[0]['image_goal'], state[0]['goal'] = self.sample_replay_goal(obs)
      else:
        state[0]['image_goal'] = tf.cast(obs['image_goal'], self._float) / 255.0 - 0.5
      state[0]['skill'] = self.get_one_time_skill()

      # Toggle exploration
      self._should_expl_ep()

    # TODO double check everything
    #return super()._policy(obs, state, training, reset, should_expl=self._should_expl_ep.value)
    #if not training:
    return super()._policy(obs, state, training, reset, should_expl=self._should_expl_ep.value)

  def sample_replay_goal(self, obs):
    """ Sample goals from replay buffer """
    assert self._config.gc_input != 'state'
    random_batch = next(self._dataset)
    random_batch = self._wm.preprocess(random_batch)
    
    images = random_batch['image']
    states = random_batch['state']
    if self._config.labelled_env_multiplexing:
      assert obs['env_idx'].shape[0] == 1
      env_ids = random_batch['env_idx'][:, 0]
      if tf.reduce_any(env_ids == obs['env_idx']):
        ids = np.nonzero(env_ids == obs['env_idx'])[0]
        images = tf.gather(images, ids)
        states = tf.gather(states, ids)
    
    random_goals = tf.reshape(images, (-1,) + tuple(images.shape[2:]))
    random_goal_states = tf.reshape(states, (-1,) + tuple(states.shape[2:]))
    # random_goals = tf.random.shuffle(random_goals)
    return random_goals[:obs['image_goal'].shape[0]], random_goal_states[:obs['image_goal'].shape[0]]


def process_eps_data(eps_data):
  keys = eps_data[0].keys()
  new_data = {}
  for key in keys:
    new_data[key] = np.array([eps_data[i][key] for i in range(len(eps_data))]).squeeze()
  return new_data

def main(logdir, config):
  logdir, logger = setup_dreamer(config, logdir)
  eval_envs, eval_eps, train_envs, train_eps, acts = create_envs(config, logger)

  prefill = max(0, config.prefill - count_steps(config.traindir))
  print(f'Prefill dataset ({prefill} steps).')
  random_agent = lambda o, d, s: ([acts.sample() for _ in d], s)
  tools.simulate(random_agent, train_envs, prefill)
  if count_steps(config.evaldir) == 0:
    tools.simulate(random_agent, eval_envs, episodes=1)
  logger.step = config.action_repeat * count_steps(config.traindir)

  print('Simulate agent.')
  train_dataset = make_dataset(train_eps, config)
  eval_dataset = iter(make_dataset(eval_eps, config))
  agent = GCDreamer(config, logger, train_dataset)
  if (logdir / 'variables.pkl').exists():
    agent.load(logdir / 'variables.pkl')
    agent._should_pretrain._once = False

  pathlib.Path(logdir / "distance_func_logs_trained_model").mkdir(parents=True, exist_ok = True)

  state = None
  assert len(eval_envs) == 1
  while agent._step.numpy().item() < config.steps:
    logger.write()
    print('Start gc evaluation.')
    executions = []
    goals = []
    #rews_across_goals = []
    num_goals = min(100, len(eval_envs[0].get_goals()))
    all_eps_data = []
    num_eval_eps = 1
    for ep_idx in range(num_eval_eps):
      ep_data_across_goals = []
      for idx in range(num_goals):
        eval_envs[0].set_goal_idx(idx)
        eval_policy = functools.partial(agent, training=False)
        sim_out = tools.simulate(eval_policy, eval_envs, episodes=1)
        obs, eps_data = sim_out[4], sim_out[6]

        ep_data_across_goals.append(process_eps_data(eps_data))
        video = eval_envs[0]._convert([t['image'] for t in eval_envs[0]._episode])
        executions.append(video[None])
        goals.append(obs[0]['image_goal'][None])

      all_eps_data.append(ep_data_across_goals)

    if ep_idx == 0:
      executions = np.concatenate(executions, 0)
      goals = np.stack(goals, 0)
      goals = np.repeat(goals, executions.shape[1], 1)
      gc_video = np.concatenate([goals, executions], -3)
      agent._logger.video(f'eval_gc_policy', gc_video)
      logger.write()

    with pathlib.Path(logdir / ("distance_func_logs_trained_model/step_"+str(logger.step)+".pkl") ).open('wb') as f:
      pickle.dump(all_eps_data, f)

    if config.sync_s3:
      os.system('aws s3 sync '+str(logdir)+ ' s3://goalexp2021/research_code/goalexp_data/'+str(logdir))

    if not config.training:
        continue
    print('Start training.')
    state = tools.simulate(agent, train_envs, config.eval_every, state=state)
    agent.save(logdir / 'variables.pkl')
  for env in train_envs + eval_envs:
    try:
      env.close()
    except Exception:
      pass


if __name__ == '__main__':
  args, remaining = parse_dreamer_args()
  main(args.logdir, remaining)
# Choose the goal for the next episode # TODO this would be much more elegant if it was implemented in the training loop (can simulate a single episode) # TODO this modifies the underlying dict, is that fine? # The world model adds the observation to the state in this case # Replace the goal in the agent state at new episode # The actor always takes goal from the state, not the observation # Toggle exploration # TODO double check everything #return super()._policy(obs, state, training, reset, should_expl=self._should_expl_ep.value) #if not training: Sample goals from replay buffer # random_goals = tf.random.shuffle(random_goals) #rews_across_goals = []