wassname/cr-dv3
1
0
Fork 0
mirror of https://github.com/wassname/cr-dv3.git synced 2026-09-21 12:50:20 +08:00
Code Issues Packages Projects Releases Wiki Activity

First commit.

Browse Source
This commit is contained in:
Isaac Kauvar
2023-06-22 09:54:27 -04:00
commit 2c4a2c08cf
86 changed files with 10667 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+7
View File
@@ -0,0 +1,7 @@
.pytest_cache
dist
__pycache__/
*.py[cod]
*.egg-info
MUJOCO_LOG.TXT
;
LICENSE
+19
View File
@@ -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.
MANIFEST.in
+1
View File
@@ -0,0 +1 @@
include requirements.txt
README.md
+71
View File
@@ -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
dreamerv3/Dockerfile
+63
View File
@@ -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
dreamerv3/__init__.py
+6
View File
@@ -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
dreamerv3/agent.py
+396
View File
@@ -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]
dreamerv3/behaviors.py
+102
View File
@@ -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 {}
dreamerv3/configs.yaml
+274
View File
@@ -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
dreamerv3/embodied/__init__.py
+11
View File
@@ -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
dreamerv3/embodied/core/__init__.py
+29
View File
@@ -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
dreamerv3/embodied/core/base.py
+119
View File
@@ -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
dreamerv3/embodied/core/basics.py
+144
View File
@@ -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
dreamerv3/embodied/core/batch.py
+45
View File
@@ -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
dreamerv3/embodied/core/batcher.py
+101
View File
@@ -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
dreamerv3/embodied/core/checkpoint.py
+93
View File
@@ -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()
dreamerv3/embodied/core/config.py
+191
View File
@@ -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__)
dreamerv3/embodied/core/counter.py
+44
View File
@@ -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
dreamerv3/embodied/core/distr.py
+221
View File
@@ -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)
dreamerv3/embodied/core/driver.py
+78
View File
@@ -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
dreamerv3/embodied/core/flags.py
+102
View File
@@ -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)
dreamerv3/embodied/core/logger.py
+323
View File
@@ -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
dreamerv3/embodied/core/metrics.py
+42
View File
@@ -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()
dreamerv3/embodied/core/parallel.py
+51
View File
@@ -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
dreamerv3/embodied/core/path.py
+223
View File
@@ -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),
]
dreamerv3/embodied/core/random.py
+14
View File
@@ -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
dreamerv3/embodied/core/space.py
+103
View File
@@ -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)
dreamerv3/embodied/core/timer.py
+61
View File
@@ -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)))
dreamerv3/embodied/core/uuid.py
+74
View File
@@ -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
dreamerv3/embodied/core/when.py
+88
View File
@@ -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
dreamerv3/embodied/core/worker.py
+241
View File
@@ -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
dreamerv3/embodied/core/wrappers.py
+364
View File
@@ -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)
dreamerv3/embodied/envs/atari.py
+140
View File
@@ -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()
dreamerv3/embodied/envs/cdmc.py
+64
View File
@@ -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)
dreamerv3/embodied/envs/crafter.py
+73
View File
@@ -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()
dreamerv3/embodied/envs/ddmc.py
+85
View File
@@ -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)
dreamerv3/embodied/envs/dmc.py
+65
View File
@@ -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)
dreamerv3/embodied/envs/dmlab.py
+142
View File
@@ -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}')
dreamerv3/embodied/envs/dummy.py
+54
View File
@@ -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,
)
dreamerv3/embodied/envs/from_dm.py
+85
View File
@@ -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)
dreamerv3/embodied/envs/from_gym.py
+118
View File
@@ -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)
dreamerv3/embodied/envs/loconav.py
+230
View File
@@ -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': (
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
'. . . . . . . . . . . . . . . . .',
),
}
dreamerv3/embodied/envs/loconav_quadruped.py
+132
View File
@@ -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'))
dreamerv3/embodied/envs/loconav_quadruped.xml
+311
View File
@@ -0,0 +1,311 @@
<mujoco model="quadruped">
<visual>
<quality shadowsize="2048"/>
<rgba rangefinder="1 1 0.1 0.1"/>
</visual>
<asset>
<texture name="grid" type="2d" builtin="checker" rgb1=".1 .2 .3" rgb2=".2 .3 .4" width="300" height="300" mark="edge" markrgb=".2 .3 .4"/>
<material name="grid" texture="grid" texrepeat="1 1" texuniform="true" reflectance=".2"/>
<material name="self" rgba=".7 .5 .3 1"/>
<material name="self_default" rgba=".7 .5 .3 1"/>
<material name="self_highlight" rgba="0 .5 .3 1"/>
<material name="effector" rgba=".7 .4 .2 1"/>
<material name="effector_default" rgba=".7 .4 .2 1"/>
<material name="effector_highlight" rgba="0 .5 .3 1"/>
<material name="decoration" rgba=".3 .5 .7 1"/>
<material name="eye" rgba="0 .2 1 1"/>
<material name="target" rgba=".6 .3 .3 1"/>
<material name="target_default" rgba=".6 .3 .3 1"/>
<material name="target_highlight" rgba=".6 .3 .3 .4"/>
<material name="site" rgba=".5 .5 .5 .3"/>
<hfield name="terrain" ncol="201" nrow="201" size="30 30 5 .1"/>
</asset>
<option timestep=".005"/>
<default>
<geom solimp=".9 .99 .003" solref=".01 1"/>
<default class="body">
<geom type="capsule" size=".08" condim="1" material="self" density="500"/>
<joint type="hinge" damping="30" armature=".01"
limited="true" solimplimit="0 .99 .01"/>
<default class="hip">
<default class="yaw">
<joint axis="0 0 1" range="-50 50"/>
</default>
<default class="pitch">
<joint axis="0 1 0" range="-20 60"/>
</default>
<geom fromto="0 0 0 .3 0 .11"/>
</default>
<default class="knee">
<joint axis="0 1 0" range="-60 50"/>
<geom size=".065" fromto="0 0 0 .25 0 -.25"/>
</default>
<default class="ankle">
<joint axis="0 1 0" range="-45 55"/>
<geom size=".055" fromto="0 0 0 0 0 -.25"/>
</default>
<default class="toe">
<geom type="sphere" size=".08" material="effector" friction="1.5"/>
<site type="sphere" size=".084" material="site" group="4"/>
</default>
</default>
<default class="rangefinder">
<site type="capsule" size=".005 .1" material="site" group="4"/>
</default>
<default class="coupling">
<equality solimp="0.95 0.99 0.01" solref=".005 .5"/>
</default>
<general ctrllimited="true" gainprm="1000" biasprm="0 -1000" biastype="affine" dyntype="filter" dynprm=".1"/>
<default class="yaw_act">
<general ctrlrange="-1 1"/>
</default>
<default class="lift_act">
<general ctrlrange="-1 1.1"/>
</default>
<default class="extend_act">
<general ctrlrange="-.8 .8"/>
</default>
</default>
<worldbody>
<camera name="sideon" pos="0 -10 5" fovy="45" mode="targetbody" target="torso" />
<camera name="float_far" pos="-4 0 2" xyaxes="0 -1 0 .5 0 1" mode="trackcom" fovy="90"/>
<body name="torso" childclass="body" pos="0 0 .57">
<camera name="x" pos="-1.7 0 1" xyaxes="0 -1 0 .75 0 1" mode="trackcom"/>
<camera name="y" pos="0 4 2" xyaxes="-1 0 0 0 -.5 1" mode="trackcom"/>
<camera name="egocentric" pos=".3 0 .11" xyaxes="0 -1 0 .4 0 1" fovy="60"/>
<light name="light" pos="0 0 4" mode="trackcom"/>
<geom name="eye_r" type="cylinder" size=".05" fromto=".1 -.07 .12 .31 -.07 .08" mass="0"/>
<site name="pupil_r" type="sphere" size=".033" pos=".3 -.07 .08" zaxis="1 0 0" material="eye"/>
<geom name="eye_l" type="cylinder" size=".05" fromto=".1 .07 .12 .31 .07 .08" mass="0"/>
<site name="pupil_l" type="sphere" size=".033" pos=".3 .07 .08" zaxis="1 0 0" material="eye"/>
<site name="workspace" type="sphere" size=".3 .3 .3" material="site" pos=".8 0 -.2" group="3"/>
<site name="rf_00" class="rangefinder" fromto=".41 -.02 .11 .34 0 .115"/>
<site name="rf_01" class="rangefinder" fromto=".41 -.01 .11 .34 0 .115"/>
<site name="rf_02" class="rangefinder" fromto=".41 0 .11 .34 0 .115"/>
<site name="rf_03" class="rangefinder" fromto=".41 .01 .11 .34 0 .115"/>
<site name="rf_04" class="rangefinder" fromto=".41 .02 .11 .34 0 .115"/>
<site name="rf_10" class="rangefinder" fromto=".41 -.02 .1 .36 0 .11"/>
<site name="rf_11" class="rangefinder" fromto=".41 -.02 .1 .36 0 .11"/>
<site name="rf_12" class="rangefinder" fromto=".41 0 .1 .36 0 .11"/>
<site name="rf_13" class="rangefinder" fromto=".41 .01 .1 .36 0 .11"/>
<site name="rf_14" class="rangefinder" fromto=".41 .02 .1 .36 0 .11"/>
<site name="rf_20" class="rangefinder" fromto=".41 -.02 .09 .38 0 .105"/>
<site name="rf_21" class="rangefinder" fromto=".41 -.01 .09 .38 0 .105"/>
<site name="rf_22" class="rangefinder" fromto=".41 0 .09 .38 0 .105"/>
<site name="rf_23" class="rangefinder" fromto=".41 .01 .09 .38 0 .105"/>
<site name="rf_24" class="rangefinder" fromto=".41 .02 .09 .38 0 .105"/>
<site name="rf_30" class="rangefinder" fromto=".41 -.02 .08 .4 0 .1"/>
<site name="rf_31" class="rangefinder" fromto=".41 -.01 .08 .4 0 .1"/>
<site name="rf_32" class="rangefinder" fromto=".41 0 .08 .4 0 .1"/>
<site name="rf_33" class="rangefinder" fromto=".41 .01 .08 .4 0 .1"/>
<site name="rf_34" class="rangefinder" fromto=".41 .02 .08 .4 0 .1"/>
<geom name="torso" type="ellipsoid" size=".3 .27 .2" density="1000"/>
<site name="torso_touch" type="box" size=".26 .26 .26" rgba="0 0 1 0"/>
<site name="torso" size=".05" rgba="1 0 0 1" />
<body name="hip_front_left" pos=".2 .2 0" euler="0 0 45" childclass="hip">
<joint name="yaw_front_left" class="yaw"/>
<joint name="pitch_front_left" class="pitch"/>
<geom name="thigh_front_left"/>
<body name="knee_front_left" pos=".3 0 .11" childclass="knee">
<joint name="knee_front_left"/>
<geom name="shin_front_left"/>
<body name="ankle_front_left" pos=".25 0 -.25" childclass="ankle">
<joint name="ankle_front_left"/>
<geom name="foot_front_left"/>
<body name="toe_front_left" pos="0 0 -.3" childclass="toe">
<geom name="toe_front_left"/>
<site name="toe_front_left"/>
</body>
</body>
</body>
</body>
<body name="hip_front_right" pos=".2 -.2 0" euler="0 0 -45" childclass="hip">
<joint name="yaw_front_right" class="yaw"/>
<joint name="pitch_front_right" class="pitch"/>
<geom name="thigh_front_right"/>
<body name="knee_front_right" pos=".3 0 .11" childclass="knee">
<joint name="knee_front_right"/>
<geom name="shin_front_right"/>
<body name="ankle_front_right" pos=".25 0 -.25" childclass="ankle">
<joint name="ankle_front_right"/>
<geom name="foot_front_right"/>
<body name="toe_front_right" pos="0 0 -.3" childclass="toe">
<geom name="toe_front_right"/>
<site name="toe_front_right"/>
</body>
</body>
</body>
</body>
<body name="hip_back_right" pos="-.2 -.2 0" euler="0 0 -135" childclass="hip">
<joint name="yaw_back_right" class="yaw"/>
<joint name="pitch_back_right" class="pitch"/>
<geom name="thigh_back_right"/>
<body name="knee_back_right" pos=".3 0 .11" childclass="knee">
<joint name="knee_back_right"/>
<geom name="shin_back_right"/>
<body name="ankle_back_right" pos=".25 0 -.25" childclass="ankle">
<joint name="ankle_back_right"/>
<geom name="foot_back_right"/>
<body name="toe_back_right" pos="0 0 -.3" childclass="toe">
<geom name="toe_back_right"/>
<site name="toe_back_right"/>
</body>
</body>
</body>
</body>
<body name="hip_back_left" pos="-.2 .2 0" euler="0 0 135" childclass="hip">
<joint name="yaw_back_left" class="yaw"/>
<joint name="pitch_back_left" class="pitch"/>
<geom name="thigh_back_left"/>
<body name="knee_back_left" pos=".3 0 .11" childclass="knee">
<joint name="knee_back_left"/>
<geom name="shin_back_left"/>
<body name="ankle_back_left" pos=".25 0 -.25" childclass="ankle">
<joint name="ankle_back_left"/>
<geom name="foot_back_left"/>
<body name="toe_back_left" pos="0 0 -.3" childclass="toe">
<geom name="toe_back_left"/>
<site name="toe_back_left"/>
</body>
</body>
</body>
</body>
</body>
</worldbody>
<tendon>
<fixed name="coupling_front_left">
<joint joint="pitch_front_left" coef=".333"/>
<joint joint="knee_front_left" coef=".333"/>
<joint joint="ankle_front_left" coef=".333"/>
</fixed>
<fixed name="coupling_front_right">
<joint joint="pitch_front_right" coef=".333"/>
<joint joint="knee_front_right" coef=".333"/>
<joint joint="ankle_front_right" coef=".333"/>
</fixed>
<fixed name="coupling_back_right">
<joint joint="pitch_back_right" coef=".333"/>
<joint joint="knee_back_right" coef=".333"/>
<joint joint="ankle_back_right" coef=".333"/>
</fixed>
<fixed name="coupling_back_left">
<joint joint="pitch_back_left" coef=".333"/>
<joint joint="knee_back_left" coef=".333"/>
<joint joint="ankle_back_left" coef=".333"/>
</fixed>
<fixed name="extend_front_left">
<joint joint="pitch_front_left" coef=".25"/>
<joint joint="knee_front_left" coef="-.5"/>
<joint joint="ankle_front_left" coef=".25"/>
</fixed>
<fixed name="lift_front_left">
<joint joint="pitch_front_left" coef=".5"/>
<joint joint="ankle_front_left" coef="-.5"/>
</fixed>
<fixed name="extend_front_right">
<joint joint="pitch_front_right" coef=".25"/>
<joint joint="knee_front_right" coef="-.5"/>
<joint joint="ankle_front_right" coef=".25"/>
</fixed>
<fixed name="lift_front_right">
<joint joint="pitch_front_right" coef=".5"/>
<joint joint="ankle_front_right" coef="-.5"/>
</fixed>
<fixed name="extend_back_right">
<joint joint="pitch_back_right" coef=".25"/>
<joint joint="knee_back_right" coef="-.5"/>
<joint joint="ankle_back_right" coef=".25"/>
</fixed>
<fixed name="lift_back_right">
<joint joint="pitch_back_right" coef=".5"/>
<joint joint="ankle_back_right" coef="-.5"/>
</fixed>
<fixed name="extend_back_left">
<joint joint="pitch_back_left" coef=".25"/>
<joint joint="knee_back_left" coef="-.5"/>
<joint joint="ankle_back_left" coef=".25"/>
</fixed>
<fixed name="lift_back_left">
<joint joint="pitch_back_left" coef=".5"/>
<joint joint="ankle_back_left" coef="-.5"/>
</fixed>
</tendon>
<equality>
<tendon name="coupling_front_left" tendon1="coupling_front_left" class="coupling"/>
<tendon name="coupling_front_right" tendon1="coupling_front_right" class="coupling"/>
<tendon name="coupling_back_right" tendon1="coupling_back_right" class="coupling"/>
<tendon name="coupling_back_left" tendon1="coupling_back_left" class="coupling"/>
</equality>
<actuator>
<general name="yaw_front_left" class="yaw_act" joint="yaw_front_left"/>
<general name="lift_front_left" class="lift_act" tendon="lift_front_left"/>
<general name="extend_front_left" class="extend_act" tendon="extend_front_left"/>
<general name="yaw_front_right" class="yaw_act" joint="yaw_front_right"/>
<general name="lift_front_right" class="lift_act" tendon="lift_front_right"/>
<general name="extend_front_right" class="extend_act" tendon="extend_front_right"/>
<general name="yaw_back_right" class="yaw_act" joint="yaw_back_right"/>
<general name="lift_back_right" class="lift_act" tendon="lift_back_right"/>
<general name="extend_back_right" class="extend_act" tendon="extend_back_right"/>
<general name="yaw_back_left" class="yaw_act" joint="yaw_back_left"/>
<general name="lift_back_left" class="lift_act" tendon="lift_back_left"/>
<general name="extend_back_left" class="extend_act" tendon="extend_back_left"/>
</actuator>
<sensor>
<accelerometer name="imu_accel" site="torso"/>
<gyro name="imu_gyro" site="torso"/>
<velocimeter name="velocimeter" site="torso"/>
<force name="force_toe_front_left" site="toe_front_left"/>
<force name="force_toe_front_right" site="toe_front_right"/>
<force name="force_toe_back_right" site="toe_back_right"/>
<force name="force_toe_back_left" site="toe_back_left"/>
<torque name="torque_toe_front_left" site="toe_front_left"/>
<torque name="torque_toe_front_right" site="toe_front_right"/>
<torque name="torque_toe_back_right" site="toe_back_right"/>
<torque name="torque_toe_back_left" site="toe_back_left"/>
<subtreecom name="center_of_mass" body="torso"/>
<rangefinder name="rf_00" site="rf_00"/>
<rangefinder name="rf_01" site="rf_01"/>
<rangefinder name="rf_02" site="rf_02"/>
<rangefinder name="rf_03" site="rf_03"/>
<rangefinder name="rf_04" site="rf_04"/>
<rangefinder name="rf_10" site="rf_10"/>
<rangefinder name="rf_11" site="rf_11"/>
<rangefinder name="rf_12" site="rf_12"/>
<rangefinder name="rf_13" site="rf_13"/>
<rangefinder name="rf_14" site="rf_14"/>
<rangefinder name="rf_20" site="rf_20"/>
<rangefinder name="rf_21" site="rf_21"/>
<rangefinder name="rf_22" site="rf_22"/>
<rangefinder name="rf_23" site="rf_23"/>
<rangefinder name="rf_24" site="rf_24"/>
<rangefinder name="rf_30" site="rf_30"/>
<rangefinder name="rf_31" site="rf_31"/>
<rangefinder name="rf_32" site="rf_32"/>
<rangefinder name="rf_33" site="rf_33"/>
<rangefinder name="rf_34" site="rf_34"/>
</sensor>
</mujoco>
dreamerv3/embodied/envs/minecraft.py
+156
View File
@@ -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'),
}
dreamerv3/embodied/envs/minecraft_base.py
+189
View File
@@ -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
dreamerv3/embodied/envs/minecraft_minerl.py
+150
View File
@@ -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 '<BreakSpeedMultiplier>{{multiplier}}</BreakSpeedMultiplier>'
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 '<GammaSetting>{{gamma}}</GammaSetting>'
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',
)
dreamerv3/embodied/envs/pinpad.py
+220
View File
@@ -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')
dreamerv3/embodied/envs/robodesk.py
+37
View File
@@ -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
dreamerv3/embodied/replay/__init__.py
+10
View File
@@ -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
dreamerv3/embodied/replay/adversarial.py
+12
View File
@@ -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'])
dreamerv3/embodied/replay/base_prioritized_reverb.py
+205
View File
@@ -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
dreamerv3/embodied/replay/chunk.py
+78
View File
@@ -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)
dreamerv3/embodied/replay/count_based.py
+12
View File
@@ -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']
dreamerv3/embodied/replay/curious_replay.py
+25
View File
@@ -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)
dreamerv3/embodied/replay/generic.py
+162
View File
@@ -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)
dreamerv3/embodied/replay/limiters.py
+108
View File
@@ -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'
dreamerv3/embodied/replay/naive_chunks.py
+82
View File
@@ -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}
dreamerv3/embodied/replay/prioritized_experience_replay.py
+20
View File
@@ -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)
dreamerv3/embodied/replay/replays.py
+26
View File
@@ -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,
)
dreamerv3/embodied/replay/reverb.py
+112
View File
@@ -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
dreamerv3/embodied/replay/saver.py
+63
View File
@@ -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
dreamerv3/embodied/replay/selectors.py
+45
View File
@@ -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
dreamerv3/embodied/run/__init__.py
+6
View File
@@ -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
dreamerv3/embodied/run/eval_only.py
+61
View File
@@ -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()
dreamerv3/embodied/run/parallel.py
+164
View File
@@ -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
dreamerv3/embodied/run/train.py
+120
View File
@@ -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()
dreamerv3/embodied/run/train_eval.py
+124
View File
@@ -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()
dreamerv3/embodied/run/train_holdout.py
+127
View File
@@ -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()
dreamerv3/embodied/run/train_save.py
+130
View File
@@ -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()
dreamerv3/embodied/scripts/install-atari.sh
+17
View File
@@ -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
dreamerv3/embodied/scripts/install-dmlab.sh
Executable
+48
View File
@@ -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
dreamerv3/embodied/scripts/install-minecraft.sh
+12
View File
@@ -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
dreamerv3/embodied/scripts/plot.py
+523
View File
@@ -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()
dreamerv3/embodied/scripts/xvfb_run.sh
+2
View File
@@ -0,0 +1,2 @@
xvfb-run -a -s "-screen 0 1024x768x24 -ac +extension GLX +render -noreset" "$@"
# xvfb-run "$@"
dreamerv3/expl.py
+37
View File
@@ -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()
dreamerv3/jaxagent.py
+240
View File
@@ -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
dreamerv3/jaxutils.py
+476
View File
@@ -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)
dreamerv3/nets.py
+699
View File
@@ -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)
dreamerv3/ninjax.py
+502
View File
@@ -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)}
dreamerv3/plot_crafter.py
+126
View File
@@ -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()
dreamerv3/train.py
+216
View File
@@ -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()
example.py
+53
View File
@@ -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()
requirements.txt
+13
View File
@@ -0,0 +1,13 @@
cloudpickle
crafter
gym==0.26.2
dm_control
jax
jaxlib
numpy
optax
rich
ruamel.yaml
tensorflow-cpu
tensorflow_probability
tensorflow
run_local.py
+37
View File
@@ -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()
setup.py
+22
View File
@@ -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',
],
)