tpc.py

# This code implements TPC algorithm and the training procedures

import argparse
import collections
import functools
import json
import os
import pathlib
import time
import numpy as np
import tensorflow as tf
from tensorflow.keras.mixed_precision import experimental as prec
from tensorflow_probability import distributions as tfd
import losses
import models
import tools
import wrappers

os.environ["TF_CPP_MIN_LOG_LEVEL"] = "3"
os.environ["MUJOCO_GL"] = "egl"
tf.get_logger().setLevel("ERROR")
# sys.path.append(str(pathlib.Path(__file__).parent))


def define_config():
    config = tools.AttrDict()
    # General.
    config.logdir = pathlib.Path(".")
    config.seed = 0
    config.steps = 1e6
    config.eval_every = 1e4
    config.log_every = 1e3
    config.log_scalars = True
    config.log_images = True
    config.gpu_growth = True
    config.precision = 16
    # Environment.
    config.task = "dmc_walker_walk"
    config.bg_train_path = "./videos/train"
    config.bg_test_path = "./videos/test"
    config.img_source_type = "none"  # use 'video' instead of None to use natural backgrounds
    config.random_bg = False  # set to True will randomize the background for each env step
    config.max_videos = 100
    config.envs = 1
    config.parallel = "none"
    config.action_repeat = 2
    config.time_limit = 1000
    config.prefill = 5000
    config.eval_noise = 0.0
    config.clip_rewards = "none"
    # Model.
    config.deter_size = 200
    config.stoch_size = 30
    config.num_units = 400
    config.dense_act = "elu"
    config.cnn_act = "relu"
    config.cnn_depth = 32
    config.weight_decay = 0.0
    config.weight_decay_pattern = r".*"
    # Training.
    config.cons_coef = 0.1
    config.reward_coef = 1.0
    config.batch_size = 250
    config.batch_length = 50
    config.warm_up = 5  # only compute loss for time step t > 5
    config.train_every = 1000
    config.train_steps = 100
    config.pretrain = 2000
    config.target_update_freq = 100  # update the target value model every 100 steps
    config.model_lr = 6e-4
    config.value_lr = 8e-5
    config.actor_lr = 8e-5
    config.grad_clip = 100.0
    config.dataset_balance = False
    # Behavior.
    config.discount = 0.99
    config.disclam = 0.95
    config.horizon = 15
    config.action_dist = "tanh_normal"
    config.action_init_std = 5.0
    config.expl = "additive_gaussian"
    config.expl_amount = 0.3
    config.expl_decay = 0.0
    config.expl_min = 0.0

    return config


