Initial commit

LICENSE

@@ -0,0 +1,21 @@
+MIT License
+
+Copyright (c) Pietro Mazzaglia
+
+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.

README.md

@@ -1,13 +1,177 @@
----
-title: Genrl
-emoji: 💻
-colorFrom: blue
-colorTo: pink
-sdk: gradio
-sdk_version: 4.37.1
-app_file: app.py
-pinned: false
-license: mit
----
-
-Check out the configuration reference at https://huggingface.co/docs/hub/spaces-config-reference
+# GenRL: Multimodal foundation world models for generalist embodied agents
+
+<p align="center">
+    <img src='assets/GenRL_fig1.png' width=90%>
+</p>
+
+<p align="center">
+         <a href="https://mazpie.github.io/genrl">Website</a>&nbsp | <a href="https://huggingface.co/mazpie/genrl_models"> Models 🤗</a>&nbsp | <a href="https://huggingface.co/datasets/mazpie/genrl_datasets"> Datasets 🤗</a>&nbsp | <a href="./demo/"> Gradio demo</a>&nbsp | <a href="./notebooks/"> Notebooks</a>&nbsp
+<br>
+
+## Get started
+
+### Creating the environment
+
+We recommend using `conda` to create the environment
+
+```
+conda create --name genrl python=3.10
+
+conda activate genrl
+
+pip install -r requirements.txt
+```
+
+### Downloading InternVideo2
+
+Download InternVideo 2 [[here]](https://huggingface.co/OpenGVLab/InternVideo2-Stage2_1B-224p-f4/blob/main/InternVideo2-stage2_1b-224p-f4.pt).
+
+Place in the `models` folder.
+
+Note: the file access is restricted, so you'll need an HuggingFace account to request access to the file.
+
+Note: By default, the code expects the model to be placed in the `models` folder. The variable `MODELS_ROOT_PATH` indicating where the model should be place is set in `tools/genrl_utils.py`.
+
+## Data
+
+### Download datasets
+
+The datasets used to pre-trained the models can be downloaded [[here]](https://huggingface.co/datasets/mazpie/genrl_datasets).
+
+The file are `tar.gz` and can be extracted using the `tar` utility on Linux. For example:
+
+```
+tar -zxvf walker_data.tar.gz
+```
+
+### Collecting and pre-processing data
+
+If you don't want to download our datasets, you collect and pre-process the data on your own.
+
+Data can be collected running a DreamerV3 agent on a task, by running:
+
+```
+python3 collect_data.py agent=dreamer task=stickman_walk
+```
+
+or the Plan2Explore agent, by running:
+
+```
+python3 collect_data.py agent=plan2explore conf/defaults=dreamer_v2 task=stickman_walk
+```
+
+A repo for the experiment will be created under the directory `exp_local`, such as: `exp_local/YYYY.MM.DD/HHMMSS_agentname`. The data can then be found in the `buffer` subdirectory.
+
+
+After obtaining the data, it should be processed to obtain the video embeddings for each frame sequence in the episodes. The processing can be done by running:
+
+```
+python3 process_dataset.py dataset_dir=data/stickman_example
+```
+
+where `data/stickman_example` is replaced by the folder of the data you want to process.
+
+## Agents 
+
+### Downloading pre-trained models
+
+If you want to test our work, without having to pre-train the models, you can do this by using our pre-trained models.
+
+Pretrained models can be found [[here]](https://huggingface.co/mazpie/genrl_models)
+
+Here's a snippet to download them easily:
+
+```
+import os
+from huggingface_hub import hf_hub_download
+
+def download_model(model_folder, model_filename):
+    REPO_ID = 'mazpie/genrl_models'
+    filename_list = [model_filename]
+    if not os.path.exists(model_folder):
+        os.makedirs(model_folder)
+    for filename in filename_list:
+        local_file = os.path.join(model_folder, filename)
+        if not os.path.exists(local_file):
+            hf_hub_download(repo_id=REPO_ID, filename=filename, local_dir=model_folder, local_dir_use_symlinks=False)
+
+download_model('models', 'genrl_stickman_500k_2.pt')
+```
+
+Pre-trained models can be used by setting `snapshot_load_dir=...` when running `train.py`.
+
+Note: the pre-trained models are not trained to solve any tasks. They only contain a pre-trained multimodal foundation world model (world model + connector and aligner).
+
+### Training multimodal foundation world models
+
+In order to train a multimodal foundation world model from data, you should run something like:
+
+```
+# Note: frames = update steps
+
+
+python3 train.py task=stickman_walk replay_load_dir=data/stickman_example num_train_frames=500_010 visual_every_frames=25_000 train_world_model=True train_connector=True reset_world_model=True reset_connector=True
+```
+
+### Behavior learning
+
+After pre-training a model, you can train the behavior for a task using:
+
+```
+python3 train.py task=stickman_walk snapshot_load_dir=models/genrl_stickman_500k_2.pt num_train_frames=50_010 batch_size=32 batch_length=32 agent.imag_reward_fn=video_text_reward eval_modality=task_imag 
+```
+
+Data-free RL can be performed by additionaly passing the option:
+
+`train_from_data=False`
+
+The prompts for each task can be found and edited in `tools/genrl_utils.py`. However, you can also pass a custom prompt for a task by passing the option:
+
+`+agent.imag_reward_args.task_prompt=custom_prompt`
+
+## Other utilities
+
+### Gradio demo
+
+There's a gradio demo that can be found at `demo/app.py`.
+
+If launching demo like a standard Python program with:
+
+```
+python3 demo/app.py
+```
+
+it will return a local endpoint (e.g. http://127.0.0.1:7860) where to access a dashboard to play with GenRL.
+
+<p align="center">
+    <img src='assets/dashboard.png' width=75%>
+</p>
+
+### Notebooks
+
+You can find several notebooks to test our code in the `notebooks` directory.
+
+`demo_videoclip` : can be used to test the correct functioning of the InternVideo2 component
+
+`text2video` : utility to generate video reconstructions from text prompts
+
+`video2video` : utility to generate video reconstructions from video prompts
+
+`visualize_dataset_episodes` : utility to generate videos from the episodes in a given dataset
+
+`visualize_env` : used to play with the environment and, for instance, understand how the reward function of each task works
+
+### Stickman environment
+
+We introduced the Stickman environment as a simplified 2D version of the Humanoid environment.
+
+This can be found in the `envs/custom_dmc_tasks` folder. You will find an `.xml` model and a `.py` files containing the tasks.
+
+## Acknowledgments
+
+We would like to thank the authors of the following repositories for their useful code and models:
+
+* [InternVideo2](https://github.com/OpenGVLab/InternVideo)
+* [Franka Kitchen](https://github.com/google-research/relay-policy-learning)
+* [DreamerV3](https://github.com/danijar/dreamerv3)
+* [DreamerV3-torch](https://github.com/NM512/dreamerv3-torch)

agent/dreamer.py

@@ -0,0 +1,462 @@
+import torch.nn as nn
+import torch
+
+import tools.utils as utils
+import agent.dreamer_utils as common
+from collections import OrderedDict
+import numpy as np
+
+from tools.genrl_utils import *
+
+def stop_gradient(x):
+  return x.detach()
+
+Module = nn.Module 
+
+def env_reward(agent, seq):
+  return agent.wm.heads['reward'](seq['feat']).mean
+
+class DreamerAgent(Module):
+
+  def __init__(self, 
+                name, cfg, obs_space, act_spec, **kwargs):
+    super().__init__()
+    self.name = name
+    self.cfg = cfg
+    self.cfg.update(**kwargs)
+    self.obs_space = obs_space
+    self.act_spec = act_spec
+    self._use_amp = (cfg.precision == 16)
+    self.device = cfg.device
+    self.act_dim = act_spec.shape[0]
+    self.wm = WorldModel(cfg, obs_space, self.act_dim,)
+    self.instantiate_acting_behavior()
+
+    self.to(cfg.device)
+    self.requires_grad_(requires_grad=False)
+
+  def instantiate_acting_behavior(self,):
+    self._acting_behavior = ActorCritic(self.cfg, self.act_spec, self.wm.inp_size).to(self.device)
+    
+  def act(self, obs, meta, step, eval_mode, state):
+    if self.cfg.only_random_actions:
+      return np.random.uniform(-1, 1, self.act_dim,).astype(self.act_spec.dtype), (None, None)
+    obs = {k : torch.as_tensor(np.copy(v), device=self.device).unsqueeze(0) for k, v in obs.items()}
+    if state is None:
+      latent = self.wm.rssm.initial(len(obs['reward']))
+      action = torch.zeros((len(obs['reward']),) + self.act_spec.shape, device=self.device)
+    else:
+      latent, action = state
+    embed = self.wm.encoder(self.wm.preprocess(obs))
+    should_sample = (not eval_mode) or (not self.cfg.eval_state_mean)
+    latent, _ = self.wm.rssm.obs_step(latent, action, embed, obs['is_first'], should_sample)
+    feat = self.wm.rssm.get_feat(latent)
+    if eval_mode:
+      actor = self._acting_behavior.actor(feat)
+      try:
+        action = actor.mean 
+      except:
+        action = actor._mean
+    else:
+      actor = self._acting_behavior.actor(feat)
+      action = actor.sample()
+    new_state = (latent, action)
+    return action.cpu().numpy()[0], new_state
+
+  def update_wm(self, data, step):
+    metrics = {}
+    state, outputs, mets = self.wm.update(data, state=None)
+    outputs['is_terminal'] = data['is_terminal']
+    metrics.update(mets)
+    return state, outputs, metrics
+
+  def update_acting_behavior(self, state=None, outputs=None, metrics={}, data=None, reward_fn=None):
+    if self.cfg.only_random_actions:
+      return {}, metrics
+    if outputs is not None:
+      post = outputs['post']
+      is_terminal = outputs['is_terminal']
+    else:
+      data = self.wm.preprocess(data)
+      embed = self.wm.encoder(data)
+      post, _ = self.wm.rssm.observe(
+          embed, data['action'], data['is_first'])
+      is_terminal = data['is_terminal']
+    #
+    start = {k: stop_gradient(v) for k,v in post.items()}
+    if reward_fn is None:
+      acting_reward_fn = lambda seq: globals()[self.cfg.acting_reward_fn](self, seq) #.mode()
+    else:
+      acting_reward_fn = lambda seq: reward_fn(self, seq) #.mode()
+    metrics.update(self._acting_behavior.update(self.wm, start, is_terminal, acting_reward_fn))
+    return start, metrics
+
+  def update(self, data, step):
+    state, outputs, metrics = self.update_wm(data, step)
+    start, metrics = self.update_acting_behavior(state, outputs, metrics, data)
+    return state, metrics
+
+  def report(self, data):
+    report = {}
+    data = self.wm.preprocess(data)
+    for key in self.wm.heads['decoder'].cnn_keys:
+      name = key.replace('/', '_')
+      report[f'openl_{name}'] = self.wm.video_pred(data, key)
+    for fn in getattr(self.cfg, 'additional_report_fns', []):
+      call_fn = globals()[fn]
+      additional_report = call_fn(self, data)
+      report.update(additional_report)
+    return report
+
+  def get_meta_specs(self):
+    return tuple()
+
+  def init_meta(self):
+    return OrderedDict()
+
+  def update_meta(self, meta, global_step, time_step, finetune=False):
+    return meta
+
+class WorldModel(Module):
+  def __init__(self, config, obs_space, act_dim,):
+    super().__init__()
+    shapes = {k: tuple(v.shape) for k, v in obs_space.items()}
+    self.shapes = shapes
+    self.cfg = config
+    self.device = config.device
+    self.encoder = common.Encoder(shapes, **config.encoder)
+    # Computing embed dim
+    with torch.no_grad():
+      zeros = {k: torch.zeros( (1,) + v) for k, v in shapes.items()}
+      outs = self.encoder(zeros)
+      embed_dim = outs.shape[1]
+    self.embed_dim = embed_dim
+    self.rssm = common.EnsembleRSSM(**config.rssm, action_dim=act_dim, embed_dim=embed_dim, device=self.device,)
+    self.heads = {}
+    self._use_amp = (config.precision == 16)
+    self.inp_size = self.rssm.get_feat_size()
+    self.decoder_input_fn = getattr(self.rssm, f'get_{config.decoder_inputs}')
+    self.decoder_input_size = getattr(self.rssm, f'get_{config.decoder_inputs}_size')()
+    self.heads['decoder'] = common.Decoder(shapes, **config.decoder, embed_dim=self.decoder_input_size, image_dist=config.image_dist)
+    self.heads['reward'] = common.MLP(self.inp_size, (1,), **config.reward_head)
+    # zero init
+    with torch.no_grad():
+      for p in self.heads['reward']._out.parameters():
+        p.data = p.data * 0
+    #
+    if config.pred_discount:
+      self.heads['discount'] = common.MLP(self.inp_size, (1,), **config.discount_head)
+    for name in config.grad_heads:
+      assert name in self.heads, name
+    self.grad_heads = config.grad_heads
+    self.heads = nn.ModuleDict(self.heads)
+    self.model_opt = common.Optimizer('model', self.parameters(), **config.model_opt, use_amp=self._use_amp)
+    self.e2e_update_fns = {}
+    self.detached_update_fns = {}
+    self.eval()
+
+  def add_module_to_update(self, name, module, update_fn, detached=False):
+    self.add_module(name, module)
+    if detached:
+      self.detached_update_fns[name] = update_fn
+    else:
+      self.e2e_update_fns[name] = update_fn
+    self.model_opt = common.Optimizer('model', self.parameters(), **self.cfg.model_opt, use_amp=self._use_amp)
+
+  def update(self, data, state=None):
+    self.train()
+    with common.RequiresGrad(self):
+      with torch.cuda.amp.autocast(enabled=self._use_amp):
+        if getattr(self.cfg, "freeze_decoder", False):
+          self.heads['decoder'].requires_grad_(False)
+        if getattr(self.cfg, "freeze_post", False) or getattr(self.cfg, "freeze_model", False):
+          self.heads['decoder'].requires_grad_(False)
+          self.encoder.requires_grad_(False)
+          # Updating only prior
+          self.grad_heads = []
+          self.rssm.requires_grad_(False)
+          if not getattr(self.cfg, "freeze_model", False):
+            self.rssm._ensemble_img_out.requires_grad_(True)
+            self.rssm._ensemble_img_dist.requires_grad_(True)
+        model_loss, state, outputs, metrics = self.loss(data, state)
+        model_loss, metrics = self.update_additional_e2e_modules(data, outputs, model_loss, metrics)
+      metrics.update(self.model_opt(model_loss, self.parameters())) 
+    if len(self.detached_update_fns) > 0:
+      detached_loss, metrics = self.update_additional_detached_modules(data, outputs, metrics)
+    self.eval()
+    return state, outputs, metrics
+
+  def update_additional_detached_modules(self, data, outputs, metrics):
+    # additional detached losses
+    detached_loss = 0
+    for k in self.detached_update_fns:
+      detached_module = getattr(self, k)
+      with common.RequiresGrad(detached_module):
+        with torch.cuda.amp.autocast(enabled=self._use_amp):
+          add_loss, add_metrics = self.detached_update_fns[k](self, k, data, outputs, metrics)
+          metrics.update(add_metrics)
+          opt_metrics = self.model_opt(add_loss, detached_module.parameters())
+          metrics.update({ f'{k}_{m}' : opt_metrics[m] for m in opt_metrics})
+    return detached_loss, metrics
+
+  def update_additional_e2e_modules(self, data, outputs, model_loss, metrics):
+    # additional e2e losses
+    for k in self.e2e_update_fns:
+      add_loss, add_metrics = self.e2e_update_fns[k](self, k, data, outputs, metrics)
+      model_loss += add_loss
+      metrics.update(add_metrics)
+    return model_loss, metrics
+
+  def observe_data(self, data, state=None):
+    data = self.preprocess(data)
+    embed = self.encoder(data)
+    post, prior = self.rssm.observe(
+        embed, data['action'], data['is_first'], state)
+    kl_loss, kl_value = self.rssm.kl_loss(post, prior, **self.cfg.kl)
+    outs = dict(embed=embed, post=post, prior=prior, is_terminal=data['is_terminal'])
+    return outs, { 'model_kl' : kl_value.mean() }
+
+  def loss(self, data, state=None):
+    data = self.preprocess(data)
+    embed = self.encoder(data)
+    post, prior = self.rssm.observe(
+        embed, data['action'], data['is_first'], state)
+    kl_loss, kl_value = self.rssm.kl_loss(post, prior, **self.cfg.kl)
+    assert len(kl_loss.shape) == 0 or (len(kl_loss.shape) == 1 and kl_loss.shape[0] == 1), kl_loss.shape
+    likes = {}
+    losses = {'kl': kl_loss}
+    feat = self.rssm.get_feat(post)
+    for name, head in self.heads.items():
+      grad_head = (name in self.grad_heads)
+      if name == 'decoder':
+        inp = self.decoder_input_fn(post)
+      else:
+        inp = feat
+      inp = inp if grad_head else stop_gradient(inp)
+      out = head(inp)
+      dists = out if isinstance(out, dict) else {name: out}
+      for key, dist in dists.items():
+        like = dist.log_prob(data[key]) 
+        likes[key] = like
+        losses[key] = -like.mean()
+    model_loss = sum(
+        self.cfg.loss_scales.get(k, 1.0) * v for k, v in losses.items())
+    outs = dict(
+        embed=embed, feat=feat, post=post,
+        prior=prior, likes=likes, kl=kl_value)
+    metrics = {f'{name}_loss': value for name, value in losses.items()}
+    metrics['model_kl'] = kl_value.mean()
+    metrics['prior_ent'] = self.rssm.get_dist(prior).entropy().mean()
+    metrics['post_ent'] = self.rssm.get_dist(post).entropy().mean()
+    last_state = {k: v[:, -1] for k, v in post.items()}
+    return model_loss, last_state, outs, metrics
+
+  def imagine(self, policy, start, is_terminal, horizon, task_cond=None, eval_policy=False):
+    flatten = lambda x: x.reshape([-1] + list(x.shape[2:]))
+    start = {k: flatten(v) for k, v in start.items()}
+    start['feat'] = self.rssm.get_feat(start)
+    inp = start['feat'] if task_cond is None else torch.cat([start['feat'], task_cond], dim=-1)
+    policy_dist = policy(inp)
+    start['action'] = torch.zeros_like(policy_dist.sample(), device=self.device) #.mode())
+    seq = {k: [v] for k, v in start.items()}
+    if task_cond is not None: seq['task'] = [task_cond]
+    for _ in range(horizon):
+      inp = seq['feat'][-1] if task_cond is None else torch.cat([seq['feat'][-1], task_cond], dim=-1)
+      policy_dist = policy(stop_gradient(inp))
+      action = policy_dist.sample() if not eval_policy else policy_dist.mean
+      state = self.rssm.img_step({k: v[-1] for k, v in seq.items()}, action)
+      feat = self.rssm.get_feat(state)
+      for key, value in {**state, 'action': action, 'feat': feat}.items():
+        seq[key].append(value)
+      if task_cond is not None: seq['task'].append(task_cond)
+    # shape will be (T, B, *DIMS)
+    seq = {k: torch.stack(v, 0) for k, v in seq.items()}
+    if 'discount' in self.heads:
+      disc = self.heads['discount'](seq['feat']).mean()
+      if is_terminal is not None:
+        # Override discount prediction for the first step with the true
+        # discount factor from the replay buffer.
+        true_first = 1.0 - flatten(is_terminal) 
+        disc = torch.cat([true_first[None], disc[1:]], 0)
+    else:
+      disc = torch.ones(list(seq['feat'].shape[:-1]) + [1], device=self.device)
+    seq['discount'] = disc * self.cfg.discount
+    # Shift discount factors because they imply whether the following state
+    # will be valid, not whether the current state is valid.
+    seq['weight'] = torch.cumprod(torch.cat([torch.ones_like(disc[:1], device=self.device), disc[:-1]], 0), 0)
+    return seq
+
+  def preprocess(self, obs):
+    obs = obs.copy()
+    for key, value in obs.items():
+      if key.startswith('log_'):
+        continue
+      if value.dtype in [np.uint8, torch.uint8]:
+        value = value / 255.0 - 0.5 
+      obs[key] = value
+    obs['reward'] = {
+        'identity': nn.Identity(),
+        'sign': torch.sign,
+        'tanh': torch.tanh,
+    }[self.cfg.clip_rewards](obs['reward'])
+    obs['discount'] = (1.0 - obs['is_terminal'].float())
+    if len(obs['discount'].shape) < len(obs['reward'].shape):
+      obs['discount'] = obs['discount'].unsqueeze(-1)
+    return obs
+
+  def video_pred(self, data, key, nvid=8):
+    decoder = self.heads['decoder'] # B, T, C, H, W
+    truth = data[key][:nvid] + 0.5
+    embed = self.encoder(data)
+    states, _ = self.rssm.observe(
+        embed[:nvid, :5], data['action'][:nvid, :5], data['is_first'][:nvid, :5])
+    recon = decoder(self.decoder_input_fn(states))[key].mean[:nvid] # mode
+    init = {k: v[:, -1] for k, v in states.items()}
+    prior = self.rssm.imagine(data['action'][:nvid, 5:], init)
+    prior_recon = decoder(self.decoder_input_fn(prior))[key].mean # mode
+    model = torch.clip(torch.cat([recon[:, :5] + 0.5, prior_recon + 0.5], 1), 0, 1)
+    error = (model - truth + 1) / 2
+    video = torch.cat([truth, model, error], 3)
+    B, T, C, H, W = video.shape
+    return video 
+
+class ActorCritic(Module):
+  def __init__(self, config, act_spec, feat_size, name=''):
+    super().__init__()
+    self.name = name
+    self.cfg = config
+    self.act_spec = act_spec
+    self._use_amp = (config.precision == 16)
+    self.device = config.device
+    
+    if getattr(self.cfg, 'discrete_actions', False):
+      self.cfg.actor.dist = 'onehot'
+
+    self.actor_grad = getattr(self.cfg, f'{self.name}_actor_grad'.strip('_'))
+    
+    inp_size = feat_size
+    self.actor = common.MLP(inp_size, act_spec.shape[0], **self.cfg.actor)
+    self.critic = common.MLP(inp_size, (1,), **self.cfg.critic)
+    if self.cfg.slow_target:
+      self._target_critic = common.MLP(inp_size, (1,), **self.cfg.critic)
+      self._updates = 0 # tf.Variable(0, tf.int64)
+    else:
+      self._target_critic = self.critic
+    self.actor_opt = common.Optimizer('actor', self.actor.parameters(), **self.cfg.actor_opt, use_amp=self._use_amp)
+    self.critic_opt = common.Optimizer('critic', self.critic.parameters(), **self.cfg.critic_opt, use_amp=self._use_amp)
+    
+    if self.cfg.reward_ema:
+        # register ema_vals to nn.Module for enabling torch.save and torch.load
+        self.register_buffer("ema_vals", torch.zeros((2,)).to(self.device))
+        self.reward_ema = common.RewardEMA(device=self.device)
+        self.rewnorm = common.StreamNorm(momentum=1, scale=1.0, device=self.device)
+    else:
+        self.rewnorm = common.StreamNorm(**self.cfg.reward_norm, device=self.device)
+
+    # zero init
+    with torch.no_grad():
+      for p in self.critic._out.parameters():
+        p.data = p.data * 0
+    # hard copy critic initial params
+    for s, d in zip(self.critic.parameters(), self._target_critic.parameters()):
+      d.data = s.data
+    #
+
+
+  def update(self, world_model, start, is_terminal, reward_fn):
+    metrics = {}
+    hor = self.cfg.imag_horizon
+    # The weights are is_terminal flags for the imagination start states.
+    # Technically, they should multiply the losses from the second trajectory
+    # step onwards, which is the first imagined step. However, we are not
+    # training the action that led into the first step anyway, so we can use
+    # them to scale the whole sequence.
+    with common.RequiresGrad(self.actor):
+      with torch.cuda.amp.autocast(enabled=self._use_amp):
+        seq = world_model.imagine(self.actor, start, is_terminal, hor)
+        reward = reward_fn(seq)
+        seq['reward'], mets1 = self.rewnorm(reward)
+        mets1 = {f'reward_{k}': v for k, v in mets1.items()}
+        target, mets2, baseline = self.target(seq)
+        actor_loss, mets3 = self.actor_loss(seq, target, baseline)
+      metrics.update(self.actor_opt(actor_loss, self.actor.parameters()))
+    with common.RequiresGrad(self.critic):
+      with torch.cuda.amp.autocast(enabled=self._use_amp):
+        seq = {k: stop_gradient(v) for k,v in seq.items()}
+        critic_loss, mets4 = self.critic_loss(seq, target)
+      metrics.update(self.critic_opt(critic_loss, self.critic.parameters()))
+    metrics.update(**mets1, **mets2, **mets3, **mets4)
+    self.update_slow_target()  # Variables exist after first forward pass.
+    return { f'{self.name}_{k}'.strip('_') : v for k,v in metrics.items() }
+
+  def actor_loss(self, seq, target, baseline): #, step):
+    # Two state-actions are lost at the end of the trajectory, one for the boostrap
+    # value prediction and one because the corresponding action does not lead
+    # anywhere anymore. One target is lost at the start of the trajectory
+    # because the initial state comes from the replay buffer.
+    policy = self.actor(stop_gradient(seq['feat'][:-2])) # actions are the ones in [1:-1]
+
+    metrics = {}
+    if self.cfg.reward_ema:
+      offset, scale = self.reward_ema(target, self.ema_vals)
+      normed_target = (target - offset) / scale
+      normed_baseline = (baseline - offset) / scale
+      # adv = normed_target - normed_baseline
+      metrics['normed_target_mean'] = normed_target.mean()
+      metrics['normed_target_std'] = normed_target.std()
+      metrics["reward_ema_005"] = self.ema_vals[0]
+      metrics["reward_ema_095"] = self.ema_vals[1]
+    else:
+      normed_target = target
+      normed_baseline = baseline
+    
+    if self.actor_grad == 'dynamics':
+      objective = normed_target[1:]
+    elif self.actor_grad == 'reinforce':
+      advantage = normed_target[1:] - normed_baseline[1:]
+      objective = policy.log_prob(stop_gradient(seq['action'][1:-1]))[:,:,None] * advantage
+    else:
+      raise NotImplementedError(self.actor_grad)
+    
+    ent = policy.entropy()[:,:,None]
+    ent_scale = self.cfg.actor_ent
+    objective += ent_scale * ent
+    metrics['actor_ent'] = ent.mean()
+    metrics['actor_ent_scale'] = ent_scale
+    
+    weight = stop_gradient(seq['weight'])
+    actor_loss = -(weight[:-2] * objective).mean() 
+    return actor_loss, metrics
+
+  def critic_loss(self, seq, target):
+    feat = seq['feat'][:-1]
+    target = stop_gradient(target)
+    weight = stop_gradient(seq['weight'])
+    dist = self.critic(feat)
+    critic_loss = -(dist.log_prob(target)[:,:,None] * weight[:-1]).mean()
+    metrics = {'critic': dist.mean.mean() } 
+    return critic_loss, metrics
+
+  def target(self, seq):
+    reward = seq['reward'] 
+    disc = seq['discount'] 
+    value = self._target_critic(seq['feat']).mean 
+    # Skipping last time step because it is used for bootstrapping.
+    target = common.lambda_return(
+        reward[:-1], value[:-1], disc[:-1],
+        bootstrap=value[-1],
+        lambda_=self.cfg.discount_lambda,
+        axis=0)
+    metrics = {}
+    metrics['critic_slow'] = value.mean()
+    metrics['critic_target'] = target.mean()
+    return target, metrics, value[:-1]
+
+  def update_slow_target(self):
+    if self.cfg.slow_target:
+      if self._updates % self.cfg.slow_target_update == 0:
+        mix = 1.0 if self._updates == 0 else float(
+            self.cfg.slow_target_fraction)
+        for s, d in zip(self.critic.parameters(), self._target_critic.parameters()):
+          d.data = mix * s.data + (1 - mix) * d.data
+      self._updates += 1 

agent/dreamer.yaml

@@ -0,0 +1,9 @@
+# @package agent
+_target_: agent.dreamer.DreamerAgent
+name: dreamer
+cfg: ???
+obs_space: ???
+act_spec: ???
+grad_heads: [decoder, reward]
+reward_norm: {momentum: 1.0, scale: 1.0, eps: 1e-8}
+actor_ent: 3e-4

agent/dreamer_utils.py

@@ -0,0 +1,1040 @@
+import re
+
+import numpy as np
+
+import tools.utils as utils
+import torch.nn as nn
+import torch
+import torch.distributions as D
+import torch.nn.functional as F
+
+Module = nn.Module 
+
+def symlog(x):
+  return torch.sign(x) * torch.log(torch.abs(x) + 1.0)
+
+def symexp(x):
+  return torch.sign(x) * (torch.exp(torch.abs(x)) - 1.0)
+
+def signed_hyperbolic(x: torch.Tensor, eps: float = 1e-3) -> torch.Tensor:
+    """Signed hyperbolic transform, inverse of signed_parabolic."""
+    return torch.sign(x) * (torch.sqrt(torch.abs(x) + 1) - 1) + eps * x
+
+def signed_parabolic(x: torch.Tensor, eps: float = 1e-3) -> torch.Tensor:
+    """Signed parabolic transform, inverse of signed_hyperbolic."""
+    z = torch.sqrt(1 + 4 * eps * (eps + 1 + torch.abs(x))) / 2 / eps - 1 / 2 / eps
+    return torch.sign(x) * (torch.square(z) - 1)
+
+class SampleDist:
+    def __init__(self, dist: D.Distribution, samples=100):
+        self._dist = dist
+        self._samples = samples
+
+    @property
+    def name(self):
+        return 'SampleDist'
+
+    def __getattr__(self, name):
+        return getattr(self._dist, name)
+
+    @property
+    def mean(self):
+        sample = self._dist.rsample((self._samples,))
+        return torch.mean(sample, 0)
+
+    def mode(self):
+        dist = self._dist.expand((self._samples, *self._dist.batch_shape))
+        sample = dist.rsample()
+        logprob = dist.log_prob(sample)
+        batch_size = sample.size(1)
+        feature_size = sample.size(2)
+        indices = torch.argmax(logprob, dim=0).reshape(1, batch_size, 1).expand(1, batch_size, feature_size)
+        return torch.gather(sample, 0, indices).squeeze(0)
+
+    def entropy(self):
+        sample = self._dist.rsample((self._samples,))
+        logprob = self._dist.log_prob(sample)
+        return -torch.mean(logprob, 0)
+
+    def sample(self):
+        return self._dist.rsample()
+
+class MSEDist:
+    def __init__(self, mode, agg="sum"):
+        self._mode = mode
+        self._agg = agg
+
+    @property
+    def mean(self):
+        return self._mode
+
+    def mode(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(list(range(len(distance.shape)))[2:])
+        elif self._agg == "sum":
+            loss = distance.sum(list(range(len(distance.shape)))[2:])
+        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 = torch.where(distance < self._tol, torch.tensor([0.], dtype=distance.dtype, device=distance.device), distance)
+    elif self._dist == 'abs':
+      distance = torch.abs(self._mode - symlog(value))
+      distance = torch.where(distance < self._tol, torch.tensor([0.], dtype=distance.dtype, device=distance.device), 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 TwoHotDist:
+    def __init__(
+        self,
+        logits,
+        low=-20.0,
+        high=20.0,
+        transfwd=symlog,
+        transbwd=symexp,
+    ):
+        assert logits.shape[-1] == 255
+        self.logits = logits
+        self.probs = torch.softmax(logits, -1)
+        self.buckets = torch.linspace(low, high, steps=255).to(logits.device)
+        self.width = (self.buckets[-1] - self.buckets[0]) / 255
+        self.transfwd = transfwd
+        self.transbwd = transbwd
+
+    @property
+    def mean(self):
+        _mean = self.probs * self.buckets
+        return self.transbwd(torch.sum(_mean, dim=-1, keepdim=True))
+
+    @property
+    def mode(self):
+        return self.mean
+
+    # Inside OneHotCategorical, log_prob is calculated using only max element in targets
+    def log_prob(self, x):
+        x = self.transfwd(x)
+        # x(time, batch, 1)
+        below = torch.sum((self.buckets <= x[..., None]).to(torch.int32), dim=-1) - 1
+        above = len(self.buckets) - torch.sum(
+            (self.buckets > x[..., None]).to(torch.int32), dim=-1
+        )
+        # this is implemented using clip at the original repo as the gradients are not backpropagated for the out of limits.
+        below = torch.clip(below, 0, len(self.buckets) - 1)
+        above = torch.clip(above, 0, len(self.buckets) - 1)
+        equal = below == above
+
+        dist_to_below = torch.where(equal, 1, torch.abs(self.buckets[below] - x))
+        dist_to_above = torch.where(equal, 1, torch.abs(self.buckets[above] - x))
+        total = dist_to_below + dist_to_above
+        weight_below = dist_to_above / total
+        weight_above = dist_to_below / total
+        target = (
+            F.one_hot(below, num_classes=len(self.buckets)) * weight_below[..., None]
+            + F.one_hot(above, num_classes=len(self.buckets)) * weight_above[..., None]
+        )
+        log_pred = self.logits - torch.logsumexp(self.logits, -1, keepdim=True)
+        target = target.squeeze(-2)
+
+        return (target * log_pred).sum(-1)
+
+    def log_prob_target(self, target):
+        log_pred = super().logits - torch.logsumexp(super().logits, -1, keepdim=True)
+        return (target * log_pred).sum(-1)
+
+class OneHotDist(D.OneHotCategorical):
+
+  def __init__(self, logits=None, probs=None, unif_mix=0.99):
+    super().__init__(logits=logits, probs=probs)
+    probs = super().probs
+    probs = unif_mix * probs + (1 - unif_mix) * torch.ones_like(probs, device=probs.device) / probs.shape[-1]
+    super().__init__(probs=probs)
+
+  def mode(self):
+    _mode = F.one_hot(torch.argmax(super().logits, axis=-1), super().logits.shape[-1])
+    return _mode.detach() + super().logits - super().logits.detach()
+
+  def sample(self, sample_shape=(), seed=None):
+    if seed is not None:
+      raise ValueError('need to check')
+    sample = super().sample(sample_shape)
+    probs = super().probs
+    while len(probs.shape) < len(sample.shape):
+      probs = probs[None]
+    sample += probs - probs.detach() # ST-gradients
+    return sample
+
+class BernoulliDist(D.Bernoulli):
+  def __init__(self, logits=None, probs=None):
+    super().__init__(logits=logits, probs=probs)
+
+  def sample(self, sample_shape=(), seed=None):
+    if seed is not None:
+      raise ValueError('need to check')
+    sample = super().sample(sample_shape)
+    probs = super().probs
+    while len(probs.shape) < len(sample.shape):
+      probs = probs[None]
+    sample += probs - probs.detach() # ST-gradients
+    return sample
+
+def static_scan_for_lambda_return(fn, inputs, start):
+  last = start
+  indices = range(inputs[0].shape[0])
+  indices = reversed(indices)
+  flag = True
+  for index in indices:
+    inp = lambda x: (_input[x].unsqueeze(0) for _input in inputs)
+    last = fn(last, *inp(index))
+    if flag:
+      outputs = last
+      flag = False
+    else:
+      outputs = torch.cat([last, outputs], dim=0) 
+  return outputs
+
+def lambda_return(
+    reward, value, pcont, bootstrap, lambda_, axis):
+  # Setting lambda=1 gives a discounted Monte Carlo return.
+  # Setting lambda=0 gives a fixed 1-step return.
+  #assert reward.shape.ndims == value.shape.ndims, (reward.shape, value.shape)
+  assert len(reward.shape) == len(value.shape), (reward.shape, value.shape)
+  if isinstance(pcont, (int, float)):
+    pcont = pcont * torch.ones_like(reward, device=reward.device)
+  dims = list(range(len(reward.shape)))
+  dims = [axis] + dims[1:axis] + [0] + dims[axis + 1:]
+  if axis != 0:
+    reward = reward.permute(dims)
+    value = value.permute(dims)
+    pcont = pcont.permute(dims)
+  if bootstrap is None:
+    bootstrap = torch.zeros_like(value[-1], device=reward.device)
+  if len(bootstrap.shape) < len(value.shape):
+    bootstrap = bootstrap[None]
+  next_values = torch.cat([value[1:], bootstrap], 0)
+  inputs = reward + pcont * next_values * (1 - lambda_)
+  returns = static_scan_for_lambda_return(
+      lambda agg, cur0, cur1: cur0 + cur1 * lambda_ * agg,
+      (inputs, pcont), bootstrap)
+  if axis != 0:
+    returns = returns.permute(dims)
+  return returns
+
+def static_scan(fn, inputs, start, reverse=False, unpack=False):
+  last = start
+  indices = range(inputs[0].shape[0])
+  flag = True
+  for index in indices:
+    inp = lambda x: (_input[x] for _input in inputs)
+    if unpack:
+      last = fn(last, *[inp[index] for inp in inputs]) 
+    else:
+      last = fn(last, inp(index)) 
+    if flag:
+      if type(last) == type({}):
+        outputs = {key: [value] for key, value in last.items()}
+      else:
+        outputs = []
+        for _last in last:
+          if type(_last) == type({}):
+            outputs.append({key: [value] for key, value in _last.items()})
+          else:
+            outputs.append([_last])
+      flag = False
+    else:
+      if type(last) == type({}):
+        for key in last.keys():
+          outputs[key].append(last[key]) 
+      else:
+        for j in range(len(outputs)):
+          if type(last[j]) == type({}):
+            for key in last[j].keys():
+              outputs[j][key].append(last[j][key]) 
+          else:
+            outputs[j].append(last[j]) 
+  # Stack everything at the end
+  if type(last) == type({}):
+    for key in last.keys():
+      outputs[key] = torch.stack(outputs[key], dim=0)
+  else:
+    for j in range(len(outputs)):
+      if type(last[j]) == type({}):
+        for key in last[j].keys():
+          outputs[j][key] = torch.stack(outputs[j][key], dim=0)
+      else:
+        outputs[j] = torch.stack(outputs[j], dim=0)
+  if type(last) == type({}):
+    outputs = [outputs]
+  return outputs
+
+class EnsembleRSSM(Module):
+
+  def __init__(
+      self, ensemble=5, stoch=30, deter=200, hidden=200, discrete=False,
+      act='SiLU', norm='none', std_act='softplus', min_std=0.1, action_dim=None, embed_dim=1536, device='cuda', 
+      single_obs_posterior=False, cell_input='stoch', cell_type='gru',):
+    super().__init__()
+    assert action_dim is not None 
+    self.device = device
+    self._embed_dim = embed_dim
+    self._action_dim = action_dim
+    self._ensemble = ensemble
+    self._stoch = stoch
+    self._deter = deter
+    self._hidden = hidden
+    self._discrete = discrete
+    self._act = get_act(act)
+    self._norm = norm
+    self._std_act = std_act
+    self._min_std = min_std
+    self._cell_type = cell_type
+    self.cell_input = cell_input
+    if cell_type == 'gru':
+      self._cell = GRUCell(self._hidden, self._deter, norm=True, device=self.device)
+    else:
+      raise NotImplementedError(f"{cell_type} not implemented")
+    self.single_obs_posterior = single_obs_posterior
+
+    if discrete:
+      self._ensemble_img_dist = nn.ModuleList([ nn.Linear(hidden, stoch*discrete) for _ in range(ensemble)])
+      self._obs_dist = nn.Linear(hidden, stoch*discrete)
+    else:
+      self._ensemble_img_dist = nn.ModuleList([ nn.Linear(hidden, 2*stoch) for _ in range(ensemble)])
+      self._obs_dist = nn.Linear(hidden, 2*stoch)
+
+    # Layer that projects (stoch, input) to cell_state space
+    cell_state_input_size = getattr(self, f'get_{self.cell_input}_size')()
+    self._img_in = nn.Sequential(nn.Linear(cell_state_input_size + action_dim, hidden, bias=norm != 'none'), NormLayer(norm, hidden))
+    # Layer that project deter -> hidden [before projecting hidden -> stoch]
+    self._ensemble_img_out = nn.ModuleList([ nn.Sequential(nn.Linear(self.get_deter_size(), hidden, bias=norm != 'none'), NormLayer(norm, hidden)) for _ in range(ensemble)])
+
+    if self.single_obs_posterior:
+      self._obs_out = nn.Sequential(nn.Linear(embed_dim, hidden, bias=norm != 'none'), NormLayer(norm, hidden))
+    else:
+      self._obs_out = nn.Sequential(nn.Linear(deter + embed_dim, hidden, bias=norm != 'none'), NormLayer(norm, hidden))
+
+  def initial(self, batch_size):
+    if self._discrete:
+      state = dict(
+          logit=torch.zeros([batch_size, self._stoch, self._discrete], device=self.device), 
+          stoch=torch.zeros([batch_size, self._stoch, self._discrete], device=self.device), 
+          deter=self._cell.get_initial_state(None, batch_size)) 
+    else:
+      state = dict(
+          mean=torch.zeros([batch_size, self._stoch], device=self.device), 
+          std=torch.zeros([batch_size, self._stoch], device=self.device),
+          stoch=torch.zeros([batch_size, self._stoch], device=self.device), 
+          deter=self._cell.get_initial_state(None, batch_size)) 
+    return state
+
+  def observe(self, embed, action, is_first, state=None):
+    swap = lambda x: x.permute([1, 0] + list(range(2, len(x.shape))))
+    if state is None: state = self.initial(action.shape[0])
+
+    post, prior = static_scan(
+        lambda prev, inputs: self.obs_step(prev[0], *inputs),
+        (swap(action), swap(embed), swap(is_first)), (state, state))
+    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, sample=True):
+    swap = lambda x: x.permute([1, 0] + list(range(2, len(x.shape))))
+    if state is None:
+      state = self.initial(action.shape[0])
+    assert isinstance(state, dict), state
+    action = swap(action)
+    prior = static_scan(self.img_step, [action, float(sample) + torch.zeros(action.shape[0])], state, unpack=True)[0] 
+    prior = {k: swap(v) for k, v in prior.items()}
+    return prior
+
+  def get_stoch_size(self,):
+    if self._discrete:
+      return self._stoch * self._discrete
+    else:
+      return self._stoch
+
+  def get_deter_size(self,):
+      return self._cell.state_size
+
+  def get_feat_size(self,):
+    return self.get_deter_size() + self.get_stoch_size()
+
+  def get_stoch(self, state):
+    stoch = state['stoch'] 
+    if self._discrete:
+      shape = list(stoch.shape[:-2]) + [self._stoch * self._discrete]
+      stoch = stoch.reshape(shape)
+    return stoch
+
+  def get_deter(self, state):
+      return state['deter']
+
+  def get_feat(self, state):
+    deter = self.get_deter(state)
+    stoch = self.get_stoch(state)
+    return torch.cat([stoch, deter], -1)
+
+  def get_dist(self, state, ensemble=False):
+    if ensemble:
+      state = self._suff_stats_ensemble(state['deter'])
+    if self._discrete:
+      logit = state['logit']
+      dist = D.Independent(OneHotDist(logit.float()), 1)
+    else:
+      mean, std = state['mean'], state['std']
+      dist = D.Independent(D.Normal(mean, std), 1)
+      dist.sample = dist.rsample
+    return dist
+
+  def get_unif_dist(self, state):
+    if self._discrete:
+      logit = state['logit']
+      dist = D.Independent(OneHotDist(torch.ones_like(logit, device=logit.device)), 1)
+    else:
+      mean, std = state['mean'], state['std']
+      dist = D.Independent(D.Normal(torch.zeros_like(mean, device=mean.device), torch.ones_like(std, device=std.device)), 1)
+      dist.sample = dist.rsample
+    return dist
+
+  def obs_step(self, prev_state, prev_action, embed, is_first, should_sample=True):
+    if is_first.any():
+      prev_state = { k: torch.einsum('b,b...->b...', 1.0 - is_first.float(), x) for k, x in prev_state.items() }
+      prev_action = torch.einsum('b,b...->b...', 1.0 - is_first.float(), prev_action)
+    #
+    prior = self.img_step(prev_state, prev_action, should_sample)
+    stoch, stats = self.get_post_stoch(embed, prior, should_sample)
+    post = {'stoch': stoch, 'deter': prior['deter'], **stats}
+    return post, prior
+
+  def get_post_stoch(self, embed, prior, should_sample=True):
+    if self.single_obs_posterior:
+      x = embed
+    else:
+      x = torch.cat([prior['deter'], embed], -1)
+    x = self._obs_out(x)
+    x = self._act(x)
+  
+    bs = list(x.shape[:-1])
+    x = x.reshape([-1, x.shape[-1]])
+    stats = self._suff_stats_layer('_obs_dist', x)
+    stats = { k: v.reshape( bs + list(v.shape[1:])) for k, v in stats.items()}
+    
+    dist = self.get_dist(stats)
+    stoch = dist.sample() if should_sample else dist.mode() 
+    return stoch, stats
+
+  def img_step(self, prev_state, prev_action, sample=True,):
+    prev_state_input = getattr(self, f'get_{self.cell_input}')(prev_state)
+    x = torch.cat([prev_state_input, prev_action], -1)
+    x = self._img_in(x)
+    x = self._act(x)
+    deter = prev_state['deter']
+    if self._cell_type == 'gru':
+      x, deter = self._cell(x, [deter])
+      temp_state = {'deter' : deter[0] }
+    else:
+      raise NotImplementedError(f"no {self._cell_type} cell method")
+    deter = deter[0]  # It's wrapped in a list.
+    stoch, stats = self.get_stoch_stats_from_deter_state(temp_state, sample)
+    prior = {'stoch': stoch, 'deter': deter, **stats}
+    return prior
+
+  def get_stoch_stats_from_deter_state(self, temp_state, sample=True):
+    stats = self._suff_stats_ensemble(self.get_deter(temp_state))
+    index = torch.randint(0, self._ensemble, ()) 
+    stats = {k: v[index] for k, v in stats.items()}
+    dist = self.get_dist(stats)
+    if sample:
+      stoch = dist.sample()
+    else:
+      try:
+        stoch = dist.mode()
+      except:
+        stoch = dist.mean
+    return stoch, stats
+
+  def _suff_stats_ensemble(self, inp):
+    bs = list(inp.shape[:-1])
+    inp = inp.reshape([-1, inp.shape[-1]])
+    stats = []
+    for k in range(self._ensemble):
+      x = self._ensemble_img_out[k](inp)
+      x = self._act(x)
+      stats.append(self._suff_stats_layer('_ensemble_img_dist', x, k=k))
+    stats = {
+        k: torch.stack([x[k] for x in stats], 0)
+        for k, v in stats[0].items()}
+    stats = {
+        k: v.reshape([v.shape[0]] + bs + list(v.shape[2:]))
+        for k, v in stats.items()}
+    return stats
+
+  def _suff_stats_layer(self, name, x, k=None):
+    layer = getattr(self, name)
+    if k is not None:
+      layer = layer[k]
+    x = layer(x)
+    if self._discrete:
+      logit = x.reshape(list(x.shape[:-1]) + [self._stoch, self._discrete])
+      return {'logit': logit}
+    else:
+      mean, std = torch.chunk(x, 2, -1)
+      std = {
+          'softplus': lambda: F.softplus(std),
+          'sigmoid': lambda: torch.sigmoid(std),
+          'sigmoid2': lambda: 2 * torch.sigmoid(std / 2),
+      }[self._std_act]()
+      std = std + self._min_std
+      return {'mean': mean, 'std': std}
+
+  def vq_loss(self, post, prior, balance):
+    dim_repr = prior['output'].shape[-1]
+    # Vectors and codes are the same, but vectors have gradients
+    dyn_loss = balance * F.mse_loss(prior['output'], post['vectors'].detach()) + (1 - balance) * F.mse_loss(prior['output'].detach(), post['vectors'])
+    dyn_loss += balance * F.mse_loss(prior['output'], post['codes'].detach()) + (1 - balance) * F.mse_loss(prior['output'].detach(), post['codes']) 
+    dyn_loss /= 2
+    vq_loss = 0.25 * F.mse_loss(post['output'], post['codes'].detach()) + F.mse_loss(post['output'].detach(), post['codes'])
+
+    loss = vq_loss + dyn_loss 
+    return loss * dim_repr, dyn_loss * dim_repr
+  
+  def kl_loss(self, post, prior, forward, balance, free, free_avg,):
+    kld = D.kl_divergence
+    sg = lambda x: {k: v.detach() for k, v in x.items()} 
+    lhs, rhs = (prior, post) if forward else (post, prior)
+    mix = balance if forward else (1 - balance)
+    dtype = post['stoch'].dtype
+    device = post['stoch'].device
+    free_tensor = torch.tensor([free], dtype=dtype, device=device)
+    if balance == 0.5:
+      value = kld(self.get_dist(lhs), self.get_dist(rhs))
+      loss = torch.maximum(value, free_tensor).mean()
+    else:
+      value_lhs = value = kld(self.get_dist(lhs), self.get_dist(sg(rhs)))
+      value_rhs = kld(self.get_dist(sg(lhs)), self.get_dist(rhs))
+      if free_avg:
+        loss_lhs = torch.maximum(value_lhs.mean(), free_tensor)
+        loss_rhs = torch.maximum(value_rhs.mean(), free_tensor)
+      else:
+        loss_lhs = torch.maximum(value_lhs, free_tensor).mean()
+        loss_rhs = torch.maximum(value_rhs, free_tensor).mean()
+      loss = mix * loss_lhs + (1 - mix) * loss_rhs
+    return loss, value
+
+
+class Encoder(Module):
+
+  def __init__(
+      self, shapes, cnn_keys=r'.*', mlp_keys=r'.*', act='SiLU', norm='none',
+      cnn_depth=48, cnn_kernels=(4, 4, 4, 4), mlp_layers=[400, 400, 400, 400], symlog_inputs=False,):
+    super().__init__()
+    self.shapes = shapes
+    self.cnn_keys = [
+        k for k, v in shapes.items() if re.match(cnn_keys, k) and len(v) == 3]
+    self.mlp_keys = [
+        k for k, v in shapes.items() if re.match(mlp_keys, k) and len(v) == 1]
+    print('Encoder CNN inputs:', list(self.cnn_keys))
+    print('Encoder MLP inputs:', list(self.mlp_keys))
+    self._act = get_act(act)
+    self._norm = norm
+    self._cnn_depth = cnn_depth
+    self._cnn_kernels = cnn_kernels
+    self._mlp_layers = mlp_layers
+    self._symlog_inputs = symlog_inputs
+
+    if len(self.cnn_keys) > 0:
+      self._conv_model = []
+      for i, kernel in enumerate(self._cnn_kernels):
+        if i == 0:
+          prev_depth = 3
+        else:
+          prev_depth = 2 ** (i-1) * self._cnn_depth  
+        depth = 2 ** i * self._cnn_depth
+        self._conv_model.append(nn.Conv2d(prev_depth, depth, kernel, stride=2))
+        self._conv_model.append(ImgChLayerNorm(depth) if norm == 'layer' else NormLayer(norm,depth))
+        self._conv_model.append(self._act)
+      self._conv_model = nn.Sequential(*self._conv_model)
+    if len(self.mlp_keys) > 0:
+      self._mlp_model = []
+      for i, width in enumerate(self._mlp_layers):
+        if i == 0:
+          prev_width = np.sum([shapes[k] for k in self.mlp_keys]) 
+        else:
+          prev_width = self._mlp_layers[i-1]
+        self._mlp_model.append(nn.Linear(prev_width, width, bias=norm != 'none'))
+        self._mlp_model.append(NormLayer(norm, width))
+        self._mlp_model.append(self._act)
+      if len(self._mlp_model) == 0:
+        self._mlp_model.append(nn.Identity())
+      self._mlp_model = nn.Sequential(*self._mlp_model)
+
+  def forward(self, data):
+    key, shape = list(self.shapes.items())[0]
+    batch_dims = data[key].shape[:-len(shape)]
+    data = {
+        k: v.reshape((-1,) + tuple(v.shape)[len(batch_dims):])
+        for k, v in data.items()}
+    outputs = []
+    if self.cnn_keys:
+      outputs.append(self._cnn({k: data[k] for k in self.cnn_keys}))
+    if self.mlp_keys:
+      outputs.append(self._mlp({k: data[k] for k in self.mlp_keys}))
+    output = torch.cat(outputs, -1)
+    return output.reshape(batch_dims + output.shape[1:])
+
+  def _cnn(self, data):
+    x = torch.cat(list(data.values()), -1)
+    x = self._conv_model(x)
+    return x.reshape(tuple(x.shape[:-3]) + (-1,))
+
+  def _mlp(self, data):
+    x = torch.cat(list(data.values()), -1)
+    if self._symlog_inputs:
+      x = symlog(x)
+    x = self._mlp_model(x)
+    return x
+
+
+class Decoder(Module):
+
+  def __init__(
+      self, shapes, cnn_keys=r'.*', mlp_keys=r'.*', act='SiLU', norm='none',
+      cnn_depth=48, cnn_kernels=(4, 4, 4, 4), mlp_layers=[400, 400, 400, 400], embed_dim=1024, mlp_dist='mse', image_dist='mse'):
+    super().__init__()
+    self._embed_dim = embed_dim
+    self._shapes = shapes
+    self.cnn_keys = [
+        k for k, v in shapes.items() if re.match(cnn_keys, k) and len(v) == 3]
+    self.mlp_keys = [
+        k for k, v in shapes.items() if re.match(mlp_keys, k) and len(v) == 1]
+    print('Decoder CNN outputs:', list(self.cnn_keys))
+    print('Decoder MLP outputs:', list(self.mlp_keys))
+    self._act = get_act(act)
+    self._norm = norm
+    self._cnn_depth = cnn_depth
+    self._cnn_kernels = cnn_kernels
+    self._mlp_layers = mlp_layers
+    self.channels = {k: self._shapes[k][0] for k in self.cnn_keys}
+    self._mlp_dist = mlp_dist
+    self._image_dist = image_dist
+
+    if len(self.cnn_keys) > 0:
+
+      self._conv_in = nn.Sequential(nn.Linear(embed_dim, 32*self._cnn_depth))
+      self._conv_model = []
+      for i, kernel in enumerate(self._cnn_kernels):
+        if i == 0:
+          prev_depth = 32*self._cnn_depth
+        else:
+          prev_depth = 2 ** (len(self._cnn_kernels) - (i - 1) - 2) * self._cnn_depth
+        depth = 2 ** (len(self._cnn_kernels) - i - 2) * self._cnn_depth
+        act, norm = self._act, self._norm
+        # Last layer is dist layer 
+        if i == len(self._cnn_kernels) - 1:
+          depth, act, norm = sum(self.channels.values()), nn.Identity(), 'none'
+        self._conv_model.append(nn.ConvTranspose2d(prev_depth, depth, kernel, stride=2))
+        self._conv_model.append(ImgChLayerNorm(depth) if norm == 'layer' else NormLayer(norm, depth))
+        self._conv_model.append(act)
+      self._conv_model = nn.Sequential(*self._conv_model)
+    if len(self.mlp_keys) > 0:
+      self._mlp_model = []
+      for i, width in enumerate(self._mlp_layers):
+        if i == 0:
+          prev_width = embed_dim
+        else:
+          prev_width = self._mlp_layers[i-1]
+        self._mlp_model.append(nn.Linear(prev_width, width, bias=self._norm != 'none'))
+        self._mlp_model.append(NormLayer(self._norm, width))
+        self._mlp_model.append(self._act)
+      self._mlp_model = nn.Sequential(*self._mlp_model)
+      for key, shape in { k : shapes[k] for k in self.mlp_keys }.items():
+        self.add_module(f'dense_{key}', DistLayer(width, shape, dist=self._mlp_dist))
+
+  def forward(self, features):
+    outputs = {}
+    
+    if self.cnn_keys:
+      outputs.update(self._cnn(features))
+    if self.mlp_keys:
+      outputs.update(self._mlp(features))
+    return outputs
+
+  def _cnn(self, features):
+    x = self._conv_in(features)
+    x = x.reshape([-1, 32 * self._cnn_depth, 1, 1,])
+    x = self._conv_model(x)
+    x = x.reshape(list(features.shape[:-1]) + list(x.shape[1:])) 
+    if len(x.shape) == 5:
+      means = torch.split(x, list(self.channels.values()), 2)
+    else:
+      means = torch.split(x, list(self.channels.values()), 1)
+    image_dist = dict(mse=lambda x : MSEDist(x), normal_unit_std=lambda x : D.Independent(D.Normal(x, 1.0), 3))[self._image_dist]
+    dists = { key: image_dist(mean) for (key, shape), mean in zip(self.channels.items(), means)}
+    return dists
+
+  def _mlp(self, features):
+    shapes = {k: self._shapes[k] for k in self.mlp_keys}
+    x = features
+    x = self._mlp_model(x)
+    dists = {}
+    for key, shape in shapes.items():
+      dists[key] = getattr(self, f'dense_{key}')(x)
+    return dists
+
+
+class MLP(Module):
+
+  def __init__(self, in_shape, shape, layers, units, act='SiLU', norm='none', **out):
+    super().__init__()
+    self._in_shape = in_shape
+    if out['dist'] == 'twohot':
+      shape = 255
+    self._shape = (shape,) if isinstance(shape, int) else shape
+    self._layers = layers
+    self._units = units
+    self._norm = norm
+    self._act = get_act(act)
+    self._out = out
+    
+    last_units = in_shape
+    for index in range(self._layers):
+      self.add_module(f'dense{index}', nn.Linear(last_units, units, bias=norm != 'none'))
+      self.add_module(f'norm{index}', NormLayer(norm, units))
+      last_units = units
+    self._out = DistLayer(units, shape, **out)
+
+  def forward(self, features):
+    x = features 
+    x = x.reshape([-1, x.shape[-1]])
+    for index in range(self._layers):
+      x = getattr(self, f'dense{index}')(x)
+      x = getattr(self, f'norm{index}')(x)
+      x = self._act(x)
+    x = x.reshape(list(features.shape[:-1]) + [x.shape[-1]])
+    return self._out(x)
+
+
+class GRUCell(Module):
+
+  def __init__(self, inp_size, size, norm=False, act='Tanh', update_bias=-1, device='cuda', **kwargs):
+    super().__init__()
+    self._inp_size = inp_size
+    self._size = size
+    self._act = get_act(act)
+    self._norm = norm
+    self._update_bias = update_bias
+    self.device = device
+    self._layer = nn.Linear(inp_size + size, 3 * size, bias=(not norm), **kwargs)
+    if norm:
+      self._norm = nn.LayerNorm(3*size)
+
+  def get_initial_state(self, inputs=None, batch_size=None, dtype=None):
+    return torch.zeros((batch_size), self._size, device=self.device)
+
+  @property
+  def state_size(self):
+    return self._size
+
+  def forward(self, inputs, deter_state):
+    """
+      inputs : non-linear combination of previous stoch and action 
+      deter_state : prev hidden state of the cell
+    """
+    deter_state = deter_state[0]  # State is wrapped in a list.
+    parts = self._layer(torch.cat([inputs, deter_state], -1))
+    if self._norm:
+      parts = self._norm(parts)
+    reset, cand, update = torch.chunk(parts, 3, -1)
+    reset = torch.sigmoid(reset)
+    cand = self._act(reset * cand)
+    update = torch.sigmoid(update + self._update_bias)
+    output = update * cand + (1 - update) * deter_state
+    return output, [output]
+
+class DistLayer(Module):
+
+  def __init__(
+      self, in_dim, shape, dist='mse', min_std=0.1, max_std=1.0, init_std=0.0, bias=True):
+    super().__init__()
+    self._in_dim = in_dim
+    self._shape = shape if type(shape) in [list,tuple] else [shape]
+    self._dist = dist
+    self._min_std = min_std
+    self._init_std = init_std
+    self._max_std = max_std
+    self._out = nn.Linear(in_dim, int(np.prod(shape)) , bias=bias)
+    if dist in ('normal', 'tanh_normal', 'trunc_normal'):
+      self._std = nn.Linear(in_dim, int(np.prod(shape)) )
+
+  def forward(self, inputs):
+    out = self._out(inputs)
+    out = out.reshape(list(inputs.shape[:-1]) + list(self._shape)) 
+    if self._dist in ('normal', 'tanh_normal', 'trunc_normal'):
+      std = self._std(inputs) 
+      std = std.reshape(list(inputs.shape[:-1]) + list(self._shape)) 
+    if self._dist == 'mse':
+      return MSEDist(out,)
+    if self._dist == 'normal_unit_std':
+      dist = D.Normal(out, 1.0)
+      dist.sample = dist.rsample
+      return D.Independent(dist, len(self._shape))
+    if self._dist == 'normal':
+      mean = torch.tanh(out)
+      std = (self._max_std - self._min_std) * torch.sigmoid(std + 2.0) + self._min_std
+      dist = D.Normal(mean, std)
+      dist.sample = dist.rsample
+      return D.Independent(dist, len(self._shape))
+    if self._dist == 'binary':
+      out = torch.sigmoid(out)
+      dist = BernoulliDist(out)
+      return D.Independent(dist, len(self._shape))
+    if self._dist == 'tanh_normal':
+      mean = 5 * torch.tanh(out / 5)
+      std = F.softplus(std + self._init_std) + self._min_std
+      dist = utils.SquashedNormal(mean, std)
+      dist = D.Independent(dist, len(self._shape))
+      return SampleDist(dist)
+    if self._dist == 'trunc_normal':
+      mean = torch.tanh(out)
+      std = 2 * torch.sigmoid((std + self._init_std) / 2) + self._min_std
+      dist = utils.TruncatedNormal(mean, std)
+      return D.Independent(dist, 1)
+    if self._dist == 'onehot':
+      return OneHotDist(out.float()) 
+    if self._dist == 'twohot':
+      return TwoHotDist(out.float())
+    if self._dist == 'symlog_mse':
+      return SymlogDist(out, len(self._shape), 'mse') 
+    raise NotImplementedError(self._dist)
+
+
+class NormLayer(Module):
+
+  def __init__(self, name, dim=None):
+    super().__init__()
+    if name == 'none':
+      self._layer = None
+    elif name == 'layer':
+      assert dim != None
+      self._layer = nn.LayerNorm(dim)
+    else:
+      raise NotImplementedError(name)
+
+  def forward(self, features):
+    if self._layer is None:
+      return features
+    return self._layer(features)
+
+
+def get_act(name):
+  if name == 'none':
+    return nn.Identity()
+  elif hasattr(nn, name):
+    return getattr(nn, name)()
+  else:
+    raise NotImplementedError(name)
+
+
+class Optimizer:
+
+  def __init__(
+      self, name, parameters, lr, eps=1e-4, clip=None, wd=None,
+      opt='adam', wd_pattern=r'.*', use_amp=False):
+    assert 0 <= wd < 1
+    assert not clip or 1 <= clip
+    self._name = name
+    self._clip = clip
+    self._wd = wd
+    self._wd_pattern = wd_pattern
+    self._opt = {
+        'adam': lambda: torch.optim.Adam(parameters, lr, eps=eps),
+        'nadam': lambda: torch.optim.Nadam(parameters, lr, eps=eps),
+        'adamax': lambda: torch.optim.Adamax(parameters, lr, eps=eps),
+        'sgd': lambda: torch.optim.SGD(parameters, lr),
+        'momentum': lambda: torch.optim.SGD(lr, momentum=0.9),
+    }[opt]()
+    self._scaler = torch.cuda.amp.GradScaler(enabled=use_amp)
+    self._once = True
+
+  def __call__(self, loss, params):
+    params = list(params)
+    assert len(loss.shape) == 0 or (len(loss.shape) == 1 and loss.shape[0] == 1), (self._name, loss.shape)
+    metrics = {}
+
+    # Count parameters.
+    if self._once:
+      count = sum(p.numel() for p in params if p.requires_grad) 
+      print(f'Found {count} {self._name} parameters.')
+      self._once = False
+
+    # Check loss.
+    metrics[f'{self._name}_loss'] = loss.detach().cpu().numpy()
+
+    # Compute scaled gradient.
+    self._scaler.scale(loss).backward()
+    self._scaler.unscale_(self._opt)
+
+    # Gradient clipping.
+    if self._clip:
+      norm = torch.nn.utils.clip_grad_norm_(params, self._clip)
+      metrics[f'{self._name}_grad_norm'] = norm.item()
+  
+    # Weight decay.
+    if self._wd:
+      self._apply_weight_decay(params)
+    
+    # # Apply gradients.
+    self._scaler.step(self._opt)
+    self._scaler.update()
+    
+    self._opt.zero_grad() 
+    return metrics
+
+  def _apply_weight_decay(self, varibs):
+    nontrivial = (self._wd_pattern != r'.*')
+    if nontrivial:
+      raise NotImplementedError('Non trivial weight decay')
+    else:
+      for var in varibs:
+        var.data = (1 - self._wd) * var.data
+
+class StreamNorm:
+
+  def __init__(self, shape=(), momentum=0.99, scale=1.0, eps=1e-8, device='cuda'):
+    # Momentum of 0 normalizes only based on the current batch.
+    # Momentum of 1 disables normalization.
+    self.device = device
+    self._shape = tuple(shape)
+    self._momentum = momentum
+    self._scale = scale
+    self._eps = eps
+    self.mag = None # torch.ones(shape).to(self.device) 
+
+    self.step = 0
+    self.mean = None # torch.zeros(shape).to(self.device)
+    self.square_mean = None # torch.zeros(shape).to(self.device)
+
+  def reset(self):
+    self.step = 0
+    self.mag = None # torch.ones_like(self.mag).to(self.device)
+    self.mean = None # torch.zeros_like(self.mean).to(self.device)
+    self.square_mean = None # torch.zeros_like(self.square_mean).to(self.device)
+
+  def __call__(self, inputs):
+    metrics = {}
+    self.update(inputs)
+    metrics['mean'] = inputs.mean()
+    metrics['std'] = inputs.std()
+    outputs = self.transform(inputs)
+    metrics['normed_mean'] = outputs.mean()
+    metrics['normed_std'] = outputs.std()
+    return outputs, metrics
+
+  def update(self, inputs):
+    self.step += 1
+    batch = inputs.reshape((-1,) + self._shape)
+    
+    mag = torch.abs(batch).mean(0) 
+    if self.mag is not None:
+      self.mag.data = self._momentum * self.mag.data + (1 - self._momentum) * mag 
+    else:
+      self.mag =  mag.clone().detach()
+    
+    mean = torch.mean(batch)
+    if self.mean is not None:
+      self.mean.data = self._momentum * self.mean.data + (1 - self._momentum) * mean 
+    else:
+      self.mean = mean.clone().detach()
+    
+    square_mean = torch.mean(batch * batch)
+    if self.square_mean is not None:
+      self.square_mean.data = self._momentum * self.square_mean.data + (1 - self._momentum) * square_mean 
+    else:
+      self.square_mean = square_mean.clone().detach()
+
+  def transform(self, inputs):
+    if self._momentum == 1:
+      return inputs
+    values = inputs.reshape((-1,) + self._shape)
+    values /= self.mag[None] + self._eps 
+    values *= self._scale
+    return values.reshape(inputs.shape)
+
+  def corrected_mean_var_std(self,):
+    corr = 1 # 1 - self._momentum ** self.step # NOTE: this led to exploding values for first few iterations
+    corr_mean = self.mean / corr 
+    corr_var = (self.square_mean / corr) - self.mean ** 2
+    corr_std = torch.sqrt(torch.maximum(corr_var, torch.zeros_like(corr_var, device=self.device)) + self._eps)
+    return corr_mean, corr_var, corr_std
+
+class RequiresGrad:
+
+  def __init__(self, model):
+    self._model = model
+
+  def __enter__(self):
+    self._model.requires_grad_(requires_grad=True)
+
+  def __exit__(self, *args):
+    self._model.requires_grad_(requires_grad=False)
+
+class RewardEMA:
+    """running mean and std"""
+
+    def __init__(self, device, alpha=1e-2):
+        self.device = device
+        self.alpha = alpha
+        self.range = torch.tensor([0.05, 0.95]).to(device)
+
+    def __call__(self, x, ema_vals):
+        flat_x = torch.flatten(x.detach())
+        x_quantile = torch.quantile(input=flat_x, q=self.range)
+        # this should be in-place operation
+        ema_vals[:] = self.alpha * x_quantile + (1 - self.alpha) * ema_vals
+        scale = torch.clip(ema_vals[1] - ema_vals[0], min=1.0)
+        offset = ema_vals[0]
+        return offset.detach(), scale.detach()
+
+class ImgChLayerNorm(nn.Module):
+    def __init__(self, ch, eps=1e-03):
+        super(ImgChLayerNorm, self).__init__()
+        self.norm = torch.nn.LayerNorm(ch, eps=eps)
+
+    def forward(self, x):
+        x = x.permute(0, 2, 3, 1)
+        x = self.norm(x)
+        x = x.permute(0, 3, 1, 2)
+        return x

agent/genrl.py

@@ -0,0 +1,124 @@
+import torch
+from agent.dreamer import DreamerAgent, ActorCritic, stop_gradient, env_reward
+import agent.dreamer_utils as common
+import agent.video_utils as video_utils
+from tools.genrl_utils import *
+
+def connector_update_fn(self, module_name, data, outputs, metrics):
+    connector = getattr(self, module_name)
+    n_frames = connector.n_frames
+    B, T = data['observation'].shape[:2]
+
+    # video embed are actions
+    if getattr(self.cfg, "viclip_encode", False):
+      video_embed = data['clip_video']
+    else:
+      # Obtaining video embed
+      with torch.no_grad():
+        viclip_model = getattr(self, 'viclip_model')
+        processed_obs = viclip_model.preprocess_transf(data['observation'].reshape(B*T, *data['observation'].shape[2:]) / 255)
+        reshaped_obs = processed_obs.reshape(B * (T // n_frames), n_frames, 3,224,224)
+        video_embed = viclip_model.get_vid_features(reshaped_obs.to(viclip_model.device))
+      
+    # Get posterior states from original model
+    wm_post = outputs['post']
+    return connector.update(video_embed, wm_post)
+
+class GenRLAgent(DreamerAgent):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        
+        self.n_frames = 8 # NOTE: this should become an hyperparam if changing the model
+        self.viclip_emb_dim =  512 # NOTE: this should become an hyperparam if changing the model
+        
+        assert self.cfg.batch_length % self.n_frames == 0, "Fix batch length param"
+        
+        if 'clip_video' in self.obs_space:
+          self.viclip_emb_dim = self.obs_space['clip_video'].shape[0]
+        
+        connector = video_utils.VideoSSM(**self.cfg.connector, **self.cfg.connector_rssm, connector_kl=self.cfg.connector_kl, 
+                                          n_frames=self.n_frames, action_dim=self.viclip_emb_dim + self.n_frames, 
+                                          clip_add_noise=self.cfg.clip_add_noise, clip_lafite_noise=self.cfg.clip_lafite_noise,
+                                          device=self.device, cell_input='stoch') 
+        
+        connector.to(self.device)
+
+        self.wm.add_module_to_update('connector', connector, connector_update_fn, detached=self.cfg.connector.detached_post)
+    
+        if getattr(self.cfg, 'imag_reward_fn', None) is not None:
+          self.instantiate_imag_behavior()
+    
+    def instantiate_imag_behavior(self):
+      self._imag_behavior = ActorCritic(self.cfg, self.act_spec, self.wm.inp_size, name='imag').to(self.device) 
+      self._imag_behavior.rewnorm = common.StreamNorm(**self.cfg.imag_reward_norm, device=self.device)    
+    
+    def finetune_mode(self,):
+      self._acting_behavior = self._imag_behavior
+      self.wm.detached_update_fns = {}
+      self.wm.e2e_update_fns = {}
+      self.wm.grad_heads.append('reward')
+
+    def update_wm(self, data, step):
+      return super().update_wm(data, step)
+
+    def report(self, data, key='observation', nvid=8):
+      # Redefine data with trim
+      n_frames = self.wm.connector.n_frames
+      obs = data['observation'][:nvid, n_frames:]
+      B, T = obs.shape[:2]
+
+      report_data = super().report(data)
+      wm = self.wm
+      n_frames = wm.connector.n_frames
+      
+      # Init is same as Dreamer for reporting
+      truth = data[key][:nvid] / 255
+      decoder = wm.heads['decoder'] # B, T, C, H, W
+      preprocessed_data = self.wm.preprocess(data)
+
+      embed = wm.encoder(preprocessed_data)
+      states, _ = wm.rssm.observe(embed[:nvid, :n_frames], data['action'][:nvid, :n_frames], data['is_first'][:nvid, :n_frames])
+      recon = decoder(wm.decoder_input_fn(states))[key].mean[:nvid] # mode
+      dreamer_init = {k: v[:, -1] for k, v in states.items()}
+
+      # video embed are actions
+      if getattr(self.cfg, "viclip_encode", False):
+        video_embed = data['clip_video'][:nvid,n_frames*2-1::n_frames]
+      else:
+        # Obtain embed
+        processed_obs = wm.viclip_model.preprocess_transf(obs.reshape(B*T, *obs.shape[2:]) / 255)
+        reshaped_obs = processed_obs.reshape(B * (T // n_frames), n_frames, 3,224,224)
+        video_embed = wm.viclip_model.get_vid_features(reshaped_obs.to(wm.viclip_model.device))
+      
+      video_embed = video_embed.to(self.device)
+
+      # Get actions
+      video_embed = video_embed.reshape(B, T // n_frames, -1).unsqueeze(2).repeat(1,1,n_frames, 1).reshape(B, T, -1)
+      prior = wm.connector.video_imagine(video_embed, dreamer_init, reset_every_n_frames=False)
+      prior_recon = decoder(wm.decoder_input_fn(prior))[key].mean # mode
+      model = torch.clip(torch.cat([recon[:, :n_frames] + 0.5, prior_recon + 0.5], 1), 0, 1)
+      error = (model - truth + 1) / 2
+      
+      # Add video to logs
+      video = torch.cat([truth, model, error], 3)
+      report_data['video_clip_pred'] = video
+      
+      return report_data
+
+    def update_imag_behavior(self, state=None, outputs=None, metrics={}, seq_data=None,):
+        if getattr(self.cfg, 'imag_reward_fn', None) is None:
+           return outputs['post'], metrics
+        if outputs is not None:
+            post = outputs['post']
+            is_terminal = outputs['is_terminal']
+        else:
+            seq_data = self.wm.preprocess(seq_data)
+            embed = self.wm.encoder(seq_data)
+            post, _ = self.wm.rssm.observe(
+                embed, seq_data['action'], seq_data['is_first'])
+            is_terminal = seq_data['is_terminal']
+        #
+        start = {k: stop_gradient(v) for k,v in post.items()}
+        imag_reward_fn = lambda seq: globals()[self.cfg.imag_reward_fn](self, seq, **self.cfg.imag_reward_args)
+        metrics.update(self._imag_behavior.update(self.wm, start, is_terminal, imag_reward_fn,))
+        return start, metrics

agent/genrl.yaml

@@ -0,0 +1,22 @@
+# @package agent
+_target_: agent.genrl.GenRLAgent
+name: genrl
+cfg: ???
+obs_space: ???
+act_spec: ???
+grad_heads: [decoder]
+reward_norm: {momentum: 1.0, scale: 1.0, eps: 1e-8}
+actor_ent: 0
+additional_report_fns: ['report_text2video']
+
+clip_add_noise: 0.0
+clip_lafite_noise: 0.5
+
+connector: { token_dropout: 0, loss_scale: 1, denoising_ae: True, detached_post: True, temporal_embeds: False, rescale_embeds: True}
+connector_rssm: {ensemble: 1, hidden: 1024, deter: 1024, stoch: 32, discrete: 32,  norm: layer, std_act: softplus, min_std: 0.1, single_obs_posterior: false, learn_initial: True } # act: elu,
+connector_kl: {free: 0.0, forward: True, balance: 0.8, free_avg: False, } # note forward is true by default
+
+imag_reward_fn: null
+imag_reward_norm: {momentum: 1.00, scale: 1.0, eps: 1e-8}
+imag_reward_args: {score_fn: 'max_cosine', sample_for_target: False, align_initial : False, weighted_align : False, align_sequence: True, skip_first_target: True }
+# +imag_reward_args.task_prompt

agent/plan2explore.py

@@ -0,0 +1,108 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from agent.dreamer import DreamerAgent, stop_gradient
+import agent.dreamer_utils as common
+
+class Disagreement(nn.Module):
+    def __init__(self, obs_dim, action_dim, hidden_dim, n_models=5, pred_dim=None):
+        super().__init__()
+        if pred_dim is None: pred_dim = obs_dim
+        self.ensemble = nn.ModuleList([
+            nn.Sequential(nn.Linear(obs_dim + action_dim, hidden_dim),
+                          nn.ReLU(), nn.Linear(hidden_dim, pred_dim))
+            for _ in range(n_models)
+        ])
+
+    def forward(self, obs, action, next_obs):
+        assert obs.shape[0] == next_obs.shape[0]
+        assert obs.shape[0] == action.shape[0]
+
+        errors = []
+        for model in self.ensemble:
+            next_obs_hat = model(torch.cat([obs, action], dim=-1))
+            model_error = torch.norm(next_obs - next_obs_hat,
+                                     dim=-1,
+                                     p=2,
+                                     keepdim=True)
+            errors.append(model_error)
+
+        return torch.cat(errors, dim=1)
+
+    def get_disagreement(self, obs, action):
+        assert obs.shape[0] == action.shape[0]
+
+        preds = []
+        for model in self.ensemble:
+            next_obs_hat = model(torch.cat([obs, action], dim=-1))
+            preds.append(next_obs_hat)
+        preds = torch.stack(preds, dim=0)
+        return torch.var(preds, dim=0).mean(dim=-1)
+
+
+class Plan2Explore(DreamerAgent):
+    def __init__(self, **kwargs):
+        super().__init__(**kwargs)
+        in_dim = self.wm.inp_size
+        pred_dim = self.wm.embed_dim
+        self.hidden_dim = pred_dim
+        self.reward_free = True
+
+        self.disagreement = Disagreement(in_dim, self.act_dim,
+                                         self.hidden_dim, pred_dim=pred_dim).to(self.device)
+
+        # optimizers
+        self.disagreement_opt = common.Optimizer('disagreement', self.disagreement.parameters(), **self.cfg.model_opt, use_amp=self._use_amp)
+        self.disagreement.train()
+        self.requires_grad_(requires_grad=False)
+
+    def update_disagreement(self, obs, action, next_obs, step):
+        metrics = dict()
+
+        error = self.disagreement(obs, action, next_obs)
+
+        loss = error.mean()
+
+        metrics.update(self.disagreement_opt(loss, self.disagreement.parameters()))
+
+        metrics['disagreement_loss'] = loss.item()
+
+        return metrics
+
+    def compute_intr_reward(self, seq):
+        obs, action = seq['feat'][:-1], stop_gradient(seq['action'][1:])
+        intr_rew = torch.zeros(list(seq['action'].shape[:-1]) + [1], device=self.device)
+        if len(action.shape) > 2:
+            B, T, _ = action.shape
+            obs = obs.reshape(B*T, -1)
+            action = action.reshape(B*T, -1)
+            reward = self.disagreement.get_disagreement(obs, action).reshape(B, T, 1)
+        else:
+            reward = self.disagreement.get_disagreement(obs, action).unsqueeze(-1)
+        intr_rew[1:] = reward
+        return intr_rew
+
+    def update(self, data, step):
+        metrics = {}
+        B, T, _ = data['action'].shape
+        state, outputs, mets = self.wm.update(data, state=None)
+        metrics.update(mets)
+        start = outputs['post']
+        start = {k: stop_gradient(v) for k,v in start.items()}
+        if self.reward_free:
+            T = T-1
+            inp = stop_gradient(outputs['feat'][:, :-1]).reshape(B*T, -1)
+            action = data['action'][:, 1:].reshape(B*T, -1)
+            out = stop_gradient(outputs['embed'][:,1:]).reshape(B*T,-1)
+            with common.RequiresGrad(self.disagreement):
+                with torch.cuda.amp.autocast(enabled=self._use_amp):
+                    metrics.update(
+                        self.update_disagreement(inp, action, out, step))
+            metrics.update(self._acting_behavior.update(
+                self.wm, start, data['is_terminal'], reward_fn=self.compute_intr_reward))
+        else:
+            reward_fn = lambda seq: self.wm.heads['reward'](seq['feat']).mean 
+            metrics.update(self._acting_behavior.update(
+                self.wm, start, data['is_terminal'], reward_fn))
+        return state, metrics

agent/plan2explore.yaml

@@ -0,0 +1,9 @@
+# @package agent
+_target_: agent.plan2explore.Plan2Explore
+name: plan2explore
+cfg: ???
+obs_space: ???
+act_spec: ???
+grad_heads: [decoder]
+reward_norm: {momentum: 0.95, scale: 1.0, eps: 1e-8}
+actor_ent: 0 

agent/video_utils.py

@@ -0,0 +1,240 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import agent.dreamer_utils as common
+from collections import defaultdict
+import numpy as np
+
+class ResidualLinear(nn.Module):
+    def __init__(self, in_channels, out_channels, norm='layer', act='SiLU', prenorm=False):
+        super().__init__()
+        self.norm_layer = common.NormLayer(norm, in_channels if prenorm else out_channels)
+        self.act = common.get_act(act)
+        self.layer = nn.Linear(in_channels, out_channels)
+        self.prenorm = prenorm
+        self.res_proj = nn.Identity() if in_channels == out_channels else nn.Linear(in_channels, out_channels)
+    
+    def forward(self, x):
+        if self.prenorm:
+            h = self.norm_layer(x)
+            h = self.layer(h)
+        else:            
+            h = self.layer(x)
+            h = self.norm_layer(h)
+        h = self.act(h)
+        return h + self.res_proj(x)
+
+class UNetDenoiser(nn.Module):
+    def __init__(self, in_channels : int, mid_channels : int, n_layers : int, norm='layer', act= 'SiLU', ):
+        super().__init__()
+        out_channels = in_channels
+        self.down = nn.ModuleList()
+        for i in range(n_layers):
+            if i == (n_layers - 1):
+                self.down.append(ResidualLinear(in_channels, mid_channels, norm=norm, act=act))
+            else:
+                self.down.append(ResidualLinear(in_channels, in_channels, norm=norm, act=act))
+
+        self.mid = nn.ModuleList()
+        for i in range(n_layers):
+            self.mid.append(ResidualLinear(mid_channels, mid_channels, norm=norm, act=act))
+        
+        self.up = nn.ModuleList()
+        for i in range(n_layers):
+            if i == 0:
+                self.up.append(ResidualLinear(mid_channels * 2, out_channels, norm='none', act='Identity'))
+            else:
+                self.up.append(ResidualLinear(out_channels * 2, out_channels, norm=norm, act=act))
+        
+    def forward(self, x):
+        down_res = []
+        for down_layer in self.down:
+            x = down_layer(x)
+            down_res.append(x)
+        
+        for mid_layer in self.mid:
+            x = mid_layer(x)
+        
+        down_res.reverse()
+        for up_layer, res in zip(self.up, down_res):
+            x = up_layer(torch.cat([x, res], dim=-1))
+        return x
+
+
+class VideoSSM(common.EnsembleRSSM):
+    def __init__(self, *args, 
+                 connector_kl={}, temporal_embeds=False, detached_post=True, n_frames=8, 
+                 token_dropout=0.,  loss_scale=1, clip_add_noise=0, clip_lafite_noise=0,
+                 rescale_embeds=False, denoising_ae=False, learn_initial=True, **kwargs,):
+        super().__init__(*args, **kwargs)
+        #
+        self.n_frames = n_frames
+        # by default, adding the n_frames in actions (doesn't hurt and easier to test whether it's useful or not)
+        self.viclip_emb_dim = kwargs['action_dim'] - self.n_frames
+        #
+        self.temporal_embeds = temporal_embeds
+        self.detached_post = detached_post
+        self.connector_kl = connector_kl
+        self.token_dropout = token_dropout
+        self.loss_scale = loss_scale
+        self.rescale_embeds = rescale_embeds
+        self.clip_add_noise = clip_add_noise
+        self.clip_lafite_noise = clip_lafite_noise
+        self.clip_const = np.sqrt(self.viclip_emb_dim).item()
+        self.denoising_ae = denoising_ae
+        if self.denoising_ae:
+            self.aligner = UNetDenoiser(self.viclip_emb_dim, self.viclip_emb_dim // 2, n_layers=2, norm='layer', act='SiLU')
+        self.learn_initial = learn_initial
+        if self.learn_initial:
+            self.initial_state_pred = nn.Sequential(
+                nn.Linear(kwargs['action_dim'], kwargs['hidden']),
+                common.NormLayer(kwargs['norm'],kwargs['hidden']), common.get_act('SiLU'),
+                nn.Linear(kwargs['hidden'], kwargs['hidden']),
+                common.NormLayer(kwargs['norm'],kwargs['hidden']), common.get_act('SiLU'),
+                nn.Linear(kwargs['hidden'], kwargs['deter'])
+            )
+        # Deleting non-useful models
+        del self._obs_out
+        del self._obs_dist
+
+    def initial(self, batch_size, init_embed=None, ignore_learned=False):
+        init = super().initial(batch_size)
+        if self.learn_initial and not ignore_learned and hasattr(self, 'initial_state_pred'):
+            assert init_embed is not None
+            # patcher to avoid edge cases
+            if init_embed.shape[-1] == self.viclip_emb_dim:
+                patcher = torch.zeros((*init_embed.shape[:-1], 8), device=self.device)
+                init_embed = torch.cat([init_embed, patcher], dim=-1)
+            init['deter'] = self.initial_state_pred(init_embed)
+            stoch, stats = self.get_stoch_stats_from_deter_state(init)
+            init['stoch'] = stoch
+            init.update(stats)
+        return init
+
+    def get_action(self, video_embed):
+        n_frames = self.n_frames 
+        B, T = video_embed.shape[:2]
+
+        if self.rescale_embeds:
+            video_embed = video_embed * self.clip_const
+        
+        temporal_embeds = F.one_hot(torch.arange(T).to(video_embed.device) % n_frames, n_frames).reshape(1, T, n_frames,).repeat(B, 1, 1,)
+        if not self.temporal_embeds:
+            temporal_embeds *= 0
+        
+        return torch.cat([video_embed, temporal_embeds],dim=-1)
+
+    def update(self, video_embed, wm_post):
+        n_frames = self.n_frames 
+        B, T = video_embed.shape[:2]
+        loss = 0
+        metrics = {}
+
+        # NOVEL
+        video_embed = video_embed[:,n_frames-1::n_frames] # tested
+        video_embed = video_embed.to(self.device)
+        video_embed = video_embed.reshape(B, T // n_frames, 1, -1).repeat(1,1, n_frames, 1).reshape(B, T, -1)
+
+        orig_video_embed = video_embed
+
+        if self.clip_add_noise > 0:
+            video_embed = video_embed  + torch.randn_like(video_embed, device=video_embed.device) * self.clip_add_noise
+            video_embed = nn.functional.normalize(video_embed, dim=-1)
+        if self.clip_lafite_noise > 0:
+            normed_noise = F.normalize(torch.randn_like(video_embed, device=video_embed.device), dim=-1)
+            video_embed = (1 - self.clip_lafite_noise) * video_embed  + self.clip_lafite_noise * normed_noise
+            video_embed = nn.functional.normalize(video_embed, dim=-1)
+
+        if self.denoising_ae:
+            assert (self.clip_lafite_noise + self.clip_add_noise) > 0, "Nothing to denoise"
+            denoised_embed = self.aligner(video_embed)
+            denoised_embed = F.normalize(denoised_embed, dim=-1)
+            denoising_loss = 1 - F.cosine_similarity(denoised_embed, orig_video_embed, dim=-1).mean() # works same as F.mse_loss(denoised_embed, orig_video_embed).mean()
+            loss += denoising_loss
+            metrics['aligner_cosine_distance'] = denoising_loss
+            # if using a denoiser, it's the denoiser's duty to denoise the video embed
+            video_embed = orig_video_embed # could also be denoised_embed for e2e training
+
+        embed_actions = self.get_action(video_embed)
+        
+        if self.detached_post:
+            wm_post = { k : v.reshape(B, T, *v.shape[2:]).detach() for k,v in wm_post.items() } 
+        else:
+            wm_post = { k : v.reshape(B, T, *v.shape[2:]) for k,v in wm_post.items() } 
+        
+        # Get prior states
+        prior_states = defaultdict(list)
+        for t in range(T):
+            # Get video action
+            action = embed_actions[:, t]
+
+            if t == 0: 
+                prev_state = self.initial(batch_size=wm_post['stoch'].shape[0], init_embed=action)
+            else:
+                # Get deter from prior, get stoch from wm_post
+                prev_state = prior
+                prev_state[self.cell_input] = wm_post[self.cell_input][:, t-1]
+
+            if self.token_dropout > 0:
+                prev_state['stoch'] = torch.einsum('b...,b->b...', prev_state['stoch'], (torch.rand(B, device=action.device) > self.token_dropout).float() )
+
+            prior = self.img_step(prev_state, action)
+            for k in prior: 
+                prior_states[k].append(prior[k])
+
+        # Aggregate
+        for k in prior_states:
+            prior_states[k] = torch.stack(prior_states[k], dim=1)
+
+        # Compute loss
+        prior = prior_states
+            
+        kl_loss, kl_value = self.kl_loss(wm_post, prior, **self.connector_kl)
+        video_loss = self.loss_scale * kl_loss 
+        metrics['connector_kl'] = kl_value.mean()
+        loss += video_loss
+        
+        # Compute initial KL
+        video_embed = video_embed.reshape(B, T // n_frames, n_frames, -1)[:,1:,0].reshape(B * (T//n_frames-1), 1, -1) # taking only one (0) and skipping first temporal step
+        embed_actions = self.get_action(video_embed)
+        wm_post = { k : v.reshape(B, T // n_frames, n_frames, *v.shape[2:])[:,1:,0].reshape(B * (T//n_frames-1), *v.shape[2:]) for k,v in wm_post.items() } 
+        action = embed_actions[:, 0]
+        prev_state = self.initial(batch_size=wm_post['stoch'].shape[0], init_embed=action)
+        prior = self.img_step(prev_state, action)
+        kl_loss, kl_value = self.kl_loss(wm_post, prior, **self.connector_kl)
+        metrics['connector_initial_kl'] = kl_value.mean()
+
+        return loss, metrics
+            
+    def video_imagine(self, video_embed, dreamer_init=None, sample=True, reset_every_n_frames=True, denoise=False):
+        n_frames = self.n_frames 
+        B, T = video_embed.shape[:2]
+
+        if self.denoising_ae and denoise:
+            denoised_embed = self.aligner(video_embed)
+            video_embed = F.normalize(denoised_embed, dim=-1)
+
+        action = self.get_action(video_embed)
+        # Imagine
+        init = self.initial(batch_size=B, init_embed=action[:, 0]) # -> this ensures only stoch is used from the current frame
+        if dreamer_init is not None:
+            init[self.cell_input] = dreamer_init[self.cell_input]
+
+        if reset_every_n_frames:
+            prior_states = defaultdict(list)
+            for action_chunk in torch.chunk(action, T // n_frames, dim=1):
+                prior = self.imagine(action_chunk, init, sample=sample)
+                for k in prior: 
+                    prior_states[k].append(prior[k])
+        
+                # -> this ensures only stoch is used from the current frame
+                init = self.initial(batch_size=B, ignore_learned=True) 
+                init[self.cell_input] = prior[self.cell_input][:, -1]
+            
+            # Agg
+            for k in prior_states:
+                prior_states[k] = torch.cat(prior_states[k], dim=1)
+            prior = prior_states
+        else:
+            prior = self.imagine(action, init, sample=sample)
+        return prior

app.py

@@ -0,0 +1,80 @@
+import os
+import sys
+import gradio as gr
+
+# prototyping
+# from demo_test import Text2Video, Video2Video
+
+from demo.t2v import Text2Video
+
+t2v_examples = [
+    ['walk fast clean',16,],
+    ['run fast clean',16,],
+    ['standing up',16],
+    ['doing the splits',16],
+    ['doing backflips',16],
+    ['a headstand',16],
+    ['karate kick',16],
+    ['crunch abs',16],
+    ['doing push ups',16],
+]
+
+def do_nothing():
+    return
+
+def videocrafter_demo(result_dir='./tmp/'):
+    text2video = Text2Video(result_dir)
+    # video2video = Video2Video(result_dir)
+
+    # tex
+    with gr.Blocks(analytics_enabled=False) as videocrafter_iface:
+        gr.Markdown("<div align='center'> <h2> GenRL: Multimodal foundation world models for generalist embodied agents </span> </h2> \
+                     <a style='font-size:18px;' href='https://github.com/mazpie/genrl'> [Github] \
+                    &nbsp; &nbsp; \
+                     <a style='font-size:18px;' href='https://huggingface.co/mazpie/genrl_models'> [Models] </div> \
+                    &nbsp; &nbsp; \
+                     <a style='font-size:18px;' href='https://huggingface.co/mazpie/genrl_models'> [Models] </div>")
+        
+        gr.Markdown("<b> Notes: </b>")
+        gr.Markdown("<b> - Low quality of the videos generated is expected, as the work focuses on visual-language alignment for behavior learning, not on video generation quality.</b>")
+        gr.Markdown("<b> - The model is trained on small 64x64 images, and the videos are generated only from a small 512-dimensional embedding. </b>")
+        gr.Markdown("<b> - Some prompts require styling instructions, e.g. fast, clean, in order to work well. See some of the examples. </b>")
+        
+        #######t2v#######
+        with gr.Tab(label="Text2Video"):
+            with gr.Column():
+                with gr.Row(): # .style(equal_height=False)
+                    with gr.Column():
+                        input_text = gr.Text(label='prompt')
+                        duration = gr.Slider(minimum=8, maximum=32, elem_id=f"duration", label="duration", value=16, step=8)
+                        send_btn = gr.Button("Send")
+                    with gr.Column(): # label='result',
+                        pass
+                    with gr.Column(): # label='result',
+                        output_video_1 =  gr.Video(autoplay=True, width=256, height=256)
+                with gr.Row():
+                    gr.Examples(examples=t2v_examples,
+                                inputs=[input_text,duration],
+                                outputs=[output_video_1],
+                                fn=text2video.get_prompt,
+                                cache_examples=False)
+                            #cache_examples=os.getenv('SYSTEM') == 'spaces')
+            send_btn.click(
+                fn=text2video.get_prompt, 
+                inputs=[input_text,duration],
+                outputs=[output_video_1],
+            )
+            input_text.submit(
+                fn=text2video.get_prompt, 
+                inputs=[input_text,duration],
+                outputs=[output_video_1],
+            )
+
+    return videocrafter_iface
+
+if __name__ == "__main__":
+    result_dir = os.path.join('./', 'results')
+    videocrafter_iface = videocrafter_demo(result_dir)
+    videocrafter_iface.queue() # concurrency_count=1, max_size=10
+    videocrafter_iface.launch()
+    # videocrafter_iface.launch(server_name='0.0.0.0', server_port=80)

collect_data.py

@@ -0,0 +1,326 @@
+import warnings
+
+warnings.filterwarnings('ignore', category=DeprecationWarning)
+
+import os
+
+os.environ['MKL_SERVICE_FORCE_INTEL'] = '1'
+
+from pathlib import Path
+
+import hydra
+import numpy as np
+import torch
+import wandb
+from dm_env import specs
+
+import tools.utils as utils
+from tools.logger import Logger
+from tools.replay import ReplayBuffer, make_replay_loader
+
+torch.backends.cudnn.benchmark = True
+
+# os.environ['WANDB_API_KEY'] = 'local-1b6c1e2a2fd8d4c98b8c049eb2914dbceccd4b7c' # local-1b6c1e2a2fd8d4c98b8c049eb2914dbceccd4b7c
+# os.environ['WANDB_BASE_URL'] = 'https://192.168.170.90:443'
+# os.environ['REQUESTS_CA_BUNDLE'] = '/etc/ssl/certs/ca-certificates.crt'
+
+def make_agent(obs_type, obs_spec, action_spec, num_expl_steps, cfg):
+    cfg.obs_type = obs_type
+    cfg.obs_shape = obs_spec.shape
+    cfg.action_shape = action_spec.shape
+    cfg.num_expl_steps = num_expl_steps
+    return hydra.utils.instantiate(cfg)
+
+
+def make_dreamer_agent(obs_space, action_spec, cur_config, cfg):
+    from copy import deepcopy
+    cur_config = deepcopy(cur_config)
+    del cur_config.agent
+    return hydra.utils.instantiate(cfg, cfg=cur_config, obs_space=obs_space, act_spec=action_spec)
+
+class Workspace:
+    def __init__(self, cfg, savedir=None, workdir=None):
+        self.workdir = Path.cwd() if workdir is None else workdir
+        print(f'workspace: {self.workdir}')
+        self.cfg = cfg
+
+        utils.set_seed_everywhere(cfg.seed)
+        self.device = torch.device(cfg.device)
+
+        # create logger
+        self.logger = Logger(self.workdir,
+                             use_tb=cfg.use_tb,
+                             use_wandb=cfg.use_wandb)
+        # create envs
+        self.task = task = cfg.task
+        img_size = cfg.img_size
+
+        import envs.main as envs
+        self.train_env = envs.make(task, cfg.obs_type, cfg.action_repeat, cfg.seed, img_size=img_size,  viclip_encode=cfg.viclip_encode, clip_hd_rendering=cfg.clip_hd_rendering) 
+
+        # # create agent 
+        self.agent = make_dreamer_agent(self.train_env.obs_space, self.train_env.act_space['action'], cfg, cfg.agent)
+        
+        # get meta specs
+        meta_specs = self.agent.get_meta_specs()
+        # create replay buffer
+        data_specs = (self.train_env.obs_space,
+                      self.train_env.act_space,
+                      specs.Array((1,), np.float32, 'reward'),
+                      specs.Array((1,), np.float32, 'discount'))
+
+        # create data storage
+        self.replay_storage = ReplayBuffer(data_specs, meta_specs,
+                                                  self.workdir / 'buffer',
+                                                  length=cfg.batch_length, **cfg.replay,
+                                                  device=cfg.device)
+
+        # create replay buffer
+        self.replay_loader = make_replay_loader(self.replay_storage,
+                                                cfg.batch_size,)
+        self._replay_iter = None
+
+        self.timer = utils.Timer()
+        self._global_step = 0
+        self._global_episode = 0
+
+    @property
+    def global_step(self):
+        return self._global_step
+
+    @property
+    def global_episode(self):
+        return self._global_episode
+
+    @property
+    def global_frame(self):
+        return self.global_step * self.cfg.action_repeat
+
+    @property
+    def replay_iter(self):
+        if self._replay_iter is None:
+            self._replay_iter = iter(self.replay_loader)
+        return self._replay_iter
+
+    def eval(self):
+        import envs.main as envs
+        eval_env = envs.make(self.task, self.cfg.obs_type, self.cfg.action_repeat, self.cfg.seed, img_size=64,)
+        step, episode, total_reward = 0, 0, 0
+        eval_until_episode = utils.Until(self.cfg.num_eval_episodes)
+        meta = self.agent.init_meta()
+        while eval_until_episode(episode):
+            time_step, dreamer_obs = eval_env.reset()
+            agent_state = None
+            while not time_step.last():
+                with torch.no_grad(), utils.eval_mode(self.agent):
+                    action, agent_state = self.agent.act(dreamer_obs, 
+                                                meta,
+                                                self.global_step,
+                                                eval_mode=True,
+                                                state=agent_state)
+                time_step, dreamer_obs = eval_env.step(action)
+                total_reward += time_step.reward
+                step += 1
+
+            episode += 1
+
+        with self.logger.log_and_dump_ctx(self.global_frame, ty='eval') as log:
+            log('episode_reward', total_reward / episode)
+            log('episode_length', step * self.cfg.action_repeat / episode)
+            log('episode', self.global_episode)
+            log('step', self.global_step)
+
+    def eval_imag_behavior(self,):
+        self.agent._backup_acting_behavior = self.agent._acting_behavior
+        self.agent._acting_behavior = self.agent._imag_behavior
+        self.eval()
+        self.agent._acting_behavior = self.agent._backup_acting_behavior
+
+    def train(self):
+        # predicates
+        train_until_step = utils.Until(self.cfg.num_train_frames, self.cfg.action_repeat)
+        seed_until_step = utils.Until(self.cfg.num_seed_frames, self.cfg.action_repeat)
+        eval_every_step = utils.Every(self.cfg.eval_every_frames, self.cfg.action_repeat)
+        train_every_n_steps = max(self.cfg.train_every_actions // self.cfg.action_repeat, 1) 
+        should_train_step = utils.Every(train_every_n_steps * self.cfg.action_repeat, self.cfg.action_repeat)
+        should_log_scalars = utils.Every(self.cfg.log_every_frames, self.cfg.action_repeat)
+        should_log_visual = utils.Every(self.cfg.visual_every_frames, self.cfg.action_repeat)
+        should_save_model = utils.Every(self.cfg.save_every_frames, self.cfg.action_repeat)
+
+        episode_step, episode_reward = 0, 0
+        time_step, dreamer_obs = self.train_env.reset()
+        agent_state = None
+        meta = self.agent.init_meta()
+        data = dreamer_obs
+        self.replay_storage.add(data, meta) 
+        metrics = None
+        while train_until_step(self.global_step):
+            if time_step.last():
+                self._global_episode += 1
+                # wait until all the metrics schema is populated
+                if metrics is not None:
+                    # log stats
+                    elapsed_time, total_time = self.timer.reset()
+                    episode_frame = episode_step * self.cfg.action_repeat
+                    with self.logger.log_and_dump_ctx(self.global_frame,
+                                                    ty='train') as log:
+                        log('fps', episode_frame / elapsed_time)
+                        log('total_time', total_time)
+                        log('episode_reward', episode_reward)
+                        log('episode_length', episode_frame)
+                        log('episode', self.global_episode)
+                        log('buffer_size', len(self.replay_storage)) 
+                        log('step', self.global_step)
+                if should_save_model(self.global_step):
+                    # save last model
+                    self.save_last_model()
+
+                # reset env
+                time_step, dreamer_obs = self.train_env.reset()
+                # Updating agent
+                agent_state = None # Resetting agent's latent state
+                meta = self.agent.init_meta()
+                data = dreamer_obs
+                self.replay_storage.add(data, meta) 
+                episode_step = 0
+                episode_reward = 0
+
+            # try to evaluate
+            if eval_every_step(self.global_step):
+                if self.cfg.eval_modality == 'task':
+                    self.eval()
+                if self.cfg.eval_modality == 'task_imag':
+                    self.eval_imag_behavior()
+                if self.cfg.eval_modality == 'from_text':
+                    self.logger.log('eval_total_time', self.timer.total_time(),
+                                    self.global_frame)
+                    self.eval_from_text()
+
+            meta = self.agent.update_meta(meta, self.global_step, time_step)
+            # sample action
+            with torch.no_grad(), utils.eval_mode(self.agent):
+                if  seed_until_step(self.global_step):
+                    action =  self.train_env.act_space['action'].sample()
+                    if getattr(self.cfg, 'discrete_actions', False):
+                        action = (action == np.max(action)).astype(np.float32) # one-hot
+                else:
+                    action, agent_state = self.agent.act(dreamer_obs, # time_step.observation
+                                            meta,
+                                            self.global_step,
+                                            eval_mode=False,
+                                            state=agent_state)
+
+            # try to update the agent
+            if not seed_until_step(self.global_step):
+                if should_train_step(self.global_step):
+                    # prof.step()
+                    # Sampling data
+                    batch_data = next(self.replay_iter)
+                    if hasattr(self.agent, ' update_wm'):
+                        state, outputs, metrics = self.agent.update_wm(batch_data, self.global_step)
+                        if hasattr(self.agent, "update_acting_behavior"):
+                            metrics = self.agent.update_acting_behavior(state=state, outputs=outputs, metrics=metrics, data=batch_data)[1]
+                        if hasattr(self.agent, "update_imag_behavior"):
+                            metrics.update(self.agent.update_imag_behavior(state=state, outputs=outputs, metrics=metrics, seq_data=batch_data,)[1])
+                    else:
+                        outputs, metrics = self.agent.update(batch_data, self.global_step)
+
+                if should_log_scalars(self.global_step):
+                    self.logger.log_metrics(metrics, self.global_frame, ty='train')
+                if self.global_step > 0 and should_log_visual(self.global_step):
+                    if hasattr(self.agent, 'report'):
+                        with torch.no_grad(), utils.eval_mode(self.agent):
+                            videos = self.agent.report(next(self.replay_iter))
+                            self.logger.log_visual(videos, self.global_frame)
+
+            # take env step
+            time_step, dreamer_obs = self.train_env.step(action)
+            episode_reward += time_step.reward
+            data = dreamer_obs
+            if time_step.last():
+                if getattr(self.train_env, "accumulate", False):
+                    assert not self.replay_storage._ongoing
+                    # NOTE: this is ok as it comes right after adding to the repl
+                    accumulated_data, accumulated_key = self.train_env.process_accumulate()
+                    data[accumulated_key] = accumulated_data[-1]
+                    self.replay_storage._ongoing_eps[0][accumulated_key][-len(accumulated_data[:-1]):] = accumulated_data[:-1]
+            self.replay_storage.add(data, meta) 
+            episode_step += 1
+            self._global_step += 1
+
+    @utils.retry
+    def save_snapshot(self):
+        snapshot = self.get_snapshot_dir() / f'snapshot_{self.global_frame}.pt'
+        keys_to_save = ['agent', '_global_step', '_global_episode']
+        payload = {k: self.__dict__[k] for k in keys_to_save}
+        with snapshot.open('wb') as f:
+            torch.save(payload, f)
+
+    def setup_wandb(self):
+        cfg = self.cfg
+        exp_name = '_'.join([
+            cfg.experiment, cfg.agent.name, cfg.task, cfg.obs_type,
+            str(cfg.seed)
+        ])
+        wandb.init(project=cfg.project_name, group=cfg.agent.name, name=exp_name)
+        flat_cfg = utils.flatten_dict(cfg)
+        wandb.config.update(flat_cfg)
+        self.wandb_run_id = wandb.run.id
+
+    @utils.retry
+    def save_last_model(self):
+        snapshot = self.root_dir / 'last_snapshot.pt'
+        if snapshot.is_file():
+            temp = Path(str(snapshot).replace("last_snapshot.pt", "second_last_snapshot.pt"))
+            os.replace(snapshot, temp)
+        keys_to_save = ['agent', '_global_step', '_global_episode']
+        if self.cfg.use_wandb: 
+            keys_to_save.append('wandb_run_id')
+        payload = {k: self.__dict__[k] for k in keys_to_save}
+        with snapshot.open('wb') as f:
+            torch.save(payload, f)
+
+    def load_snapshot(self, snapshot_dir):
+        try:
+            snapshot = snapshot_dir / 'last_snapshot.pt'
+            with snapshot.open('rb') as f:
+                payload = torch.load(f)
+        except:
+            snapshot = snapshot_dir / 'second_last_snapshot.pt'
+            with snapshot.open('rb') as f:
+                payload = torch.load(f)
+        for k,v in payload.items():
+            setattr(self, k, v)
+            if k == 'wandb_run_id':
+                assert wandb.run is None
+                cfg = self.cfg
+                exp_name = '_'.join([
+                    cfg.experiment, cfg.agent.name, cfg.task, cfg.obs_type,
+                    str(cfg.seed)
+                ])
+                wandb.init(project=cfg.project_name, group=cfg.agent.name, name=exp_name, id=v, resume="must")
+
+    def get_snapshot_dir(self):
+        snap_dir = self.cfg.snapshot_dir
+        snapshot_dir = self.workdir / Path(snap_dir)
+        snapshot_dir.mkdir(exist_ok=True, parents=True)
+        return snapshot_dir 
+
+@hydra.main(config_path='.', config_name='collect_data')
+def main(cfg):
+    from collect_data import Workspace as W
+    root_dir = Path.cwd()
+    cfg.workdir = str(root_dir)
+    workspace = W(cfg)
+    workspace.root_dir = root_dir
+    snapshot = workspace.root_dir / 'last_snapshot.pt'
+    if snapshot.exists():
+        print(f'resuming: {snapshot}')
+        workspace.load_snapshot(workspace.root_dir)
+    if cfg.use_wandb and wandb.run is None:
+        # otherwise it was resumed
+        workspace.setup_wandb()   
+    workspace.train()
+
+if __name__ == '__main__':
+    main()

collect_data.yaml

@@ -0,0 +1,54 @@
+defaults:
+  - agent: dreamer
+  - conf/env: dmc_pixels
+  - conf/defaults: dreamer_v3
+  - override hydra/launcher: submitit_local
+
+# mode
+label: default
+# task settings
+task: stickman_walk
+# train settings
+num_train_frames: 2000010
+num_seed_frames: 4000
+# eval
+eval_every_frames: 100000
+eval_modality: null
+num_eval_episodes: 3
+# snapshot
+snapshot_dir: ../../../trained_models/${obs_type}/${task}/${agent.name}/${seed}
+save_every_frames: 10_000
+# misc
+seed: 1
+device: cuda:0
+use_tb: true
+use_wandb: true
+
+# Clip stuff
+viclip_encode: false
+viclip_model: internvideo2
+clip_hd_rendering: false
+
+# experiment
+experiment: data
+project_name: genrl
+
+# log settings
+log_every_frames: 1000
+visual_every_frames: 100000000 # edit for debug
+workdir: ???
+
+hydra:
+  run:
+    dir: ./exp_local/${now:%Y.%m.%d}/${now:%H%M%S}_${agent.name}
+  sweep:
+    dir: ./exp_sweep/${now:%Y.%m.%d}/${now:%H%M}_${agent.name}_${experiment}
+    subdir: ${hydra.job.num}
+  launcher:
+    timeout_min: 4300
+    cpus_per_task: 10
+    gpus_per_node: 1
+    tasks_per_node: 1
+    mem_gb: 160
+    nodes: 1
+    submitit_folder: ./exp_sweep/${now:%Y.%m.%d}/${now:%H%M}_${agent.name}_${experiment}/.slurm

conf/defaults/dreamer_v2.yaml

@@ -0,0 +1,38 @@
+# @package _global_
+img_size: 64
+
+# Dreamer defaults
+rssm: {ensemble: 1, hidden: 512, deter: 512, stoch: 32, discrete: 32,  norm: none, std_act: softplus, min_std: 0.1, single_obs_posterior: false, } # act: elu,
+discount_head: {layers: 4, units: 512, norm: none, dist: binary} # act: elu
+reward_head: {layers: 4, units: 512, norm: none, dist: mse} # act: elu
+kl: {free: 1.0, forward: False, balance: 0.8, free_avg: False, }
+loss_scales: {kl: 1.0, reward: 1.0, discount: 1.0, proprio: 1.0}
+model_opt: {opt: adam, lr: 3e-4, eps: 1e-5, clip: 1000, wd: 1e-6}
+replay: {capacity: 2e6, ongoing: False, minlen: 50, maxlen: 50, prioritize_ends: False}
+decoder_inputs: feat
+image_dist: normal_unit_std
+
+actor: {layers: 4, units: 512, norm: none, dist: trunc_normal, min_std: 0.1 } # act: elu 
+critic: {layers: 4, units: 512, norm: none, dist: mse} # act: elu, 
+actor_opt: {opt: adam, lr: 8e-5, eps: 1e-5, clip: 100, wd: 1e-6}
+critic_opt: {opt: adam, lr: 8e-5, eps: 1e-5, clip: 100, wd: 1e-6}
+discount: 0.99
+discount_lambda: 0.95
+slow_target: True
+slow_target_update: 100
+slow_target_fraction: 1
+slow_baseline: True
+reward_ema: False
+
+acting_reward_fn: env_reward
+clip_rewards: identity
+
+batch_size: 50 
+batch_length: 50 
+imag_horizon: 15
+eval_state_mean: False
+
+precision: 16
+train_every_actions: 10
+only_random_actions: False
+# 

conf/defaults/dreamer_v3.yaml

@@ -0,0 +1,38 @@
+# @package _global_
+img_size: 64
+
+# Dreamer defaults
+rssm: {ensemble: 1, hidden: 512, deter: 512, stoch: 32, discrete: 32,  norm: layer, std_act: softplus, min_std: 0.1,  single_obs_posterior: false, } # act: elu,
+discount_head: {layers: 4, units: 512, norm: layer, dist: binary} # act: elu
+reward_head: {layers: 4, units: 512, norm: layer, dist: twohot} # act: elu
+kl: { free: 1.0, forward: False, balance: 0.85, free_avg: False,}
+loss_scales: {kl: 0.6, reward: 1.0, discount: 1.0, proprio: 1.0}
+model_opt: {opt: adam, lr: 1e-4, eps: 1e-8, clip: 1000, wd: 1e-6}
+replay: {capacity: 2e6, ongoing: False, minlen: 50, maxlen: 50, prioritize_ends: False}
+decoder_inputs: feat
+image_dist: mse
+# Actor Critic
+actor: {layers: 4, units: 512, norm: layer, dist: normal, min_std: 0.1 } # act: elu 
+critic: {layers: 4, units: 512, norm: layer, dist: twohot } # act: elu, 
+actor_opt: {opt: adam, lr: 3e-5, eps: 1e-5, clip: 100, wd: 1e-6}
+critic_opt: {opt: adam, lr: 3e-5, eps: 1e-5, clip: 100, wd: 1e-6}
+discount: 0.99
+discount_lambda: 0.95
+slow_target: True
+slow_target_update: 100
+slow_target_fraction: 1
+slow_baseline: True
+reward_ema: True
+
+acting_reward_fn: env_reward
+clip_rewards: identity
+
+batch_size: 50 
+batch_length: 50 
+imag_horizon: 15
+eval_state_mean: False
+
+precision: 16
+train_every_actions: 10
+only_random_actions: False
+# 

conf/defaults/genrl.yaml

@@ -0,0 +1,37 @@
+# @package _global_
+img_size: 64
+
+# Dreamer defaults
+rssm: {ensemble: 1, hidden: 1024, deter: 1024, stoch: 32, discrete: 32,  norm: layer, std_act: softplus, min_std: 0.1,  single_obs_posterior: true, } # act: elu,
+discount_head: {layers: 4, units: 512, norm: none, dist: binary} # act: elu
+reward_head: {layers: 4, units: 1024, norm: layer, dist: twohot} # act: elu
+kl: { free: 1.0, forward: False, balance: 0.85, free_avg: False, }
+loss_scales: {kl: 0.6, reward: 1.0, discount: 1.0, proprio: 1.0}
+model_opt: {opt: adam, lr: 1e-4, eps: 1e-8, clip: 1000, wd: 1e-6}
+replay: {capacity: 20e6, ongoing: False, minlen: 48, maxlen: 48, prioritize_ends: False}
+decoder_inputs: stoch
+image_dist: mse
+# Actor Critic
+actor: {layers: 4, units: 1024, norm: layer, dist: normal, min_std: 0.1 } # act: elu 
+critic: {layers: 4, units: 1024, norm: layer, dist: twohot } # act: elu, 
+actor_opt: {opt: adam, lr: 3e-5, eps: 1e-5, clip: 100, wd: 1e-6}
+critic_opt: {opt: adam, lr: 3e-5, eps: 1e-5, clip: 100, wd: 1e-6}
+discount: 0.99
+discount_lambda: 0.95
+slow_target: True
+slow_target_update: 100
+slow_target_fraction: 1
+slow_baseline: True
+reward_ema: True
+
+acting_reward_fn: env_reward
+clip_rewards: identity
+
+batch_size: 48 
+batch_length: 48 
+imag_horizon: 16
+eval_state_mean: False
+
+precision: 16
+train_every_actions: 10
+only_random_actions: False

conf/env/dmc_pixels.yaml

@@ -0,0 +1,8 @@
+# @package _global_
+obs_type: pixels
+action_repeat: 2 
+encoder: {mlp_keys: '$^', cnn_keys: 'observation', norm: layer, cnn_depth: 48, cnn_kernels: [4, 4, 4, 4], mlp_layers: [400, 400, 400, 400]} # act: elu
+decoder: {mlp_keys: '$^', cnn_keys: 'observation', norm: layer, cnn_depth: 48, cnn_kernels: [5, 5, 6, 6], mlp_layers: [400, 400, 400, 400], } # act: elu
+pred_discount: False
+imag_actor_grad: dynamics
+actor_grad: dynamics

conf/train_mode/train_behavior.yaml

@@ -0,0 +1,5 @@
+num_train_frames: 500_010
+batch_size: 32
+batch_length: 32
+agent.imag_reward_fn: video_text_reward
+eval_modality: task_imag

conf/train_mode/train_model.yaml

@@ -0,0 +1,6 @@
+num_train_frames: 5_000_010
+visual_every_frames: 250_000
+train_world_model: True
+train_connector: True
+reset_world_model: True
+reset_connector: True

demo/demo_test.py

@@ -0,0 +1,23 @@
+import os
+from pathlib import Path
+VIDEO_PATH = Path(os.path.abspath('')) / 'assets' / 'video_samples'
+
+class Text2Video():
+    def __init__(self, result_dir='./tmp/') -> None:
+        pass
+
+    def get_prompt(self, input_text, steps=50, cfg_scale=15.0, eta=1.0, fps=16):
+
+        return str(VIDEO_PATH / 'headstand.mp4')
+    
+class Video2Video:
+    def __init__(self, result_dir='./tmp/') -> None:
+        pass
+
+    def get_image(self, input_image, input_prompt, i2v_steps=50, i2v_cfg_scale=15.0, i2v_eta=1.0, i2v_fps=16):
+
+        return str(VIDEO_PATH / 'dancing.mp4')
+        
+if __name__ == '__main__':
+    t2v = Text2Video()
+    print(t2v.get_prompt('test'))

demo/t2v.py

@@ -0,0 +1,115 @@
+from pathlib import Path 
+import os
+import sys
+sys.path.append(str(Path(os.path.abspath(''))))
+
+import torch
+import numpy as np
+from tools.genrl_utils import ViCLIPGlobalInstance
+
+import time
+import torchvision
+from huggingface_hub import hf_hub_download
+
+def save_videos(batch_tensors, savedir, filenames, fps=10):
+    # b,samples,c,t,h,w
+    n_samples = batch_tensors.shape[1]
+    for idx, vid_tensor in enumerate(batch_tensors):
+        video = vid_tensor.detach().cpu()
+        video = torch.clamp(video.float(), 0., 1.)
+        video = video.permute(1, 0, 2, 3, 4) # t,n,c,h,w
+        frame_grids = [torchvision.utils.make_grid(framesheet, nrow=int(n_samples)) for framesheet in video] #[3, 1*h, n*w]
+        grid = torch.stack(frame_grids, dim=0) # stack in temporal dim [t, 3, n*h, w]
+        grid = (grid * 255).to(torch.uint8).permute(0, 2, 3, 1)
+        savepath = os.path.join(savedir, f"{filenames[idx]}.mp4")
+        torchvision.io.write_video(savepath, grid, fps=fps, video_codec='h264', options={'crf': '10'})
+
+class Text2Video():
+    def __init__(self,result_dir='./tmp/',gpu_num=1) -> None:
+        model_folder = str(Path(os.path.abspath('')) / 'models')
+        model_filename = 'genrl_stickman_500k_2.pt'
+        
+        if not os.path.isfile(os.path.join(model_folder, model_filename)):
+            self.download_model(model_folder, model_filename)
+        if not os.path.isfile(os.path.join(model_folder, 'InternVideo2-stage2_1b-224p-f4.pt')):
+            self.download_internvideo2(model_folder)
+        self.agent = torch.load(os.path.join(model_folder, model_filename))
+        model_name = 'internvideo2'
+
+        # Get ViCLIP
+        viclip_global_instance = ViCLIPGlobalInstance(model_name)
+        if not viclip_global_instance._instantiated:
+            print("Instantiating InternVideo2")
+            viclip_global_instance.instantiate()
+        self.clip = viclip_global_instance.viclip
+        self.tokenizer = viclip_global_instance.viclip_tokenizer
+
+        self.result_dir = result_dir
+        if not os.path.exists(self.result_dir):
+            os.mkdir(self.result_dir)
+
+    def get_prompt(self, prompt, duration):
+        torch.cuda.empty_cache()
+        print('start:', prompt, time.strftime('%Y-%m-%d %H:%M:%S',time.localtime(time.time())))
+        start = time.time()
+
+        prompt_str = prompt.replace("/", "_slash_") if "/" in prompt else prompt
+        prompt_str = prompt_str.replace(" ", "_") if " " in prompt else prompt_str
+
+        labels_list = [prompt_str]
+        with torch.no_grad():
+            wm = world_model = self.agent.wm
+            connector = self.agent.wm.connector
+            decoder = world_model.heads['decoder']
+            n_frames = connector.n_frames
+            
+            # Get text(video) embed
+            text_feat = []
+            for text in labels_list:
+                with torch.no_grad():
+                    text_feat.append(self.clip.get_txt_feat(text,))
+            text_feat = torch.stack(text_feat, dim=0).to(self.clip.device)
+
+            video_embed = text_feat
+
+            B = video_embed.shape[0]
+            T = 1
+
+            # Get actions
+            video_embed = video_embed.repeat(1, duration, 1)
+            with torch.no_grad():
+                # Imagine
+                prior = wm.connector.video_imagine(video_embed, None, sample=False, reset_every_n_frames=False, denoise=True)
+                # Decode
+                prior_recon = decoder(wm.decoder_input_fn(prior))['observation'].mean + 0.5
+
+        save_videos(prior_recon.unsqueeze(0), self.result_dir, filenames=[prompt_str], fps=15)
+        print(f"Saved in {prompt_str}.mp4. Time used: {(time.time() - start):.2f} seconds")
+        return os.path.join(self.result_dir, f"{prompt_str}.mp4")
+    
+    def download_model(self, model_folder, model_filename):
+        REPO_ID = 'mazpie/genrl_models'
+        filename_list = [model_filename]
+        if not os.path.exists(model_folder):
+            os.makedirs(model_folder)
+        for filename in filename_list:
+            local_file = os.path.join(model_folder, filename)
+
+            if not os.path.exists(local_file):
+                hf_hub_download(repo_id=REPO_ID, filename=filename, local_dir=model_folder, local_dir_use_symlinks=False)
+    
+    def download_internvideo2(self, model_folder):
+        REPO_ID = 'OpenGVLab/InternVideo2-Stage2_1B-224p-f4'
+        filename_list = ['InternVideo2-stage2_1b-224p-f4.pt']
+        if not os.path.exists(model_folder):
+            os.makedirs(model_folder)
+        for filename in filename_list:
+            local_file = os.path.join(model_folder, filename)
+
+            if not os.path.exists(local_file):
+                hf_hub_download(repo_id=REPO_ID, filename=filename, local_dir=model_folder, local_dir_use_symlinks=False)
+
+if __name__ == '__main__':
+    t2v = Text2Video()
+    video_path = t2v.get_prompt('a black swan swims on the pond', 8)
+    print('done', video_path)

envs/custom_dmc_tasks/__init__.py

@@ -0,0 +1,13 @@
+from . import cheetah
+from . import walker
+from . import quadruped
+from . import jaco
+from . import stickman
+from dm_control import suite
+
+suite._DOMAINS['stickman'] = stickman
+suite.ALL_TASKS = suite.ALL_TASKS + suite._get_tasks('custom')
+suite.TASKS_BY_DOMAIN = suite._get_tasks_by_domain(suite.ALL_TASKS)
+    
+def make_jaco(task, obs_type, seed, img_size, ):
+    return jaco.make(task, obs_type, seed, img_size, )

envs/custom_dmc_tasks/cheetah.py

@@ -0,0 +1,247 @@
+# Copyright 2017 The dm_control Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Cheetah Domain."""
+
+import collections
+import os
+
+from dm_control.suite import cheetah
+from dm_control import mujoco
+from dm_control.rl import control
+from dm_control.suite import base
+from dm_control.suite import common
+from dm_control.utils import containers
+from dm_control.utils import rewards
+from dm_control.utils import io as resources
+
+# How long the simulation will run, in seconds.
+_DEFAULT_TIME_LIMIT = 10
+
+_DOWN_HEIGHT = 0.15
+_HIGH_HEIGHT = 1.00
+_MID_HEIGHT = 0.45
+
+
+# Running speed above which reward is 1.
+_RUN_SPEED = 10
+_SPIN_SPEED = 5
+
+def make(task,
+         task_kwargs=None,
+         environment_kwargs=None,
+         visualize_reward=False):
+    task_kwargs = task_kwargs or {}
+    if environment_kwargs is not None:
+        task_kwargs = task_kwargs.copy()
+        task_kwargs['environment_kwargs'] = environment_kwargs
+    env = SUITE[task](**task_kwargs)
+    env.task.visualize_reward = visualize_reward
+    return env
+
+
+def get_model_and_assets():
+    """Returns a tuple containing the model XML string and a dict of assets."""
+    root_dir = os.path.dirname(os.path.dirname(__file__))
+    xml = resources.GetResource(
+        os.path.join(root_dir, 'custom_dmc_tasks', 'cheetah.xml'))
+    return xml, common.ASSETS
+
+
+@cheetah.SUITE.add('custom')
+def flipping(time_limit=_DEFAULT_TIME_LIMIT,
+                 random=None,
+                 environment_kwargs=None):
+    """Returns the run task."""
+    physics = Physics.from_xml_string(*get_model_and_assets())
+    task = Cheetah(forward=False, flip=False, random=random, goal='flipping')
+    environment_kwargs = environment_kwargs or {}
+    return control.Environment(physics,
+                               task,
+                               time_limit=time_limit,
+                               **environment_kwargs)
+
+@cheetah.SUITE.add('custom')
+def standing(time_limit=_DEFAULT_TIME_LIMIT,
+                 random=None,
+                 environment_kwargs=None):
+    """Returns the run task."""
+    physics = Physics.from_xml_string(*get_model_and_assets())
+    task = Cheetah(forward=False, flip=False, random=random, goal='standing')
+    environment_kwargs = environment_kwargs or {}
+    return control.Environment(physics,
+                               task,
+                               time_limit=time_limit,
+                               **environment_kwargs)
+
+
+@cheetah.SUITE.add('custom')
+def lying_down(time_limit=_DEFAULT_TIME_LIMIT,
+                 random=None,
+                 environment_kwargs=None):
+    """Returns the run task."""
+    physics = Physics.from_xml_string(*get_model_and_assets())
+    task = Cheetah(forward=False, flip=False, random=random, goal='lying_down')
+    environment_kwargs = environment_kwargs or {}
+    return control.Environment(physics,
+                               task,
+                               time_limit=time_limit,
+                               **environment_kwargs)
+
+
+@cheetah.SUITE.add('custom')
+def run_backward(time_limit=_DEFAULT_TIME_LIMIT,
+                 random=None,
+                 environment_kwargs=None):
+    """Returns the run task."""
+    physics = Physics.from_xml_string(*get_model_and_assets())
+    task = Cheetah(forward=False, flip=False, random=random, goal='run_backward')
+    environment_kwargs = environment_kwargs or {}
+    return control.Environment(physics,
+                               task,
+                               time_limit=time_limit,
+                               **environment_kwargs)
+
+
+@cheetah.SUITE.add('custom')
+def flip(time_limit=_DEFAULT_TIME_LIMIT,
+                 random=None,
+                 environment_kwargs=None):
+    """Returns the run task."""
+    physics = Physics.from_xml_string(*get_model_and_assets())
+    task = Cheetah(forward=True, flip=True, random=random, goal='flip')
+    environment_kwargs = environment_kwargs or {}
+    return control.Environment(physics,
+                               task,
+                               time_limit=time_limit,
+                               **environment_kwargs)
+
+
+@cheetah.SUITE.add('custom')
+def flip_backward(time_limit=_DEFAULT_TIME_LIMIT,
+                  random=None,
+                  environment_kwargs=None):
+    """Returns the run task."""
+    physics = Physics.from_xml_string(*get_model_and_assets())
+    task = Cheetah(forward=False, flip=True, random=random, goal='flip_backward')
+    environment_kwargs = environment_kwargs or {}
+    return control.Environment(physics,
+                               task,
+                               time_limit=time_limit,
+                               **environment_kwargs)
+
+
+class Physics(mujoco.Physics):
+    """Physics simulation with additional features for the Cheetah domain."""
+    def speed(self):
+        """Returns the horizontal speed of the Cheetah."""
+        return self.named.data.sensordata['torso_subtreelinvel'][0]
+
+    def angmomentum(self):
+        """Returns the angular momentum of torso of the Cheetah about Y axis."""
+        return self.named.data.subtree_angmom['torso'][1]
+
+
+class Cheetah(base.Task):
+    """A `Task` to train a running Cheetah."""
+    def __init__(self, goal=None, forward=True, flip=False, random=None):
+        self._forward = 1 if forward else -1
+        self._flip = flip
+        self._goal = goal
+        super(Cheetah, self).__init__(random=random)
+
+    def initialize_episode(self, physics):
+        """Sets the state of the environment at the start of each episode."""
+        # The indexing below assumes that all joints have a single DOF.
+        assert physics.model.nq == physics.model.njnt
+        is_limited = physics.model.jnt_limited == 1
+        lower, upper = physics.model.jnt_range[is_limited].T
+        physics.data.qpos[is_limited] = self.random.uniform(lower, upper)
+
+        # Stabilize the model before the actual simulation.
+        for _ in range(200):
+            physics.step()
+
+        physics.data.time = 0
+        self._timeout_progress = 0
+        super().initialize_episode(physics)
+
+    def _get_lying_down_reward(self, physics):
+        torso = physics.named.data.xpos['torso', 'z']
+        
+        torso_down = rewards.tolerance(torso,
+                                        bounds=(-float('inf'), _DOWN_HEIGHT),
+                                        margin=_DOWN_HEIGHT * 1.5,)
+
+        feet = physics.named.data.xpos['bfoot', 'z'] + physics.named.data.xpos['ffoot', 'z']
+
+        feet_up = rewards.tolerance(feet,
+                                        bounds=(_MID_HEIGHT, float('inf')),
+                                        margin=_MID_HEIGHT / 2,)
+        return (torso_down + feet_up) / 2
+
+    def _get_standing_reward(self, physics):
+        bfoot = physics.named.data.xpos['bfoot', 'z']
+        ffoot = physics.named.data.xpos['ffoot', 'z']
+        max_foot = bfoot if bfoot > ffoot else ffoot
+        min_foot = bfoot if bfoot <= ffoot else ffoot
+        
+        low_foot_low = rewards.tolerance(min_foot,
+                                        bounds=(-float('inf'), _DOWN_HEIGHT),
+                                        margin=_DOWN_HEIGHT * 1.5,)
+
+        high_foot_high = rewards.tolerance(max_foot,
+                                        bounds=(_HIGH_HEIGHT, float('inf')),
+                                        margin=_HIGH_HEIGHT / 2,)
+        return high_foot_high * low_foot_low
+
+    def _get_flip_reward(self, physics):
+        return rewards.tolerance(self._forward * physics.angmomentum(),
+                                        bounds=(_SPIN_SPEED, float('inf')),
+                                        margin=_SPIN_SPEED,
+                                        value_at_margin=0,
+                                        sigmoid='linear')
+    
+    def get_observation(self, physics):
+        """Returns an observation of the state, ignoring horizontal position."""
+        obs = collections.OrderedDict()
+        # Ignores horizontal position to maintain translational invariance.
+        obs['position'] = physics.data.qpos[1:].copy()
+        obs['velocity'] = physics.velocity()
+        return obs
+
+    def get_reward(self, physics):
+        """Returns a reward to the agent."""
+        if self._goal in ['run', 'flip', 'run_backward', 'flip_backward']:
+            if self._flip:
+                return self._get_flip_reward(physics)
+            else:
+                reward = rewards.tolerance(self._forward * physics.speed(),
+                                        bounds=(_RUN_SPEED, float('inf')),
+                                        margin=_RUN_SPEED,
+                                        value_at_margin=0,
+                                        sigmoid='linear')
+                return reward
+        elif self._goal == 'lying_down':
+            return self._get_lying_down_reward(physics)
+        elif self._goal == 'flipping':
+            self._forward = True
+            fwd_reward = self._get_flip_reward(physics)
+            self._forward = False
+            back_reward = self._get_flip_reward(physics)
+            return max(fwd_reward, back_reward)
+        elif self._goal == 'standing':
+            return self._get_standing_reward(physics)
+        else:
+            raise NotImplementedError(self._goal)

envs/custom_dmc_tasks/cheetah.xml

@@ -0,0 +1,74 @@
+<mujoco model="cheetah">
+  <include file="./common/skybox.xml"/>
+  <include file="./common/visual.xml"/>
+  <include file="./common/materials.xml"/>
+
+
+  <compiler settotalmass="14"/>
+
+  <default>
+    <default class="cheetah">
+      <joint limited="true" damping=".01" armature=".1" stiffness="8" type="hinge" axis="0 1 0"/>
+      <geom contype="1" conaffinity="1" condim="3" friction=".4 .1 .1" material="self"/>
+    </default>
+    <default class="free">
+      <joint limited="false" damping="0" armature="0" stiffness="0"/>
+    </default>
+    <motor ctrllimited="true" ctrlrange="-1 1"/>
+  </default>
+
+  <statistic center="0 0 .7" extent="2"/>
+
+  <option timestep="0.01"/>
+
+  <worldbody>
+    <geom name="ground" type="plane" conaffinity="1" pos="98 0 0" size="200 .8 .5" material="grid"/>
+    <body name="torso" pos="0 0 .7" childclass="cheetah">
+      <light name="light" pos="0 0 2" mode="trackcom"/>
+      <camera name="side" pos="0 -3 0" quat="0.707 0.707 0 0" mode="trackcom"/>
+      <camera name="back" pos="-1.8 -1.3 0.8" xyaxes="0.45 -0.9 0 0.3 0.15 0.94" mode="trackcom"/>
+      <joint name="rootx" type="slide" axis="1 0 0" class="free"/>
+      <joint name="rootz" type="slide" axis="0 0 1" class="free"/>
+      <joint name="rooty" type="hinge" axis="0 1 0" class="free"/>
+      <geom name="torso" type="capsule" fromto="-.5 0 0 .5 0 0" size="0.046"/>
+      <geom name="head" type="capsule" pos=".6 0 .1" euler="0 50 0" size="0.046 .15"/>
+      <body name="bthigh" pos="-.5 0 0">
+        <joint name="bthigh" range="-30 60" stiffness="240" damping="6"/>
+        <geom name="bthigh" type="capsule" pos=".1 0 -.13" euler="0 -218 0" size="0.046 .145"/>
+        <body name="bshin" pos=".16 0 -.25">
+          <joint name="bshin" range="-50 50" stiffness="180" damping="4.5"/>
+          <geom name="bshin" type="capsule" pos="-.14 0 -.07" euler="0 -116 0" size="0.046 .15"/>
+          <body name="bfoot" pos="-.28 0 -.14">
+            <joint name="bfoot" range="-230 50" stiffness="120" damping="3"/>
+            <geom name="bfoot" type="capsule" pos=".03 0 -.097" euler="0 -15 0" size="0.046 .094"/>
+          </body>
+        </body>
+      </body>
+      <body name="fthigh" pos=".5 0 0">
+        <joint name="fthigh" range="-57 .40" stiffness="180" damping="4.5"/>
+        <geom name="fthigh" type="capsule" pos="-.07 0 -.12" euler="0 30 0" size="0.046 .133"/>
+        <body name="fshin" pos="-.14 0 -.24">
+          <joint name="fshin" range="-70 50" stiffness="120" damping="3"/>
+          <geom name="fshin" type="capsule" pos=".065 0 -.09" euler="0 -34 0" size="0.046 .106"/>
+          <body name="ffoot" pos=".13 0 -.18">
+            <joint name="ffoot" range="-28 28" stiffness="60" damping="1.5"/>
+            <geom name="ffoot" type="capsule" pos=".045 0 -.07" euler="0 -34 0" size="0.046 .07"/>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+
+  <sensor>
+    <subtreelinvel name="torso_subtreelinvel" body="torso"/>
+  </sensor>
+
+  <actuator>
+    <motor name="bthigh" joint="bthigh" gear="120" />
+    <motor name="bshin" joint="bshin" gear="90" />
+    <motor name="bfoot" joint="bfoot" gear="60" />
+    <motor name="fthigh" joint="fthigh" gear="90" />
+    <motor name="fshin" joint="fshin" gear="60" />
+    <motor name="ffoot" joint="ffoot" gear="30" />
+  </actuator>
+</mujoco>

envs/custom_dmc_tasks/jaco.py

@@ -0,0 +1,222 @@
+# Copyright 2019 The dm_control Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+"""A task where the goal is to move the hand close to a target prop or site."""
+
+import collections
+
+from dm_control import composer
+from dm_control.composer import initializers
+from dm_control.composer.observation import observable
+from dm_control.composer.variation import distributions
+from dm_control.entities import props
+from dm_control.manipulation.shared import arenas
+from dm_control.manipulation.shared import cameras
+from dm_control.manipulation.shared import constants
+from dm_control.manipulation.shared import observations
+from dm_control.manipulation.shared import registry
+from dm_control.manipulation.shared import robots
+from dm_control.manipulation.shared import tags
+from dm_control.manipulation.shared import workspaces
+from dm_control.utils import rewards
+import numpy as np
+
+
+_ReachWorkspace = collections.namedtuple(
+    '_ReachWorkspace', ['target_bbox', 'tcp_bbox', 'arm_offset'])
+
+# Ensures that the props are not touching the table before settling.
+_PROP_Z_OFFSET = 0.001
+
+_DUPLO_WORKSPACE = _ReachWorkspace(
+    target_bbox=workspaces.BoundingBox(
+        lower=(-0.1, -0.1, _PROP_Z_OFFSET),
+        upper=(0.1, 0.1, _PROP_Z_OFFSET)),
+    tcp_bbox=workspaces.BoundingBox(
+        lower=(-0.1, -0.1, 0.2),
+        upper=(0.1, 0.1, 0.4)),
+    arm_offset=robots.ARM_OFFSET)
+
+_SITE_WORKSPACE = _ReachWorkspace(
+    target_bbox=workspaces.BoundingBox(
+        lower=(-0.2, -0.2, 0.02),
+        upper=(0.2, 0.2, 0.4)),
+    tcp_bbox=workspaces.BoundingBox(
+        lower=(-0.2, -0.2, 0.02),
+        upper=(0.2, 0.2, 0.4)),
+    arm_offset=robots.ARM_OFFSET)
+
+_TARGET_RADIUS = 0.05
+_TIME_LIMIT = 10
+
+TASKS = {
+    'reach_top_left':  workspaces.BoundingBox(
+        lower=(-0.09, 0.09, _PROP_Z_OFFSET),
+        upper=(-0.09, 0.09, _PROP_Z_OFFSET)),
+    'reach_top_right': workspaces.BoundingBox(
+        lower=(0.09, 0.09, _PROP_Z_OFFSET),
+        upper=(0.09, 0.09, _PROP_Z_OFFSET)),
+    'reach_bottom_left': workspaces.BoundingBox(
+        lower=(-0.09, -0.09, _PROP_Z_OFFSET),
+        upper=(-0.09, -0.09, _PROP_Z_OFFSET)),
+    'reach_bottom_right': workspaces.BoundingBox(
+        lower=(0.09, -0.09, _PROP_Z_OFFSET),
+        upper=(0.09, -0.09, _PROP_Z_OFFSET)),
+}
+
+
+def make(task_id, obs_type, seed, img_size=64,):
+  obs_settings = observations.VISION if obs_type == 'pixels' else observations.PERFECT_FEATURES
+  obs_settings = obs_settings._replace(camera=obs_settings[-1]._replace(width=img_size))
+  obs_settings = obs_settings._replace(camera=obs_settings[-1]._replace(height=img_size))
+  if obs_type == 'states':
+    global _TIME_LIMIT
+    _TIME_LIMIT = 10.04
+    # Note: Adding this fixes the problem of having 249 steps with action repeat = 1
+  task = _reach(task_id, obs_settings=obs_settings, use_site=False)
+  return composer.Environment(task, time_limit=_TIME_LIMIT, random_state=seed)
+
+
+class MTReach(composer.Task):
+  """Bring the hand close to a target prop or site."""
+
+  def __init__(
+      self, task_id, arena, arm, hand, prop, obs_settings, workspace, control_timestep):
+    """Initializes a new `Reach` task.
+
+    Args:
+      arena: `composer.Entity` instance.
+      arm: `robot_base.RobotArm` instance.
+      hand: `robot_base.RobotHand` instance.
+      prop: `composer.Entity` instance specifying the prop to reach to, or None
+        in which case the target is a fixed site whose position is specified by
+        the workspace.
+      obs_settings: `observations.ObservationSettings` instance.
+      workspace: `_ReachWorkspace` specifying the placement of the prop and TCP.
+      control_timestep: Float specifying the control timestep in seconds.
+    """
+    self._task_id = task_id
+    self._arena = arena
+    self._arm = arm
+    self._hand = hand
+    self._arm.attach(self._hand)
+    self._arena.attach_offset(self._arm, offset=workspace.arm_offset)
+    self.control_timestep = control_timestep
+    self._tcp_initializer = initializers.ToolCenterPointInitializer(
+        self._hand, self._arm,
+        position=distributions.Uniform(*workspace.tcp_bbox),
+        quaternion=workspaces.DOWN_QUATERNION)
+
+    # Add custom camera observable.
+    self._task_observables = cameras.add_camera_observables(
+        arena, obs_settings, cameras.FRONT_CLOSE)
+
+    target_pos_distribution = distributions.Uniform(*TASKS[task_id])
+    self._prop = prop
+    if prop:
+      # The prop itself is used to visualize the target location.
+      self._make_target_site(parent_entity=prop, visible=False)
+      self._target = self._arena.add_free_entity(prop)
+      self._prop_placer = initializers.PropPlacer(
+          props=[prop],
+          position=target_pos_distribution,
+          quaternion=workspaces.uniform_z_rotation,
+          settle_physics=True)
+    else:
+      self._target = self._make_target_site(parent_entity=arena, visible=True)
+      self._target_placer = target_pos_distribution
+
+      # Commented to match EXORL
+      # obs = observable.MJCFFeature('pos', self._target)
+      # obs.configure(**obs_settings.prop_pose._asdict())
+      # self._task_observables['target_position'] = obs
+
+    # Add sites for visualizing the prop and target bounding boxes.
+    workspaces.add_bbox_site(
+        body=self.root_entity.mjcf_model.worldbody,
+        lower=workspace.tcp_bbox.lower, upper=workspace.tcp_bbox.upper,
+        rgba=constants.GREEN, name='tcp_spawn_area')
+    workspaces.add_bbox_site(
+        body=self.root_entity.mjcf_model.worldbody,
+        lower=workspace.target_bbox.lower, upper=workspace.target_bbox.upper,
+        rgba=constants.BLUE, name='target_spawn_area')
+
+  def _make_target_site(self, parent_entity, visible):
+    return workspaces.add_target_site(
+        body=parent_entity.mjcf_model.worldbody,
+        radius=_TARGET_RADIUS, visible=visible,
+        rgba=constants.RED, name='target_site')
+
+  @property
+  def root_entity(self):
+    return self._arena
+
+  @property
+  def arm(self):
+    return self._arm
+
+  @property
+  def hand(self):
+    return self._hand
+
+  @property
+  def task_observables(self):
+    return self._task_observables
+
+  def get_reward(self, physics):
+    hand_pos = physics.bind(self._hand.tool_center_point).xpos
+    target_pos = physics.bind(self._target).xpos
+    # This was used exceptionally for the PT reward predictor experiments 
+    # target_pos = distributions.Uniform(*TASKS[self._task_id])()
+    distance = np.linalg.norm(hand_pos - target_pos)
+    return rewards.tolerance(
+        distance, bounds=(0, _TARGET_RADIUS), margin=_TARGET_RADIUS)
+
+  def initialize_episode(self, physics, random_state):
+    self._hand.set_grasp(physics, close_factors=random_state.uniform())
+    self._tcp_initializer(physics, random_state)
+    if self._prop:
+      self._prop_placer(physics, random_state)
+    else:
+      physics.bind(self._target).pos = (
+          self._target_placer(random_state=random_state))
+
+
+def _reach(task_id, obs_settings, use_site):
+  """Configure and instantiate a `Reach` task.
+
+  Args:
+    obs_settings: An `observations.ObservationSettings` instance.
+    use_site: Boolean, if True then the target will be a fixed site, otherwise
+      it will be a moveable Duplo brick.
+
+  Returns:
+    An instance of `reach.Reach`.
+  """
+  arena = arenas.Standard()
+  arm = robots.make_arm(obs_settings=obs_settings)
+  hand = robots.make_hand(obs_settings=obs_settings)
+  if use_site:
+    workspace = _SITE_WORKSPACE
+    prop = None
+  else:
+    workspace = _DUPLO_WORKSPACE
+    prop = props.Duplo(observable_options=observations.make_options(
+        obs_settings, observations.FREEPROP_OBSERVABLES))
+  task = MTReach(task_id, arena=arena, arm=arm, hand=hand, prop=prop,
+               obs_settings=obs_settings,
+               workspace=workspace,
+               control_timestep=constants.CONTROL_TIMESTEP)
+  return task

envs/custom_dmc_tasks/quadruped.py

@@ -0,0 +1,683 @@
+# Copyright 2019 The dm_control Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+
+"""Quadruped Domain."""
+
+import collections
+
+from dm_control.suite import quadruped
+from dm_control import mujoco
+from dm_control.mujoco.wrapper import mjbindings
+from dm_control.rl import control
+from dm_control.suite import base
+from dm_control.suite import common
+from dm_control.utils import containers
+from dm_control.utils import rewards
+from dm_control.utils import xml_tools
+from dm_control.utils import io as resources
+from lxml import etree
+import numpy as np
+from scipy import ndimage
+import os 
+
+enums = mjbindings.enums
+mjlib = mjbindings.mjlib
+
+
+_DEFAULT_TIME_LIMIT = 20
+_CONTROL_TIMESTEP = .02
+
+# Horizontal speeds above which the move reward is 1.
+_RUN_SPEED = 5
+_WALK_SPEED = 0.5
+
+_JUMP_HEIGHT = 1.0 # -also good for foot up
+_LIE_DOWN_HEIGHT = 0.2
+_FOOT_DOWN_HEIGHT = 0.2
+_FOOT_UP_HEIGHT = 0.8
+
+# Constants related to terrain generation.
+_HEIGHTFIELD_ID = 0
+_TERRAIN_SMOOTHNESS = 0.15  # 0.0: maximally bumpy; 1.0: completely smooth.
+_TERRAIN_BUMP_SCALE = 2  # Spatial scale of terrain bumps (in meters).
+
+# Named model elements.
+_TOES = ['toe_front_left', 'toe_back_left', 'toe_back_right', 'toe_front_right']
+_WALLS = ['wall_px', 'wall_py', 'wall_nx', 'wall_ny']
+
+def make(task,
+         task_kwargs=None,
+         environment_kwargs=None,
+         visualize_reward=False):
+    task_kwargs = task_kwargs or {}
+    if environment_kwargs is not None:
+        task_kwargs = task_kwargs.copy()
+        task_kwargs['environment_kwargs'] = environment_kwargs
+    env = SUITE[task](**task_kwargs)
+    env.task.visualize_reward = visualize_reward
+    return env
+
+def get_model_and_assets():
+    """Returns a tuple containing the model XML string and a dict of assets."""
+    root_dir = os.path.dirname(os.path.dirname(__file__))
+    xml = resources.GetResource(
+        os.path.join(root_dir, 'custom_dmc_tasks', 'quadruped.xml'))
+    return xml, common.ASSETS
+
+
+def make_model(floor_size=None, terrain=False, rangefinders=False,
+               walls_and_ball=False):
+  """Returns the model XML string."""
+  root_dir = os.path.dirname(os.path.dirname(__file__))
+  xml_string = common.read_model(os.path.join(root_dir, 'custom_dmc_tasks', 'quadruped.xml'))
+  parser = etree.XMLParser(remove_blank_text=True)
+  mjcf = etree.XML(xml_string, parser)
+
+  # Set floor size.
+  if floor_size is not None:
+    floor_geom = mjcf.find('.//geom[@name=\'floor\']')
+    floor_geom.attrib['size'] = f'{floor_size} {floor_size} .5'
+
+  # Remove walls, ball and target.
+  if not walls_and_ball:
+    for wall in _WALLS:
+      wall_geom = xml_tools.find_element(mjcf, 'geom', wall)
+      wall_geom.getparent().remove(wall_geom)
+
+    # Remove ball.
+    ball_body = xml_tools.find_element(mjcf, 'body', 'ball')
+    ball_body.getparent().remove(ball_body)
+
+    # Remove target.
+    target_site = xml_tools.find_element(mjcf, 'site', 'target')
+    target_site.getparent().remove(target_site)
+
+  # Remove terrain.
+  if not terrain:
+    terrain_geom = xml_tools.find_element(mjcf, 'geom', 'terrain')
+    terrain_geom.getparent().remove(terrain_geom)
+
+  # Remove rangefinders if they're not used, as range computations can be
+  # expensive, especially in a scene with heightfields.
+  if not rangefinders:
+    rangefinder_sensors = mjcf.findall('.//rangefinder')
+    for rf in rangefinder_sensors:
+      rf.getparent().remove(rf)
+
+  return etree.tostring(mjcf, pretty_print=True)
+
+
+@quadruped.SUITE.add('custom')
+def lie_down(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Walk task."""
+  xml_string = make_model(floor_size=_DEFAULT_TIME_LIMIT * _WALK_SPEED)
+  physics = Physics.from_xml_string(xml_string, common.ASSETS)
+  task = Stand(goal='lie_down', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(physics, task, time_limit=time_limit,
+                             control_timestep=_CONTROL_TIMESTEP,
+                             **environment_kwargs)
+
+
+@quadruped.SUITE.add('custom')
+def two_legs(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Walk task."""
+  xml_string = make_model(floor_size=_DEFAULT_TIME_LIMIT * _WALK_SPEED)
+  physics = Physics.from_xml_string(xml_string, common.ASSETS)
+  task = Stand(goal='two_legs', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(physics, task, time_limit=time_limit,
+                             control_timestep=_CONTROL_TIMESTEP,
+                             **environment_kwargs)
+
+
+@quadruped.SUITE.add('custom')
+def stand(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Walk task."""
+  xml_string = make_model(floor_size=_DEFAULT_TIME_LIMIT * _WALK_SPEED)
+  physics = Physics.from_xml_string(xml_string, common.ASSETS)
+  task = Stand(goal='stand', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(physics, task, time_limit=time_limit,
+                             control_timestep=_CONTROL_TIMESTEP,
+                             **environment_kwargs)
+
+@quadruped.SUITE.add('custom')
+def jump(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Walk task."""
+  xml_string = make_model(floor_size=_DEFAULT_TIME_LIMIT * _WALK_SPEED)
+  physics = Physics.from_xml_string(xml_string, common.ASSETS)
+  task = Jump(desired_height=_JUMP_HEIGHT, random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(physics, task, time_limit=time_limit,
+                             control_timestep=_CONTROL_TIMESTEP,
+                             **environment_kwargs)
+
+@quadruped.SUITE.add('custom')
+def roll(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Walk task."""
+  xml_string = make_model(floor_size=_DEFAULT_TIME_LIMIT * _WALK_SPEED)
+  physics = Physics.from_xml_string(xml_string, common.ASSETS)
+  task = Roll(desired_speed=_WALK_SPEED, random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(physics, task, time_limit=time_limit,
+                             control_timestep=_CONTROL_TIMESTEP,
+                             **environment_kwargs)
+
+@quadruped.SUITE.add('custom')
+def roll_fast(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Walk task."""
+  xml_string = make_model(floor_size=_DEFAULT_TIME_LIMIT * _WALK_SPEED)
+  physics = Physics.from_xml_string(xml_string, common.ASSETS)
+  task = Roll(desired_speed=_RUN_SPEED, random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(physics, task, time_limit=time_limit,
+                             control_timestep=_CONTROL_TIMESTEP,
+                             **environment_kwargs)
+
+class Physics(mujoco.Physics):
+  """Physics simulation with additional features for the Quadruped domain."""
+
+  def _reload_from_data(self, data):
+    super()._reload_from_data(data)
+    # Clear cached sensor names when the physics is reloaded.
+    self._sensor_types_to_names = {}
+    self._hinge_names = []
+
+  def _get_sensor_names(self, *sensor_types):
+    try:
+      sensor_names = self._sensor_types_to_names[sensor_types]
+    except KeyError:
+      [sensor_ids] = np.where(np.in1d(self.model.sensor_type, sensor_types))
+      sensor_names = [self.model.id2name(s_id, 'sensor') for s_id in sensor_ids]
+      self._sensor_types_to_names[sensor_types] = sensor_names
+    return sensor_names
+
+  def torso_upright(self):
+    """Returns the dot-product of the torso z-axis and the global z-axis."""
+    return np.asarray(self.named.data.xmat['torso', 'zz'])
+
+  def torso_velocity(self):
+    """Returns the velocity of the torso, in the local frame."""
+    return self.named.data.sensordata['velocimeter'].copy()
+  
+  def com_height(self):
+    return self.named.data.sensordata['center_of_mass'].copy()[2]
+
+  def egocentric_state(self):
+    """Returns the state without global orientation or position."""
+    if not self._hinge_names:
+      [hinge_ids] = np.nonzero(self.model.jnt_type ==
+                               enums.mjtJoint.mjJNT_HINGE)
+      self._hinge_names = [self.model.id2name(j_id, 'joint')
+                           for j_id in hinge_ids]
+    return np.hstack((self.named.data.qpos[self._hinge_names],
+                      self.named.data.qvel[self._hinge_names],
+                      self.data.act))
+
+  def toe_positions(self):
+    """Returns toe positions in egocentric frame."""
+    torso_frame = self.named.data.xmat['torso'].reshape(3, 3)
+    torso_pos = self.named.data.xpos['torso']
+    torso_to_toe = self.named.data.xpos[_TOES] - torso_pos
+    return torso_to_toe.dot(torso_frame)
+
+  def force_torque(self):
+    """Returns scaled force/torque sensor readings at the toes."""
+    force_torque_sensors = self._get_sensor_names(enums.mjtSensor.mjSENS_FORCE,
+                                                  enums.mjtSensor.mjSENS_TORQUE)
+    return np.arcsinh(self.named.data.sensordata[force_torque_sensors])
+
+  def imu(self):
+    """Returns IMU-like sensor readings."""
+    imu_sensors = self._get_sensor_names(enums.mjtSensor.mjSENS_GYRO,
+                                         enums.mjtSensor.mjSENS_ACCELEROMETER)
+    return self.named.data.sensordata[imu_sensors]
+
+  def rangefinder(self):
+    """Returns scaled rangefinder sensor readings."""
+    rf_sensors = self._get_sensor_names(enums.mjtSensor.mjSENS_RANGEFINDER)
+    rf_readings = self.named.data.sensordata[rf_sensors]
+    no_intersection = -1.0
+    return np.where(rf_readings == no_intersection, 1.0, np.tanh(rf_readings))
+
+  def origin_distance(self):
+    """Returns the distance from the origin to the workspace."""
+    return np.asarray(np.linalg.norm(self.named.data.site_xpos['workspace']))
+
+  def origin(self):
+    """Returns origin position in the torso frame."""
+    torso_frame = self.named.data.xmat['torso'].reshape(3, 3)
+    torso_pos = self.named.data.xpos['torso']
+    return -torso_pos.dot(torso_frame)
+
+  def ball_state(self):
+    """Returns ball position and velocity relative to the torso frame."""
+    data = self.named.data
+    torso_frame = data.xmat['torso'].reshape(3, 3)
+    ball_rel_pos = data.xpos['ball'] - data.xpos['torso']
+    ball_rel_vel = data.qvel['ball_root'][:3] - data.qvel['root'][:3]
+    ball_rot_vel = data.qvel['ball_root'][3:]
+    ball_state = np.vstack((ball_rel_pos, ball_rel_vel, ball_rot_vel))
+    return ball_state.dot(torso_frame).ravel()
+
+  def target_position(self):
+    """Returns target position in torso frame."""
+    torso_frame = self.named.data.xmat['torso'].reshape(3, 3)
+    torso_pos = self.named.data.xpos['torso']
+    torso_to_target = self.named.data.site_xpos['target'] - torso_pos
+    return torso_to_target.dot(torso_frame)
+
+  def ball_to_target_distance(self):
+    """Returns horizontal distance from the ball to the target."""
+    ball_to_target = (self.named.data.site_xpos['target'] -
+                      self.named.data.xpos['ball'])
+    return np.linalg.norm(ball_to_target[:2])
+
+  def self_to_ball_distance(self):
+    """Returns horizontal distance from the quadruped workspace to the ball."""
+    self_to_ball = (self.named.data.site_xpos['workspace']
+                    -self.named.data.xpos['ball'])
+    return np.linalg.norm(self_to_ball[:2])
+
+
+def _find_non_contacting_height(physics, orientation, x_pos=0.0, y_pos=0.0):
+  """Find a height with no contacts given a body orientation.
+  Args:
+    physics: An instance of `Physics`.
+    orientation: A quaternion.
+    x_pos: A float. Position along global x-axis.
+    y_pos: A float. Position along global y-axis.
+  Raises:
+    RuntimeError: If a non-contacting configuration has not been found after
+    10,000 attempts.
+  """
+  z_pos = 0.0  # Start embedded in the floor.
+  num_contacts = 1
+  num_attempts = 0
+  # Move up in 1cm increments until no contacts.
+  while num_contacts > 0:
+    try:
+      with physics.reset_context():
+        physics.named.data.qpos['root'][:3] = x_pos, y_pos, z_pos
+        physics.named.data.qpos['root'][3:] = orientation
+    except control.PhysicsError:
+      # We may encounter a PhysicsError here due to filling the contact
+      # buffer, in which case we simply increment the height and continue.
+      pass
+    num_contacts = physics.data.ncon
+    z_pos += 0.01
+    num_attempts += 1
+    if num_attempts > 10000:
+      raise RuntimeError('Failed to find a non-contacting configuration.')
+
+
+def _common_observations(physics):
+  """Returns the observations common to all tasks."""
+  obs = collections.OrderedDict()
+  obs['egocentric_state'] = physics.egocentric_state()
+  obs['torso_velocity'] = physics.torso_velocity()
+  obs['torso_upright'] = physics.torso_upright()
+  obs['imu'] = physics.imu()
+  obs['force_torque'] = physics.force_torque()
+  return obs
+
+def _lie_down_reward(physics, deviation_angle=0):
+  """Returns a reward proportional to how upright the torso is.
+  Args:
+    physics: an instance of `Physics`.
+    deviation_angle: A float, in degrees. The reward is 0 when the torso is
+      exactly upside-down and 1 when the torso's z-axis is less than
+      `deviation_angle` away from the global z-axis.
+  """
+  torso = physics.named.data.xpos['torso', 'z']
+  return rewards.tolerance(
+                  torso,
+                  bounds=(-float('inf'), _LIE_DOWN_HEIGHT),
+                  margin=_LIE_DOWN_HEIGHT * 1.5)
+
+
+def _two_legs_reward(physics, deviation_angle=0):
+  """Returns a reward proportional to how upright the torso is.
+  Args:
+    physics: an instance of `Physics`.
+    deviation_angle: A float, in degrees. The reward is 0 when the torso is
+      exactly upside-down and 1 when the torso's z-axis is less than
+      `deviation_angle` away from the global z-axis.
+  """
+  toes = []
+  for t in ['toe_front_left', 'toe_front_right', 'toe_back_left', 'toe_back_right']:
+    toe = physics.named.data.xpos[t, 'z']
+    toes.append(toe)
+  toes = sorted(toes)
+  min_toes = sum(toes[:2]) / 2
+  max_toes = sum(toes[2:]) / 2
+  toes_up = rewards.tolerance(
+                  max_toes,
+                  bounds=(_FOOT_UP_HEIGHT, float('inf')),
+                  margin=_FOOT_UP_HEIGHT // 2)
+  toes_down = rewards.tolerance(
+                                min_toes,
+                                bounds=(-float('inf'), _FOOT_DOWN_HEIGHT),
+                                margin=_FOOT_DOWN_HEIGHT * 1.5)
+  return toes_down * toes_up
+
+
+def _upright_reward(physics, deviation_angle=0):
+  """Returns a reward proportional to how upright the torso is.
+  Args:
+    physics: an instance of `Physics`.
+    deviation_angle: A float, in degrees. The reward is 0 when the torso is
+      exactly upside-down and 1 when the torso's z-axis is less than
+      `deviation_angle` away from the global z-axis.
+  """
+  deviation = np.cos(np.deg2rad(deviation_angle))
+  return rewards.tolerance(
+      physics.torso_upright(),
+      bounds=(deviation, float('inf')),
+      sigmoid='linear',
+      margin=1 + deviation,
+      value_at_margin=0)
+
+
+class Move(base.Task):
+  """A quadruped task solved by moving forward at a designated speed."""
+
+  def __init__(self, desired_speed, random=None):
+    """Initializes an instance of `Move`.
+    Args:
+      desired_speed: A float. If this value is zero, reward is given simply
+        for standing upright. Otherwise this specifies the horizontal velocity
+        at which the velocity-dependent reward component is maximized.
+      random: Optional, either a `numpy.random.RandomState` instance, an
+        integer seed for creating a new `RandomState`, or None to select a seed
+        automatically (default).
+    """
+    self._desired_speed = desired_speed
+    super().__init__(random=random)
+
+  def initialize_episode(self, physics):
+    """Sets the state of the environment at the start of each episode.
+    Args:
+      physics: An instance of `Physics`.
+    """
+    # Initial configuration.
+    orientation = self.random.randn(4)
+    orientation /= np.linalg.norm(orientation)
+    _find_non_contacting_height(physics, orientation)
+    super().initialize_episode(physics)
+
+  def get_observation(self, physics):
+    """Returns an observation to the agent."""
+    return _common_observations(physics)
+
+  def get_reward(self, physics):
+    """Returns a reward to the agent."""
+
+    # Move reward term.
+    move_reward = rewards.tolerance(
+        physics.torso_velocity()[0],
+        bounds=(self._desired_speed, float('inf')),
+        margin=self._desired_speed,
+        value_at_margin=0.5,
+        sigmoid='linear')
+
+    return _upright_reward(physics) * move_reward
+
+
+class Stand(base.Task):
+  """A quadruped task solved by moving forward at a designated speed."""
+
+  def __init__(self, random=None, goal='stand'):
+    """Initializes an instance of `Move`.
+    Args:
+      desired_speed: A float. If this value is zero, reward is given simply
+        for standing upright. Otherwise this specifies the horizontal velocity
+        at which the velocity-dependent reward component is maximized.
+      random: Optional, either a `numpy.random.RandomState` instance, an
+        integer seed for creating a new `RandomState`, or None to select a seed
+        automatically (default).
+    """
+    super().__init__(random=random)
+    self._goal = goal
+
+  def initialize_episode(self, physics):
+    """Sets the state of the environment at the start of each episode.
+    Args:
+      physics: An instance of `Physics`.
+    """
+    # Initial configuration.
+    orientation = self.random.randn(4)
+    orientation /= np.linalg.norm(orientation)
+    _find_non_contacting_height(physics, orientation)
+    super().initialize_episode(physics)
+
+  def get_observation(self, physics):
+    """Returns an observation to the agent."""
+    return _common_observations(physics)
+
+  def get_reward(self, physics):
+    """Returns a reward to the agent."""
+    if self._goal == 'stand':
+      return _upright_reward(physics) 
+    elif self._goal == 'lie_down':
+          return _lie_down_reward(physics) 
+    elif self._goal == 'two_legs':
+      return _two_legs_reward(physics) 
+
+class Jump(base.Task):
+  """A quadruped task solved by moving forward at a designated speed."""
+
+  def __init__(self, desired_height, random=None):
+    """Initializes an instance of `Move`.
+    Args:
+      desired_speed: A float. If this value is zero, reward is given simply
+        for standing upright. Otherwise this specifies the horizontal velocity
+        at which the velocity-dependent reward component is maximized.
+      random: Optional, either a `numpy.random.RandomState` instance, an
+        integer seed for creating a new `RandomState`, or None to select a seed
+        automatically (default).
+    """
+    self._desired_height = desired_height
+    super().__init__(random=random)
+
+  def initialize_episode(self, physics):
+    """Sets the state of the environment at the start of each episode.
+    Args:
+      physics: An instance of `Physics`.
+    """
+    # Initial configuration.
+    orientation = self.random.randn(4)
+    orientation /= np.linalg.norm(orientation)
+    _find_non_contacting_height(physics, orientation)
+    super().initialize_episode(physics)
+
+  def get_observation(self, physics):
+    """Returns an observation to the agent."""
+    return _common_observations(physics)
+
+  def get_reward(self, physics):
+    """Returns a reward to the agent."""
+
+    # Move reward term.
+    jump_up = rewards.tolerance(
+        physics.com_height(),
+        bounds=(self._desired_height, float('inf')),
+        margin=self._desired_height,
+        value_at_margin=0.5,
+        sigmoid='linear')
+
+    return _upright_reward(physics) * jump_up
+
+
+class Roll(base.Task):
+  """A quadruped task solved by moving forward at a designated speed."""
+
+  def __init__(self, desired_speed, random=None):
+    """Initializes an instance of `Move`.
+    Args:
+      desired_speed: A float. If this value is zero, reward is given simply
+        for standing upright. Otherwise this specifies the horizontal velocity
+        at which the velocity-dependent reward component is maximized.
+      random: Optional, either a `numpy.random.RandomState` instance, an
+        integer seed for creating a new `RandomState`, or None to select a seed
+        automatically (default).
+    """
+    self._desired_speed = desired_speed
+    super().__init__(random=random)
+
+  def initialize_episode(self, physics):
+    """Sets the state of the environment at the start of each episode.
+    Args:
+      physics: An instance of `Physics`.
+    """
+    # Initial configuration.
+    orientation = self.random.randn(4)
+    orientation /= np.linalg.norm(orientation)
+    _find_non_contacting_height(physics, orientation)
+    super().initialize_episode(physics)
+
+  def get_observation(self, physics):
+    """Returns an observation to the agent."""
+    return _common_observations(physics)
+
+  def get_reward(self, physics):
+    """Returns a reward to the agent."""
+    # Move reward term.
+    move_reward = rewards.tolerance(
+        np.linalg.norm(physics.torso_velocity()),
+        bounds=(self._desired_speed, float('inf')),
+        margin=self._desired_speed,
+        value_at_margin=0.5,
+        sigmoid='linear')
+
+    return _upright_reward(physics) * move_reward
+
+
+class Escape(base.Task):
+  """A quadruped task solved by escaping a bowl-shaped terrain."""
+
+  def initialize_episode(self, physics):
+    """Sets the state of the environment at the start of each episode.
+    Args:
+      physics: An instance of `Physics`.
+    """
+    # Get heightfield resolution, assert that it is square.
+    res = physics.model.hfield_nrow[_HEIGHTFIELD_ID]
+    assert res == physics.model.hfield_ncol[_HEIGHTFIELD_ID]
+    # Sinusoidal bowl shape.
+    row_grid, col_grid = np.ogrid[-1:1:res*1j, -1:1:res*1j]
+    radius = np.clip(np.sqrt(col_grid**2 + row_grid**2), .04, 1)
+    bowl_shape = .5 - np.cos(2*np.pi*radius)/2
+    # Random smooth bumps.
+    terrain_size = 2 * physics.model.hfield_size[_HEIGHTFIELD_ID, 0]
+    bump_res = int(terrain_size / _TERRAIN_BUMP_SCALE)
+    bumps = self.random.uniform(_TERRAIN_SMOOTHNESS, 1, (bump_res, bump_res))
+    smooth_bumps = ndimage.zoom(bumps, res / float(bump_res))
+    # Terrain is elementwise product.
+    terrain = bowl_shape * smooth_bumps
+    start_idx = physics.model.hfield_adr[_HEIGHTFIELD_ID]
+    physics.model.hfield_data[start_idx:start_idx+res**2] = terrain.ravel()
+    super().initialize_episode(physics)
+
+    # If we have a rendering context, we need to re-upload the modified
+    # heightfield data.
+    if physics.contexts:
+      with physics.contexts.gl.make_current() as ctx:
+        ctx.call(mjlib.mjr_uploadHField,
+                 physics.model.ptr,
+                 physics.contexts.mujoco.ptr,
+                 _HEIGHTFIELD_ID)
+
+    # Initial configuration.
+    orientation = self.random.randn(4)
+    orientation /= np.linalg.norm(orientation)
+    _find_non_contacting_height(physics, orientation)
+
+  def get_observation(self, physics):
+    """Returns an observation to the agent."""
+    obs = _common_observations(physics)
+    obs['origin'] = physics.origin()
+    obs['rangefinder'] = physics.rangefinder()
+    return obs
+
+  def get_reward(self, physics):
+    """Returns a reward to the agent."""
+
+    # Escape reward term.
+    terrain_size = physics.model.hfield_size[_HEIGHTFIELD_ID, 0]
+    escape_reward = rewards.tolerance(
+        physics.origin_distance(),
+        bounds=(terrain_size, float('inf')),
+        margin=terrain_size,
+        value_at_margin=0,
+        sigmoid='linear')
+
+    return _upright_reward(physics, deviation_angle=20) * escape_reward
+
+
+class Fetch(base.Task):
+  """A quadruped task solved by bringing a ball to the origin."""
+
+  def initialize_episode(self, physics):
+    """Sets the state of the environment at the start of each episode.
+    Args:
+      physics: An instance of `Physics`.
+    """
+    # Initial configuration, random azimuth and horizontal position.
+    azimuth = self.random.uniform(0, 2*np.pi)
+    orientation = np.array((np.cos(azimuth/2), 0, 0, np.sin(azimuth/2)))
+    spawn_radius = 0.9 * physics.named.model.geom_size['floor', 0]
+    x_pos, y_pos = self.random.uniform(-spawn_radius, spawn_radius, size=(2,))
+    _find_non_contacting_height(physics, orientation, x_pos, y_pos)
+
+    # Initial ball state.
+    physics.named.data.qpos['ball_root'][:2] = self.random.uniform(
+        -spawn_radius, spawn_radius, size=(2,))
+    physics.named.data.qpos['ball_root'][2] = 2
+    physics.named.data.qvel['ball_root'][:2] = 5*self.random.randn(2)
+    super().initialize_episode(physics)
+
+  def get_observation(self, physics):
+    """Returns an observation to the agent."""
+    obs = _common_observations(physics)
+    obs['ball_state'] = physics.ball_state()
+    obs['target_position'] = physics.target_position()
+    return obs
+
+  def get_reward(self, physics):
+    """Returns a reward to the agent."""
+
+    # Reward for moving close to the ball.
+    arena_radius = physics.named.model.geom_size['floor', 0] * np.sqrt(2)
+    workspace_radius = physics.named.model.site_size['workspace', 0]
+    ball_radius = physics.named.model.geom_size['ball', 0]
+    reach_reward = rewards.tolerance(
+        physics.self_to_ball_distance(),
+        bounds=(0, workspace_radius+ball_radius),
+        sigmoid='linear',
+        margin=arena_radius, value_at_margin=0)
+
+    # Reward for bringing the ball to the target.
+    target_radius = physics.named.model.site_size['target', 0]
+    fetch_reward = rewards.tolerance(
+        physics.ball_to_target_distance(),
+        bounds=(0, target_radius),
+        sigmoid='linear',
+        margin=arena_radius, value_at_margin=0)
+
+    reach_then_fetch = reach_reward * (0.5 + 0.5*fetch_reward)
+
+    return _upright_reward(physics) * reach_then_fetch

envs/custom_dmc_tasks/quadruped.xml

@@ -0,0 +1,328 @@
+<mujoco model="quadruped">
+
+  <include file="./common/skybox.xml"/>
+  <include file="./common/visual.xml"/>
+  <include file="./common/materials.xml"/>
+
+  <visual>
+    <rgba rangefinder="1 1 0.1 0.1"/>
+    <map znear=".005" zfar="20"/>
+  </visual>
+
+  <asset>
+    <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="wall">
+      <geom type="plane" material="decoration"/>
+    </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>
+
+  <asset>
+    <texture name="ball" builtin="checker" mark="cross" width="151" height="151"
+             rgb1="0.1 0.1 0.1" rgb2="0.9 0.9 0.9" markrgb="1 1 1"/>
+    <material name="ball" texture="ball" />
+  </asset>
+
+
+  <worldbody>
+    <geom name="floor" type="plane" size="15 15 .5" material="grid"/>
+    <geom name="wall_px" class="wall" pos="-15.7 0 .7" zaxis="1 0 1"  size="1 15 .5"/>
+    <geom name="wall_py" class="wall" pos="0 -15.7 .7" zaxis="0 1 1"  size="15 1 .5"/>
+    <geom name="wall_nx" class="wall" pos="15.7 0 .7" zaxis="-1 0 1"  size="1 15 .5"/>
+    <geom name="wall_ny" class="wall" pos="0 15.7 .7" zaxis="0 -1 1"  size="15 1 .5"/>
+    <site name="target" type="cylinder" size=".4 .06" pos="0 0 .05" material="target"/>
+
+    <geom name="terrain" type="hfield" hfield="terrain" rgba=".2 .3 .4 1" pos="0 0 -.01"/>
+
+    <camera name="global"  pos="-10 10 10" xyaxes="-1 -1 0 1 0 1" mode="trackcom"/>
+    <body name="torso" childclass="body" pos="0 0 .57">
+      <freejoint name="root"/>
+
+      <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>
+
+    <body name="ball" pos="0 0 3">
+      <freejoint name="ball_root"/>
+      <geom name="ball" size=".15" material="ball" priority="1" condim="6" friction=".7 .005 .005"
+            solref="-10000 -30"/>
+      <light name="ball_light" pos="0 0 4" mode="trackcom"/>
+    </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>

envs/custom_dmc_tasks/stickman.py

@@ -0,0 +1,647 @@
+# Copyright 2017 The dm_control Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#    http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or  implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ============================================================================
+"""Stickman Domain."""
+
+from __future__ import absolute_import
+from __future__ import division
+from __future__ import print_function
+
+import collections
+import os
+import numpy as np
+import types
+
+from dm_control import mujoco
+from dm_control.rl import control
+from dm_control.suite import base
+from dm_control.suite import common
+from dm_control.suite.utils import randomizers
+from dm_control.utils import containers
+from dm_control.utils import rewards
+from dm_control.utils import io as resources
+from dm_control import suite
+
+class StickmanYogaPoses:
+    lie_back   = [ -1.2 ,  0. ,  -1.57,  0, 0. , 0.0, 0, -0.,  0.0]
+    lie_front  = [-1.2, -0, 1.57, 0, 0, 0, 0, 0., 0.]
+    legs_up    = [ -1.24 ,  0. ,  -1.57,  1.57, 0. , 0.0,  1.57, -0.,  0.0]
+
+    kneel      = [ -0.5 ,  0. ,  0,  0, -1.57, -0.8,  1.57, -1.57,  0.0]
+    side_angle = [ -0.3 ,  0. ,  0.9,  0, 0, -0.7,  1.87, -1.07,  0.0]
+    stand_up   = [-0.15, 0., 0.34, 0.74, -1.34, -0., 1.1, -0.66, -0.1]
+
+    lean_back  = [-0.27, 0., -0.45, 0.22, -1.5, 0.86, 0.6, -0.8, -0.4]
+    boat       = [ -1.04 ,  0. ,  -0.8,  1.6, 0. , 0.0, 1.6, -0.,  0.0]
+    bridge     = [-1.1, 0., -2.2, -0.3, -1.5, 0., -0.3, -0.8, -0.4]
+
+    head_stand = [-1, 0., -3, 0.6, -1, -0.3, 0.9, -0.5, 0.3]
+    one_feet   = [-0.2, 0., 0, 0.7, -1.34, 0.5, 1.5, -0.6, 0.1]
+    arabesque  = [-0.34, 0., 1.57, 1.57, 0, 0., 0, -0., 0.]
+
+    # new
+    high_kick = [-0.165, 3.3  , 5.55 , 1.35 ,-0, +0.5 , -0.7, 0. , 0.2,]
+    splits    = [-0.7, 0., 0.5, -0.7, -1. , 0, 1.75, 0., -0.45 ]
+    sit_knees = [-0.6, -0.2,  0.2, 0.95, -2.5, 0 , 0.95, -2.5, 0 ]
+
+
+_DEFAULT_TIME_LIMIT = 25
+_CONTROL_TIMESTEP = .025
+
+# Minimal height of torso over foot above which stand reward is 1.
+_STAND_HEIGHT = 1.15
+
+# Horizontal speeds (meters/second) above which move reward is 1.
+_WALK_SPEED = 1
+_RUN_SPEED = 8
+
+# Copied from walker:
+_YOGA_HANDS_UP_HEIGHT = 1.75
+_YOGA_STAND_HEIGHT = 1.0 # lower than stan height = 1.2
+_YOGA_LIE_DOWN_HEIGHT = 0.1
+_YOGA_LEGS_UP_HEIGHT = 1.1
+
+_YOGA_FEET_UP_HEIGHT = 0.5
+_YOGA_FEET_UP_LIE_DOWN_HEIGHT = 0.35
+
+_YOGA_KNEE_HEIGHT = 0.25
+_YOGA_KNEESTAND_HEIGHT = 0.75
+
+_YOGA_SITTING_HEIGHT = 0.55
+_YOGA_SITTING_LEGS_HEIGHT = 0.15
+
+# speed from: https://github.com/rll-research/url_benchmark/blob/710c3eb/custom_dmc_tasks/py
+_SPIN_SPEED = 5.0
+#
+_PUNCH_SPEED = 5.0
+_PUNCH_DIST = 0.29
+
+
+SUITE = containers.TaggedTasks()
+
+def make(task,
+         task_kwargs=None,
+         environment_kwargs=None,
+         visualize_reward=False):
+    task_kwargs = task_kwargs or {}
+    if environment_kwargs is not None:
+        task_kwargs = task_kwargs.copy()
+        task_kwargs['environment_kwargs'] = environment_kwargs
+    env = SUITE[task](**task_kwargs)
+    env.task.visualize_reward = visualize_reward
+    return env
+
+def get_model_and_assets():
+    """Returns a tuple containing the model XML string and a dict of assets."""
+    root_dir = os.path.dirname(os.path.dirname(__file__))
+    xml = resources.GetResource(os.path.join(root_dir, 'custom_dmc_tasks', 'stickman.xml'))
+    return xml, common.ASSETS
+
+@SUITE.add('custom')
+def hands_up(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the hands_up task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(goal='hands_up', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+@SUITE.add('custom')
+def boxing(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the boxing task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(goal='boxing', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+@SUITE.add('custom')
+def arabesque(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Arabesque task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(goal='arabesque', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+@SUITE.add('custom')
+def lying_down(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Lie Down task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(goal='lying_down', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+@SUITE.add('custom')
+def legs_up(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Legs Up task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(goal='legs_up', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+@SUITE.add('custom')
+def high_kick(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the High Kick task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(goal='high_kick', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+@SUITE.add('custom')
+def one_foot(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the High Kick task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(goal='one_foot', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+@SUITE.add('custom')
+def lunge_pose(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the High Kick task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(goal='lunge_pose', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+@SUITE.add('custom')
+def sit_knees(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the High Kick task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(goal='sit_knees', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+@SUITE.add('custom')
+def headstand(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Headstand task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(goal='flip', move_speed=0, random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+@SUITE.add('custom')
+def urlb_flip(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Flip task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(goal='urlb_flip', move_speed=_SPIN_SPEED, random=random) 
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+@SUITE.add('custom')
+def flipping(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Flipping task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(goal='flipping', move_speed=2 * _RUN_SPEED, random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+@SUITE.add('custom')
+def flip(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Flip task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(goal='flip', move_speed=2 * _RUN_SPEED, random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+@SUITE.add('custom')
+def backflip(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Backflip task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(goal='flip', move_speed=-2 * _RUN_SPEED, random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+@SUITE.add('custom')
+def stand(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Stand task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(move_speed=0, goal='stand', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+@SUITE.add('custom')
+def walk(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Walk task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(move_speed=_WALK_SPEED, goal='walk', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+@SUITE.add('custom')
+def run(time_limit=_DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Run task."""
+  physics = Physics.from_xml_string(*get_model_and_assets())
+  task = Stickman(move_speed=_RUN_SPEED, goal='run', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=_CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+class Physics(mujoco.Physics):
+    """Physics simulation with additional features for the stickman domain."""
+    def torso_upright(self):
+        """Returns projection from z-axes of torso to the z-axes of world."""
+        return self.named.data.xmat['torso', 'zz']
+
+    def torso_height(self):
+        """Returns the height of the torso."""
+        return self.named.data.xpos['torso', 'z']
+
+    def horizontal_velocity(self):
+        """Returns the horizontal velocity of the center-of-mass."""
+        return self.named.data.sensordata['torso_subtreelinvel'][0]
+
+    def orientations(self):
+        """Returns planar orientations of all bodies."""
+        return self.named.data.xmat[1:, ['xx', 'xz']].ravel()
+
+    def angmomentum(self):
+        """Returns the angular momentum of torso of the stickman about Y axis."""
+        return self.named.data.subtree_angmom['torso'][1]
+
+
+class Stickman(base.Task):
+    """A planar stickman task."""
+    def __init__(self, move_speed=0., goal='walk', forward=True, random=None):
+        """Initializes an instance of `Stickman`.
+
+    Args:
+      move_speed: A float. If this value is zero, reward is given simply for
+        standing up. Otherwise this specifies a target horizontal velocity for
+        the walking task.
+      random: Optional, either a `numpy.random.RandomState` instance, an
+        integer seed for creating a new `RandomState`, or None to select a seed
+        automatically (default).
+    """
+        self._move_speed = move_speed
+        self._forward = 1 if forward else -1
+        self._goal = goal
+        super().__init__(random=random)
+
+    def _hands_up_reward(self, physics):
+        standing = self._stand_reward(physics)
+        left_hand_height = physics.named.data.xpos['left_hand', 'z']
+        right_hand_height = physics.named.data.xpos['right_hand', 'z']
+        
+        hand_height = (left_hand_height + right_hand_height) / 2
+
+        hands_up = rewards.tolerance(hand_height,
+                                bounds=(_YOGA_HANDS_UP_HEIGHT, float('inf')),
+                                margin=_YOGA_HANDS_UP_HEIGHT/2)
+        return standing * hands_up
+
+    def _boxing_reward(self, physics):
+        # torso up, but lower than standing
+        # foot up, higher than torso
+        # foot down
+        standing = self._stand_reward(physics)
+
+        left_hand_velocity = abs(physics.named.data.subtree_linvel['left_hand'][0])
+        right_hand_velocity = abs(physics.named.data.subtree_linvel['right_hand'][0])
+        punch_reward = rewards.tolerance(
+                max(left_hand_velocity, right_hand_velocity),
+                bounds=(_PUNCH_SPEED, float('inf')),
+                margin=_PUNCH_SPEED / 2,
+                value_at_margin=0.5,
+                sigmoid='linear')
+
+        # left_hand_dist = physics.named.data.xpos['left_hand', 'x'] - physics.named.data.xpos['torso', 'x']
+        # right_hand_dist = physics.named.data.xpos['right_hand', 'x'] - physics.named.data.xpos['torso', 'x']
+        # punch_reward = rewards.tolerance(
+        #         max(left_hand_dist, right_hand_dist),
+        #         bounds=(_PUNCH_DIST, float('inf')),
+        #         margin=_PUNCH_DIST / 2,)
+
+        return standing * punch_reward
+
+    def _arabesque_reward(self, physics):
+        # standing horizontal
+        # one foot up, same height as torso
+        # one foot down
+        standing = rewards.tolerance(physics.torso_height(),
+                                bounds=(_YOGA_STAND_HEIGHT, float('inf')),
+                                margin=_YOGA_STAND_HEIGHT/2)
+        
+        left_foot_height = physics.named.data.xpos['left_foot', 'z']
+        right_foot_height = physics.named.data.xpos['right_foot', 'z']
+        
+        max_foot = 'right_foot' if right_foot_height > left_foot_height else 'left_foot'
+        min_foot = 'right_foot' if right_foot_height <= left_foot_height else 'left_foot'
+
+        min_foot_height = physics.named.data.xpos[min_foot, 'z']
+        max_foot_height = physics.named.data.xpos[max_foot, 'z']
+
+        min_foot_down = rewards.tolerance(min_foot_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        max_foot_up = rewards.tolerance(max_foot_height,
+                                bounds=(_YOGA_STAND_HEIGHT, float('inf')),
+                                margin=_YOGA_STAND_HEIGHT/2)
+        
+        min_foot_x = physics.named.data.xpos[min_foot, 'x']
+        max_foot_x = physics.named.data.xpos[max_foot, 'x']
+        
+        correct_foot_pose = 0.1 if max_foot_x > min_foot_x else 1.0
+ 
+        feet_pose = (min_foot_down + max_foot_up * 2) / 3
+        return standing * feet_pose * correct_foot_pose
+    
+    def _lying_down_reward(self, physics):
+        # torso down and horizontal
+        # thigh and feet down
+        torso_down = rewards.tolerance(physics.torso_height(),
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        horizontal = 1 - abs(physics.torso_upright())
+        
+        thigh_height = (physics.named.data.xpos['left_thigh', 'z'] + physics.named.data.xpos['right_thigh', 'z']) / 2
+        thigh_down = rewards.tolerance(thigh_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        leg_height = (physics.named.data.xpos['left_leg', 'z'] + physics.named.data.xpos['right_leg', 'z']) / 2
+        leg_down = rewards.tolerance(leg_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        feet_height = (physics.named.data.xpos['left_foot', 'z'] + physics.named.data.xpos['right_foot', 'z']) / 2
+        feet_down = rewards.tolerance(feet_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        return (3*torso_down + horizontal + thigh_down + feet_down + leg_down) / 7
+    
+    def _legs_up_reward(self, physics):
+        # torso down and horizontal
+        # legs up with thigh down
+        torso_down = rewards.tolerance(physics.torso_height(),
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        horizontal = 1 - abs(physics.torso_upright())
+        torso_down = (3*torso_down +horizontal) / 4
+        
+        feet_height = (physics.named.data.xpos['left_foot', 'z'] + physics.named.data.xpos['right_foot', 'z']) / 2
+        feet_up = rewards.tolerance(feet_height,
+                                bounds=(_YOGA_FEET_UP_LIE_DOWN_HEIGHT, float('inf')),
+                                margin=_YOGA_FEET_UP_LIE_DOWN_HEIGHT/2)
+
+        return torso_down * feet_up
+    
+    def _high_kick_reward(self, physics):
+        # torso up, but lower than standing
+        # foot up, higher than torso
+        # foot down
+        standing = rewards.tolerance(physics.torso_height(),
+                                    bounds=(_YOGA_STAND_HEIGHT, float('inf')),
+                                    margin=_YOGA_STAND_HEIGHT/2)
+
+        left_foot_height = physics.named.data.xpos['left_foot', 'z']
+        right_foot_height = physics.named.data.xpos['right_foot', 'z']
+        
+        min_foot_height = min(left_foot_height, right_foot_height)
+        max_foot_height = max(left_foot_height, right_foot_height)
+
+        min_foot_down = rewards.tolerance(min_foot_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        max_foot_up = rewards.tolerance(max_foot_height,
+                                bounds=(_STAND_HEIGHT, float('inf')),
+                                margin=_STAND_HEIGHT/2)
+        
+        feet_pose = (3 * max_foot_up + min_foot_down) / 4
+
+        return standing * feet_pose
+    
+    def _one_foot_reward(self, physics):
+        # torso up, standing
+        # foot up higher than foot down
+        standing = rewards.tolerance(physics.torso_height(),
+                                    bounds=(_YOGA_STAND_HEIGHT, float('inf')),
+                                    margin=_YOGA_STAND_HEIGHT/2)
+
+        left_foot_height = physics.named.data.xpos['left_foot', 'z']
+        right_foot_height = physics.named.data.xpos['right_foot', 'z']
+        
+        min_foot_height = min(left_foot_height, right_foot_height)
+        max_foot_height = max(left_foot_height, right_foot_height)
+
+        min_foot_down = rewards.tolerance(min_foot_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        max_foot_up = rewards.tolerance(max_foot_height,
+                                bounds=(_YOGA_FEET_UP_HEIGHT, float('inf')),
+                                margin=_YOGA_FEET_UP_HEIGHT/2)
+        
+        return standing * max_foot_up * min_foot_down
+
+    def _lunge_pose_reward(self, physics):
+        # torso up, standing, but lower
+        # leg up higher than leg down
+        # horiontal thigh and leg
+        standing = rewards.tolerance(physics.torso_height(),
+                                    bounds=(_YOGA_KNEESTAND_HEIGHT, float('inf')),
+                                    margin=_YOGA_KNEESTAND_HEIGHT/2)
+        upright = (1 + physics.torso_upright()) / 2
+        torso = (3*standing + upright) / 4
+
+        left_leg_height = physics.named.data.xpos['left_leg', 'z']
+        right_leg_height = physics.named.data.xpos['right_leg', 'z']
+        
+        min_leg_height = min(left_leg_height, right_leg_height)
+        max_leg_height = max(left_leg_height, right_leg_height)
+
+        min_leg_down = rewards.tolerance(min_leg_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        max_leg_up = rewards.tolerance(max_leg_height,
+                                bounds=(_YOGA_KNEE_HEIGHT, float('inf')),
+                                margin=_YOGA_KNEE_HEIGHT / 2)
+        
+        max_thigh = 'left_thigh' if max_leg_height == left_leg_height else 'right_thigh'
+        min_leg = 'left_leg' if min_leg_height == left_leg_height else 'right_leg'
+
+        max_thigh_horiz = 1 - abs(physics.named.data.xmat[max_thigh, 'zz'])
+        min_leg_horiz = 1 - abs(physics.named.data.xmat[min_leg, 'zz'])
+        
+        legs = (min_leg_down + max_leg_up + max_thigh_horiz + min_leg_horiz) / 4
+
+        return torso * legs
+
+    def _sit_knees_reward(self, physics):
+        # torso up, standing, but lower
+        # foot up higher than foot down
+        standing = rewards.tolerance(physics.torso_height(),
+                                    bounds=(_YOGA_SITTING_HEIGHT, float('inf')),
+                                    margin=_YOGA_SITTING_HEIGHT/2)
+        upright = (1 + physics.torso_upright()) / 2
+        torso_up = (3*standing + upright) / 4
+
+        legs_height = (physics.named.data.xpos['left_leg', 'z'] + physics.named.data.xpos['right_leg', 'z']) / 2
+        legs_down = rewards.tolerance(legs_height,
+                                bounds=(-float('inf'), _YOGA_SITTING_LEGS_HEIGHT),
+                                margin=_YOGA_SITTING_LEGS_HEIGHT*1.5)
+        feet_height = (physics.named.data.xpos['left_foot', 'z'] + physics.named.data.xpos['right_foot', 'z']) / 2
+        feet_down = rewards.tolerance(feet_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        
+        l_thigh_foot_distance = max(0.1, abs(physics.named.data.xpos['left_foot', 'x'] - physics.named.data.xpos['left_thigh', 'x'])) - 0.1
+        r_thigh_foot_distance = max(0.1, abs(physics.named.data.xpos['right_foot', 'x'] - physics.named.data.xpos['right_thigh', 'x'])) - 0.1
+        close = np.exp(-(l_thigh_foot_distance + r_thigh_foot_distance)/2)
+        
+        legs = (3 * legs_down + feet_down) / 4
+        return torso_up * legs * close
+
+    def _urlb_flip_reward(self, physics):
+        standing = rewards.tolerance(physics.torso_height(),
+                                     bounds=(_STAND_HEIGHT, float('inf')),
+                                     margin=_STAND_HEIGHT / 2)
+        upright = (1 + physics.torso_upright()) / 2
+        stand_reward = (3 * standing + upright) / 4
+        move_reward = rewards.tolerance(self._forward *
+                                        physics.named.data.subtree_angmom['torso'][1], # physics.angmomentum(),
+                                        bounds=(_SPIN_SPEED, float('inf')),
+                                        margin=_SPIN_SPEED,
+                                        value_at_margin=0,
+                                        sigmoid='linear')
+        return stand_reward * (5 * move_reward + 1) / 6
+
+    def _flip_reward(self, physics):
+        thigh_height = (physics.named.data.xpos['left_thigh', 'z'] + physics.named.data.xpos['right_thigh', 'z']) / 2
+        thigh_up = rewards.tolerance(thigh_height,
+                                bounds=(_YOGA_STAND_HEIGHT, float('inf')),
+                                margin=_YOGA_STAND_HEIGHT/2)
+        feet_height = (physics.named.data.xpos['left_foot', 'z'] + physics.named.data.xpos['right_foot', 'z']) / 2
+        legs_up = rewards.tolerance(feet_height,
+                                bounds=(_YOGA_LEGS_UP_HEIGHT, float('inf')),
+                                margin=_YOGA_LEGS_UP_HEIGHT/2)
+        upside_down_reward = (3*legs_up + 2*thigh_up) / 5
+        if self._move_speed == 0:
+            return upside_down_reward
+        move_reward = rewards.tolerance(physics.named.data.subtree_angmom['torso'][1], # physics.angmomentum(),
+                                    bounds=(self._move_speed, float('inf')) if self._move_speed > 0 else (-float('inf'), self._move_speed),
+                                    margin=abs(self._move_speed)/2,
+                                    value_at_margin=0.5,
+                                    sigmoid='linear')
+        return upside_down_reward * (5*move_reward + 1) / 6
+
+    
+    def _stand_reward(self, physics):
+        standing = rewards.tolerance(physics.torso_height(),
+                                    bounds=(_STAND_HEIGHT, float('inf')),
+                                    margin=_STAND_HEIGHT / 2)
+        upright = (1 + physics.torso_upright()) / 2
+        return (3 * standing + upright) / 4
+
+    def initialize_episode(self, physics):
+        """Sets the state of the environment at the start of each episode.
+
+    In 'standing' mode, use initial orientation and small velocities.
+    In 'random' mode, randomize joint angles and let fall to the floor.
+
+    Args:
+      physics: An instance of `Physics`.
+
+    """
+        randomizers.randomize_limited_and_rotational_joints(physics, self.random)
+        super().initialize_episode(physics)
+
+    def get_observation(self, physics):
+        """Returns an observation of body orientations, height and velocites."""
+        obs = collections.OrderedDict()
+        obs['orientations'] = physics.orientations()
+        obs['height'] = physics.torso_height()
+        obs['velocity'] = physics.velocity()
+        return obs
+
+    def get_reward(self, physics):
+        """Returns a reward to the agent."""
+        if self._goal in ['stand', 'walk', 'run']:
+            stand_reward = self._stand_reward(physics)
+            move_reward = rewards.tolerance(
+                self._forward * physics.horizontal_velocity(),
+                bounds=(self._move_speed, float('inf')),
+                margin=self._move_speed / 2,
+                value_at_margin=0.5,
+                sigmoid='linear')
+            return stand_reward * (5 * move_reward + 1) / 6
+        if self._goal == 'flipping':
+            self._move_speed = abs(self._move_speed)
+            pos_rew = self._flip_reward(physics)
+            self._move_speed = -abs(self._move_speed)
+            neg_rew = self._flip_reward(physics)
+            return max(pos_rew, neg_rew)
+        try:
+            reward_fn = getattr(self, f'_{self._goal}_reward')
+            return reward_fn(physics)
+        except Exception as e:
+            print(e)
+            raise NotImplementedError(f'Goal {self._goal} or function "_{self._goal}_reward" not implemented.')
+
+if __name__ == '__main__':
+    from dm_control import viewer
+    import numpy as np
+    
+    env = boxing()
+    env.task.visualize_reward = True
+
+    action_spec = env.action_spec()
+
+    def zero_policy(time_step):
+        print(time_step.reward)
+        return np.zeros(action_spec.shape)
+    
+    ts = env.reset()
+    while True:
+       ts = env.step(zero_policy(ts))
+
+    viewer.launch(env, policy=zero_policy)
+    
+    # obs = env.reset()
+    # next_obs, reward, done, info = env.step(np.zeros(6))

envs/custom_dmc_tasks/stickman.xml

@@ -0,0 +1,108 @@
+<mujoco model="stickman">
+  <include file="./common/skybox.xml"/>
+  <include file="./common/visual.xml"/>
+  <include file="./common/materials.xml"/>
+
+
+  <option timestep="0.0025"/>
+
+  <statistic extent="2" center="0 0 1"/>
+
+  <default>
+    <joint damping=".1" armature="0.01" limited="true" solimplimit="0 .99 .01"/>
+    <geom contype="1" conaffinity="0" friction=".7 .1 .1"/>
+    <motor ctrlrange="-1 1" ctrllimited="true"/>
+    <site size="0.01"/>
+    <default class="stickman">
+      <geom material="self" type="capsule"/>
+      <joint axis="0 -1 0"/>
+    </default>
+  </default>
+
+  <worldbody>
+    <geom name="floor" type="plane" conaffinity="1" pos="248 0 0" size="250 .8 .2" material="grid" zaxis="0 0 1"/>
+    <body name="torso" pos="0 0 1.25" childclass="stickman">
+      <light name="light" pos="0 0 2" mode="trackcom"/>
+      <camera name="side" pos="0 -2.25 1.0" euler="60 0 0" mode="trackcom"/>
+      <camera name="back" pos="-2 0 .5" xyaxes="0 -1 0 1 0 3" mode="trackcom"/>
+      <joint name="rootz" axis="0 0 1" type="slide" limited="false" armature="0" damping="0"/>
+      <joint name="rootx" axis="1 0 0" type="slide" limited="false" armature="0" damping="0"/>
+      <joint name="rooty" axis="0 1 0" type="hinge" limited="false" armature="0" damping="0"/>
+      <geom name="torso" size="0.06 0.25"/>
+
+      <body name="neck" pos="0 .0 +0.35">
+        <geom name="neck" zaxis="0 0 1" size="0.045 0.045"/>
+        <body name="head" pos="0 .0 +0.1">
+          <geom name="head" type="sphere" size=".1"/>
+        </body>
+      </body>
+
+      <body name="right_thigh" pos="0 -.05 -0.25">
+        <joint name="right_hip" range="-20 100"/>
+        <geom name="right_thigh" pos="0 0 -0.225" size="0.05 0.225"/>
+        <body name="right_leg" pos="0 0 -0.7">
+          <joint name="right_knee" pos="0 0 0.25" range="-150 0"/>
+          <geom name="right_leg" size="0.04 0.25"/>
+          <body name="right_foot" pos="0.06 0 -0.25">
+            <joint name="right_ankle" pos="-0.06 0 0" range="-45 45"/>
+            <geom name="right_foot" zaxis="1 0 0" size="0.04 0.1"/>
+          </body>
+        </body>
+      </body>
+      <body name="left_thigh" pos="0 .05 -0.25" >
+        <joint name="left_hip" range="-20 100"/>
+        <geom name="left_thigh" pos="0 0 -0.225" size="0.05 0.225"/>
+        <body name="left_leg" pos="0 0 -0.7">
+          <joint name="left_knee" pos="0 0 0.25" range="-150 0"/>
+          <geom name="left_leg" size="0.04 0.25"/>
+          <body name="left_foot" pos="0.06 0 -0.25">
+            <joint name="left_ankle" pos="-0.06 0 0" range="-45 45"/>
+            <geom name="left_foot" zaxis="1 0 0" size="0.04 0.1"/>
+          </body>
+        </body>
+      </body>
+
+      <body name="left_arm" pos="0 .05 +0.2">
+        <joint name="left_shoulder" range="-20 100"/>
+        <geom name="left_arm" pos="0 0 -0.135" size="0.04 0.135"/>
+        <body name="left_forearm" pos="0 0 -0.45">
+          <joint name="left_elbow" pos="0 0 0.15" range="0 150"/>
+          <geom name="left_forearm" size="0.035 0.15"/>
+          <body name="left_hand" pos="0.0 0 -0.15">
+            <geom name="left_hand" type="sphere" size=".05"/>
+          </body>
+        </body>
+      </body>
+
+      <body name="right_arm" pos="0 -.05 +0.2">
+        <joint name="right_shoulder" range="-20 100"/>
+        <geom name="right_arm" pos="0 0 -0.135" size="0.04 0.135"/>
+        <body name="right_forearm" pos="0 0 -0.45">
+          <joint name="right_elbow" pos="0 0 0.15" range="0 150"/>
+          <geom name="right_forearm" size="0.035 0.15"/>
+          <body name="right_hand" pos="0.0 0 -0.15">
+            <geom name="right_hand" type="sphere" size=".05"/>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+
+  <sensor>
+    <subtreelinvel name="torso_subtreelinvel" body="torso"/>
+  </sensor>
+
+  <actuator>    
+    <motor name="right_hip" joint="right_hip" gear="100"/>
+    <motor name="right_knee" joint="right_knee" gear="50"/>
+    <motor name="right_ankle" joint="right_ankle" gear="20"/>
+    <motor name="left_hip" joint="left_hip" gear="100"/>
+    <motor name="left_knee" joint="left_knee" gear="50"/>
+    <motor name="left_ankle" joint="left_ankle" gear="20"/>
+
+    <motor name="left_shoulder" joint="left_shoulder" gear="100"/>
+    <motor name="left_elbow" joint="left_elbow" gear="50"/>
+    <motor name="right_shoulder" joint="right_shoulder" gear="100"/>
+    <motor name="right_elbow" joint="right_elbow" gear="50"/>
+  </actuator>
+</mujoco>

envs/custom_dmc_tasks/walker.py

@@ -0,0 +1,489 @@
+import os
+
+import numpy as np
+from dm_control.rl import control
+from dm_control.suite import common
+from dm_control.suite import walker
+from dm_control.utils import rewards
+from dm_control.utils import io as resources
+
+_TASKS_DIR = os.path.join(os.path.dirname(os.path.dirname(__file__)), 'custom_dmc_tasks')
+
+_YOGA_STAND_HEIGHT = 1.0 # lower than stan height = 1.2
+_YOGA_LIE_DOWN_HEIGHT = 0.1
+_YOGA_LEGS_UP_HEIGHT = 1.1
+
+_YOGA_FEET_UP_HEIGHT = 0.5
+_YOGA_FEET_UP_LIE_DOWN_HEIGHT = 0.35
+
+_YOGA_KNEE_HEIGHT = 0.25
+_YOGA_KNEESTAND_HEIGHT = 0.75
+
+_YOGA_SITTING_HEIGHT = 0.55
+_YOGA_SITTING_LEGS_HEIGHT = 0.15
+
+# speed from: https://github.com/rll-research/url_benchmark/blob/710c3eb/custom_dmc_tasks/walker.py
+_SPIN_SPEED = 5.0
+#
+
+class WalkerYogaPoses:
+    """
+    Joint positions for some yoga poses
+    """
+    lie_back   = [ -1.2 ,  0. ,  -1.57,  0, 0. , 0.0, 0, -0.,  0.0]
+    lie_front  = [-1.2,   -0,      1.57, 0, -0.2, 0, 0, -0.2, 0.]
+    legs_up    = [ -1.24 ,  0. ,  -1.57,  1.57, 0. , 0.0,  1.57, -0.,  0.0]
+
+    kneel      = [ -0.5 ,  0. ,  0,  0, -1.57, -0.8,  1.57, -1.57,  0.0]
+    side_angle = [ -0.3 ,  0. ,  0.9,  0, 0, -0.7,  1.87, -1.07,  0.0]
+    stand_up   = [-0.15, 0., 0.34, 0.74, -1.34, -0., 1.1, -0.66, -0.1]
+
+    lean_back  = [-0.27, 0., -0.45, 0.22, -1.5, 0.86, 0.6, -0.8, -0.4]
+    boat       = [ -1.04 ,  0. ,  -0.8,  1.6, 0. , 0.0, 1.6, -0.,  0.0]
+    bridge     = [-1.1, 0., -2.2, -0.3, -1.5, 0., -0.3, -0.8, -0.4]
+
+    head_stand = [-1, 0., -3, 0.6, -1, -0.3, 0.9, -0.5, 0.3]
+    one_foot   = [-0.2, 0., 0, 0.7, -1.34, 0.5, 1.5, -0.6, 0.1]
+
+    arabesque  = [-0.34, 0., 1.57, 1.57, 0, 0., 0, -0., 0.]
+
+    # new
+    high_kick = [-0.165, 3.3  , 5.55 , 1.35 ,-0, +0.5 , -0.7, 0. , 0.2,]
+    splits    = [-0.7, 0., 0.5, -0.7, -1. , 0, 1.75, 0., -0.45 ]
+
+
+def get_model_and_assets():
+    """Returns a tuple containing the model XML string and a dict of assets."""
+    return resources.GetResource(os.path.join(_TASKS_DIR, 'walker.xml')), common.ASSETS
+
+
+@walker.SUITE.add('custom')
+def walk_backwards(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Walk Backwards task."""
+  physics = walker.Physics.from_xml_string(*get_model_and_assets())
+  task = BackwardsPlanarWalker(move_speed=walker._WALK_SPEED, random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+@walker.SUITE.add('custom')
+def run_backwards(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Run Backwards task."""
+  physics = walker.Physics.from_xml_string(*get_model_and_assets())
+  task = BackwardsPlanarWalker(move_speed=walker._RUN_SPEED, random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+@walker.SUITE.add('custom')
+def arabesque(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Arabesque task."""
+  physics = walker.Physics.from_xml_string(*get_model_and_assets())
+  task = YogaPlanarWalker(goal='arabesque', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+@walker.SUITE.add('custom')
+def lying_down(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Lie Down task."""
+  physics = walker.Physics.from_xml_string(*get_model_and_assets())
+  task = YogaPlanarWalker(goal='lying_down', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+@walker.SUITE.add('custom')
+def legs_up(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Legs Up task."""
+  physics = walker.Physics.from_xml_string(*get_model_and_assets())
+  task = YogaPlanarWalker(goal='legs_up', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+@walker.SUITE.add('custom')
+def high_kick(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the High Kick task."""
+  physics = walker.Physics.from_xml_string(*get_model_and_assets())
+  task = YogaPlanarWalker(goal='high_kick', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+@walker.SUITE.add('custom')
+def one_foot(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the High Kick task."""
+  physics = walker.Physics.from_xml_string(*get_model_and_assets())
+  task = YogaPlanarWalker(goal='one_foot', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+@walker.SUITE.add('custom')
+def lunge_pose(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the High Kick task."""
+  physics = walker.Physics.from_xml_string(*get_model_and_assets())
+  task = YogaPlanarWalker(goal='lunge_pose', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+@walker.SUITE.add('custom')
+def sit_knees(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the High Kick task."""
+  physics = walker.Physics.from_xml_string(*get_model_and_assets())
+  task = YogaPlanarWalker(goal='sit_knees', random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+@walker.SUITE.add('custom')
+def headstand(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Headstand task."""
+  physics = walker.Physics.from_xml_string(*get_model_and_assets())
+  task = YogaPlanarWalker(goal='flip', move_speed=0, random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+@walker.SUITE.add('custom')
+def urlb_flip(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Flip task."""
+  physics = walker.Physics.from_xml_string(*get_model_and_assets())
+  task = YogaPlanarWalker(goal='urlb_flip', move_speed=_SPIN_SPEED, random=random) 
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+@walker.SUITE.add('custom')
+def flipping(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the flipping task."""
+  physics = walker.Physics.from_xml_string(*get_model_and_assets())
+  task = YogaPlanarWalker(goal='flipping', move_speed=2* walker._RUN_SPEED, random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+@walker.SUITE.add('custom')
+def flip(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Flip task."""
+  physics = walker.Physics.from_xml_string(*get_model_and_assets())
+  task = YogaPlanarWalker(goal='flip', move_speed=2* walker._RUN_SPEED, random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+@walker.SUITE.add('custom')
+def backflip(time_limit=walker._DEFAULT_TIME_LIMIT, random=None, environment_kwargs=None):
+  """Returns the Backflip task."""
+  physics = walker.Physics.from_xml_string(*get_model_and_assets())
+  task = YogaPlanarWalker(goal='flip', move_speed=-2 * walker._RUN_SPEED, random=random)
+  environment_kwargs = environment_kwargs or {}
+  return control.Environment(
+      physics, task, time_limit=time_limit, control_timestep=walker._CONTROL_TIMESTEP,
+      **environment_kwargs)
+
+
+class BackwardsPlanarWalker(walker.PlanarWalker):
+    """Backwards PlanarWalker task."""
+    def __init__(self, move_speed, random=None):
+        super().__init__(move_speed, random)
+    
+    def get_reward(self, physics):
+        standing = rewards.tolerance(physics.torso_height(),
+                                 bounds=(_YOGA_STAND_HEIGHT, float('inf')),
+                                 margin=_YOGA_STAND_HEIGHT/2)
+        upright = (1 + physics.torso_upright()) / 2
+        stand_reward = (3*standing + upright) / 4
+        if self._move_speed == 0:
+            return stand_reward
+        else:
+            move_reward = rewards.tolerance(physics.horizontal_velocity(),
+                                            bounds=(-float('inf'), -self._move_speed),
+                                            margin=self._move_speed/2,
+                                            value_at_margin=0.5,
+                                            sigmoid='linear')
+            return stand_reward * (5*move_reward + 1) / 6
+
+
+class YogaPlanarWalker(walker.PlanarWalker):
+    """Yoga PlanarWalker tasks."""
+    
+    def __init__(self, goal='arabesque', move_speed=0, random=None):
+        super().__init__(0, random)
+        self._goal = goal
+        self._move_speed = move_speed
+    
+    def _arabesque_reward(self, physics):
+        # standing horizontal
+        # one foot up, same height as torso
+        # one foot down
+        standing = rewards.tolerance(physics.torso_height(),
+                                bounds=(_YOGA_STAND_HEIGHT, float('inf')),
+                                margin=_YOGA_STAND_HEIGHT/2)
+        
+        left_foot_height = physics.named.data.xpos['left_foot', 'z']
+        right_foot_height = physics.named.data.xpos['right_foot', 'z']
+        
+        max_foot = 'right_foot' if right_foot_height > left_foot_height else 'left_foot'
+        min_foot = 'right_foot' if right_foot_height <= left_foot_height else 'left_foot'
+
+        min_foot_height = physics.named.data.xpos[min_foot, 'z']
+        max_foot_height = physics.named.data.xpos[max_foot, 'z']
+
+        min_foot_down = rewards.tolerance(min_foot_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        max_foot_up = rewards.tolerance(max_foot_height,
+                                bounds=(_YOGA_STAND_HEIGHT, float('inf')),
+                                margin=_YOGA_STAND_HEIGHT/2)
+        
+        min_foot_x = physics.named.data.xpos[min_foot, 'x']
+        max_foot_x = physics.named.data.xpos[max_foot, 'x']
+        
+        correct_foot_pose = 0.1 if max_foot_x > min_foot_x else 1.0
+ 
+        feet_pose = (min_foot_down + max_foot_up * 2) / 3
+        return standing * feet_pose * correct_foot_pose
+    
+    def _lying_down_reward(self, physics):
+        # torso down and horizontal
+        # thigh and feet down
+        torso_down = rewards.tolerance(physics.torso_height(),
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        horizontal = 1 - abs(physics.torso_upright())
+        
+        thigh_height = (physics.named.data.xpos['left_thigh', 'z'] + physics.named.data.xpos['right_thigh', 'z']) / 2
+        thigh_down = rewards.tolerance(thigh_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        leg_height = (physics.named.data.xpos['left_leg', 'z'] + physics.named.data.xpos['right_leg', 'z']) / 2
+        leg_down = rewards.tolerance(leg_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        feet_height = (physics.named.data.xpos['left_foot', 'z'] + physics.named.data.xpos['right_foot', 'z']) / 2
+        feet_down = rewards.tolerance(feet_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        return (3*torso_down + horizontal + thigh_down + feet_down + leg_down) / 7
+    
+    def _legs_up_reward(self, physics):
+        # torso down and horizontal
+        # legs up with thigh down
+        torso_down = rewards.tolerance(physics.torso_height(),
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        horizontal = 1 - abs(physics.torso_upright())
+        torso_down = (3*torso_down +horizontal) / 4
+        
+        feet_height = (physics.named.data.xpos['left_foot', 'z'] + physics.named.data.xpos['right_foot', 'z']) / 2
+        feet_up = rewards.tolerance(feet_height,
+                                bounds=(_YOGA_FEET_UP_LIE_DOWN_HEIGHT, float('inf')),
+                                margin=_YOGA_FEET_UP_LIE_DOWN_HEIGHT/2)
+
+        return torso_down * feet_up
+    
+    def _high_kick_reward(self, physics):
+        # torso up, but lower than standing
+        # foot up, higher than torso
+        # foot down
+        standing = rewards.tolerance(physics.torso_height(),
+                                    bounds=(_YOGA_STAND_HEIGHT, float('inf')),
+                                    margin=_YOGA_STAND_HEIGHT/2)
+
+        left_foot_height = physics.named.data.xpos['left_foot', 'z']
+        right_foot_height = physics.named.data.xpos['right_foot', 'z']
+        
+        min_foot_height = min(left_foot_height, right_foot_height)
+        max_foot_height = max(left_foot_height, right_foot_height)
+
+        min_foot_down = rewards.tolerance(min_foot_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        max_foot_up = rewards.tolerance(max_foot_height,
+                                bounds=(walker._STAND_HEIGHT, float('inf')),
+                                margin=walker._STAND_HEIGHT/2)
+        
+        feet_pose = (3 * max_foot_up + min_foot_down) / 4
+
+        return standing * feet_pose
+    
+    def _one_foot_reward(self, physics):
+        # torso up, standing
+        # foot up higher than foot down
+        standing = rewards.tolerance(physics.torso_height(),
+                                    bounds=(_YOGA_STAND_HEIGHT, float('inf')),
+                                    margin=_YOGA_STAND_HEIGHT/2)
+
+        left_foot_height = physics.named.data.xpos['left_foot', 'z']
+        right_foot_height = physics.named.data.xpos['right_foot', 'z']
+        
+        min_foot_height = min(left_foot_height, right_foot_height)
+        max_foot_height = max(left_foot_height, right_foot_height)
+
+        min_foot_down = rewards.tolerance(min_foot_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        max_foot_up = rewards.tolerance(max_foot_height,
+                                bounds=(_YOGA_FEET_UP_HEIGHT, float('inf')),
+                                margin=_YOGA_FEET_UP_HEIGHT/2)
+        
+        return standing * max_foot_up * min_foot_down
+
+    def _lunge_pose_reward(self, physics):
+        # torso up, standing, but lower
+        # leg up higher than leg down
+        # horiontal thigh and leg
+        standing = rewards.tolerance(physics.torso_height(),
+                                    bounds=(_YOGA_KNEESTAND_HEIGHT, float('inf')),
+                                    margin=_YOGA_KNEESTAND_HEIGHT/2)
+        upright = (1 + physics.torso_upright()) / 2
+        torso = (3*standing + upright) / 4
+
+        left_leg_height = physics.named.data.xpos['left_leg', 'z']
+        right_leg_height = physics.named.data.xpos['right_leg', 'z']
+        
+        min_leg_height = min(left_leg_height, right_leg_height)
+        max_leg_height = max(left_leg_height, right_leg_height)
+
+        min_leg_down = rewards.tolerance(min_leg_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        max_leg_up = rewards.tolerance(max_leg_height,
+                                bounds=(_YOGA_KNEE_HEIGHT, float('inf')),
+                                margin=_YOGA_KNEE_HEIGHT / 2)
+        
+        max_thigh = 'left_thigh' if max_leg_height == left_leg_height else 'right_thigh'
+        min_leg = 'left_leg' if min_leg_height == left_leg_height else 'right_leg'
+
+        max_thigh_horiz = 1 - abs(physics.named.data.xmat[max_thigh, 'zz'])
+        min_leg_horiz = 1 - abs(physics.named.data.xmat[min_leg, 'zz'])
+        
+        legs = (min_leg_down + max_leg_up + max_thigh_horiz + min_leg_horiz) / 4
+
+        return torso * legs
+
+    def _sit_knees_reward(self, physics):
+        # torso up, standing, but lower
+        # foot up higher than foot down
+        standing = rewards.tolerance(physics.torso_height(),
+                                    bounds=(_YOGA_SITTING_HEIGHT, float('inf')),
+                                    margin=_YOGA_SITTING_HEIGHT/2)
+        upright = (1 + physics.torso_upright()) / 2
+        torso_up = (3*standing + upright) / 4
+
+        legs_height = (physics.named.data.xpos['left_leg', 'z'] + physics.named.data.xpos['right_leg', 'z']) / 2
+        legs_down = rewards.tolerance(legs_height,
+                                bounds=(-float('inf'), _YOGA_SITTING_LEGS_HEIGHT),
+                                margin=_YOGA_SITTING_LEGS_HEIGHT*1.5)
+        feet_height = (physics.named.data.xpos['left_foot', 'z'] + physics.named.data.xpos['right_foot', 'z']) / 2
+        feet_down = rewards.tolerance(feet_height,
+                                bounds=(-float('inf'), _YOGA_LIE_DOWN_HEIGHT),
+                                margin=_YOGA_LIE_DOWN_HEIGHT*1.5)
+        
+        l_thigh_foot_distance = max(0.1, abs(physics.named.data.xpos['left_foot', 'x'] - physics.named.data.xpos['left_thigh', 'x'])) - 0.1
+        r_thigh_foot_distance = max(0.1, abs(physics.named.data.xpos['right_foot', 'x'] - physics.named.data.xpos['right_thigh', 'x'])) - 0.1
+        close = np.exp(-(l_thigh_foot_distance + r_thigh_foot_distance)/2)
+        
+        legs = (3 * legs_down + feet_down) / 4
+        return torso_up * legs * close
+
+    def _urlb_flip_reward(self, physics):
+        standing = rewards.tolerance(physics.torso_height(),
+                                     bounds=(walker._STAND_HEIGHT, float('inf')),
+                                     margin=walker._STAND_HEIGHT / 2)
+        upright = (1 + physics.torso_upright()) / 2
+        stand_reward = (3 * standing + upright) / 4
+        move_reward = rewards.tolerance(physics.named.data.subtree_angmom['torso'][1], # physics.angmomentum(),
+                                        bounds=(_SPIN_SPEED, float('inf')),
+                                        margin=_SPIN_SPEED,
+                                        value_at_margin=0,
+                                        sigmoid='linear')
+        return stand_reward * (5 * move_reward + 1) / 6
+
+    def _flip_reward(self, physics):
+        thigh_height = (physics.named.data.xpos['left_thigh', 'z'] + physics.named.data.xpos['right_thigh', 'z']) / 2
+        thigh_up = rewards.tolerance(thigh_height,
+                                bounds=(_YOGA_STAND_HEIGHT, float('inf')),
+                                margin=_YOGA_STAND_HEIGHT/2)
+        feet_height = (physics.named.data.xpos['left_foot', 'z'] + physics.named.data.xpos['right_foot', 'z']) / 2
+        legs_up = rewards.tolerance(feet_height,
+                                bounds=(_YOGA_LEGS_UP_HEIGHT, float('inf')),
+                                margin=_YOGA_LEGS_UP_HEIGHT/2)
+        upside_down_reward = (3*legs_up + 2*thigh_up) / 5
+        if self._move_speed == 0:
+            return upside_down_reward
+        move_reward = rewards.tolerance(physics.named.data.subtree_angmom['torso'][1], # physics.angmomentum(),
+                                    bounds=(self._move_speed, float('inf')) if self._move_speed > 0 else (-float('inf'), self._move_speed),
+                                    margin=abs(self._move_speed)/2,
+                                    value_at_margin=0.5,
+                                    sigmoid='linear')
+        return upside_down_reward * (5*move_reward + 1) / 6
+    
+    def get_reward(self, physics):
+        if self._goal == 'arabesque':
+            return self._arabesque_reward(physics)
+        elif self._goal == 'lying_down':
+            return self._lying_down_reward(physics)
+        elif self._goal == 'legs_up':
+            return self._legs_up_reward(physics)
+        elif self._goal == 'flip':
+            return self._flip_reward(physics)
+        elif self._goal == 'flipping':
+            self._move_speed = abs(self._move_speed)
+            pos_rew = self._flip_reward(physics)
+            self._move_speed = -abs(self._move_speed)
+            neg_rew = self._flip_reward(physics)
+            return max(pos_rew, neg_rew)
+        elif self._goal == 'high_kick':
+            return self._high_kick_reward(physics)
+        elif self._goal == 'one_foot':
+            return self._one_foot_reward(physics)
+        elif self._goal == 'lunge_pose':
+            return self._lunge_pose_reward(physics)
+        elif self._goal == 'sit_knees':
+            return self._sit_knees_reward(physics)
+        elif self._goal == 'urlb_flip':
+            return self._urlb_flip_reward(physics)
+        else:
+            raise NotImplementedError(f'Goal {self._goal} is not implemented.')
+
+
+if __name__ == '__main__':
+    from dm_control import viewer
+    import numpy as np
+    
+    env = sit_knees()
+    env.task.visualize_reward = True
+
+    action_spec = env.action_spec()
+
+    def zero_policy(time_step):
+        print(time_step.reward)
+        return np.zeros(action_spec.shape)
+    viewer.launch(env, policy=zero_policy)
+    
+    # obs = env.reset()
+    # next_obs, reward, done, info = env.step(np.zeros(6))

envs/custom_dmc_tasks/walker.xml

@@ -0,0 +1,71 @@
+<mujoco model="planar walker">
+  <include file="./common/skybox.xml"/>
+  <include file="./common/visual.xml"/>
+  <include file="./common/materials.xml"/>
+
+
+  <option timestep="0.0025"/>
+
+  <statistic extent="2" center="0 0 1"/>
+
+  <default>
+    <joint damping=".1" armature="0.01" limited="true" solimplimit="0 .99 .01"/>
+    <geom contype="1" conaffinity="0" friction=".7 .1 .1"/>
+    <motor ctrlrange="-1 1" ctrllimited="true"/>
+    <site size="0.01"/>
+    <default class="walker">
+      <geom material="self" type="capsule"/>
+      <joint axis="0 -1 0"/>
+    </default>
+  </default>
+
+  <worldbody>
+    <geom name="floor" type="plane" conaffinity="1" pos="248 0 0" size="250 .8 .2" material="grid" zaxis="0 0 1"/>
+    <body name="torso" pos="0 0 1.3" childclass="walker">
+      <light name="light" pos="0 0 2" mode="trackcom"/>
+      <camera name="side" pos="0 -2 .7" euler="60 0 0" mode="trackcom"/>
+      <camera name="back" pos="-2 0 .5" xyaxes="0 -1 0 1 0 3" mode="trackcom"/>
+      <joint name="rootz" axis="0 0 1" type="slide" limited="false" armature="0" damping="0"/>
+      <joint name="rootx" axis="1 0 0" type="slide" limited="false" armature="0" damping="0"/>
+      <joint name="rooty" axis="0 1 0" type="hinge" limited="false" armature="0" damping="0"/>
+      <geom name="torso" size="0.07 0.3"/>
+      <body name="right_thigh" pos="0 -.05 -0.3">
+        <joint name="right_hip" range="-20 100"/>
+        <geom name="right_thigh" pos="0 0 -0.225" size="0.05 0.225"/>
+        <body name="right_leg" pos="0 0 -0.7">
+          <joint name="right_knee" pos="0 0 0.25" range="-150 0"/>
+          <geom name="right_leg" size="0.04 0.25"/>
+          <body name="right_foot" pos="0.06 0 -0.25">
+            <joint name="right_ankle" pos="-0.06 0 0" range="-45 45"/>
+            <geom name="right_foot" zaxis="1 0 0" size="0.05 0.1"/>
+          </body>
+        </body>
+      </body>
+      <body name="left_thigh" pos="0 .05 -0.3" >
+        <joint name="left_hip" range="-20 100"/>
+        <geom name="left_thigh" pos="0 0 -0.225" size="0.05 0.225"/>
+        <body name="left_leg" pos="0 0 -0.7">
+          <joint name="left_knee" pos="0 0 0.25" range="-150 0"/>
+          <geom name="left_leg" size="0.04 0.25"/>
+          <body name="left_foot" pos="0.06 0 -0.25">
+            <joint name="left_ankle" pos="-0.06 0 0" range="-45 45"/>
+            <geom name="left_foot" zaxis="1 0 0" size="0.05 0.1"/>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+
+  <sensor>
+    <subtreelinvel name="torso_subtreelinvel" body="torso"/>
+  </sensor>
+
+  <actuator>
+    <motor name="right_hip" joint="right_hip" gear="100"/>
+    <motor name="right_knee" joint="right_knee" gear="50"/>
+    <motor name="right_ankle" joint="right_ankle" gear="20"/>
+    <motor name="left_hip" joint="left_hip" gear="100"/>
+    <motor name="left_knee" joint="left_knee" gear="50"/>
+    <motor name="left_ankle" joint="left_ankle" gear="20"/>
+  </actuator>
+</mujoco>

envs/kitchen_extra.py

@@ -0,0 +1,299 @@
+"""Environments using kitchen and Franka robot."""
+import logging
+import sys
+from pathlib import Path
+sys.path.append((Path(__file__).parent.parent / 'third_party' / 'relay-policy-learning' / 'adept_envs').__str__())  
+import adept_envs
+from adept_envs.franka.kitchen_multitask_v0 import KitchenTaskRelaxV1
+import os
+import numpy as np
+from dm_control.mujoco import engine
+
+OBS_ELEMENT_INDICES = {
+    "bottom burner": np.array([11, 12]),
+    "top burner": np.array([15, 16]),
+    "light switch": np.array([17, 18]),
+    "slide cabinet": np.array([19]),
+    "hinge cabinet": np.array([20, 21]),
+    "microwave": np.array([22]),
+    "kettle": np.array([23, 24, 25, 26, 27, 28, 29]),
+}
+OBS_ELEMENT_GOALS = {
+    "bottom burner": np.array([-0.88, -0.01]),
+    "top burner": np.array([-0.92, -0.01]),
+    "light switch": np.array([-0.69, -0.05]),
+    "slide cabinet": np.array([0.37]),
+    "hinge cabinet": np.array([0.0, 1.45]),
+    "microwave": np.array([-0.75]),
+    "kettle": np.array([-0.23, 0.75, 1.62, 0.99, 0.0, 0.0, -0.06]),
+}
+BONUS_THRESH = 0.3
+
+logging.basicConfig(
+    level="INFO",
+    format="%(asctime)s [%(levelname)s] %(message)s",
+    filemode="w",
+)
+logger = logging.getLogger()
+
+XPOS_NAMES = {
+    "light switch" : "lightswitchroot",
+    "slide cabinet" : "slidelink",
+    "microwave" : "microdoorroot",
+    "kettle" : "kettle",
+}
+
+class KitchenBase(KitchenTaskRelaxV1):
+    # A string of element names. The robot's task is then to modify each of
+    # these elements appropriately.
+    TASK_ELEMENTS = []
+    ALL_TASKS = [
+        "bottom burner",
+        "top burner",
+        "light switch",
+        "slide cabinet",
+        "hinge cabinet",
+        "microwave",
+        "kettle",
+    ]
+    REMOVE_TASKS_WHEN_COMPLETE = True
+    TERMINATE_ON_TASK_COMPLETE = True
+    TERMINATE_ON_WRONG_COMPLETE = False
+    COMPLETE_IN_ANY_ORDER = (
+        True  # This allows for the tasks to be completed in arbitrary order.
+    )
+    GRIPPER_DISTANCE_REW = False
+
+    def __init__(
+        self, dense=True, dataset_url=None, ref_max_score=None, ref_min_score=None, **kwargs
+    ):
+        self.tasks_to_complete = list(self.TASK_ELEMENTS)
+        self.goal_masking = True
+        self.dense = dense
+        self.use_grasp_rewards = False
+
+        super(KitchenBase, self).__init__(**kwargs)
+
+    def set_goal_masking(self, goal_masking=True):
+        """Sets goal masking for goal-conditioned approaches (like RPL)."""
+        self.goal_masking = goal_masking
+
+    def _get_task_goal(self, task=None, actually_return_goal=False):
+        if task is None:
+            task = ["microwave", "kettle", "bottom burner", "light switch"]
+        new_goal = np.zeros_like(self.goal)
+        if self.goal_masking and not actually_return_goal:
+            return new_goal
+        for element in task:
+            element_idx = OBS_ELEMENT_INDICES[element]
+            element_goal = OBS_ELEMENT_GOALS[element]
+            new_goal[element_idx] = element_goal
+
+        return new_goal
+
+    def reset_model(self):
+        self.tasks_to_complete = list(self.TASK_ELEMENTS)
+        return super(KitchenBase, self).reset_model()
+
+    def _get_reward_n_score(self, obs_dict):
+        reward_dict, score = super(KitchenBase, self)._get_reward_n_score(obs_dict)
+        next_q_obs = obs_dict["qp"]
+        next_obj_obs = obs_dict["obj_qp"]
+        idx_offset = len(next_q_obs)
+        completions = []
+        dense = 0
+        if self.GRIPPER_DISTANCE_REW:
+            assert len(self.tasks_to_complete) == 1
+            element = next(iter(self.tasks_to_complete))
+            gripper_pos = (self.sim.named.data.xpos['panda0_leftfinger'] + self.sim.named.data.xpos['panda0_rightfinger']) / 2
+            object_pos = self.sim.named.data.xpos[XPOS_NAMES[element]]
+            gripper_obj_dist = np.linalg.norm(object_pos - gripper_pos)
+            if self.dense:
+                reward_dict["bonus"] = -gripper_obj_dist
+                reward_dict["r_total"] = -gripper_obj_dist
+                score = -gripper_obj_dist
+            else:
+                reward_dict["bonus"] = gripper_obj_dist < 0.15 
+                reward_dict["r_total"] = gripper_obj_dist < 0.15
+                score = gripper_obj_dist < 0.15               
+            return reward_dict, score
+        for element in self.tasks_to_complete:
+            element_idx = OBS_ELEMENT_INDICES[element]
+            distance = np.linalg.norm(
+                next_obj_obs[..., element_idx - idx_offset] - OBS_ELEMENT_GOALS[element]
+            )
+            dense += -1 * distance  # reward must be negative distance for RL
+            is_grasped = True
+            if not self.initializing and self.use_grasp_rewards:
+                if element == "slide cabinet":
+                    is_grasped = False
+                    for i in range(1, 6):
+                        obj_pos = self.get_site_xpos("schandle{}".format(i))
+                        left_pad = self.get_site_xpos("leftpad")
+                        right_pad = self.get_site_xpos("rightpad")
+                        within_sphere_left = np.linalg.norm(obj_pos - left_pad) < 0.07
+                        within_sphere_right = np.linalg.norm(obj_pos - right_pad) < 0.07
+                        right = right_pad[0] < obj_pos[0]
+                        left = obj_pos[0] < left_pad[0]
+                        if (
+                            right
+                            and left
+                            and within_sphere_right
+                            and within_sphere_left
+                        ):
+                            is_grasped = True
+                if element == "top left burner":
+                    is_grasped = False
+                    obj_pos = self.get_site_xpos("tlbhandle")
+                    left_pad = self.get_site_xpos("leftpad")
+                    right_pad = self.get_site_xpos("rightpad")
+                    within_sphere_left = np.linalg.norm(obj_pos - left_pad) < 0.035
+                    within_sphere_right = np.linalg.norm(obj_pos - right_pad) < 0.04
+                    right = right_pad[0] < obj_pos[0]
+                    left = obj_pos[0] < left_pad[0]
+                    if within_sphere_right and within_sphere_left and right and left:
+                        is_grasped = True
+                if element == "microwave":
+                    is_grasped = False
+                    for i in range(1, 6):
+                        obj_pos = self.get_site_xpos("mchandle{}".format(i))
+                        left_pad = self.get_site_xpos("leftpad")
+                        right_pad = self.get_site_xpos("rightpad")
+                        within_sphere_left = np.linalg.norm(obj_pos - left_pad) < 0.05
+                        within_sphere_right = np.linalg.norm(obj_pos - right_pad) < 0.05
+                        if (
+                            right_pad[0] < obj_pos[0]
+                            and obj_pos[0] < left_pad[0]
+                            and within_sphere_right
+                            and within_sphere_left
+                        ):
+                            is_grasped = True
+                if element == "hinge cabinet":
+                    is_grasped = False
+                    for i in range(1, 6):
+                        obj_pos = self.get_site_xpos("hchandle{}".format(i))
+                        left_pad = self.get_site_xpos("leftpad")
+                        right_pad = self.get_site_xpos("rightpad")
+                        within_sphere_left = np.linalg.norm(obj_pos - left_pad) < 0.06
+                        within_sphere_right = np.linalg.norm(obj_pos - right_pad) < 0.06
+                        if (
+                            right_pad[0] < obj_pos[0]
+                            and obj_pos[0] < left_pad[0]
+                            and within_sphere_right
+                        ):
+                            is_grasped = True
+                if element == "light switch":
+                    is_grasped = False
+                    for i in range(1, 4):
+                        obj_pos = self.get_site_xpos("lshandle{}".format(i))
+                        left_pad = self.get_site_xpos("leftpad")
+                        right_pad = self.get_site_xpos("rightpad")
+                        within_sphere_left = np.linalg.norm(obj_pos - left_pad) < 0.045
+                        within_sphere_right = np.linalg.norm(obj_pos - right_pad) < 0.03
+                        if within_sphere_right and within_sphere_left:
+                            is_grasped = True
+            complete = distance < BONUS_THRESH #  and is_grasped
+            if complete:
+                completions.append(element)
+        if self.REMOVE_TASKS_WHEN_COMPLETE:
+            [self.tasks_to_complete.remove(element) for element in completions]
+        bonus = float(len(completions))
+        reward_dict["bonus"] = bonus
+        reward_dict["r_total"] = bonus
+        if self.dense:
+            reward_dict["r_total"] = dense
+        score = bonus
+        return reward_dict, score
+
+    def step(self, a, b=None):
+        obs, reward, done, env_info = super(KitchenBase, self).step(a, b=b)
+        if self.TERMINATE_ON_TASK_COMPLETE:
+            done = not self.tasks_to_complete
+        if self.TERMINATE_ON_WRONG_COMPLETE:
+            all_goal = self._get_task_goal(task=self.ALL_TASKS)
+            for wrong_task in list(set(self.ALL_TASKS) - set(self.TASK_ELEMENTS)):
+                element_idx = OBS_ELEMENT_INDICES[wrong_task]
+                distance = np.linalg.norm(obs[..., element_idx] - all_goal[element_idx])
+                complete = distance < BONUS_THRESH
+                if complete:
+                    done = True
+                    break
+        env_info["completed_tasks"] = set(self.TASK_ELEMENTS) - set(
+            self.tasks_to_complete
+        )
+        return obs, reward, done, env_info
+
+    def get_goal(self):
+        """Loads goal state from dataset for goal-conditioned approaches (like RPL)."""
+        raise NotImplementedError
+
+    def _split_data_into_seqs(self, data):
+        """Splits dataset object into list of sequence dicts."""
+        seq_end_idxs = np.where(data["terminals"])[0]
+        start = 0
+        seqs = []
+        for end_idx in seq_end_idxs:
+            seqs.append(
+                dict(
+                    states=data["observations"][start : end_idx + 1],
+                    actions=data["actions"][start : end_idx + 1],
+                )
+            )
+            start = end_idx + 1
+        return seqs
+    
+    def render(self, mode='rgb_array', resolution=(64,64)):
+        if mode =='rgb_array':
+            camera = engine.MovableCamera(self.sim, *resolution)
+            camera.set_pose(distance=2.2, lookat=[-0.2, .5, 2.], azimuth=70, elevation=-35)
+            img = camera.render()
+            return img
+        else:
+            super(KitchenTaskRelaxV1, self).render()
+
+
+class KitchenSlideV0(KitchenBase):
+    TASK_ELEMENTS = ["slide cabinet",]
+    COMPLETE_IN_ANY_ORDER = False
+
+class KitchenHingeV0(KitchenBase):
+    TASK_ELEMENTS = ["hinge cabinet",]
+    COMPLETE_IN_ANY_ORDER = False
+
+class KitchenLightV0(KitchenBase):
+    TASK_ELEMENTS = ["light switch",]
+    COMPLETE_IN_ANY_ORDER = False
+
+class KitchenKettleV0(KitchenBase):
+    TASK_ELEMENTS = ["kettle",]
+    COMPLETE_IN_ANY_ORDER = False
+
+class KitchenMicrowaveV0(KitchenBase):
+    TASK_ELEMENTS = ["microwave",]
+    COMPLETE_IN_ANY_ORDER = False
+
+class KitchenBurnerV0(KitchenBase):
+    TASK_ELEMENTS = ["bottom burner",]
+    COMPLETE_IN_ANY_ORDER = False
+
+class KitchenTopBurnerV0(KitchenBase):
+    TASK_ELEMENTS = ["top burner",]
+    COMPLETE_IN_ANY_ORDER = False
+
+class KitchenMicrowaveKettleBottomBurnerLightV0(KitchenBase):
+    TASK_ELEMENTS = ["microwave", "kettle", "bottom burner", "light switch"]
+    COMPLETE_IN_ANY_ORDER = False
+
+
+class KitchenMicrowaveKettleLightSliderV0(KitchenBase):
+    TASK_ELEMENTS = ["microwave", "kettle", "light switch", "slide cabinet"]
+    COMPLETE_IN_ANY_ORDER = False
+
+
+class KitchenKettleMicrowaveLightSliderV0(KitchenBase):
+    TASK_ELEMENTS = ["kettle", "microwave", "light switch", "slide cabinet"]
+    COMPLETE_IN_ANY_ORDER = False
+
+
+class KitchenAllV0(KitchenBase):
+    TASK_ELEMENTS = KitchenBase.ALL_TASKS

envs/main.py

@@ -0,0 +1,743 @@
+from collections import OrderedDict, deque
+from typing import Any, NamedTuple
+import os
+
+import dm_env
+import numpy as np
+from dm_env import StepType, specs
+
+import gym
+import torch
+
+class ExtendedTimeStep(NamedTuple):
+    step_type: Any
+    reward: Any
+    discount: Any
+    observation: Any
+    action: Any
+
+    def first(self):
+        return self.step_type == StepType.FIRST
+
+    def mid(self):
+        return self.step_type == StepType.MID
+
+    def last(self):
+        return self.step_type == StepType.LAST
+
+    def __getitem__(self, attr):
+        return getattr(self, attr)
+
+
+class FlattenJacoObservationWrapper(dm_env.Environment):
+    def __init__(self, env):
+        self._env = env
+        self._obs_spec = OrderedDict()
+        wrapped_obs_spec = env.observation_spec().copy()
+        if 'front_close' in wrapped_obs_spec:
+            spec = wrapped_obs_spec['front_close']
+            # drop batch dim
+            self._obs_spec['pixels'] = specs.BoundedArray(shape=spec.shape[1:],
+                                                          dtype=spec.dtype,
+                                                          minimum=spec.minimum,
+                                                          maximum=spec.maximum,
+                                                          name='pixels')
+            wrapped_obs_spec.pop('front_close')
+
+        for key, spec in wrapped_obs_spec.items():
+            assert spec.dtype == np.float64
+            assert type(spec) == specs.Array
+        dim = np.sum(
+            np.fromiter((int(np.prod(spec.shape))
+                         for spec in wrapped_obs_spec.values()), np.int32))
+
+        self._obs_spec['observations'] = specs.Array(shape=(dim,),
+                                                     dtype=np.float32,
+                                                     name='observations')
+
+    def _transform_observation(self, time_step):
+        obs = OrderedDict()
+
+        if 'front_close' in time_step.observation:
+            pixels = time_step.observation['front_close']
+            time_step.observation.pop('front_close')
+            pixels = np.squeeze(pixels)
+            obs['pixels'] = pixels
+
+        features = []
+        for feature in time_step.observation.values():
+            features.append(feature.ravel())
+        obs['observations'] = np.concatenate(features, axis=0)
+        return time_step._replace(observation=obs)
+
+    def reset(self):
+        time_step = self._env.reset()
+        return self._transform_observation(time_step)
+
+    def step(self, action):
+        time_step = self._env.step(action)
+        return self._transform_observation(time_step)
+
+    def observation_spec(self):
+        return self._obs_spec
+
+    def action_spec(self):
+        return self._env.action_spec()
+
+    def __getattr__(self, name):
+        return getattr(self._env, name)
+
+
+class ActionRepeatWrapper(dm_env.Environment):
+    def __init__(self, env, num_repeats):
+        self._env = env
+        self._num_repeats = num_repeats
+
+    def step(self, action):
+        reward = 0.0
+        discount = 1.0
+        for i in range(self._num_repeats):
+            time_step = self._env.step(action)
+            reward += (time_step.reward or 0.0) * discount
+            discount *= time_step.discount
+            if time_step.last():
+                break
+
+        return time_step._replace(reward=reward, discount=discount)
+
+    def observation_spec(self):
+        return self._env.observation_spec()
+
+    def action_spec(self):
+        return self._env.action_spec()
+
+    def reset(self):
+        return self._env.reset()
+
+    def __getattr__(self, name):
+        return getattr(self._env, name)
+
+
+class FramesWrapper(dm_env.Environment):
+    def __init__(self, env, num_frames=1, pixels_key='pixels'):
+        self._env = env
+        self._num_frames = num_frames
+        self._frames = deque([], maxlen=num_frames)
+        self._pixels_key = pixels_key
+
+        wrapped_obs_spec = env.observation_spec()
+        assert pixels_key in wrapped_obs_spec
+
+        pixels_shape = wrapped_obs_spec[pixels_key].shape
+        # remove batch dim
+        if len(pixels_shape) == 4:
+            pixels_shape = pixels_shape[1:]
+        self._obs_spec = specs.BoundedArray(shape=np.concatenate(
+            [[pixels_shape[2] * num_frames], pixels_shape[:2]], axis=0),
+                                            dtype=np.uint8,
+                                            minimum=0,
+                                            maximum=255,
+                                            name='observation')
+
+    def _transform_observation(self, time_step):
+        assert len(self._frames) == self._num_frames
+        obs = np.concatenate(list(self._frames), axis=0)
+        return time_step._replace(observation=obs)
+
+    def _extract_pixels(self, time_step):
+        pixels = time_step.observation[self._pixels_key]
+        # remove batch dim
+        if len(pixels.shape) == 4:
+            pixels = pixels[0]
+        return pixels.transpose(2, 0, 1).copy()
+
+    def reset(self):
+        time_step = self._env.reset()
+        pixels = self._extract_pixels(time_step)
+        for _ in range(self._num_frames):
+            self._frames.append(pixels)
+        return self._transform_observation(time_step)
+
+    def step(self, action):
+        time_step = self._env.step(action)
+        pixels = self._extract_pixels(time_step)
+        self._frames.append(pixels)
+        return self._transform_observation(time_step)
+
+    def observation_spec(self):
+        return self._obs_spec
+
+    def action_spec(self):
+        return self._env.action_spec()
+
+    def __getattr__(self, name):
+        return getattr(self._env, name)
+
+class OneHotAction(gym.Wrapper):
+    def __init__(self, env):
+        assert isinstance(env.action_space, gym.spaces.Discrete)
+        super().__init__(env)
+        self._random = np.random.RandomState()
+        shape = (self.env.action_space.n,)
+        space = gym.spaces.Box(low=0, high=1, shape=shape, dtype=np.float32)
+        space.discrete = True
+        self.action_space = space
+
+    def step(self, action):
+        index = np.argmax(action).astype(int)
+        reference = np.zeros_like(action)
+        reference[index] = 1
+        if not np.allclose(reference, action):
+            raise ValueError(f"Invalid one-hot action:\n{action}")
+        return self.env.step(index)
+
+    def reset(self):
+        return self.env.reset()
+
+    def _sample_action(self):
+        actions = self.env.action_space.n
+        index = self._random.randint(0, actions)
+        reference = np.zeros(actions, dtype=np.float32)
+        reference[index] = 1.0
+        return reference
+
+class ActionDTypeWrapper(dm_env.Environment):
+    def __init__(self, env, dtype):
+        self._env = env
+        wrapped_action_spec = env.action_spec()
+        self._action_spec = specs.BoundedArray(wrapped_action_spec.shape,
+                                               dtype,
+                                               wrapped_action_spec.minimum,
+                                               wrapped_action_spec.maximum,
+                                               'action')
+
+    def step(self, action):
+        action = action.astype(self._env.action_spec().dtype)
+        return self._env.step(action)
+
+    def observation_spec(self):
+        return self._env.observation_spec()
+
+    def action_spec(self):
+        return self._action_spec
+
+    def reset(self):
+        return self._env.reset()
+
+    def __getattr__(self, name):
+        return getattr(self._env, name)
+
+
+class ObservationDTypeWrapper(dm_env.Environment):
+    def __init__(self, env, dtype):
+        self._env = env
+        self._dtype = dtype
+        wrapped_obs_spec = env.observation_spec()['observations']
+        self._obs_spec = specs.Array(wrapped_obs_spec.shape, dtype,
+                                     'observation')
+
+    def _transform_observation(self, time_step):
+        obs = time_step.observation['observations'].astype(self._dtype)
+        return time_step._replace(observation=obs)
+
+    def reset(self):
+        time_step = self._env.reset()
+        return self._transform_observation(time_step)
+
+    def step(self, action):
+        time_step = self._env.step(action)
+        return self._transform_observation(time_step)
+
+    def observation_spec(self):
+        return self._obs_spec
+
+    def action_spec(self):
+        return self._env.action_spec()
+
+    def __getattr__(self, name):
+        return getattr(self._env, name)
+
+
+class ExtendedTimeStepWrapper(dm_env.Environment):
+    def __init__(self, env):
+        self._env = env
+
+    def reset(self):
+        time_step = self._env.reset()
+        return self._augment_time_step(time_step)
+
+    def step(self, action):
+        time_step = self._env.step(action)
+        return self._augment_time_step(time_step, action)
+
+    def _augment_time_step(self, time_step, action=None):
+        if action is None:
+            action_spec = self.action_spec()
+            action = np.zeros(action_spec.shape, dtype=action_spec.dtype)
+        return ExtendedTimeStep(observation=time_step.observation,
+                                step_type=time_step.step_type,
+                                action=action,
+                                reward=time_step.reward or 0.0,
+                                discount=time_step.discount or 1.0)
+
+    def observation_spec(self):
+        return self._env.observation_spec()
+
+    def action_spec(self):
+        return self._env.action_spec()
+
+    def __getattr__(self, name):
+        return getattr(self._env, name)
+
+class DMC:
+  def __init__(self, env):
+    self._env = env 
+    self._ignored_keys = []
+
+  def step(self, action):
+    time_step = self._env.step(action)
+    assert time_step.discount in (0, 1)
+    obs = {
+        'reward': time_step.reward,
+        'is_first': False,
+        'is_last': time_step.last(),
+        'is_terminal': time_step.discount == 0,
+        'observation': time_step.observation,
+        'action' : action,
+        'discount': time_step.discount
+    }
+    return time_step, obs 
+
+  def reset(self):
+    time_step = self._env.reset()
+    obs = {
+        'reward': 0.0,
+        'is_first': True,
+        'is_last': False,
+        'is_terminal': False,
+        'observation': time_step.observation,
+        'action' : np.zeros_like(self.act_space['action'].sample()),
+        'discount': time_step.discount
+    }
+    return time_step, obs
+
+  def __getattr__(self, name):
+    if name == 'obs_space':
+        obs_spaces = {
+            'observation': self._env.observation_spec(), 
+            'is_first': gym.spaces.Box(0, 1, (), dtype=bool),
+            'is_last': gym.spaces.Box(0, 1, (), dtype=bool),
+            'is_terminal': gym.spaces.Box(0, 1, (), dtype=bool),
+        }
+        return obs_spaces
+    if name == 'act_space':
+        spec = self._env.action_spec()
+        action = gym.spaces.Box((spec.minimum)*spec.shape[0], (spec.maximum)*spec.shape[0], shape=spec.shape, dtype=np.float32)
+        act_space = {'action': action}
+        return act_space
+    return getattr(self._env, name)
+
+
+class OneHotAction(gym.Wrapper):
+    def __init__(self, env):
+        assert isinstance(env.action_space, gym.spaces.Discrete)
+        super().__init__(env)
+        self._random = np.random.RandomState()
+        shape = (self.env.action_space.n,)
+        space = gym.spaces.Box(low=0, high=1, shape=shape, dtype=np.float32)
+        space.discrete = True
+        self.action_space = space
+
+    def step(self, action):
+        index = np.argmax(action).astype(int)
+        reference = np.zeros_like(action)
+        reference[index] = 1
+        if not np.allclose(reference, action):
+            raise ValueError(f"Invalid one-hot action:\n{action}")
+        return self.env.step(index)
+
+    def reset(self):
+        return self.env.reset()
+
+    def _sample_action(self):
+        actions = self.env.action_space.n
+        index = self._random.randint(0, actions)
+        reference = np.zeros(actions, dtype=np.float32)
+        reference[index] = 1.0
+        return reference
+
+class KitchenWrapper:
+    def __init__(
+        self,
+        name,
+        seed=0,
+        action_repeat=1,
+        size=(64, 64),
+    ):
+        import envs.kitchen_extra as kitchen_extra
+        self._env  = {
+            'microwave' : kitchen_extra.KitchenMicrowaveV0,
+            'kettle' : kitchen_extra.KitchenKettleV0,
+            'burner' : kitchen_extra.KitchenBurnerV0,
+            'light'  : kitchen_extra.KitchenLightV0,
+            'hinge'  : kitchen_extra.KitchenHingeV0,
+            'slide'  : kitchen_extra.KitchenSlideV0,
+            'top_burner' : kitchen_extra.KitchenTopBurnerV0,
+        }[name]()
+            
+        self._size = size
+        self._action_repeat = action_repeat
+        self._seed = seed
+        self._eval = False
+
+    def eval_mode(self,):
+        self._env.dense = False
+        self._eval = True
+
+    @property
+    def obs_space(self):
+        spaces = {
+            "observation": gym.spaces.Box(0, 255, (3,) + self._size, dtype=np.uint8),
+            "is_first": gym.spaces.Box(0, 1, (), dtype=bool),
+            "is_last": gym.spaces.Box(0, 1, (), dtype=bool),
+            "is_terminal": gym.spaces.Box(0, 1, (), dtype=bool),
+            "state": self._env.observation_space,
+        }
+        return spaces
+
+    @property
+    def act_space(self):
+        action = self._env.action_space
+        return {"action": action}
+
+    def step(self, action):
+        # assert np.isfinite(action["action"]).all(), action["action"]
+        reward = 0.0
+        for _ in range(self._action_repeat):
+            state, rew, done, info = self._env.step(action.copy())
+            reward += rew 
+        obs = {
+            "reward": reward,
+            "is_first": False,
+            "is_last": False,  # will be handled by timelimit wrapper
+            "is_terminal": False,  # will be handled by per_episode function
+            "observation": info['images'].transpose(2, 0, 1).copy(),
+            "state": state.astype(np.float32),
+            'action' : action,
+            'discount' : 1
+        }
+        if self._eval:
+            obs['reward'] = min(obs['reward'], 1)
+            if obs['reward'] > 0:
+                obs['is_last'] = True
+        return dm_env.TimeStep(
+                step_type=dm_env.StepType.MID if not obs['is_last'] else dm_env.StepType.LAST, 
+                reward=obs['reward'],
+                discount=1,
+                observation=obs['observation']), obs
+
+    def reset(self,):
+        state = self._env.reset()
+        obs = {
+            "reward": 0.0,
+            "is_first": True,
+            "is_last": False,
+            "is_terminal": False,
+            "observation": self.get_visual_obs(self._size),
+            "state": state.astype(np.float32),
+            'action' : np.zeros_like(self.act_space['action'].sample()),
+            'discount' : 1
+        }
+        return dm_env.TimeStep(
+                step_type=dm_env.StepType.FIRST,
+                reward=None,
+                discount=None,
+                observation=obs['observation']), obs
+
+    def __getattr__(self, name):
+        if name == 'obs_space':
+            return self.obs_space
+        if name == 'act_space':
+            return self.act_space
+        return getattr(self._env, name)
+    
+    def get_visual_obs(self, resolution):
+        img = self._env.render(resolution=resolution,).transpose(2, 0, 1).copy()
+        return img
+
+class ViClipWrapper:
+    def __init__(self, env, hd_rendering=False, device='cuda'):
+        self._env = env
+        try:
+            from tools.genrl_utils import viclip_global_instance
+        except:
+            from tools.genrl_utils import ViCLIPGlobalInstance
+            viclip_global_instance = ViCLIPGlobalInstance()
+
+        if not viclip_global_instance._instantiated:
+            viclip_global_instance.instantiate(device)
+        self.viclip_model = viclip_global_instance.viclip
+        self.n_frames = self.viclip_model.n_frames
+        self.viclip_emb_dim = viclip_global_instance.viclip_emb_dim
+        self.n_frames = self.viclip_model.n_frames
+        self.buffer = deque(maxlen=self.n_frames)
+        # NOTE: these are hardcoded for now, as they are the best settings
+        self.accumulate = True
+        self.accumulate_buffer = []
+        self.anticipate_conv1 = False
+        self.hd_rendering = hd_rendering
+
+    def hd_render(self, obs):
+        if not self.hd_rendering:
+            return obs['observation']
+        if self._env._domain_name in ['mw', 'kitchen', 'mujoco']:
+            return self.get_visual_obs((224,224,))
+        else:
+            render_kwargs = {**getattr(self, '_render_kwargs', {})}
+            render_kwargs.update({'width' : 224, 'height' : 224})
+            return self._env.physics.render(**render_kwargs).transpose(2,0,1)
+
+    def preprocess(self, x):
+        return x
+
+    def process_accumulate(self, process_at_once=4): # NOTE: this could be varied for increasing FPS, depending on the size of the GPU
+        self.accumulate = False
+        x = np.stack(self.accumulate_buffer, axis=0)
+        # Splitting in chunks
+        chunks = []
+        chunk_idxs = list(range(0, x.shape[0] + 1, process_at_once))
+        if chunk_idxs[-1] != x.shape[0]:
+            chunk_idxs.append(x.shape[0])
+        start = 0
+        for end in chunk_idxs[1:]:
+            embeds = self.clip_process(x[start:end], bypass=True)
+            chunks.append(embeds.cpu())
+            start = end
+        embeds = torch.cat(chunks, dim=0)
+        assert embeds.shape[0] == len(self.accumulate_buffer)
+        self.accumulate = True
+        self.accumulate_buffer = []
+        return [*embeds.cpu().numpy()], 'clip_video'
+    
+    def process_episode(self, obs, process_at_once=8):
+        self.accumulate = False
+        sequences = []
+        for j in range(obs.shape[0] - self.n_frames + 1):
+            sequences.append(obs[j:j+self.n_frames].copy())
+        sequences = np.stack(sequences, axis=0)
+
+        idx_start = 0
+        clip_vid = []
+        for idx_end in range(process_at_once, sequences.shape[0] + process_at_once, process_at_once):
+            x = sequences[idx_start:idx_end]
+            with torch.no_grad(): # , torch.cuda.amp.autocast():
+                x = self.clip_process(x, bypass=True) 
+            clip_vid.append(x)
+            idx_start = idx_end
+        if len(clip_vid) == 1: # process all at once
+            embeds = clip_vid[0]
+        else:
+            embeds = torch.cat(clip_vid, dim=0)
+        pad = torch.zeros( (self.n_frames - 1, *embeds.shape[1:]), device=embeds.device, dtype=embeds.dtype)
+        embeds = torch.cat([pad, embeds], dim=0)
+        assert embeds.shape[0] == obs.shape[0], f"Shapes are different {embeds.shape[0]} {obs.shape[0]}"
+        return embeds.cpu().numpy()
+
+    def get_sequence(self,):
+        return np.expand_dims(np.stack(self.buffer, axis=0), axis=0)
+    
+    def clip_process(self, x, bypass=False):
+        if len(self.buffer) == self.n_frames or bypass:
+            if self.accumulate:
+                self.accumulate_buffer.append(self.preprocess(x)[0])
+                return torch.zeros(self.viclip_emb_dim)
+            with torch.no_grad():
+                B, n_frames, C, H, W = x.shape
+                obs = torch.from_numpy(x.copy().reshape(B * n_frames, C, H, W)).to(self.viclip_model.device)
+                processed_obs = self.viclip_model.preprocess_transf(obs / 255)
+                reshaped_obs = processed_obs.reshape(B, n_frames, 3,processed_obs.shape[-2],processed_obs.shape[-1])
+                video_embed = self.viclip_model.get_vid_features(reshaped_obs)
+            return video_embed.detach()
+        else:
+            return torch.zeros(self.viclip_emb_dim)
+
+    def step(self, action):
+        ts, obs = self._env.step(action)
+        self.buffer.append(self.hd_render(obs))
+        obs['clip_video'] = self.clip_process(self.get_sequence()).cpu().numpy()
+        return ts, obs
+
+    def reset(self,):
+        # Important to reset the buffer        
+        self.buffer = deque(maxlen=self.n_frames)
+
+        ts, obs = self._env.reset()
+        self.buffer.append(self.hd_render(obs))
+        obs['clip_video'] = self.clip_process(self.get_sequence()).cpu().numpy()
+        return ts, obs
+
+    def __getattr__(self, name):
+        if name == 'obs_space':
+            space = self._env.obs_space
+            space['clip_video'] = gym.spaces.Box(-np.inf, np.inf, (self.viclip_emb_dim,), dtype=np.float32)  
+            return space
+        return getattr(self._env, name)
+
+class TimeLimit:
+
+  def __init__(self, env, duration):
+    self._env = env
+    self._duration = duration
+    self._step = None
+
+  def __getattr__(self, name):
+    if name.startswith('__'):
+      raise AttributeError(name)
+    return getattr(self._env, name)
+
+  def step(self, action):
+    assert self._step is not None, 'Must reset environment.'
+    ts, obs = self._env.step(action)
+    self._step += 1
+    if self._duration and self._step >= self._duration:
+      ts = dm_env.TimeStep(dm_env.StepType.LAST, ts.reward, ts.discount, ts.observation)
+      obs['is_last'] = True
+      self._step = None
+    return ts, obs
+
+  def reset(self):
+    self._step = 0
+    return self._env.reset()
+
+  def reset_with_task_id(self, task_id):
+    self._step = 0
+    return self._env.reset_with_task_id(task_id)
+  
+class ClipActionWrapper:
+
+  def __init__(self, env, low=-1.0, high=1.0):
+    self._env = env
+    self._low = low
+    self._high = high
+
+  def __getattr__(self, name):
+    if name.startswith('__'):
+      raise AttributeError(name)
+    return getattr(self._env, name)
+
+  def step(self, action):
+    clipped_action = np.clip(action, self._low, self._high)
+    return self._env.step(clipped_action)
+
+  def reset(self):
+    self._step = 0
+    return self._env.reset()
+
+  def reset_with_task_id(self, task_id):
+    self._step = 0
+    return self._env.reset_with_task_id(task_id)
+
+class NormalizeAction:
+
+  def __init__(self, env, key='action'):
+    self._env = env
+    self._key = key
+    space = env.act_space[key]
+    self._mask = np.isfinite(space.low) & np.isfinite(space.high)
+    self._low = np.where(self._mask, space.low, -1)
+    self._high = np.where(self._mask, space.high, 1)
+
+  def __getattr__(self, name):
+    if name.startswith('__'):
+      raise AttributeError(name)
+    try:
+      return getattr(self._env, name)
+    except AttributeError:
+      raise ValueError(name)
+
+  @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 = gym.spaces.Box(low, high, dtype=np.float32)
+    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})
+
+def _make_jaco(obs_type, domain, task, action_repeat, seed, img_size,):
+    import envs.custom_dmc_tasks as cdmc
+    env = cdmc.make_jaco(task, obs_type, seed, img_size,)
+    env = ActionDTypeWrapper(env, np.float32)
+    env = ActionRepeatWrapper(env, action_repeat)
+    env = FlattenJacoObservationWrapper(env)
+    env._size = (img_size, img_size)
+    return env
+
+
+def _make_dmc(obs_type, domain, task, action_repeat, seed, img_size,):
+    visualize_reward = False
+    from dm_control import manipulation, suite
+    import envs.custom_dmc_tasks as cdmc
+
+    if (domain, task) in suite.ALL_TASKS:
+        env = suite.load(domain,
+                         task,
+                         task_kwargs=dict(random=seed),
+                         environment_kwargs=dict(flat_observation=True),
+                         visualize_reward=visualize_reward)
+    else:
+        env = cdmc.make(domain,
+                        task,
+                        task_kwargs=dict(random=seed),
+                        environment_kwargs=dict(flat_observation=True),
+                        visualize_reward=visualize_reward)
+    env = ActionDTypeWrapper(env, np.float32)
+    env = ActionRepeatWrapper(env, action_repeat)
+    if obs_type == 'pixels':
+        from dm_control.suite.wrappers import pixels
+        # zoom in camera for quadruped
+        camera_id = dict(locom_rodent=1,quadruped=2).get(domain, 0)
+        render_kwargs = dict(height=img_size, width=img_size, camera_id=camera_id)
+        env = pixels.Wrapper(env,
+                             pixels_only=True,
+                             render_kwargs=render_kwargs)
+        env._size = (img_size, img_size)
+        env._camera = camera_id
+    return env
+
+
+def make(name, obs_type, action_repeat, seed, img_size=64, viclip_encode=False, clip_hd_rendering=False, device='cuda'):
+    assert obs_type in ['states', 'pixels']
+    domain, task = name.split('_', 1)
+    if domain == 'kitchen':
+        env = TimeLimit(KitchenWrapper(task, seed=seed, action_repeat=action_repeat, size=(img_size,img_size)), 280 // action_repeat)
+    else:
+        os.environ['PYOPENGL_PLATFORM'] = 'egl' 
+        os.environ['MUJOCO_GL'] = 'egl'
+
+        domain = dict(cup='ball_in_cup', point='point_mass').get(domain, domain)
+
+        make_fn = _make_jaco if domain == 'jaco' else _make_dmc
+        env = make_fn(obs_type, domain, task, action_repeat, seed, img_size,)
+
+        if obs_type == 'pixels':
+            env = FramesWrapper(env,)
+        else:
+            env = ObservationDTypeWrapper(env, np.float32)
+
+        from dm_control.suite.wrappers import action_scale
+        env = action_scale.Wrapper(env, minimum=-1.0, maximum=+1.0)
+        env = ExtendedTimeStepWrapper(env)
+
+        env =  DMC(env)
+    env._domain_name = domain
+    
+    if isinstance(env.act_space['action'], gym.spaces.Box):
+        env = ClipActionWrapper(env,)
+
+    if viclip_encode:
+        env = ViClipWrapper(env, hd_rendering=clip_hd_rendering, device=device)
+    return env

notebooks/demo_videoclip.ipynb

@@ -0,0 +1,124 @@
+{
+ "cells": [
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "# InternVideo 2 demo\n",
+    "\n",
+    "It can be used to test the capabilities of InternVideo2 and to verify that the models are loaded correctly"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {
+    "scrolled": true
+   },
+   "outputs": [],
+   "source": [
+    "import pathlib\n",
+    "import sys\n",
+    "import os\n",
+    "sys.path.append(str(pathlib.Path(os.path.abspath('')).parent))\n",
+    "\n",
+    "from tools.genrl_utils import viclip_global_instance\n",
+    "viclip_global_instance.instantiate()"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "import cv2\n",
+    "import numpy as np\n",
+    "import torch\n",
+    "from tools.genrl_utils import INTERNVIDEO_PATH\n",
+    "\n",
+    "def _frame_from_video(video):\n",
+    "    while video.isOpened():\n",
+    "        success, frame = video.read()\n",
+    "        if success:\n",
+    "            yield frame\n",
+    "        else:\n",
+    "            break\n",
+    "\n",
+    "ASSET_PATH = pathlib.Path(os.path.abspath('')).parent / 'assets'\n",
+    "\n",
+    "# 83 % -  A man in a gray sweater plays fetch with his dog in the snowy yard, throwing a toy and watching it run.\n",
+    "video = cv2.VideoCapture( str(INTERNVIDEO_PATH / 'InternVideo2/multi_modality/demo/example1.mp4') )\n",
+    "# # 99 % -  A karate kick\n",
+    "# video = cv2.VideoCapture( str( ASSET_PATH / 'video_samples/karate_kick.mp4') ) \n",
+    "# # 99 % -  A headstand\n",
+    "# video = cv2.VideoCapture( str( ASSET_PATH / 'video_samples/headstand.mp4') ) \n",
+    "\n",
+    "frames = [x for x in _frame_from_video(video)]\n",
+    "processed_frames = viclip_global_instance.viclip.preprocess_transf(torch.from_numpy(np.stack(frames[:8], axis=0)).permute(0,3,1,2) / 255)\n",
+    "frames_tensor = processed_frames.reshape(1, 8, 3, 224,224)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "text_candidates = [\"A playful dog and its owner wrestle in the snowy yard, chasing each other with joyous abandon.\",\n",
+    "                   \"A man in a gray coat walks through the snowy landscape, pulling a sleigh loaded with toys.\",\n",
+    "                   \"A person dressed in a blue jacket shovels the snow-covered pavement outside their house.\",\n",
+    "                   \"A pet dog excitedly runs through the snowy yard, chasing a toy thrown by its owner.\",\n",
+    "                   \"A person stands on the snowy floor, pushing a sled loaded with blankets, preparing for a fun-filled ride.\",\n",
+    "                   \"A man in a gray hat and coat walks through the snowy yard, carefully navigating around the trees.\",\n",
+    "                   \"A playful dog slides down a snowy hill, wagging its tail with delight.\",\n",
+    "                   \"A person in a blue jacket walks their pet on a leash, enjoying a peaceful winter walk among the trees.\",\n",
+    "                   \"A man in a gray sweater plays fetch with his dog in the snowy yard, throwing a toy and watching it run.\",\n",
+    "                   \"A person bundled up in a blanket walks through the snowy landscape, enjoying the serene winter scenery.\",\n",
+    "                   \"A person playing with a kid in the street\",\n",
+    "                   \"A group of friends playing bowling.\",\n",
+    "                   \"A japanese girl eating noodles\",\n",
+    "                   \"A painting by Monet\",\n",
+    "                   \"A karate kick\",\n",
+    "                   \"A headstand\"]"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "text_feat = viclip_global_instance.viclip.get_txt_feat(text_candidates)\n",
+    "video_feat = viclip_global_instance.viclip.get_vid_features(frames_tensor.to(viclip_global_instance.viclip.device))\n",
+    "\n",
+    "sorted_probs, sorted_idxs = (100.0 * video_feat @ text_feat.T).softmax(dim=-1)[0].topk(len(text_feat))\n",
+    "\n",
+    "for p, i in zip(sorted_probs, sorted_idxs):\n",
+    "    if p > 0.01:\n",
+    "        print(int(p * 100), '% - ', text_candidates[i])"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3 (ipykernel)",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 4
+}

notebooks/text2video.ipynb

@@ -0,0 +1,161 @@
+{
+ "cells": [
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from pathlib import Path \n",
+    "import os\n",
+    "import sys\n",
+    "sys.path.append(str(Path(os.path.abspath('')).parent))\n",
+    "\n",
+    "import torch\n",
+    "import numpy as np\n",
+    "\n",
+    "import matplotlib.pyplot as plt\n",
+    "import matplotlib.animation as animation\n",
+    "\n",
+    "agent_path = Path(os.path.abspath('')).parent / 'models' / 'genrl_stickman_500k_2.pt'\n",
+    "print(\"Model path\", agent_path)\n",
+    "\n",
+    "agent = torch.load(agent_path)"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "from tools.genrl_utils import ViCLIPGlobalInstance, DOMAIN2PREDICATES\n",
+    "model_name = getattr(agent.cfg, 'viclip_model', 'viclip')\n",
+    "# Get ViCLIP\n",
+    "if 'viclip_global_instance' not in locals() or model_name != viclip_global_instance._model:\n",
+    "    viclip_global_instance = ViCLIPGlobalInstance(model_name)\n",
+    "    if not viclip_global_instance._instantiated:\n",
+    "        print(\"Instantiating\")\n",
+    "        viclip_global_instance.instantiate()\n",
+    "    clip = viclip_global_instance.viclip\n",
+    "    tokenizer = viclip_global_instance.viclip_tokenizer"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": null,
+   "metadata": {},
+   "outputs": [],
+   "source": [
+    "SAVE = True\n",
+    "DENOISE = True\n",
+    "REVERSE = False\n",
+    "REPEAT_TIME = 2 # standard is n_frames for = 1 \n",
+    "TEXT_OVERLAY = True\n",
+    "\n",
+    "domain = agent.cfg.task.split('_')\n",
+    "\n",
+    "labels_list = ['high kick', 'stand up straight', 'doing splits']\n",
+    "\n",
+    "with torch.no_grad():\n",
+    "    wm = world_model = agent.wm\n",
+    "    connector = agent.wm.connector\n",
+    "    decoder = world_model.heads['decoder']\n",
+    "    n_frames = connector.n_frames\n",
+    "    \n",
+    "    # Get text(video) embed\n",
+    "    text_feat = []\n",
+    "    for text in labels_list:\n",
+    "        with torch.no_grad():\n",
+    "            text_feat.append(clip.get_txt_feat(text,))\n",
+    "    text_feat = torch.stack(text_feat, dim=0).to(clip.device)\n",
+    "\n",
+    "    video_embed = text_feat\n",
+    "\n",
+    "    B = video_embed.shape[0]\n",
+    "    T = 1\n",
+    "\n",
+    "    # Get initial state\n",
+    "    init = connector.initial(B, init_embed=video_embed)\n",
+    "\n",
+    "    # Get actions\n",
+    "    video_embed = video_embed.repeat(1,n_frames, 1)\n",
+    "    action = wm.connector.get_action(video_embed)\n",
+    "\n",
+    "    with torch.no_grad():\n",
+    "        # Imagine\n",
+    "        prior = wm.connector.video_imagine(video_embed, None, sample=False, reset_every_n_frames=False, denoise=DENOISE)\n",
+    "        # Decode\n",
+    "        prior_recon = decoder(wm.decoder_input_fn(prior))['observation'].mean + 0.5\n",
+    "\n",
+    "    # Plotting video\n",
+    "    R = int(np.sqrt(B))\n",
+    "    C = min((B + (R-1)) // R, B) \n",
+    "\n",
+    "    fig, axes = plt.subplots(R, C, figsize=(3.5 * C, 4 * R))\n",
+    "    fig.subplots_adjust(top = 1, bottom = 0, right = 1, left = 0, hspace = 0, wspace = 0)\n",
+    "    fig.set_size_inches(4,4)\n",
+    "    \n",
+    "    if B == 1:\n",
+    "        axes = [[axes]]\n",
+    "    elif R == 1:\n",
+    "        axes = [axes] \n",
+    "    axes = [ a for row in axes for a in row]\n",
+    "\n",
+    "    file_path = f'temp_text2video.gif'\n",
+    "\n",
+    "    if SAVE:\n",
+    "        ims = []\n",
+    "        for t in range(prior_recon.shape[1]):\n",
+    "            if t == 0 :\n",
+    "                continue\n",
+    "            toadd = []\n",
+    "            for b in range(prior_recon.shape[0]):\n",
+    "                ax = axes[b]\n",
+    "                ax.set_axis_off()\n",
+    "                img = np.clip(prior_recon[b, t if not REVERSE else -t].cpu().permute(1,2,0), 0, 1)\n",
+    "                frame =  ax.imshow(img)\n",
+    "                if TEXT_OVERLAY:   \n",
+    "                    test = ax.text(0,5, labels_list[b], color='white')\n",
+    "                toadd.append(frame) # add both the image and the text to the list of artists \n",
+    "            ims.append(toadd)\n",
+    "\n",
+    "        # Save GIFs\n",
+    "        anim = animation.ArtistAnimation(fig, ims, interval=700, blit=True, repeat_delay=700)\n",
+    "        writer = animation.PillowWriter(fps=15, metadata=dict(artist='Me'), bitrate=1800)\n",
+    "        domain = agent.cfg.task.split('_')[0]\n",
+    "        os.makedirs(f'videos/{domain}/text2video', exist_ok=True)\n",
+    "        file_path = f'videos/{domain}/text2video/{\"_\".join(labels_list).replace(\" \",\"_\")}.gif'\n",
+    "        print(\"GIF path: \", Path(os.path.abspath('')) / file_path)\n",
+    "        anim.save(file_path, writer=writer)"
+   ]
+  }
+ ],
+ "metadata": {
+  "kernelspec": {
+   "display_name": "Python 3.8.10 ('base')",
+   "language": "python",
+   "name": "python3"
+  },
+  "language_info": {
+   "codemirror_mode": {
+    "name": "ipython",
+    "version": 3
+   },
+   "file_extension": ".py",
+   "mimetype": "text/x-python",
+   "name": "python",
+   "nbconvert_exporter": "python",
+   "pygments_lexer": "ipython3",
+   "version": "3.10.14"
+  },
+  "orig_nbformat": 4,
+  "vscode": {
+   "interpreter": {
+    "hash": "3d597f4c481aa0f25dceb95d2a0067e73c0966dcbd003d741d821a7208527ecf"
+   }
+  }
+ },
+ "nbformat": 4,
+ "nbformat_minor": 2
+}