commit 2c4a2c08cfdb1338fc1d0356d6ed0d4ccb2b19b4 Author: Isaac Kauvar Date: Thu Jun 22 09:54:27 2023 -0400 First commit. diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..738a527 --- /dev/null +++ b/.gitignore @@ -0,0 +1,7 @@ +.pytest_cache +dist +__pycache__/ +*.py[cod] +*.egg-info +MUJOCO_LOG.TXT +; diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..039183b --- /dev/null +++ b/LICENSE @@ -0,0 +1,19 @@ +Copyright (c) 2023 Danijar Hafner + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/MANIFEST.in b/MANIFEST.in new file mode 100644 index 0000000..f9bd145 --- /dev/null +++ b/MANIFEST.in @@ -0,0 +1 @@ +include requirements.txt diff --git a/README.md b/README.md new file mode 100644 index 0000000..28b61ec --- /dev/null +++ b/README.md @@ -0,0 +1,71 @@ +Fork of https://github.com/danijar/dreamerv3 on February 27, 2023 + +# Install instructions on a fresh Ubuntu 22.04 (x86) install +```bash +sudo apt install build-essential -y + +# Replace ubuntu2204 with your ubuntu version if it's different +wget https://developer.download.nvidia.com/compute/cuda/repos/ubuntu2204/x86_64/cuda-keyring_1.0-1_all.deb + +sudo dpkg -i cuda-keyring_1.0-1_all.deb +sudo apt update +sudo apt install cuda-11-8 -y + +echo 'export CUDA_HOME=/usr/local/cuda' >> ~/.bashrc +echo 'export LD_LIBRARY_PATH=$LD_LIBRARY_PATH:/usr/local/cuda/lib64:/usr/local/cuda/extras/CUPTI/lib64' >> ~/.bashrc +echo 'export PATH=$PATH:$CUDA_HOME/bin' >> ~/.bashrc + + +# reload the bashrc to set the cuda path +source ~/.bashrc + +sudo apt-get install libcudnn8=8.8.0.121-1+cuda11.8 +sudo apt-get install libcudnn8-dev=8.8.0.121-1+cuda11.8 + +mkdir src +cd src +# (optional) git config --global credential.helper store +git clone https://github.com/AutonomousAgentsLab/curiousreplay-dv3.git + +cd curiousreplay-dv3 +git checkout release-working-ik + +sudo apt install python-is-python3 python3.10-venv ffmpeg -y + +# Create and activate a virtual environment +python -m venv ~/src/envs/dv3 +source ~/src/envs/dv3/bin/activate + +pip install --upgrade pip +pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html +pip install -r requirements.txt +pip install dm-reverb +``` + +# Run Curious Replay +```bash +# Run curious-replay crafter +python dreamerv3/train.py --logdir ~/logdir/crafter-dv3-cr_1 \ +--env.crafter.outdir ~/logdir/crafter-dv3-cr_1 --configs crafter --replay curious-replay + +# Run curious-replay DMC +python dreamerv3/train.py --logdir ~/logdir/dmc_vision-dv3-cr_1 \ +--configs dmc_vision --replay curious-replay --envs.amount 1 --task dmc_walker_walk + +# Run baseline crafter +python dreamerv3/train.py --logdir ~/logdir/crafter-dv3_1 \ +--env.crafter.outdir ~/logdir/crafter-dv3_1 --configs crafter + +# Run tensorboard +tensorboard --logdir ~/logdir/crafter-dv3-cr_1 + +# Summarize crafter results +pip install pandas matplotlib +python dreamerv3/plot_crafter.py +``` + +# Limitations + +* No support for parallel environments, so it may need to be run with `--envs.amount 1` flag to override the default number of envs. +* No support for resuming runs + diff --git a/dreamerv3/Dockerfile b/dreamerv3/Dockerfile new file mode 100644 index 0000000..458ae9a --- /dev/null +++ b/dreamerv3/Dockerfile @@ -0,0 +1,63 @@ +# 1. Test setup: +# docker run -it --rm --gpus all nvidia/cuda:11.4.2-cudnn8-runtime-ubuntu20.04 nvidia-smi +# +# If the above does not work, try adding the --privileged flag +# and changing the command to `sh -c 'ldconfig -v && nvidia-smi'`. +# +# 2. Start training: +# docker build -f dreamerv3/Dockerfile -t img . && \ +# docker run -it --rm --gpus all -v ~/logdir:/logdir img \ +# sh scripts/xvfb_run.sh python3 dreamerv3/train.py \ +# --logdir "/logdir/$(date +%Y%m%d-%H%M%S)" \ +# --configs dmc_vision --task dmc_walker_walk +# +# 3. See results: +# tensorboard --logdir ~/logdir + +# System +FROM nvidia/cuda:11.4.2-cudnn8-devel-ubuntu20.04 +ARG DEBIAN_FRONTEND=noninteractive +ENV TZ=America/San_Francisco +ENV PYTHONUNBUFFERED 1 +ENV PIP_DISABLE_PIP_VERSION_CHECK 1 +ENV PIP_NO_CACHE_DIR 1 +RUN apt-get update && apt-get install -y \ + ffmpeg git python3-pip vim libglew-dev \ + x11-xserver-utils xvfb \ + && apt-get clean +RUN pip3 install --upgrade pip + +# Envs +ENV MUJOCO_GL egl +ENV DMLAB_DATASET_PATH /dmlab_data +COPY scripts scripts +RUN sh scripts/install-dmlab.sh +RUN sh scripts/install-atari.sh +RUN sh scripts/install-minecraft.sh +ENV NUMBA_CACHE_DIR=/tmp +RUN pip3 install crafter +RUN pip3 install dm_control +RUN pip3 install robodesk +RUN pip3 install bsuite + +# Agent +RUN pip3 install jax[cuda11_cudnn82] -f https://storage.googleapis.com/jax-releases/jax_cuda_releases.html +RUN pip3 install jaxlib +RUN pip3 install tensorflow_probability +RUN pip3 install optax +RUN pip3 install tensorflow-cpu +ENV XLA_PYTHON_CLIENT_MEM_FRACTION 0.8 + +# Google Cloud DNS cache (optional) +ENV GCS_RESOLVE_REFRESH_SECS=60 +ENV GCS_REQUEST_CONNECTION_TIMEOUT_SECS=300 +ENV GCS_METADATA_REQUEST_TIMEOUT_SECS=300 +ENV GCS_READ_REQUEST_TIMEOUT_SECS=300 +ENV GCS_WRITE_REQUEST_TIMEOUT_SECS=600 + +# Embodied +RUN pip3 install numpy cloudpickle ruamel.yaml rich zmq msgpack +COPY . /embodied +RUN chown -R 1000:root /embodied && chmod -R 775 /embodied + +WORKDIR embodied diff --git a/dreamerv3/__init__.py b/dreamerv3/__init__.py new file mode 100644 index 0000000..f0573bf --- /dev/null +++ b/dreamerv3/__init__.py @@ -0,0 +1,6 @@ +import sys, pathlib +sys.path.append(str(pathlib.Path(__file__).parent)) + +from .agent import Agent +configs = Agent.configs +from .train import wrap_env diff --git a/dreamerv3/agent.py b/dreamerv3/agent.py new file mode 100644 index 0000000..235c410 --- /dev/null +++ b/dreamerv3/agent.py @@ -0,0 +1,396 @@ +import embodied +import jax +import jax.numpy as jnp +import ruamel.yaml as yaml +tree_map = jax.tree_util.tree_map +sg = lambda x: tree_map(jax.lax.stop_gradient, x) + +import logging +logger = logging.getLogger() +class CheckTypesFilter(logging.Filter): + def filter(self, record): + return 'check_types' not in record.getMessage() +logger.addFilter(CheckTypesFilter()) + +from . import behaviors +from . import jaxagent +from . import jaxutils +from . import nets +from . import ninjax as nj + + +@jaxagent.Wrapper +class Agent(nj.Module): + + configs = yaml.YAML(typ='safe').load( + (embodied.Path(__file__).parent / 'configs.yaml').read()) + + def __init__(self, obs_space, act_space, step, config): + self.config = config + self.obs_space = obs_space + self.act_space = act_space['action'] + self.step = step + self.wm = WorldModel(obs_space, act_space, config, name='wm') + self.task_behavior = getattr(behaviors, config.task_behavior)( + self.wm, self.act_space, self.config, name='task_behavior') + if config.expl_behavior == 'None': + self.expl_behavior = self.task_behavior + else: + self.expl_behavior = getattr(behaviors, config.expl_behavior)( + self.wm, self.act_space, self.config, name='expl_behavior') + + def policy_initial(self, batch_size): + return ( + self.wm.initial(batch_size), + self.task_behavior.initial(batch_size), + self.expl_behavior.initial(batch_size)) + + def train_initial(self, batch_size): + return self.wm.initial(batch_size) + + def policy(self, obs, state, mode='train'): + self.config.jax.jit and print('Tracing policy function.') + obs = self.preprocess(obs) + (prev_latent, prev_action), task_state, expl_state = state + embed = self.wm.encoder(obs) + latent, _ = self.wm.rssm.obs_step( + prev_latent, prev_action, embed, obs['is_first']) + self.expl_behavior.policy(latent, expl_state) + task_outs, task_state = self.task_behavior.policy(latent, task_state) + expl_outs, expl_state = self.expl_behavior.policy(latent, expl_state) + if mode == 'eval': + outs = task_outs + outs['action'] = outs['action'].sample(seed=nj.rng()) + outs['log_entropy'] = jnp.zeros(outs['action'].shape[:1]) + elif mode == 'explore': + outs = expl_outs + outs['log_entropy'] = outs['action'].entropy() + outs['action'] = outs['action'].sample(seed=nj.rng()) + elif mode == 'train': + outs = task_outs + outs['log_entropy'] = outs['action'].entropy() + outs['action'] = outs['action'].sample(seed=nj.rng()) + state = ((latent, outs['action']), task_state, expl_state) + return outs, state + + def train(self, data, state): + self.config.jax.jit and print('Tracing train function.') + metrics = {} + data = self.preprocess(data) + state, wm_outs, mets = self.wm.train(data, state) + metrics.update(mets) + context = {**data, **wm_outs['post']} + start = tree_map(lambda x: x.reshape([-1] + list(x.shape[2:])), context) + _, mets = self.task_behavior.train(self.wm.imagine, start, context) + metrics.update(mets) + if self.config.expl_behavior != 'None': + _, mets = self.expl_behavior.train(self.wm.imagine, start, context) + metrics.update({'expl_' + key: value for key, value in mets.items()}) + + if 'keyA' in data.keys(): + outs = {'key': data['key'], + 'env_step': data['env_step'], + 'model_loss': metrics['model_loss_raw'].copy(), + 'td_error': metrics['td_error'].copy()} + + else: + outs = {} + + # Don't need the full model_loss_raw or td_error after the priority calculation, summarize it. + metrics.update({'model_loss_raw': metrics['model_loss_raw'].mean()}) + metrics.update({'td_error': metrics['td_error'].mean()}) + + return outs, state, metrics + + def report(self, data): + self.config.jax.jit and print('Tracing report function.') + data = self.preprocess(data) + report = {} + report.update(self.wm.report(data)) + mets = self.task_behavior.report(data) + report.update({f'task_{k}': v for k, v in mets.items()}) + if self.expl_behavior is not self.task_behavior: + mets = self.expl_behavior.report(data) + report.update({f'expl_{k}': v for k, v in mets.items()}) + return report + + def preprocess(self, obs): + obs = obs.copy() + for key, value in obs.items(): + if key.startswith('log_') or key in ('key', 'env_step'): + continue + if len(value.shape) > 3 and value.dtype == jnp.uint8: + value = jaxutils.cast_to_compute(value) / 255.0 + else: + value = value.astype(jnp.float32) + obs[key] = value + obs['cont'] = 1.0 - obs['is_terminal'].astype(jnp.float32) + return obs + + +class WorldModel(nj.Module): + + def __init__(self, obs_space, act_space, config): + self.obs_space = obs_space + self.act_space = act_space['action'] + self.config = config + shapes = {k: tuple(v.shape) for k, v in obs_space.items()} + shapes = {k: v for k, v in shapes.items() if not k.startswith('log_')} + self.encoder = nets.MultiEncoder(shapes, **config.encoder, name='enc') + self.rssm = nets.RSSM(**config.rssm, name='rssm') + self.heads = { + 'decoder': nets.MultiDecoder(shapes, **config.decoder, name='dec'), + 'reward': nets.MLP((), **config.reward_head, name='rew'), + 'cont': nets.MLP((), **config.cont_head, name='cont')} + self.opt = jaxutils.Optimizer(name='model_opt', **config.model_opt) + scales = self.config.loss_scales.copy() + image, vector = scales.pop('image'), scales.pop('vector') + scales.update({k: image for k in self.heads['decoder'].cnn_shapes}) + scales.update({k: vector for k in self.heads['decoder'].mlp_shapes}) + self.scales = scales + + def initial(self, batch_size): + prev_latent = self.rssm.initial(batch_size) + prev_action = jnp.zeros((batch_size, *self.act_space.shape)) + return prev_latent, prev_action + + def train(self, data, state): + modules = [self.encoder, self.rssm, *self.heads.values()] + mets, (state, outs, metrics) = self.opt( + modules, self.loss, data, state, has_aux=True) + metrics.update(mets) + return state, outs, metrics + + def loss(self, data, state): + embed = self.encoder(data) + prev_latent, prev_action = state + prev_actions = jnp.concatenate([ + prev_action[:, None], data['action'][:, :-1]], 1) + post, prior = self.rssm.observe( + embed, prev_actions, data['is_first'], prev_latent) + dists = {} + feats = {**post, 'embed': embed} + for name, head in self.heads.items(): + out = head(feats if name in self.config.grad_heads else sg(feats)) + out = out if isinstance(out, dict) else {name: out} + dists.update(out) + losses = {} + losses['dyn'] = self.rssm.dyn_loss(post, prior, **self.config.dyn_loss) + losses['rep'] = self.rssm.rep_loss(post, prior, **self.config.rep_loss) + for key, dist in dists.items(): + loss = -dist.log_prob(data[key].astype(jnp.float32)) + assert loss.shape == embed.shape[:2], (key, loss.shape) + losses[key] = loss + scaled = {k: v * self.scales[k] for k, v in losses.items()} + model_loss = sum(scaled.values()) + out = {'embed': embed, 'post': post, 'prior': prior} + out.update({f'{k}_loss': v for k, v in losses.items()}) + last_latent = {k: v[:, -1] for k, v in post.items()} + last_action = data['action'][:, -1] + state = last_latent, last_action + metrics = self._metrics(data, dists, post, prior, losses, model_loss) + metrics['model_loss_raw'] = model_loss # Store model loss for Curious Replay prioritization + return model_loss.mean(), (state, out, metrics) + + def imagine(self, policy, start, horizon): + first_cont = (1.0 - start['is_terminal']).astype(jnp.float32) + keys = list(self.rssm.initial(1).keys()) + start = {k: v for k, v in start.items() if k in keys} + start['action'] = policy(start) + def step(prev, _): + prev = prev.copy() + state = self.rssm.img_step(prev, prev.pop('action')) + return {**state, 'action': policy(state)} + traj = jaxutils.scan( + step, jnp.arange(horizon), start, self.config.imag_unroll) + traj = { + k: jnp.concatenate([start[k][None], v], 0) for k, v in traj.items()} + cont = self.heads['cont'](traj).mode() + traj['cont'] = jnp.concatenate([first_cont[None], cont[1:]], 0) + discount = 1 - 1 / self.config.horizon + traj['weight'] = jnp.cumprod(discount * traj['cont'], 0) / discount + return traj + + def report(self, data): + state = self.initial(len(data['is_first'])) + report = {} + report.update(self.loss(data, state)[-1][-1]) + context, _ = self.rssm.observe( + self.encoder(data)[:6, :5], data['action'][:6, :5], + data['is_first'][:6, :5]) + start = {k: v[:, -1] for k, v in context.items()} + recon = self.heads['decoder'](context) + openl = self.heads['decoder']( + self.rssm.imagine(data['action'][:6, 5:], start)) + for key in self.heads['decoder'].cnn_shapes.keys(): + truth = data[key][:6].astype(jnp.float32) + model = jnp.concatenate([recon[key].mode()[:, :5], openl[key].mode()], 1) + error = (model - truth + 1) / 2 + video = jnp.concatenate([truth, model, error], 2) + report[f'openl_{key}'] = jaxutils.video_grid(video) + return report + + def _metrics(self, data, dists, post, prior, losses, model_loss): + entropy = lambda feat: self.rssm.get_dist(feat).entropy() + metrics = {} + metrics.update(jaxutils.tensorstats(entropy(prior), 'prior_ent')) + metrics.update(jaxutils.tensorstats(entropy(post), 'post_ent')) + metrics.update({f'{k}_loss_mean': v.mean() for k, v in losses.items()}) + metrics.update({f'{k}_loss_std': v.std() for k, v in losses.items()}) + metrics['model_loss_mean'] = model_loss.mean() + metrics['model_loss_std'] = model_loss.std() + metrics['reward_max_data'] = jnp.abs(data['reward']).max() + metrics['reward_max_pred'] = jnp.abs(dists['reward'].mean()).max() + if 'reward' in dists and not self.config.jax.debug_nans: + stats = jaxutils.balance_stats(dists['reward'], data['reward'], 0.1) + metrics.update({f'reward_{k}': v for k, v in stats.items()}) + if 'cont' in dists and not self.config.jax.debug_nans: + stats = jaxutils.balance_stats(dists['cont'], data['cont'], 0.5) + metrics.update({f'cont_{k}': v for k, v in stats.items()}) + return metrics + + +class ImagActorCritic(nj.Module): + + def __init__(self, critics, scales, act_space, config): + critics = {k: v for k, v in critics.items() if scales[k]} + for key, scale in scales.items(): + assert not scale or key in critics, key + self.critics = {k: v for k, v in critics.items() if scales[k]} + self.scales = scales + self.act_space = act_space + self.config = config + disc = act_space.discrete + self.grad = config.actor_grad_disc if disc else config.actor_grad_cont + self.actor = nets.MLP( + name='actor', dims='deter', shape=act_space.shape, **config.actor, + dist=config.actor_dist_disc if disc else config.actor_dist_cont) + self.retnorms = { + k: jaxutils.Moments(**config.retnorm, name=f'retnorm_{k}') + for k in critics} + self.opt = jaxutils.Optimizer(name='actor_opt', **config.actor_opt) + + def initial(self, batch_size): + return {} + + def policy(self, state, carry): + return {'action': self.actor(state)}, carry + + def train(self, imagine, start, context): + def loss(start): + policy = lambda s: self.actor(sg(s)).sample(seed=nj.rng()) + traj = imagine(policy, start, self.config.imag_horizon) + loss, metrics = self.loss(traj) + return loss, (traj, metrics) + mets, (traj, metrics) = self.opt(self.actor, loss, start, has_aux=True) + metrics.update(mets) + for key, critic in self.critics.items(): + mets = critic.train(traj, self.actor) + metrics.update({f'{key}_critic_{k}': v for k, v in mets.items()}) + return traj, metrics + + def loss(self, traj): + metrics = {} + advs = [] + total = sum(self.scales[k] for k in self.critics) + for key, critic in self.critics.items(): + rew, ret, base = critic.score(traj, self.actor) + offset, invscale = self.retnorms[key](ret) + normed_ret = (ret - offset) / invscale + normed_base = (base - offset) / invscale + advs.append((normed_ret - normed_base) * self.scales[key] / total) + metrics.update(jaxutils.tensorstats(rew, f'{key}_reward')) + metrics.update(jaxutils.tensorstats(ret, f'{key}_return_raw')) + metrics.update(jaxutils.tensorstats(normed_ret, f'{key}_return_normed')) + metrics[f'{key}_return_rate'] = (jnp.abs(ret) >= 0.5).mean() + + if len(self.critics) != 1: + raise NotImplementedError('Must have exactly one critic for TD error calculation.') + + r = jnp.reshape(rew[0], (self.config.batch_size, self.config.batch_length)) + v = jnp.reshape(base[0], (self.config.batch_size, self.config.batch_length)) + disc = (jnp.reshape(traj['cont'][0], (self.config.batch_size, self.config.batch_length)) * + (1 - 1 / self.config.horizon)) + td_error = r[:, :-1] + disc[:, 1:] * v[:, 1:] - v[:, :-1] + metrics['td_error'] = td_error # Store TD error for PER prioritization + + adv = jnp.stack(advs).sum(0) + policy = self.actor(sg(traj)) + logpi = policy.log_prob(sg(traj['action']))[:-1] + loss = {'backprop': -adv, 'reinforce': -logpi * sg(adv)}[self.grad] + ent = policy.entropy()[:-1] + loss -= self.config.actent * ent + loss *= sg(traj['weight'])[:-1] + loss *= self.config.loss_scales.actor + metrics.update(self._metrics(traj, policy, logpi, ent, adv)) + return loss.mean(), metrics + + def _metrics(self, traj, policy, logpi, ent, adv): + metrics = {} + ent = policy.entropy()[:-1] + rand = (ent - policy.minent) / (policy.maxent - policy.minent) + rand = rand.mean(range(2, len(rand.shape))) + act = traj['action'] + act = jnp.argmax(act, -1) if self.act_space.discrete else act + metrics.update(jaxutils.tensorstats(act, 'action')) + metrics.update(jaxutils.tensorstats(rand, 'policy_randomness')) + metrics.update(jaxutils.tensorstats(ent, 'policy_entropy')) + metrics.update(jaxutils.tensorstats(logpi, 'policy_logprob')) + metrics.update(jaxutils.tensorstats(adv, 'adv')) + metrics['imag_weight_dist'] = jaxutils.subsample(traj['weight']) + return metrics + + +class VFunction(nj.Module): + + def __init__(self, rewfn, config): + self.rewfn = rewfn + self.config = config + self.net = nets.MLP((), name='net', dims='deter', **self.config.critic) + self.slow = nets.MLP((), name='slow', dims='deter', **self.config.critic) + self.updater = jaxutils.SlowUpdater( + self.net, self.slow, + self.config.slow_critic_fraction, + self.config.slow_critic_update) + self.opt = jaxutils.Optimizer(name='critic_opt', **self.config.critic_opt) + + def train(self, traj, actor): + target = sg(self.score(traj)[1]) + mets, metrics = self.opt(self.net, self.loss, traj, target, has_aux=True) + metrics.update(mets) + self.updater() + return metrics + + def loss(self, traj, target): + metrics = {} + traj = {k: v[:-1] for k, v in traj.items()} + dist = self.net(traj) + loss = -dist.log_prob(sg(target)) + if self.config.critic_slowreg == 'logprob': + reg = -dist.log_prob(sg(self.slow(traj).mean())) + elif self.config.critic_slowreg == 'xent': + reg = -jnp.einsum( + '...i,...i->...', + sg(self.slow(traj).probs), + jnp.log(dist.probs)) + else: + raise NotImplementedError(self.config.critic_slowreg) + loss += self.config.loss_scales.slowreg * reg + loss = (loss * sg(traj['weight'])).mean() + loss *= self.config.loss_scales.critic + metrics = jaxutils.tensorstats(dist.mean()) + return loss, metrics + + def score(self, traj, actor=None): + rew = self.rewfn(traj) + assert len(rew) == len(traj['action']) - 1, ( + 'should provide rewards for all but last action') + discount = 1 - 1 / self.config.horizon + disc = traj['cont'][1:] * discount + value = self.net(traj).mean() + vals = [value[-1]] + interm = rew + disc * value[1:] * (1 - self.config.return_lambda) + for t in reversed(range(len(disc))): + vals.append(interm[t] + disc[t] * self.config.return_lambda * vals[-1]) + ret = jnp.stack(list(reversed(vals))[:-1]) + return rew, ret, value[:-1] diff --git a/dreamerv3/behaviors.py b/dreamerv3/behaviors.py new file mode 100644 index 0000000..8a373a6 --- /dev/null +++ b/dreamerv3/behaviors.py @@ -0,0 +1,102 @@ +import jax.numpy as jnp +from tensorflow_probability.substrates import jax as tfp +tfd = tfp.distributions + +from . import agent +from . import expl +from . import ninjax as nj +from . import jaxutils + + +class Greedy(nj.Module): + + def __init__(self, wm, act_space, config): + rewfn = lambda s: wm.heads['reward'](s).mean()[1:] + if config.critic_type == 'vfunction': + critics = {'extr': agent.VFunction(rewfn, config, name='critic')} + else: + raise NotImplementedError(config.critic_type) + self.ac = agent.ImagActorCritic( + critics, {'extr': 1.0}, act_space, config, name='ac') + + def initial(self, batch_size): + return self.ac.initial(batch_size) + + def policy(self, latent, state): + return self.ac.policy(latent, state) + + def train(self, imagine, start, data): + return self.ac.train(imagine, start, data) + + def report(self, data): + return {} + + +class Random(nj.Module): + + def __init__(self, wm, act_space, config): + self.config = config + self.act_space = act_space + + def initial(self, batch_size): + return jnp.zeros(batch_size) + + def policy(self, latent, state): + batch_size = len(state) + shape = (batch_size,) + self.act_space.shape + if self.act_space.discrete: + dist = jaxutils.OneHotDist(jnp.zeros(shape)) + else: + dist = tfd.Uniform(-jnp.ones(shape), jnp.ones(shape)) + dist = tfd.Independent(dist, 1) + return {'action': dist}, state + + def train(self, imagine, start, data): + return None, {} + + def report(self, data): + return {} + + +class Explore(nj.Module): + + REWARDS = { + 'disag': expl.Disag, + } + + def __init__(self, wm, act_space, config): + self.config = config + self.rewards = {} + critics = {} + for key, scale in config.expl_rewards.items(): + if not scale: + continue + if key == 'extr': + rewfn = lambda s: wm.heads['reward'](s).mean()[1:] + critics[key] = agent.VFunction(rewfn, config, name=key) + else: + rewfn = self.REWARDS[key]( + wm, act_space, config, name=key + '_reward') + critics[key] = agent.VFunction(rewfn, config, name=key) + self.rewards[key] = rewfn + scales = {k: v for k, v in config.expl_rewards.items() if v} + self.ac = agent.ImagActorCritic( + critics, scales, act_space, config, name='ac') + + def initial(self, batch_size): + return self.ac.initial(batch_size) + + def policy(self, latent, state): + return self.ac.policy(latent, state) + + def train(self, imagine, start, data): + metrics = {} + for key, rewfn in self.rewards.items(): + mets = rewfn.train(data) + metrics.update({f'{key}_k': v for k, v in mets.items()}) + traj, mets = self.ac.train(imagine, start, data) + metrics.update(mets) + return traj, metrics + + def report(self, data): + return {} diff --git a/dreamerv3/configs.yaml b/dreamerv3/configs.yaml new file mode 100644 index 0000000..d595cd2 --- /dev/null +++ b/dreamerv3/configs.yaml @@ -0,0 +1,274 @@ +defaults: + + seed: 0 + method: name + task: dummy_disc + logdir: /dev/null + replay: uniform + replay_size: 1e6 + replay_online: False + eval_dir: '' + filter: '.*' + + jax: + platform: gpu + jit: True + precision: float16 + prealloc: True + debug_nans: False + logical_cpus: 0 + debug: False + policy_devices: [0] + train_devices: [0] + metrics_every: 10 + + run: + script: train + steps: 1e10 + expl_until: 0 + log_every: 300 + save_every: 900 + eval_every: 1e6 + eval_initial: True + eval_eps: 1 + eval_samples: 1 + train_ratio: 32.0 + train_fill: 0 + eval_fill: 0 + log_zeros: False + log_keys_video: [image] + log_keys_sum: '^$' + log_keys_mean: '(log_entropy)' + log_keys_max: '^$' + from_checkpoint: '' + sync_every: 10 + # actor_addr: 'tcp://127.0.0.1:5551' + actor_addr: 'ipc:///tmp/5551' + actor_batch: 32 + + envs: {amount: 4, parallel: process, length: 0, reset: True, restart: True, discretize: 0, checks: False} + wrapper: {length: 0, reset: True, discretize: 0, checks: False} + env: + atari: {size: [64, 64], repeat: 4, sticky: True, gray: False, actions: all, lives: unused, noops: 0, resize: opencv} + dmlab: {size: [64, 64], repeat: 4, episodic: True} + minecraft: {size: [64, 64], break_speed: 100.0} + dmc: {size: [64, 64], repeat: 2, camera: -1} + loconav: {size: [64, 64], repeat: 2, camera: -1} + crafter: {outdir: /tmp/crafter} + ddmc: {size: [64, 64], repeat: 2, camera: -1} + cdmc: {size: [64, 64], repeat: 2, camera: -1} + + # Agent + task_behavior: Greedy + expl_behavior: None + batch_size: 16 + batch_length: 64 + data_loaders: 8 + + # World Model + grad_heads: [decoder, reward, cont] + rssm: {deter: 4096, units: 1024, stoch: 32, classes: 32, act: silu, norm: layer, initial: learned, unimix: 0.01, unroll: False, action_clip: 1.0, winit: normal, fan: avg} + encoder: {mlp_keys: '.*', cnn_keys: '.*', act: silu, norm: layer, mlp_layers: 5, mlp_units: 1024, cnn: resnet, cnn_depth: 96, cnn_blocks: 0, resize: stride, winit: normal, fan: avg, symlog_inputs: True, minres: 4} + decoder: {mlp_keys: '.*', cnn_keys: '.*', act: silu, norm: layer, mlp_layers: 5, mlp_units: 1024, cnn: resnet, cnn_depth: 96, cnn_blocks: 0, image_dist: mse, vector_dist: symlog_mse, inputs: [deter, stoch], resize: stride, winit: normal, fan: avg, outscale: 1.0, minres: 4, cnn_sigmoid: False} + reward_head: {layers: 5, units: 1024, act: silu, norm: layer, dist: symlog_disc, outscale: 0.0, outnorm: False, inputs: [deter, stoch], winit: normal, fan: avg, bins: 255} + cont_head: {layers: 5, units: 1024, act: silu, norm: layer, dist: binary, outscale: 1.0, outnorm: False, inputs: [deter, stoch], winit: normal, fan: avg} + loss_scales: {image: 1.0, vector: 1.0, reward: 1.0, cont: 1.0, dyn: 0.5, rep: 0.1, actor: 1.0, critic: 1.0, slowreg: 1.0} + dyn_loss: {impl: kl, free: 1.0} + rep_loss: {impl: kl, free: 1.0} + model_opt: {opt: adam, lr: 1e-4, eps: 1e-8, clip: 1000.0, wd: 0.0, warmup: 0, lateclip: 0.0} + + # Actor Critic + actor: {layers: 5, units: 1024, act: silu, norm: layer, minstd: 0.1, maxstd: 1.0, outscale: 1.0, outnorm: False, unimix: 0.01, inputs: [deter, stoch], winit: normal, fan: avg, symlog_inputs: False} + critic: {layers: 5, units: 1024, act: silu, norm: layer, dist: symlog_disc, outscale: 0.0, outnorm: False, inputs: [deter, stoch], winit: normal, fan: avg, bins: 255, symlog_inputs: False} + actor_opt: {opt: adam, lr: 3e-5, eps: 1e-5, clip: 100.0, wd: 0.0, warmup: 0, lateclip: 0.0} + critic_opt: {opt: adam, lr: 3e-5, eps: 1e-5, clip: 100.0, wd: 0.0, warmup: 0, lateclip: 0.0} + actor_dist_disc: onehot + actor_dist_cont: normal + actor_grad_disc: reinforce + actor_grad_cont: backprop + critic_type: vfunction + imag_horizon: 15 + imag_unroll: False + horizon: 333 + return_lambda: 0.95 + critic_slowreg: logprob + slow_critic_update: 1 + slow_critic_fraction: 0.02 + retnorm: {impl: perc_ema, decay: 0.99, max: 1.0, perclo: 5.0, perchi: 95.0} + actent: 3e-4 + + # Exploration + expl_rewards: {extr: 1.0, disag: 0.1} + expl_opt: {opt: adam, lr: 1e-4, eps: 1e-5, clip: 100.0, wd: 0.0, warmup: 0} + disag_head: {layers: 5, units: 1024, act: silu, norm: layer, dist: mse, outscale: 1.0, inputs: [deter, stoch, action], winit: normal, fan: avg} + disag_target: [stoch] + disag_models: 8 + + # Replay Configuration + replay_hyper: {initial_priority: 1e5, c: 1e4, beta: 0.7, epsilon: 0.01, alpha: 0.7, key_find_priority: 1e7} + +minecraft: + + task: minecraft_diamond + envs.amount: 16 + run: + script: train_save + eval_fill: 1e5 + train_ratio: 16 + log_keys_max: '^log_inventory.*' + encoder: {mlp_keys: 'inventory|inventory_max|equipped|health|hunger|breath|reward', cnn_keys: 'image'} + decoder: {mlp_keys: 'inventory|inventory_max|equipped|health|hunger|breath', cnn_keys: 'image'} + +dmlab: + + task: dmlab_explore_goal_locations_small + envs.amount: 8 + encoder: {mlp_keys: '$^', cnn_keys: 'image'} + decoder: {mlp_keys: '$^', cnn_keys: 'image'} + run.train_ratio: 64 + +atari: + + task: atari_pong + envs.amount: 8 + run: + steps: 5.5e7 + eval_eps: 10 + train_ratio: 64 + encoder: {mlp_keys: '$^', cnn_keys: 'image'} + decoder: {mlp_keys: '$^', cnn_keys: 'image'} + +atari100k: + + task: atari_pong + envs: {amount: 1} + env.atari: {gray: False, repeat: 4, sticky: False, noops: 30, actions: needed} + run: + script: train_eval + steps: 1.5e5 + eval_every: 1e5 + eval_initial: False + eval_eps: 100 + train_ratio: 1024 + jax.precision: float32 + rssm.deter: 512 + .*\.cnn_depth: 32 + .*\.layers: 2 + .*\.units$: 512 + actor_eval_sample: True + encoder: {mlp_keys: '$^', cnn_keys: 'image'} + decoder: {mlp_keys: '$^', cnn_keys: 'image'} + +crafter: + + task: crafter_reward + envs.amount: 1 + run: + log_keys_max: '^log_achievement_.*' + log_keys_sum: '^log_reward$' + run.train_ratio: 512 + encoder: {mlp_keys: '$^', cnn_keys: 'image'} + decoder: {mlp_keys: '$^', cnn_keys: 'image'} + +dmc_vision: + + task: dmc_walker_walk + run.train_ratio: 512 + rssm.deter: 512 + .*\.cnn_depth: 32 + .*\.layers: 2 + .*\.units: 512 + encoder: {mlp_keys: '$^', cnn_keys: 'image'} + decoder: {mlp_keys: '$^', cnn_keys: 'image'} + +dmc_proprio: + + task: dmc_walker_walk + run.train_ratio: 512 + rssm.deter: 512 + .*\.cnn_depth: 32 + .*\.layers: 2 + .*\.units: 512 + encoder: {mlp_keys: '.*', cnn_keys: '$^'} + decoder: {mlp_keys: '.*', cnn_keys: '$^'} + +bsuite: + + task: bsuite_mnist/0 + envs: {amount: 1, parallel: none} + run: + script: train + train_ratio: 1024 # 128 for cartpole + rssm.deter: 512 + .*\.cnn_depth: 32 + .*\.layers: 2 + .*\.units: 512 + +loconav: + + task: loconav_ant_maze_m + env.loconav.repeat: 2 + run: + train_ratio: 512 + log_keys_max: '^log_.*' + encoder: {mlp_keys: '.*', cnn_keys: 'image'} + decoder: {mlp_keys: '.*', cnn_keys: 'image'} + +small: + rssm.deter: 512 + .*\.cnn_depth: 32 + .*\.units: 512 + .*\.layers: 2 + +medium: + rssm.deter: 1024 + .*\.cnn_depth: 48 + .*\.units: 640 + .*\.layers: 3 + +large: + rssm.deter: 2048 + .*\.cnn_depth: 64 + .*\.units: 768 + .*\.layers: 4 + +xlarge: + rssm.deter: 4096 + .*\.cnn_depth: 96 + .*\.units: 1024 + .*\.layers: 5 + +multicpu: + + jax: + logical_cpus: 8 + policy_devices: [0, 1] + train_devices: [2, 3, 4, 5, 6, 7] + run: + actor_batch: 4 + envs: + amount: 8 + batch_size: 12 + batch_length: 10 + +debug: + + jax: {jit: True, prealloc: False, debug: True, platform: cpu} + envs: {restart: False, amount: 3} + wrapper: {length: 100, checks: True} + run: + eval_every: 1000 + log_every: 5 + save_every: 10 + train_ratio: 32 + actor_batch: 2 + batch_size: 8 + batch_length: 12 + replay_size: 1e5 + encoder.cnn_depth: 8 + decoder.cnn_depth: 8 + rssm: {deter: 32, units: 16, stoch: 4, classes: 4} + .*unroll: False + .*\.layers: 2 + .*\.units: 16 + .*\.wd$: 0.0 diff --git a/dreamerv3/embodied/__init__.py b/dreamerv3/embodied/__init__.py new file mode 100644 index 0000000..530f8ac --- /dev/null +++ b/dreamerv3/embodied/__init__.py @@ -0,0 +1,11 @@ +try: + import rich.traceback + rich.traceback.install() +except ImportError: + pass + +from .core import * + +from . import envs +from . import replay +from . import run diff --git a/dreamerv3/embodied/core/__init__.py b/dreamerv3/embodied/core/__init__.py new file mode 100644 index 0000000..5131255 --- /dev/null +++ b/dreamerv3/embodied/core/__init__.py @@ -0,0 +1,29 @@ +from .base import Agent, Env, Wrapper, Replay + +from .basics import convert, treemap, pack, unpack +from .basics import print_ as print +from .basics import format_ as format + +from .space import Space +from .path import Path +from .checkpoint import Checkpoint +from .config import Config +from .counter import Counter +from .driver import Driver +from .flags import Flags +from .logger import Logger +from .parallel import Parallel +from .timer import Timer +from .worker import Worker +from .batcher import Batcher +from .metrics import Metrics +from .uuid import uuid + +from .batch import BatchEnv +from .random import RandomAgent +from .distr import Client, Server, BatchServer + +from . import logger +from . import when +from . import wrappers +from . import distr diff --git a/dreamerv3/embodied/core/base.py b/dreamerv3/embodied/core/base.py new file mode 100644 index 0000000..6e8d559 --- /dev/null +++ b/dreamerv3/embodied/core/base.py @@ -0,0 +1,119 @@ +class Agent: + + configs = {} # dict of dicts + + def __init__(self, obs_space, act_space, step, config): + pass + + def dataset(self, generator_fn): + raise NotImplementedError( + 'dataset(generator_fn) -> generator_fn') + + def policy(self, obs, state=None, mode='train'): + raise NotImplementedError( + "policy(obs, state=None, mode='train') -> act, state") + + def train(self, data, state=None): + raise NotImplementedError( + 'train(data, state=None) -> outs, state, metrics') + + def report(self, data): + raise NotImplementedError( + 'report(data) -> metrics') + + def save(self): + raise NotImplementedError('save() -> data') + + def load(self, data): + raise NotImplementedError('load(data) -> None') + + def sync(self): + # This method allows the agent to sync parameters from its training devices + # to its policy devices in the case of a multi-device agent. + pass + + +class Env: + + def __len__(self): + return 0 # Return positive integer for batched envs. + + def __bool__(self): + return True # Env is always truthy, despite length zero. + + def __repr__(self): + return ( + f'{self.__class__.__name__}(' + f'len={len(self)}, ' + f'obs_space={self.obs_space}, ' + f'act_space={self.act_space})') + + @property + def obs_space(self): + # The observation space must contain the keys is_first, is_last, and + # is_terminal. Commonly, it also contains the keys reward and image. By + # convention, keys starting with log_ are not consumed by the agent. + raise NotImplementedError('Returns: dict of spaces') + + @property + def act_space(self): + # The observation space must contain the keys action and reset. This + # restriction may be lifted in the future. + raise NotImplementedError('Returns: dict of spaces') + + def step(self, action): + raise NotImplementedError('Returns: dict') + + def render(self): + raise NotImplementedError('Returns: array') + + def close(self): + pass + + +class Wrapper: + + def __init__(self, env): + self.env = env + + def __len__(self): + return len(self.env) + + def __bool__(self): + return bool(self.env) + + def __getattr__(self, name): + if name.startswith('__'): + raise AttributeError(name) + try: + return getattr(self.env, name) + except AttributeError: + raise ValueError(name) + + +class Replay: + + def __len__(self): + raise NotImplementedError('Returns: total number of steps') + + @property + def stats(self): + raise NotImplementedError('Returns: metrics') + + def add(self, transition, worker=0): + raise NotImplementedError('Returns: None') + + def add_traj(self, trajectory): + raise NotImplementedError('Returns: None') + + def dataset(self): + raise NotImplementedError('Yields: trajectory') + + def prioritize(self, keys, priorities): + pass + + def save(self): + pass + + def load(self, data): + pass diff --git a/dreamerv3/embodied/core/basics.py b/dreamerv3/embodied/core/basics.py new file mode 100644 index 0000000..2b2abea --- /dev/null +++ b/dreamerv3/embodied/core/basics.py @@ -0,0 +1,144 @@ +import builtins +import pickle + +import numpy as np + +from . import space as spacelib + +try: + import rich.console + console = rich.console.Console() +except ImportError: + console = None + + +CONVERSION = { + np.floating: np.float32, + np.signedinteger: np.int64, + np.uint8: np.uint8, + bool: bool, +} + + +def convert(value): + value = np.asarray(value) + if value.dtype not in CONVERSION.values(): + for src, dst in CONVERSION.items(): + if np.issubdtype(value.dtype, src): + if value.dtype != dst: + value = value.astype(dst) + break + else: + raise TypeError(f"Object '{value}' has unsupported dtype: {value.dtype}") + return value + + +def print_(value, color=None): + global console + value = format_(value) + if console: + if color: + value = f'[{color}]{value}[/{color}]' + console.print(value) + else: + builtins.print(value) + + +def format_(value): + if isinstance(value, dict): + if value and all(isinstance(x, spacelib.Space) for x in value.values()): + return '\n'.join(f' {k:<16} {v}' for k, v in value.items()) + items = [f'{format_(k)}: {format_(v)}' for k, v in value.items()] + return '{' + ', '.join(items) + '}' + if isinstance(value, list): + return '[' + ', '.join(f'{format_(x)}' for x in value) + ']' + if isinstance(value, tuple): + return '(' + ', '.join(f'{format_(x)}' for x in value) + ')' + if hasattr(value, 'shape') and hasattr(value, 'dtype'): + shape = ','.join(str(x) for x in value.shape) + dtype = value.dtype.name + for long, short in {'float': 'f', 'uint': 'u', 'int': 'i'}.items(): + dtype = dtype.replace(long, short) + return f'{dtype}[{shape}]' + if isinstance(value, bytes): + value = '0x' + value.hex() if r'\x' in str(value) else str(value) + if len(value) > 32: + value = value[:32 - 3] + '...' + return str(value) + + +def treemap(fn, *trees, isleaf=None): + assert trees, 'Provide one or more nested Python structures' + kw = dict(isleaf=isleaf) + first = trees[0] + assert all(isinstance(x, type(first)) for x in trees) + if isleaf and isleaf(trees): + return fn(*trees) + if isinstance(first, list): + assert all(len(x) == len(first) for x in trees), format_(trees) + return [treemap( + fn, *[t[i] for t in trees], **kw) for i in range(len(first))] + if isinstance(first, tuple): + assert all(len(x) == len(first) for x in trees), format_(trees) + return tuple([treemap( + fn, *[t[i] for t in trees], **kw) for i in range(len(first))]) + if isinstance(first, dict): + assert all(set(x.keys()) == set(first.keys()) for x in trees), ( + format_(trees)) + return {k: treemap(fn, *[t[k] for t in trees], **kw) for k in first} + return fn(*trees) + + +def pack(data): + return pickle.dumps(data) + # import msgpack + # def fn(data): + # if isinstance(data, np.ndarray): + # return [b'type_numpy', list(data.shape), data.dtype.name, data.tobytes()] + # if isinstance(data, bytes): + # return [b'type_bytes', data] + # if isinstance(data, tuple): + # return [b'type_tuple', *[fn(x) for x in data]] + # if isinstance(data, list): + # return [fn(x) for x in data] + # if isinstance(data, str): + # return data.encode('utf-8') + # if isinstance(data, dict): + # return {k: fn(v) for k, v in data.items()} + # if allow_pickle: + # primitives = (type(None), bool, int, float, str, bytes) + # if not isinstance(data, primitives): + # return [b'type_pickle', pickle.dumps(data)] + # return data + # data = fn(data) + # # print(format_(data)) + # data = msgpack.packb( + # data, use_single_float=True, use_bin_type=True, strict_types=True) + # return data + + +def unpack(buffer): + return pickle.loads(buffer) + # import msgpack + # import pickle + # def fn(data): + # if isinstance(data, list) and data and data[0] == b'type_numpy': + # return np.frombuffer(data[3], data[2].decode('utf-8')).reshape(data[1]) + # if isinstance(data, list) and data and data[0] == b'type_bytes': + # return data[1] + # if isinstance(data, list) and data and data[0] == b'type_tuple': + # return tuple([fn(x) for x in data[1:]]) + # if isinstance(data, list) and data and data[0] == b'type_pickle': + # assert allow_pickle, 'Buffer contains pickled Python objects.' + # return pickle.loads(data[1]) + # if isinstance(data, list): + # return [fn(x) for x in data] + # if isinstance(data, str): + # return data.decode('utf-8') + # if isinstance(data, dict): + # return {k.decode('utf-8'): fn(v) for k, v in data.items()} + # return data + # data = msgpack.unpackb(buffer, raw=True, use_list=True) + # data = fn(data) + # # print(format_(data)) + # return data diff --git a/dreamerv3/embodied/core/batch.py b/dreamerv3/embodied/core/batch.py new file mode 100644 index 0000000..1c76e30 --- /dev/null +++ b/dreamerv3/embodied/core/batch.py @@ -0,0 +1,45 @@ +import numpy as np + +from . import base + + +class BatchEnv(base.Env): + + def __init__(self, envs, parallel): + assert all(len(env) == 0 for env in envs) + assert len(envs) > 0 + self._envs = envs + self._parallel = parallel + self._keys = list(self.obs_space.keys()) + + @property + def obs_space(self): + return self._envs[0].obs_space + + @property + def act_space(self): + return self._envs[0].act_space + + def __len__(self): + return len(self._envs) + + def step(self, action): + assert all(len(v) == len(self._envs) for v in action.values()), ( + len(self._envs), {k: v.shape for k, v in action.items()}) + obs = [] + for i, env in enumerate(self._envs): + act = {k: v[i] for k, v in action.items()} + obs.append(env.step(act)) + if self._parallel: + obs = [ob() for ob in obs] + return {k: np.array([ob[k] for ob in obs]) for k in obs[0]} + + def render(self): + return np.stack([env.render() for env in self._envs]) + + def close(self): + for env in self._envs: + try: + env.close() + except Exception: + pass diff --git a/dreamerv3/embodied/core/batcher.py b/dreamerv3/embodied/core/batcher.py new file mode 100644 index 0000000..8f93f9f --- /dev/null +++ b/dreamerv3/embodied/core/batcher.py @@ -0,0 +1,101 @@ +import queue as queuelib +import sys +import threading +import time +import traceback + +import numpy as np + + +class Batcher: + + def __init__( + self, sources, workers=0, postprocess=None, + prefetch_source=4, prefetch_batch=2): + self._workers = workers + self._postprocess = postprocess + if workers: + # Round-robin assign sources to workers. + self._running = True + self._threads = [] + self._queues = [] + assignments = [([], []) for _ in range(workers)] + for index, source in enumerate(sources): + queue = queuelib.Queue(prefetch_source) + self._queues.append(queue) + assignments[index % workers][0].append(source) + assignments[index % workers][1].append(queue) + for args in assignments: + creator = threading.Thread( + target=self._creator, args=args, daemon=True) + creator.start() + self._threads.append(creator) + self._batches = queuelib.Queue(prefetch_batch) + batcher = threading.Thread( + target=self._batcher, args=(self._queues, self._batches), + daemon=True) + batcher.start() + self._threads.append(batcher) + else: + self._iterators = [source() for source in sources] + self._once = False + + def close(self): + if self._workers: + self._running = False + for thread in self._threads: + thread.close() + + def __iter__(self): + if self._once: + raise RuntimeError( + 'You can only create one iterator per Batcher object to ensure that ' + 'data is consumed in order. Create another Batcher object instead.') + self._once = True + return self + + def __call__(self): + return self.__iter__() + + def __next__(self): + if self._workers: + batch = self._batches.get() + else: + elems = [next(x) for x in self._iterators] + batch = {k: np.stack([x[k] for x in elems], 0) for k in elems[0]} + if isinstance(batch, Exception): + raise batch + return batch + + def _creator(self, sources, outputs): + try: + iterators = [source() for source in sources] + while self._running: + waiting = True + for iterator, queue in zip(iterators, outputs): + if queue.full(): + continue + queue.put(next(iterator)) + waiting = False + if waiting: + time.sleep(0.001) + except Exception as e: + e.stacktrace = ''.join(traceback.format_exception(*sys.exc_info())) + outputs[0].put(e) + raise + + def _batcher(self, sources, output): + try: + while self._running: + elems = [x.get() for x in sources] + for elem in elems: + if isinstance(elem, Exception): + raise elem + batch = {k: np.stack([x[k] for x in elems], 0) for k in elems[0]} + if self._postprocess: + batch = self._postprocess(batch) + output.put(batch) # Will wait here if the queue is full. + except Exception as e: + e.stacktrace = ''.join(traceback.format_exception(*sys.exc_info())) + output.put(e) + raise diff --git a/dreamerv3/embodied/core/checkpoint.py b/dreamerv3/embodied/core/checkpoint.py new file mode 100644 index 0000000..e1437b0 --- /dev/null +++ b/dreamerv3/embodied/core/checkpoint.py @@ -0,0 +1,93 @@ +import concurrent.futures +import time + +from . import basics +from . import path + + +class Checkpoint: + + def __init__(self, filename=None, log=True, parallel=True): + self._filename = filename and path.Path(filename) + self._log = log + self._values = {} + self._parallel = parallel + if self._parallel: + self._worker = concurrent.futures.ThreadPoolExecutor(1) + self._promise = None + + def __setattr__(self, name, value): + if name in ('exists', 'save', 'load'): + return super().__setattr__(name, value) + if name.startswith('_'): + return super().__setattr__(name, value) + has_load = hasattr(value, 'load') and callable(value.load) + has_save = hasattr(value, 'save') and callable(value.save) + if not (has_load and has_save): + message = f"Checkpoint entry '{name}' must implement save() and load()." + raise ValueError(message) + self._values[name] = value + + def __getattr__(self, name): + if name.startswith('_'): + raise AttributeError(name) + try: + return getattr(self._values, name) + except AttributeError: + raise ValueError(name) + + def exists(self, filename=None): + assert self._filename or filename + filename = path.Path(filename or self._filename) + exists = self._filename.exists() + self._log and exists and print('Found existing checkpoint.') + self._log and not exists and print('Did not find any checkpoint.') + return exists + + def save(self, filename=None, keys=None): + assert self._filename or filename + filename = path.Path(filename or self._filename) + self._log and print(f'Writing checkpoint: {filename}') + if self._parallel: + self._promise and self._promise.result() + self._promise = self._worker.submit(self._save, filename, keys) + else: + self._save(filename, keys) + + def _save(self, filename, keys): + keys = tuple(self._values.keys() if keys is None else keys) + assert all([not k.startswith('_') for k in keys]), keys + data = {k: self._values[k].save() for k in keys} + data['_timestamp'] = time.time() + if filename.exists(): + old = filename.parent / (filename.name + '.old') + filename.copy(old) + filename.write(basics.pack(data), mode='wb') + old.remove() + else: + filename.write(basics.pack(data), mode='wb') + self._log and print(f'Wrote checkpoint: {filename}') + + def load(self, filename=None, keys=None): + assert self._filename or filename + filename = path.Path(filename or self._filename) + self._log and print(f'Loading checkpoint: {filename}') + data = basics.unpack(filename.read('rb')) + keys = tuple(data.keys() if keys is None else keys) + for key in keys: + if key.startswith('_'): + continue + try: + self._values[key].load(data[key]) + except Exception: + print(f'Error loading {key} from checkpoint.') + raise + if self._log: + age = time.time() - data['_timestamp'] + print(f'Loaded checkpoint from {age:.0f} seconds ago.') + + def load_or_save(self): + if self.exists(): + self.load() + else: + self.save() diff --git a/dreamerv3/embodied/core/config.py b/dreamerv3/embodied/core/config.py new file mode 100644 index 0000000..0c600ed --- /dev/null +++ b/dreamerv3/embodied/core/config.py @@ -0,0 +1,191 @@ +import io +import json +import re + +from . import path + + +class Config(dict): + + SEP = '.' + IS_PATTERN = re.compile(r'.*[^A-Za-z0-9_.-].*') + + def __init__(self, *args, **kwargs): + mapping = dict(*args, **kwargs) + mapping = self._flatten(mapping) + mapping = self._ensure_keys(mapping) + mapping = self._ensure_values(mapping) + self._flat = mapping + self._nested = self._nest(mapping) + # Need to assign the values to the base class dictionary so that + # conversion to dict does not lose the content. + super().__init__(self._nested) + + @property + def flat(self): + return self._flat.copy() + + def save(self, filename): + filename = path.Path(filename) + if filename.suffix == '.json': + filename.write(json.dumps(dict(self))) + elif filename.suffix in ('.yml', '.yaml'): + import ruamel.yaml as yaml + with io.StringIO() as stream: + yaml.safe_dump(dict(self), stream) + filename.write(stream.getvalue()) + else: + raise NotImplementedError(filename.suffix) + + @classmethod + def load(cls, filename): + filename = path.Path(filename) + if filename.suffix == '.json': + return cls(json.loads(filename.read_text())) + elif filename.suffix in ('.yml', '.yaml'): + import ruamel.yaml as yaml + return cls(yaml.safe_load(filename.read_text())) + else: + raise NotImplementedError(filename.suffix) + + def __contains__(self, name): + try: + self[name] + return True + except KeyError: + return False + + def __getattr__(self, name): + if name.startswith('_'): + return super().__getattr__(name) + try: + return self[name] + except KeyError: + raise AttributeError(name) + + def __getitem__(self, name): + result = self._nested + for part in name.split(self.SEP): + try: + result = result[part] + except TypeError: + raise KeyError + if isinstance(result, dict): + result = type(self)(result) + return result + + def __setattr__(self, key, value): + if key.startswith('_'): + return super().__setattr__(key, value) + message = f"Tried to set key '{key}' on immutable config. Use update()." + raise AttributeError(message) + + def __setitem__(self, key, value): + if key.startswith('_'): + return super().__setitem__(key, value) + message = f"Tried to set key '{key}' on immutable config. Use update()." + raise AttributeError(message) + + def __reduce__(self): + return (type(self), (dict(self),)) + + def __str__(self): + lines = ['\nConfig:'] + keys, vals, typs = [], [], [] + for key, val in self.flat.items(): + keys.append(key + ':') + vals.append(self._format_value(val)) + typs.append(self._format_type(val)) + max_key = max(len(k) for k in keys) if keys else 0 + max_val = max(len(v) for v in vals) if vals else 0 + for key, val, typ in zip(keys, vals, typs): + key = key.ljust(max_key) + val = val.ljust(max_val) + lines.append(f'{key} {val} ({typ})') + return '\n'.join(lines) + + def update(self, *args, **kwargs): + result = self._flat.copy() + inputs = self._flatten(dict(*args, **kwargs)) + for key, new in inputs.items(): + if self.IS_PATTERN.match(key): + pattern = re.compile(key) + keys = {k for k in result if pattern.match(k)} + else: + keys = [key] + if not keys: + raise KeyError(f'Unknown key or pattern {key}.') + for key in keys: + old = result[key] + try: + if isinstance(old, int) and isinstance(new, float): + if float(int(new)) != new: + message = f"Cannot convert fractional float {new} to int." + raise ValueError(message) + result[key] = type(old)(new) + except (ValueError, TypeError): + raise TypeError( + f"Cannot convert '{new}' to type '{type(old).__name__}' " + + f"for key '{key}' with previous value '{old}'.") + return type(self)(result) + + def _flatten(self, mapping): + result = {} + for key, value in mapping.items(): + if isinstance(value, dict): + for k, v in self._flatten(value).items(): + if self.IS_PATTERN.match(key) or self.IS_PATTERN.match(k): + combined = f'{key}\\{self.SEP}{k}' + else: + combined = f'{key}{self.SEP}{k}' + result[combined] = v + else: + result[key] = value + return result + + def _nest(self, mapping): + result = {} + for key, value in mapping.items(): + parts = key.split(self.SEP) + node = result + for part in parts[:-1]: + if part not in node: + node[part] = {} + node = node[part] + node[parts[-1]] = value + return result + + def _ensure_keys(self, mapping): + for key in mapping: + assert not self.IS_PATTERN.match(key), key + return mapping + + def _ensure_values(self, mapping): + result = json.loads(json.dumps(mapping)) + for key, value in result.items(): + if isinstance(value, list): + value = tuple(value) + if isinstance(value, tuple): + if len(value) == 0: + message = 'Empty lists are disallowed because their type is unclear.' + raise TypeError(message) + if not isinstance(value[0], (str, float, int, bool)): + message = 'Lists can only contain strings, floats, ints, bools' + message += f' but not {type(value[0])}' + raise TypeError(message) + if not all(isinstance(x, type(value[0])) for x in value[1:]): + message = 'Elements of a list must all be of the same type.' + raise TypeError(message) + result[key] = value + return result + + def _format_value(self, value): + if isinstance(value, (list, tuple)): + return '[' + ', '.join(self._format_value(x) for x in value) + ']' + return str(value) + + def _format_type(self, value): + if isinstance(value, (list, tuple)): + assert len(value) > 0, value + return self._format_type(value[0]) + 's' + return str(type(value).__name__) diff --git a/dreamerv3/embodied/core/counter.py b/dreamerv3/embodied/core/counter.py new file mode 100644 index 0000000..f6e2683 --- /dev/null +++ b/dreamerv3/embodied/core/counter.py @@ -0,0 +1,44 @@ +import functools + + +@functools.total_ordering +class Counter: + + def __init__(self, initial=0): + self.value = initial + + def __repr__(self): + return f'Counter({self.value})' + + def __int__(self): + return int(self.value) + + def __eq__(self, other): + return int(self) == other + + def __ne__(self, other): + return int(self) != other + + def __lt__(self, other): + return int(self) < other + + def __add__(self, other): + return int(self) + other + + def __radd__(self, other): + return other - int(self) + + def __sub__(self, other): + return int(self) - other + + def __rsub__(self, other): + return other - int(self) + + def increment(self, amount=1): + self.value += amount + + def save(self): + return self.value + + def load(self, value): + self.value = value diff --git a/dreamerv3/embodied/core/distr.py b/dreamerv3/embodied/core/distr.py new file mode 100644 index 0000000..ee2b9c9 --- /dev/null +++ b/dreamerv3/embodied/core/distr.py @@ -0,0 +1,221 @@ +import ctypes +import sys +import threading +import time +import traceback +import uuid + +import numpy as np + +from . import basics + + +class Client: + + def __init__(self, address, timeout_ms=-1, ipv6=False): + import zmq + addresses = [address] if isinstance(address, str) else address + context = zmq.Context.instance() + self.socket = context.socket(zmq.REQ) + self.socket.setsockopt(zmq.IDENTITY, uuid.uuid4().bytes) + self.socket.RCVTIMEO = timeout_ms + for address in addresses: + basics.print_(f'Client connecting to {address}', color='green') + ipv6 and self.socket.setsockopt(zmq.IPV6, 1) + self.socket.connect(address) + self.result = True + + def __call__(self, data): + assert isinstance(data, dict), type(data) + if self.result is None: + self._receive() + self.result = None + self.socket.send(basics.pack(data)) + return self._receive + + def _receive(self): + try: + recieved = self.socket.recv() + except Exception as e: + raise RuntimeError(f'Failed to receive data from server: {e}') + self.result = basics.unpack(recieved) + if self.result.get('type', 'data') == 'error': + msg = self.result.get('message', None) + raise RuntimeError(f'Server responded with an error: {msg}') + return self.result + + +class Server: + + def __init__(self, address, function, ipv6=False): + import zmq + context = zmq.Context.instance() + self.socket = context.socket(zmq.REP) + basics.print_(f'Server listening at {address}', color='green') + ipv6 and self.socket.setsockopt(zmq.IPV6, 1) + self.socket.bind(address) + self.function = function + + def run(self): + while True: + payload = self.socket.recv() + inputs = basics.unpack(payload) + assert isinstance(inputs, dict), type(inputs) + try: + result = self.function(inputs) + assert isinstance(result, dict), type(result) + except Exception as e: + result = {'type': 'error', 'message': str(e)} + self.socket.send(basics.pack(payload)) + raise + payload = basics.pack(result) + self.socket.send(payload) + + +class BatchServer: + + def __init__(self, address, batch, function, ipv6=False): + import zmq + context = zmq.Context.instance() + self.socket = context.socket(zmq.ROUTER) + basics.print_(f'BatchServer listening at {address}', color='green') + ipv6 and self.socket.setsockopt(zmq.IPV6, 1) + self.socket.bind(address) + self.function = function + self.batch = batch + + def run(self): + inputs = None + while True: + addresses = [] + for i in range(self.batch): + address, empty, payload = self.socket.recv_multipart() + data = basics.unpack(payload) + assert isinstance(data, dict), type(data) + if inputs is None: + inputs = { + k: np.empty((self.batch, *v.shape), v.dtype) + for k, v in data.items() if not isinstance(v, str)} + for key, value in data.items(): + inputs[key][i] = value + addresses.append(address) + try: + results = self.function(inputs, [x.hex() for x in addresses]) + assert isinstance(results, dict), type(results) + for key, value in results.items(): + if not isinstance(value, str): + assert len(value) == self.batch, (key, value.shape) + except Exception as e: + results = {'type': 'error', 'message': str(e)} + self._respond(addresses, results) + raise + self._respond(addresses, results) + + def _respond(self, addresses, results): + for i, address in enumerate(addresses): + payload = basics.pack({ + k: v if isinstance(v, str) else v[i] + for k, v in results.items()}) + self.socket.send_multipart([address, b'', payload]) + + +class Thread(threading.Thread): + + lock = threading.Lock() + + def __init__(self, fn, *args, name=None): + self.fn = fn + self.exitcode = None + name = name or fn.__name__ + super().__init__(target=self._wrapper, args=args, name=name, daemon=True) + + def _wrapper(self, *args): + try: + self.fn(*args) + except Exception: + with self.lock: + print('-' * 79) + print(f'Exception in worker: {self.name}') + print('-' * 79) + print(''.join(traceback.format_exception(*sys.exc_info()))) + self.exitcode = 1 + raise + self.exitcode = 0 + + def terminate(self): + if not self.is_alive(): + return + if hasattr(self, '_thread_id'): + thread_id = self._thread_id + else: + thread_id = [k for k, v in threading._active.items() if v is self][0] + result = ctypes.pythonapi.PyThreadState_SetAsyncExc( + ctypes.c_long(thread_id), ctypes.py_object(SystemExit)) + if result > 1: + ctypes.pythonapi.PyThreadState_SetAsyncExc( + ctypes.c_long(thread_id), None) + print('Shut down worker:', self.name) + + +class Process: + + lock = None + initializers = [] + + def __init__(self, fn, *args, name=None): + import multiprocessing + import cloudpickle + mp = multiprocessing.get_context('spawn') + if Process.lock is None: + Process.lock = mp.Lock() + name = name or fn.__name__ + initializers = cloudpickle.dumps(self.initializers) + args = (initializers,) + args + self._process = mp.Process( + target=self._wrapper, args=(Process.lock, fn, *args), + name=name) + + def start(self): + self._process.start() + + @property + def name(self): + return self._process.name + + @property + def exitcode(self): + return self._process.exitcode + + def terminate(self): + self._process.terminate() + print('Shut down worker:', self.name) + + def _wrapper(self, lock, fn, *args): + try: + import cloudpickle + initializers, *args = args + for initializer in cloudpickle.loads(initializers): + initializer() + fn(*args) + except Exception: + with lock: + print('-' * 79) + print(f'Exception in worker: {self.name}') + print('-' * 79) + print(''.join(traceback.format_exception(*sys.exc_info()))) + raise + + +def run(workers): + [x.start() for x in workers] + while True: + if all(x.exitcode == 0 for x in workers): + print('All workers terminated successfully.') + return + for worker in workers: + if worker.exitcode not in (None, 0): + # Wait for everybody who wants to print their error messages. + time.sleep(1) + [x.terminate() for x in workers if x is not worker] + raise RuntimeError(f'Stopped workers due to crash in {worker.name}.') + time.sleep(0.1) diff --git a/dreamerv3/embodied/core/driver.py b/dreamerv3/embodied/core/driver.py new file mode 100644 index 0000000..aa4d66e --- /dev/null +++ b/dreamerv3/embodied/core/driver.py @@ -0,0 +1,78 @@ +import collections + +import numpy as np + +from .basics import convert + + +class Driver: + + _CONVERSION = { + np.floating: np.float32, + np.signedinteger: np.int32, + np.uint8: np.uint8, + bool: bool, + } + + def __init__(self, env, **kwargs): + assert len(env) > 0 + self._env = env + self._kwargs = kwargs + self._on_steps = [] + self._on_episodes = [] + self.reset() + + def reset(self): + self._acts = { + k: convert(np.zeros((len(self._env),) + v.shape, v.dtype)) + for k, v in self._env.act_space.items()} + self._acts['reset'] = np.ones(len(self._env), bool) + self._eps = [collections.defaultdict(list) for _ in range(len(self._env))] + self._state = None + + def on_step(self, callback): + self._on_steps.append(callback) + + def on_episode(self, callback): + self._on_episodes.append(callback) + + def __call__(self, policy, steps=0, episodes=0): + step, episode = 0, 0 + while step < steps or episode < episodes: + step, episode = self._step(policy, step, episode) + + def _step(self, policy, step, episode): + assert all(len(x) == len(self._env) for x in self._acts.values()) + acts = {k: v for k, v in self._acts.items() if not k.startswith('log_')} + obs = self._env.step(acts) + obs = {k: convert(v) for k, v in obs.items()} + assert all(len(x) == len(self._env) for x in obs.values()), obs + acts, self._state = policy(obs, self._state, **self._kwargs) + acts = {k: convert(v) for k, v in acts.items()} + if obs['is_last'].any(): + mask = 1 - obs['is_last'] + acts = {k: v * self._expand(mask, len(v.shape)) for k, v in acts.items()} + acts['reset'] = obs['is_last'].copy() + self._acts = acts + trns = {**obs, **acts} + if obs['is_first'].any(): + for i, first in enumerate(obs['is_first']): + if first: + self._eps[i].clear() + for i in range(len(self._env)): + trn = {k: v[i] for k, v in trns.items()} + [self._eps[i][k].append(v) for k, v in trn.items()] + [fn(trn, i, **self._kwargs) for fn in self._on_steps] + step += 1 + if obs['is_last'].any(): + for i, done in enumerate(obs['is_last']): + if done: + ep = {k: convert(v) for k, v in self._eps[i].items()} + [fn(ep.copy(), i, **self._kwargs) for fn in self._on_episodes] + episode += 1 + return step, episode + + def _expand(self, value, dims): + while len(value.shape) < dims: + value = value[..., None] + return value diff --git a/dreamerv3/embodied/core/flags.py b/dreamerv3/embodied/core/flags.py new file mode 100644 index 0000000..785fbcd --- /dev/null +++ b/dreamerv3/embodied/core/flags.py @@ -0,0 +1,102 @@ +import re +import sys + +from . import config + + +class Flags: + + def __init__(self, *args, **kwargs): + self._config = config.Config(*args, **kwargs) + + def parse(self, argv=None, help_exists=True): + parsed, remaining = self.parse_known(argv) + for flag in remaining: + if flag.startswith('--'): + raise ValueError(f"Flag '{flag}' did not match any config keys.") + assert not remaining, remaining + return parsed + + def parse_known(self, argv=None, help_exists=False): + if argv is None: + argv = sys.argv[1:] + if '--help' in argv: + print('\nHelp:') + lines = str(self._config).split('\n')[2:] + print('\n'.join('--' + re.sub(r'[:,\[\]]', '', x) for x in lines)) + help_exists and sys.exit() + parsed = {} + remaining = [] + key = None + vals = None + for arg in argv: + if arg.startswith('--'): + if key: + self._submit_entry(key, vals, parsed, remaining) + if '=' in arg: + key, val = arg.split('=', 1) + vals = [val] + else: + key, vals = arg, [] + else: + if key: + vals.append(arg) + else: + remaining.append(arg) + self._submit_entry(key, vals, parsed, remaining) + parsed = self._config.update(parsed) + return parsed, remaining + + def _submit_entry(self, key, vals, parsed, remaining): + if not key and not vals: + return + if not key: + vals = ', '.join(f"'{x}'" for x in vals) + raise ValueError(f"Values {vals} were not preceded by any flag.") + name = key[len('--'):] + if '=' in name: + remaining.extend([key] + vals) + return + if self._config.IS_PATTERN.fullmatch(name): + pattern = re.compile(name) + keys = {k for k in self._config.flat if pattern.fullmatch(k)} + elif name in self._config: + keys = [name] + else: + keys = [] + if not keys: + remaining.extend([key] + vals) + return + if not vals: + raise ValueError(f"Flag '{key}' was not followed by any values.") + for key in keys: + parsed[key] = self._parse_flag_value(self._config[key], vals, key) + + def _parse_flag_value(self, default, value, key): + value = value if isinstance(value, (tuple, list)) else (value,) + if isinstance(default, (tuple, list)): + if len(value) == 1 and ',' in value[0]: + value = value[0].split(',') + return tuple(self._parse_flag_value(default[0], [x], key) for x in value) + assert len(value) == 1, value + value = str(value[0]) + if default is None: + return value + if isinstance(default, bool): + try: + return bool(['False', 'True'].index(value)) + except ValueError: + message = f"Expected bool but got '{value}' for key '{key}'." + raise TypeError(message) + if isinstance(default, int): + try: + value = float(value) # Allow scientific notation for integers. + assert float(int(value)) == value + except (TypeError, AssertionError): + message = f"Expected int but got float '{value}' for key '{key}'." + raise TypeError(message) + return int(value) + if isinstance(default, dict): + raise TypeError( + f"Key '{key}' refers to a whole dict. Please speicfy a subkey.") + return type(default)(value) diff --git a/dreamerv3/embodied/core/logger.py b/dreamerv3/embodied/core/logger.py new file mode 100644 index 0000000..2a2d068 --- /dev/null +++ b/dreamerv3/embodied/core/logger.py @@ -0,0 +1,323 @@ +import collections +import concurrent.futures +import datetime +import json +import os +import re +import time + +import numpy as np + +from . import path + + +class Logger: + + def __init__(self, step, outputs, multiplier=1): + assert outputs, 'Provide a list of logger outputs.' + self.step = step + self.outputs = outputs + self.multiplier = multiplier + self._last_step = None + self._last_time = None + self._metrics = [] + + def add(self, mapping, prefix=None): + step = int(self.step) * self.multiplier + for name, value in dict(mapping).items(): + name = f'{prefix}/{name}' if prefix else name + value = np.asarray(value) + if len(value.shape) not in (0, 1, 2, 3, 4): + raise ValueError( + f"Shape {value.shape} for name '{name}' cannot be " + "interpreted as scalar, histogram, image, or video.") + self._metrics.append((step, name, value)) + + def scalar(self, name, value): + self.add({name: value}) + + def image(self, name, value): + self.add({name: value}) + + def video(self, name, value): + self.add({name: value}) + + def write(self, fps=False): + if fps: + value = self._compute_fps() + if value is not None: + self.scalar('fps', value) + if not self._metrics: + return + for output in self.outputs: + output(tuple(self._metrics)) + self._metrics.clear() + + def _compute_fps(self): + step = int(self.step) * self.multiplier + if self._last_step is None: + self._last_time = time.time() + self._last_step = step + return None + steps = step - self._last_step + duration = time.time() - self._last_time + self._last_time += duration + self._last_step = step + return steps / duration + + +class AsyncOutput: + + def __init__(self, callback, parallel=True): + self._callback = callback + self._parallel = parallel + if parallel: + self._executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) + self._future = None + + def __call__(self, summaries): + if self._parallel: + self._future and self._future.result() + self._future = self._executor.submit(self._callback, summaries) + else: + self._callback(summaries) + + +class TerminalOutput: + + def __init__(self, pattern=r'.*', name=None): + self._pattern = re.compile(pattern) + self._name = name + try: + import rich.console + self._console = rich.console.Console() + except ImportError: + self._console = None + + def __call__(self, summaries): + step = max(s for s, _, _, in summaries) + scalars = {k: float(v) for _, k, v in summaries if len(v.shape) == 0} + scalars = {k: v for k, v in scalars.items() if self._pattern.search(k)} + formatted = {k: self._format_value(v) for k, v in scalars.items()} + if self._console: + if self._name: + self._console.rule(f'[green bold]{self._name} (Step {step})') + else: + self._console.rule(f'[green bold]Step {step}') + self._console.print(' [blue]/[/blue] '.join( + f'{k} {v}' for k, v in formatted.items())) + print('') + else: + message = ' / '.join(f'{k} {v}' for k, v in formatted.items()) + message = f'[{step}] {message}' + if self._name: + message = f'[{self._name}] {message}' + print(message, flush=True) + + def _format_value(self, value): + value = float(value) + if value == 0: + return '0' + elif 0.01 < abs(value) < 10000: + value = f'{value:.2f}' + value = value.rstrip('0') + value = value.rstrip('0') + value = value.rstrip('.') + return value + else: + value = f'{value:.1e}' + value = value.replace('.0e', 'e') + value = value.replace('+0', '') + value = value.replace('+', '') + value = value.replace('-0', '-') + return value + + +class JSONLOutput(AsyncOutput): + + def __init__( + self, logdir, filename='metrics.jsonl', pattern=r'.*', parallel=True): + super().__init__(self._write, parallel) + self._filename = filename + self._pattern = re.compile(pattern) + self._logdir = path.Path(logdir) + self._logdir.mkdirs() + + def _write(self, summaries): + bystep = collections.defaultdict(dict) + for step, name, value in summaries: + if len(value.shape) == 0 and self._pattern.search(name): + bystep[step][name] = float(value) + lines = ''.join([ + json.dumps({'step': step, **scalars}) + '\n' + for step, scalars in bystep.items()]) + with (self._logdir / self._filename).open('a') as f: + f.write(lines) + + +class TensorBoardOutput(AsyncOutput): + + def __init__(self, logdir, fps=20, maxsize=1e9, parallel=True): + super().__init__(self._write, parallel) + self._logdir = str(logdir) + if self._logdir.startswith('/gcs/'): + self._logdir = self._logdir.replace('/gcs/', 'gs://') + self._fps = fps + self._writer = None + self._maxsize = self._logdir.startswith('gs://') and maxsize + if self._maxsize: + self._checker = concurrent.futures.ThreadPoolExecutor(max_workers=1) + self._promise = None + + def _write(self, summaries): + import tensorflow as tf + reset = False + if self._maxsize: + result = self._promise and self._promise.result() + # print('Current TensorBoard event file size:', result) + reset = (self._promise and result >= self._maxsize) + self._promise = self._checker.submit(self._check) + if not self._writer or reset: + print('Creating new TensorBoard event file writer.') + self._writer = tf.summary.create_file_writer( + self._logdir, flush_millis=1000, max_queue=10000) + self._writer.set_as_default() + for step, name, value in summaries: + try: + if len(value.shape) == 0: + tf.summary.scalar(name, value, step) + elif len(value.shape) == 1: + if len(value) > 1024: + value = value.copy() + np.random.shuffle(value) + value = value[:1024] + tf.summary.histogram(name, value, step) + elif len(value.shape) == 2: + tf.summary.image(name, value, step) + elif len(value.shape) == 3: + tf.summary.image(name, value, step) + elif len(value.shape) == 4: + self._video_summary(name, value, step) + except Exception: + print('Error writing summary:', name) + raise + self._writer.flush() + + def _check(self): + import tensorflow as tf + events = tf.io.gfile.glob(self._logdir.rstrip('/') + '/events.out.*') + return tf.io.gfile.stat(sorted(events)[-1]).length if events else 0 + + def _video_summary(self, name, video, step): + import tensorflow as tf + import tensorflow.compat.v1 as tf1 + name = name if isinstance(name, str) else name.decode('utf-8') + if np.issubdtype(video.dtype, np.floating): + video = np.clip(255 * video, 0, 255).astype(np.uint8) + try: + T, H, W, C = video.shape + summary = tf1.Summary() + image = tf1.Summary.Image(height=H, width=W, colorspace=C) + image.encoded_image_string = _encode_gif(video, self._fps) + summary.value.add(tag=name, image=image) + tf.summary.experimental.write_raw_pb(summary.SerializeToString(), step) + except (IOError, OSError) as e: + print('GIF summaries require ffmpeg in $PATH.', e) + tf.summary.image(name, video, step) + + +class WandBOutput: + + def __init__(self, pattern, logdir, config): + self._pattern = re.compile(pattern) + import wandb + wandb.init( + project="dreamerv3", + name=logdir.name, + # sync_tensorboard=True,, + entity='word-bots', + config=dict(config), + ) + self._wandb = wandb + + def __call__(self, summaries): + bystep = collections.defaultdict(dict) + wandb = self._wandb + for step, name, value in summaries: + if len(value.shape) == 0 and self._pattern.search(name): + bystep[step][name] = float(value) + elif len(value.shape) == 1: + bystep[step][name] = wandb.Histogram(value) + elif len(value.shape) == 2: + value = np.clip(255 * value, 0, 255).astype(np.uint8) + value = np.transpose(value, [2, 0, 1]) + bystep[step][name] = wandb.Image(value) + elif len(value.shape) == 3: + value = np.clip(255 * value, 0, 255).astype(np.uint8) + value = np.transpose(value, [2, 0, 1]) + bystep[step][name] = wandb.Image(value) + elif len(value.shape) == 4: + # Sanity check that the channeld dimension is last + assert value.shape[3] in [1, 3, 4], f"Invalid shape: {value.shape}" + value = np.transpose(value, [0, 3, 1, 2]) + # If the video is a float, convert it to uint8 + if np.issubdtype(value.dtype, np.floating): + value = np.clip(255 * value, 0, 255).astype(np.uint8) + bystep[step][name] = wandb.Video(value) + + for step, metrics in bystep.items(): + self._wandb.log(metrics, step=step) + + +class MLFlowOutput: + + def __init__(self, run_name=None, resume_id=None, config=None, prefix=None): + import mlflow + self._mlflow = mlflow + self._prefix = prefix + self._setup(run_name, resume_id, config) + + def __call__(self, summaries): + bystep = collections.defaultdict(dict) + for step, name, value in summaries: + if len(value.shape) == 0 and self._pattern.search(name): + name = f'{self._prefix}/{name}' if self._prefix else name + bystep[step][name] = float(value) + for step, metrics in bystep.items(): + self._mlflow.log_metrics(metrics, step=step) + + def _setup(self, run_name, resume_id, config): + tracking_uri = os.environ.get('MLFLOW_TRACKING_URI', 'local') + run_name = run_name or os.environ.get('MLFLOW_RUN_NAME') + resume_id = resume_id or os.environ.get('MLFLOW_RESUME_ID') + print('MLFlow Tracking URI:', tracking_uri) + print('MLFlow Run Name: ', run_name) + print('MLFlow Resume ID: ', resume_id) + if resume_id: + runs = self._mlflow.search_runs(None, f'tags.resume_id="{resume_id}"') + assert len(runs), ('No runs to resume found.', resume_id) + self._mlflow.start_run(run_name=run_name, run_id=runs['run_id'].iloc[0]) + for key, value in config.items(): + self._mlflow.log_param(key, value) + else: + tags = {'resume_id': resume_id or ''} + self._mlflow.start_run(run_name=run_name, tags=tags) + + +def _encode_gif(frames, fps): + from subprocess import Popen, PIPE + h, w, c = frames[0].shape + pxfmt = {1: 'gray', 3: 'rgb24'}[c] + cmd = ' '.join([ + 'ffmpeg -y -f rawvideo -vcodec rawvideo', + f'-r {fps:.02f} -s {w}x{h} -pix_fmt {pxfmt} -i - -filter_complex', + '[0:v]split[x][z];[z]palettegen[y];[x]fifo[x];[x][y]paletteuse', + f'-r {fps:.02f} -f gif -']) + proc = Popen(cmd.split(' '), stdin=PIPE, stdout=PIPE, stderr=PIPE) + for image in frames: + proc.stdin.write(image.tobytes()) + out, err = proc.communicate() + if proc.returncode: + raise IOError('\n'.join([' '.join(cmd), err.decode('utf8')])) + del proc + return out diff --git a/dreamerv3/embodied/core/metrics.py b/dreamerv3/embodied/core/metrics.py new file mode 100644 index 0000000..db6e9bc --- /dev/null +++ b/dreamerv3/embodied/core/metrics.py @@ -0,0 +1,42 @@ +import collections +import warnings + +import numpy as np + + +class Metrics: + + def __init__(self): + self._scalars = collections.defaultdict(list) + self._lasts = {} + + def scalar(self, key, value): + self._scalars[key].append(value) + + def image(self, key, value): + self._lasts[key].append(value) + + def video(self, key, value): + self._lasts[key].append(value) + + def add(self, mapping, prefix=None): + for key, value in mapping.items(): + key = prefix + '/' + key if prefix else key + if hasattr(value, 'shape') and len(value.shape) > 0: + self._lasts[key] = value + else: + self._scalars[key].append(value) + + def result(self, reset=True): + result = {} + result.update(self._lasts) + with warnings.catch_warnings(): # Ignore empty slice warnings. + warnings.simplefilter('ignore', category=RuntimeWarning) + for key, values in self._scalars.items(): + result[key] = np.nanmean(values, dtype=np.float64) + reset and self.reset() + return result + + def reset(self): + self._scalars.clear() + self._lasts.clear() diff --git a/dreamerv3/embodied/core/parallel.py b/dreamerv3/embodied/core/parallel.py new file mode 100644 index 0000000..919b6bc --- /dev/null +++ b/dreamerv3/embodied/core/parallel.py @@ -0,0 +1,51 @@ +import enum +from functools import partial as bind + +from . import worker + + +class Parallel: + + def __init__(self, ctor, strategy): + self.worker = worker.Worker( + bind(self._respond, ctor), strategy, state=True) + self.callables = {} + + def __getattr__(self, name): + if name.startswith('_'): + raise AttributeError(name) + try: + if name not in self.callables: + self.callables[name] = self.worker(Message.CALLABLE, name)() + if self.callables[name]: + return bind(self.worker, Message.CALL, name) + else: + return self.worker(Message.READ, name)() + except AttributeError: + raise ValueError(name) + + def __len__(self): + return self.worker(Message.CALL, '__len__')() + + def close(self): + self.worker.close() + + @staticmethod + def _respond(ctor, state, message, name, *args, **kwargs): + state = state or ctor() + if message == Message.CALLABLE: + assert not args and not kwargs, (args, kwargs) + result = callable(getattr(state, name)) + elif message == Message.CALL: + result = getattr(state, name)(*args, **kwargs) + elif message == Message.READ: + assert not args and not kwargs, (args, kwargs) + result = getattr(state, name) + return state, result + + +class Message(enum.Enum): + + CALLABLE = 2 + CALL = 3 + READ = 4 diff --git a/dreamerv3/embodied/core/path.py b/dreamerv3/embodied/core/path.py new file mode 100644 index 0000000..82b431c --- /dev/null +++ b/dreamerv3/embodied/core/path.py @@ -0,0 +1,223 @@ +import contextlib +import glob +import os +import re +import shutil + + +class Path: + + filesystems = [] + + def __new__(cls, path): + path = str(path) + for impl, pred in cls.filesystems: + if pred(path): + obj = super().__new__(impl) + obj.__init__(path) + return obj + raise NotImplementedError(f'No filesystem supports: {path}') + + def __getnewargs__(self): + return (self._path,) + + def __init__(self, path): + assert isinstance(path, str) + path = re.sub(r'^\./*', '', path) # Remove leading dot or dot slashes. + path = re.sub(r'(?<=[^/])/$', '', path) # Remove single trailing slash. + path = path or '.' # Empty path is represented by a dot. + self._path = path + + def __truediv__(self, part): + sep = '' if self._path.endswith('/') else '/' + return type(self)(f'{self._path}{sep}{str(part)}') + + def __repr__(self): + return f'Path({str(self)})' + + def __fspath__(self): + return str(self) + + def __eq__(self, other): + return self._path == other._path + + def __lt__(self, other): + return self._path < other._path + + def __str__(self): + return self._path + + @property + def parent(self): + if '/' not in self._path: + return type(self)('.') + parent = self._path.rsplit('/', 1)[0] + parent = parent or ('/' if self._path.startswith('/') else '.') + return type(self)(parent) + + @property + def name(self): + if '/' not in self._path: + return self._path + return self._path.rsplit('/', 1)[1] + + @property + def stem(self): + return self.name.split('.', 1)[0] if '.' in self.name else self.name + + @property + def suffix(self): + return ('.' + self.name.split('.', 1)[1]) if '.' in self.name else '' + + def read(self, mode='r'): + assert mode in 'r rb'.split(), mode + with self.open(mode) as f: + return f.read() + + def write(self, content, mode='w'): + assert mode in 'w a wb ab'.split(), mode + with self.open(mode) as f: + f.write(content) + + @contextlib.contextmanager + def open(self, mode='r'): + raise NotImplementedError + + def absolute(self): + raise NotImplementedError + + def glob(self, pattern): + raise NotImplementedError + + def exists(self): + raise NotImplementedError + + def isfile(self): + raise NotImplementedError + + def isdir(self): + raise NotImplementedError + + def mkdirs(self): + raise NotImplementedError + + def remove(self): + raise NotImplementedError + + def rmtree(self): + raise NotImplementedError + + def copy(self, dest): + raise NotImplementedError + + def move(self, dest): + self.copy(dest) + self.remove() + + +class LocalPath(Path): + + def __init__(self, path): + super().__init__(os.path.expanduser(str(path))) + + @contextlib.contextmanager + def open(self, mode='r'): + with open(str(self), mode=mode) as f: + yield f + + def absolute(self): + return type(self)(os.path.absolute(str(self))) + + def glob(self, pattern): + for path in glob.glob(f'{str(self)}/{pattern}'): + yield type(self)(path) + + def exists(self): + return os.path.exists(str(self)) + + def isfile(self): + return os.path.isfile(str(self)) + + def isdir(self): + return os.path.isdir(str(self)) + + def mkdirs(self): + os.makedirs(str(self), exist_ok=True) + + def remove(self): + os.rmdir(str(self)) if self.isdir() else os.remove(str(self)) + + def rmtree(self): + shutil.rmtree(self) + + def copy(self, dest): + if self.isfile(): + shutil.copy(self, type(self)(dest)) + else: + shutil.copytree(self, type(self)(dest), dirs_exist_ok=True) + + def move(self, dest): + shutil.move(self, dest) + + +class GFilePath(Path): + + def __init__(self, path): + path = str(path) + if not (path.startswith('/') or '://' in path): + path = os.path.abspath(os.path.expanduser(path)) + super().__init__(path) + import tensorflow as tf + self._gfile = tf.io.gfile + + @contextlib.contextmanager + def open(self, mode='r'): + path = str(self) + if 'a' in mode and path.startswith('/cns/'): + path += '%r=3.2' + if mode.startswith('x') and self.exists(): + raise FileExistsError(path) + mode = mode.replace('x', 'w') + with self._gfile.GFile(path, mode) as f: + yield f + + def absolute(self): + return self + + def glob(self, pattern): + for path in self._gfile.glob(f'{str(self)}/{pattern}'): + yield type(self)(path) + + def exists(self): + return self._gfile.exists(str(self)) + + def isfile(self): + return self.exists() and not self.isdir() + + def isdir(self): + return self._gfile.isdir(str(self)) + + def mkdirs(self): + self._gfile.makedirs(str(self)) + + def remove(self): + self._gfile.remove(str(self)) + + def rmtree(self): + self._gfile.rmtree(str(self)) + + def copy(self, dest): + self._gfile.copy(str(self), str(dest), overwrite=True) + + def move(self, dest): + dest = Path(dest) + if dest.isdir(): + dest.rmtree() + self._gfile.rename(self, str(dest), overwrite=True) + + +Path.filesystems = [ + (GFilePath, lambda path: path.startswith('gs://')), + (GFilePath, lambda path: path.startswith('/cns/')), + (LocalPath, lambda path: True), +] diff --git a/dreamerv3/embodied/core/random.py b/dreamerv3/embodied/core/random.py new file mode 100644 index 0000000..2e59f9d --- /dev/null +++ b/dreamerv3/embodied/core/random.py @@ -0,0 +1,14 @@ +import numpy as np + + +class RandomAgent: + + def __init__(self, act_space): + self.act_space = act_space + + def policy(self, obs, state=None, mode='train'): + batch_size = len(next(iter(obs.values()))) + act = { + k: np.stack([v.sample() for _ in range(batch_size)]) + for k, v in self.act_space.items() if k != 'reset'} + return act, state diff --git a/dreamerv3/embodied/core/space.py b/dreamerv3/embodied/core/space.py new file mode 100644 index 0000000..9447cd7 --- /dev/null +++ b/dreamerv3/embodied/core/space.py @@ -0,0 +1,103 @@ +import numpy as np + + +class Space: + + def __init__(self, dtype, shape=(), low=None, high=None): + # For integer types, high is the excluside upper bound. + shape = (shape,) if isinstance(shape, int) else shape + self._dtype = np.dtype(dtype) + assert self._dtype is not object, self._dtype + assert isinstance(shape, tuple), shape + self._low = self._infer_low(dtype, shape, low, high) + self._high = self._infer_high(dtype, shape, low, high) + self._shape = self._infer_shape(dtype, shape, low, high) + self._discrete = ( + np.issubdtype(self.dtype, np.integer) or self.dtype == bool) + self._random = np.random.RandomState() + + @property + def dtype(self): + return self._dtype + + @property + def shape(self): + return self._shape + + @property + def low(self): + return self._low + + @property + def high(self): + return self._high + + @property + def discrete(self): + return self._discrete + + def __repr__(self): + return ( + f'Space(dtype={self.dtype.name}, ' + f'shape={self.shape}, ' + f'low={self.low.min()}, ' + f'high={self.high.max()})') + + def __contains__(self, value): + value = np.asarray(value) + if value.shape != self.shape: + return False + if (value > self.high).any(): + return False + if (value < self.low).any(): + return False + if (value.astype(self.dtype).astype(value.dtype) != value).any(): + return False + return True + + def sample(self): + low, high = self.low, self.high + if np.issubdtype(self.dtype, np.floating): + low = np.maximum(np.ones(self.shape) * np.finfo(self.dtype).min, low) + high = np.minimum(np.ones(self.shape) * np.finfo(self.dtype).max, high) + return self._random.uniform(low, high, self.shape).astype(self.dtype) + + def _infer_low(self, dtype, shape, low, high): + if low is not None: + try: + return np.broadcast_to(low, shape) + except ValueError: + raise ValueError(f'Cannot broadcast {low} to shape {shape}') + elif np.issubdtype(dtype, np.floating): + return -np.inf * np.ones(shape) + elif np.issubdtype(dtype, np.integer): + return np.iinfo(dtype).min * np.ones(shape, dtype) + elif np.issubdtype(dtype, bool): + return np.zeros(shape, bool) + else: + raise ValueError('Cannot infer low bound from shape and dtype.') + + def _infer_high(self, dtype, shape, low, high): + if high is not None: + try: + return np.broadcast_to(high, shape) + except ValueError: + raise ValueError(f'Cannot broadcast {high} to shape {shape}') + elif np.issubdtype(dtype, np.floating): + return np.inf * np.ones(shape) + elif np.issubdtype(dtype, np.integer): + return np.iinfo(dtype).max * np.ones(shape, dtype) + elif np.issubdtype(dtype, bool): + return np.ones(shape, bool) + else: + raise ValueError('Cannot infer high bound from shape and dtype.') + + def _infer_shape(self, dtype, shape, low, high): + if shape is None and low is not None: + shape = low.shape + if shape is None and high is not None: + shape = high.shape + if not hasattr(shape, '__len__'): + shape = (shape,) + assert all(dim and dim > 0 for dim in shape), shape + return tuple(shape) diff --git a/dreamerv3/embodied/core/timer.py b/dreamerv3/embodied/core/timer.py new file mode 100644 index 0000000..5858987 --- /dev/null +++ b/dreamerv3/embodied/core/timer.py @@ -0,0 +1,61 @@ +import collections +import contextlib +import time + +import numpy as np + + +class Timer: + + def __init__(self, columns=('frac', 'min', 'avg', 'max', 'count', 'total')): + available = ('frac', 'avg', 'min', 'max', 'count', 'total') + assert all(x in available for x in columns), columns + self._columns = columns + self._durations = collections.defaultdict(list) + self._start = time.time() + + def reset(self): + for timings in self._durations.values(): + timings.clear() + self._start = time.time() + + @contextlib.contextmanager + def scope(self, name): + start = time.time() + yield + stop = time.time() + self._durations[name].append(stop - start) + + def wrap(self, name, obj, methods): + for method in methods: + decorator = self.scope(f'{name}.{method}') + setattr(obj, method, decorator(getattr(obj, method))) + + def stats(self, reset=True, log=False): + metrics = {} + metrics['duration'] = time.time() - self._start + for name, durs in self._durations.items(): + available = {} + available['count'] = len(durs) + available['total'] = np.sum(durs) + available['frac'] = np.sum(durs) / metrics['duration'] + if len(durs): + available['avg'] = np.mean(durs) + available['min'] = np.min(durs) + available['max'] = np.max(durs) + for key, value in available.items(): + if key in self._columns: + metrics[f'{name}_{key}'] = value + if log: + self._log(metrics) + if reset: + self.reset() + return metrics + + def _log(self, metrics): + names = self._durations.keys() + names = sorted(names, key=lambda k: -metrics[f'{k}_frac']) + print('Timer:'.ljust(20), ' '.join(x.rjust(8) for x in self._columns)) + for name in names: + values = [metrics[f'{name}_{col}'] for col in self._columns] + print(f'{name.ljust(20)}', ' '.join((f'{x:8.4f}' for x in values))) diff --git a/dreamerv3/embodied/core/uuid.py b/dreamerv3/embodied/core/uuid.py new file mode 100644 index 0000000..cf518f4 --- /dev/null +++ b/dreamerv3/embodied/core/uuid.py @@ -0,0 +1,74 @@ +import string +import uuid as uuidlib + +import numpy as np + + +class uuid: + """UUID that is stored as 16 byte string and can be converted to and from + int, string, and array types.""" + + DEBUG_ID = None + BASE62 = string.digits + string.ascii_letters + BASE62REV = {x: i for i, x in enumerate(BASE62)} + + @classmethod + def reset(cls, *, debug): + cls.DEBUG_ID = 0 if debug else None + + def __init__(self, value=None): + if value is None: + if self.DEBUG_ID is None: + self.value = uuidlib.uuid4().bytes + else: + type(self).DEBUG_ID += 1 + self.value = self.DEBUG_ID.to_bytes(16, 'big') + elif isinstance(value, uuid): + self.value = value.value + elif isinstance(value, int): + self.value = value.to_bytes(16, 'big') + elif isinstance(value, str): + if self.DEBUG_ID is None: + integer = 0 + for index, char in enumerate(value[::-1]): + integer += (62 ** index) * self.BASE62REV[char] + self.value = integer.to_bytes(16, 'big') + else: + self.value = int(value).to_bytes(16, 'big') + elif isinstance(value, np.ndarray): + self.value = value.tobytes() + else: + raise ValueError(value) + assert type(self.value) == bytes, type(self.value) + assert len(self.value) == 16, len(self.value) + self._hash = hash(self.value) + + def __int__(self): + return int.from_bytes(self.value, 'big') + + def __str__(self): + if self.DEBUG_ID is not None: + return str(int(self)) + chars = [] + integer = int(self) + while integer != 0: + chars.append(self.BASE62[integer % 62]) + integer //= 62 + while len(chars) < 22: + chars.append('0') + return ''.join(chars[::-1]) + + def __array__(self): + return np.frombuffer(self.value, np.uint8) + + def __getitem__(self, index): + return self.__array__()[index] + + def __repr__(self): + return str(self) + + def __eq__(self, other): + return self.value == other.value + + def __hash__(self): + return self._hash diff --git a/dreamerv3/embodied/core/when.py b/dreamerv3/embodied/core/when.py new file mode 100644 index 0000000..2604eda --- /dev/null +++ b/dreamerv3/embodied/core/when.py @@ -0,0 +1,88 @@ +import time + + +class Every: + + def __init__(self, every, initial=True): + self._every = every + self._initial = initial + self._prev = None + + def __call__(self, step): + step = int(step) + if self._every < 0: + return True + if self._every == 0: + return False + if self._prev is None: + self._prev = (step // self._every) * self._every + return self._initial + if step >= self._prev + self._every: + self._prev += self._every + return True + return False + + +class Ratio: + + def __init__(self, ratio): + assert ratio >= 0, ratio + self._ratio = ratio + self._prev = None + + def __call__(self, step): + step = int(step) + if self._ratio == 0: + return 0 + if self._prev is None: + self._prev = step + return 1 + repeats = int((step - self._prev) * self._ratio) + self._prev += repeats / self._ratio + return repeats + + +class Once: + + def __init__(self): + self._once = True + + def __call__(self): + if self._once: + self._once = False + return True + return False + + +class Until: + + def __init__(self, until): + self._until = until + + def __call__(self, step): + step = int(step) + if not self._until: + return True + return step < self._until + + +class Clock: + + def __init__(self, every): + self._every = every + self._prev = None + + def __call__(self, step=None): + if self._every < 0: + return True + if self._every == 0: + return False + now = time.time() + if self._prev is None: + self._prev = now + return True + if now >= self._prev + self._every: + # self._prev += self._every + self._prev = now + return True + return False diff --git a/dreamerv3/embodied/core/worker.py b/dreamerv3/embodied/core/worker.py new file mode 100644 index 0000000..f3ec90c --- /dev/null +++ b/dreamerv3/embodied/core/worker.py @@ -0,0 +1,241 @@ +import atexit +import concurrent.futures +import enum +import os +import sys +import time +import traceback +from functools import partial as bind + + +class Worker: + + initializers = [] + + def __init__(self, fn, strategy='thread', state=False): + if not state: + fn = lambda s, *args, fn=fn, **kwargs: (s, fn(*args, **kwargs)) + inits = self.initializers + self.impl = { + 'blocking': BlockingWorker, + 'thread': ThreadWorker, + 'process': bind(ProcessPipeWorker, initializers=inits), + 'daemon': bind(ProcessPipeWorker, initializers=inits, daemon=True), + 'process_slow': bind(ProcessWorker, initializers=inits), + }[strategy](fn) + self.promise = None + + def __call__(self, *args, **kwargs): + self.promise and self.promise() # Raise previous exception if any. + self.promise = self.impl(*args, **kwargs) + return self.promise + + def wait(self): + return self.impl.wait() + + def close(self): + self.impl.close() + + +class BlockingWorker: + + def __init__(self, fn): + self.fn = fn + self.state = None + + def __call__(self, *args, **kwargs): + self.state, result = self.fn(self.state, *args, **kwargs) + # return lambda: result + return lambda result=result: result + + def wait(self): + pass + + def close(self): + pass + + +class ThreadWorker: + + def __init__(self, fn): + self.fn = fn + self.state = None + self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=1) + self.futures = [] + + def __call__(self, *args, **kwargs): + future = self.executor.submit(self._worker, *args, **kwargs) + self.futures.append(future) + future.add_done_callback(lambda f: self.futures.remove(f)) + return future.result + + def wait(self): + concurrent.futures.wait(self.futures) + + def close(self): + self.executor.shutdown(wait=False, cancel_futures=True) + + def _worker(self, *args, **kwargs): + self.state, output = self.fn(self.state, *args, **kwargs) + return output + + +class ProcessWorker: + + def __init__(self, fn, initializers=()): + import cloudpickle + import multiprocessing + fn = cloudpickle.dumps(fn) + initializers = cloudpickle.dumps(initializers) + self.executor = concurrent.futures.ProcessPoolExecutor( + max_workers=1, mp_context=multiprocessing.get_context('spawn'), + initializer=self._initializer, initargs=(fn, initializers)) + self.futures = [] + + def __call__(self, *args, **kwargs): + future = self.executor.submit(self._worker, *args, **kwargs) + self.futures.append(future) + future.add_done_callback(lambda f: self.futures.remove(f)) + return future.result + + def wait(self): + concurrent.futures.wait(self.futures) + + def close(self): + self.executor.shutdown(wait=False, cancel_futures=True) + + @staticmethod + def _initializer(fn, initializers): + global _FN, _STATE + import cloudpickle + _FN = cloudpickle.loads(fn) + _STATE = None + for initializer in cloudpickle.loads(initializers): + initializers() + + @staticmethod + def _worker(*args, **kwargs): + global _FN, _STATE + _STATE, output = _FN(_STATE, *args, **kwargs) + return output + + +class ProcessPipeWorker: + + def __init__(self, fn, initializers=(), daemon=False): + import multiprocessing + import cloudpickle + self._context = multiprocessing.get_context('spawn') + self._pipe, pipe = self._context.Pipe() + fn = cloudpickle.dumps(fn) + initializers = cloudpickle.dumps(initializers) + self._process = self._context.Process( + target=self._loop, + args=(pipe, fn, initializers), + daemon=daemon) + self._process.start() + self._nextid = 0 + self._results = {} + assert self._submit(Message.OK)() + atexit.register(self.close) + + def __call__(self, *args, **kwargs): + return self._submit(Message.RUN, (args, kwargs)) + + def wait(self): + pass + + def close(self): + try: + self._pipe.send((Message.STOP, self._nextid, None)) + self._pipe.close() + except (AttributeError, IOError): + pass # The connection was already closed. + try: + self._process.join(0.1) + if self._process.exitcode is None: + try: + os.kill(self._process.pid, 9) + time.sleep(0.1) + except Exception: + pass + except (AttributeError, AssertionError): + pass + + def _submit(self, message, payload=None): + callid = self._nextid + self._nextid += 1 + self._pipe.send((message, callid, payload)) + return Future(self._receive, callid) + + def _receive(self, callid): + while callid not in self._results: + try: + message, callid, payload = self._pipe.recv() + except (OSError, EOFError): + raise RuntimeError('Lost connection to worker.') + if message == Message.ERROR: + raise Exception(payload) + assert message == Message.RESULT, message + self._results[callid] = payload + return self._results.pop(callid) + + @staticmethod + def _loop(pipe, function, initializers): + try: + callid = None + state = None + import cloudpickle + initializers = cloudpickle.loads(initializers) + function = cloudpickle.loads(function) + [fn() for fn in initializers] + while True: + if not pipe.poll(0.1): + continue # Wake up for keyboard interrupts. + message, callid, payload = pipe.recv() + if message == Message.OK: + pipe.send((Message.RESULT, callid, True)) + elif message == Message.STOP: + return + elif message == Message.RUN: + args, kwargs = payload + state, result = function(state, *args, **kwargs) + pipe.send((Message.RESULT, callid, result)) + else: + raise KeyError(f'Invalid message: {message}') + except (EOFError, KeyboardInterrupt): + return + except Exception: + stacktrace = ''.join(traceback.format_exception(*sys.exc_info())) + print(f'Error inside process worker: {stacktrace}.', flush=True) + pipe.send((Message.ERROR, callid, stacktrace)) + return + finally: + try: + pipe.close() + except Exception: + pass + + +class Future: + + def __init__(self, receive, callid): + self._receive = receive + self._callid = callid + self._result = None + self._complete = False + + def __call__(self): + if not self._complete: + self._result = self._receive(self._callid) + self._complete = True + return self._result + + +class Message(enum.Enum): + + OK = 1 + RUN = 2 + RESULT = 3 + STOP = 4 + ERROR = 5 diff --git a/dreamerv3/embodied/core/wrappers.py b/dreamerv3/embodied/core/wrappers.py new file mode 100644 index 0000000..e1c7e46 --- /dev/null +++ b/dreamerv3/embodied/core/wrappers.py @@ -0,0 +1,364 @@ +import functools +import time + +import numpy as np + +from . import base +from . import space as spacelib + + +class TimeLimit(base.Wrapper): + + def __init__(self, env, duration, reset=True): + super().__init__(env) + self._duration = duration + self._reset = reset + self._step = 0 + self._done = False + + def step(self, action): + if action['reset'] or self._done: + self._step = 0 + self._done = False + if self._reset: + action.update(reset=True) + return self.env.step(action) + else: + action.update(reset=False) + obs = self.env.step(action) + obs['is_first'] = True + return obs + self._step += 1 + obs = self.env.step(action) + if self._duration and self._step >= self._duration: + obs['is_last'] = True + self._done = obs['is_last'] + return obs + + +class ActionRepeat(base.Wrapper): + + def __init__(self, env, repeat): + super().__init__(env) + self._repeat = repeat + self._done = False + + def step(self, action): + if action['reset'] or self._done: + return self.env.step(action) + reward = 0.0 + for _ in range(self._repeat): + obs = self.env.step(action) + reward += obs['reward'] + if obs['is_last'] or obs['is_terminal']: + break + obs['reward'] = np.float32(reward) + self._done = obs['is_last'] + return obs + + +class ClipAction(base.Wrapper): + + def __init__(self, env, key='action', low=-1, high=1): + super().__init__(env) + self._key = key + self._low = low + self._high = high + + def step(self, action): + clipped = np.clip(action[self._key], self._low, self._high) + return self.env.step({**action, self._key: clipped}) + + +class NormalizeAction(base.Wrapper): + + def __init__(self, env, key='action'): + super().__init__(env) + self._key = key + self._space = env.act_space[key] + self._mask = np.isfinite(self._space.low) & np.isfinite(self._space.high) + self._low = np.where(self._mask, self._space.low, -1) + self._high = np.where(self._mask, self._space.high, 1) + + @functools.cached_property + def act_space(self): + low = np.where(self._mask, -np.ones_like(self._low), self._low) + high = np.where(self._mask, np.ones_like(self._low), self._high) + space = spacelib.Space(np.float32, self._space.shape, low, high) + return {**self.env.act_space, self._key: space} + + def step(self, action): + orig = (action[self._key] + 1) / 2 * (self._high - self._low) + self._low + orig = np.where(self._mask, orig, action[self._key]) + return self.env.step({**action, self._key: orig}) + + +class OneHotAction(base.Wrapper): + + def __init__(self, env, key='action'): + super().__init__(env) + self._count = int(env.act_space[key].high) + self._key = key + + @functools.cached_property + def act_space(self): + shape = (self._count,) + space = spacelib.Space(np.float32, shape, 0, 1) + space.sample = functools.partial(self._sample_action, self._count) + space._discrete = True + return {**self.env.act_space, self._key: space} + + def step(self, action): + if not action['reset']: + assert action[self._key].min() == 0.0, action + assert action[self._key].max() == 1.0, action + assert action[self._key].sum() == 1.0, action + index = np.argmax(action[self._key]) + return self.env.step({**action, self._key: index}) + + @staticmethod + def _sample_action(count): + index = np.random.randint(0, count) + action = np.zeros(count, dtype=np.float32) + action[index] = 1.0 + return action + + +class ExpandScalars(base.Wrapper): + + def __init__(self, env): + super().__init__(env) + self._obs_expanded = [] + self._obs_space = {} + for key, space in self.env.obs_space.items(): + if space.shape == () and key != 'reward' and not space.discrete: + space = spacelib.Space(space.dtype, (1,), space.low, space.high) + self._obs_expanded.append(key) + self._obs_space[key] = space + self._act_expanded = [] + self._act_space = {} + for key, space in self.env.act_space.items(): + if space.shape == () and not space.discrete: + space = spacelib.Space(space.dtype, (1,), space.low, space.high) + self._act_expanded.append(key) + self._act_space[key] = space + + @functools.cached_property + def obs_space(self): + return self._obs_space + + @functools.cached_property + def act_space(self): + return self._act_space + + def step(self, action): + action = { + key: np.squeeze(value, 0) if key in self._act_expanded else value + for key, value in action.items()} + obs = self.env.step(action) + obs = { + key: np.expand_dims(value, 0) if key in self._obs_expanded else value + for key, value in obs.items()} + return obs + + +class FlattenTwoDimObs(base.Wrapper): + + def __init__(self, env): + super().__init__(env) + self._keys = [] + self._obs_space = {} + for key, space in self.env.obs_space.items(): + if len(space.shape) == 2: + space = spacelib.Space( + space.dtype, + (int(np.prod(space.shape)),), + space.low.flatten(), + space.high.flatten()) + self._keys.append(key) + self._obs_space[key] = space + + @functools.cached_property + def obs_space(self): + return self._obs_space + + def step(self, action): + obs = self.env.step(action).copy() + for key in self._keys: + obs[key] = obs[key].flatten() + return obs + + +class FlattenTwoDimActions(base.Wrapper): + + def __init__(self, env): + super().__init__(env) + self._origs = {} + self._act_space = {} + for key, space in self.env.act_space.items(): + if len(space.shape) == 2: + space = spacelib.Space( + space.dtype, + (int(np.prod(space.shape)),), + space.low.flatten(), + space.high.flatten()) + self._origs[key] = space.shape + self._act_space[key] = space + + @functools.cached_property + def act_space(self): + return self._act_space + + def step(self, action): + action = action.copy() + for key, shape in self._origs.items(): + action[key] = action[key].reshape(shape) + return self.env.step(action) + + +class CheckSpaces(base.Wrapper): + + def __init__(self, env): + super().__init__(env) + + def step(self, action): + for key, value in action.items(): + self._check(value, self.env.act_space[key], key) + obs = self.env.step(action) + for key, value in obs.items(): + self._check(value, self.env.obs_space[key], key) + return obs + + def _check(self, value, space, key): + if not isinstance(value, ( + np.ndarray, np.generic, list, tuple, int, float, bool)): + raise TypeError(f'Invalid type {type(value)} for key {key}.') + if value in space: + return + dtype = np.array(value).dtype + shape = np.array(value).shape + lowest, highest = np.min(value), np.max(value) + raise ValueError( + f"Value for '{key}' with dtype {dtype}, shape {shape}, " + f"lowest {lowest}, highest {highest} is not in {space}.") + + +class DiscretizeAction(base.Wrapper): + + def __init__(self, env, key='action', bins=5): + super().__init__(env) + self._dims = np.squeeze(env.act_space[key].shape, 0).item() + self._values = np.linspace(-1, 1, bins) + self._key = key + + @functools.cached_property + def act_space(self): + shape = (self._dims, len(self._values)) + space = spacelib.Space(np.float32, shape, 0, 1) + space.sample = functools.partial( + self._sample_action, self._dims, self._values) + space._discrete = True + return {**self.env.act_space, self._key: space} + + def step(self, action): + if not action['reset']: + assert (action[self._key].min(-1) == 0.0).all(), action + assert (action[self._key].max(-1) == 1.0).all(), action + assert (action[self._key].sum(-1) == 1.0).all(), action + indices = np.argmax(action[self._key], axis=-1) + continuous = np.take(self._values, indices) + return self.env.step({**action, self._key: continuous}) + + @staticmethod + def _sample_action(dims, values): + indices = np.random.randint(0, len(values), dims) + action = np.zeros((dims, len(values)), dtype=np.float32) + action[np.arange(dims), indices] = 1.0 + return action + + +class ResizeImage(base.Wrapper): + + def __init__(self, env, size=(64, 64)): + super().__init__(env) + self._size = size + self._keys = [ + k for k, v in env.obs_space.items() + if len(v.shape) > 1 and v.shape[:2] != size] + print(f'Resizing keys {",".join(self._keys)} to {self._size}.') + if self._keys: + from PIL import Image + self._Image = Image + + @functools.cached_property + def obs_space(self): + spaces = self.env.obs_space + for key in self._keys: + shape = self._size + spaces[key].shape[2:] + spaces[key] = spacelib.Space(np.uint8, shape) + return spaces + + def step(self, action): + obs = self.env.step(action) + for key in self._keys: + obs[key] = self._resize(obs[key]) + return obs + + def _resize(self, image): + image = self._Image.fromarray(image) + image = image.resize(self._size, self._Image.NEAREST) + image = np.array(image) + return image + + +class RenderImage(base.Wrapper): + + def __init__(self, env, key='image'): + super().__init__(env) + self._key = key + self._shape = self.env.render().shape + + @functools.cached_property + def obs_space(self): + spaces = self.env.obs_space + spaces[self._key] = spacelib.Space(np.uint8, self._shape) + return spaces + + def step(self, action): + obs = self.env.step(action) + obs[self._key] = self.env.render() + return obs + + +class RestartOnException(base.Wrapper): + + def __init__( + self, ctor, exceptions=(Exception,), window=300, maxfails=2, wait=20): + if not isinstance(exceptions, (tuple, list)): + exceptions = [exceptions] + self._ctor = ctor + self._exceptions = tuple(exceptions) + self._window = window + self._maxfails = maxfails + self._wait = wait + self._last = time.time() + self._fails = 0 + super().__init__(self._ctor()) + + def step(self, action): + try: + return self.env.step(action) + except self._exceptions as e: + if time.time() > self._last + self._window: + self._last = time.time() + self._fails = 1 + else: + self._fails += 1 + if self._fails > self._maxfails: + raise RuntimeError('The env crashed too many times.') + message = f'Restarting env after crash with {type(e).__name__}: {e}' + print(message, flush=True) + time.sleep(self._wait) + self.env = self._ctor() + action['reset'] = np.ones_like(action['reset']) + return self.env.step(action) diff --git a/dreamerv3/embodied/envs/atari.py b/dreamerv3/embodied/envs/atari.py new file mode 100644 index 0000000..cd085a4 --- /dev/null +++ b/dreamerv3/embodied/envs/atari.py @@ -0,0 +1,140 @@ +import embodied +import numpy as np + + +class Atari(embodied.Env): + + LOCK = None + + def __init__( + self, name, repeat=4, size=(84, 84), gray=True, noops=0, lives='unused', + sticky=True, actions='all', length=108000, resize='opencv', seed=None): + assert size[0] == size[1] + assert lives in ('unused', 'discount', 'reset'), lives + assert actions in ('all', 'needed'), actions + assert resize in ('opencv', 'pillow'), resize + if self.LOCK is None: + import multiprocessing as mp + mp = mp.get_context('spawn') + self.LOCK = mp.Lock() + self._resize = resize + if self._resize == 'opencv': + import cv2 + self._cv2 = cv2 + if self._resize == 'pillow': + from PIL import Image + self._image = Image + import gym.envs.atari + if name == 'james_bond': + name = 'jamesbond' + self._repeat = repeat + self._size = size + self._gray = gray + self._noops = noops + self._lives = lives + self._sticky = sticky + self._length = length + self._random = np.random.RandomState(seed) + with self.LOCK: + self._env = gym.envs.atari.AtariEnv( + game=name, + obs_type='image', + frameskip=1, repeat_action_probability=0.25 if sticky else 0.0, + full_action_space=(actions == 'all')) + assert self._env.unwrapped.get_action_meanings()[0] == 'NOOP' + shape = self._env.observation_space.shape + self._buffer = [np.zeros(shape, np.uint8) for _ in range(2)] + self._ale = self._env.unwrapped.ale + self._last_lives = None + self._done = True + self._step = 0 + + @property + def obs_space(self): + shape = self._size + (1 if self._gray else 3,) + return { + 'image': embodied.Space(np.uint8, shape), + 'reward': embodied.Space(np.float32), + 'is_first': embodied.Space(bool), + 'is_last': embodied.Space(bool), + 'is_terminal': embodied.Space(bool), + } + + @property + def act_space(self): + return { + 'action': embodied.Space(np.int32, (), 0, self._env.action_space.n), + 'reset': embodied.Space(bool), + } + + def step(self, action): + if action['reset'] or self._done: + with self.LOCK: + self._reset() + self._done = False + self._step = 0 + return self._obs(0.0, is_first=True) + total = 0.0 + dead = False + for repeat in range(self._repeat): + _, reward, over, info = self._env.step(action['action']) + self._step += 1 + total += reward + if repeat == self._repeat - 2: + self._screen(self._buffer[1]) + if over: + break + if self._lives != 'unused': + current = self._ale.lives() + if current < self._last_lives: + dead = True + self._last_lives = current + break + if not self._repeat: + self._buffer[1][:] = self._buffer[0][:] + self._screen(self._buffer[0]) + self._done = over or (self._length and self._step >= self._length) + return self._obs( + total, + is_last=self._done or (dead and self._lives == 'reset'), + is_terminal=dead or over) + + def _reset(self): + self._env.reset() + if self._noops: + for _ in range(self._random.randint(self._noops)): + _, _, dead, _ = self._env.step(0) + if dead: + self._env.reset() + self._last_lives = self._ale.lives() + self._screen(self._buffer[0]) + self._buffer[1].fill(0) + + def _obs(self, reward, is_first=False, is_last=False, is_terminal=False): + np.maximum(self._buffer[0], self._buffer[1], out=self._buffer[0]) + image = self._buffer[0] + if image.shape[:2] != self._size: + if self._resize == 'opencv': + image = self._cv2.resize( + image, self._size, interpolation=self._cv2.INTER_AREA) + if self._resize == 'pillow': + image = self._image.fromarray(image) + image = image.resize(self._size, self._image.NEAREST) + image = np.array(image) + if self._gray: + weights = [0.299, 0.587, 1 - (0.299 + 0.587)] + image = np.tensordot(image, weights, (-1, 0)).astype(image.dtype) + image = image[:, :, None] + return dict( + image=image, + reward=reward, + is_first=is_first, + is_last=is_last, + is_terminal=is_last, + ) + + def _screen(self, array): + self._ale.getScreenRGB2(array) + + def close(self): + return self._env.close() diff --git a/dreamerv3/embodied/envs/cdmc.py b/dreamerv3/embodied/envs/cdmc.py new file mode 100644 index 0000000..ab004b8 --- /dev/null +++ b/dreamerv3/embodied/envs/cdmc.py @@ -0,0 +1,64 @@ +import functools +import os + +import embodied +import numpy as np + + +class CDMC(embodied.Env): + + DEFAULT_CAMERAS = dict( + locom_rodent=1, + quadruped=2, + ) + + def __init__(self, env, repeat=1, render=True, size=(64, 64), camera=-1, + unconstrain_at_step=5e5): + + print("CDMC repeat", repeat) + unconstrain_at_step = unconstrain_at_step // repeat + + # This env variable is meant for headless GPU machines but may fail on CPU-only machines. + if 'MUJOCO_GL' not in os.environ: + os.environ['MUJOCO_GL'] = 'egl' + if isinstance(env, str): + domain, task = env.split('_', 1) + if camera == -1: + camera = self.DEFAULT_CAMERAS.get(domain, 0) + if domain == 'cup': # Only domain with multiple words. + domain = 'ball_in_cup' + + from adaptgym.envs.cdmc import suite + self._dmenv = suite.load(domain, task) + self._dmenv.task._unconstrain_at_step = unconstrain_at_step + + from . import from_dm + self._env = from_dm.FromDM(self._dmenv) + self._env = embodied.wrappers.ExpandScalars(self._env) + self._env = embodied.wrappers.ActionRepeat(self._env, repeat) + self._render = render + self._size = size + self._camera = camera + + @functools.cached_property + def obs_space(self): + spaces = self._env.obs_space.copy() + if self._render: + spaces['image'] = embodied.Space(np.uint8, self._size + (3,)) + return spaces + + @functools.cached_property + def act_space(self): + return self._env.act_space + + def step(self, action): + for key, space in self.act_space.items(): + if not space.discrete: + assert np.isfinite(action[key]).all(), (key, action[key]) + obs = self._env.step(action) + if self._render: + obs['image'] = self.render() + return obs + + def render(self): + return self._dmenv.physics.render(*self._size, camera_id=self._camera) diff --git a/dreamerv3/embodied/envs/crafter.py b/dreamerv3/embodied/envs/crafter.py new file mode 100644 index 0000000..ba85949 --- /dev/null +++ b/dreamerv3/embodied/envs/crafter.py @@ -0,0 +1,73 @@ +import embodied +import numpy as np + + +class Crafter(embodied.Env): + + def __init__(self, task, size=(64, 64), outdir=None, seed=None): + assert task in ('reward', 'noreward') + import crafter + self._env = crafter.Env(size=size, reward=(task == 'reward'), seed=seed) + if outdir: + outdir = embodied.Path(outdir) + self._env = crafter.Recorder( + self._env, outdir, + save_stats=True, + save_video=False, + save_episode=False, + ) + self._achievements = crafter.constants.achievements.copy() + self._done = True + + @property + def obs_space(self): + spaces = { + 'image': embodied.Space(np.uint8, self._env.observation_space.shape), + 'reward': embodied.Space(np.float32), + 'is_first': embodied.Space(bool), + 'is_last': embodied.Space(bool), + 'is_terminal': embodied.Space(bool), + 'log_reward': embodied.Space(np.float32), + } + spaces.update({ + f'log_achievement_{k}': embodied.Space(np.int32) + for k in self._achievements}) + return spaces + + @property + def act_space(self): + return { + 'action': embodied.Space(np.int32, (), 0, self._env.action_space.n), + 'reset': embodied.Space(bool), + } + + def step(self, action): + if action['reset'] or self._done: + self._done = False + image = self._env.reset() + return self._obs(image, 0.0, {}, is_first=True) + image, reward, self._done, info = self._env.step(action['action']) + reward = np.float32(reward) + return self._obs( + image, reward, info, + is_last=self._done, + is_terminal=info['discount'] == 0) + + def _obs( + self, image, reward, info, + is_first=False, is_last=False, is_terminal=False): + log_achievements = { + f'log_achievement_{k}': info['achievements'][k] if info else 0 + for k in self._achievements} + return dict( + image=image, + reward=reward, + is_first=is_first, + is_last=is_last, + is_terminal=is_terminal, + log_reward=np.float32(info['reward'] if info else 0.0), + **log_achievements, + ) + + def render(self): + return self._env.render() diff --git a/dreamerv3/embodied/envs/ddmc.py b/dreamerv3/embodied/envs/ddmc.py new file mode 100644 index 0000000..0d371f0 --- /dev/null +++ b/dreamerv3/embodied/envs/ddmc.py @@ -0,0 +1,85 @@ +import functools +import os + +import embodied +import numpy as np + + +class DDMC(embodied.Env): + + DEFAULT_CAMERAS = dict( + locom_rodent=1, + quadruped=2, + ) + + def __init__(self, env, repeat=1, render=True, size=(64, 64), camera=-1): + # This env variable is meant for headless GPU machines but may fail on CPU-only machines. + if 'MUJOCO_GL' not in os.environ: + os.environ['MUJOCO_GL'] = 'egl' + if isinstance(env, str): + domain, task = env.split('_', 1) + if camera == -1: + camera = self.DEFAULT_CAMERAS.get(domain, 0) + if domain == 'cup': # Only domain with multiple words. + domain = 'ball_in_cup' + + from adaptgym.envs.distracting_control import suite + + print('Default DDMC config.') + dynamic = False + num_videos = 3 + randomize_background = 0 + shuffle_background = 0 + do_color_change = 0 + ground_plane_alpha = 0.1 + background_dataset_videos = ['boat', 'bmx-bumps', 'flamingo'] + continuous_video_frames = True + do_just_background = True + difficulty = 'easy' + specify_background = '0,0,1e6;1,1e6,2e6;0,2e6,1e9' # ABA + + self._dmenv = suite.load(domain, task, difficulty=difficulty, + pixels_only=False, do_just_background=do_just_background, + do_color_change=do_color_change, + background_dataset_videos=background_dataset_videos, + background_kwargs=dict(num_videos=num_videos, + dynamic=dynamic, + randomize_background=randomize_background, + shuffle_buffer_size=shuffle_background * 500, + seed=1, + ground_plane_alpha=ground_plane_alpha, + continuous_video_frames=continuous_video_frames, + specify_background=specify_background, + divide_step_count_by=repeat, + )) + + from . import from_dm + self._env = from_dm.FromDM(self._dmenv) + self._env = embodied.wrappers.ExpandScalars(self._env) + self._env = embodied.wrappers.ActionRepeat(self._env, repeat) + self._render = render + self._size = size + self._camera = camera + + @functools.cached_property + def obs_space(self): + spaces = self._env.obs_space.copy() + if self._render: + spaces['image'] = embodied.Space(np.uint8, self._size + (3,)) + return spaces + + @functools.cached_property + def act_space(self): + return self._env.act_space + + def step(self, action): + for key, space in self.act_space.items(): + if not space.discrete: + assert np.isfinite(action[key]).all(), (key, action[key]) + obs = self._env.step(action) + if self._render: + obs['image'] = self.render() + return obs + + def render(self): + return self._dmenv.physics.render(*self._size, camera_id=self._camera) diff --git a/dreamerv3/embodied/envs/dmc.py b/dreamerv3/embodied/envs/dmc.py new file mode 100644 index 0000000..d1c5437 --- /dev/null +++ b/dreamerv3/embodied/envs/dmc.py @@ -0,0 +1,65 @@ +import functools +import os + +import embodied +import numpy as np + + +class DMC(embodied.Env): + + DEFAULT_CAMERAS = dict( + locom_rodent=1, + quadruped=2, + ) + + def __init__(self, env, repeat=1, render=True, size=(64, 64), camera=-1): + # TODO: This env variable is meant for headless GPU machines but may fail + # on CPU-only machines. + if 'MUJOCO_GL' not in os.environ: + os.environ['MUJOCO_GL'] = 'egl' + if isinstance(env, str): + domain, task = env.split('_', 1) + if camera == -1: + camera = self.DEFAULT_CAMERAS.get(domain, 0) + if domain == 'cup': # Only domain with multiple words. + domain = 'ball_in_cup' + if domain == 'manip': + from dm_control import manipulation + env = manipulation.load(task + '_vision') + elif domain == 'locom': + from dm_control.locomotion.examples import basic_rodent_2020 + env = getattr(basic_rodent_2020, task)() + else: + from dm_control import suite + env = suite.load(domain, task) + self._dmenv = env + from . import from_dm + self._env = from_dm.FromDM(self._dmenv) + self._env = embodied.wrappers.ExpandScalars(self._env) + self._env = embodied.wrappers.ActionRepeat(self._env, repeat) + self._render = render + self._size = size + self._camera = camera + + @functools.cached_property + def obs_space(self): + spaces = self._env.obs_space.copy() + if self._render: + spaces['image'] = embodied.Space(np.uint8, self._size + (3,)) + return spaces + + @functools.cached_property + def act_space(self): + return self._env.act_space + + def step(self, action): + for key, space in self.act_space.items(): + if not space.discrete: + assert np.isfinite(action[key]).all(), (key, action[key]) + obs = self._env.step(action) + if self._render: + obs['image'] = self.render() + return obs + + def render(self): + return self._dmenv.physics.render(*self._size, camera_id=self._camera) diff --git a/dreamerv3/embodied/envs/dmlab.py b/dreamerv3/embodied/envs/dmlab.py new file mode 100644 index 0000000..2bc2edf --- /dev/null +++ b/dreamerv3/embodied/envs/dmlab.py @@ -0,0 +1,142 @@ +import embodied +import numpy as np + + +class DMLab(embodied.Env): + + # Small action set used by IMPALA. + IMPALA_ACTION_SET = ( + ( 0, 0, 0, 1, 0, 0, 0), # Forward + ( 0, 0, 0, -1, 0, 0, 0), # Backward + ( 0, 0, -1, 0, 0, 0, 0), # Strafe Left + ( 0, 0, 1, 0, 0, 0, 0), # Strafe Right + (-20, 0, 0, 0, 0, 0, 0), # Look Left + ( 20, 0, 0, 0, 0, 0, 0), # Look Right + (-20, 0, 0, 1, 0, 0, 0), # Look Left + Forward + ( 20, 0, 0, 1, 0, 0, 0), # Look Right + Forward + ( 0, 0, 0, 0, 1, 0, 0), # Fire + ) + + # Large action set used by PopArt and R2D2. + POPART_ACTION_SET = [ + ( 0, 0, 0, 1, 0, 0, 0), # FW + ( 0, 0, 0, -1, 0, 0, 0), # BW + ( 0, 0, -1, 0, 0, 0, 0), # Strafe Left + ( 0, 0, 1, 0, 0, 0, 0), # Strafe Right + (-10, 0, 0, 0, 0, 0, 0), # Small LL + ( 10, 0, 0, 0, 0, 0, 0), # Small LR + (-60, 0, 0, 0, 0, 0, 0), # Large LL + ( 60, 0, 0, 0, 0, 0, 0), # Large LR + ( 0, 10, 0, 0, 0, 0, 0), # Look Down + ( 0, -10, 0, 0, 0, 0, 0), # Look Up + (-10, 0, 0, 1, 0, 0, 0), # FW + Small LL + ( 10, 0, 0, 1, 0, 0, 0), # FW + Small LR + (-60, 0, 0, 1, 0, 0, 0), # FW + Large LL + ( 60, 0, 0, 1, 0, 0, 0), # FW + Large LR + ( 0, 0, 0, 0, 1, 0, 0), # Fire + ] + + def __init__( + self, level, repeat=4, size=(64, 64), mode='train', + action_set=IMPALA_ACTION_SET, episodic=True, seed=None): + import deepmind_lab + cache = None + # path = os.environ.get('DMLAB_CACHE', None) + # if path: + # cache = Cache(path) + self._size = size + self._repeat = repeat + self._action_set = action_set + self._episodic = episodic + self._random = np.random.RandomState(seed) + config = dict(height=size[0], width=size[1], logLevel='WARN') + if mode == 'train': + if level.endswith('_test'): + level = level.replace('_test', '_train') + elif mode == 'eval': + config.update(allowHoldOutLevels='true', mixerSeed=0x600D5EED) + else: + raise NotImplementedError(mode) + config = {k: str(v) for k, v in config.items()} + self._env = deepmind_lab.Lab( + level='contributed/dmlab30/' + level, + observations=['RGB_INTERLEAVED'], + level_cache=cache, config=config) + self._prev_image = None + self._done = True + + @property + def obs_space(self): + return { + 'image': embodied.Space(np.uint8, self._size + (3,)), + 'reward': embodied.Space(np.float32), + 'is_first': embodied.Space(bool), + 'is_last': embodied.Space(bool), + 'is_terminal': embodied.Space(bool), + } + + @property + def act_space(self): + return { + 'action': embodied.Space(np.int32, (), 0, len(self._action_set)), + 'reset': embodied.Space(bool), + } + + def step(self, action): + if action['reset'] or self._done: + self._env.reset(seed=self._random.randint(0, 2 ** 31 - 1)) + self._done = False + return self._obs(0.0, is_first=True) + raw_action = np.array(self._action_set[action['action']], np.intc) + reward = self._env.step(raw_action, num_steps=self._repeat) + self._done = not self._env.is_running() + return self._obs(reward, is_last=self._done) + + def _obs(self, reward, is_first=False, is_last=False): + return dict( + image=self.render(), + reward=reward, + is_first=is_first, + is_last=is_last, + is_terminal=is_last if self._episodic else False, + ) + + def render(self): + if not self._done: + self._prev_image = self._env.observations()['RGB_INTERLEAVED'] + return self._prev_image + + def close(self): + self._env.close() + + +class Cache: + + def __init__(self, cache_dir): + self._cache_dir = cache_dir + + def get_path(self, key): + import hashlib, os + key = hashlib.md5(key.encode('utf-8')).hexdigest() + dir_, filename = key[:3], key[3:] + return os.path.join(self._cache_dir, dir_, filename) + + def fetch(self, key, pk3_path): + import tensorflow as tf + path = self.get_path(key) + try: + tf.io.gfile.copy(path, pk3_path, overwrite=True) + return True + except tf.errors.OpError: + return False + + def write(self, key, pk3_path): + import os + import tensorflow as tf + path = self.get_path(key) + try: + if not tf.io.gfile.exists(path): + tf.io.gfile.makedirs(os.path.dirname(path)) + tf.io.gfile.copy(pk3_path, path) + except Exception as e: + print(f'Could to store level: {e}') diff --git a/dreamerv3/embodied/envs/dummy.py b/dreamerv3/embodied/envs/dummy.py new file mode 100644 index 0000000..5165a38 --- /dev/null +++ b/dreamerv3/embodied/envs/dummy.py @@ -0,0 +1,54 @@ +import embodied +import numpy as np + + +class Dummy(embodied.Env): + + def __init__(self, task, size=(64, 64), length=100): + assert task in ('cont', 'disc') + self._task = task + self._size = size + self._length = length + self._step = 0 + self._done = False + + @property + def obs_space(self): + return { + 'image': embodied.Space(np.uint8, self._size + (3,)), + 'vector': embodied.Space(np.float32, (7,)), + 'step': embodied.Space(np.int32, (), 0, self._length), + 'reward': embodied.Space(np.float32), + 'is_first': embodied.Space(bool), + 'is_last': embodied.Space(bool), + 'is_terminal': embodied.Space(bool), + } + + @property + def act_space(self): + if self._task == 'cont': + space = embodied.Space(np.float32, (6,)) + else: + space = embodied.Space(np.int32, (), 0, 5) + return {'action': space, 'reset': embodied.Space(bool)} + + def step(self, action): + if action['reset'] or self._done: + self._step = 0 + self._done = False + return self._obs(0.0, is_first=True) + action = action['action'] + self._step += 1 + self._done = (self._step >= self._length) + return self._obs(1.0, is_last=self._done, is_terminal=self._done) + + def _obs(self, reward, is_first=False, is_last=False, is_terminal=False): + return dict( + image=np.zeros(self._size + (3,), np.uint8), + vector=np.zeros(7, np.float32), + step=self._step, + reward=reward, + is_first=is_first, + is_last=is_last, + is_terminal=is_terminal, + ) diff --git a/dreamerv3/embodied/envs/from_dm.py b/dreamerv3/embodied/envs/from_dm.py new file mode 100644 index 0000000..19231a2 --- /dev/null +++ b/dreamerv3/embodied/envs/from_dm.py @@ -0,0 +1,85 @@ +import functools +import os + +import embodied +import numpy as np + + +class FromDM(embodied.Env): + + def __init__(self, env, obs_key='observation', act_key='action'): + self._env = env + obs_spec = self._env.observation_spec() + act_spec = self._env.action_spec() + self._obs_dict = isinstance(obs_spec, dict) + self._act_dict = isinstance(act_spec, dict) + self._obs_key = not self._obs_dict and obs_key + self._act_key = not self._act_dict and act_key + self._obs_empty = [] + self._done = True + + @functools.cached_property + def obs_space(self): + spec = self._env.observation_spec() + spec = spec if self._obs_dict else {self._obs_key: spec} + if 'reward' in spec: + spec['obs_reward'] = spec.pop('reward') + for key, value in spec.copy().items(): + if int(np.prod(value.shape)) == 0: + self._obs_empty.append(key) + del spec[key] + return { + 'reward': embodied.Space(np.float32), + 'is_first': embodied.Space(bool), + 'is_last': embodied.Space(bool), + 'is_terminal': embodied.Space(bool), + **{k or self._obs_key: self._convert(v) for k, v in spec.items()}, + } + + @functools.cached_property + def act_space(self): + spec = self._env.action_spec() + spec = spec if self._act_dict else {self._act_key: spec} + return { + 'reset': embodied.Space(bool), + **{k or self._act_key: self._convert(v) for k, v in spec.items()}, + } + + def step(self, action): + action = action.copy() + reset = action.pop('reset') + if reset or self._done: + time_step = self._env.reset() + else: + action = action if self._act_dict else action[self._act_key] + time_step = self._env.step(action) + self._done = time_step.last() + return self._obs(time_step) + + def _obs(self, time_step): + if not time_step.first(): + assert time_step.discount in (0, 1), time_step.discount + obs = time_step.observation + obs = dict(obs) if self._obs_dict else {self._obs_key: obs} + if 'reward' in obs: + obs['obs_reward'] = obs.pop('reward') + for key in self._obs_empty: + del obs[key] + return dict( + reward=np.float32(0.0 if time_step.first() else time_step.reward), + is_first=time_step.first(), + is_last=time_step.last(), + is_terminal=False if time_step.first() else time_step.discount == 0, + **obs, + ) + + def _convert(self, space): + if hasattr(space, 'num_values'): + return embodied.Space(space.dtype, (), 0, space.num_values) + elif hasattr(space, 'minimum'): + assert np.isfinite(space.minimum).all(), space.minimum + assert np.isfinite(space.maximum).all(), space.maximum + return embodied.Space( + space.dtype, space.shape, space.minimum, space.maximum) + else: + return embodied.Space(space.dtype, space.shape, None, None) diff --git a/dreamerv3/embodied/envs/from_gym.py b/dreamerv3/embodied/envs/from_gym.py new file mode 100644 index 0000000..d997827 --- /dev/null +++ b/dreamerv3/embodied/envs/from_gym.py @@ -0,0 +1,118 @@ +import functools + +import embodied +import gym +import numpy as np + + +class FromGym(embodied.Env): + + def __init__(self, env, obs_key='image', act_key='action', **kwargs): + if isinstance(env, str): + self._env = gym.make(env, **kwargs) + else: + assert not kwargs, kwargs + self._env = env + self._obs_dict = hasattr(self._env.observation_space, 'spaces') + self._act_dict = hasattr(self._env.action_space, 'spaces') + self._obs_key = obs_key + self._act_key = act_key + self._done = True + self._info = None + + @property + def info(self): + return self._info + + @functools.cached_property + def obs_space(self): + if self._obs_dict: + spaces = self._flatten(self._env.observation_space.spaces) + else: + spaces = {self._obs_key: self._env.observation_space} + spaces = {k: self._convert(v) for k, v in spaces.items()} + return { + **spaces, + 'reward': embodied.Space(np.float32), + 'is_first': embodied.Space(bool), + 'is_last': embodied.Space(bool), + 'is_terminal': embodied.Space(bool), + } + + @functools.cached_property + def act_space(self): + if self._act_dict: + spaces = self._flatten(self._env.action_space.spaces) + else: + spaces = {self._act_key: self._env.action_space} + spaces = {k: self._convert(v) for k, v in spaces.items()} + spaces['reset'] = embodied.Space(bool) + return spaces + + def step(self, action): + if action['reset'] or self._done: + self._done = False + obs = self._env.reset() + return self._obs(obs, 0.0, is_first=True) + if self._act_dict: + action = self._unflatten(action) + else: + action = action[self._act_key] + obs, reward, self._done, self._info = self._env.step(action) + return self._obs( + obs, reward, + is_last=bool(self._done), + is_terminal=bool(self._info.get('is_terminal', self._done))) + + def _obs( + self, obs, reward, is_first=False, is_last=False, is_terminal=False): + if not self._obs_dict: + obs = {self._obs_key: obs} + obs = self._flatten(obs) + obs = {k: np.asarray(v) for k, v in obs.items()} + obs.update( + reward=np.float32(reward), + is_first=is_first, + is_last=is_last, + is_terminal=is_terminal) + return obs + + def render(self): + image = self._env.render('rgb_array') + assert image is not None + return image + + def close(self): + try: + self._env.close() + except Exception: + pass + + def _flatten(self, nest, prefix=None): + result = {} + for key, value in nest.items(): + key = prefix + '/' + key if prefix else key + if isinstance(value, gym.spaces.Dict): + value = value.spaces + if isinstance(value, dict): + result.update(self._flatten(value, key)) + else: + result[key] = value + return result + + def _unflatten(self, flat): + result = {} + for key, value in flat.items(): + parts = key.split('/') + node = result + for part in parts[:-1]: + if part not in node: + node[part] = {} + node = node[part] + node[parts[-1]] = value + return result + + def _convert(self, space): + if hasattr(space, 'n'): + return embodied.Space(np.int32, (), 0, space.n) + return embodied.Space(space.dtype, space.shape, space.low, space.high) diff --git a/dreamerv3/embodied/envs/loconav.py b/dreamerv3/embodied/envs/loconav.py new file mode 100644 index 0000000..d34513a --- /dev/null +++ b/dreamerv3/embodied/envs/loconav.py @@ -0,0 +1,230 @@ +import functools +import os +import warnings + +import embodied +import numpy as np + + +class LocoNav(embodied.Env): + + DEFAULT_CAMERAS = dict( + ant=4, + quadruped=5, + ) + + def __init__( + self, name, repeat=1, size=(64, 64), camera=-1, again=False, + termination=False, weaker=1.0): + # TODO: This env variable is meant for headless GPU machines but may fail + # on CPU-only machines. + if name.endswith('hz'): + name, freq = name.rsplit('_', 1) + freq = int(freq.strip('hz')) + else: + freq = 50 + if 'MUJOCO_GL' not in os.environ: + os.environ['MUJOCO_GL'] = 'egl' + from dm_control import composer + from dm_control.locomotion.props import target_sphere + from dm_control.locomotion.tasks import random_goal_maze + walker, arena = name.split('_', 1) + if camera == -1: + camera = self.DEFAULT_CAMERAS.get(walker, 0) + self._walker = self._make_walker(walker) + arena = self._make_arena(arena) + target = target_sphere.TargetSphere(radius=1.2, height_above_ground=0.0) + task = random_goal_maze.RepeatSingleGoalMaze( + walker=self._walker, maze_arena=arena, target=target, + max_repeats=1000 if again else 1, + randomize_spawn_rotation=True, + target_reward_scale=1.0, + aliveness_threshold=-0.5 if termination else -1.0, + contact_termination=False, + physics_timestep=min(1 / freq / 4, 0.02), + control_timestep=1 / freq) + if not again: + def after_step(self, physics, random_state): + super(random_goal_maze.RepeatSingleGoalMaze, self).after_step( + physics, random_state) + self._rewarded_this_step = self._target.activated + self._targets_obtained = int(self._target.activated) + task.after_step = functools.partial(after_step, task) + env = composer.Environment( + time_limit=60, task=task, random_state=None, + strip_singleton_obs_buffer_dim=True) + from . import dmc + self._env = dmc.DMC(env, repeat, size=size, camera=camera) + self._visited = None + self._weaker = weaker + + @property + def obs_space(self): + return { + **self._env.obs_space, + 'log_coverage': embodied.Space(np.int64, low=0), + } + + @property + def act_space(self): + return self._env.act_space + + def step(self, action): + with warnings.catch_warnings(): + warnings.filterwarnings('ignore', '.*is a deprecated alias for.*') + action = action.copy() + action['action'] *= self._weaker + obs = self._env.step(action) + if obs['is_first']: + self._visited = set() + global_pos = self._walker.get_pose( + self._env._dmenv._physics)[0].reshape(-1) + self._visited.add(tuple(np.round(global_pos[:2]).astype(int).tolist())) + obs['log_coverage'] = len(self._visited) + return obs + + def _make_walker(self, name): + if name == 'ant': + from dm_control.locomotion.walkers import ant + return ant.Ant() + elif name == 'quadruped': + from . import loconav_quadruped + return loconav_quadruped.Quadruped() + else: + raise NotImplementedError(name) + + def _make_arena(self, name): + import labmaze + from dm_control import mjcf + from dm_control.locomotion.arenas import labmaze_textures + from dm_control.locomotion.arenas import mazes + import matplotlib.pyplot as plt + class WallTexture(labmaze_textures.WallTextures): + def _build(self, color=[0.8, 0.8, 0.8], model='labmaze_style_01'): + self._mjcf_root = mjcf.RootElement(model=model) + self._textures = [self._mjcf_root.asset.add( + 'texture', type='2d', name='wall', builtin='flat', + rgb1=color, width=100, height=100)] + wall_textures = {'*': WallTexture([0.8, 0.8, 0.8])} + cmap = plt.get_cmap('tab10') + for index in range(9): + wall_textures[str(index + 1)] = WallTexture(cmap(index)[:3]) + layout = ''.join([ + line[::2].replace('.', ' ') + '\n' for line in MAPS[name]]) + maze = labmaze.FixedMazeWithRandomGoals( + entity_layer=layout, + num_spawns=1, num_objects=1, random_state=None) + arena = mazes.MazeWithTargets( + maze, xy_scale=1.2, z_height=2.0, aesthetic='default', + wall_textures=wall_textures, name='maze') + return arena + + +MAPS = { + + 'maze_s': ( + ' 6 6 6 6 6', + ' 6 . . . 6', + ' 6 . G . 6', + ' 6 . . . 6', + ' 5 . . . 4', + ' 5 . . . 4', + '1 1 1 1 5 5 5 . . . 4', + '1 . . . . . . . . . 3', + '1 . P . . . . . . . 3', + '1 . . . . . . . . . 3', + '1 1 1 1 2 2 2 3 3 3 3', + ), + + 'maze_m': ( + '6 6 6 6 8 8 8 7 7 7 7', + '6 . . . . . . . . . 7', + '6 . G . . . . . . . 7', + '6 . . . . . . . . . 7', + '6 6 6 5 5 5 5 . . . 4', + ' 5 . . . 4', + '1 1 1 1 5 5 5 . . . 4', + '1 . . . . . . . . . 3', + '1 . P . . . . . . . 3', + '1 . . . . . . . . . 3', + '1 1 1 1 2 2 2 3 3 3 3', + ), + + 'maze_l': ( + '8 8 8 8 7 7 7 6 6 6 6 . . .', + '8 . . . . . . . . . 6 . . .', + '8 . G . . . . . . . 6 . . .', + '8 . . . . . . . . . 6 5 5 5', + '8 8 8 8 7 7 7 . . . . . . 5', + '. . . . . . 7 . . . . . . 5', + '1 1 1 1 1 . 7 . . . . . . 5', + '1 . . . 1 . 7 9 9 9 . . . 5', + '1 . . . 1 . . . . 9 . . . 5', + '1 . . . 1 1 1 9 9 9 . . . 5', + '2 . . . . . . . . . . . . 4', + '2 . . . . P . . . . . . . 4', + '2 . . . . . . . . . . . . 4', + '2 2 2 2 3 3 3 3 3 3 4 4 4 4', + ), + + 'maze_xl': ( + '9 9 9 9 9 9 9 8 8 8 8 . 4 4 4 4 4', + '9 . . . . . . . . . 8 . 4 . . . 4', + '9 . . . . . . . G . 8 . 4 . . . 4', + '9 . . . . . . . . . 8 . 4 . . . 4', + '6 . . . 7 7 7 8 8 8 8 . 5 . . . 3', + '6 . . . 7 . . . . . . . 5 . . . 3', + '6 . . . 7 7 7 5 5 5 5 5 5 . . . 3', + '5 . . . . . . . . . . . . . . . 3', + '5 . . . . . . . . . . . . . . . 3', + '5 . . . . . . . . . . . . . . . 3', + '5 5 5 5 4 4 4 . . . 6 6 6 . . . 3', + '. . . . . . 4 . . . 6 . 6 . . . 3', + '1 1 1 1 4 4 4 . . . 6 . 6 . . . 3', + '1 . . . . . . . . . 2 . 1 . . . 1', + '1 . P . . . . . . . 2 . 1 . . . 1', + '1 . . . . . . . . . 2 . 1 . . . 1', + '1 1 1 1 1 1 1 2 2 2 2 . 1 1 1 1 1', + ), + + 'maze_xxl': ( + '7 7 7 7 * * * 6 6 6 * * * 9 9 9 9', + '7 . . . . . . . . . . . . . . . 9', + '7 . . . . . . . . . . . . . G . 9', + '7 . . . . . . . . . . . . . . . 9', + '* . . . 5 5 5 * * * * * * 9 9 9 9', + '* . . . 5 . . . . . . . . . . . .', + '* . . . 5 5 5 * * * * * * 3 3 3 3', + '8 . . . . . . . . . . . . . . . 3', + '8 . . . . . . . . . . . . . . . 3', + '8 . . . . . . . . . . . . . . . 3', + '8 8 8 8 * * * * * * 4 4 4 . . . *', + '. . . . . . . . . . . . 4 . . . *', + '1 1 1 1 * * * * * * 4 4 4 . . . *', + '1 . . . . . . . . . . . . . . . 2', + '1 . P . . . . . . . . . . . . . 2', + '1 . . . . . . . . . . . . . . . 2', + '1 1 1 1 * * * 6 6 6 * * * 2 2 2 2', + ), + + 'empty': ( + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + '. . . . . . . . . . . . . . . . .', + ), + +} diff --git a/dreamerv3/embodied/envs/loconav_quadruped.py b/dreamerv3/embodied/envs/loconav_quadruped.py new file mode 100644 index 0000000..5c5701d --- /dev/null +++ b/dreamerv3/embodied/envs/loconav_quadruped.py @@ -0,0 +1,132 @@ +import os + +from dm_control import composer +from dm_control import mjcf +from dm_control.composer.observation import observable +from dm_control.locomotion.walkers import base +from dm_control.locomotion.walkers import legacy_base +from dm_control.mujoco.wrapper import mjbindings +import numpy as np + +enums = mjbindings.enums +mjlib = mjbindings.mjlib + + +class Quadruped(legacy_base.Walker): + + def _build(self, name='walker', initializer=None): + super()._build(initializer=initializer) + self._mjcf_root = mjcf.from_path( + os.path.join(os.path.dirname(__file__), 'loconav_quadruped.xml')) + if name: + self._mjcf_root.model = name + self._prev_action = np.zeros( + self.action_spec.shape, self.action_spec.dtype) + + def initialize_episode(self, physics, random_state): + self._prev_action = np.zeros_like(self._prev_action) + + def apply_action(self, physics, action, random_state): + super().apply_action(physics, action, random_state) + self._prev_action[:] = action + + def _build_observables(self): + return QuadrupedObservables(self) + + @property + def mjcf_model(self): + return self._mjcf_root + + @property + def upright_pose(self): + return base.WalkerPose() + + @composer.cached_property + def actuators(self): + return self._mjcf_root.find_all('actuator') + + @composer.cached_property + def root_body(self): + return self._mjcf_root.find('body', 'torso') + + @composer.cached_property + def bodies(self): + return tuple(self.mjcf_model.find_all('body')) + + @composer.cached_property + def mocap_tracking_bodies(self): + return tuple(self.mjcf_model.find_all('body')) + + @property + def mocap_joints(self): + return self.mjcf_model.find_all('joint') + + @property + def _foot_bodies(self): + return ( + self._mjcf_root.find('body', 'toe_front_left'), + self._mjcf_root.find('body', 'toe_front_right'), + self._mjcf_root.find('body', 'toe_back_right'), + self._mjcf_root.find('body', 'toe_back_left'), + ) + + @composer.cached_property + def end_effectors(self): + return self._foot_bodies + + @composer.cached_property + def observable_joints(self): + return self._mjcf_root.find_all('joint') + + @composer.cached_property + def egocentric_camera(self): + return self._mjcf_root.find('camera', 'egocentric') + + def aliveness(self, physics): + return (physics.bind(self.root_body).xmat[-1] - 1.) / 2. + + @composer.cached_property + def ground_contact_geoms(self): + foot_geoms = [] + for foot in self._foot_bodies: + foot_geoms.extend(foot.find_all('geom')) + return tuple(foot_geoms) + + @property + def prev_action(self): + return self._prev_action + + +class QuadrupedObservables(legacy_base.WalkerObservables): + + @composer.observable + def actuator_activations(self): + def actuator_activations_in_egocentric_frame(physics): + return physics.data.act + return observable.Generic(actuator_activations_in_egocentric_frame) + + @composer.observable + def root_global_pos(self): + def root_pos(physics): + root_xpos, _ = self._entity.get_pose(physics) + return np.reshape(root_xpos, -1) + return observable.Generic(root_pos) + + @composer.observable + def torso_global_pos(self): + def torso_pos(physics): + root_body = self._entity.root_body + root_body_xpos = physics.bind(root_body).xpos + return np.reshape(root_body_xpos, -1) + return observable.Generic(torso_pos) + + @property + def proprioception(self): + return ([ + self.joints_pos, self.joints_vel, self.actuator_activations, + self.sensors_accelerometer, self.sensors_gyro, + self.sensors_velocimeter, + self.sensors_force, self.sensors_torque, + self.world_zaxis, + self.root_global_pos, self.torso_global_pos, + ] + self._collect_from_attachments('proprioception')) diff --git a/dreamerv3/embodied/envs/loconav_quadruped.xml b/dreamerv3/embodied/envs/loconav_quadruped.xml new file mode 100644 index 0000000..c556dae --- /dev/null +++ b/dreamerv3/embodied/envs/loconav_quadruped.xml @@ -0,0 +1,311 @@ + + + + + + + + + + + + + + + + + + + + + + + + + + + diff --git a/dreamerv3/embodied/envs/minecraft.py b/dreamerv3/embodied/envs/minecraft.py new file mode 100644 index 0000000..512ae7f --- /dev/null +++ b/dreamerv3/embodied/envs/minecraft.py @@ -0,0 +1,156 @@ +import embodied +import numpy as np + +from . import minecraft_base + + +class Minecraft(embodied.Wrapper): + + def __init__(self, task, *args, **kwargs): + super().__init__({ + 'wood': MinecraftWood, + 'climb': MinecraftClimb, + 'diamond': MinecraftDiamond, + }[task](*args, **kwargs)) + + +class MinecraftWood(embodied.Wrapper): + + def __init__(self, *args, **kwargs): + actions = BASIC_ACTIONS + self.rewards = [ + CollectReward('log', repeated=1), + HealthReward(), + ] + length = kwargs.pop('length', 36000) + env = minecraft_base.MinecraftBase(actions, *args, **kwargs) + env = embodied.wrappers.TimeLimit(env, length) + super().__init__(env) + + def step(self, action): + obs = self.env.step(action) + obs['reward'] = sum([fn(obs, self.env.inventory) for fn in self.rewards]) + return obs + + +class MinecraftClimb(embodied.Wrapper): + + def __init__(self, *args, **kwargs): + actions = BASIC_ACTIONS + length = kwargs.pop('length', 36000) + env = minecraft_base.MinecraftBase(actions, *args, **kwargs) + env = embodied.wrappers.TimeLimit(env, length) + super().__init__(env) + self._previous = None + self._health_reward = HealthReward() + + def step(self, action): + obs = self.env.step(action) + x, y, z = obs['log_player_pos'] + height = np.float32(y) + if obs['is_first']: + self._previous = height + obs['reward'] = height - self._previous + obs['reward'] += self._health_reward(obs) + self._previous = height + return obs + + +class MinecraftDiamond(embodied.Wrapper): + + def __init__(self, *args, **kwargs): + actions = { + **BASIC_ACTIONS, + 'craft_planks': dict(craft='planks'), + 'craft_stick': dict(craft='stick'), + 'craft_crafting_table': dict(craft='crafting_table'), + 'place_crafting_table': dict(place='crafting_table'), + 'craft_wooden_pickaxe': dict(nearbyCraft='wooden_pickaxe'), + 'craft_stone_pickaxe': dict(nearbyCraft='stone_pickaxe'), + 'craft_iron_pickaxe': dict(nearbyCraft='iron_pickaxe'), + 'equip_stone_pickaxe': dict(equip='stone_pickaxe'), + 'equip_wooden_pickaxe': dict(equip='wooden_pickaxe'), + 'equip_iron_pickaxe': dict(equip='iron_pickaxe'), + 'craft_furnace': dict(nearbyCraft='furnace'), + 'place_furnace': dict(place='furnace'), + 'smelt_iron_ingot': dict(nearbySmelt='iron_ingot'), + } + self.rewards = [ + CollectReward('log', once=1), + CollectReward('planks', once=1), + CollectReward('stick', once=1), + CollectReward('crafting_table', once=1), + CollectReward('wooden_pickaxe', once=1), + CollectReward('cobblestone', once=1), + CollectReward('stone_pickaxe', once=1), + CollectReward('iron_ore', once=1), + CollectReward('furnace', once=1), + CollectReward('iron_ingot', once=1), + CollectReward('iron_pickaxe', once=1), + CollectReward('diamond', once=1), + HealthReward(), + ] + length = kwargs.pop('length', 36000) + env = minecraft_base.MinecraftBase(actions, *args, **kwargs) + env = embodied.wrappers.TimeLimit(env, length) + super().__init__(env) + + def step(self, action): + obs = self.env.step(action) + obs['reward'] = sum([fn(obs, self.env.inventory) for fn in self.rewards]) + return obs + + +class CollectReward: + + def __init__(self, item, once=0, repeated=0): + self.item = item + self.once = once + self.repeated = repeated + self.previous = 0 + self.maximum = 0 + + def __call__(self, obs, inventory): + current = inventory[self.item] + if obs['is_first']: + self.previous = current + self.maximum = current + return 0 + reward = self.repeated * max(0, current - self.previous) + if self.maximum == 0 and current > 0: + reward += self.once + self.previous = current + self.maximum = max(self.maximum, current) + return reward + + +class HealthReward: + + def __init__(self, scale=0.01): + self.scale = scale + self.previous = None + + def __call__(self, obs, inventory=None): + health = obs['health'] + if obs['is_first']: + self.previous = health + return 0 + reward = self.scale * (health - self.previous) + self.previous = health + return np.float32(reward) + + +BASIC_ACTIONS = { + 'noop': dict(), + 'attack': dict(attack=1), + 'turn_up': dict(camera=(-15, 0)), + 'turn_down': dict(camera=(15, 0)), + 'turn_left': dict(camera=(0, -15)), + 'turn_right': dict(camera=(0, 15)), + 'forward': dict(forward=1), + 'back': dict(back=1), + 'left': dict(left=1), + 'right': dict(right=1), + 'jump': dict(jump=1, forward=1), + 'place_dirt': dict(place='dirt'), +} diff --git a/dreamerv3/embodied/envs/minecraft_base.py b/dreamerv3/embodied/envs/minecraft_base.py new file mode 100644 index 0000000..3d8f48c --- /dev/null +++ b/dreamerv3/embodied/envs/minecraft_base.py @@ -0,0 +1,189 @@ +import logging +import threading + +import embodied +import numpy as np + + +class MinecraftBase(embodied.Env): + + _LOCK = threading.Lock() + + def __init__( + self, actions, + repeat=1, + size=(64, 64), + break_speed=100.0, + gamma=10.0, + sticky_attack=30, + sticky_jump=10, + pitch_limit=(-60, 60), + logs=True, # TODO + ): + if logs: + logging.basicConfig(level=logging.DEBUG) + self._repeat = repeat + self._size = size + if break_speed != 1.0: + sticky_attack = 0 + + # Make env + with self._LOCK: + from .import minecraft_minerl + self._gymenv = minecraft_minerl.MineRLEnv(size, break_speed, gamma).make() + from . import from_gym + self._env = from_gym.FromGym(self._gymenv) + self._inventory = {} + + # Observations + self._inv_keys = [ + k for k in self._env.obs_space if k.startswith('inventory/') + if k != 'inventory/log2'] + self._step = 0 + self._max_inventory = None + self._equip_enum = self._gymenv.observation_space[ + 'equipped_items']['mainhand']['type'].values.tolist() + self._obs_space = self.obs_space + + # Actions + self._noop_action = minecraft_minerl.NOOP_ACTION + actions = self._insert_defaults(actions) + self._action_names = tuple(actions.keys()) + self._action_values = tuple(actions.values()) + message = f'Minecraft action space ({len(self._action_values)}):' + print(message, ', '.join(self._action_names)) + self._sticky_attack_length = sticky_attack + self._sticky_attack_counter = 0 + self._sticky_jump_length = sticky_jump + self._sticky_jump_counter = 0 + self._pitch_limit = pitch_limit + self._pitch = 0 + + @property + def obs_space(self): + return { + 'image': embodied.Space(np.uint8, self._size + (3,)), + 'inventory': embodied.Space(np.float32, len(self._inv_keys), 0), + 'inventory_max': embodied.Space(np.float32, len(self._inv_keys), 0), + 'equipped': embodied.Space(np.float32, len(self._equip_enum), 0, 1), + 'reward': embodied.Space(np.float32), + 'health': embodied.Space(np.float32), + 'hunger': embodied.Space(np.float32), + 'breath': embodied.Space(np.float32), + 'is_first': embodied.Space(bool), + 'is_last': embodied.Space(bool), + 'is_terminal': embodied.Space(bool), + **{f'log_{k}': embodied.Space(np.int64) for k in self._inv_keys}, + 'log_player_pos': embodied.Space(np.float32, 3), + } + + @property + def act_space(self): + return { + 'action': embodied.Space(np.int64, (), 0, len(self._action_values)), + 'reset': embodied.Space(bool), + } + + def step(self, action): + action = action.copy() + index = action.pop('action') + action.update(self._action_values[index]) + action = self._action(action) + if action['reset']: + obs = self._reset() + else: + following = self._noop_action.copy() + for key in ('attack', 'forward', 'back', 'left', 'right'): + following[key] = action[key] + for act in [action] + ([following] * (self._repeat - 1)): + obs = self._env.step(act) + if 'error' in self._env.info: + obs = self._reset() + break + obs = self._obs(obs) + self._step += 1 + assert 'pov' not in obs, list(obs.keys()) + return obs + + @property + def inventory(self): + return self._inventory + + def _reset(self): + with self._LOCK: + obs = self._env.step({'reset': True}) + self._step = 0 + self._max_inventory = None + self._sticky_attack_counter = 0 + self._sticky_jump_counter = 0 + self._pitch = 0 + self._inventory = {} + return obs + + def _obs(self, obs): + obs['inventory/log'] += obs.pop('inventory/log2') + self._inventory = { + k.split('/', 1)[1]: obs[k] for k in self._inv_keys + if k != 'inventory/air'} + inventory = np.array([obs[k] for k in self._inv_keys], np.float32) + if self._max_inventory is None: + self._max_inventory = inventory + else: + self._max_inventory = np.maximum(self._max_inventory, inventory) + index = self._equip_enum.index(obs['equipped_items/mainhand/type']) + equipped = np.zeros(len(self._equip_enum), np.float32) + equipped[index] = 1.0 + player_x = obs['location_stats/xpos'] + player_y = obs['location_stats/ypos'] + player_z = obs['location_stats/zpos'] + obs = { + 'image': obs['pov'], + 'inventory': inventory, + 'inventory_max': self._max_inventory.copy(), + 'equipped': equipped, + 'health': np.float32(obs['life_stats/life'] / 20), + 'hunger': np.float32(obs['life_stats/food'] / 20), + 'breath': np.float32(obs['life_stats/air'] / 300), + 'reward': 0.0, + 'is_first': obs['is_first'], + 'is_last': obs['is_last'], + 'is_terminal': obs['is_terminal'], + **{f'log_{k}': np.int64(obs[k]) for k in self._inv_keys}, + 'log_player_pos': np.array([player_x, player_y, player_z], np.float32), + } + for key, value in obs.items(): + space = self._obs_space[key] + if not isinstance(value, np.ndarray): + value = np.array(value) + assert value in space, (key, value, value.dtype, value.shape, space) + return obs + + def _action(self, action): + if self._sticky_attack_length: + if action['attack']: + self._sticky_attack_counter = self._sticky_attack_length + if self._sticky_attack_counter > 0: + action['attack'] = 1 + action['jump'] = 0 + self._sticky_attack_counter -= 1 + if self._sticky_jump_length: + if action['jump']: + self._sticky_jump_counter = self._sticky_jump_length + if self._sticky_jump_counter > 0: + action['jump'] = 1 + action['forward'] = 1 + self._sticky_jump_counter -= 1 + if self._pitch_limit and action['camera'][0]: + lo, hi = self._pitch_limit + if not (lo <= self._pitch + action['camera'][0] <= hi): + action['camera'] = (0, action['camera'][1]) + self._pitch += action['camera'][0] + return action + + def _insert_defaults(self, actions): + actions = {name: action.copy() for name, action in actions.items()} + for key, default in self._noop_action.items(): + for action in actions.values(): + if key not in action: + action[key] = default + return actions diff --git a/dreamerv3/embodied/envs/minecraft_minerl.py b/dreamerv3/embodied/envs/minecraft_minerl.py new file mode 100644 index 0000000..b412218 --- /dev/null +++ b/dreamerv3/embodied/envs/minecraft_minerl.py @@ -0,0 +1,150 @@ +from minerl.herobraine.env_spec import EnvSpec +from minerl.herobraine.hero import handler +from minerl.herobraine.hero import handlers +from minerl.herobraine.hero import mc +from minerl.herobraine.hero.mc import INVERSE_KEYMAP + + +def edit_options(**kwargs): + import os, pathlib, re + for word in os.popen('pip3 --version').read().split(' '): + if '-packages/pip' in word: + break + else: + raise RuntimeError('Could not found python package directory.') + packages = pathlib.Path(word).parent + filename = packages / 'minerl/Malmo/Minecraft/run/options.txt' + options = filename.read_text() + if 'fovEffectScale:' not in options: + options += 'fovEffectScale:1.0\n' + if 'simulationDistance:' not in options: + options += 'simulationDistance:12\n' + for key, value in kwargs.items(): + assert f'{key}:' in options, key + assert isinstance(value, str), (value, type(value)) + options = re.sub(f'{key}:.*\n', f'{key}:{value}\n', options) + filename.write_text(options) + + +edit_options( + difficulty='2', + renderDistance='6', + simulationDistance='6', + fovEffectScale='0.0', + ao='1', + gamma='5.0', +) + + +class MineRLEnv(EnvSpec): + + def __init__(self, resolution=(64, 64), break_speed=50, gamma=10.0): + self.resolution = resolution + self.break_speed = break_speed + self.gamma = gamma + super().__init__(name='MineRLEnv-v1') + + def create_agent_start(self): + return [ + BreakSpeedMultiplier(self.break_speed), + ] + + def create_agent_handlers(self): + return [] + + def create_server_world_generators(self): + return [handlers.DefaultWorldGenerator(force_reset=True)] + + def create_server_quit_producers(self): + return [handlers.ServerQuitWhenAnyAgentFinishes()] + + def create_server_initial_conditions(self): + return [ + handlers.TimeInitialCondition( + allow_passage_of_time=True, + start_time=0, + ), + handlers.SpawningInitialCondition( + allow_spawning=True, + ) + ] + + def create_observables(self): + return [ + handlers.POVObservation(self.resolution), + handlers.FlatInventoryObservation(mc.ALL_ITEMS), + handlers.EquippedItemObservation( + mc.ALL_ITEMS, _default='air', _other='other'), + handlers.ObservationFromCurrentLocation(), + handlers.ObservationFromLifeStats(), + ] + + def create_actionables(self): + kw = dict(_other='none', _default='none') + return [ + handlers.KeybasedCommandAction('forward', INVERSE_KEYMAP['forward']), + handlers.KeybasedCommandAction('back', INVERSE_KEYMAP['back']), + handlers.KeybasedCommandAction('left', INVERSE_KEYMAP['left']), + handlers.KeybasedCommandAction('right', INVERSE_KEYMAP['right']), + handlers.KeybasedCommandAction('jump', INVERSE_KEYMAP['jump']), + handlers.KeybasedCommandAction('sneak', INVERSE_KEYMAP['sneak']), + handlers.KeybasedCommandAction('attack', INVERSE_KEYMAP['attack']), + handlers.CameraAction(), + handlers.PlaceBlock(['none'] + mc.ALL_ITEMS, **kw), + handlers.EquipAction(['none'] + mc.ALL_ITEMS, **kw), + handlers.CraftAction(['none'] + mc.ALL_ITEMS, **kw), + handlers.CraftNearbyAction(['none'] + mc.ALL_ITEMS, **kw), + handlers.SmeltItemNearby(['none'] + mc.ALL_ITEMS, **kw), + ] + + def is_from_folder(self, folder): + return folder == 'none' + + def get_docstring(self): + return '' + + def determine_success_from_rewards(self, rewards): + return True + + def create_rewardables(self): + return [] + + def create_server_decorators(self): + return [] + + def create_mission_handlers(self): + return [] + + def create_monitors(self): + return [] + + +class BreakSpeedMultiplier(handler.Handler): + + def __init__(self, multiplier=1.0): + self.multiplier = multiplier + + def to_string(self): + return f'break_speed({self.multiplier})' + + def xml_template(self): + return '{{multiplier}}' + + +class Gamma(handler.Handler): + + def __init__(self, gamma=2.0): + self.gamma = gamma + + def to_string(self): + return f'gamma({self.gamma})' + + def xml_template(self): + return '{{gamma}}' + + +NOOP_ACTION = dict( + camera=(0, 0), forward=0, back=0, left=0, right=0, attack=0, sprint=0, + jump=0, sneak=0, craft='none', nearbyCraft='none', nearbySmelt='none', + place='none', equip='none', +) diff --git a/dreamerv3/embodied/envs/pinpad.py b/dreamerv3/embodied/envs/pinpad.py new file mode 100644 index 0000000..da0e008 --- /dev/null +++ b/dreamerv3/embodied/envs/pinpad.py @@ -0,0 +1,220 @@ +import collections + +import embodied +import numpy as np + + +class PinPad(embodied.Env): + + COLORS = { + '1': (255, 0, 0), + '2': ( 0, 255, 0), + '3': ( 0, 0, 255), + '4': (255, 255, 0), + '5': (255, 0, 255), + '6': ( 0, 255, 255), + '7': (128, 0, 128), + '8': ( 0, 128, 128), + } + + def __init__(self, task, length=10000): + assert length > 0 + layout = { + 'three': LAYOUT_THREE, + 'four': LAYOUT_FOUR, + 'five': LAYOUT_FIVE, + 'six': LAYOUT_SIX, + 'seven': LAYOUT_SEVEN, + 'eight': LAYOUT_EIGHT, + }[task] + self.layout = np.array([list(line) for line in layout.split('\n')]).T + assert self.layout.shape == (16, 14), self.layout.shape + self.length = length + self.random = np.random.RandomState() + self.pads = set(self.layout.flatten().tolist()) - set('* #\n') + self.target = tuple(sorted(self.pads)) + self.spawns = [] + for (x, y), char in np.ndenumerate(self.layout): + if char != '#': + self.spawns.append((x, y)) + print(f'Created PinPad env with sequence: {"->".join(self.target)}') + self.sequence = collections.deque(maxlen=len(self.target)) + self.player = None + self.steps = None + self.done = None + self.countdown = None + + @property + def act_space(self): + return { + 'action': embodied.Space(np.int64, (), 0, 5), + 'reset': embodied.Space(bool), + } + + @property + def obs_space(self): + return { + 'image': embodied.Space(np.uint8, (64, 64, 3)), + 'reward': embodied.Space(np.float32), + 'is_first': embodied.Space(bool), + 'is_last': embodied.Space(bool), + 'is_terminal': embodied.Space(bool), + } + + def step(self, action): + if self.done or action['reset']: + self.player = self.spawns[self.random.randint(len(self.spawns))] + self.sequence.clear() + self.steps = 0 + self.done = False + self.countdown = 0 + return self._obs(reward=0.0, is_first=True) + if self.countdown: + self.countdown -= 1 + if self.countdown == 0: + self.player = self.spawns[self.random.randint(len(self.spawns))] + self.sequence.clear() + reward = 0.0 + move = [(0, 0), (0, 1), (0, -1), (1, 0), (-1, 0)][action['action']] + x = np.clip(self.player[0] + move[0], 0, 15) + y = np.clip(self.player[1] + move[1], 0, 13) + tile = self.layout[x][y] + if tile != '#': + self.player = (x, y) + if tile in self.pads: + if not self.sequence or self.sequence[-1] != tile: + self.sequence.append(tile) + if tuple(self.sequence) == self.target and not self.countdown: + reward += 10.0 + self.countdown = 10 + self.steps += 1 + self.done = self.done or (self.steps >= self.length) + return self._obs(reward=reward, is_last=self.done) + + def render(self): + grid = np.zeros((16, 16, 3), np.uint8) + 255 + white = np.array([255, 255, 255]) + if self.countdown: + grid[:] = (223, 255, 223) + current = self.layout[self.player[0]][self.player[1]] + for (x, y), char in np.ndenumerate(self.layout): + if char == '#': + grid[x, y] = (192, 192, 192) + elif char in self.pads: + color = np.array(self.COLORS[char]) + color = color if char == current else (10 * color + 90 * white) / 100 + grid[x, y] = color + grid[self.player] = (0, 0, 0) + grid[:, -2:] = (192, 192, 192) + for i, char in enumerate(self.sequence): + grid[2 * i + 1, -2] = self.COLORS[char] + image = np.repeat(np.repeat(grid, 4, 0), 4, 1) + return image.transpose((1, 0, 2)) + + def _obs(self, reward, is_first=False, is_last=False, is_terminal=False): + return dict( + image=self.render(), reward=reward, is_first=is_first, is_last=is_last, + is_terminal=is_terminal) + + +LAYOUT_THREE = """ +################ +#1111 3333# +#1111 3333# +#1111 3333# +#1111 3333# +# # +# # +# # +# # +# 2222 # +# 2222 # +# 2222 # +# 2222 # +################ +""".strip('\n') + +LAYOUT_FOUR = """ +################ +#1111 4444# +#1111 4444# +#1111 4444# +#1111 4444# +# # +# # +# # +# # +#3333 2222# +#3333 2222# +#3333 2222# +#3333 2222# +################ +""".strip('\n') + +LAYOUT_FIVE = """ +################ +# 4444# +#111 4444# +#111 4444# +#111 # +#111 555# +# 555# +# 555# +#333 555# +#333 # +#333 2222# +#333 2222# +# 2222# +################ +""".strip('\n') + +LAYOUT_SIX = """ +################ +#111 555# +#111 555# +#111 555# +# # +#33 66# +#33 66# +#33 66# +#33 66# +# # +#444 222# +#444 222# +#444 222# +################ +""".strip('\n') + +LAYOUT_SEVEN = """ +################ +#111 444# +#111 444# +#11 44# +# # +#33 55# +#33 55# +#33 55# +#33 55# +# # +#66 22# +#666 7777 222# +#666 7777 222# +################ +""".strip('\n') + +LAYOUT_EIGHT = """ +################ +#111 8888 444# +#111 8888 444# +#11 44# +# # +#33 55# +#33 55# +#33 55# +#33 55# +# # +#66 22# +#666 7777 222# +#666 7777 222# +################ +""".strip('\n') diff --git a/dreamerv3/embodied/envs/robodesk.py b/dreamerv3/embodied/envs/robodesk.py new file mode 100644 index 0000000..3226147 --- /dev/null +++ b/dreamerv3/embodied/envs/robodesk.py @@ -0,0 +1,37 @@ +import os + +import embodied + + +class RoboDesk(embodied.Env): + + def __init__(self, task, mode, repeat=1, length=500, resets=True): + assert mode in ('train', 'eval') + # TODO: This env variable is meant for headless GPU machines but may fail + # on CPU-only machines. + if 'MUJOCO_GL' not in os.environ: + os.environ['MUJOCO_GL'] = 'egl' + try: + from robodesk import robodesk + except ImportError: + import robodesk + task, reward = task.rsplit('_', 1) + if mode == 'eval': + reward = 'success' + assert reward in ('dense', 'sparse', 'success'), reward + self._gymenv = robodesk.RoboDesk(task, reward, repeat, length) + from . import from_gym + self._env = from_gym.FromGym(self._gymenv) + + @property + def obs_space(self): + return self._env.obs_space + + @property + def act_space(self): + return self._env.act_space + + def step(self, action): + obs = self._env.step(action) + obs['is_terminal'] = False + return obs diff --git a/dreamerv3/embodied/replay/__init__.py b/dreamerv3/embodied/replay/__init__.py new file mode 100644 index 0000000..647c776 --- /dev/null +++ b/dreamerv3/embodied/replay/__init__.py @@ -0,0 +1,10 @@ +from .generic import Generic +from .reverb import Reverb +from .replays import Uniform +from .naive_chunks import NaiveChunks +from .curious_replay import CuriousReplay +from .prioritized_experience_replay import PrioritizedExperienceReplay +from .count_based import CountBasedReplay +from .adversarial import AdversarialReplay +from . import selectors +from . import limiters diff --git a/dreamerv3/embodied/replay/adversarial.py b/dreamerv3/embodied/replay/adversarial.py new file mode 100644 index 0000000..29a3930 --- /dev/null +++ b/dreamerv3/embodied/replay/adversarial.py @@ -0,0 +1,12 @@ +import numpy as np +from dreamerv3.embodied.replay import CuriousReplay + + +class AdversarialReplay(CuriousReplay): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.should_track_visit_counts = False + + @staticmethod + def _calculate_priority_score(model_loss, visit_count, hyper): + return np.power((model_loss + hyper['epsilon']), hyper['alpha']) diff --git a/dreamerv3/embodied/replay/base_prioritized_reverb.py b/dreamerv3/embodied/replay/base_prioritized_reverb.py new file mode 100644 index 0000000..d3b752a --- /dev/null +++ b/dreamerv3/embodied/replay/base_prioritized_reverb.py @@ -0,0 +1,205 @@ +import pickle +from abc import abstractmethod +from collections import defaultdict, deque +from functools import partial as bind + +import embodied +import numpy as np +import tensorflow as tf + + +class BasePrioritizedReverb: + + def __init__( + self, length, capacity=None, directory=None, chunks=None, flush=100, hyper=None): + del chunks + import reverb + self.length = length + self.capacity = capacity + self.directory = directory and embodied.Path(directory) + self.checkpointer = None + self.server = None + self.client = None + self.writers = None + self.counters = None + self.signature = None + self.flush = flush + + self.hyper = hyper + self.should_track_visit_counts = False + + # Constants + self.priority_scalar = 10.0 # Used to scale all priorities. Avoids reverb precision issue. + self.maximum_attempts_to_find_key = 10000 + max_steps = int(self.capacity * 2) + + self.step_to_keyA = np.zeros((max_steps, ), dtype=np.uint32) + self.step_to_keyB = np.zeros((max_steps, ), dtype=np.uint32) + self.visit_count = np.zeros((max_steps, ), dtype=np.uint32) + + self.env_step_count = defaultdict(int) + self.queue = deque(maxlen=2 * flush) + + if self.directory: + self.directory.mkdirs() + path = str(self.directory) + try: + self.checkpointer = reverb.checkpointers.DefaultCheckpointer(path) + except AttributeError: + self.checkpointer = reverb.checkpointers.RecordIOCheckpointer(path) + self.sigpath = self.directory.parent / (self.directory.name + '_sig.pkl') + if self.directory and self.sigpath.exists(): + with self.sigpath.open('rb') as file: + self.signature = pickle.load(file) + self._create_server() + + def _create_server(self): + import reverb + import tensorflow as tf + self.server = reverb.Server(tables=[reverb.Table( + name='table', + sampler=reverb.selectors.Prioritized(1.0), + remover=reverb.selectors.Fifo(), + max_size=int(self.capacity), + rate_limiter=reverb.rate_limiters.MinSize(1), + signature={ + key: tf.TensorSpec(shape, dtype) + for key, (shape, dtype) in self.signature.items()}, + )], port=None, checkpointer=self.checkpointer) + self.client = reverb.Client(f'localhost:{self.server.port}') + self.writers = defaultdict(bind( + self.client.trajectory_writer, self.length)) + self.counters = defaultdict(int) + + def __len__(self): + if not self.client: + return 0 + return self.client.server_info()['table'].current_size + + @property + def stats(self): + return {'size': len(self)} + + def add(self, step, worker=0): + step = {k: v for k, v in step.items() if not k.startswith('log_')} + step = {k: embodied.convert(v) for k, v in step.items()} + step['id'] = np.asarray(embodied.uuid(step.get('id'))) + step['env_step'] = np.asarray(self.env_step_count[worker]) + step['worker'] = np.asarray(worker) + if not self.server: + self.signature = { + k: ((self.length, *v.shape), v.dtype) for k, v in step.items()} + self._create_server() + + step = {k: v for k, v in step.items() if not k.startswith('log_')} + writer = self.writers[worker] + self.queue.append(step) + + if (self.env_step_count[worker] + 1) < self.length: + writer.append(self.queue.popleft()) + + else: + self.counters[worker] += 1 + if self.counters[worker] >= self.flush: + for i in range(self.flush): + writer.append(self.queue.popleft()) + seq = {k: v[-self.length:] for k, v in writer.history.items()} + writer.create_item('table', priority=self.hyper['key_find_priority'], trajectory=seq) + self.counters[worker] = 0 + writer.flush() + self._find_keys_up_to_step(step['env_step']) + + self.env_step_count[worker] += 1 + + def _find_keys_up_to_step(self, fill_to_step): + """Find the key for all steps just created in the table so that we can set their priorities later. + The keys are likely to be sampled because they are given key_find_priority initially. This is set to the + initial_priority after the keys are found.""" + + import reverb + + dataset = reverb.TrajectoryDataset.from_table_signature( + server_address=f'localhost:{self.server.port}', + table='table', + max_in_flight_samples_per_worker=10, + ) + + found_so_far = np.zeros((int(self.flush), ), dtype=np.uint8) + fill_start_step = fill_to_step - self.flush + 1 + + priorities_to_set = {} + attempts = 0 + + for sample in dataset: + seq = sample.data + step_sampled = int(seq['env_step'][-1]) + if step_sampled >= fill_start_step: + key = sample.info.key + self.step_to_keyA[step_sampled], self.step_to_keyB[step_sampled] = self._split_key(key) + priorities_to_set[int(key)] = self.hyper['initial_priority'] / self.priority_scalar + found_so_far[step_sampled - fill_start_step] = 1 + + if np.all(found_so_far): + break + + attempts += 1 + if attempts > self.maximum_attempts_to_find_key: + raise Exception(f'dreamerv3/embodied/replay/reverb.py: _fill_step_to_key -> ' + f'did not find env_step in {self.maximum_attempts_to_find_key} attempts') + + self.client.mutate_priorities('table', priorities_to_set) + + def dataset(self): + import reverb + dataset = reverb.TrajectoryDataset.from_table_signature( + server_address=f'localhost:{self.server.port}', + table='table', + max_in_flight_samples_per_worker=1, + num_workers_per_iterator=1, + max_samples_per_stream=1,) + for sample in dataset: + seq = sample.data + seq = {k: embodied.convert(v) for k, v in seq.items()} + seq['keyA'], seq['keyB'] = self._split_key(sample.info.key) + seq['key'] = (seq['keyA'], seq['keyB']) + seq['probability'] = sample.info.probability + seq['priority'] = sample.info.priority + seq['times_sampled'] = sample.info.times_sampled + + if 'is_first' in seq: + seq['is_first'] = np.array(seq['is_first']) + seq['is_first'][0] = True + + yield seq + + def _split_key(self, key): + """Split the uint64 key into two 32 bit ints""" + keyA_tf = key // tf.constant(2 ** 32, dtype=tf.uint64) + keyB_tf = key % tf.constant(2 ** 32, dtype=tf.uint64) + return np.uint32(keyA_tf), np.uint32(keyB_tf) + + def _combine_key(self, keyA, keyB) -> tf.uint64: + """Combine the two 32bit ints into a single 64bit int""" + keyA_tf = tf.convert_to_tensor(keyA, dtype=tf.uint64) + keyB_tf = tf.convert_to_tensor(keyB, dtype=tf.uint64) + + return keyA_tf * tf.constant(2 ** 32, dtype=tf.uint64) + keyB_tf + + def update_visit_count(self, env_steps): + flat_env_steps = env_steps.flatten() + self.visit_count[flat_env_steps] += 1 + + @abstractmethod + def prioritize(self, key, env_steps, losses, td_error): + pass + + def save(self, wait=False): + for writer in self.writers.values(): + writer.flush() + with self.sigpath.open('wb') as file: + file.write(pickle.dumps(self.signature)) + if self.directory: + self.client.checkpoint() + + def load(self, data=None): + pass diff --git a/dreamerv3/embodied/replay/chunk.py b/dreamerv3/embodied/replay/chunk.py new file mode 100644 index 0000000..d9e14a8 --- /dev/null +++ b/dreamerv3/embodied/replay/chunk.py @@ -0,0 +1,78 @@ +import io +from datetime import datetime + +import embodied +import numpy as np + + +class Chunk: + + def __init__(self, size, successor=None): + now = datetime.now() + self.time = now.strftime("%Y%m%dT%H%M%S") + f'F{now.microsecond:06d}' + self.uuid = str(embodied.uuid()) + self.successor = successor + self.size = size + self.data = None + self.length = 0 + + def __repr__(self): + succ = self.successor or str(embodied.uuid(0)) + succ = succ.uuid if isinstance(succ, type(self)) else succ + return ( + f'Chunk(uuid={self.uuid}, ' + f'succ={succ}, ' + f'len={self.length})') + + def __len__(self): + return self.length + + def __bool__(self): + return True + + def append(self, step): + if not self.data: + example = {k: embodied.convert(v) for k, v in step.items()} + self.data = { + k: np.empty((self.size,) + v.shape, v.dtype) + for k, v in example.items()} + for key, value in step.items(): + self.data[key][self.length] = value + self.length += 1 + + def save(self, directory): + succ = self.successor or str(embodied.uuid(0)) + succ = succ.uuid if isinstance(succ, type(self)) else succ + filename = f'{self.time}-{self.uuid}-{succ}-{self.length}.npz' + filename = embodied.Path(directory) / filename + data = {k: embodied.convert(v) for k, v in self.data.items()} + with io.BytesIO() as stream: + np.savez_compressed(stream, **data) + stream.seek(0) + filename.write(stream.read(), mode='wb') + print(f'Saved chunk: {filename.name}') + + @classmethod + def load(cls, filename): + length = int(filename.stem.split('-')[3]) + with embodied.Path(filename).open('rb') as f: + data = np.load(f) + data = {k: data[k] for k in data.keys()} + chunk = cls(length) + chunk.time = filename.stem.split('-')[0] + chunk.uuid = filename.stem.split('-')[1] + chunk.successor = filename.stem.split('-')[2] + chunk.length = length + chunk.data = data + return chunk + + @classmethod + def scan(cls, directory, capacity=None, shorten=0): + directory = embodied.Path(directory) + filenames, total = [], 0 + for filename in reversed(sorted(directory.glob('*.npz'))): + if capacity and total >= capacity: + break + filenames.append(filename) + total += max(0, int(filename.stem.split('-')[3]) - shorten) + return sorted(filenames) diff --git a/dreamerv3/embodied/replay/count_based.py b/dreamerv3/embodied/replay/count_based.py new file mode 100644 index 0000000..1fd0a53 --- /dev/null +++ b/dreamerv3/embodied/replay/count_based.py @@ -0,0 +1,12 @@ +import numpy as np +from dreamerv3.embodied.replay import CuriousReplay + + +class CountBasedReplay(CuriousReplay): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.should_track_visit_counts = True + + @staticmethod + def _calculate_priority_score(model_loss, visit_count, hyper): + return hyper['c'] * np.power(hyper['beta'], visit_count) + hyper['epsilon'] diff --git a/dreamerv3/embodied/replay/curious_replay.py b/dreamerv3/embodied/replay/curious_replay.py new file mode 100644 index 0000000..03b83a6 --- /dev/null +++ b/dreamerv3/embodied/replay/curious_replay.py @@ -0,0 +1,25 @@ +import numpy as np +from dreamerv3.embodied.replay.base_prioritized_reverb import BasePrioritizedReverb + + +class CuriousReplay(BasePrioritizedReverb): + + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.should_track_visit_counts = True + + @staticmethod + def _calculate_priority_score(model_loss, visit_count, hyper): + return (hyper['c'] * np.power(hyper['beta'], visit_count)) \ + + np.power((model_loss + hyper['epsilon']), hyper['alpha']) + + def prioritize(self, key, env_steps, losses, td_error): + flat_steps = env_steps.flatten() + flat_losses = losses.flatten() + flat_count = self.visit_count[flat_steps] + flat_priority = self._calculate_priority_score(flat_losses, + flat_count, + self.hyper) / self.priority_scalar + flat_keys = self._combine_key(self.step_to_keyA[flat_steps], self.step_to_keyB[flat_steps]) + flat_updates = {int(k): p for k, p in zip(flat_keys, flat_priority)} + self.client.mutate_priorities('table', flat_updates) diff --git a/dreamerv3/embodied/replay/generic.py b/dreamerv3/embodied/replay/generic.py new file mode 100644 index 0000000..3ae64d5 --- /dev/null +++ b/dreamerv3/embodied/replay/generic.py @@ -0,0 +1,162 @@ +import time +from collections import defaultdict, deque +from functools import partial as bind + +import embodied +import numpy as np + +from . import saver + + +class Generic: + + def __init__( + self, length, capacity, remover, sampler, limiter, directory, + overlap=None, online=False, chunks=1024): + assert capacity is None or 1 <= capacity + self.length = length + self.capacity = capacity + self.remover = remover + self.sampler = sampler + self.limiter = limiter + self.stride = 1 if overlap is None else length - overlap + self.streams = defaultdict(bind(deque, maxlen=length)) + self.counters = defaultdict(int) + self.table = {} + self.online = online + if self.online: + self.online_queue = deque() + self.online_stride = length + self.online_counters = defaultdict(int) + self.saver = directory and saver.Saver(directory, chunks) + self.metrics = { + 'samples': 0, + 'sample_wait_dur': 0, + 'sample_wait_count': 0, + 'inserts': 0, + 'insert_wait_dur': 0, + 'insert_wait_count': 0, + } + self.load() + + def __len__(self): + return len(self.table) + + @property + def stats(self): + ratio = lambda x, y: x / y if y else np.nan + m = self.metrics + stats = { + 'size': len(self), + 'inserts': m['inserts'], + 'samples': m['samples'], + 'insert_wait_avg': ratio(m['insert_wait_dur'], m['inserts']), + 'insert_wait_frac': ratio(m['insert_wait_count'], m['inserts']), + 'sample_wait_avg': ratio(m['sample_wait_dur'], m['samples']), + 'sample_wait_frac': ratio(m['sample_wait_count'], m['samples']), + } + for key in self.metrics: + self.metrics[key] = 0 + return stats + + def add(self, step, worker=0, load=False): + step = {k: v for k, v in step.items() if not k.startswith('log_')} + step['id'] = np.asarray(embodied.uuid(step.get('id'))) + stream = self.streams[worker] + stream.append(step) + self.saver and self.saver.add(step, worker) + self.counters[worker] += 1 + if self.online: + self.online_counters[worker] += 1 + if len(stream) >= self.length and ( + self.online_counters[worker] >= self.online_stride): + self.online_queue.append(tuple(stream)) + self.online_counters[worker] = 0 + if len(stream) < self.length or self.counters[worker] < self.stride: + return + self.counters[worker] = 0 + key = embodied.uuid() + seq = tuple(stream) + if load: + assert self.limiter.want_load()[0] + else: + dur = wait(self.limiter.want_insert, 'Replay insert is waiting') + self.metrics['inserts'] += 1 + self.metrics['insert_wait_dur'] += dur + self.metrics['insert_wait_count'] += int(dur > 0) + self.table[key] = seq + self.remover[key] = seq + self.sampler[key] = seq + while self.capacity and len(self) > self.capacity: + self._remove(self.remover()) + + def _sample(self): + dur = wait(self.limiter.want_sample, 'Replay sample is waiting') + self.metrics['samples'] += 1 + self.metrics['sample_wait_dur'] += dur + self.metrics['sample_wait_count'] += int(dur > 0) + if self.online: + try: + seq = self.online_queue.popleft() + except IndexError: + seq = self.table[self.sampler()] + else: + seq = self.table[self.sampler()] + seq = {k: [step[k] for step in seq] for k in seq[0]} + seq = {k: embodied.convert(v) for k, v in seq.items()} + if 'is_first' in seq: + seq['is_first'][0] = True + return seq + + def _remove(self, key): + wait(self.limiter.want_remove, 'Replay remove is waiting') + del self.table[key] + del self.remover[key] + del self.sampler[key] + + def dataset(self): + while True: + yield self._sample() + + def prioritize(self, ids, prios): + if hasattr(self.sampler, 'prioritize'): + self.sampler.prioritize(ids, prios) + + def save(self, wait=False): + if not self.saver: + return + self.saver.save(wait) + # return { + # 'saver': self.saver.save(wait), + # # 'remover': self.remover.save(wait), + # # 'sampler': self.sampler.save(wait), + # # 'limiter': self.limiter.save(wait), + # } + + def load(self, data=None): + if not self.saver: + return + workers = set() + for step, worker in self.saver.load(self.capacity, self.length): + workers.add(worker) + self.add(step, worker, load=True) + for worker in workers: + del self.streams[worker] + del self.counters[worker] + # self.remover.load(data['remover']) + # self.sampler.load(data['sampler']) + # self.limiter.load(data['limiter']) + + +def wait(predicate, message, sleep=0.001, notify=1.0): + start = time.time() + notified = False + while True: + allowed, detail = predicate() + duration = time.time() - start + if allowed: + return duration + if not notified and duration >= notify: + print(f'{message} ({detail})') + notified = True + time.sleep(sleep) diff --git a/dreamerv3/embodied/replay/limiters.py b/dreamerv3/embodied/replay/limiters.py new file mode 100644 index 0000000..73a2ee2 --- /dev/null +++ b/dreamerv3/embodied/replay/limiters.py @@ -0,0 +1,108 @@ +import threading + + +class MinSize: + + def __init__(self, minimum): + assert 1 <= minimum, minimum + self.minimum = minimum + self.size = 0 + self.lock = threading.Lock() + + def want_load(self): + with self.lock: + self.size += 1 + return True, 'ok' + + def want_insert(self): + with self.lock: + self.size += 1 + return True, 'ok' + + def want_remove(self): + with self.lock: + if self.size < 1: + return False, 'is empty' + self.size -= 1 + return True, 'ok' + + def want_sample(self): + if self.size < self.minimum: + return False, f'too empty: {self.size} < {self.minimum}' + return True, 'ok' + + +class SamplesPerInsert: + + def __init__(self, samples_per_insert, tolerance, minimum=1): + assert 1 <= minimum + self.samples_per_insert = samples_per_insert + self.minimum = minimum + self.avail = -minimum + self.min_avail = -tolerance + self.max_avail = tolerance * samples_per_insert + self.size = 0 + self.lock = threading.Lock() + + def want_load(self): + with self.lock: + self.size += 1 + return True, 'ok' + + def want_insert(self): + with self.lock: + if self.avail >= self.max_avail: + return False, f'rate limited: {self.avail:.3f} >= {self.max_avail:.3f}' + self.avail += self.samples_per_insert + self.size += 1 + return True, 'ok' + + def want_remove(self): + with self.lock: + if self.size < 1: + return False, 'is empty' + self.size -= 1 + return True, 'ok' + + def want_sample(self): + with self.lock: + if self.size < self.minimum: + return False, f'too empty: {self.size} < {self.minimum}' + if self.avail <= self.min_avail: + return False, f'rate limited: {self.avail:.3f} <= {self.min_avail:.3f}' + self.avail -= 1 + return True, 'ok' + + +class Queue: + + def __init__(self, capacity): + assert 1 <= capacity + self.capacity = capacity + self.size = 0 + self.lock = threading.Lock() + + def want_load(self): + with self.lock: + self.size += 1 + return True, 'ok' + + def want_insert(self): + with self.lock: + if self.size >= self.capacity: + return False, f'is full: {self.size} >= {self.capacity}' + self.size += 1 + return True, 'ok' + + def want_remove(self): + with self.lock: + if self.size < 1: + return False, 'is empty' + self.size -= 1 + return True, 'ok' + + def want_sample(self): + if self.size < 1: + return False, 'is empty' + else: + return True, 'ok' diff --git a/dreamerv3/embodied/replay/naive_chunks.py b/dreamerv3/embodied/replay/naive_chunks.py new file mode 100644 index 0000000..f98dbdf --- /dev/null +++ b/dreamerv3/embodied/replay/naive_chunks.py @@ -0,0 +1,82 @@ +import concurrent.futures +import threading +import time +import uuid +from collections import deque, defaultdict +from functools import partial as bind + +import numpy as np +import embodied + +from . import chunk as chunklib + + +class NaiveChunks(embodied.Replay): + + def __init__(self, length, capacity=None, directory=None, chunks=1024, seed=0): + assert 1 <= length <= chunks + self.length = length + self.capacity = capacity + self.directory = directory and embodied.Path(directory) + self.chunks = chunks + self.buffers = {} + self.rng = np.random.default_rng(seed) + self.ongoing = defaultdict(bind(chunklib.Chunk, chunks)) + if directory: + self.directory.mkdirs() + self.workers = concurrent.futures.ThreadPoolExecutor(16) + self.promises = deque() + + def __len__(self): + return len(self.buffers) * self.length + + @property + def stats(self): + return {'size': len(self), 'chunks': len(self.buffers)} + + def add(self, step, worker=0): + chunk = self.ongoing[worker] + chunk.append(step) + if len(chunk) >= self.chunks: + self.buffers[chunk.uuid] = self.ongoing.pop(worker) + self.promises.append(self.workers.submit(chunk.save, self.directory)) + for promise in [x for x in self.promises if x.done()]: + promise.result() + self.promises.remove(promise) + while len(self) > self.capacity: + del self.buffers[next(iter(self.buffers.keys()))] + + def _sample(self): + counter = 0 + while not self.buffers: + if counter % 100 == 0: + print('Replay sample is waiting') + time.sleep(0.1) + counter += 1 + keys = tuple(self.buffers.keys()) + chunk = self.buffers[keys[self.rng.integers(0, len(keys))]] + idx = self.rng.integers(0, len(chunk) - self.length + 1) + seq = {k: chunk.data[k][idx: idx + self.length] for k in chunk.data.keys()} + seq['is_first'][0] = True + return seq + + def dataset(self): + while True: + yield self._sample() + + def save(self, wait=False): + for chunk in self.ongoing.values(): + if chunk.length: + self.promises.append(self.workers.submit(chunk.save, self.directory)) + if wait: + [x.result() for x in self.promises] + self.promises.clear() + + def load(self, data=None): + filenames = chunklib.Chunk.scan(self.directory, capacity) + if not filenames: + return + threads = min(len(filenames), 32) + with concurrent.futures.ThreadPoolExecutor(threads) as executor: + chunks = list(executor.map(chunklib.Chunk.load, filenames)) + self.buffers = {chunk.uuid: chunk for chunk in chunks} diff --git a/dreamerv3/embodied/replay/prioritized_experience_replay.py b/dreamerv3/embodied/replay/prioritized_experience_replay.py new file mode 100644 index 0000000..b5a681b --- /dev/null +++ b/dreamerv3/embodied/replay/prioritized_experience_replay.py @@ -0,0 +1,20 @@ +import numpy as np +from dreamerv3.embodied.replay.base_prioritized_reverb import BasePrioritizedReverb + + +class PrioritizedExperienceReplay(BasePrioritizedReverb): + def __init__(self, *args, **kwargs): + super().__init__(*args, **kwargs) + self.should_track_visit_counts = False + + @staticmethod + def _calculate_priority_score(td_error, hyper): + return np.power(np.abs(td_error) + hyper['epsilon'], hyper['alpha']) + + def prioritize(self, key, env_steps, losses, td_error): + flat_steps = env_steps[:, 1:].flatten() + flat_td = td_error.flatten() + flat_priority = self._calculate_priority_score(flat_td, self.hyper) / self.priority_scalar + flat_keys = self._combine_key(self.step_to_keyA[flat_steps], self.step_to_keyB[flat_steps]) + flat_updates = {int(k): p for k, p in zip(flat_keys, flat_priority)} + self.client.mutate_priorities('table', flat_updates) diff --git a/dreamerv3/embodied/replay/replays.py b/dreamerv3/embodied/replay/replays.py new file mode 100644 index 0000000..9aef203 --- /dev/null +++ b/dreamerv3/embodied/replay/replays.py @@ -0,0 +1,26 @@ +from . import generic +from . import selectors +from . import limiters + + +class Uniform(generic.Generic): + + def __init__( + self, length, capacity=None, directory=None, online=False, chunks=1024, + min_size=1, samples_per_insert=None, tolerance=1e4, seed=0): + if samples_per_insert: + limiter = limiters.SamplesPerInsert( + samples_per_insert, tolerance, min_size) + else: + limiter = limiters.MinSize(min_size) + assert not capacity or min_size <= capacity + super().__init__( + length=length, + capacity=capacity, + remover=selectors.Fifo(), + sampler=selectors.Uniform(seed), + limiter=limiter, + directory=directory, + online=online, + chunks=chunks, + ) diff --git a/dreamerv3/embodied/replay/reverb.py b/dreamerv3/embodied/replay/reverb.py new file mode 100644 index 0000000..86bd35f --- /dev/null +++ b/dreamerv3/embodied/replay/reverb.py @@ -0,0 +1,112 @@ +import pickle +from collections import defaultdict +from functools import partial as bind + +import embodied +import numpy as np + + +class Reverb: + + def __init__( + self, length, capacity=None, directory=None, chunks=None, flush=100): + del chunks + import reverb + self.length = length + self.capacity = capacity + self.directory = directory and embodied.Path(directory) + self.checkpointer = None + self.server = None + self.client = None + self.writers = None + self.counters = None + self.signature = None + self.flush = flush + if self.directory: + self.directory.mkdirs() + path = str(self.directory) + try: + self.checkpointer = reverb.checkpointers.DefaultCheckpointer(path) + except AttributeError: + self.checkpointer = reverb.checkpointers.RecordIOCheckpointer(path) + self.sigpath = self.directory.parent / (self.directory.name + '_sig.pkl') + if self.directory and self.sigpath.exists(): + with self.sigpath.open('rb') as file: + self.signature = pickle.load(file) + self._create_server() + + def _create_server(self): + import reverb + import tensorflow as tf + self.server = reverb.Server(tables=[reverb.Table( + name='table', + sampler=reverb.selectors.Uniform(), + remover=reverb.selectors.Fifo(), + max_size=int(self.capacity), + rate_limiter=reverb.rate_limiters.MinSize(1), + signature={ + key: tf.TensorSpec(shape, dtype) + for key, (shape, dtype) in self.signature.items()}, + )], port=None, checkpointer=self.checkpointer) + self.client = reverb.Client(f'localhost:{self.server.port}') + self.writers = defaultdict(bind( + self.client.trajectory_writer, self.length)) + self.counters = defaultdict(int) + + def __len__(self): + if not self.client: + return 0 + return self.client.server_info()['table'].current_size + + @property + def stats(self): + return {'size': len(self)} + + def add(self, step, worker=0): + step = {k: v for k, v in step.items() if not k.startswith('log_')} + step = {k: embodied.convert(v) for k, v in step.items()} + step['id'] = np.asarray(embodied.uuid(step.get('id'))) + if not self.server: + self.signature = { + k: ((self.length, *v.shape), v.dtype) for k, v in step.items()} + self._create_server() + step = {k: v for k, v in step.items() if not k.startswith('log_')} + writer = self.writers[worker] + writer.append(step) + if len(next(iter(writer.history.values()))) >= self.length: + seq = {k: v[-self.length:] for k, v in writer.history.items()} + writer.create_item('table', priority=1.0, trajectory=seq) + self.counters[worker] += 1 + if self.counters[worker] > self.flush: + self.counters[worker] = 0 + writer.flush() + + def dataset(self): + import reverb + dataset = reverb.TrajectoryDataset.from_table_signature( + server_address=f'localhost:{self.server.port}', + table='table', + max_in_flight_samples_per_worker=10) + for sample in dataset: + seq = sample.data + seq = {k: embodied.convert(v) for k, v in seq.items()} + # seq['key'] = sample.info.key # uint64 + # seq['prob'] = sample.info.probability + if 'is_first' in seq: + seq['is_first'] = np.array(seq['is_first']) + seq['is_first'][0] = True + yield seq + + def prioritize(self, ids, prios): + raise NotImplementedError + + def save(self, wait=False): + for writer in self.writers.values(): + writer.flush() + with self.sigpath.open('wb') as file: + file.write(pickle.dumps(self.signature)) + if self.directory: + self.client.checkpoint() + + def load(self, data=None): + pass diff --git a/dreamerv3/embodied/replay/saver.py b/dreamerv3/embodied/replay/saver.py new file mode 100644 index 0000000..1b7958a --- /dev/null +++ b/dreamerv3/embodied/replay/saver.py @@ -0,0 +1,63 @@ +import concurrent.futures +from collections import defaultdict, deque +from functools import partial as bind + +import embodied + +from . import chunk as chunklib + + +class Saver: + + def __init__(self, directory, chunks=1024): + self.directory = embodied.Path(directory) + self.directory.mkdirs() + self.chunks = chunks + self.buffers = defaultdict(bind(chunklib.Chunk, chunks)) + self.workers = concurrent.futures.ThreadPoolExecutor(16) + self.promises = deque() + self.loading = False + + def add(self, step, worker): + if self.loading: + return + buffer = self.buffers[worker] + buffer.append(step) + if buffer.length >= self.chunks: + self.buffers[worker] = buffer.successor = chunklib.Chunk(self.chunks) + self.promises.append(self.workers.submit(buffer.save, self.directory)) + for promise in [x for x in self.promises if x.done()]: + promise.result() + self.promises.remove(promise) + + def save(self, wait=False): + for buffer in self.buffers.values(): + if buffer.length: + self.promises.append(self.workers.submit(buffer.save, self.directory)) + if wait: + [x.result() for x in self.promises] + self.promises.clear() + + def load(self, capacity, length): + filenames = chunklib.Chunk.scan(self.directory, capacity, length - 1) + if not filenames: + return + threads = min(len(filenames), 32) + with concurrent.futures.ThreadPoolExecutor(threads) as executor: + chunks = list(executor.map(chunklib.Chunk.load, filenames)) + streamids = {} + for chunk in reversed(sorted(chunks, key=lambda x: x.time)): + if chunk.successor not in streamids: + streamids[chunk.uuid] = int(embodied.uuid()) + else: + streamids[chunk.uuid] = streamids[chunk.successor] + self.loading = True + for i, chunk in enumerate(chunks): + stream = streamids[chunk.uuid] + for index in range(chunk.length): + step = {k: v[index] for k, v in chunk.data.items()} + yield step, stream + # Free memory early to not require twice the replay capacity. + chunks[i] = None + del chunk + self.loading = False diff --git a/dreamerv3/embodied/replay/selectors.py b/dreamerv3/embodied/replay/selectors.py new file mode 100644 index 0000000..51700c9 --- /dev/null +++ b/dreamerv3/embodied/replay/selectors.py @@ -0,0 +1,45 @@ +from collections import deque + +import numpy as np + + +class Fifo: + + def __init__(self): + self.queue = deque() + + def __call__(self): + return self.queue[0] + + def __setitem__(self, key, steps): + self.queue.append(key) + + def __delitem__(self, key): + if self.queue[0] == key: + self.queue.popleft() + else: + # TODO: This branch is unused but very slow. + self.queue.remove(key) + + +class Uniform: + + def __init__(self, seed=0): + self.indices = {} + self.keys = [] + self.rng = np.random.default_rng(seed) + + def __call__(self): + index = self.rng.integers(0, len(self.keys)).item() + return self.keys[index] + + def __setitem__(self, key, steps): + self.indices[key] = len(self.keys) + self.keys.append(key) + + def __delitem__(self, key): + index = self.indices.pop(key) + last = self.keys.pop() + if index != len(self.keys): + self.keys[index] = last + self.indices[last] = index diff --git a/dreamerv3/embodied/run/__init__.py b/dreamerv3/embodied/run/__init__.py new file mode 100644 index 0000000..c6a11dd --- /dev/null +++ b/dreamerv3/embodied/run/__init__.py @@ -0,0 +1,6 @@ +from .eval_only import eval_only +from .parallel import parallel +from .train import train +from .train_eval import train_eval +from .train_holdout import train_holdout +from .train_save import train_save diff --git a/dreamerv3/embodied/run/eval_only.py b/dreamerv3/embodied/run/eval_only.py new file mode 100644 index 0000000..2db3050 --- /dev/null +++ b/dreamerv3/embodied/run/eval_only.py @@ -0,0 +1,61 @@ +import re + +import embodied +import numpy as np + + +def eval_only(agent, env, logger, args): + + logdir = embodied.Path(args.logdir) + logdir.mkdirs() + print('Logdir', logdir) + should_log = embodied.when.Clock(args.log_every) + step = logger.step + metrics = embodied.Metrics() + print('Observation space:', env.obs_space) + print('Action space:', env.act_space) + + timer = embodied.Timer() + timer.wrap('agent', agent, ['policy']) + timer.wrap('env', env, ['step']) + timer.wrap('logger', logger, ['write']) + + nonzeros = set() + def per_episode(ep): + length = len(ep['reward']) - 1 + score = float(ep['reward'].astype(np.float64).sum()) + logger.add({'length': length, 'score': score}, prefix='episode') + print(f'Episode has {length} steps and return {score:.1f}.') + stats = {} + for key in args.log_keys_video: + if key in ep: + stats[f'policy_{key}'] = ep[key] + for key, value in ep.items(): + if not args.log_zeros and key not in nonzeros and (value == 0).all(): + continue + nonzeros.add(key) + if re.match(args.log_keys_sum, key): + stats[f'sum_{key}'] = ep[key].sum() + if re.match(args.log_keys_mean, key): + stats[f'mean_{key}'] = ep[key].mean() + if re.match(args.log_keys_max, key): + stats[f'max_{key}'] = ep[key].max(0).mean() + metrics.add(stats, prefix='stats') + + driver = embodied.Driver(env) + driver.on_episode(lambda ep, worker: per_episode(ep)) + driver.on_step(lambda tran, _: step.increment()) + + checkpoint = embodied.Checkpoint() + checkpoint.agent = agent + checkpoint.load(args.from_checkpoint, keys=['agent']) + + print('Start evaluation loop.') + policy = lambda *args: agent.policy(*args, mode='eval') + while step < args.steps: + driver(policy, steps=100) + if should_log(step): + logger.add(metrics.result()) + logger.add(timer.stats(), prefix='timer') + logger.write(fps=True) + logger.write() diff --git a/dreamerv3/embodied/run/parallel.py b/dreamerv3/embodied/run/parallel.py new file mode 100644 index 0000000..1adab12 --- /dev/null +++ b/dreamerv3/embodied/run/parallel.py @@ -0,0 +1,164 @@ +import sys +import time +from collections import defaultdict + +import embodied +import numpy as np + + +def parallel(agent, replay, logger, make_env, num_envs, args): + step = logger.step + timer = embodied.Timer() + timer.wrap('agent', agent, ['policy', 'train', 'report', 'save']) + timer.wrap('replay', replay, ['add', 'save']) + timer.wrap('logger', logger, ['write']) + workers = [] + workers.append(embodied.distr.Thread( + actor, step, agent, replay, logger, args.actor_addr, args)) + workers.append(embodied.distr.Thread( + learner, step, agent, replay, logger, timer, args)) + if num_envs == 1: + workers.append(embodied.distr.Thread( + env, make_env, args.actor_addr, 0, args, timer)) + else: + for i in range(num_envs): + workers.append(embodied.distr.Process( + env, make_env, args.actor_addr, i, args)) + embodied.distr.run(workers) + + +def actor(step, agent, replay, logger, actor_addr, args): + metrics = embodied.Metrics() + scalars = defaultdict(lambda: defaultdict(list)) + videos = defaultdict(lambda: defaultdict(list)) + should_log = embodied.when.Clock(args.log_every) + + _, initial = agent.policy(dummy_data( + agent.agent.obs_space, (args.actor_batch,))) + initial = embodied.treemap(lambda x: x[0], initial) + allstates = defaultdict(lambda: initial) + agent.sync() + + def callback(obs, env_addrs): + states = [allstates[a] for a in env_addrs] + states = embodied.treemap(lambda *xs: list(xs), *states) + act, states = agent.policy(obs, states) + act['reset'] = obs['is_last'].copy() + for i, a in enumerate(env_addrs): + allstates[a] = embodied.treemap(lambda x: x[i], states) + + trans = {**obs, **act} + for i, a in enumerate(env_addrs): + tran = {k: v[i].copy() for k, v in trans.items()} + replay.add(tran.copy(), worker=a) + [scalars[a][k].append(v) for k, v in tran.items() if v.size == 1] + [videos[a][k].append(tran[k]) for k in args.log_keys_video] + step.increment(args.actor_batch) + + for i, a in enumerate(env_addrs): + if not trans['is_last'][i]: + continue + ep = {**scalars.pop(a), **videos.pop(a)} + ep = {k: embodied.convert(v) for k, v in ep.items()} + logger.add({ + 'length': len(ep['reward']) - 1, + 'score': sum(ep['reward']), + }, prefix='episode') + stats = {} + for key in args.log_keys_video: + stats[f'policy_{key}'] = ep[key] + metrics.add(stats, prefix='stats') + + if should_log(): + logger.add(metrics.result()) + + return act + + print('[actor] Start server') + embodied.BatchServer(actor_addr, args.actor_batch, callback).run() + + +def learner(step, agent, replay, logger, timer, args): + logdir = embodied.Path(args.logdir) + metrics = embodied.Metrics() + should_log = embodied.when.Clock(args.log_every) + should_save = embodied.when.Clock(args.save_every) + should_sync = embodied.when.Every(args.sync_every) + updates = embodied.Counter() + + checkpoint = embodied.Checkpoint(logdir / 'checkpoint.ckpt') + checkpoint.step = step + checkpoint.agent = agent + checkpoint.replay = replay + if args.from_checkpoint: + checkpoint.load(args.from_checkpoint) + checkpoint.load_or_save() + + dataset = agent.dataset(replay.dataset) + state = None + stats = dict(last_time=time.time(), last_step=int(step), batch_entries=0) + while True: + batch = next(dataset) + outs, state, mets = agent.train(batch, state) + metrics.add(mets) + updates.increment() + stats['batch_entries'] += batch['is_first'].size + + if should_sync(updates): + agent.sync() + + if should_log(): + train = metrics.result() + report = agent.report(batch) + report = {k: v for k, v in report.items() if 'train/' + k not in train} + logger.add(train, prefix='train') + logger.add(report, prefix='report') + logger.add(timer.stats(), prefix='timer') + logger.add(replay.stats, prefix='replay') + + duration = time.time() - stats['last_time'] + actor_fps = (int(step) - stats['last_step']) / duration + learner_fps = stats['batch_entries'] / duration + logger.add({ + 'actor_fps': actor_fps, + 'learner_fps': learner_fps, + 'train_ratio': learner_fps / actor_fps if actor_fps else np.inf, + }, prefix='parallel') + stats = dict(last_time=time.time(), last_step=int(step), batch_entries=0) + + logger.write(fps=True) + + if should_save(): + checkpoint.save() + + +def env(make_env, actor_addr, i, args, timer=None): + # TODO: Optionally write NPZ episodes. + print(f'[env{i}] Make env') + env = make_env() + if timer: + timer.wrap('env', env, ['step']) + actor = embodied.Client(actor_addr) + act = {k: v.sample() for k, v in env.act_space.items()} + done = False + while True: + act['reset'] = done + obs = env.step(act) + obs = {k: np.asarray(v) for k, v in obs.items()} + done = obs['is_last'] + promise = actor(obs) + try: + act = promise() + except RuntimeError: + sys.exit(0) + act = {k: v for k, v in act.items() if not k.startswith('log_')} + + +def dummy_data(spaces, batch_dims): + # TODO: Get rid of this function by adding initial_policy_state() and + # initial_train_state() to the agent API. + spaces = list(spaces.items()) + data = {k: np.zeros(v.shape, v.dtype) for k, v in spaces} + for dim in reversed(batch_dims): + data = {k: np.repeat(v[None], dim, axis=0) for k, v in data.items()} + return data diff --git a/dreamerv3/embodied/run/train.py b/dreamerv3/embodied/run/train.py new file mode 100644 index 0000000..53898f6 --- /dev/null +++ b/dreamerv3/embodied/run/train.py @@ -0,0 +1,120 @@ +import re + +import embodied +import jax +import numpy as np + + +def train(agent, env, replay, logger, args): + + logdir = embodied.Path(args.logdir) + logdir.mkdirs() + print('Logdir', logdir) + should_expl = embodied.when.Until(args.expl_until) + should_train = embodied.when.Ratio(args.train_ratio / args.batch_steps) + should_log = embodied.when.Clock(args.log_every) + should_save = embodied.when.Clock(args.save_every) + should_sync = embodied.when.Every(args.sync_every) + step = logger.step + updates = embodied.Counter() + metrics = embodied.Metrics() + print('Observation space:', embodied.format(env.obs_space), sep='\n') + print('Action space:', embodied.format(env.act_space), sep='\n') + + timer = embodied.Timer() + timer.wrap('agent', agent, ['policy', 'train', 'report', 'save']) + timer.wrap('env', env, ['step']) + timer.wrap('replay', replay, ['add', 'save']) + timer.wrap('logger', logger, ['write']) + + nonzeros = set() + def per_episode(ep): + length = len(ep['reward']) - 1 + score = float(ep['reward'].astype(np.float64).sum()) + sum_abs_reward = float(np.abs(ep['reward']).astype(np.float64).sum()) + logger.add({ + 'length': length, + 'score': score, + 'sum_abs_reward': sum_abs_reward, + 'reward_rate': (np.abs(ep['reward']) >= 0.5).mean(), + }, prefix='episode') + print(f'Episode has {length} steps and return {score:.1f}.') + stats = {} + for key in args.log_keys_video: + if key in ep: + stats[f'policy_{key}'] = ep[key] + for key, value in ep.items(): + if not args.log_zeros and key not in nonzeros and (value == 0).all(): + continue + nonzeros.add(key) + if re.match(args.log_keys_sum, key): + stats[f'sum_{key}'] = ep[key].sum() + if re.match(args.log_keys_mean, key): + stats[f'mean_{key}'] = ep[key].mean() + if re.match(args.log_keys_max, key): + stats[f'max_{key}'] = ep[key].max(0).mean() + metrics.add(stats, prefix='stats') + + driver = embodied.Driver(env) + driver.on_episode(lambda ep, worker: per_episode(ep)) + driver.on_step(lambda tran, _: step.increment()) + driver.on_step(replay.add) + + print('Prefill train dataset.') + random_agent = embodied.RandomAgent(env.act_space) + while len(replay) < max(args.batch_steps, args.train_fill): + driver(random_agent.policy, steps=100) + logger.add(metrics.result()) + logger.write() + + dataset = agent.dataset(replay.dataset) + state = [None] # To be writable from train step function below. + batch = [None] + def train_step(tran, worker): + for _ in range(should_train(step)): + with timer.scope('dataset'): + batch[0] = next(dataset) + outs, state[0], mets = agent.train(batch[0], state[0]) + metrics.add(mets, prefix='train') + + if getattr(replay, 'update_visit_count', False): + replay.update_visit_count(jax.device_get(batch[0]['env_step'])) + + if 'key' in outs: + replay.prioritize(outs['key'], + outs['env_step'], + outs['model_loss'], + outs['td_error']) + + updates.increment() + if should_sync(updates): + agent.sync() + if should_log(step): + agg = metrics.result() + report = agent.report(batch[0]) + report = {k: v for k, v in report.items() if 'train/' + k not in agg} + logger.add(agg) + logger.add(report, prefix='report') + logger.add(replay.stats, prefix='replay') + logger.add(timer.stats(), prefix='timer') + logger.write(fps=True) + driver.on_step(train_step) + + checkpoint = embodied.Checkpoint(logdir / 'checkpoint.ckpt') + timer.wrap('checkpoint', checkpoint, ['save', 'load']) + checkpoint.step = step + checkpoint.agent = agent + checkpoint.replay = replay + if args.from_checkpoint: + checkpoint.load(args.from_checkpoint) + checkpoint.load_or_save() + should_save(step) # Register that we jused saved. + + print('Start training loop.') + policy = lambda *args: agent.policy( + *args, mode='explore' if should_expl(step) else 'train') + while step < args.steps: + driver(policy, steps=100) + if should_save(step): + checkpoint.save() + logger.write() diff --git a/dreamerv3/embodied/run/train_eval.py b/dreamerv3/embodied/run/train_eval.py new file mode 100644 index 0000000..0de337e --- /dev/null +++ b/dreamerv3/embodied/run/train_eval.py @@ -0,0 +1,124 @@ +import re + +import embodied +import numpy as np + + +def train_eval( + agent, train_env, eval_env, train_replay, eval_replay, logger, args): + + logdir = embodied.Path(args.logdir) + logdir.mkdirs() + print('Logdir', logdir) + should_expl = embodied.when.Until(args.expl_until) + should_train = embodied.when.Ratio(args.train_ratio / args.batch_steps) + should_log = embodied.when.Clock(args.log_every) + should_save = embodied.when.Clock(args.save_every) + should_eval = embodied.when.Every(args.eval_every, args.eval_initial) + should_sync = embodied.when.Every(args.sync_every) + step = logger.step + updates = embodied.Counter() + metrics = embodied.Metrics() + print('Observation space:', embodied.format(train_env.obs_space), sep='\n') + print('Action space:', embodied.format(train_env.act_space), sep='\n') + + timer = embodied.Timer() + timer.wrap('agent', agent, ['policy', 'train', 'report', 'save']) + timer.wrap('env', train_env, ['step']) + if hasattr(train_replay, '_sample'): + timer.wrap('replay', train_replay, ['_sample']) + + nonzeros = set() + def per_episode(ep, mode): + length = len(ep['reward']) - 1 + score = float(ep['reward'].astype(np.float64).sum()) + logger.add({ + 'length': length, 'score': score, + 'reward_rate': (ep['reward'] - ep['reward'].min() >= 0.1).mean(), + }, prefix=('episode' if mode == 'train' else f'{mode}_episode')) + print(f'Episode has {length} steps and return {score:.1f}.') + stats = {} + for key in args.log_keys_video: + if key in ep: + stats[f'policy_{key}'] = ep[key] + for key, value in ep.items(): + if not args.log_zeros and key not in nonzeros and (value == 0).all(): + continue + nonzeros.add(key) + if re.match(args.log_keys_sum, key): + stats[f'sum_{key}'] = ep[key].sum() + if re.match(args.log_keys_mean, key): + stats[f'mean_{key}'] = ep[key].mean() + if re.match(args.log_keys_max, key): + stats[f'max_{key}'] = ep[key].max(0).mean() + metrics.add(stats, prefix=f'{mode}_stats') + + driver_train = embodied.Driver(train_env) + driver_train.on_episode(lambda ep, worker: per_episode(ep, mode='train')) + driver_train.on_step(lambda tran, _: step.increment()) + driver_train.on_step(train_replay.add) + driver_eval = embodied.Driver(eval_env) + driver_eval.on_step(eval_replay.add) + driver_eval.on_episode(lambda ep, worker: per_episode(ep, mode='eval')) + + random_agent = embodied.RandomAgent(train_env.act_space) + print('Prefill train dataset.') + while len(train_replay) < max(args.batch_steps, args.train_fill): + driver_train(random_agent.policy, steps=100) + print('Prefill eval dataset.') + while len(eval_replay) < max(args.batch_steps, args.eval_fill): + driver_eval(random_agent.policy, steps=100) + logger.add(metrics.result()) + logger.write() + + dataset_train = agent.dataset(train_replay.dataset) + dataset_eval = agent.dataset(eval_replay.dataset) + state = [None] # To be writable from train step function below. + batch = [None] + def train_step(tran, worker): + for _ in range(should_train(step)): + with timer.scope('dataset_train'): + batch[0] = next(dataset_train) + outs, state[0], mets = agent.train(batch[0], state[0]) + metrics.add(mets, prefix='train') + if 'priority' in outs: + train_replay.prioritize(outs['key'], outs['priority']) + updates.inc() + if should_sync(updates): + agent.sync() + if should_log(step): + logger.add(metrics.result()) + logger.add(agent.report(batch[0]), prefix='report') + with timer.scope('dataset_eval'): + eval_batch = next(dataset_eval) + logger.add(agent.report(eval_batch), prefix='eval') + logger.add(train_replay.stats, prefix='replay') + logger.add(eval_replay.stats, prefix='eval_replay') + logger.add(timer.stats(), prefix='timer') + logger.write(fps=True) + driver_train.on_step(train_step) + + checkpoint = embodied.Checkpoint(logdir / 'checkpoint.ckpt') + checkpoint.step = step + checkpoint.agent = agent + checkpoint.train_replay = train_replay + checkpoint.eval_replay = eval_replay + if args.from_checkpoint: + checkpoint.load(args.from_checkpoint) + checkpoint.load_or_save() + should_save(step) # Register that we jused saved. + + print('Start training loop.') + policy_train = lambda *args: agent.policy( + *args, mode='explore' if should_expl(step) else 'train') + policy_eval = lambda *args: agent.policy(*args, mode='eval') + while step < args.steps: + if should_eval(step): + print('Starting evaluation at step', int(step)) + driver_eval.reset() + driver_eval(policy_eval, episodes=max(len(eval_env), args.eval_eps)) + driver_train(policy_train, steps=100) + if should_save(step): + checkpoint.save() + logger.write() + logger.write() diff --git a/dreamerv3/embodied/run/train_holdout.py b/dreamerv3/embodied/run/train_holdout.py new file mode 100644 index 0000000..35b3a01 --- /dev/null +++ b/dreamerv3/embodied/run/train_holdout.py @@ -0,0 +1,127 @@ +import re + +import embodied +import numpy as np + + +def train_holdout(agent, env, train_replay, eval_replay, logger, args): + + logdir = embodied.Path(args.logdir) + logdir.mkdirs() + print('Logdir', logdir) + should_expl = embodied.when.Until(args.expl_until) + should_train = embodied.when.Ratio(args.train_ratio / args.batch_steps) + should_log = embodied.when.Clock(args.log_every) + should_save = embodied.when.Clock(args.save_every) + should_sync = embodied.when.Every(args.sync_every) + step = logger.step + updates = embodied.Counter() + metrics = embodied.Metrics() + print('Observation space:', embodied.format(env.obs_space), sep='\n') + print('Action space:', embodied.format(env.act_space), sep='\n') + + timer = embodied.Timer() + timer.wrap('agent', agent, ['policy', 'train', 'report', 'save']) + timer.wrap('env', env, ['step']) + if hasattr(train_replay, '_sample'): + timer.wrap('replay', train_replay, ['_sample']) + + nonzeros = set() + def per_episode(ep): + length = len(ep['reward']) - 1 + score = float(ep['reward'].astype(np.float64).sum()) + logger.add({ + 'length': length, 'score': score, + 'reward_rate': (ep['reward'] - ep['reward'].min() >= 0.1).mean(), + }, prefix='episode') + print(f'Episode has {length} steps and return {score:.1f}.') + stats = {} + for key in args.log_keys_video: + if key in ep: + stats[f'policy_{key}'] = ep[key] + for key, value in ep.items(): + if not args.log_zeros and key not in nonzeros and (value == 0).all(): + continue + nonzeros.add(key) + if re.match(args.log_keys_sum, key): + stats[f'sum_{key}'] = ep[key].sum() + if re.match(args.log_keys_mean, key): + stats[f'mean_{key}'] = ep[key].mean() + if re.match(args.log_keys_max, key): + stats[f'max_{key}'] = ep[key].max(0).mean() + metrics.add(stats, prefix='stats') + + driver = embodied.Driver(env) + driver.on_episode(lambda ep, worker: per_episode(ep)) + driver.on_step(lambda tran, _: step.increment()) + driver.on_step(train_replay.add) + + print('Fill eval dataset.') + driver_eval = embodied.Driver(env) + driver_eval.on_step(eval_replay.add) + random_agent = embodied.RandomAgent(env.act_space) + while len(eval_replay) < max(args.batch_steps, args.eval_fill): + print(len(eval_replay), max(args.batch_steps, args.eval_fill)) + driver_eval(random_agent.policy, steps=100) + del driver_eval + print('Prefill train dataset.') + while len(train_replay) < max(args.batch_steps, args.train_fill): + print(len(train_replay), max(args.batch_steps, args.train_fill)) + driver(random_agent.policy, steps=100) + logger.add(metrics.result()) + logger.write() + + dataset_train = agent.dataset(train_replay.dataset) + dataset_eval = agent.dataset(eval_replay.dataset) + state = [None] # To be writable from train step function below. + batch = [None] + def train_step(tran, worker): + for _ in range(should_train(step)): + with timer.scope('dataset_train'): + batch[0] = next(dataset_train) + outs, state[0], mets = agent.train(batch[0], state[0]) + metrics.add(mets, prefix='train') + if 'priority' in outs: + train_replay.prioritize(outs['key'], outs['priority']) + updates.increment() + if should_sync(updates): + agent.sync() + if should_log(step): + logger.add(metrics.result()) + logger.add(agent.report(batch[0]), prefix='report') + with timer.scope('dataset_eval'): + eval_batch = next(dataset_eval) + logger.add(agent.report(eval_batch), prefix='eval') + logger.add(train_replay.stats, prefix='replay') + logger.add(eval_replay.stats, prefix='eval_replay') + logger.add(timer.stats(), prefix='timer') + logger.write(fps=True) + driver.on_step(train_step) + + checkpoint = embodied.Checkpoint(logdir / 'checkpoint.ckpt') + checkpoint.step = step + checkpoint.agent = agent + checkpoint.train_replay = train_replay + checkpoint.eval_replay = eval_replay + if args.from_checkpoint: + checkpoint.load(args.from_checkpoint) + checkpoint.load_or_save() + should_save(step) # Register that we jused saved. + + print('Start training loop.') + policy = lambda *args: agent.policy( + *args, mode='explore' if should_expl(step) else 'train') + while step < args.steps: + # scalars = collections.defaultdict(list) + # for _ in range(args.eval_samples): + # for key, value in agent.report(next(dataset_eval)).items(): + # if value.shape == (): + # scalars[key].append(value) + # for name, values in scalars.items(): + # logger.scalar(f'eval/{name}', np.array(values, np.float64).mean()) + # logger.write() + driver(policy, steps=100) + if should_save(step): + checkpoint.save() + logger.write() + logger.write() diff --git a/dreamerv3/embodied/run/train_save.py b/dreamerv3/embodied/run/train_save.py new file mode 100644 index 0000000..1e59e2c --- /dev/null +++ b/dreamerv3/embodied/run/train_save.py @@ -0,0 +1,130 @@ +import io +import re +from datetime import datetime + +import embodied +import numpy as np + + +def train_save(agent, env, replay, logger, args): + + logdir = embodied.Path(args.logdir) + logdir.mkdirs() + print('Logdir:', logdir) + should_expl = embodied.when.Until(args.expl_until) + should_train = embodied.when.Ratio(args.train_ratio / args.batch_steps) + should_log = embodied.when.Clock(args.log_every) + should_save = embodied.when.Clock(args.save_every) + should_sync = embodied.when.Every(args.sync_every) + step = logger.step + updates = embodied.Counter() + metrics = embodied.Metrics() + print('Observation space:', embodied.format(env.obs_space), sep='\n') + print('Action space:', embodied.format(env.act_space), sep='\n') + + timer = embodied.Timer() + timer.wrap('agent', agent, ['policy', 'train', 'report', 'save']) + timer.wrap('env', env, ['step']) + timer.wrap('replay', replay, ['add', 'save']) + timer.wrap('logger', logger, ['write']) + + nonzeros = set() + def per_episode(ep): + length = len(ep['reward']) - 1 + score = float(ep['reward'].astype(np.float64).sum()) + sum_abs_reward = float(np.abs(ep['reward']).astype(np.float64).sum()) + logger.add({ + 'length': length, + 'score': score, + 'sum_abs_reward': sum_abs_reward, + 'reward_rate': (np.abs(ep['reward']) >= 0.5).mean(), + }, prefix='episode') + print(f'Episode has {length} steps and return {score:.1f}.') + stats = {} + for key in args.log_keys_video: + if key in ep: + stats[f'policy_{key}'] = ep[key] + for key, value in ep.items(): + if not args.log_zeros and key not in nonzeros and (value == 0).all(): + continue + nonzeros.add(key) + if re.match(args.log_keys_sum, key): + stats[f'sum_{key}'] = ep[key].sum() + if re.match(args.log_keys_mean, key): + stats[f'mean_{key}'] = ep[key].mean() + if re.match(args.log_keys_max, key): + stats[f'max_{key}'] = ep[key].max(0).mean() + metrics.add(stats, prefix='stats') + + epsdir = embodied.Path(args.logdir) / 'saved_episodes' + epsdir.mkdirs() + print('Saving episodes:', epsdir) + def save(ep): + time = datetime.now().strftime("%Y%m%dT%H%M%S") + uuid = str(embodied.uuid()) + score = str(np.round(ep['reward'].sum(), 1)).replace('-', 'm') + length = len(ep['reward']) + filename = epsdir / f'{time}-{uuid}-len{length}-rew{score}.npz' + with io.BytesIO() as stream: + np.savez_compressed(stream, **ep) + stream.seek(0) + filename.write(stream.read(), mode='wb') + print('Saved episode:', filename) + saver = embodied.Worker(save, 'thread') + + driver = embodied.Driver(env) + driver.on_episode(lambda ep, worker: per_episode(ep)) + driver.on_episode(lambda ep, worker: saver(ep)) + driver.on_step(lambda tran, _: step.increment()) + driver.on_step(replay.add) + + print('Prefill train dataset.') + random_agent = embodied.RandomAgent(env.act_space) + while len(replay) < max(args.batch_steps, args.train_fill): + driver(random_agent.policy, steps=100) + logger.add(metrics.result()) + logger.write() + + dataset = agent.dataset(replay.dataset) + state = [None] # To be writable from train step function below. + batch = [None] + def train_step(tran, worker): + for _ in range(should_train(step)): + with timer.scope('dataset'): + batch[0] = next(dataset) + outs, state[0], mets = agent.train(batch[0], state[0]) + metrics.add(mets, prefix='train') + if 'priority' in outs: + replay.prioritize(outs['key'], outs['priority']) + updates.increment() + if should_sync(updates): + agent.sync() + if should_log(step): + agg = metrics.result() + report = agent.report(batch[0]) + report = {k: v for k, v in report.items() if 'train/' + k not in agg} + logger.add(agg) + logger.add(report, prefix='report') + logger.add(replay.stats, prefix='replay') + logger.add(timer.stats(), prefix='timer') + logger.write(fps=True) + driver.on_step(train_step) + + checkpoint = embodied.Checkpoint(logdir / 'checkpoint.ckpt') + timer.wrap('checkpoint', checkpoint, ['save', 'load']) + checkpoint.step = step + checkpoint.agent = agent + checkpoint.replay = replay + if args.from_checkpoint: + checkpoint.load(args.from_checkpoint) + checkpoint.load_or_save() + should_save(step) # Register that we jused saved. + + print('Start training loop.') + policy = lambda *args: agent.policy( + *args, mode='explore' if should_expl(step) else 'train') + while step < args.steps: + driver(policy, steps=100) + if should_save(step): + checkpoint.save() + logger.write() diff --git a/dreamerv3/embodied/scripts/install-atari.sh b/dreamerv3/embodied/scripts/install-atari.sh new file mode 100644 index 0000000..bed4d2a --- /dev/null +++ b/dreamerv3/embodied/scripts/install-atari.sh @@ -0,0 +1,17 @@ +#!/bin/sh +set -eu + +apt-get update +apt-get install -y wget +apt-get install -y unrar +apt-get clean + +pip3 install gym==0.19.0 +pip3 install atari-py==0.2.9 +pip3 install opencv-python + +mkdir roms && cd roms +wget -L -nv http://www.atarimania.com/roms/Roms.rar +unrar x -o+ Roms.rar +python3 -m atari_py.import_roms ROMS +cd .. && rm -rf roms diff --git a/dreamerv3/embodied/scripts/install-dmlab.sh b/dreamerv3/embodied/scripts/install-dmlab.sh new file mode 100755 index 0000000..96a8f4d --- /dev/null +++ b/dreamerv3/embodied/scripts/install-dmlab.sh @@ -0,0 +1,48 @@ +#!/bin/sh +set -eu + +# Dependencies +apt-get update && apt-get install -y \ + build-essential curl freeglut3 gettext git libffi-dev libglu1-mesa \ + libglu1-mesa-dev libjpeg-dev liblua5.1-0-dev libosmesa6-dev \ + libsdl2-dev lua5.1 pkg-config python-setuptools python3-dev \ + software-properties-common unzip zip zlib1g-dev g++ +pip3 install numpy + +# Bazel +apt-get install -y apt-transport-https curl gnupg +curl -fsSL https://bazel.build/bazel-release.pub.gpg | gpg --dearmor > bazel.gpg +mv bazel.gpg /etc/apt/trusted.gpg.d/ +echo "deb [arch=amd64] https://storage.googleapis.com/bazel-apt stable jdk1.8" | tee /etc/apt/sources.list.d/bazel.list +apt-get update && apt-get install -y bazel + +# Build +git clone https://github.com/deepmind/lab.git +cd lab +echo 'build --cxxopt=-std=c++17' > .bazelrc +bazel build -c opt //python/pip_package:build_pip_package +./bazel-bin/python/pip_package/build_pip_package /tmp/dmlab_pkg +pip3 install --force-reinstall /tmp/dmlab_pkg/deepmind_lab-*.whl +cd .. +rm -rf lab + +# Dataset +mkdir dmlab_data +cd dmlab_data +pip3 install Pillow +curl https://bradylab.ucsd.edu/stimuli/ObjectsAll.zip -o ObjectsAll.zip +unzip ObjectsAll.zip +cd OBJECTSALL +python3 << EOM +import os +from PIL import Image +files = [f for f in os.listdir('.') if f.lower().endswith('jpg')] +for i, file in enumerate(sorted(files)): + print(file) + im = Image.open(file) + im.save('../%04d.png' % (i+1)) +EOM +cd .. +rm -rf __MACOSX OBJECTSALL ObjectsAll.zip + +apt-get clean diff --git a/dreamerv3/embodied/scripts/install-minecraft.sh b/dreamerv3/embodied/scripts/install-minecraft.sh new file mode 100644 index 0000000..02f9954 --- /dev/null +++ b/dreamerv3/embodied/scripts/install-minecraft.sh @@ -0,0 +1,12 @@ +#!/bin/sh +set -eu + +apt-get update +apt-get install -y libgl1-mesa-dev +apt-get install -y libx11-6 +apt-get install -y openjdk-8-jdk +apt-get install -y x11-xserver-utils +apt-get install -y xvfb +apt-get clean + +pip3 install minerl==0.4.4 diff --git a/dreamerv3/embodied/scripts/plot.py b/dreamerv3/embodied/scripts/plot.py new file mode 100644 index 0000000..34297bf --- /dev/null +++ b/dreamerv3/embodied/scripts/plot.py @@ -0,0 +1,523 @@ +import argparse +import collections +import functools +import gzip +import json +import multiprocessing as mp +import pathlib +import re +import subprocess +import warnings + +import matplotlib.pyplot as plt +import matplotlib.ticker as ticker +import numpy as np +import pandas as pd +import rich.console +import tqdm + +TITLES = { + 'dmlab_explore_goal_locations_small': 'DMLab Goals Small', + 'crafter_reward': 'Crafter', + 'pinpad2_three': 'Pin Pad Three', + 'pinpad2_four': 'Pin Pad Four', + 'pinpad2_five': 'Pin Pad Five', + 'pinpad2_six': 'Pin Pad Six', + 'pinpad2_eight': 'Pin Pad Eight', + 'loconav_ant_maze_s_50hz': 'Ant Maze S', + 'loconav_ant_maze_m_50hz': 'Ant Maze M', + 'loconav_ant_maze_l_50hz': 'Ant Maze L', + 'loconav_ant_maze_xl_50hz': 'Ant Maze XL', +} + +COLORS = { + 'contrast': ( + '#0022ff', '#33aa00', '#ff0011', '#ddaa00', '#cc44dd', '#0088aa', + '#001177', '#117700', '#990022', '#885500', '#553366', '#006666'), + 'gradient': ( + '#a0da39', '#4ac16d', '#277f8e', '#365c8d', '#46327e', '#440154'), + 'gradient_more': ( + '#fde725', '#a0da39', '#4ac16d', '#1fa187', '#277f8e', '#365c8d', + '#46327e', '#440154'), +} + + +def main(): + console = rich.console.Console() + args = parse_args() + runs = [] + for directory in args.indirs: + seed_prefix = len(args.indirs) > 1 and directory.name + method_prefix = args.prefix and directory.name + runs += load_metrics( + directory, args.pattern, args.xaxis, args.yaxis, args.yaxis2, + seed_prefix, method_prefix, args.tasks, args.methods, args.workers) + tasks = [] + for regex in args.tasks: + found = [x['task'] for x in runs if re.search(regex, x['task'])] + [tasks.append(x) for x in natsort(found) if x not in tasks] + methods = [] + for regex in args.methods: + found = [x['method'] for x in runs if re.search(regex, x['method'])] + [methods.append(x) for x in natsort(found) if x not in methods] + seeds = natsort(set(run['seed'] for run in runs)) + console.print(f'Tasks ({len(tasks)}): [cyan]{", ".join(tasks)}[/cyan]') + console.print(f'Methods ({len(methods)}): [cyan]{", ".join(methods)}[/cyan]') + console.print(f'Seed ({len(seeds)}): [cyan]{", ".join(seeds)}[/cyan]') + if not runs: + console.print('Nothing to plot!', style='red') + return + args.outdir.mkdir(parents=True, exist_ok=True) + + if args.stats: + print('Computing stats...', flush=True) + len(tasks) == 1 and 'mean' in args.stats and args.stats.remove('mean') + len(tasks) == 1 and 'median' in args.stats and args.stats.remove('median') + extra_runs, extra_tasks = compute_stats(runs, args.stats, args.bins) + runs += extra_runs + tasks += extra_tasks + + print('Binning runs...', flush=True) + if args.bins: + maxs = collections.defaultdict(list) + for run in runs: + maxs[(run['task'], run['method'])].append(run['xs'].max()) + maxs = {k: max(vs) for k, vs in maxs.items()} + for run in runs: + if run['task'].startswith('stats_'): + continue + max_ = maxs[(run['task'], run['method'])] + 1e-8 + max_ = min(max_, args.xlim[1]) if args.xlim else max_ + step = max(1e-8, max_ / 30) if args.bins < 0 else args.bins + borders = np.arange(0, max_, step) + xs, ys = binning(run['xs'], run['ys'], borders, np.nanmean, fill='nan') + run['xs'], run['ys'] = xs, ys + + print('Saving runs...', flush=True) + filename = args.outdir / 'runs.json.gz' + with gzip.open(filename, 'w') as f: + f.write(json.dumps([ + {**run, 'xs': run['xs'].tolist(), 'ys': run['ys'].tolist()} + for run in runs]).encode('utf-8')) + console.print(f'Saved [green]{filename}[/green]') + + print('Plotting...', flush=True) + fig, axes = plots(len(tasks), args.cols, args.size) + for task, ax in zip(tasks, axes): + title = TITLES.get(task, task.split('_', 1)[1].replace('_', ' ').title()) + ax.set_title(title) + if not task.startswith('stats_'): + args.xlim and ax.set_xlim(*args.xlim) + args.ylim and ax.set_ylim(*args.ylim) + args.xticks and ax.set_xticks(args.xticks) + ax.xaxis.set_major_formatter(smart_format) + # ax.tick_params(axis='both', labelsize=7) # TOFO + for task, ax in zip(tasks, axes): + for i, method in enumerate(methods): + relevant = [ + run for run in runs + if run['task'] == task and run['method'] == method] + if not relevant: + console.print(f'Missing {method} on {task}!', style='red') + continue + if args.bins and args.agg: + groups = [relevant] + else: + groups = [[run] for run in relevant] + for group in groups: + xs = group[0]['xs'] + ys = np.stack([run['ys'] for run in group], 0) + mean = reduce(ys, np.nanmean, 0) + std = reduce(ys, np.nanstd, 0) + curve( + ax, xs, mean, mean - std, mean + std, + label=args.labels.get(method, method), + order=i, color=args.colors(i)) + legendcols = args.legendcols or min(4, args.cols, len(axes)) + legend(fig, adjust=True, ncol=legendcols) + if args.stats: + for ax in axes[-len(extra_tasks):]: + ax.set_facecolor((0.9, 0.9, 0.9)) + save(fig, args.outdir / 'curves.png') + save(fig, args.outdir / 'curves.pdf') + + +def compute_stats(runs, stats, bins): + extra_runs = [] + select = lambda baselines, name: { + k: v[name] for k, v in baselines.items() if name in v} + for stats in stats: + if stats == 'tasks': + extra_runs += stats_num_tasks(runs, bins) + elif stats == 'mean': + extra_runs += stats_self_norm(runs, bins, 'mean', np.nanmean) + elif stats == 'median': + extra_runs += stats_self_norm(runs, bins, 'median', np.nanmedian) + elif stats == 'atari_mean': + path = pathlib.Path('~/scores/atari_baselines.json').expanduser() + baselines = json.loads(path.read_text()) + mins = select(baselines, 'random') + maxs = select(baselines, 'human_gamer') + extra_runs += stats_fixed_norm( + runs, bins, mins, maxs, 'gamer_mean', np.nanmean) + elif stats == 'atari_median': + path = pathlib.Path('~/scores/atari_baselines.json').expanduser() + baselines = json.loads(path.read_text()) + mins = select(baselines, 'random') + maxs = select(baselines, 'human_gamer') + extra_runs += stats_fixed_norm( + runs, bins, mins, maxs, 'gamer_median', np.nanmedian) + elif stats == 'atari_record': + path = pathlib.Path('~/scores/atari_baselines.json').expanduser() + baselines = json.loads(path.read_text()) + mins = select(baselines, 'random') + maxs = select(baselines, 'human_record') + extra_runs += stats_fixed_norm( + runs, bins, mins, maxs, 'record_mean', np.nanmean) + elif stats == 'atari_record_clip': + path = pathlib.Path('~/scores/atari_baselines.json').expanduser() + baselines = json.loads(path.read_text()) + mins = select(baselines, 'random') + maxs = select(baselines, 'human_record') + extra_runs += stats_fixed_norm( + runs, bins, mins, maxs, 'record_mean_clip', + lambda x, a: np.nanmean(np.minimum(x, 1), a)) + elif stats == 'dmlab_mean': + path = pathlib.Path('~/scores/dmlab_baselines.json').expanduser() + baselines = json.loads(path.read_text()) + mins = select(baselines, 'random') + maxs = select(baselines, 'human') + extra_runs += stats_fixed_norm( + runs, bins, mins, maxs, 'human_mean', + lambda vals, axis: np.nanmean(np.minimum(vals, 1), axis)) + else: + raise NotImplementedError(stats) + extra_tasks = natsort(set(run['task'] for run in extra_runs)) + return extra_runs, extra_tasks + + +def stats_self_norm(runs, bins, name='mean', aggregator=np.nanmean): + methods = natsort(set(run['method'] for run in runs)) + seeds = natsort(set(run['seed'] for run in runs)) + lengths, mins, maxs = {}, {}, {} + for run in runs: + lengths[run['task']] = max(lengths.get(run['task'], 0), max(run['xs'])) + mins[run['task']] = min(mins.get(run['task'], np.inf), min(run['ys'])) + maxs[run['task']] = max(maxs.get(run['task'], -np.inf), max(run['ys'])) + if bins <= 0: + borders = { + task: np.linspace(0, length + 1e-8, 30) + for task, length in lengths.items()} + else: + border = np.arange(0, max(lengths.values()) + 1e-8, bins) + borders = {task: border for task, length in lengths.items()} + extra_runs = [] + for method in methods: + for seed in seeds: + scores = [] + for run in runs: + if not (run['method'] == method and run['seed'] == seed): + continue + task = run['task'] + if np.isclose(mins[task], maxs[task]): + continue + _, ys = binning( + run['xs'], run['ys'], borders[task], np.nanmean, fill='last') + scores.append((ys - mins[task]) / (maxs[task] - mins[task])) + if scores: + scores = np.array(scores) + xs = np.linspace(0, 1, len(scores[0])) + extra_runs.append({ + 'task': f'stats_normalized_{name}', 'method': method, 'seed': seed, + 'xs': xs, 'ys': reduce(scores, aggregator, 0)}) + return extra_runs + + +def stats_fixed_norm( + runs, bins, mins, maxs, name='mean', aggregator=np.nanmean): + methods = natsort(set(run['method'] for run in runs)) + seeds = natsort(set(run['seed'] for run in runs)) + lengths = {} + for run in runs: + lengths[run['task']] = max(lengths.get(run['task'], 0), max(run['xs'])) + if bins <= 0: + borders = { + task: np.linspace(0, length + 1e-8, 30) + for task, length in lengths.items()} + else: + border = np.arange(0, max(lengths.values()) + 1e-8, bins) + borders = {task: border for task, length in lengths.items()} + extra_runs = [] + for method in methods: + for seed in seeds: + scores = [] + for run in runs: + if not (run['method'] == method and run['seed'] == seed): + continue + task = run['task'] + _, ys = binning( + run['xs'], run['ys'], borders[task], np.nanmean, fill='last') + if task == 'atari_jamesbond' and 'atari_james_bond' in mins: + task = 'atari_james_bond' + scores.append((ys - mins[task]) / (maxs[task] - mins[task])) + if scores: + xs = np.linspace(0, 1, len(scores[0])) + extra_runs.append({ + 'task': f'stats_{name}', 'method': method, 'seed': seed, + 'xs': xs, 'ys': reduce(scores, aggregator, 0)}) + return extra_runs + + +def stats_num_tasks(runs, bins): + methods = natsort(set(run['method'] for run in runs)) + seeds = natsort(set(run['seed'] for run in runs)) + lengths = {} + for run in runs: + lengths[run['task']] = max(lengths.get(run['task'], 0), max(run['xs'])) + if bins <= 0: + borders = { + task: np.linspace(0, length + 1e-8, 30) + for task, length in lengths.items()} + else: + border = np.arange(0, max(lengths.values()) + 1e-8, bins) + borders = {task: border for task, length in lengths.items()} + extra_runs = [] + for method in methods: + for seed in seeds: + nonempty = [] + for run in runs: + if not (run['method'] == method and run['seed'] == seed): + continue + task = run['task'] + _, ys = binning( + run['xs'], run['ys'], borders[task], np.nanmean, fill='nan') + nonempty.append(np.isfinite(ys)) + if nonempty: + xs = np.linspace(0, 1, len(nonempty[0])) + extra_runs.append({ + 'task': 'stats_number_of_tasks', 'method': method, 'seed': seed, + 'xs': xs, 'ys': np.sum(nonempty, 0)}) + return extra_runs + + +def load_metrics( + directory, pattern, xaxis, yaxis, yaxis2, seed_prefix=None, + method_prefix=None, tasks=(r'.*',), methods=(r'.*',), workers=1): + console = rich.console.Console() + directory = directory.expanduser().resolve() + tasks = [re.compile(regex) for regex in tasks] + methods = [re.compile(regex) for regex in methods] + runs = [] + for filename in directory.glob(pattern): + task, method, seed = filename.parts[-4:-1] + if not any(p.search(task) for p in tasks): + continue + if not any(p.search(method) for p in methods): + continue + if seed_prefix: + seed = f'{seed_prefix}_{seed}' + if method_prefix: + method = f'{method_prefix}_{method}' + runs.append({ + 'task': task, 'method': method, 'seed': seed, 'filename': filename}) + console.print(f'Loading {len(runs)} runs from [green]{directory}[/green]...') + jobs = [ + functools.partial(load_run, run, xaxis, yaxis, yaxis2) for run in runs] + if workers > 1: + with mp.Pool(workers) as pool: + promises = [pool.apply_async(j) for j in jobs] + runs = [promise.get() for promise in tqdm.tqdm(promises)] + else: + runs = [job() for job in tqdm.tqdm(jobs)] + runs = [r for r in runs if r is not None] + return runs + + +def load_run(run, xaxis, yaxis, yaxis2): + try: + console = rich.console.Console() + filename = run.pop('filename') + try: + df = pd.read_json(filename, lines=True) + except ValueError: + records = [] + for i, line in enumerate(pathlib.Path(filename).read_text().split('\n')): + if not line: + continue + try: + records.append(json.loads(line)) + except ValueError: + print(f'Skipping invalid JSON line {i} in {filename}.') + df = pd.DataFrame(records) + yaxis = yaxis if yaxis in df.columns else yaxis2 + df = df[[xaxis, yaxis]].dropna() + run['xs'] = df[xaxis].to_numpy() + run['ys'] = df[yaxis].to_numpy() + return run + except Exception as e: + console.print( + f'Exception loading {run["method"]} on {run["task"]}:\n {e}', + style='red') + return None + + +def plots( + amount, cols=4, size=(2, 2.3), xticks=4, yticks=5, grid=(1, 1), **kwargs): + cols = min(cols, amount) + rows = int(np.ceil(amount / cols)) + size = (cols * size[0], rows * size[1]) + fig, axes = plt.subplots(rows, cols, figsize=size, squeeze=False, **kwargs) + axes = axes.flatten() + for ax in axes: + ax.xaxis.set_major_locator(ticker.MaxNLocator(xticks)) + ax.yaxis.set_major_locator(ticker.MaxNLocator(yticks)) + if grid: + grid = (grid, grid) if not hasattr(grid, '__len__') else grid + ax.grid(which='both', color='#eeeeee') + ax.xaxis.set_minor_locator(ticker.AutoMinorLocator(int(grid[0]))) + ax.yaxis.set_minor_locator(ticker.AutoMinorLocator(int(grid[1]))) + ax.tick_params(which='minor', length=0) + for ax in axes[amount:]: + ax.axis('off') + axes = axes[:amount] + return fig, axes + + +def curve(ax, xs, ys, low=None, high=None, label=None, order=0, **kwargs): + finite = np.isfinite(ys) + ax.plot( + xs[finite], ys[finite], + label=label, zorder=1000 - order, **kwargs) + if low is not None and finite.sum() > 1: + ax.fill_between( + xs[finite], low[finite], high[finite], + zorder=100 - order, alpha=0.2, lw=0, **kwargs) + + +def legend(fig, mapping=None, adjust=False, **kwargs): + options = dict( + fontsize='medium', numpoints=1, labelspacing=0, columnspacing=1.2, + handlelength=1.5, handletextpad=0.5, ncol=4, loc='lower center') + options.update(kwargs) + entries = {} + for ax in fig.axes: + for handle, label in zip(*ax.get_legend_handles_labels()): + if mapping and label in mapping: + label = mapping[label] + entries[label] = handle + leg = fig.legend(entries.values(), entries.keys(), **options) + leg.get_frame().set_edgecolor('white') + if adjust is not False: + pad = adjust if isinstance(adjust, (int, float)) else 0.5 + extent = leg.get_window_extent(fig.canvas.get_renderer()) + extent = extent.transformed(fig.transFigure.inverted()) + yloc, xloc = options['loc'].split() + y0 = dict(lower=extent.y1, center=0, upper=0)[yloc] + y1 = dict(lower=1, center=1, upper=extent.y0)[yloc] + x0 = dict(left=extent.x1, center=0, right=0)[xloc] + x1 = dict(left=1, center=1, right=extent.x0)[xloc] + fig.tight_layout(rect=[x0, y0, x1, y1], h_pad=pad, w_pad=pad) + + +def smart_format(x, pos=None): + if abs(x) < 1e3: + if float(int(x)) == float(x): + return str(int(x)) + return str(round(x, 10)).rstrip('0') + if abs(x) < 1e6: + return f'{x/1e3:.0f}K' if x == x // 1e3 * 1e3 else f'{x/1e3:.1f}K' + if abs(x) < 1e9: + return f'{x/1e6:.0f}M' if x == x // 1e6 * 1e6 else f'{x/1e6:.1f}M' + return f'{x/1e9:.0f}B' if x == x // 1e9 * 1e9 else f'{x/1e9:.1f}B' + + +def save(fig, filename): + console = rich.console.Console() + filename = pathlib.Path(filename).expanduser() + filename.parent.mkdir(parents=True, exist_ok=True) + fig.savefig(filename) + console.print(f'Saved [green]{filename}[/green]') + if filename.suffix == '.pdf': + try: + subprocess.call(['pdfcrop', str(filename), str(filename)]) + except FileNotFoundError: + print('Install LaTeX to crop PDF outputs.') + + +def binning(xs, ys, borders, reducer=np.nanmean, fill='nan'): + xs = xs if isinstance(xs, np.ndarray) else np.array(xs) + ys = ys if isinstance(ys, np.ndarray) else np.array(ys) + order = np.argsort(xs) + xs, ys = xs[order], ys[order] + binned = [] + for start, stop in zip(borders[:-1], borders[1:]): + left = (xs <= start).sum() + right = (xs <= stop).sum() + if left < right: + value = reduce(ys[left:right], reducer) + elif binned: + value = {'nan': np.nan, 'last': binned[-1]}[fill] + else: + value = np.nan + binned.append(value) + return borders[1:], np.array(binned) + + +def reduce(values, reducer=np.nanmean, *args, **kwargs): + with warnings.catch_warnings(): # Buckets can be empty. + warnings.simplefilter('ignore', category=RuntimeWarning) + return reducer(values, *args, **kwargs) + + +def natsort(sequence): + pattern = re.compile(r'([0-9]+)') + return sorted(sequence, key=lambda x: [ + (int(y) if y.isdigit() else y) for y in pattern.split(x)]) + + +def parse_args(argv=None): + boolean = lambda x: bool(['False', 'True'].index(x)) + parser = argparse.ArgumentParser() + parser.add_argument('--indirs', nargs='+', type=pathlib.Path, required=True) + parser.add_argument('--outdir', type=pathlib.Path, required=True) + parser.add_argument('--pattern', type=str, default='**/scores.jsonl') + parser.add_argument('--prefix', type=boolean, default=False) + parser.add_argument('--xaxis', type=str, default='step') + parser.add_argument('--yaxis', type=str, default='episode/score') + parser.add_argument('--yaxis2', type=str, default='eval_episode/score') + parser.add_argument('--tasks', nargs='+', default=[r'.*']) + parser.add_argument('--methods', nargs='+', default=[r'.*']) + parser.add_argument('--bins', type=float, default=-1) + parser.add_argument('--agg', type=boolean, default=True) + parser.add_argument('--size', nargs=2, type=float, default=[2.5, 2.3]) + parser.add_argument('--cols', type=int, default=6) + parser.add_argument('--legendcols', type=int, default=0) + parser.add_argument('--xlim', nargs=2, type=float, default=None) + parser.add_argument('--ylim', nargs=2, type=float, default=None) + parser.add_argument('--xticks', nargs='+', type=float, default=None) + parser.add_argument('--labels', nargs='+', default=[]) + parser.add_argument('--colors', type=str, nargs='+', default=['contrast']) + parser.add_argument('--workers', type=int, default=12) + parser.add_argument('--stats', type=str, nargs='*', default=[ + 'mean', 'median', 'tasks']) + args = parser.parse_args(argv) + args.indirs = tuple([x.expanduser() for x in args.indirs]) + args.outdir = args.outdir.expanduser() / args.indirs[0].stem + assert len(args.labels) % 2 == 0 + args.labels = {k: v for k, v in zip(args.labels[:-1], args.labels[1:])} + if len(args.colors) == 1: + try: + args.colors = plt.get_cmap(args.colors[0]) + except ValueError: + if args.colors[0] in COLORS: + cmap = COLORS[args.colors[0]] + else: + cmap = args.colors + args.colors = lambda i: cmap[i % len(cmap)] + if args.stats == ['none']: + args.stats = [] + return args + + +if __name__ == '__main__': + main() diff --git a/dreamerv3/embodied/scripts/xvfb_run.sh b/dreamerv3/embodied/scripts/xvfb_run.sh new file mode 100644 index 0000000..4a1e06e --- /dev/null +++ b/dreamerv3/embodied/scripts/xvfb_run.sh @@ -0,0 +1,2 @@ +xvfb-run -a -s "-screen 0 1024x768x24 -ac +extension GLX +render -noreset" "$@" +# xvfb-run "$@" diff --git a/dreamerv3/expl.py b/dreamerv3/expl.py new file mode 100644 index 0000000..6df6445 --- /dev/null +++ b/dreamerv3/expl.py @@ -0,0 +1,37 @@ +import jax +import jax.numpy as jnp +tree_map = jax.tree_util.tree_map +sg = lambda x: tree_map(jax.lax.stop_gradient, x) + +from . import nets +from . import jaxutils +from . import ninjax as nj + + +class Disag(nj.Module): + + def __init__(self, wm, act_space, config): + self.config = config.update({'disag_head.inputs': ['tensor']}) + self.opt = jaxutils.Optimizer(name='disag_opt', **config.expl_opt) + self.inputs = nets.Input(config.disag_head.inputs, dims='deter') + self.target = nets.Input(self.config.disag_target, dims='deter') + self.nets = [ + nets.MLP(shape=None, **self.config.disag_head, name=f'disag{i}') + for i in range(self.config.disag_models)] + + def __call__(self, traj): + inp = self.inputs(traj) + preds = jnp.array([net(inp).mode() for net in self.nets]) + return preds.std(0).mean(-1)[1:] + + def train(self, data): + return self.opt(self.nets, self.loss, data) + + def loss(self, data): + inp = sg(self.inputs(data)[:, :-1]) + tar = sg(self.target(data)[:, 1:]) + losses = [] + for net in self.nets: + net._shape = tar.shape[2:] + losses.append(-net(inp).log_prob(tar).mean()) + return jnp.array(losses).sum() diff --git a/dreamerv3/jaxagent.py b/dreamerv3/jaxagent.py new file mode 100644 index 0000000..096da0c --- /dev/null +++ b/dreamerv3/jaxagent.py @@ -0,0 +1,240 @@ +import os + +import embodied +import jax +import jax.numpy as jnp +import numpy as np + +from . import jaxutils +from . import ninjax as nj + +tree_map = jax.tree_util.tree_map +tree_flatten = jax.tree_util.tree_flatten + + +def Wrapper(agent_cls): + class Agent(JAXAgent): + configs = agent_cls.configs + inner = agent_cls + def __init__(self, *args, **kwargs): + super().__init__(agent_cls, *args, **kwargs) + return Agent + + +class JAXAgent(embodied.Agent): + + def __init__(self, agent_cls, obs_space, act_space, step, config): + self.config = config.jax + self.batch_size = config.batch_size + self.batch_length = config.batch_length + self.data_loaders = config.data_loaders + self._setup() + self.agent = agent_cls(obs_space, act_space, step, config, name='agent') + self.rng = np.random.default_rng(config.seed) + + available = jax.devices(self.config.platform) + self.policy_devices = [available[i] for i in self.config.policy_devices] + self.train_devices = [available[i] for i in self.config.train_devices] + self.single_device = (self.policy_devices == self.train_devices) and ( + len(self.policy_devices) == 1) + print(f'JAX devices ({jax.local_device_count()}):', available) + print('Policy devices:', ', '.join([str(x) for x in self.policy_devices])) + print('Train devices: ', ', '.join([str(x) for x in self.train_devices])) + + self._once = True + self._updates = embodied.Counter() + self._should_metrics = embodied.when.Every(self.config.metrics_every) + self._transform() + self.varibs = self._init_varibs(obs_space, act_space) + self.sync() + + def policy(self, obs, state=None, mode='train'): + obs = obs.copy() + obs = self._convert_inps(obs, self.policy_devices) + rng = self._next_rngs(self.policy_devices) + varibs = self.varibs if self.single_device else self.policy_varibs + if state is None: + state, _ = self._init_policy(varibs, rng, obs['is_first']) + else: + state = tree_map( + np.asarray, state, is_leaf=lambda x: isinstance(x, list)) + state = self._convert_inps(state, self.policy_devices) + (outs, state), _ = self._policy(varibs, rng, obs, state, mode=mode) + outs = self._convert_outs(outs, self.policy_devices) + # TODO: Consider keeping policy states in accelerator memory. + state = self._convert_outs(state, self.policy_devices) + return outs, state + + def train(self, data, state=None): + rng = self._next_rngs(self.train_devices) + if state is None: + state, self.varibs = self._init_train(self.varibs, rng, data['is_first']) + (outs, state, mets), self.varibs = self._train( + self.varibs, rng, data, state) + outs = self._convert_outs(outs, self.train_devices) + self._updates.increment() + if self._should_metrics(self._updates): + mets = self._convert_mets(mets, self.train_devices) + else: + mets = {} + if self._once: + self._once = False + assert jaxutils.Optimizer.PARAM_COUNTS + for name, count in jaxutils.Optimizer.PARAM_COUNTS.items(): + mets[f'params_{name}'] = float(count) + return outs, state, mets + + def report(self, data): + rng = self._next_rngs(self.train_devices) + mets, _ = self._report(self.varibs, rng, data) + mets = self._convert_mets(mets, self.train_devices) + return mets + + def dataset(self, generator): + batcher = embodied.Batcher( + sources=[generator] * self.batch_size, + workers=self.data_loaders, + postprocess=lambda x: self._convert_inps(x, self.train_devices), + prefetch_source=4, prefetch_batch=1) + return batcher() + + def save(self): + if len(self.train_devices) > 1: + varibs = tree_map(lambda x: x[0], self.varibs) + else: + varibs = self.varibs + varibs = jax.device_get(varibs) + data = tree_map(np.asarray, varibs) + return data + + def load(self, state): + if len(self.train_devices) == 1: + self.varibs = jax.device_put(state, self.train_devices[0]) + else: + self.varibs = jax.device_put_replicated(state, self.train_devices) + self.sync() + + def sync(self): + if self.single_device: + return + if len(self.train_devices) == 1: + varibs = self.varibs + else: + varibs = tree_map(lambda x: x[0].device_buffer, self.varibs) + if len(self.policy_devices) == 1: + self.policy_varibs = jax.device_put(varibs, self.policy_devices[0]) + else: + self.policy_varibs = jax.device_put_replicated( + varibs, self.policy_devices) + + def _setup(self): + try: + import tensorflow as tf + tf.config.set_visible_devices([], 'GPU') + tf.config.set_visible_devices([], 'TPU') + except Exception as e: + print('Could not disable TensorFlow devices:', e) + if not self.config.prealloc: + os.environ["XLA_PYTHON_CLIENT_PREALLOCATE"] = "false" + os.environ['XLA_PYTHON_CLIENT_MEM_FRACTION'] = '0.8' + xla_flags = [] + if self.config.logical_cpus: + count = self.config.logical_cpus + xla_flags.append(f'--xla_force_host_platform_device_count={count}') + if xla_flags: + os.environ['XLA_FLAGS'] = ' '.join(xla_flags) + jax.config.update('jax_platform_name', self.config.platform) + jax.config.update('jax_disable_jit', not self.config.jit) + jax.config.update('jax_debug_nans', self.config.debug_nans) + jax.config.update('jax_transfer_guard', 'disallow') + if self.config.platform == 'cpu': + jax.config.update('jax_disable_most_optimizations', self.config.debug) + jaxutils.COMPUTE_DTYPE = getattr(jnp, self.config.precision) + + def _transform(self): + self._init_policy = nj.pure(lambda x: self.agent.policy_initial(len(x))) + self._init_train = nj.pure(lambda x: self.agent.train_initial(len(x))) + self._policy = nj.pure(self.agent.policy) + self._train = nj.pure(self.agent.train) + self._report = nj.pure(self.agent.report) + if len(self.train_devices) == 1: + kw = dict(device=self.train_devices[0]) + self._init_train = nj.jit(self._init_train, **kw) + self._train = nj.jit(self._train, **kw) + self._report = nj.jit(self._report, **kw) + else: + kw = dict(devices=self.train_devices) + self._init_train = nj.pmap(self._init_train, 'i', **kw) + self._train = nj.pmap(self._train, 'i', **kw) + self._report = nj.pmap(self._report, 'i', **kw) + if len(self.policy_devices) == 1: + kw = dict(device=self.policy_devices[0]) + self._init_policy = nj.jit(self._init_policy, **kw) + self._policy = nj.jit(self._policy, static=['mode'], **kw) + else: + kw = dict(devices=self.policy_devices) + self._init_policy = nj.pmap(self._init_policy, 'i', **kw) + self._policy = nj.pmap(self._policy, 'i', static=['mode'], **kw) + + def _convert_inps(self, value, devices): + if len(devices) == 1: + value = jax.device_put(value, devices[0]) + else: + check = tree_map(lambda x: len(x) % len(devices) == 0, value) + if not all(jax.tree_util.tree_leaves(check)): + shapes = tree_map(lambda x: x.shape, value) + raise ValueError( + f'Batch must by divisible by {len(devices)} devices: {shapes}') + # TODO: Avoid the reshape? + value = tree_map( + lambda x: x.reshape((len(devices), -1) + x.shape[1:]), value) + shards = [] + for i in range(len(devices)): + shards.append(tree_map(lambda x: x[i], value)) + value = jax.device_put_sharded(shards, devices) + return value + + def _convert_outs(self, value, devices): + value = jax.device_get(value) + value = tree_map(np.asarray, value) + if len(devices) > 1: + value = tree_map(lambda x: x.reshape((-1,) + x.shape[2:]), value) + return value + + def _convert_mets(self, value, devices): + value = jax.device_get(value) + value = tree_map(np.asarray, value) + if len(devices) > 1: + value = tree_map(lambda x: x[0], value) + return value + + def _next_rngs(self, devices, mirror=False, high=2 ** 63 - 1): + if len(devices) == 1: + return jax.device_put(self.rng.integers(high), devices[0]) + elif mirror: + return jax.device_put_replicated( + self.rng.integers(high), devices) + else: + return jax.device_put_sharded( + list(self.rng.integers(high, size=len(devices))), devices) + + def _init_varibs(self, obs_space, act_space): + varibs = {} + rng = self._next_rngs(self.train_devices, mirror=True) + dims = (self.batch_size, self.batch_length) + data = self._dummy_batch({**obs_space, **act_space}, dims) + data = self._convert_inps(data, self.train_devices) + state, varibs = self._init_train(varibs, rng, data['is_first']) + varibs = self._train(varibs, rng, data, state, init_only=True) + # obs = self._dummy_batch(obs_space, (1,)) + # state, varibs = self._init_policy(varibs, rng, obs['is_first']) + # varibs = self._policy( + # varibs, rng, obs, state, mode='train', init_only=True) + return varibs + + def _dummy_batch(self, spaces, batch_dims): + spaces = list(spaces.items()) + data = {k: np.zeros(v.shape, v.dtype) for k, v in spaces} + for dim in reversed(batch_dims): + data = {k: np.repeat(v[None], dim, axis=0) for k, v in data.items()} + return data diff --git a/dreamerv3/jaxutils.py b/dreamerv3/jaxutils.py new file mode 100644 index 0000000..6ee5d15 --- /dev/null +++ b/dreamerv3/jaxutils.py @@ -0,0 +1,476 @@ +import re + +import jax +import jax.numpy as jnp +import numpy as np +import optax +from tensorflow_probability.substrates import jax as tfp + +from . import ninjax as nj + +tfd = tfp.distributions +tree_map = jax.tree_util.tree_map +sg = lambda x: tree_map(jax.lax.stop_gradient, x) +COMPUTE_DTYPE = jnp.float32 + + +def cast_to_compute(values): + return tree_map(lambda x: x.astype(COMPUTE_DTYPE), values) + + +def parallel(): + try: + jax.lax.axis_index('i') + return True + except NameError: + return False + + +def tensorstats(tensor, prefix=None): + metrics = { + 'mean': tensor.mean(), + 'std': tensor.std(), + 'mag': jnp.abs(tensor).max(), + 'min': tensor.min(), + 'max': tensor.max(), + 'dist': subsample(tensor), + } + if prefix: + metrics = {f'{prefix}_{k}': v for k, v in metrics.items()} + return metrics + + +def subsample(values, amount=1024): + values = values.flatten() + if len(values) > amount: + values = jax.random.permutation(nj.rng(), values)[:amount] + return values + + +def scan(fn, inputs, start, unroll=True, modify=False): + fn2 = lambda carry, inp: (fn(carry, inp),) * 2 + if not unroll: + return nj.scan(fn2, start, inputs, modify=modify)[1] + length = len(jax.tree_util.tree_leaves(inputs)[0]) + carrydef = jax.tree_util.tree_structure(start) + carry = start + outs = [] + for index in range(length): + carry, out = fn2(carry, tree_map(lambda x: x[index], inputs)) + flat, treedef = jax.tree_util.tree_flatten(out) + assert treedef == carrydef, (treedef, carrydef) + outs.append(flat) + outs = [ + jnp.stack([carry[i] for carry in outs], 0) + for i in range(len(outs[0]))] + return carrydef.unflatten(outs) + + +def symlog(x): + return jnp.sign(x) * jnp.log(1 + jnp.abs(x)) + + +def symexp(x): + return jnp.sign(x) * (jnp.exp(jnp.abs(x)) - 1) + + +class OneHotDist(tfd.OneHotCategorical): + + def __init__(self, logits=None, probs=None, dtype=jnp.float32): + super().__init__(logits, probs, dtype) + + @classmethod + def _parameter_properties(cls, dtype, num_classes=None): + return super()._parameter_properties(dtype) + + def sample(self, sample_shape=(), seed=None): + sample = sg(super().sample(sample_shape, seed)) + probs = self._pad(super().probs_parameter(), sample.shape) + return sg(sample) + (probs - sg(probs)).astype(sample.dtype) + + def _pad(self, tensor, shape): + while len(tensor.shape) < len(shape): + tensor = tensor[None] + return tensor + + +class MSEDist: + + def __init__(self, mode, dims, agg='sum'): + self._mode = mode + self._dims = tuple([-x for x in range(1, dims + 1)]) + self._agg = agg + self.batch_shape = mode.shape[:len(mode.shape) - dims] + self.event_shape = mode.shape[len(mode.shape) - dims:] + + def mode(self): + return self._mode + + def mean(self): + return self._mode + + def log_prob(self, value): + assert self._mode.shape == value.shape, (self._mode.shape, value.shape) + distance = ((self._mode - value) ** 2) + if self._agg == 'mean': + loss = distance.mean(self._dims) + elif self._agg == 'sum': + loss = distance.sum(self._dims) + else: + raise NotImplementedError(self._agg) + return -loss + + +class SymlogDist: + + def __init__(self, mode, dims, dist='mse', agg='sum', tol=1e-8): + self._mode = mode + self._dims = tuple([-x for x in range(1, dims + 1)]) + self._dist = dist + self._agg = agg + self._tol = tol + self.batch_shape = mode.shape[:len(mode.shape) - dims] + self.event_shape = mode.shape[len(mode.shape) - dims:] + + def mode(self): + return symexp(self._mode) + + def mean(self): + return symexp(self._mode) + + def log_prob(self, value): + assert self._mode.shape == value.shape, (self._mode.shape, value.shape) + if self._dist == 'mse': + distance = (self._mode - symlog(value)) ** 2 + distance = jnp.where(distance < self._tol, 0, distance) + elif self._dist == 'abs': + distance = jnp.abs(self._mode - symlog(value)) + distance = jnp.where(distance < self._tol, 0, distance) + else: + raise NotImplementedError(self._dist) + if self._agg == 'mean': + loss = distance.mean(self._dims) + elif self._agg == 'sum': + loss = distance.sum(self._dims) + else: + raise NotImplementedError(self._agg) + return -loss + + +class DiscDist: + + def __init__( + self, logits, dims=0, low=-20, high=20, + transfwd=symlog, transbwd=symexp): + self.logits = logits + self.probs = jax.nn.softmax(logits) + self.dims = tuple([-x for x in range(1, dims + 1)]) + self.bins = jnp.linspace(low, high, logits.shape[-1]) + self.low = low + self.high = high + self.transfwd = transfwd + self.transbwd = transbwd + self.batch_shape = logits.shape[:len(logits.shape) - dims - 1] + self.event_shape = logits.shape[len(logits.shape) - dims: -1] + + def mean(self): + return self.transbwd((self.probs * self.bins).sum(-1)) + + def mode(self): + return self.transbwd((self.probs * self.bins).sum(-1)) + + def log_prob(self, x): + x = self.transfwd(x) + below = (self.bins <= x[..., None]).astype(jnp.int32).sum(-1) - 1 + above = len(self.bins) - ( + self.bins > x[..., None]).astype(jnp.int32).sum(-1) + below = jnp.clip(below, 0, len(self.bins) - 1) + above = jnp.clip(above, 0, len(self.bins) - 1) + equal = (below == above) + dist_to_below = jnp.where(equal, 1, jnp.abs(self.bins[below] - x)) + dist_to_above = jnp.where(equal, 1, jnp.abs(self.bins[above] - x)) + total = dist_to_below + dist_to_above + weight_below = dist_to_above / total + weight_above = dist_to_below / total + target = ( + jax.nn.one_hot(below, len(self.bins)) * weight_below[..., None] + + jax.nn.one_hot(above, len(self.bins)) * weight_above[..., None]) + log_pred = self.logits - jax.scipy.special.logsumexp( + self.logits, -1, keepdims=True) + return (target * log_pred).sum(-1).sum(self.dims) + + +def video_grid(video): + B, T, H, W, C = video.shape + return video.transpose((1, 2, 0, 3, 4)).reshape((T, H, B * W, C)) + + +def balance_stats(dist, target, thres): + # Values are NaN when there are no positives or negatives in the current + # batch, which means they will be ignored when aggregating metrics via + # np.nanmean() later, as they should. + pos = (target.astype(jnp.float32) > thres).astype(jnp.float32) + neg = (target.astype(jnp.float32) <= thres).astype(jnp.float32) + pred = (dist.mean().astype(jnp.float32) > thres).astype(jnp.float32) + loss = -dist.log_prob(target) + return dict( + pos_loss=(loss * pos).sum() / pos.sum(), + neg_loss=(loss * neg).sum() / neg.sum(), + pos_acc=(pred * pos).sum() / pos.sum(), + neg_acc=((1 - pred) * neg).sum() / neg.sum(), + rate=pos.mean(), + avg=target.astype(jnp.float32).mean(), + pred=dist.mean().astype(jnp.float32).mean(), + ) + + +class Moments(nj.Module): + + def __init__( + self, impl='mean_std', decay=0.99, max=1e8, eps=0.0, perclo=5, + perchi=95): + self.impl = impl + self.decay = decay + self.max = max + self.eps = eps + self.perclo = perclo + self.perchi = perchi + if self.impl == 'off': + pass + elif self.impl == 'mean_std': + self.step = nj.Variable(jnp.zeros, (), jnp.int32, name='step') + self.mean = nj.Variable(jnp.zeros, (), jnp.float32, name='mean') + self.sqrs = nj.Variable(jnp.zeros, (), jnp.float32, name='sqrs') + elif self.impl == 'min_max': + self.low = nj.Variable(jnp.zeros, (), jnp.float32, name='low') + self.high = nj.Variable(jnp.zeros, (), jnp.float32, name='high') + elif self.impl == 'perc_ema': + self.low = nj.Variable(jnp.zeros, (), jnp.float32, name='low') + self.high = nj.Variable(jnp.zeros, (), jnp.float32, name='high') + elif self.impl == 'perc_ema_corr': + self.step = nj.Variable(jnp.zeros, (), jnp.int32, name='step') + self.low = nj.Variable(jnp.zeros, (), jnp.float32, name='low') + self.high = nj.Variable(jnp.zeros, (), jnp.float32, name='high') + elif self.impl == 'mean_mag': + self.mag = nj.Variable(jnp.zeros, (), jnp.float32, name='mag') + elif self.impl == 'max_mag': + self.mag = nj.Variable(jnp.zeros, (), jnp.float32, name='mag') + else: + raise NotImplementedError(self.impl) + + def __call__(self, x): + self.update(x) + return self.stats() + + def update(self, x): + if parallel(): + mean = lambda x: jax.lax.pmean(x.mean(), 'i') + min_ = lambda x: jax.lax.pmin(x.min(), 'i') + max_ = lambda x: jax.lax.pmax(x.max(), 'i') + per = lambda x, q: jnp.percentile(jax.lax.all_gather(x, 'i'), q) + else: + mean = jnp.mean + min_ = jnp.min + max_ = jnp.max + per = jnp.percentile + x = sg(x.astype(jnp.float32)) + m = self.decay + if self.impl == 'off': + pass + elif self.impl == 'mean_std': + self.step.write(self.step.read() + 1) + self.mean.write(m * self.mean.read() + (1 - m) * mean(x)) + self.sqrs.write(m * self.sqrs.read() + (1 - m) * mean(x * x)) + elif self.impl == 'min_max': + low, high = min_(x), max_(x) + self.low.write(m * jnp.minimum(self.low.read(), low) + (1 - m) * low) + self.high.write(m * jnp.maximum(self.high.read(), high) + (1 - m) * high) + elif self.impl == 'perc_ema': + low, high = per(x, self.perclo), per(x, self.perchi) + self.low.write(m * self.low.read() + (1 - m) * low) + self.high.write(m * self.high.read() + (1 - m) * high) + elif self.impl == 'perc_ema_corr': + self.step.write(self.step.read() + 1) + low, high = per(x, self.perclo), per(x, self.perchi) + self.low.write(m * self.low.read() + (1 - m) * low) + self.high.write(m * self.high.read() + (1 - m) * high) + elif self.impl == 'mean_mag': + curr = mean(jnp.abs(x)) + self.mag.write(m * self.mag.read() + (1 - m) * curr) + elif self.impl == 'max_mag': + curr = max_(jnp.abs(x)) + self.mag.write(m * jnp.maximum(self.mag.read(), curr) + (1 - m) * curr) + else: + raise NotImplementedError(self.impl) + + def stats(self): + if self.impl == 'off': + return 0.0, 1.0 + elif self.impl == 'mean_std': + corr = 1 - self.decay ** self.step.read().astype(jnp.float32) + mean = self.mean.read() / corr + var = (self.sqrs.read() / corr) - self.mean.read() ** 2 + std = jnp.sqrt(jnp.maximum(var, 1 / self.max ** 2) + self.eps) + return sg(mean), sg(std) + elif self.impl == 'min_max': + offset = self.low.read() + invscale = jnp.maximum(1 / self.max, self.high.read() - self.low.read()) + return sg(offset), sg(invscale) + elif self.impl == 'perc_ema': + offset = self.low.read() + invscale = jnp.maximum(1 / self.max, self.high.read() - self.low.read()) + return sg(offset), sg(invscale) + elif self.impl == 'perc_ema_corr': + corr = 1 - self.decay ** self.step.read().astype(jnp.float32) + lo = self.low.read() / corr + hi = self.high.read() / corr + invscale = jnp.maximum(1 / self.max, hi - lo) + return sg(lo), sg(invscale) + elif self.impl == 'mean_mag': + offset = jnp.array(0) + invscale = jnp.maximum(1 / self.max, self.mag.read()) + return sg(offset), sg(invscale) + elif self.impl == 'max_mag': + offset = jnp.array(0) + invscale = jnp.maximum(1 / self.max, self.mag.read()) + return sg(offset), sg(invscale) + else: + raise NotImplementedError(self.impl) + + +class Optimizer(nj.Module): + + PARAM_COUNTS = {} + + def __init__( + self, lr, opt='adam', eps=1e-5, clip=100.0, warmup=0, wd=0.0, + wd_pattern=r'/(w|kernel)$', lateclip=0.0): + assert opt in ('adam', 'belief', 'yogi') + assert wd_pattern[0] not in ('0', '1') + # assert self.path not in self.PARAM_COUNTS + self.PARAM_COUNTS[self.path] = None + wd_pattern = re.compile(wd_pattern) + chain = [] + if clip: + chain.append(optax.clip_by_global_norm(clip)) + if opt == 'adam': + chain.append(optax.scale_by_adam(eps=eps)) + else: + raise NotImplementedError(opt) + if lateclip: + chain.append(late_grad_clip(lateclip)) + if wd: + chain.append(optax.additive_weight_decay(wd, lambda params: ( + tree_map(lambda k: bool(wd_pattern.search(k)), tree_keys(params))))) + if warmup: + schedule = optax.linear_schedule(0.0, -lr, warmup) + chain.append(optax.inject_hyperparams(optax.scale)(schedule)) + else: + chain.append(optax.scale(-lr)) + self.opt = optax.chain(*chain) + self.step = nj.Variable(jnp.array, 0, jnp.int32, name='step') + self.scaling = (COMPUTE_DTYPE == jnp.float16) + if self.scaling: + self.opt = optax.apply_if_finite(self.opt, max_consecutive_errors=1000) + self.grad_scale = nj.Variable( + jnp.array, 1e4, jnp.float32, name='grad_scale') + self.good_steps = nj.Variable( + jnp.array, 0, jnp.int32, name='good_steps') + + def __call__(self, modules, lossfn, *args, has_aux=False, **kwargs): + def wrapped(*args, **kwargs): + outs = lossfn(*args, **kwargs) + loss, aux = outs if has_aux else (outs, None) + assert loss.dtype == jnp.float32, (self.name, loss.dtype) + assert loss.shape == (), (self.name, loss.shape) + if self.scaling: + loss *= sg(self.grad_scale.read()) + return loss, aux + metrics = {} + loss, params, grads, aux = nj.grad( + wrapped, modules, has_aux=True)(*args, **kwargs) + if not self.PARAM_COUNTS[self.path]: + count = sum([np.prod(x.shape) for x in params.values()]) + print(f'Optimizer {self.name} has {count:,} variables.') + self.PARAM_COUNTS[self.path] = count + if parallel(): + grads = tree_map(lambda x: jax.lax.pmean(x, 'i'), grads) + if self.scaling: + grads = tree_map(lambda x: x / self.grad_scale.read(), grads) + finite = self._update_scale(grads) + metrics[f'{self.name}_grad_scale'] = self.grad_scale.read() + metrics[f'{self.name}_grad_overflow'] = (~finite).astype(jnp.float32) + optstate = self.get('state', self.opt.init, params) + updates, optstate = self.opt.update(grads, optstate, params) + self.put('state', optstate) + nj.context().update(optax.apply_updates(params, updates)) + norm = optax.global_norm(grads) + if self.scaling: + norm = jnp.where(jnp.isfinite(norm), norm, jnp.nan) + self.step.write(self.step.read() + jnp.isfinite(norm).astype(jnp.int32)) + metrics['loss'] = loss.mean() + metrics['grad_norm'] = norm + metrics['grad_steps'] = self.step.read() + metrics = {f'{self.name}_{k}': v for k, v in metrics.items()} + return (metrics, aux) if has_aux else metrics + + def _update_scale(self, grads): + finite = jnp.array([ + jnp.isfinite(x).all() for x in jax.tree_util.tree_leaves(grads)]).all() + keep = (finite & (self.good_steps.read() < 1000)) + incr = (finite & (self.good_steps.read() >= 1000)) + decr = ~finite + self.good_steps.write( + keep.astype(jnp.int32) * (self.good_steps.read() + 1)) + self.grad_scale.write(jnp.clip( + keep.astype(jnp.float32) * self.grad_scale.read() + + incr.astype(jnp.float32) * self.grad_scale.read() * 2 + + decr.astype(jnp.float32) * self.grad_scale.read() / 2, + 1e-4, 1e4)) + return finite + + +def late_grad_clip(value=1.0): + def init_fn(params): + return () + def update_fn(updates, state, params): + updates = tree_map(lambda x: jnp.clip(x, -value, value), updates) + return updates, () + return optax.GradientTransformation(init_fn, update_fn) + + +def tree_keys(params, prefix=''): + if hasattr(params, 'items'): + return type(params)({ + k: tree_keys(v, prefix + '/' + k.lstrip('/')) + for k, v in params.items()}) + elif isinstance(params, (tuple, list)): + return [tree_keys(x, prefix) for x in params] + elif isinstance(params, jnp.ndarray): + return prefix + else: + raise TypeError(type(params)) + + +class SlowUpdater: + + def __init__(self, src, dst, fraction=1.0, period=1): + self.src = src + self.dst = dst + self.fraction = fraction + self.period = period + self.updates = nj.Variable(jnp.zeros, (), jnp.int32, name='updates') + + def __call__(self): + assert self.src.getm() + updates = self.updates.read() + need_init = (updates == 0).astype(jnp.float32) + need_update = (updates % self.period == 0).astype(jnp.float32) + mix = jnp.clip(1.0 * need_init + self.fraction * need_update, 0, 1) + source = { + k.replace(f'/{self.src.name}/', f'/{self.dst.name}/'): v + for k, v in self.src.getm().items()} + self.dst.putm(tree_map( + lambda s, d: mix * s + (1 - mix) * d, + source, self.dst.getm())) + self.updates.write(updates + 1) diff --git a/dreamerv3/nets.py b/dreamerv3/nets.py new file mode 100644 index 0000000..d11adda --- /dev/null +++ b/dreamerv3/nets.py @@ -0,0 +1,699 @@ +import re + +import jax +import jax.numpy as jnp +import numpy as np +from tensorflow_probability.substrates import jax as tfp +f32 = jnp.float32 +tfd = tfp.distributions +tree_map = jax.tree_util.tree_map +sg = lambda x: tree_map(jax.lax.stop_gradient, x) + +from . import jaxutils +from . import ninjax as nj +cast = jaxutils.cast_to_compute + + +class RSSM(nj.Module): + + def __init__( + self, deter=1024, stoch=32, classes=32, unroll=False, initial='learned', + unimix=0.01, action_clip=1.0, **kw): + self._deter = deter + self._stoch = stoch + self._classes = classes + self._unroll = unroll + self._initial = initial + self._unimix = unimix + self._action_clip = action_clip + self._kw = kw + + def initial(self, bs): + if self._classes: + state = dict( + deter=jnp.zeros([bs, self._deter], f32), + logit=jnp.zeros([bs, self._stoch, self._classes], f32), + stoch=jnp.zeros([bs, self._stoch, self._classes], f32)) + else: + state = dict( + deter=jnp.zeros([bs, self._deter], f32), + mean=jnp.zeros([bs, self._stoch], f32), + std=jnp.ones([bs, self._stoch], f32), + stoch=jnp.zeros([bs, self._stoch], f32)) + if self._initial == 'zeros': + return cast(state) + elif self._initial == 'learned': + deter = self.get('initial', jnp.zeros, state['deter'][0].shape, f32) + state['deter'] = jnp.repeat(jnp.tanh(deter)[None], bs, 0) + state['stoch'] = self.get_stoch(cast(state['deter'])) + return cast(state) + else: + raise NotImplementedError(self._initial) + + def observe(self, embed, action, is_first, state=None): + swap = lambda x: x.transpose([1, 0] + list(range(2, len(x.shape)))) + if state is None: + state = self.initial(action.shape[0]) + step = lambda prev, inputs: self.obs_step(prev[0], *inputs) + inputs = swap(action), swap(embed), swap(is_first) + start = state, state + post, prior = jaxutils.scan(step, inputs, start, self._unroll) + post = {k: swap(v) for k, v in post.items()} + prior = {k: swap(v) for k, v in prior.items()} + return post, prior + + def imagine(self, action, state=None): + swap = lambda x: x.transpose([1, 0] + list(range(2, len(x.shape)))) + state = self.initial(action.shape[0]) if state is None else state + assert isinstance(state, dict), state + action = swap(action) + prior = jaxutils.scan(self.img_step, action, state, self._unroll) + prior = {k: swap(v) for k, v in prior.items()} + return prior + + def get_dist(self, state, argmax=False): + if self._classes: + logit = state['logit'].astype(f32) + return tfd.Independent(jaxutils.OneHotDist(logit), 1) + else: + mean = state['mean'].astype(f32) + std = state['std'].astype(f32) + return tfp.MultivariateNormalDiag(mean, std) + + def obs_step(self, prev_state, prev_action, embed, is_first): + is_first = cast(is_first) + prev_action = cast(prev_action) + if self._action_clip > 0.0: + prev_action *= sg(self._action_clip / jnp.maximum( + self._action_clip, jnp.abs(prev_action))) + prev_state, prev_action = jax.tree_util.tree_map( + lambda x: self._mask(x, 1.0 - is_first), (prev_state, prev_action)) + prev_state = jax.tree_util.tree_map( + lambda x, y: x + self._mask(y, is_first), + prev_state, self.initial(len(is_first))) + prior = self.img_step(prev_state, prev_action) + x = jnp.concatenate([prior['deter'], embed], -1) + x = self.get('obs_out', Linear, **self._kw)(x) + stats = self._stats('obs_stats', x) + dist = self.get_dist(stats) + stoch = dist.sample(seed=nj.rng()) + post = {'stoch': stoch, 'deter': prior['deter'], **stats} + return cast(post), cast(prior) + + def img_step(self, prev_state, prev_action): + prev_stoch = prev_state['stoch'] + prev_action = cast(prev_action) + if self._action_clip > 0.0: + prev_action *= sg(self._action_clip / jnp.maximum( + self._action_clip, jnp.abs(prev_action))) + if self._classes: + shape = prev_stoch.shape[:-2] + (self._stoch * self._classes,) + prev_stoch = prev_stoch.reshape(shape) + if len(prev_action.shape) > len(prev_stoch.shape): # 2D actions. + shape = prev_action.shape[:-2] + (np.prod(prev_action.shape[-2:]),) + prev_action = prev_action.reshape(shape) + x = jnp.concatenate([prev_stoch, prev_action], -1) + x = self.get('img_in', Linear, **self._kw)(x) + x, deter = self._gru(x, prev_state['deter']) + x = self.get('img_out', Linear, **self._kw)(x) + stats = self._stats('img_stats', x) + dist = self.get_dist(stats) + stoch = dist.sample(seed=nj.rng()) + prior = {'stoch': stoch, 'deter': deter, **stats} + return cast(prior) + + def get_stoch(self, deter): + x = self.get('img_out', Linear, **self._kw)(deter) + stats = self._stats('img_stats', x) + dist = self.get_dist(stats) + return cast(dist.mode()) + + def _gru(self, x, deter): + x = jnp.concatenate([deter, x], -1) + kw = {**self._kw, 'act': 'none', 'units': 3 * self._deter} + x = self.get('gru', Linear, **kw)(x) + reset, cand, update = jnp.split(x, 3, -1) + reset = jax.nn.sigmoid(reset) + cand = jnp.tanh(reset * cand) + update = jax.nn.sigmoid(update - 1) + deter = update * cand + (1 - update) * deter + return deter, deter + + def _stats(self, name, x): + if self._classes: + x = self.get(name, Linear, self._stoch * self._classes)(x) + logit = x.reshape(x.shape[:-1] + (self._stoch, self._classes)) + if self._unimix: + probs = jax.nn.softmax(logit, -1) + uniform = jnp.ones_like(probs) / probs.shape[-1] + probs = (1 - self._unimix) * probs + self._unimix * uniform + logit = jnp.log(probs) + stats = {'logit': logit} + return stats + else: + x = self.get(name, Linear, 2 * self._stoch)(x) + mean, std = jnp.split(x, 2, -1) + std = 2 * jax.nn.sigmoid(std / 2) + 0.1 + return {'mean': mean, 'std': std} + + def _mask(self, value, mask): + return jnp.einsum('b...,b->b...', value, mask.astype(value.dtype)) + + def dyn_loss(self, post, prior, impl='kl', free=1.0): + if impl == 'kl': + loss = self.get_dist(sg(post)).kl_divergence(self.get_dist(prior)) + elif impl == 'logprob': + loss = -self.get_dist(prior).log_prob(sg(post['stoch'])) + else: + raise NotImplementedError(impl) + if free: + loss = jnp.maximum(loss, free) + return loss + + def rep_loss(self, post, prior, impl='kl', free=1.0): + if impl == 'kl': + loss = self.get_dist(post).kl_divergence(self.get_dist(sg(prior))) + elif impl == 'uniform': + uniform = jax.tree_util.tree_map(lambda x: jnp.zeros_like(x), prior) + loss = self.get_dist(post).kl_divergence(self.get_dist(uniform)) + elif impl == 'entropy': + loss = -self.get_dist(post).entropy() + elif impl == 'none': + loss = jnp.zeros(post['deter'].shape[:-1]) + else: + raise NotImplementedError(impl) + if free: + loss = jnp.maximum(loss, free) + return loss + + +class MultiEncoder(nj.Module): + + def __init__( + self, shapes, cnn_keys=r'.*', mlp_keys=r'.*', mlp_layers=4, + mlp_units=512, cnn='resize', cnn_depth=48, + cnn_blocks=2, resize='stride', + symlog_inputs=False, minres=4, **kw): + excluded = ('is_first', 'is_last') + shapes = {k: v for k, v in shapes.items() if ( + k not in excluded and not k.startswith('log_'))} + self.cnn_shapes = {k: v for k, v in shapes.items() if ( + len(v) == 3 and re.match(cnn_keys, k))} + self.mlp_shapes = {k: v for k, v in shapes.items() if ( + len(v) in (1, 2) and re.match(mlp_keys, k))} + self.shapes = {**self.cnn_shapes, **self.mlp_shapes} + print('Encoder CNN shapes:', self.cnn_shapes) + print('Encoder MLP shapes:', self.mlp_shapes) + cnn_kw = {**kw, 'minres': minres, 'name': 'cnn'} + mlp_kw = {**kw, 'symlog_inputs': symlog_inputs, 'name': 'mlp'} + if cnn == 'resnet': + self._cnn = ImageEncoderResnet(cnn_depth, cnn_blocks, resize, **cnn_kw) + else: + raise NotImplementedError(cnn) + if self.mlp_shapes: + self._mlp = MLP(None, mlp_layers, mlp_units, dist='none', **mlp_kw) + + def __call__(self, data): + some_key, some_shape = list(self.shapes.items())[0] + batch_dims = data[some_key].shape[:-len(some_shape)] + data = { + k: v.reshape((-1,) + v.shape[len(batch_dims):]) + for k, v in data.items()} + outputs = [] + if self.cnn_shapes: + inputs = jnp.concatenate([data[k] for k in self.cnn_shapes], -1) + output = self._cnn(inputs) + output = output.reshape((output.shape[0], -1)) + outputs.append(output) + if self.mlp_shapes: + inputs = [ + data[k][..., None] if len(self.shapes[k]) == 0 else data[k] + for k in self.mlp_shapes] + inputs = jnp.concatenate([x.astype(f32) for x in inputs], -1) + inputs = jaxutils.cast_to_compute(inputs) + outputs.append(self._mlp(inputs)) + outputs = jnp.concatenate(outputs, -1) + outputs = outputs.reshape(batch_dims + outputs.shape[1:]) + return outputs + + +class MultiDecoder(nj.Module): + + def __init__( + self, shapes, inputs=['tensor'], cnn_keys=r'.*', mlp_keys=r'.*', + mlp_layers=4, mlp_units=512, cnn='resize', cnn_depth=48, cnn_blocks=2, + image_dist='mse', vector_dist='mse', resize='stride', bins=255, + outscale=1.0, minres=4, cnn_sigmoid=False, **kw): + excluded = ('is_first', 'is_last', 'is_terminal', 'reward') + shapes = {k: v for k, v in shapes.items() if k not in excluded} + self.cnn_shapes = { + k: v for k, v in shapes.items() + if re.match(cnn_keys, k) and len(v) == 3} + self.mlp_shapes = { + k: v for k, v in shapes.items() + if re.match(mlp_keys, k) and len(v) == 1} + self.shapes = {**self.cnn_shapes, **self.mlp_shapes} + print('Decoder CNN shapes:', self.cnn_shapes) + print('Decoder MLP shapes:', self.mlp_shapes) + cnn_kw = {**kw, 'minres': minres, 'sigmoid': cnn_sigmoid} + mlp_kw = {**kw, 'dist': vector_dist, 'outscale': outscale, 'bins': bins} + if self.cnn_shapes: + shapes = list(self.cnn_shapes.values()) + assert all(x[:-1] == shapes[0][:-1] for x in shapes) + shape = shapes[0][:-1] + (sum(x[-1] for x in shapes),) + if cnn == 'resnet': + self._cnn = ImageDecoderResnet( + shape, cnn_depth, cnn_blocks, resize, **cnn_kw, name='cnn') + else: + raise NotImplementedError(cnn) + if self.mlp_shapes: + self._mlp = MLP( + self.mlp_shapes, mlp_layers, mlp_units, **mlp_kw, name='mlp') + self._inputs = Input(inputs, dims='deter') + self._image_dist = image_dist + + def __call__(self, inputs, drop_loss_indices=None): + features = self._inputs(inputs) + dists = {} + if self.cnn_shapes: + feat = features + if drop_loss_indices is not None: + feat = feat[:, drop_loss_indices] + flat = feat.reshape([-1, feat.shape[-1]]) + output = self._cnn(flat) + output = output.reshape(feat.shape[:-1] + output.shape[1:]) + split_indices = np.cumsum([v[-1] for v in self.cnn_shapes.values()][:-1]) + means = jnp.split(output, split_indices, -1) + dists.update({ + key: self._make_image_dist(key, mean) + for (key, shape), mean in zip(self.cnn_shapes.items(), means)}) + if self.mlp_shapes: + dists.update(self._mlp(features)) + return dists + + def _make_image_dist(self, name, mean): + mean = mean.astype(f32) + if self._image_dist == 'normal': + return tfd.Independent(tfd.Normal(mean, 1), 3) + if self._image_dist == 'mse': + return jaxutils.MSEDist(mean, 3, 'sum') + raise NotImplementedError(self._image_dist) + + +class ImageEncoderResnet(nj.Module): + + def __init__(self, depth, blocks, resize, minres, **kw): + self._depth = depth + self._blocks = blocks + self._resize = resize + self._minres = minres + self._kw = kw + + def __call__(self, x): + stages = int(np.log2(x.shape[-2]) - np.log2(self._minres)) + depth = self._depth + x = jaxutils.cast_to_compute(x) - 0.5 + # print(x.shape) + for i in range(stages): + kw = {**self._kw, 'preact': False} + if self._resize == 'stride': + x = self.get(f's{i}res', Conv2D, depth, 4, 2, **kw)(x) + elif self._resize == 'stride3': + s = 2 if i else 3 + k = 5 if i else 4 + x = self.get(f's{i}res', Conv2D, depth, k, s, **kw)(x) + elif self._resize == 'mean': + N, H, W, D = x.shape + x = self.get(f's{i}res', Conv2D, depth, 3, 1, **kw)(x) + x = x.reshape((N, H // 2, W // 2, 4, D)).mean(-2) + elif self._resize == 'max': + x = self.get(f's{i}res', Conv2D, depth, 3, 1, **kw)(x) + x = jax.lax.reduce_window( + x, -jnp.inf, jax.lax.max, (1, 3, 3, 1), (1, 2, 2, 1), 'same') + else: + raise NotImplementedError(self._resize) + for j in range(self._blocks): + skip = x + kw = {**self._kw, 'preact': True} + x = self.get(f's{i}b{j}conv1', Conv2D, depth, 3, **kw)(x) + x = self.get(f's{i}b{j}conv2', Conv2D, depth, 3, **kw)(x) + x += skip + # print(x.shape) + depth *= 2 + if self._blocks: + x = get_act(self._kw['act'])(x) + x = x.reshape((x.shape[0], -1)) + # print(x.shape) + return x + + +class ImageDecoderResnet(nj.Module): + + def __init__(self, shape, depth, blocks, resize, minres, sigmoid, **kw): + self._shape = shape + self._depth = depth + self._blocks = blocks + self._resize = resize + self._minres = minres + self._sigmoid = sigmoid + self._kw = kw + + def __call__(self, x): + stages = int(np.log2(self._shape[-2]) - np.log2(self._minres)) + depth = self._depth * 2 ** (stages - 1) + x = jaxutils.cast_to_compute(x) + x = self.get('in', Linear, (self._minres, self._minres, depth))(x) + for i in range(stages): + for j in range(self._blocks): + skip = x + kw = {**self._kw, 'preact': True} + x = self.get(f's{i}b{j}conv1', Conv2D, depth, 3, **kw)(x) + x = self.get(f's{i}b{j}conv2', Conv2D, depth, 3, **kw)(x) + x += skip + # print(x.shape) + depth //= 2 + kw = {**self._kw, 'preact': False} + if i == stages - 1: + kw = {} + depth = self._shape[-1] + if self._resize == 'stride': + x = self.get(f's{i}res', Conv2D, depth, 4, 2, transp=True, **kw)(x) + elif self._resize == 'stride3': + s = 3 if i == stages - 1 else 2 + k = 5 if i == stages - 1 else 4 + x = self.get(f's{i}res', Conv2D, depth, k, s, transp=True, **kw)(x) + elif self._resize == 'resize': + x = jnp.repeat(jnp.repeat(x, 2, 1), 2, 2) + x = self.get(f's{i}res', Conv2D, depth, 3, 1, **kw)(x) + else: + raise NotImplementedError(self._resize) + if max(x.shape[1:-1]) > max(self._shape[:-1]): + padh = (x.shape[1] - self._shape[0]) / 2 + padw = (x.shape[2] - self._shape[1]) / 2 + x = x[:, int(np.ceil(padh)): -int(padh), :] + x = x[:, :, int(np.ceil(padw)): -int(padw)] + # print(x.shape) + assert x.shape[-3:] == self._shape, (x.shape, self._shape) + if self._sigmoid: + x = jax.nn.sigmoid(x) + else: + x = x + 0.5 + return x + + +class MLP(nj.Module): + + def __init__( + self, shape, layers, units, inputs=['tensor'], dims=None, + symlog_inputs=False, **kw): + assert shape is None or isinstance(shape, (int, tuple, dict)), shape + if isinstance(shape, int): + shape = (shape,) + self._shape = shape + self._layers = layers + self._units = units + self._inputs = Input(inputs, dims=dims) + self._symlog_inputs = symlog_inputs + distkeys = ( + 'dist', 'outscale', 'minstd', 'maxstd', 'outnorm', 'unimix', 'bins') + self._dense = {k: v for k, v in kw.items() if k not in distkeys} + self._dist = {k: v for k, v in kw.items() if k in distkeys} + + def __call__(self, inputs): + feat = self._inputs(inputs) + if self._symlog_inputs: + feat = jaxutils.symlog(feat) + x = jaxutils.cast_to_compute(feat) + x = x.reshape([-1, x.shape[-1]]) + for i in range(self._layers): + x = self.get(f'h{i}', Linear, self._units, **self._dense)(x) + x = x.reshape(feat.shape[:-1] + (x.shape[-1],)) + if self._shape is None: + return x + elif isinstance(self._shape, tuple): + return self._out('out', self._shape, x) + elif isinstance(self._shape, dict): + return {k: self._out(k, v, x) for k, v in self._shape.items()} + else: + raise ValueError(self._shape) + + def _out(self, name, shape, x): + return self.get(f'dist_{name}', Dist, shape, **self._dist)(x) + + +class Dist(nj.Module): + + def __init__( + self, shape, dist='mse', outscale=0.1, outnorm=False, minstd=1.0, + maxstd=1.0, unimix=0.0, bins=255): + assert all(isinstance(dim, int) for dim in shape), shape + self._shape = shape + self._dist = dist + self._minstd = minstd + self._maxstd = maxstd + self._unimix = unimix + self._outscale = outscale + self._outnorm = outnorm + self._bins = bins + + def __call__(self, inputs): + dist = self.inner(inputs) + assert tuple(dist.batch_shape) == tuple(inputs.shape[:-1]), ( + dist.batch_shape, dist.event_shape, inputs.shape) + return dist + + def inner(self, inputs): + kw = {} + kw['outscale'] = self._outscale + kw['outnorm'] = self._outnorm + shape = self._shape + if self._dist.endswith('_disc'): + shape = (*self._shape, self._bins) + out = self.get('out', Linear, int(np.prod(shape)), **kw)(inputs) + out = out.reshape(inputs.shape[:-1] + shape).astype(f32) + if self._dist in ('normal', 'trunc_normal'): + std = self.get('std', Linear, int(np.prod(self._shape)), **kw)(inputs) + std = std.reshape(inputs.shape[:-1] + self._shape).astype(f32) + if self._dist == 'symlog_mse': + return jaxutils.SymlogDist(out, len(self._shape), 'mse', 'sum') + if self._dist == 'symlog_disc': + return jaxutils.DiscDist( + out, len(self._shape), -20, 20, jaxutils.symlog, jaxutils.symexp) + if self._dist == 'mse': + return jaxutils.MSEDist(out, len(self._shape), 'sum') + if self._dist == 'normal': + lo, hi = self._minstd, self._maxstd + std = (hi - lo) * jax.nn.sigmoid(std + 2.0) + lo + dist = tfd.Normal(jnp.tanh(out), std) + dist = tfd.Independent(dist, len(self._shape)) + dist.minent = np.prod(self._shape) * tfd.Normal(0.0, lo).entropy() + dist.maxent = np.prod(self._shape) * tfd.Normal(0.0, hi).entropy() + return dist + if self._dist == 'binary': + dist = tfd.Bernoulli(out) + return tfd.Independent(dist, len(self._shape)) + if self._dist == 'onehot': + if self._unimix: + probs = jax.nn.softmax(out, -1) + uniform = jnp.ones_like(probs) / probs.shape[-1] + probs = (1 - self._unimix) * probs + self._unimix * uniform + out = jnp.log(probs) + dist = jaxutils.OneHotDist(out) + if len(self._shape) > 1: + dist = tfd.Independent(dist, len(self._shape) - 1) + dist.minent = 0.0 + dist.maxent = np.prod(self._shape[:-1]) * jnp.log(self._shape[-1]) + return dist + raise NotImplementedError(self._dist) + + +class Conv2D(nj.Module): + + def __init__( + self, depth, kernel, stride=1, transp=False, act='none', norm='none', + pad='same', bias=True, preact=False, winit='uniform', fan='avg'): + self._depth = depth + self._kernel = kernel + self._stride = stride + self._transp = transp + self._act = get_act(act) + self._norm = Norm(norm, name='norm') + self._pad = pad.upper() + self._bias = bias and (preact or norm == 'none') + self._preact = preact + self._winit = winit + self._fan = fan + + def __call__(self, hidden): + if self._preact: + hidden = self._norm(hidden) + hidden = self._act(hidden) + hidden = self._layer(hidden) + else: + hidden = self._layer(hidden) + hidden = self._norm(hidden) + hidden = self._act(hidden) + return hidden + + def _layer(self, x): + if self._transp: + shape = (self._kernel, self._kernel, self._depth, x.shape[-1]) + kernel = self.get('kernel', Initializer( + self._winit, fan=self._fan), shape) + kernel = jaxutils.cast_to_compute(kernel) + x = jax.lax.conv_transpose( + x, kernel, (self._stride, self._stride), self._pad, + dimension_numbers=('NHWC', 'HWOI', 'NHWC')) + else: + shape = (self._kernel, self._kernel, x.shape[-1], self._depth) + kernel = self.get('kernel', Initializer( + self._winit, fan=self._fan), shape) + kernel = jaxutils.cast_to_compute(kernel) + x = jax.lax.conv_general_dilated( + x, kernel, (self._stride, self._stride), self._pad, + dimension_numbers=('NHWC', 'HWIO', 'NHWC')) + if self._bias: + bias = self.get('bias', jnp.zeros, self._depth, np.float32) + bias = jaxutils.cast_to_compute(bias) + x += bias + return x + + +class Linear(nj.Module): + + def __init__( + self, units, act='none', norm='none', bias=True, outscale=1.0, + outnorm=False, winit='uniform', fan='avg'): + self._units = tuple(units) if hasattr(units, '__len__') else (units,) + self._act = get_act(act) + self._norm = norm + self._bias = bias and norm == 'none' + self._outscale = outscale + self._outnorm = outnorm + self._winit = winit + self._fan = fan + + def __call__(self, x): + shape = (x.shape[-1], np.prod(self._units)) + kernel = self.get('kernel', Initializer( + self._winit, self._outscale, fan=self._fan), shape) + kernel = jaxutils.cast_to_compute(kernel) + x = x @ kernel + if self._bias: + bias = self.get('bias', jnp.zeros, np.prod(self._units), np.float32) + bias = jaxutils.cast_to_compute(bias) + x += bias + if len(self._units) > 1: + x = x.reshape(x.shape[:-1] + self._units) + x = self.get('norm', Norm, self._norm)(x) + x = self._act(x) + return x + + +class Norm(nj.Module): + + def __init__(self, impl): + self._impl = impl + + def __call__(self, x): + dtype = x.dtype + if self._impl == 'none': + return x + elif self._impl == 'layer': + x = x.astype(f32) + x = jax.nn.standardize(x, axis=-1, epsilon=1e-3) + x *= self.get('scale', jnp.ones, x.shape[-1], f32) + x += self.get('bias', jnp.zeros, x.shape[-1], f32) + return x.astype(dtype) + else: + raise NotImplementedError(self._impl) + + +class Input: + + def __init__(self, keys=['tensor'], dims=None): + assert isinstance(keys, (list, tuple)), keys + self._keys = tuple(keys) + self._dims = dims or self._keys[0] + + def __call__(self, inputs): + if not isinstance(inputs, dict): + inputs = {'tensor': inputs} + inputs = inputs.copy() + for key in self._keys: + if key.startswith('softmax_'): + inputs[key] = jax.nn.softmax(inputs[key[len('softmax_'):]]) + if not all(k in inputs for k in self._keys): + needs = f'{{{", ".join(self._keys)}}}' + found = f'{{{", ".join(inputs.keys())}}}' + raise KeyError(f'Cannot find keys {needs} among inputs {found}.') + values = [inputs[k] for k in self._keys] + dims = len(inputs[self._dims].shape) + for i, value in enumerate(values): + if len(value.shape) > dims: + values[i] = value.reshape( + value.shape[:dims - 1] + (np.prod(value.shape[dims - 1:]),)) + values = [x.astype(inputs[self._dims].dtype) for x in values] + return jnp.concatenate(values, -1) + + +class Initializer: + + def __init__(self, dist='uniform', scale=1.0, fan='avg'): + self.scale = scale + self.dist = dist + self.fan = fan + + def __call__(self, shape): + if self.scale == 0.0: + value = jnp.zeros(shape, f32) + elif self.dist == 'uniform': + fanin, fanout = self._fans(shape) + denoms = {'avg': (fanin + fanout) / 2, 'in': fanin, 'out': fanout} + scale = self.scale / denoms[self.fan] + limit = np.sqrt(3 * scale) + value = jax.random.uniform( + nj.rng(), shape, f32, -limit, limit) + elif self.dist == 'normal': + fanin, fanout = self._fans(shape) + denoms = {'avg': np.mean((fanin, fanout)), 'in': fanin, 'out': fanout} + scale = self.scale / denoms[self.fan] + std = np.sqrt(scale) / 0.87962566103423978 + value = std * jax.random.truncated_normal( + nj.rng(), -2, 2, shape, f32) + elif self.dist == 'ortho': + nrows, ncols = shape[-1], np.prod(shape) // shape[-1] + matshape = (nrows, ncols) if nrows > ncols else (ncols, nrows) + mat = jax.random.normal(nj.rng(), matshape, f32) + qmat, rmat = jnp.linalg.qr(mat) + qmat *= jnp.sign(jnp.diag(rmat)) + qmat = qmat.T if nrows < ncols else qmat + qmat = qmat.reshape(nrows, *shape[:-1]) + value = self.scale * jnp.moveaxis(qmat, 0, -1) + else: + raise NotImplementedError(self.dist) + return value + + def _fans(self, shape): + if len(shape) == 0: + return 1, 1 + elif len(shape) == 1: + return shape[0], shape[0] + elif len(shape) == 2: + return shape + else: + space = int(np.prod(shape[:-2])) + return shape[-2] * space, shape[-1] * space + + +def get_act(name): + if callable(name): + return name + elif name == 'none': + return lambda x: x + elif name == 'mish': + return lambda x: x * jnp.tanh(jax.nn.softplus(x)) + elif hasattr(jax.nn, name): + return getattr(jax.nn, name) + else: + raise NotImplementedError(name) diff --git a/dreamerv3/ninjax.py b/dreamerv3/ninjax.py new file mode 100644 index 0000000..385a428 --- /dev/null +++ b/dreamerv3/ninjax.py @@ -0,0 +1,502 @@ +import contextlib +import functools +import inspect +import re +import threading +from functools import partial as bind + +import jax +import jax.numpy as jnp + +__version__ = '0.9.0' + + +############################################################################### +# State +############################################################################### + + +# When running an impure function that accesses state, it will find the state +# in this global variable. The pure() wrapper populates this global variable +# with the provided state, calls the inner function, and then the takes the +# resulting state out of the global variable to return it back to the user. +# To allow multi-threaded programs to use impure functions in parallel, the +# context is a dictionary with a slot for each thread identifier. +CONTEXT = {} + + +class Context(dict): + + def __init__(self, entries, rng, create, modify, ignore, reserve, name): + super().__init__(entries) + self.create = create # Allow creating new state entries. + self.modify = modify # Allow modifying existing state entries. + self.ignore = ignore # Ignore modifications to existing state entries. + self.rng = rng + self.reserve = reserve + self.name = name + + def update(self, entries): + for key, value in dict(entries).items(): + self[key] = value + + def __setitem__(self, key, value): + if not self.modify: + raise RuntimeError( + 'Cannot modify state entries here. If you want to modify ' + 'state inside of scan() set modify=True. ' + + f'You were trying to set {key} to shape {value.shape} and ' + + f'dtype {value.dtype}.') + if self.ignore and key in self: + return # Do not overwrite existing entries. + if not self.create and key not in self: + raise RuntimeError( + 'Can only create state entries during first call. ' + + f'You were trying to set {key} to shape {value.shape} and ' + + f'dtype {value.dtype}.') + super().__setitem__(key, value) + + +def pure(fun, nested=False): + """Wrap an impure function that uses global state to explicitly pass the + state in and out. The result is a pure function that is composable with JAX + transformation. The pure function can be used as follows: + `out, state = fun(state, rng, *args, **kwargs)`.""" + def purified( + state, rng, *args, create=None, modify=None, ignore=None, **kwargs): + context = CONTEXT.get(threading.get_ident(), None) + if context: + create = create if create is not None else context.create + modify = modify if modify is not None else context.modify + ignore = ignore if ignore is not None else context.ignore + assert context.create or not create, 'Parent context disabled create.' + assert context.modify or not modify, 'Parent context disabled modify.' + assert not context.ignore or ignore, 'Parent context enabled ignore.' + else: + create = create if create is not None else True + modify = modify if modify is not None else True + ignore = ignore if ignore is not None else False + if not isinstance(state, dict): + raise ValueError('Must provide a dict as state.') + if context and (not nested): + raise RuntimeError( + f'You are trying to call pure {fun.__name__}() inside pure ' + f'{context.name}(). Is that intentional? If you want to nest pure ' + f'functions, use pure(..., nested=True) for the inner function.') + # raise RuntimeError( + # f'If you want to nest run() calls, use nested=True. ({context})') + before = context + try: + name = fun.__name__ + if rng.shape == (): + rng = jax.random.PRNGKey(rng) + context = Context(state.copy(), rng, create, modify, ignore, [], name) + CONTEXT[threading.get_ident()] = context + out = fun(*args, **kwargs) + state = dict(context) + return out, state + finally: + CONTEXT[threading.get_ident()] = before + purified.pure = True + return purified + + +def context(): + """Access and modify the global context from within an impure function. For + advanced users only. Prefer to use module methods to access and modify state + and rng() to get the next RNG key.""" + context = CONTEXT.get(threading.get_ident(), None) + if context is None: + raise RuntimeError('Wrap impure functions in pure() before running them.') + return context + + +@jax.named_scope('rng') +def rng(amount=None, reserve=16): + """Split the global RNG key and return a new local key.""" + ctx = context() + if amount: + keys = jax.random.split(ctx.rng, amount + 1) + ctx.rng = keys[0] + return keys[1:] + else: + if not ctx.reserve: + keys = jax.random.split(ctx.rng, reserve) + ctx.rng = keys[0] + ctx.reserve = list(keys[1:]) + return ctx.reserve.pop(0) + + +def creating(): + """Indicates whether the program is currently allowed to create state + entries. Can use used for initialization logic that should be excluded from + compiled functions.""" + return context().create + + +############################################################################### +# Transformations +############################################################################### + + +@jax.named_scope('grad') +def grad(fun, keys, has_aux=False): + """Compute the gradient of an impure function with respect to the specified + state entries or modules. The transformed function returns a tuple containing + the computed value, selected state entries, their gradients, and if + applicable auxiliary outputs of the function.""" + keys = keys if hasattr(keys, '__len__') else (keys,) + if getattr(fun, 'pure', False): + raise ValueError('Use plain jax.grad() for pure functions.') + if not has_aux: + fun = lambda *args, _fun=fun, **kwargs: (_fun(*args, *kwargs), {}) + fun = pure(fun, nested=True) + def forward(x1, x2, rng, *args, **kwargs): + (y, aux), state = fun({**x1, **x2}, rng, *args, create=False, **kwargs) + return y, (aux, state) + backward = jax.value_and_grad(forward, has_aux=True) + @functools.wraps(backward) + def wrapper(*args, **kwargs): + _prerun(fun, *args, **kwargs) + assert all(isinstance(x, (str, Module)) for x in keys) + strs = [x for x in keys if isinstance(x, str)] + mods = [x for x in keys if isinstance(x, Module)] + for mod in mods: + strs += mod.getm() + x1 = {k: v for k, v in context().items() if k in strs} + x2 = {k: v for k, v in context().items() if k not in strs} + (y, (aux, state)), dx = backward(x1, x2, rng(), *args, **kwargs) + context().update(state) + return (y, x1, dx, aux) if has_aux else (y, x1, dx) + return wrapper + + +def jit(fun, static=None, **kwargs): + """Compiles a pure function for fast execution. Only the first call of the + function is allowed to create state entries.""" + if not getattr(fun, 'pure', False): + raise ValueError('Use pure() before applying jit().') + static = static or () + + @bind(jax.jit, static_argnums=[0], **kwargs) + def init(statics, rng, *args, **kw): + # Return only state so JIT can remove dead code for fast initialization. + s = fun({}, rng, *args, ignore=True, **dict(statics), **kw)[1] + return s + + @bind(jax.jit, static_argnums=[0], **kwargs) + def apply(statics, state, rng, *args, **kw): + return fun(state, rng, *args, create=False, **dict(statics), **kw) + + @functools.wraps(fun) + def wrapper(state, rng, *args, init_only=False, **kw): + if any([name not in kw for name in static]): + raise ValueError('Please pass all static arguments by keyword.') + state = state.copy() + statics = tuple(sorted([(k, v) for k, v in kw.items() if k in static])) + kw = {k: v for k, v in kw.items() if k not in static} + if not hasattr(wrapper, 'keys'): + created = init(statics, rng, *args, **kw) + wrapper.keys = set(created.keys()) + for key, value in created.items(): + if key not in state: + state[key] = value + if init_only: + return state + else: + selected = {k: v for k, v in state.items() if k in wrapper.keys} + out, updated = apply(statics, selected, rng, *args, **kw) + return out, {**state, **updated} + return wrapper + + +def pmap(fun, axis_name=None, static=None, **kwargs): + """Compiles n pure function for fast execution across multiple devices. Only + the first call of the function is allowed to create state entries.""" + if not getattr(fun, 'pure', False): + raise ValueError('Use pure() before applying jit().') + static = static or () + + @bind( + jax.pmap, axis_name=axis_name, static_broadcasted_argnums=[0], **kwargs) + def init(statics, rng, *args, **kw): + # Return only state so JIT can remove dead code for fast initialization. + return fun({}, rng, *args, ignore=True, **dict(statics), **kw)[1] + + @bind( + jax.pmap, axis_name=axis_name, static_broadcasted_argnums=[0], **kwargs) + def apply(statics, state, rng, *args, **kw): + return fun(state, rng, *args, create=False, **dict(statics), **kw) + + @functools.wraps(fun) + def wrapper(state, rng, *args, init_only=False, **kw): + if any([name not in kw for name in static]): + raise ValueError('Please pass all static arguments by keyword.') + state = state.copy() + statics = tuple(sorted([(k, v) for k, v in kw.items() if k in static])) + kw = {k: v for k, v in kw.items() if k not in static} + if not hasattr(wrapper, 'keys'): + created = init(statics, rng, *args, **kw) + wrapper.keys = set(created.keys()) + for key, value in created.items(): + if key not in state: + state[key] = value + if init_only: + return state + else: + selected = {k: v for k, v in state.items() if k in wrapper.keys} + out, updated = apply(statics, selected, rng, *args, **kw) + return out, {**state, **updated} + return wrapper + + +@jax.named_scope('cond') +def cond(pred, true_fun, false_fun, *operands): + true_fun = pure(true_fun, nested=True) + false_fun = pure(false_fun, nested=True) + _prerun(true_fun, *operands) + _prerun(false_fun, *operands) + out, state = jax.lax.cond( + pred, + lambda state, rng1, rng2, *args: true_fun(state, rng1, *args), + lambda state, rng1, rng2, *args: false_fun(state, rng2, *args), + dict(context()), *rng(2), *operands) + context().update(state) + return out + + +@jax.named_scope('scan') +def scan(fun, carry, xs, reverse=False, unroll=1, modify=False): + fun = pure(fun, nested=True) + _prerun(fun, carry, jax.tree_util.tree_map(lambda x: x[0], xs)) + length = len(jax.tree_util.tree_leaves(xs)[0]) + rngs = rng(length) + if modify: + def inner(carry, x): + carry, state = carry + x, rng = x + (carry, y), state = fun(state, rng, carry, x, create=False) + return (carry, state), y + (carry, state), ys = jax.lax.scan( + inner, (carry, dict(context())), (xs, rngs), length, reverse, unroll) + context().update(state) + else: + def inner(carry, x): + x, rng = x + (carry, y), state = fun( + dict(context()), rng, carry, x, create=False, modify=False) + return carry, y + carry, ys = jax.lax.scan(inner, carry, (xs, rngs), length, reverse, unroll) + return carry, ys + + +@jax.named_scope('_prerun') +def _prerun(fun, *args, **kwargs): + if not context().create: + return + discarded, state = fun(dict(context()), rng(), *args, ignore=True, **kwargs) + # jax.tree_util.tree_map( + # lambda x: hasattr(x, 'delete') and x.delete(), discarded) + context().update(state) + + +############################################################################### +# Modules +############################################################################### + + +SCOPE = '' + + +@contextlib.contextmanager +def scope(name, absolute=False): + """Enter a relative or absolute name scope. Name scopes are used to make + names of state entries unique.""" + global SCOPE + if SCOPE is None: + raise RuntimeError( + 'Purify stateful functions with fn = pure(fn) before running them.') + outside = SCOPE + if absolute: + SCOPE = name + elif SCOPE == '': + SCOPE = name + else: + SCOPE = outside + '/' + name + yield SCOPE + SCOPE = outside + + +class ModuleMeta(type): + + """Meta class that creates a unique path for each module instance and wraps + the methods and properties of the module to enter the name scope.""" + + def __new__(mcs, name, bases, clsdict): + """This runs once per user module class definition. It wraps the methods of + the module class to automatically enter the name scope of the module.""" + method_names = [] + for key, value in clsdict.items(): + if key.startswith('__') and key != '__call__': + continue + elif isinstance(value, property): + clsdict[key] = property( + value.fget if not value.fget else _scope_method(value.fget), + value.fset if not value.fset else _scope_method(value.fset), + value.fdel if not value.fdel else _scope_method(value.fdel), + doc=value.__doc__) + elif inspect.isfunction(value): + method_names.append(key) + cls = super(ModuleMeta, mcs).__new__(mcs, name, bases, clsdict) + for method_name in method_names: + method = getattr(cls, method_name) + method = _scope_method(method) + setattr(cls, method_name, method) + return cls + + def __call__(cls, *args, name=None, **kwargs): + """This runs once per use module instance creation. It derives a unique + name and path for the module instance.""" + if not isinstance(name, str): + raise KeyError( + "Please provide a module name via Module(..., name='example').") + if not re.match(r'[A-Za-z0-9_]+', name): + raise KeyError( + 'Only letters, numbers, and underscores are allowed in scope names.') + obj = cls.__new__(cls) + with scope(name) as path: + obj._path = path + obj._submodules = {} + init = _scope_method(cls.__init__) + init(obj, *args, **kwargs) + return obj + + +def _scope_method(method): + @functools.wraps(method) + def wrapper(self, *args, **kwargs): + with scope(self._path, absolute=True): + with jax.named_scope(self._path.split('/')[-1]): + return method(self, *args, **kwargs) + return wrapper + + +class Module(object, metaclass=ModuleMeta): + + """Base class for users to inherit their modules from. Provides automatic + name scoping via the meta class and helper functions for accessing state.""" + + def __repr__(self): + return f'{self.__class__.__name__}({self.path})' + + @property + def path(self): + """The unique name scope of this module instance as a string.""" + return self._path + + @property + def name(self): + """The name of this module instance as a string.""" + return self._path.split('/')[-1] + + def get(self, name, *args, **kwargs): + """Retrieve or create a state entry that belongs to this module.""" + path = self.path + '/' + name + if name in self._submodules: + return self._submodules[name] + if path in context(): + return context()[path] + ctor, *args = args + if 'name' in inspect.signature(ctor).parameters: + kwargs['name'] = name + value = ctor(*args, **kwargs) + flat, _ = jax.tree_util.tree_flatten(value) + if all(isinstance(x, jnp.ndarray) for x in flat): + context()[path] = value + else: + self._submodules[name] = value + return value + + def put(self, name, value): + """Update or create a single state entry that belongs to this module.""" + self.putm({self.path + '/' + name: value}) + return value + + def getm(self, pattern=r'.*', allow_empty=False): + """Read the state entries of this module, optionally filtered by regex.""" + pattern = re.compile(pattern) + prefix = self.path + '/' + results = {} + for key, value in context().items(): + if not key.startswith(prefix): + continue + if pattern.match(key[len(prefix):]): + results[key] = value + if not allow_empty and not results: + raise KeyError(f'Pattern {pattern} matched no state keys.') + return results + + def putm(self, mapping): + """Update or create multiple state entries that belong to this module.""" + prefix = self.path + '/' + for key in mapping: + if not key.startswith(prefix): + raise KeyError(f'Key {key} does not belong to module {self.path}.') + context().update(mapping) + + +class Variable(Module): + + def __init__(self, ctor, *args, **kwargs): + self.ctor = ctor + self.args = args + self.kwargs = kwargs + + def read(self): + return self.get('value', self.ctor, *self.args, **self.kwargs) + + def write(self, value): + return self.put('value', value) + + +############################################################################### +# Integrations +############################################################################### + + +class HaikuModule(Module): + + def __init__(self, ctor, *args, **kwargs): + import haiku as hk + def net(*args_, **kwargs_): + return ctor(*args, **kwargs)(*args_, **kwargs_) + self.transformed = hk.transform(net) + + def __call__(self, *args, **kwargs): + state = self.get('state', self.transformed.init, rng(), *args, **kwargs) + return self.transformed.apply(state, rng(), *args, **kwargs) + + +class FlaxModule(Module): + + def __init__(self, ctor, *args, **kwargs): + self.module = ctor(*args, **kwargs) + + def __call__(self, *args, **kwargs): + state = self.get('state', self.module.init, rng(), *args, **kwargs) + return self.module.apply(state, *args, **kwargs) + + +class OptaxModule(Module): + + def __init__(self, ctor, *args, **kwargs): + self.opt = ctor(*args, **kwargs) + + def __call__(self, loss, keys, *args, **kwargs): + import optax + loss, params, grads = grad(loss, keys)(*args, **kwargs) + optstate = self.get('state', self.opt.init, params) + updates, optstate = self.opt.update(grads, optstate) + self.put('state', optstate) + context().update(optax.apply_updates(params, updates)) + return {'loss': loss.mean(), 'grad_norm': optax.global_norm(grads)} diff --git a/dreamerv3/plot_crafter.py b/dreamerv3/plot_crafter.py new file mode 100644 index 0000000..d89c513 --- /dev/null +++ b/dreamerv3/plot_crafter.py @@ -0,0 +1,126 @@ +"""Simple script to plot results of a crafter run with dreamerv3-cr.""" +import pathlib, os +import numpy as np +import pandas as pd +import matplotlib.pyplot as plt +import collections +import json +import warnings +import argparse + +def main(): + runs = [] + + home = os.path.expanduser("~") + default_filename = f'{home}/logdir/crafter-dv3-cr_1/stats.jsonl' + + parser = argparse.ArgumentParser(description='Process some arguments') + parser.add_argument('--filename', default=default_filename, + help='The path to stats.jsonl file') + args = parser.parse_args() + + filename = args.filename + print(f"Filename: {filename}") + budget=1e6 + rewards, lengths, achievements = load_stats(pathlib.Path(filename), budget) + task, method, seed = pathlib.Path(filename).parts[-5:-2] + + print(f'Run length {sum(lengths)}: {filename}') + runs.append(dict( + task=task, + method=method, + seed=str(id), + xs=np.cumsum(lengths).tolist(), + reward=rewards, + length=lengths, + **achievements, + )) + + scores, tasks, percents, methods = print_summary(runs, budget, verbose=True) + + if len(rewards) < 100: + print('>>> Not plotting reward curve until at least 100 episodes saved out') + else: + plt.plot(np.cumsum(lengths).tolist(), pd.Series(rewards).rolling(window=100).mean().values) + plt.xlabel('Steps') + plt.ylabel('Reward') + plt.show() + +def load_stats(filename, budget): + steps = 0 + rewards = [] + lengths = [] + achievements = collections.defaultdict(list) + for line in filename.read_text().split('\n'): + if not line.strip(): + continue + episode = json.loads(line) + steps += episode['length'] + if steps > budget: + break + lengths.append(episode['length']) + for key, value in episode.items(): + if key.startswith('achievement_'): + achievements[key].append(value) + unlocks = int(np.sum([(v[-1] >= 1) for v in achievements.values()])) + health = -0.9 + rewards.append(unlocks + health) + return rewards, lengths, achievements + + +def print_summary(runs, budget, verbose): + episodes = np.array([len(x['length']) for x in runs]) + rewards = np.array([np.mean(x['reward']) for x in runs]) + lengths = np.array([np.mean(x['length']) for x in runs]) + percents, methods, seeds, tasks = compute_success_rates( + runs, budget, sortby=0) + scores = np.squeeze(compute_scores(percents)) + print(f'Score: {np.mean(scores):10.2f} ± {np.std(scores):.2f}') + print(f'Reward: {np.mean(rewards):10.2f} ± {np.std(rewards):.2f}') + print(f'Length: {np.mean(lengths):10.2f} ± {np.std(lengths):.2f}') + print(f'Episodes: {np.mean(episodes):10.2f} ± {np.std(episodes):.2f}') + if verbose: + for task, percent in zip(tasks, np.squeeze(percents).T): + name = task[len('achievement_'):].replace('_', ' ').title() + print(f'{name:<20} {np.mean(percent):6.2f}%') + return scores, tasks, percents, methods + + +def compute_success_rates(runs, budget=1e6, sortby=None): + methods = sorted(set(run['method'] for run in runs)) + seeds = sorted(set(run['seed'] for run in runs)) + tasks = sorted(key for key in runs[0] if key.startswith('achievement_')) + percents = np.empty((len(methods), len(seeds), len(tasks))) + percents[:] = np.nan + for run in runs: + episodes = (np.array(run['xs']) <= budget).sum() + i = methods.index(run['method']) + j = seeds.index(run['seed']) + for key, values in run.items(): + if key in tasks: + k = tasks.index(key) + percent = 100 * (np.array(values[:episodes]) >= 1).mean() + if np.isnan(percent): + print(percent) + percents[i, j, k] = percent + if isinstance(sortby, (str, int)): + if isinstance(sortby, str): + sortby = methods.index(sortby) + order = np.argsort(-np.nanmean(percents[sortby], 0), -1) + percents = percents[:, :, order] + tasks = np.array(tasks)[order].tolist() + return percents, methods, seeds, tasks + + +def compute_scores(percents): + # Geometric mean with an offset of 1%. + assert (0 <= percents).all() and (percents <= 100).all() + if (percents <= 1.0).all(): + print('Warning: The input may not be in the right range.') + with warnings.catch_warnings(): # Empty seeds become NaN. + warnings.simplefilter('ignore', category=RuntimeWarning) + scores = np.exp(np.nanmean(np.log(1 + percents), -1)) - 1 + return scores + +if __name__ == "__main__": + main() diff --git a/dreamerv3/train.py b/dreamerv3/train.py new file mode 100644 index 0000000..81f2e31 --- /dev/null +++ b/dreamerv3/train.py @@ -0,0 +1,216 @@ +import importlib +import pathlib +import sys +import warnings +from functools import partial as bind + +warnings.filterwarnings('ignore', '.*box bound precision lowered.*') +warnings.filterwarnings('ignore', '.*using stateful random seeds*') +warnings.filterwarnings('ignore', '.*is a deprecated alias for.*') +warnings.filterwarnings('ignore', '.*truncated to dtype int32.*') + +directory = pathlib.Path(__file__).resolve() +directory = directory.parent +sys.path.append(str(directory.parent)) +sys.path.append(str(directory.parent.parent)) +sys.path.append(str(directory.parent.parent.parent)) +__package__ = directory.name + +import embodied +from embodied import wrappers + + +def main(argv=None): + from . import agent as agt + + parsed, other = embodied.Flags(configs=['defaults']).parse_known(argv) + config = embodied.Config(agt.Agent.configs['defaults']) + for name in parsed.configs: + config = config.update(agt.Agent.configs[name]) + config = embodied.Flags(config).parse(other) + args = embodied.Config( + **config.run, logdir=config.logdir, + batch_steps=config.batch_size * config.batch_length) + print(config) + + logdir = embodied.Path(args.logdir) + logdir.mkdirs() + config.save(logdir / 'config.yaml') + step = embodied.Counter() + logger = make_logger(parsed, logdir, step, config) + + cleanup = [] + try: + + if args.script == 'train': + replay = make_replay(config, logdir / 'replay') + env = make_envs(config) + cleanup.append(env) + agent = agt.Agent(env.obs_space, env.act_space, step, config) + embodied.run.train(agent, env, replay, logger, args) + + elif args.script == 'train_save': + replay = make_replay(config, logdir / 'replay') + env = make_envs(config) + cleanup.append(env) + agent = agt.Agent(env.obs_space, env.act_space, step, config) + embodied.run.train_save(agent, env, replay, logger, args) + + elif args.script == 'train_eval': + replay = make_replay(config, logdir / 'replay') + eval_replay = make_replay(config, logdir / 'eval_replay', is_eval=True) + env = make_envs(config) + eval_env = make_envs(config) # mode='eval' + cleanup += [env, eval_env] + agent = agt.Agent(env.obs_space, env.act_space, step, config) + embodied.run.train_eval( + agent, env, eval_env, replay, eval_replay, logger, args) + + elif args.script == 'train_holdout': + replay = make_replay(config, logdir / 'replay') + if config.eval_dir: + assert not config.train.eval_fill + eval_replay = make_replay(config, config.eval_dir, is_eval=True) + else: + assert 0 < args.eval_fill <= config.replay_size // 10, args.eval_fill + eval_replay = make_replay(config, logdir / 'eval_replay', is_eval=True) + env = make_envs(config) + cleanup.append(env) + agent = agt.Agent(env.obs_space, env.act_space, step, config) + embodied.run.train_holdout( + agent, env, replay, eval_replay, logger, args) + + elif args.script == 'eval_only': + env = make_envs(config) # mode='eval' + cleanup.append(env) + agent = agt.Agent(env.obs_space, env.act_space, step, config) + embodied.run.eval_only(agent, env, logger, args) + + elif args.script == 'parallel': + assert config.run.actor_batch <= config.envs.amount, ( + config.run.actor_batch, config.envs.amount) + step = embodied.Counter() + env = make_env(config) + agent = agt.Agent(env.obs_space, env.act_space, step, config) + env.close() + replay = make_replay(config, logdir / 'replay', rate_limit=True) + embodied.run.parallel( + agent, replay, logger, bind(make_env, config), + num_envs=config.envs.amount, args=args) + + else: + raise NotImplementedError(args.script) + finally: + for obj in cleanup: + obj.close() + + +def make_logger(parsed, logdir, step, config): + multiplier = config.env.get(config.task.split('_')[0], {}).get('repeat', 1) + logger = embodied.Logger(step, [ + embodied.logger.TerminalOutput(config.filter), + embodied.logger.JSONLOutput(logdir, 'metrics.jsonl'), + embodied.logger.JSONLOutput(logdir, 'scores.jsonl', 'episode/score'), + embodied.logger.TensorBoardOutput(logdir), + # embodied.logger.WandBOutput(logdir.name, config), + # embodied.logger.MLFlowOutput(logdir.name), + ], multiplier) + return logger + + +def make_replay( + config, directory=None, is_eval=False, rate_limit=False, **kwargs): + assert config.replay == 'uniform' or not rate_limit + length = config.batch_length + size = config.replay_size // 10 if is_eval else config.replay_size + if config.replay == 'uniform' or is_eval: + kw = {'online': config.replay_online} + if rate_limit and config.run.train_ratio > 0: + kw['samples_per_insert'] = config.run.train_ratio / config.batch_length + kw['tolerance'] = 10 * config.batch_size + kw['min_size'] = config.batch_size + replay = embodied.replay.Uniform(length, size, directory, **kw) + elif config.replay == 'reverb': + replay = embodied.replay.Reverb(length, size, directory) + elif config.replay == 'curious-replay': + replay = embodied.replay.CuriousReplay(length, size, directory, hyper=config.replay_hyper) + elif config.replay == 'per': + replay = embodied.replay.PrioritizedExperienceReplay(length, size, directory, hyper=config.replay_hyper) + elif config.replay == 'count-based': + replay = embodied.replay.CountBasedReplay(length, size, directory, hyper=config.replay_hyper) + elif config.replay == 'adversarial': + replay = embodied.replay.AdversarialReplay(length, size, directory, hyper=config.replay_hyper) + elif config.replay == 'chunks': + replay = embodied.replay.NaiveChunks(length, size, directory) + else: + raise NotImplementedError(config.replay) + return replay + + +def make_envs(config, **overrides): + suite, task = config.task.split('_', 1) + ctors = [] + for index in range(config.envs.amount): + ctor = lambda: make_env(config, **overrides) + if config.envs.parallel != 'none': + ctor = bind(embodied.Parallel, ctor, config.envs.parallel) + if config.envs.restart: + ctor = bind(wrappers.RestartOnException, ctor) + ctors.append(ctor) + envs = [ctor() for ctor in ctors] + return embodied.BatchEnv(envs, parallel=(config.envs.parallel != 'none')) + + +def make_env(config, **overrides): + # You can add custom environments by creating and returning the environment + # instance here. Environments with different interfaces can be converted + # using `embodied.envs.from_gym.FromGym` and `embodied.envs.from_dm.FromDM`. + suite, task = config.task.split('_', 1) + ctor = { + 'dummy': 'embodied.envs.dummy:Dummy', + 'gym': 'embodied.envs.from_gym:FromGym', + 'dm': 'embodied.envs.from_dmenv:FromDM', + 'crafter': 'embodied.envs.crafter:Crafter', + 'dmc': 'embodied.envs.dmc:DMC', + 'atari': 'embodied.envs.atari:Atari', + 'dmlab': 'embodied.envs.dmlab:DMLab', + 'minecraft': 'embodied.envs.minecraft:Minecraft', + 'loconav': 'embodied.envs.loconav:LocoNav', + 'pinpad': 'embodied.envs.pinpad:PinPad', + 'cdmc': 'embodied.envs.cdmc:CDMC', + 'ddmc': 'embodied.envs.ddmc:DDMC', + }[suite] + if isinstance(ctor, str): + module, cls = ctor.split(':') + module = importlib.import_module(module) + ctor = getattr(module, cls) + kwargs = config.env.get(suite, {}) + kwargs.update(overrides) + env = ctor(task, **kwargs) + return wrap_env(env, config) + + +def wrap_env(env, config): + args = config.wrapper + for name, space in env.act_space.items(): + if name == 'reset': + continue + elif space.discrete: + env = wrappers.OneHotAction(env, name) + elif args.discretize: + env = wrappers.DiscretizeAction(env, name, args.discretize) + else: + env = wrappers.NormalizeAction(env, name) + env = wrappers.ExpandScalars(env) + if args.length: + env = wrappers.TimeLimit(env, args.length, args.reset) + if args.checks: + env = wrappers.CheckSpaces(env) + for name, space in env.act_space.items(): + if not space.discrete: + env = wrappers.ClipAction(env, name) + return env + + +if __name__ == '__main__': + main() diff --git a/example.py b/example.py new file mode 100644 index 0000000..50d0d46 --- /dev/null +++ b/example.py @@ -0,0 +1,53 @@ +def main(): + + import warnings + import dreamerv3 + from dreamerv3 import embodied + warnings.filterwarnings('ignore', '.*truncated to dtype int32.*') + + # See configs.yaml for all options. + config = embodied.Config(dreamerv3.configs['defaults']) + config = config.update(dreamerv3.configs['medium']) + config = config.update({ + 'logdir': '~/logdir/run1', + 'run.train_ratio': 64, + 'run.log_every': 30, # Seconds + 'batch_size': 16, + 'jax.prealloc': False, + 'encoder.mlp_keys': '$^', + 'decoder.mlp_keys': '$^', + 'encoder.cnn_keys': 'image', + 'decoder.cnn_keys': 'image', + # 'jax.platform': 'cpu', + }) + config = embodied.Flags(config).parse() + + logdir = embodied.Path(config.logdir) + step = embodied.Counter() + logger = embodied.Logger(step, [ + embodied.logger.TerminalOutput(), + embodied.logger.JSONLOutput(logdir, 'metrics.jsonl'), + embodied.logger.TensorBoardOutput(logdir), + # embodied.logger.WandBOutput(logdir.name, config), + # embodied.logger.MLFlowOutput(logdir.name), + ]) + + import crafter + from embodied.envs import from_gym + env = crafter.Env() # Replace this with your Gym env. + env = from_gym.FromGym(env) + env = dreamerv3.wrap_env(env, config) + env = embodied.BatchEnv([env], parallel=False) + + agent = dreamerv3.Agent(env.obs_space, env.act_space, step, config) + replay = embodied.replay.Uniform( + config.batch_length, config.replay_size, logdir / 'replay') + args = embodied.Config( + **config.run, logdir=config.logdir, + batch_steps=config.batch_size * config.batch_length) + embodied.run.train(agent, env, replay, logger, args) + # embodied.run.eval_only(agent, env, logger, args) + + +if __name__ == '__main__': + main() diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..e718d4a --- /dev/null +++ b/requirements.txt @@ -0,0 +1,13 @@ +cloudpickle +crafter +gym==0.26.2 +dm_control +jax +jaxlib +numpy +optax +rich +ruamel.yaml +tensorflow-cpu +tensorflow_probability +tensorflow diff --git a/run_local.py b/run_local.py new file mode 100644 index 0000000..17f4683 --- /dev/null +++ b/run_local.py @@ -0,0 +1,37 @@ +import time + +def main(): + + n = int(time.time()) + + args = \ + ['--logdir', f'/home/cd/src/aal/logdir/crafter-dv3-20230502-{n}', + '--env.crafter.outdir', f'/home/cd/src/aal/logdir/crafter-dv3-20230502-{n}', + #'--configs', 'crafter', + '--configs', 'crafter', 'small', + #'--configs', 'dmc_vision', + #'--task', 'cdmc_cartpole_swingup', + '--jax.jit', 'True', + '--replay', 'curious-replay', # curious-replay; per; count-based; adversarial + #'--replay', 'per', + '--replay_hyper.initial_priority', '1e5', + '--replay_hyper.c', '1e4', + '--replay_hyper.beta', '0.7', + '--replay_hyper.epsilon', '0.01', + '--replay_hyper.alpha', '0.7', + # '--run.script', 'train_eval', + # '--run.steps', '1.5e4', + # '--run.eval_every', '1e4', + # '--run.eval_initial', 'False', + # '--run.eval_eps', '100', + '--envs.amount', '1', + #'--batch_size', '8', + ] + + print('Local launch of Dreamer v3 🚀...') + from dreamerv3 import train + train.main(args) + + +if __name__ == '__main__': + main() diff --git a/setup.py b/setup.py new file mode 100644 index 0000000..a332e4c --- /dev/null +++ b/setup.py @@ -0,0 +1,22 @@ +import pathlib +import setuptools +from setuptools import find_namespace_packages + + +setuptools.setup( + name='dreamerv3', + version='1.5.0', + description='Mastering Diverse Domains through World Models', + url='http://github.com/danijar/dreamerv3', + long_description=pathlib.Path('README.md').read_text(), + long_description_content_type='text/markdown', + packages=find_namespace_packages(exclude=['example.py']), + include_package_data=True, + install_requires=pathlib.Path('requirements.txt').read_text().splitlines(), + classifiers=[ + 'Intended Audience :: Science/Research', + 'License :: OSI Approved :: MIT License', + 'Programming Language :: Python :: 3', + 'Topic :: Scientific/Engineering :: Artificial Intelligence', + ], +)