class TPC(tools.Module):
    def __init__(self, config, datadir, actspace, writer):
        self._c = config
        self._actspace = actspace
        self._actdim = actspace.n if hasattr(actspace, "n") else actspace.shape[0]
        self._writer = writer
        self._random = np.random.RandomState(config.seed)
        with tf.device("cpu:0"):
            self._step = tf.Variable(count_steps(datadir, config), dtype=tf.int64)
            self._anneal_cons_coef = tf.Variable(0.05, dtype=tf.float16)
        self._should_pretrain = tools.Once()
        self._should_train = tools.Every(config.train_every)
        self._should_log = tools.Every(config.log_every)
        self._last_log = None
        self._last_time = time.time()
        self._metrics = collections.defaultdict(tf.metrics.Mean)
        self._metrics["expl_amount"]  # Create variable for checkpoint.
        self._float = prec.global_policy().compute_dtype
        self._dataset = iter(load_dataset(datadir, self._c))
        self._build_model()

    def __call__(self, obs, reset, state=None, training=True):
        step = self._step.numpy().item()
        cons_coef = self._anneal_cons_coef.numpy().item()
        tf.summary.experimental.set_step(step)
        if state is not None and reset.any():
            mask = tf.cast(1 - reset, self._float)[:, None]
            state = tf.nest.map_structure(lambda x: x * mask, state)
        if self._should_train(step):
            log = self._should_log(step)
            pretrain = self._should_pretrain()
            n = self._c.pretrain if pretrain else self._c.train_steps
            print(f"Training for {n} steps.")
            for train_step in range(n):
                if pretrain:
                    if train_step % (n // 5) == 0 and cons_coef < self._c.cons_coef:
                        self._anneal_cons_coef.assign_add(0.01)
                log_images = self._c.log_images and log and train_step == 0
                if train_step % self._c.target_update_freq == 0:
                    self._target_value.update(self._value)
                self.train(next(self._dataset), log_images, pretrain)
            if log:
                self._write_summaries()
        action, state = self.policy(obs, state, training)
        if training:
            self._step.assign_add(len(reset) * self._c.action_repeat)
            if step > 10000 and step % 10000 == 0 and cons_coef < self._c.cons_coef:
                self._anneal_cons_coef.assign_add(0.05)
        return action, state

    @tf.function
    def policy(self, obs, state, training):
        if state is None:
            latent = self._dynamics.initial(len(obs["image"]))
            action = tf.zeros((len(obs["image"]), self._actdim), self._float)
        else:
            latent, action = state
        embed = self._encode(preprocess(obs, self._c))
        latent, _ = self._dynamics.obs_step(latent, action, embed)
        feat = self._dynamics.get_feat(latent)
        if training:
            action = self._actor(feat).sample()
        else:
            action = self._actor(feat).mode()
        action = self._exploration(action, training)
        state = (latent, action)
        return action, state

    def load(self, filename):
        super().load(filename)
        self._should_pretrain()

    @tf.function()
    def train(self, data, log_images=False, pretrain=False):
        self._train(data, log_images, pretrain)

    def _train(self, data, log_images, pretrain):
        with tf.GradientTape(persistent=True) as model_tape:
            embed = self._encode(data)
            post, prior = self._dynamics.observe(embed, data["action"])
            avg_std = tf.reduce_mean(prior["std"])
            avg_z_norm = tf.reduce_mean(tf.reduce_sum(post["stoch_clean"] ** 2, axis=-1))
            feat = self._dynamics.get_feat(post)
            reward_pred = self._reward(feat)
            likes = tools.AttrDict()
            post_noise = post["stoch_clean"] + tf.random.normal(
                post["stoch_clean"].shape, mean=0.0, stddev=0.2, dtype=post["stoch_clean"].dtype
            )
            likes.nce = losses.nce_seq(prior["mean"], prior["std"], post_noise, self._c.warm_up)
            likes.iid_nce = losses.nce_seq(post["stoch_clean"], 0.2 * tf.ones_like(post["stoch_clean"]), post_noise, 0)
            likes.consistency = losses.cons_seq(prior["mean"], prior["std"], post["stoch_clean"], self._c.warm_up)
            likes.center_norm = -tf.reduce_sum((tf.reduce_mean(post["stoch_clean"], axis=[0, 1])) ** 2, axis=-1)
            likes.reward = tf.reduce_mean(reward_pred.log_prob(data["reward"]))

            model_loss = -(
                likes.nce
                + self._anneal_cons_coef * likes.consistency
                + likes.iid_nce
                + self._c.reward_coef * likes.reward
                + 0.01 * likes.center_norm
            )

        with tf.GradientTape() as actor_tape:
            post_ = {k: v[:50] for k, v in post.items()}
            imag_feat = self._imagine_ahead(post_)
            reward = self._reward(imag_feat).mode()
            pcont = self._c.discount * tf.ones_like(reward)
            target_value = self._target_value(imag_feat).mode()
            target_returns = tools.lambda_return(
                reward[:-1],
                target_value[:-1],
                pcont[:-1],
                bootstrap=target_value[-1],
                lambda_=self._c.disclam,
                axis=0,
                estimate_value=True,
            )
            discount = tf.stop_gradient(tf.math.cumprod(tf.concat([tf.ones_like(pcont[:1]), pcont[:-2]], 0), 0))
            actor_loss = -tf.reduce_mean(discount * target_returns)

        with tf.GradientTape() as value_tape:
            value_pred = self._value(imag_feat)[:-1]
            target = tf.stop_gradient(target_returns)
            value_loss = -tf.reduce_mean(discount * value_pred.log_prob(target))

        model_grads, model_norm = self._model_opt.get_grad(model_tape, model_loss)
        actor_grads, actor_norm = self._actor_opt.get_grad(actor_tape, actor_loss)
        value_grads, value_norm = self._value_opt.get_grad(value_tape, value_loss)
        if pretrain:
            actor_grads = [tf.zeros_like(grad) for grad in actor_grads]
            actor_norm = tf.ones_like(actor_norm)
            value_grads = [tf.zeros_like(grad) for grad in value_grads]
            value_norm = tf.ones_like(value_norm)

        self._model_opt(model_grads, model_norm)
        self._actor_opt(actor_grads, actor_norm)
        self._value_opt(value_grads, value_norm)

        if tf.distribute.get_replica_context().replica_id_in_sync_group == 0:
            if self._c.log_scalars:
                self._scalar_summaries(
                    data,
                    feat,
                    likes,
                    model_loss,
                    value_loss,
                    actor_loss,
                    avg_z_norm,
                    avg_std,
                    model_norm,
                    value_norm,
                    actor_norm,
                )

    def _build_model(self):
        acts = dict(elu=tf.nn.elu, relu=tf.nn.relu, swish=tf.nn.swish, leaky_relu=tf.nn.leaky_relu)
        cnn_act = acts[self._c.cnn_act]
        act = acts[self._c.dense_act]
        self._encode = models.ConvEncoder(self._c.cnn_depth, cnn_act)
        self._dynamics = models.RSSM(self._c.stoch_size, self._c.deter_size, self._c.deter_size)
        self._reward = models.DenseDecoder((), 2, self._c.num_units, act=act)
        self._value = models.DenseDecoder((), 3, self._c.num_units, act=act)
        self._target_value = models.DenseDecoder((), 3, self._c.num_units, act=act)
        self._actor = models.ActionDecoder(
            self._actdim, 4, self._c.num_units, self._c.action_dist, init_std=self._c.action_init_std, act=act
        )
        model_modules = [self._encode, self._dynamics, self._reward]

        Optimizer = functools.partial(
            tools.Adam, wd=self._c.weight_decay, clip=self._c.grad_clip, wdpattern=self._c.weight_decay_pattern
        )
        self._model_opt = Optimizer("model", model_modules, self._c.model_lr)
        self._value_opt = Optimizer("value", [self._value], self._c.value_lr)
        self._actor_opt = Optimizer("actor", [self._actor], self._c.actor_lr)
        # Do a train step to initialize all variables, including optimizer
        # statistics. Ideally, we would use batch size zero, but that doesn't work
        # in multi-GPU mode.
        self.train(next(self._dataset))

    def _exploration(self, action, training):
        if training:
            amount = self._c.expl_amount
            if self._c.expl_decay:
                amount *= 0.5 ** (tf.cast(self._step, tf.float32) / self._c.expl_decay)
            if self._c.expl_min:
                amount = tf.maximum(self._c.expl_min, amount)
            self._metrics["expl_amount"].update_state(amount)
        elif self._c.eval_noise:
            amount = self._c.eval_noise
        else:
            return action
        if self._c.expl == "additive_gaussian":
            return tf.clip_by_value(tfd.Normal(action, amount).sample(), -1, 1)
        if self._c.expl == "completely_random":
            return tf.random.uniform(action.shape, -1, 1)
        if self._c.expl == "epsilon_greedy":
            indices = tfd.Categorical(0 * action).sample()
            return tf.where(
                tf.random.uniform(action.shape[:1], 0, 1) < amount,
                tf.one_hot(indices, action.shape[-1], dtype=self._float),
                action,
            )
        raise NotImplementedError(self._c.expl)

    def _imagine_ahead(self, post):
        def flatten(x):
            return tf.reshape(x, [-1] + list(x.shape[2:]))
        start = {k: flatten(v) for k, v in post.items()}

        def policy(state):
            return self._actor(tf.stop_gradient(self._dynamics.get_feat(state))).sample()

        def scan_fn(prev, _):
            return self._dynamics.img_step(prev, policy(prev))
        states = tools.static_scan(scan_fn, tf.range(self._c.horizon), start)
        imag_feat = self._dynamics.get_feat(states)
        return imag_feat

    def _scalar_summaries(
        self,
        data,
        feat,
        likes,
        model_loss,
        value_loss,
        actor_loss,
        avg_z_norm,
        avg_std,
        model_norm,
        value_norm,
        actor_norm,
    ):
        self._metrics["model_grad_norm"].update_state(model_norm)
        self._metrics["value_grad_norm"].update_state(value_norm)
        self._metrics["actor_grad_norm"].update_state(actor_norm)
        self._metrics["avg_z_norm"].update_state(avg_z_norm)
        self._metrics["avg_std"].update_state(avg_std)
        for name, logprob in likes.items():
            self._metrics[name + "_loss"].update_state(-logprob)
        self._metrics["cons_coef"].update_state(self._anneal_cons_coef)
        self._metrics["model_loss"].update_state(model_loss)
        self._metrics["value_loss"].update_state(value_loss)
        self._metrics["actor_loss"].update_state(actor_loss)
        self._metrics["action_ent"].update_state(self._actor(feat).entropy())

    def _write_summaries(self):
        step = int(self._step.numpy())
        metrics = [(k, float(v.result())) for k, v in self._metrics.items()]
        if self._last_log is not None:
            duration = time.time() - self._last_time
            self._last_time += duration
            metrics.append(("fps", (step - self._last_log) / duration))
        self._last_log = step
        [m.reset_states() for m in self._metrics.values()]
        with (self._c.logdir / "metrics.jsonl").open("a") as f:
            f.write(json.dumps({"step": step, **dict(metrics)}) + "\n")
        [tf.summary.scalar("agent/" + k, m) for k, m in metrics]
        print(f"[{step}]", " / ".join(f"{k} {v:.3f}" for k, v in metrics))
        self._writer.flush()


def preprocess(obs, config):
    dtype = prec.global_policy().compute_dtype
    obs = obs.copy()
    with tf.device("cpu:0"):
        obs["image"] = tf.cast(obs["image"], dtype) / 255.0 - 0.5
        clip_rewards = dict(none=lambda x: x, tanh=tf.tanh)[config.clip_rewards]
        obs["reward"] = clip_rewards(obs["reward"])
    return obs


def count_steps(datadir, config):
    return tools.count_episodes(datadir)[1] * config.action_repeat


def load_dataset(directory, config):
    episode = next(tools.load_episodes(directory, 1))
    types = {k: v.dtype for k, v in episode.items()}
    shapes = {k: (None,) + v.shape[1:] for k, v in episode.items()}

    def generator():
        return tools.load_episodes(directory, config.train_steps, config.batch_length, config.dataset_balance)

    dataset = tf.data.Dataset.from_generator(generator, types, shapes)
    dataset = dataset.batch(config.batch_size, drop_remainder=True)
    dataset = dataset.map(functools.partial(preprocess, config=config))
    dataset = dataset.prefetch(10)
    return dataset


def summarize_episode(episode, config, datadir, writer, prefix):
    episodes, steps = tools.count_episodes(datadir)
    length = (len(episode["reward"]) - 1) * config.action_repeat
    ret = episode["reward"].sum()
    print(f"{prefix.title()} episode of length {length} with return {ret:.1f}.")
    metrics = [
        (f"{prefix}/return", float(episode["reward"].sum())),
        (f"{prefix}/length", len(episode["reward"]) - 1),
        ("episodes", episodes),
    ]
    step = count_steps(datadir, config)
    with (config.logdir / "metrics.jsonl").open("a") as f:
        f.write(json.dumps(dict([("step", step)] + metrics)) + "\n")
    with writer.as_default():  # Env might run in a different thread.
        tf.summary.experimental.set_step(step)
        [tf.summary.scalar("sim/" + k, v) for k, v in metrics]
        if prefix == "test":
            tools.video_summary(f"sim/{prefix}/video", episode["image"][None])


def make_env(config, bg_path, writer, prefix, datadir, store):
    suite, task = config.task.split("_", 1)
    if suite == "dmc":
        env = wrappers.DeepMindControl(
            task,
            bg_path=bg_path,
            img_source=config.img_source_type,
            random_bg=config.random_bg,
            max_videos=config.max_videos,
        )
        env = wrappers.ActionRepeat(env, config.action_repeat)
        env = wrappers.NormalizeActions(env)
    else:
        raise NotImplementedError(suite)
    env = wrappers.TimeLimit(env, config.time_limit / config.action_repeat)
    callbacks = []
    if store:
        callbacks.append(lambda ep: tools.save_episodes(datadir, [ep]))
    callbacks.append(lambda ep: summarize_episode(ep, config, datadir, writer, prefix))
    env = wrappers.Collect(env, callbacks, config.precision)
    env = wrappers.RewardObs(env)
    return env


def main(config):
    if config.gpu_growth:
        for gpu in tf.config.experimental.list_physical_devices("GPU"):
            tf.config.experimental.set_memory_growth(gpu, True)
    assert config.precision in (16, 32), config.precision
    if config.precision == 16:
        prec.set_policy(prec.Policy("mixed_float16"))
    config.steps = int(config.steps)
    config.logdir.mkdir(parents=True, exist_ok=True)
    print("Logdir", config.logdir)

    # Create environments.
    datadir = config.logdir / "episodes"
    writer = tf.summary.create_file_writer(str(config.logdir), max_queue=1000, flush_millis=20000)
    writer.set_as_default()
    train_envs = [
        wrappers.Async(
            lambda: make_env(config, config.bg_train_path, writer, "train", datadir, store=True), config.parallel
        )
        for _ in range(config.envs)
    ]
    test_envs = [
        wrappers.Async(
            lambda: make_env(config, config.bg_test_path, writer, "test", datadir, store=False), config.parallel
        )
        for _ in range(config.envs)
    ]
    actspace = train_envs[0].action_space

    # Prefill dataset with random episodes.
    step = count_steps(datadir, config)
    prefill = max(0, config.prefill - step)
    print(f"Prefill dataset with {prefill} steps.")

    def random_agent(o, d, _):
        return ([actspace.sample() for _ in d], None)

    tools.simulate(random_agent, train_envs, prefill / config.action_repeat)
    writer.flush()

    # Train and regularly evaluate the agent.
    step = count_steps(datadir, config)
    print(f"Simulating agent for {config.steps-step} steps.")
    agent = TPC(config, datadir, actspace, writer)
    if (config.logdir / "variables.pkl").exists():
        print("Load checkpoint.")
        agent.load(config.logdir / "variables.pkl")
    state = None
    while step < config.steps:
        print("Start evaluation.")
        tools.simulate(functools.partial(agent, training=False), test_envs, episodes=1)
        writer.flush()
        print("Start collection.")
        steps = config.eval_every // config.action_repeat
        state = tools.simulate(agent, train_envs, steps, state=state)
        step = count_steps(datadir, config)
        agent.save(config.logdir / "variables.pkl")
    for env in train_envs + test_envs:
        env.close()


if __name__ == "__main__":
    try:
        import colored_traceback

        colored_traceback.add_hook()
    except ImportError:
        pass
    parser = argparse.ArgumentParser()
    for key, value in define_config().items():
        parser.add_argument(f"--{key}", type=tools.args_type(value), default=value)
    main(parser.parse_args())