Add files using upload-large-folder tool

VRL3/LICENSE

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

VRL3/src/cfgs_adroit/task/door.yaml

@@ -0,0 +1,4 @@
+num_train_frames: 4100000
+task_name: door-v0
+agent:
+  encoder_lr_scale: 1

VRL3/src/cfgs_adroit/task/relocate.yaml

@@ -0,0 +1,4 @@
+num_train_frames: 4100000
+task_name: relocate-v0
+agent:
+  encoder_lr_scale: 0.01

VRL3/src/logger.py

@@ -0,0 +1,182 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import csv
+import datetime
+from collections import defaultdict
+
+import numpy as np
+import torch
+import torchvision
+from termcolor import colored
+from torch.utils.tensorboard import SummaryWriter
+
+COMMON_TRAIN_FORMAT = [('frame', 'F', 'int'), ('step', 'S', 'int'),
+                       ('episode', 'E', 'int'), ('episode_length', 'L', 'int'),
+                       ('episode_reward', 'R', 'float'),
+                       ('buffer_size', 'BS', 'int'), ('fps', 'FPS', 'float'),
+                       ('total_time', 'T', 'time')]
+
+COMMON_EVAL_FORMAT = [('frame', 'F', 'int'), ('step', 'S', 'int'),
+                      ('episode', 'E', 'int'), ('episode_length', 'L', 'int'),
+                      ('episode_reward', 'R', 'float'),
+                      ('total_time', 'T', 'time')]
+
+
+class AverageMeter(object):
+    def __init__(self):
+        self._sum = 0
+        self._count = 0
+
+    def update(self, value, n=1):
+        self._sum += value
+        self._count += n
+
+    def value(self):
+        return self._sum / max(1, self._count)
+
+
+class MetersGroup(object):
+    def __init__(self, csv_file_name, formating):
+        self._csv_file_name = csv_file_name
+        self._formating = formating
+        self._meters = defaultdict(AverageMeter)
+        self._csv_file = None
+        self._csv_writer = None
+
+    def log(self, key, value, n=1):
+        self._meters[key].update(value, n)
+
+    def _prime_meters(self):
+        data = dict()
+        for key, meter in self._meters.items():
+            if key.startswith('train'):
+                key = key[len('train') + 1:]
+            else:
+                key = key[len('eval') + 1:]
+            key = key.replace('/', '_')
+            data[key] = meter.value()
+        return data
+
+    def _remove_old_entries(self, data):
+        rows = []
+        with self._csv_file_name.open('r') as f:
+            reader = csv.DictReader(f)
+            for row in reader:
+                if float(row['episode']) >= data['episode']:
+                    break
+                rows.append(row)
+        with self._csv_file_name.open('w') as f:
+            writer = csv.DictWriter(f,
+                                    fieldnames=sorted(data.keys()),
+                                    restval=0.0)
+            writer.writeheader()
+            for row in rows:
+                writer.writerow(row)
+
+    def _dump_to_csv(self, data):
+        if self._csv_writer is None:
+            should_write_header = True
+            if self._csv_file_name.exists():
+                self._remove_old_entries(data)
+                should_write_header = False
+
+            self._csv_file = self._csv_file_name.open('a')
+            self._csv_writer = csv.DictWriter(self._csv_file,
+                                              fieldnames=sorted(data.keys()),
+                                              restval=0.0)
+            if should_write_header:
+                self._csv_writer.writeheader()
+
+        self._csv_writer.writerow(data)
+        self._csv_file.flush()
+
+    def _format(self, key, value, ty):
+        if ty == 'int':
+            value = int(value)
+            return f'{key}: {value}'
+        elif ty == 'float':
+            return f'{key}: {value:.04f}'
+        elif ty == 'time':
+            value = str(datetime.timedelta(seconds=int(value)))
+            return f'{key}: {value}'
+        else:
+            raise f'invalid format type: {ty}'
+
+    def _dump_to_console(self, data, prefix):
+        prefix = colored(prefix, 'yellow' if prefix == 'train' else 'green')
+        pieces = [f'| {prefix: <14}']
+        for key, disp_key, ty in self._formating:
+            value = data.get(key, 0)
+            pieces.append(self._format(disp_key, value, ty))
+        print(' | '.join(pieces))
+
+    def dump(self, step, prefix):
+        if len(self._meters) == 0:
+            return
+        data = self._prime_meters()
+        data['frame'] = step
+        self._dump_to_csv(data)
+        self._dump_to_console(data, prefix)
+        self._meters.clear()
+
+
+class Logger(object):
+    def __init__(self, log_dir, use_tb, stage2_logger=False):
+        self._log_dir = log_dir
+        if not stage2_logger:
+            self._train_mg = MetersGroup(log_dir / 'train.csv',
+                                         formating=COMMON_TRAIN_FORMAT)
+            self._eval_mg = MetersGroup(log_dir / 'eval.csv',
+                                        formating=COMMON_EVAL_FORMAT)
+        else:
+            self._train_mg = MetersGroup(log_dir / 'train_stage2.csv',
+                                         formating=COMMON_TRAIN_FORMAT)
+            self._eval_mg = MetersGroup(log_dir / 'eval_stage2.csv',
+                                        formating=COMMON_EVAL_FORMAT)
+        if use_tb:
+            self._sw = SummaryWriter(str(log_dir / 'tb'))
+        else:
+            self._sw = None
+
+    def _try_sw_log(self, key, value, step):
+        if self._sw is not None:
+            self._sw.add_scalar(key, value, step)
+
+    def log(self, key, value, step):
+        assert key.startswith('train') or key.startswith('eval')
+        if type(value) == torch.Tensor:
+            value = value.item()
+        self._try_sw_log(key, value, step)
+        mg = self._train_mg if key.startswith('train') else self._eval_mg
+        mg.log(key, value)
+
+    def log_metrics(self, metrics, step, ty):
+        for key, value in metrics.items():
+            self.log(f'{ty}/{key}', value, step)
+
+    def dump(self, step, ty=None):
+        if ty is None or ty == 'eval':
+            self._eval_mg.dump(step, 'eval')
+        if ty is None or ty == 'train':
+            self._train_mg.dump(step, 'train')
+
+    def log_and_dump_ctx(self, step, ty):
+        return LogAndDumpCtx(self, step, ty)
+
+
+class LogAndDumpCtx:
+    def __init__(self, logger, step, ty):
+        self._logger = logger
+        self._step = step
+        self._ty = ty
+
+    def __enter__(self):
+        return self
+
+    def __call__(self, key, value):
+        self._logger.log(f'{self._ty}/{key}', value, self._step)
+
+    def __exit__(self, *args):
+        self._logger.dump(self._step, self._ty)

VRL3/src/replay_buffer.py

@@ -0,0 +1,222 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import datetime
+import io
+import random
+import traceback
+from collections import defaultdict
+
+import numpy as np
+import torch
+import torch.nn as nn
+from torch.utils.data import IterableDataset
+
+
+def episode_len(episode):
+    # subtract -1 because the dummy first transition
+    return next(iter(episode.values())).shape[0] - 1
+
+
+def save_episode(episode, fn):
+    with io.BytesIO() as bs:
+        np.savez_compressed(bs, **episode)
+        bs.seek(0)
+        with fn.open('wb') as f:
+            f.write(bs.read())
+
+
+def load_episode(fn):
+    with fn.open('rb') as f:
+        episode = np.load(f)
+        episode = {k: episode[k] for k in episode.keys()}
+        return episode
+
+
+class ReplayBufferStorage:
+    def __init__(self, data_specs, replay_dir):
+        self._data_specs = data_specs
+        self._replay_dir = replay_dir
+        replay_dir.mkdir(exist_ok=True)
+        self._current_episode = defaultdict(list)
+        self._preload()
+
+    def __len__(self):
+        return self._num_transitions
+
+    def add(self, time_step):
+        for spec in self._data_specs:
+            value = time_step[spec.name]
+            if np.isscalar(value):
+                value = np.full(spec.shape, value, spec.dtype)
+            # print(spec.name, spec.shape, spec.dtype, value.shape, value.dtype)
+            assert spec.shape == value.shape and spec.dtype == value.dtype
+            self._current_episode[spec.name].append(value)
+        if time_step.last():
+            episode = dict()
+            for spec in self._data_specs:
+                value = self._current_episode[spec.name]
+                episode[spec.name] = np.array(value, spec.dtype)
+            self._current_episode = defaultdict(list)
+            self._store_episode(episode)
+
+    def _preload(self):
+        self._num_episodes = 0
+        self._num_transitions = 0
+        for fn in self._replay_dir.glob('*.npz'):
+            _, _, eps_len = fn.stem.split('_')
+            self._num_episodes += 1
+            self._num_transitions += int(eps_len)
+
+    def _store_episode(self, episode):
+        eps_idx = self._num_episodes
+        eps_len = episode_len(episode)
+        self._num_episodes += 1
+        self._num_transitions += eps_len
+        ts = datetime.datetime.now().strftime('%Y%m%dT%H%M%S')
+        eps_fn = f'{ts}_{eps_idx}_{eps_len}.npz'
+        save_episode(episode, self._replay_dir / eps_fn)
+
+
+class ReplayBuffer(IterableDataset):
+    def __init__(self, replay_dir, max_size, num_workers, nstep, discount,
+                 fetch_every, save_snapshot, is_adroit=False, return_next_action=False):
+        self._replay_dir = replay_dir
+        self._size = 0
+        self._max_size = max_size
+        self._num_workers = max(1, num_workers)
+        self._episode_fns = []
+        self._episodes = dict()
+        self._nstep = nstep
+        self._discount = discount
+        self._fetch_every = fetch_every
+        self._samples_since_last_fetch = fetch_every
+        self._save_snapshot = save_snapshot
+        self._is_adroit = is_adroit
+        self._return_next_action = return_next_action
+
+    def set_nstep(self, nstep):
+        self._nstep = nstep
+
+    def _sample_episode(self):
+        eps_fn = random.choice(self._episode_fns)
+        return self._episodes[eps_fn]
+
+    def _store_episode(self, eps_fn):
+        try:
+            episode = load_episode(eps_fn)
+        except:
+            return False
+        eps_len = episode_len(episode)
+        while eps_len + self._size > self._max_size:
+            early_eps_fn = self._episode_fns.pop(0)
+            early_eps = self._episodes.pop(early_eps_fn)
+            self._size -= episode_len(early_eps)
+            early_eps_fn.unlink(missing_ok=True)
+        self._episode_fns.append(eps_fn)
+        self._episode_fns.sort()
+        self._episodes[eps_fn] = episode
+        self._size += eps_len
+
+        if not self._save_snapshot:
+            eps_fn.unlink(missing_ok=True)
+        return True
+
+    def _try_fetch(self):
+        if self._samples_since_last_fetch < self._fetch_every:
+            return
+        self._samples_since_last_fetch = 0
+        try:
+            worker_id = torch.utils.data.get_worker_info().id
+        except:
+            worker_id = 0
+        eps_fns = sorted(self._replay_dir.glob('*.npz'), reverse=True)
+        fetched_size = 0
+        for eps_fn in eps_fns:
+            eps_idx, eps_len = [int(x) for x in eps_fn.stem.split('_')[1:]]
+            if eps_idx % self._num_workers != worker_id:
+                continue
+            if eps_fn in self._episodes.keys():
+                break
+            if fetched_size + eps_len > self._max_size:
+                break
+            fetched_size += eps_len
+            if not self._store_episode(eps_fn):
+                break
+
+    def _sample(self):
+        try:
+            self._try_fetch()
+        except:
+            traceback.print_exc()
+        self._samples_since_last_fetch += 1
+        episode = self._sample_episode()
+        # add +1 for the first dummy transition
+        idx = np.random.randint(0, episode_len(episode) - self._nstep + 1) + 1
+        obs = episode['observation'][idx - 1]
+        action = episode['action'][idx]
+        next_obs = episode['observation'][idx + self._nstep - 1]
+        reward = np.zeros_like(episode['reward'][idx])
+        discount = np.ones_like(episode['discount'][idx])
+        for i in range(self._nstep):
+            step_reward = episode['reward'][idx + i]
+            reward += discount * step_reward
+            discount *= episode['discount'][idx + i] * self._discount
+
+        if self._return_next_action:
+            next_action = episode['action'][idx + self._nstep - 1]
+
+        if not self._is_adroit:
+            if self._return_next_action:
+                return (obs, action, reward, discount, next_obs, next_action)
+            else:
+                return (obs, action, reward, discount, next_obs)
+        else:
+            obs_sensor = episode['observation_sensor'][idx - 1]
+            obs_sensor_next = episode['observation_sensor'][idx + self._nstep - 1]
+            if self._return_next_action:
+                return (obs, action, reward, discount, next_obs, obs_sensor, obs_sensor_next, next_action)
+            else:
+                return (obs, action, reward, discount, next_obs, obs_sensor, obs_sensor_next)
+
+    def __iter__(self):
+        while True:
+            yield self._sample()
+
+
+def _worker_init_fn(worker_id):
+    seed = np.random.get_state()[1][0] + worker_id
+    np.random.seed(seed)
+    random.seed(seed)
+
+
+def make_replay_loader(replay_dir, max_size, batch_size, num_workers,
+                       save_snapshot, nstep, discount, fetch_every=1000, is_adroit=False, return_next_action=False):
+    max_size_per_worker = max_size // max(1, num_workers)
+
+    iterable = ReplayBuffer(replay_dir,
+                            max_size_per_worker,
+                            num_workers,
+                            nstep,
+                            discount,
+                            fetch_every=fetch_every,
+                            save_snapshot=save_snapshot, 
+                            is_adroit=is_adroit,
+                            return_next_action=return_next_action)
+
+    loader = torch.utils.data.DataLoader(iterable,
+                                         batch_size=batch_size,
+                                         num_workers=num_workers,
+                                         pin_memory=True,
+                                         worker_init_fn=_worker_init_fn)
+    return loader
+
+def reinit_data_loader(data_loader, batch_size, num_workers):
+    # reinit a data loader with a new batch size
+    loader = torch.utils.data.DataLoader(data_loader.dataset,
+                                         batch_size=batch_size,
+                                         num_workers=num_workers,
+                                         pin_memory=True,
+                                         worker_init_fn=_worker_init_fn)
+    return loader

VRL3/src/stage1_models.py

@@ -0,0 +1,318 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+from numpy import identity
+import torch.nn as nn
+import numpy as np
+
+"""
+most code here are modified from the TORCHVISION.MODELS.RESNET
+"""
+
+import torch
+import torch.nn as nn
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+class BasicBlock(nn.Module):
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(BasicBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.conv2 = conv3x3(planes, planes)
+        self.bn2 = norm_layer(planes)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+class OneLayerBlock(nn.Module):
+    # similar to BasicBlock, but shallower
+    expansion = 1
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(OneLayerBlock, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        if groups != 1 or base_width != 64:
+            raise ValueError('BasicBlock only supports groups=1 and base_width=64')
+        if dilation > 1:
+            raise NotImplementedError("Dilation > 1 not supported in BasicBlock")
+        # Both self.conv1 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv3x3(inplanes, planes, stride)
+        self.bn1 = norm_layer(planes)
+        self.relu = nn.ReLU(inplace=True)
+        self.stride = stride
+
+    def forward(self, x):
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+        return out
+
+class Bottleneck(nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+
+    expansion = 4
+
+    def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+                 base_width=64, dilation=1, norm_layer=None):
+        super(Bottleneck, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        width = int(planes * (base_width / 64.)) * groups
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(inplanes, width)
+        self.bn1 = norm_layer(width)
+        self.conv2 = conv3x3(width, width, stride, groups, dilation)
+        self.bn2 = norm_layer(width)
+        self.conv3 = conv1x1(width, planes * self.expansion)
+        self.bn3 = norm_layer(planes * self.expansion)
+        self.relu = nn.ReLU(inplace=True)
+        self.downsample = downsample
+        self.stride = stride
+
+    def forward(self, x):
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+def drq_weight_init(m):
+    # weight init scheme used in DrQv2
+    if isinstance(m, nn.Linear):
+        nn.init.orthogonal_(m.weight.data)
+        if hasattr(m.bias, 'data'):
+            m.bias.data.fill_(0.0)
+    elif isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
+        gain = nn.init.calculate_gain('relu')
+        nn.init.orthogonal_(m.weight.data, gain)
+        if hasattr(m.bias, 'data'):
+            m.bias.data.fill_(0.0)
+
+class Stage3ShallowEncoder(nn.Module):
+    """
+    this is the encoder architecture used in DrQv2
+    """
+    def __init__(self, obs_shape, n_channel):
+        super().__init__()
+        assert len(obs_shape) == 3
+        self.repr_dim = n_channel * 35 * 35
+        self.conv1 = nn.Conv2d(obs_shape[0], n_channel, 3, stride=2)
+        self.conv2 = nn.Conv2d(n_channel, n_channel, 3, stride=1)
+        self.conv3 = nn.Conv2d(n_channel, n_channel, 3, stride=1)
+        self.conv4 = nn.Conv2d(n_channel, n_channel, 3, stride=1)
+        self.relu = nn.ReLU(inplace=True)
+        self.apply(drq_weight_init)
+
+    def _forward_impl(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.conv2(x))
+        x = self.relu(self.conv3(x))
+        x = self.relu(self.conv4(x))
+        return x
+
+    def forward(self, obs):
+        o = obs
+        h = self._forward_impl(o)
+        h = h.view(h.shape[0], -1)
+        return h
+
+class ResNet84(nn.Module):
+    """
+    default stage 1 encoder used by VRL3, this is modified from the PyTorch standard ResNet class
+    but is more lightweight and this is much faster with 84x84 input size
+    use "layers" to specify how deep the network is
+    use "start_num_channel" to control how wide it is
+    """
+    def __init__(self, block, layers, num_classes=1000, zero_init_residual=False,
+                 groups=1, width_per_group=64, replace_stride_with_dilation=None,
+                 norm_layer=None, start_num_channel=32):
+        super(ResNet84, self).__init__()
+        if norm_layer is None:
+            norm_layer = nn.BatchNorm2d
+        self._norm_layer = norm_layer
+
+        self.start_num_channel = start_num_channel
+        self.inplanes = start_num_channel
+        self.dilation = 1
+        if replace_stride_with_dilation is None:
+            # each element in the tuple indicates if we should replace
+            # the 2x2 stride with a dilated convolution instead
+            replace_stride_with_dilation = [False, False, False]
+        if len(replace_stride_with_dilation) != 3:
+            raise ValueError("replace_stride_with_dilation should be None "
+                             "or a 3-element tuple, got {}".format(replace_stride_with_dilation))
+        self.groups = groups
+        self.base_width = width_per_group
+        self.conv1 = nn.Conv2d(3, self.inplanes, 3, stride=2)
+        self.bn1 = norm_layer(self.inplanes)
+        self.relu = nn.ReLU(inplace=True)
+        # self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+
+        self.layer1 = self._make_layer(block, start_num_channel, layers[0])
+        self.layer2 = self._make_layer(block, start_num_channel * 2, layers[1], stride=2,
+                                       dilate=replace_stride_with_dilation[0])
+        self.layer3 = self._make_layer(block, start_num_channel * 4, layers[2], stride=2,
+                                       dilate=replace_stride_with_dilation[1])
+        self.layer4 = self._make_layer(block, start_num_channel * 8, layers[3], stride=2,
+                                       dilate=replace_stride_with_dilation[2])
+        self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+        self.fc = nn.Linear(start_num_channel * 8 * block.expansion, num_classes)
+
+        for m in self.modules():
+            if isinstance(m, nn.Conv2d):
+                nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+            elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+                nn.init.constant_(m.weight, 1)
+                nn.init.constant_(m.bias, 0)
+
+        # Zero-initialize the last BN in each residual branch,
+        # so that the residual branch starts with zeros, and each residual block behaves like an identity.
+        # This improves the model by 0.2~0.3% according to https://arxiv.org/abs/1706.02677
+        if zero_init_residual:
+            for m in self.modules():
+                if isinstance(m, Bottleneck):
+                    nn.init.constant_(m.bn3.weight, 0)
+                elif isinstance(m, BasicBlock):
+                    nn.init.constant_(m.bn2.weight, 0)
+
+    def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+        # vrl3: if block is 0, allows a smaller network size
+        if blocks == 0:
+            block = OneLayerBlock
+
+        norm_layer = self._norm_layer
+        downsample = None
+        previous_dilation = self.dilation
+        if dilate:
+            self.dilation *= stride
+            stride = 1
+        if stride != 1 or self.inplanes != planes * block.expansion:
+            downsample = nn.Sequential(
+                conv1x1(self.inplanes, planes * block.expansion, stride),
+                norm_layer(planes * block.expansion),
+            )
+
+        layers = []
+        layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
+                            self.base_width, previous_dilation, norm_layer))
+        self.inplanes = planes * block.expansion
+        for _ in range(1, blocks):
+            layers.append(block(self.inplanes, planes, groups=self.groups,
+                                base_width=self.base_width, dilation=self.dilation,
+                                norm_layer=norm_layer))
+
+        return nn.Sequential(*layers)
+
+    def _forward_impl(self, x):
+        x = self.conv1(x)
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        x = self.layer1(x)
+        x = self.layer2(x)
+        x = self.layer3(x)
+        x = self.layer4(x)
+        x = self.avgpool(x)
+        x = torch.flatten(x, 1)
+        x = self.fc(x)
+
+        return x
+
+    def get_feature_size(self):
+        assert self.start_num_channel % 32 == 0
+        multiplier = self.start_num_channel // 32
+        size = 256 * multiplier
+        return size
+
+    def forward(self, x):
+        return self._forward_impl(x)
+
+    def get_features(self, x):
+        x = self.conv1(x)
+        # print("0", x.shape) # 32 x 41 x 41 = 53792
+        x = self.bn1(x)
+        x = self.relu(x)
+
+        x = self.layer1(x)
+        # print("1", x.shape) # 32 x 41 x 41= 53792
+
+        x = self.layer2(x)
+        # print("2", x.shape) # 64 x 21 x 21 = 28224
+
+        x = self.layer3(x)
+        # print("3", x.shape) # 128 x 11 x 11 = 15488
+
+        x = self.layer4(x)
+        # print("4", x.shape) # 256 x 6 x 6 = 9216
+
+        x = self.avgpool(x)
+        # print("pool", x.shape) # 256 x 1 x 1
+
+        final_out = torch.flatten(x, 1)
+        # print("flatten", x.shape) # 256
+        return final_out
+
+class Identity(nn.Module):
+    def __init__(self, input_placeholder=None):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return x
+

VRL3/src/train_stage1.py

@@ -0,0 +1,493 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+# this file is modified from the pytorch official tutorial
+# NOTE: stage 1 training code is currently being cleaned up
+
+import argparse
+import os
+import random
+import shutil
+import time
+import warnings
+
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.backends.cudnn as cudnn
+import torch.distributed as dist
+import torch.optim
+import torch.multiprocessing as mp
+import torch.utils.data
+import torch.utils.data.distributed
+import torchvision.transforms as transforms
+import torchvision.datasets as datasets
+import torchvision.models as models
+
+# TODO use another config file to indicate the location of the training data and also where to save models...
+#  should also add an option to just test the accuracy of models...
+#  we probably can test this locally ....
+
+from stage1_models import BasicBlock, ResNet84
+
+rl_model_names = ['resnet6_32channel', 'resnet10_32channel', 'resnet18_32channel',
+                  'resnet6_64channel', 'resnet10_64channel', 'resnet18_64channel',]
+model_names = sorted(name for name in models.__dict__
+    if name.islower() and not name.startswith("__")
+    and callable(models.__dict__[name])) + rl_model_names
+
+parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
+parser.add_argument('data', metavar='DIR',
+                    help='path to dataset')
+parser.add_argument('-a', '--arch', metavar='ARCH', default='resnet10_32channel',
+                    choices=model_names,
+                    help='model architecture: ' +
+                        ' | '.join(model_names) +
+                        ' (default: resnet18)')
+parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
+                    help='number of data loading workers (default: 4)')
+parser.add_argument('--epochs', default=90, type=int, metavar='N',
+                    help='number of total epochs to run')
+parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
+                    help='manual epoch number (useful on restarts)')
+parser.add_argument('-b', '--batch-size', default=256, type=int,
+                    metavar='N',
+                    help='mini-batch size (default: 256), this is the total '
+                         'batch size of all GPUs on the current node when '
+                         'using Data Parallel or Distributed Data Parallel')
+parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
+                    metavar='LR', help='initial learning rate', dest='lr')
+parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
+                    help='momentum')
+parser.add_argument('--wd', '--weight-decay', default=1e-4, type=float,
+                    metavar='W', help='weight decay (default: 1e-4)',
+                    dest='weight_decay')
+parser.add_argument('-p', '--print-freq', default=10, type=int,
+                    metavar='N', help='print frequency (default: 10)')
+parser.add_argument('--resume', default='', type=str, metavar='PATH',
+                    help='path to latest checkpoint (default: none)')
+parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
+                    help='evaluate model on validation set')
+parser.add_argument('--pretrained', dest='pretrained', action='store_true',
+                    help='use pre-trained model')
+parser.add_argument('--world-size', default=-1, type=int,
+                    help='number of nodes for distributed training')
+parser.add_argument('--rank', default=-1, type=int,
+                    help='node rank for distributed training')
+parser.add_argument('--dist-url', default='tcp://224.66.41.62:23456', type=str,
+                    help='url used to set up distributed training')
+parser.add_argument('--dist-backend', default='nccl', type=str,
+                    help='distributed backend')
+parser.add_argument('--seed', default=None, type=int,
+                    help='seed for initializing training. ')
+parser.add_argument('--gpu', default=None, type=int,
+                    help='GPU id to use.')
+parser.add_argument('--multiprocessing-distributed', action='store_true',
+                    help='Use multi-processing distributed training to launch '
+                         'N processes per node, which has N GPUs. This is the '
+                         'fastest way to use PyTorch for either single node or '
+                         'multi node data parallel training')
+parser.add_argument('--debug', default=0, type=int,
+                    help='1 for debug mode, 2 for super fast debug mode')
+
+best_acc1 = 0
+INPUT_SIZE = 84
+VAL_RESIZE = 100
+
+def main():
+    print(model_names)
+
+    args = parser.parse_args()
+
+    if args.seed is not None:
+        random.seed(args.seed)
+        torch.manual_seed(args.seed)
+        cudnn.deterministic = True
+        warnings.warn('You have chosen to seed training. '
+                      'This will turn on the CUDNN deterministic setting, '
+                      'which can slow down your training considerably! '
+                      'You may see unexpected behavior when restarting '
+                      'from checkpoints.')
+
+    if args.gpu is not None:
+        warnings.warn('You have chosen a specific GPU. This will completely '
+                      'disable data parallelism.')
+
+    if args.dist_url == "env://" and args.world_size == -1:
+        args.world_size = int(os.environ["WORLD_SIZE"])
+
+    args.distributed = args.world_size > 1 or args.multiprocessing_distributed
+
+    ngpus_per_node = torch.cuda.device_count()
+    if args.multiprocessing_distributed:
+        # Since we have ngpus_per_node processes per node, the total world_size
+        # needs to be adjusted accordingly
+        args.world_size = ngpus_per_node * args.world_size
+        # Use torch.multiprocessing.spawn to launch distributed processes: the
+        # main_worker process function
+        mp.spawn(main_worker, nprocs=ngpus_per_node, args=(ngpus_per_node, args))
+    else:
+        # Simply call main_worker function
+        main_worker(args.gpu, ngpus_per_node, args)
+
+
+def main_worker(gpu, ngpus_per_node, args):
+    global best_acc1
+    args.gpu = gpu
+
+    if args.gpu is not None:
+        print("Use GPU: {} for training".format(args.gpu))
+
+    if args.distributed:
+        if args.dist_url == "env://" and args.rank == -1:
+            args.rank = int(os.environ["RANK"])
+        if args.multiprocessing_distributed:
+            # For multiprocessing distributed training, rank needs to be the
+            # global rank among all the processes
+            args.rank = args.rank * ngpus_per_node + gpu
+        dist.init_process_group(backend=args.dist_backend, init_method=args.dist_url,
+                                world_size=args.world_size, rank=args.rank)
+    # create model
+    if args.debug > 0:
+        print("=> creating model for debug 2")
+        # model = ResNet84(BasicBlock, [1, 1, 1, 1], num_classes=5) # 1, 1, 1, 1 will make a resnet10
+        model = ResNet84(BasicBlock, [0, 0, 0, 0], num_classes=5) # 0, 0, 0, 0 make a convnet6 (5 conv layers in total lol)
+        x = torch.rand((1, 3, 84, 84)).float()
+        out = model(x)
+        print(model)
+        quit()
+
+        # model = ResNetTest2(BasicBlock, [2, 2, 2, 2])
+        #model = Drq4Encoder((3, 84, 84), n_channel, 200)
+    else:
+        if args.pretrained:
+            print("=> using pre-trained model '{}'".format(args.arch))
+            model = models.__dict__[args.arch](pretrained=True)
+        else:
+            print("=> creating model '{}'".format(args.arch))
+            if args.arch in rl_model_names:
+                if args.arch == 'resnet18_32channel':
+                    model = ResNet84(BasicBlock, [2, 2, 2, 2], start_num_channel=32)  # 1, 1, 1, 1 will make a resnet10
+                elif args.arch == 'resnet10_32channel':
+                    model = ResNet84(BasicBlock, [1, 1, 1, 1], start_num_channel=32)  # 1, 1, 1, 1 will make a resnet10
+                elif args.arch == 'resnet6_32channel':
+                    model = ResNet84(BasicBlock, [0, 0, 0, 0], start_num_channel=32)  # resnet 6 (actually not even resnet because no skip connection)
+                elif args.arch == 'resnet18_64channel':
+                    model = ResNet84(BasicBlock, [2, 2, 2, 2], start_num_channel=64)
+                elif args.arch == 'resnet10_64channel':
+                    model = ResNet84(BasicBlock, [1, 1, 1, 1], start_num_channel=64)
+                elif args.arch == 'resnet6_64channel':
+                    model = ResNet84(BasicBlock, [0, 0, 0, 0], start_num_channel=64)
+                else:
+                    print("specialized model not yet implemented")
+                    quit()
+            else:
+                model = models.__dict__[args.arch]()
+
+    if not torch.cuda.is_available():
+        print('using CPU, this will be slow')
+    elif args.distributed:
+        print("distributed")
+        # For multiprocessing distributed, DistributedDataParallel constructor
+        # should always set the single device scope, otherwise,
+        # DistributedDataParallel will use all available devices.
+        if args.gpu is not None:
+            torch.cuda.set_device(args.gpu)
+            model.cuda(args.gpu)
+            # When using a single GPU per process and per
+            # DistributedDataParallel, we need to divide the batch size
+            # ourselves based on the total number of GPUs we have
+            args.batch_size = int(args.batch_size / ngpus_per_node)
+            args.workers = int((args.workers + ngpus_per_node - 1) / ngpus_per_node)
+            model = torch.nn.parallel.DistributedDataParallel(model, device_ids=[args.gpu])
+        else:
+            model.cuda()
+            # DistributedDataParallel will divide and allocate batch_size to all
+            # available GPUs if device_ids are not set
+            model = torch.nn.parallel.DistributedDataParallel(model)
+    elif args.gpu is not None:
+        print("use gpu:", args.gpu)
+        torch.cuda.set_device(args.gpu)
+        model = model.cuda(args.gpu)
+    else:
+        print("data parallel")
+        # DataParallel will divide and allocate batch_size to all available GPUs
+        if args.arch.startswith('alexnet') or args.arch.startswith('vgg'):
+            model.features = torch.nn.DataParallel(model.features)
+            model.cuda()
+        else:
+            model = torch.nn.DataParallel(model).cuda()
+
+    # define loss function (criterion) and optimizer
+    criterion = nn.CrossEntropyLoss().cuda(args.gpu)
+
+    optimizer = torch.optim.SGD(model.parameters(), args.lr,
+                                momentum=args.momentum,
+                                weight_decay=args.weight_decay)
+
+    # optionally resume from a checkpoint
+    if args.resume:
+        if os.path.isfile(args.resume):
+            print("=> loading checkpoint '{}'".format(args.resume))
+            if args.gpu is None:
+                checkpoint = torch.load(args.resume)
+            else:
+                # Map model to be loaded to specified single gpu.
+                loc = 'cuda:{}'.format(args.gpu)
+                checkpoint = torch.load(args.resume, map_location=loc)
+            args.start_epoch = checkpoint['epoch']
+            best_acc1 = checkpoint['best_acc1']
+            if args.gpu is not None:
+                # best_acc1 may be from a checkpoint from a different GPU
+                best_acc1 = best_acc1.to(args.gpu)
+            model.load_state_dict(checkpoint['state_dict'])
+            optimizer.load_state_dict(checkpoint['optimizer'])
+            print("=> loaded checkpoint '{}' (epoch {})"
+                  .format(args.resume, checkpoint['epoch']))
+        else:
+            print("=> no checkpoint found at '{}'".format(args.resume))
+
+    cudnn.benchmark = True
+
+    # Data loading code
+    traindir = os.path.join(args.data, 'train')
+    valdir = os.path.join(args.data, 'val')
+    normalize = transforms.Normalize(mean=[0.485, 0.456, 0.406],
+                                     std=[0.229, 0.224, 0.225])
+
+    print("train directory is:", traindir)
+    print("val directory is:", valdir)
+
+    train_dataset = datasets.ImageFolder(
+        traindir,
+        transforms.Compose([
+            transforms.RandomResizedCrop(INPUT_SIZE),
+            transforms.RandomHorizontalFlip(),
+            transforms.ToTensor(),
+            normalize,
+        ]))
+
+    print("data set ready")
+
+    if args.distributed:
+        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
+    else:
+        train_sampler = None
+
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
+        num_workers=args.workers, pin_memory=True, sampler=train_sampler)
+
+    val_loader = torch.utils.data.DataLoader(
+        datasets.ImageFolder(valdir, transforms.Compose([
+            # transforms.Resize(256),
+            transforms.Resize(VAL_RESIZE),
+            transforms.CenterCrop(INPUT_SIZE),
+            transforms.ToTensor(),
+            normalize,
+        ])),
+        batch_size=args.batch_size, shuffle=False,
+        num_workers=args.workers, pin_memory=True)
+
+    if args.evaluate:
+        validate(val_loader, model, criterion, args)
+        return
+
+    for epoch in range(args.start_epoch, args.epochs):
+        print(epoch)
+        epoch_start_time = time.time()
+
+        if args.distributed:
+            train_sampler.set_epoch(epoch)
+        adjust_learning_rate(optimizer, epoch, args)
+
+        # train for one epoch
+        train(train_loader, model, criterion, optimizer, epoch, args)
+
+        # evaluate on validation set
+        acc1 = validate(val_loader, model, criterion, args)
+
+        # remember best acc@1 and save checkpoint
+        is_best = acc1 > best_acc1
+        best_acc1 = max(acc1, best_acc1)
+
+        if not args.multiprocessing_distributed or (args.multiprocessing_distributed
+                and args.rank % ngpus_per_node == 0):
+            save_checkpoint({
+                'epoch': epoch + 1,
+                'arch': args.arch,
+                'state_dict': model.state_dict(),
+                'best_acc1': best_acc1,
+                'optimizer' : optimizer.state_dict(),
+            }, is_best,
+            save_name_prefix=args.arch)
+
+        epoch_end_time = time.time() - epoch_start_time
+        print("epoch finished in %.3f hour" % (epoch_end_time/3600))
+
+def train(train_loader, model, criterion, optimizer, epoch, args):
+    batch_time = AverageMeter('Time', ':6.3f')
+    data_time = AverageMeter('Data', ':6.3f')
+    losses = AverageMeter('Loss', ':.4e')
+    top1 = AverageMeter('Acc@1', ':6.2f')
+    top5 = AverageMeter('Acc@5', ':6.2f')
+    progress = ProgressMeter(
+        len(train_loader),
+        [batch_time, data_time, losses, top1, top5],
+        prefix="Epoch: [{}]".format(epoch))
+
+    # switch to train mode
+    model.train()
+
+    end = time.time()
+    for i, (images, target) in enumerate(train_loader):
+        # measure data loading time
+        data_time.update(time.time() - end)
+
+        if args.gpu is not None:
+            images = images.cuda(args.gpu, non_blocking=True)
+        if torch.cuda.is_available():
+            target = target.cuda(args.gpu, non_blocking=True)
+
+        # compute output
+        output = model(images)
+        loss = criterion(output, target)
+
+        # measure accuracy and record loss
+        acc1, acc5 = accuracy(output, target, topk=(1, 5))
+        losses.update(loss.item(), images.size(0))
+        top1.update(acc1[0], images.size(0))
+        top5.update(acc5[0], images.size(0))
+
+        # compute gradient and do SGD step
+        optimizer.zero_grad()
+        loss.backward()
+        optimizer.step()
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        if i % args.print_freq == 0:
+            progress.display(i)
+
+
+def validate(val_loader, model, criterion, args):
+    batch_time = AverageMeter('Time', ':6.3f')
+    losses = AverageMeter('Loss', ':.4e')
+    top1 = AverageMeter('Acc@1', ':6.2f')
+    top5 = AverageMeter('Acc@5', ':6.2f')
+    progress = ProgressMeter(
+        len(val_loader),
+        [batch_time, losses, top1, top5],
+        prefix='Test: ')
+
+    # switch to evaluate mode
+    model.eval()
+
+    with torch.no_grad():
+        end = time.time()
+        for i, (images, target) in enumerate(val_loader):
+            if args.gpu is not None:
+                images = images.cuda(args.gpu, non_blocking=True)
+            if torch.cuda.is_available():
+                target = target.cuda(args.gpu, non_blocking=True)
+
+            # compute output
+            output = model(images)
+            loss = criterion(output, target)
+
+            # measure accuracy and record loss
+            acc1, acc5 = accuracy(output, target, topk=(1, 5))
+            losses.update(loss.item(), images.size(0))
+            top1.update(acc1[0], images.size(0))
+            top5.update(acc5[0], images.size(0))
+
+            # measure elapsed time
+            batch_time.update(time.time() - end)
+            end = time.time()
+
+            if i % args.print_freq == 0:
+                progress.display(i)
+
+        # TODO: this should also be done with the ProgressMeter
+        print(' * Acc@1 {top1.avg:.3f} Acc@5 {top5.avg:.3f}'
+              .format(top1=top1, top5=top5))
+
+    return top1.avg
+
+
+def save_checkpoint(state, is_best, filename='checkpoint.pth.tar', save_name_prefix = ''):
+    save_name = save_name_prefix + '_' + filename
+    torch.save(state, save_name)
+    if is_best:
+        best_model_save_name = save_name_prefix + '_' + 'model_best.pth.tar'
+        shutil.copyfile(save_name, best_model_save_name)
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+    def __init__(self, name, fmt=':f'):
+        self.name = name
+        self.fmt = fmt
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+    def __str__(self):
+        fmtstr = '{name} {val' + self.fmt + '} ({avg' + self.fmt + '})'
+        return fmtstr.format(**self.__dict__)
+
+
+class ProgressMeter(object):
+    def __init__(self, num_batches, meters, prefix=""):
+        self.batch_fmtstr = self._get_batch_fmtstr(num_batches)
+        self.meters = meters
+        self.prefix = prefix
+
+    def display(self, batch):
+        entries = [self.prefix + self.batch_fmtstr.format(batch)]
+        entries += [str(meter) for meter in self.meters]
+        print('\t'.join(entries))
+
+    def _get_batch_fmtstr(self, num_batches):
+        num_digits = len(str(num_batches // 1))
+        fmt = '{:' + str(num_digits) + 'd}'
+        return '[' + fmt + '/' + fmt.format(num_batches) + ']'
+
+
+def adjust_learning_rate(optimizer, epoch, args):
+    """Sets the learning rate to the initial LR decayed by 10 every 30 epochs"""
+    lr = args.lr * (0.1 ** (epoch // 30))
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+
+def accuracy(output, target, topk=(1,)):
+    """Computes the accuracy over the k top predictions for the specified values of k"""
+    with torch.no_grad():
+        maxk = max(topk)
+        batch_size = target.size(0)
+
+        _, pred = output.topk(maxk, 1, True, True)
+        pred = pred.t()
+        correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+        res = []
+        for k in topk:
+            correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
+            res.append(correct_k.mul_(100.0 / batch_size))
+        return res
+
+
+if __name__ == '__main__':
+    main()

VRL3/src/utils.py

@@ -0,0 +1,149 @@
+# Copyright (c) Facebook, Inc. and its affiliates. All Rights Reserved.
+#
+# This source code is licensed under the MIT license found in the
+# LICENSE file in the root directory of this source tree.
+import random
+import re
+import time
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from omegaconf import OmegaConf
+from torch import distributions as pyd
+from torch.distributions.utils import _standard_normal
+
+
+class eval_mode:
+    def __init__(self, *models):
+        self.models = models
+
+    def __enter__(self):
+        self.prev_states = []
+        for model in self.models:
+            self.prev_states.append(model.training)
+            model.train(False)
+
+    def __exit__(self, *args):
+        for model, state in zip(self.models, self.prev_states):
+            model.train(state)
+        return False
+
+
+def set_seed_everywhere(seed):
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    np.random.seed(seed)
+    random.seed(seed)
+
+
+def soft_update_params(net, target_net, tau):
+    for param, target_param in zip(net.parameters(), target_net.parameters()):
+        target_param.data.copy_(tau * param.data +
+                                (1 - tau) * target_param.data)
+
+
+def to_torch(xs, device):
+    return tuple(torch.as_tensor(x, device=device) for x in xs)
+
+
+def weight_init(m):
+    if isinstance(m, nn.Linear):
+        nn.init.orthogonal_(m.weight.data)
+        if hasattr(m.bias, 'data'):
+            m.bias.data.fill_(0.0)
+    elif isinstance(m, nn.Conv2d) or isinstance(m, nn.ConvTranspose2d):
+        gain = nn.init.calculate_gain('relu')
+        nn.init.orthogonal_(m.weight.data, gain)
+        if hasattr(m.bias, 'data'):
+            m.bias.data.fill_(0.0)
+
+
+class Until:
+    def __init__(self, until, action_repeat=1):
+        self._until = until
+        self._action_repeat = action_repeat
+
+    def __call__(self, step):
+        if self._until is None:
+            return True
+        until = self._until // self._action_repeat
+        return step < until
+
+
+class Every:
+    def __init__(self, every, action_repeat=1):
+        self._every = every
+        self._action_repeat = action_repeat
+
+    def __call__(self, step):
+        if self._every is None:
+            return False
+        every = self._every // self._action_repeat
+        if step % every == 0:
+            return True
+        return False
+
+
+class Timer:
+    def __init__(self):
+        self._start_time = time.time()
+        self._last_time = time.time()
+
+    def reset(self):
+        elapsed_time = time.time() - self._last_time
+        self._last_time = time.time()
+        total_time = time.time() - self._start_time
+        return elapsed_time, total_time
+
+    def total_time(self):
+        return time.time() - self._start_time
+
+
+class TruncatedNormal(pyd.Normal):
+    def __init__(self, loc, scale, low=-1.0, high=1.0, eps=1e-6):
+        super().__init__(loc, scale, validate_args=False)
+        self.low = low
+        self.high = high
+        self.eps = eps
+
+    def _clamp(self, x):
+        clamped_x = torch.clamp(x, self.low + self.eps, self.high - self.eps)
+        x = x - x.detach() + clamped_x.detach()
+        return x
+
+    def sample(self, clip=None, sample_shape=torch.Size()):
+        shape = self._extended_shape(sample_shape)
+        eps = _standard_normal(shape,
+                               dtype=self.loc.dtype,
+                               device=self.loc.device)
+        eps *= self.scale
+        if clip is not None:
+            eps = torch.clamp(eps, -clip, clip)
+        x = self.loc + eps
+        return self._clamp(x)
+
+
+def schedule(schdl, step):
+    try:
+        return float(schdl)
+    except ValueError:
+        match = re.match(r'linear\((.+),(.+),(.+)\)', schdl)
+        if match:
+            init, final, duration = [float(g) for g in match.groups()]
+            mix = np.clip(step / duration, 0.0, 1.0)
+            return (1.0 - mix) * init + mix * final
+        match = re.match(r'step_linear\((.+),(.+),(.+),(.+),(.+)\)', schdl)
+        if match:
+            init, final1, duration1, final2, duration2 = [
+                float(g) for g in match.groups()
+            ]
+            if step <= duration1:
+                mix = np.clip(step / duration1, 0.0, 1.0)
+                return (1.0 - mix) * init + mix * final1
+            else:
+                mix = np.clip((step - duration1) / duration2, 0.0, 1.0)
+                return (1.0 - mix) * final1 + mix * final2
+    raise NotImplementedError(schdl)

VRL3/src/vrl3_agent.py

@@ -0,0 +1,632 @@
+# Copyright (c) Microsoft Corporation.
+# Licensed under the MIT License.
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+from torchvision import datasets, models, transforms
+from transfer_util import initialize_model
+from stage1_models import BasicBlock, ResNet84
+import os
+import copy
+from PIL import Image
+import platform
+from numbers import Number
+import utils
+
+class RandomShiftsAug(nn.Module):
+    def __init__(self, pad):
+        super().__init__()
+        self.pad = pad
+
+    def forward(self, x):
+        n, c, h, w = x.size()
+        assert h == w
+        padding = tuple([self.pad] * 4)
+        x = F.pad(x, padding, 'replicate')
+        eps = 1.0 / (h + 2 * self.pad)
+        arange = torch.linspace(-1.0 + eps,
+                                1.0 - eps,
+                                h + 2 * self.pad,
+                                device=x.device,
+                                dtype=x.dtype)[:h]
+        arange = arange.unsqueeze(0).repeat(h, 1).unsqueeze(2)
+        base_grid = torch.cat([arange, arange.transpose(1, 0)], dim=2)
+        base_grid = base_grid.unsqueeze(0).repeat(n, 1, 1, 1)
+
+        shift = torch.randint(0,
+                              2 * self.pad + 1,
+                              size=(n, 1, 1, 2),
+                              device=x.device,
+                              dtype=x.dtype)
+        shift *= 2.0 / (h + 2 * self.pad)
+
+        grid = base_grid + shift
+        return F.grid_sample(x,
+                             grid,
+                             padding_mode='zeros',
+                             align_corners=False)
+
+class Identity(nn.Module):
+    def __init__(self, input_placeholder=None):
+        super(Identity, self).__init__()
+
+    def forward(self, x):
+        return x
+
+class RLEncoder(nn.Module):
+    def __init__(self, obs_shape, model_name, device):
+        super().__init__()
+        # a wrapper over a non-RL encoder model
+        self.device = device
+        assert len(obs_shape) == 3
+        self.n_input_channel = obs_shape[0]
+        assert self.n_input_channel % 3 == 0
+        self.n_images = self.n_input_channel // 3
+        self.model = self.init_model(model_name)
+        self.model.fc = Identity()
+        self.repr_dim = self.model.get_feature_size()
+
+        self.normalize_op = transforms.Normalize((0.485, 0.456, 0.406),
+                                                 (0.229, 0.224, 0.225))
+        self.channel_mismatch = True
+
+    def init_model(self, model_name):
+        # model name is e.g. resnet6_32channel
+        n_layer_string, n_channel_string = model_name.split('_')
+        layer_string_to_layer_list = {
+            'resnet6': [0, 0, 0, 0],
+            'resnet10': [1, 1, 1, 1],
+            'resnet18': [2, 2, 2, 2],
+        }
+        channel_string_to_n_channel = {
+            '32channel': 32,
+            '64channel': 64,
+        }
+        layer_list = layer_string_to_layer_list[n_layer_string]
+        start_num_channel = channel_string_to_n_channel[n_channel_string]
+        return ResNet84(BasicBlock, layer_list, start_num_channel=start_num_channel).to(self.device)
+
+    def expand_first_layer(self):
+        # convolutional channel expansion to deal with input mismatch
+        multiplier = self.n_images
+        self.model.conv1.weight.data = self.model.conv1.weight.data.repeat(1,multiplier,1,1) / multiplier
+        means = (0.485, 0.456, 0.406) * multiplier
+        stds = (0.229, 0.224, 0.225) * multiplier
+        self.normalize_op = transforms.Normalize(means, stds)
+        self.channel_mismatch = False
+
+    def freeze_bn(self):
+        # freeze batch norm layers (VRL3 ablation shows modifying how
+        # batch norm is trained does not affect performance)
+        for module in self.model.modules():
+            if isinstance(module, nn.BatchNorm2d):
+                if hasattr(module, 'weight'):
+                    module.weight.requires_grad_(False)
+                if hasattr(module, 'bias'):
+                    module.bias.requires_grad_(False)
+                module.eval()
+
+    def get_parameters_that_require_grad(self):
+        params = []
+        for name, param in self.named_parameters():
+            if param.requires_grad == True:
+                params.append(param)
+        return params
+
+    def transform_obs_tensor_batch(self, obs):
+        # transform obs batch before put into the pretrained resnet
+        new_obs = self.normalize_op(obs.float()/255)
+        return new_obs
+
+    def _forward_impl(self, x):
+        x = self.model.get_features(x)
+        return x
+
+    def forward(self, obs):
+        o = self.transform_obs_tensor_batch(obs)
+        h = self._forward_impl(o)
+        return h
+
+class Stage3ShallowEncoder(nn.Module):
+    def __init__(self, obs_shape, n_channel):
+        super().__init__()
+
+        assert len(obs_shape) == 3
+        self.repr_dim = n_channel * 35 * 35
+
+        self.n_input_channel = obs_shape[0]
+        self.conv1 = nn.Conv2d(obs_shape[0], n_channel, 3, stride=2)
+        self.conv2 = nn.Conv2d(n_channel, n_channel, 3, stride=1)
+        self.conv3 = nn.Conv2d(n_channel, n_channel, 3, stride=1)
+        self.conv4 = nn.Conv2d(n_channel, n_channel, 3, stride=1)
+        self.relu = nn.ReLU(inplace=True)
+
+        # TODO here add prediction head so we can do contrastive learning...
+
+        self.apply(utils.weight_init)
+        self.normalize_op = transforms.Normalize((0.485, 0.456, 0.406, 0.485, 0.456, 0.406, 0.485, 0.456, 0.406),
+                                                 (0.229, 0.224, 0.225, 0.229, 0.224, 0.225, 0.229, 0.224, 0.225))
+
+        self.compress = nn.Sequential(nn.Linear(self.repr_dim, 50), nn.LayerNorm(50), nn.Tanh())
+        self.pred_layer = nn.Linear(50, 50, bias=False)
+
+    def transform_obs_tensor_batch(self, obs):
+        # transform obs batch before put into the pretrained resnet
+        # correct order might be first augment, then resize, then normalize
+        # obs = F.interpolate(obs, size=self.pretrained_model_input_size)
+        new_obs = obs / 255.0 - 0.5
+        # new_obs = self.normalize_op(new_obs)
+        return new_obs
+
+    def _forward_impl(self, x):
+        x = self.relu(self.conv1(x))
+        x = self.relu(self.conv2(x))
+        x = self.relu(self.conv3(x))
+        x = self.relu(self.conv4(x))
+        return x
+
+    def forward(self, obs):
+        o = self.transform_obs_tensor_batch(obs)
+        h = self._forward_impl(o)
+        h = h.view(h.shape[0], -1)
+        return h
+
+    def get_anchor_output(self, obs, actions=None):
+        # typically go through conv and then compression layer and then a mlp
+        # used for UL update
+        conv_out = self.forward(obs)
+        compressed = self.compress(conv_out)
+        pred = self.pred_layer(compressed)
+        return pred, conv_out
+
+    def get_positive_output(self, obs):
+        # typically go through conv, compression
+        # used for UL update
+        conv_out = self.forward(obs)
+        compressed = self.compress(conv_out)
+        return compressed
+
+class Encoder(nn.Module):
+    def __init__(self, obs_shape, n_channel):
+        super().__init__()
+
+        assert len(obs_shape) == 3
+        self.repr_dim = n_channel * 35 * 35
+
+        self.convnet = nn.Sequential(nn.Conv2d(obs_shape[0], n_channel, 3, stride=2),
+                                     nn.ReLU(), nn.Conv2d(n_channel, n_channel, 3, stride=1),
+                                     nn.ReLU(), nn.Conv2d(n_channel, n_channel, 3, stride=1),
+                                     nn.ReLU(), nn.Conv2d(n_channel, n_channel, 3, stride=1),
+                                     nn.ReLU())
+
+        self.apply(utils.weight_init)
+
+    def forward(self, obs):
+        obs = obs / 255.0 - 0.5
+        h = self.convnet(obs)
+        h = h.view(h.shape[0], -1)
+        return h
+
+class IdentityEncoder(nn.Module):
+    def __init__(self, obs_shape):
+        super().__init__()
+
+        assert len(obs_shape) == 1
+        self.repr_dim = obs_shape[0]
+
+    def forward(self, obs):
+        return obs
+
+class Actor(nn.Module):
+    def __init__(self, repr_dim, action_shape, feature_dim, hidden_dim):
+        super().__init__()
+
+        self.trunk = nn.Sequential(nn.Linear(repr_dim, feature_dim),
+                                   nn.LayerNorm(feature_dim), nn.Tanh())
+
+        self.policy = nn.Sequential(nn.Linear(feature_dim, hidden_dim),
+                                    nn.ReLU(inplace=True),
+                                    nn.Linear(hidden_dim, hidden_dim),
+                                    nn.ReLU(inplace=True),
+                                    nn.Linear(hidden_dim, action_shape[0]))
+
+        self.action_shift=0
+        self.action_scale=1
+        self.apply(utils.weight_init)
+
+    def forward(self, obs, std):
+        h = self.trunk(obs)
+
+        mu = self.policy(h)
+        mu = torch.tanh(mu)
+        mu = mu * self.action_scale + self.action_shift
+        std = torch.ones_like(mu) * std
+
+        dist = utils.TruncatedNormal(mu, std)
+        return dist
+
+    def forward_with_pretanh(self, obs, std):
+        h = self.trunk(obs)
+
+        mu = self.policy(h)
+        pretanh = mu
+        mu = torch.tanh(mu)
+        mu = mu * self.action_scale + self.action_shift
+        std = torch.ones_like(mu) * std
+
+        dist = utils.TruncatedNormal(mu, std)
+        return dist, pretanh
+
+class Critic(nn.Module):
+    def __init__(self, repr_dim, action_shape, feature_dim, hidden_dim):
+        super().__init__()
+
+        self.trunk = nn.Sequential(nn.Linear(repr_dim, feature_dim),
+                                   nn.LayerNorm(feature_dim), nn.Tanh())
+
+        self.Q1 = nn.Sequential(
+            nn.Linear(feature_dim + action_shape[0], hidden_dim),
+            nn.ReLU(inplace=True), nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(inplace=True), nn.Linear(hidden_dim, 1))
+
+        self.Q2 = nn.Sequential(
+            nn.Linear(feature_dim + action_shape[0], hidden_dim),
+            nn.ReLU(inplace=True), nn.Linear(hidden_dim, hidden_dim),
+            nn.ReLU(inplace=True), nn.Linear(hidden_dim, 1))
+
+        self.apply(utils.weight_init)
+
+    def forward(self, obs, action):
+        h = self.trunk(obs)
+        h_action = torch.cat([h, action], dim=-1)
+        q1 = self.Q1(h_action)
+        q2 = self.Q2(h_action)
+
+        return q1, q2
+
+class VRL3Agent:
+    def __init__(self, obs_shape, action_shape, device, use_sensor, lr, feature_dim,
+                 hidden_dim, critic_target_tau, num_expl_steps,
+                 update_every_steps, stddev_clip, use_tb, use_data_aug, encoder_lr_scale,
+                 stage1_model_name, safe_q_target_factor, safe_q_threshold, pretanh_penalty, pretanh_threshold,
+                 stage2_update_encoder, cql_weight, cql_temp, cql_n_random, stage2_std, stage2_bc_weight,
+                 stage3_update_encoder, std0, std1, std_n_decay,
+                 stage3_bc_lam0, stage3_bc_lam1):
+        self.device = device
+        self.critic_target_tau = critic_target_tau
+        self.update_every_steps = update_every_steps
+        self.use_tb = use_tb
+        self.num_expl_steps = num_expl_steps
+
+        self.stage2_std = stage2_std
+        self.stage2_update_encoder = stage2_update_encoder
+
+        if std1 > std0:
+            std1 = std0
+        self.stddev_schedule = "linear(%s,%s,%s)" % (str(std0), str(std1), str(std_n_decay))
+
+        self.stddev_clip = stddev_clip
+        self.use_data_aug = use_data_aug
+        self.safe_q_target_factor = safe_q_target_factor
+        self.q_threshold = safe_q_threshold
+        self.pretanh_penalty = pretanh_penalty
+
+        self.cql_temp = cql_temp
+        self.cql_weight = cql_weight
+        self.cql_n_random = cql_n_random
+
+        self.pretanh_threshold = pretanh_threshold
+
+        self.stage2_bc_weight = stage2_bc_weight
+        self.stage3_bc_lam0 = stage3_bc_lam0
+        self.stage3_bc_lam1 = stage3_bc_lam1
+
+        if stage3_update_encoder and encoder_lr_scale > 0 and len(obs_shape) > 1:
+            self.stage3_update_encoder = True
+        else:
+            self.stage3_update_encoder = False
+
+        self.encoder = RLEncoder(obs_shape, stage1_model_name, device).to(device)
+
+        self.act_dim = action_shape[0]
+
+        if use_sensor:
+            downstream_input_dim = self.encoder.repr_dim + 24
+        else:
+            downstream_input_dim = self.encoder.repr_dim
+
+        self.actor = Actor(downstream_input_dim, action_shape, feature_dim,
+                           hidden_dim).to(device)
+        self.critic = Critic(downstream_input_dim, action_shape, feature_dim,
+                             hidden_dim).to(device)
+        self.critic_target = Critic(downstream_input_dim, action_shape,
+                                    feature_dim, hidden_dim).to(device)
+        self.critic_target.load_state_dict(self.critic.state_dict())
+
+        # optimizers
+        self.actor_opt = torch.optim.Adam(self.actor.parameters(), lr=lr)
+        self.critic_opt = torch.optim.Adam(self.critic.parameters(), lr=lr)
+
+        encoder_lr = lr * encoder_lr_scale
+        """ set up encoder optimizer """
+        self.encoder_opt = torch.optim.Adam(self.encoder.parameters(), lr=encoder_lr)
+        # data augmentation
+        self.aug = RandomShiftsAug(pad=4)
+        self.train()
+        self.critic_target.train()
+
+    def load_pretrained_encoder(self, model_path, verbose=True):
+        if verbose:
+            print("Trying to load pretrained model from:", model_path)
+        checkpoint = torch.load(model_path, map_location=torch.device(self.device))
+        state_dict = checkpoint['state_dict']
+
+        pretrained_dict = {}
+        # remove `module.` if model was pretrained with distributed mode
+        for k, v in state_dict.items():
+            if 'module.' in k:
+                name = k[7:]
+            else:
+                name = k
+            pretrained_dict[name] = v
+        self.encoder.model.load_state_dict(pretrained_dict, strict=False)
+        if verbose:
+            print("Pretrained model loaded!")
+
+    def switch_to_RL_stages(self, verbose=True):
+        # run convolutional channel expansion to match input shape
+        self.encoder.expand_first_layer()
+        if verbose:
+            print("Convolutional channel expansion finished: now can take in %d images as input." % self.encoder.n_images)
+
+    def train(self, training=True):
+        self.training = training
+        self.encoder.train(training)
+        self.actor.train(training)
+        self.critic.train(training)
+
+    def act(self, obs, step, eval_mode, obs_sensor=None, is_tensor_input=False, force_action_std=None):
+        """
+        obs: 3x84x84, uint8, [0,255]
+        """
+        # eval_mode should be False when taking an exploration action in stage 3
+        # eval_mode should be True when evaluate agent performance
+
+        if force_action_std == None:
+            stddev = utils.schedule(self.stddev_schedule, step)
+            if step < self.num_expl_steps and not eval_mode:
+                action = np.random.uniform(0, 1, (self.act_dim,)).astype(np.float32)
+                return action
+        else:
+            stddev = force_action_std
+
+        if is_tensor_input:
+            obs = self.encoder(obs)
+        else:
+            obs = torch.as_tensor(obs, device=self.device)
+            obs = self.encoder(obs.unsqueeze(0))
+
+        if obs_sensor is not None:
+            obs_sensor = torch.as_tensor(obs_sensor, device=self.device)
+            obs_sensor = obs_sensor.unsqueeze(0)
+            obs_combined = torch.cat([obs, obs_sensor], dim=1)
+        else:
+            obs_combined = obs
+
+        dist = self.actor(obs_combined, stddev)
+        if eval_mode:
+            action = dist.mean
+        else:
+            action = dist.sample(clip=None)
+            if step < self.num_expl_steps:
+                action.uniform_(-1.0, 1.0)
+        return action.cpu().numpy()[0]
+
+    def update(self, replay_iter, step, stage, use_sensor):
+        # for stage 2 and 3, we use the same functions but with different hyperparameters
+        assert stage in (2, 3)
+        metrics = dict()
+
+        if stage == 2:
+            update_encoder = self.stage2_update_encoder
+            stddev = self.stage2_std
+            conservative_loss_weight = self.cql_weight
+            bc_weight = self.stage2_bc_weight
+
+        if stage == 3:
+            if step % self.update_every_steps != 0:
+                return metrics
+            update_encoder = self.stage3_update_encoder
+
+            stddev = utils.schedule(self.stddev_schedule, step)
+            conservative_loss_weight = 0
+
+            # compute stage 3 BC weight
+            bc_data_per_iter = 40000
+            i_iter = step // bc_data_per_iter
+            bc_weight = self.stage3_bc_lam0 * self.stage3_bc_lam1 ** i_iter
+
+        # batch data
+        batch = next(replay_iter)
+        if use_sensor: # TODO might want to...?
+            obs, action, reward, discount, next_obs, obs_sensor, obs_sensor_next = utils.to_torch(batch, self.device)
+        else:
+            obs, action, reward, discount, next_obs = utils.to_torch(batch, self.device)
+            obs_sensor, obs_sensor_next = None, None
+
+        # augment
+        if self.use_data_aug:
+            obs = self.aug(obs.float())
+            next_obs = self.aug(next_obs.float())
+        else:
+            obs = obs.float()
+            next_obs = next_obs.float()
+
+        # encode
+        if update_encoder:
+            obs = self.encoder(obs)
+        else:
+            with torch.no_grad():
+                obs = self.encoder(obs)
+
+        with torch.no_grad():
+            next_obs = self.encoder(next_obs)
+
+        # concatenate obs with additional sensor observation if needed
+        obs_combined = torch.cat([obs, obs_sensor], dim=1) if obs_sensor is not None else obs
+        obs_next_combined = torch.cat([next_obs, obs_sensor_next], dim=1) if obs_sensor_next is not None else next_obs
+
+        # update critic
+        metrics.update(self.update_critic_vrl3(obs_combined, action, reward, discount, obs_next_combined,
+                                               stddev, update_encoder, conservative_loss_weight))
+
+        # update actor, following previous works, we do not use actor gradient for encoder update
+        metrics.update(self.update_actor_vrl3(obs_combined.detach(), action, stddev, bc_weight,
+                                              self.pretanh_penalty, self.pretanh_threshold))
+
+        metrics['batch_reward'] = reward.mean().item()
+
+        # update critic target networks
+        utils.soft_update_params(self.critic, self.critic_target, self.critic_target_tau)
+        return metrics
+
+    def update_critic_vrl3(self, obs, action, reward, discount, next_obs, stddev, update_encoder, conservative_loss_weight):
+        metrics = dict()
+        batch_size = obs.shape[0]
+
+        """
+        STANDARD Q LOSS COMPUTATION:
+        - get standard Q loss first, this is the same as in any other online RL methods
+        - except for the safe Q technique, which controls how large the Q value can be
+        """
+        with torch.no_grad():
+            dist = self.actor(next_obs, stddev)
+            next_action = dist.sample(clip=self.stddev_clip)
+            target_Q1, target_Q2 = self.critic_target(next_obs, next_action)
+            target_V = torch.min(target_Q1, target_Q2)
+            target_Q = reward + (discount * target_V)
+
+            if self.safe_q_target_factor < 1:
+                target_Q[target_Q > (self.q_threshold + 1)] = self.q_threshold + (target_Q[target_Q > (self.q_threshold+1)] - self.q_threshold) ** self.safe_q_target_factor
+
+        Q1, Q2 = self.critic(obs, action)
+        critic_loss = F.mse_loss(Q1, target_Q) + F.mse_loss(Q2, target_Q)
+
+        """
+        CONSERVATIVE Q LOSS COMPUTATION:
+        - sample random actions, actions from policy and next actions from policy, as done in CQL authors' code
+          (though this detail is not really discussed in the CQL paper)
+        - only compute this loss when conservative loss weight > 0
+        """
+        if conservative_loss_weight > 0:
+            random_actions = (torch.rand((batch_size * self.cql_n_random, self.act_dim), device=self.device) - 0.5) * 2
+
+            dist = self.actor(obs, stddev)
+            current_actions = dist.sample(clip=self.stddev_clip)
+
+            dist = self.actor(next_obs, stddev)
+            next_current_actions = dist.sample(clip=self.stddev_clip)
+
+            # now get Q values for all these actions (for both Q networks)
+            obs_repeat = obs.unsqueeze(1).repeat(1, self.cql_n_random, 1).view(obs.shape[0] * self.cql_n_random,
+                                                                               obs.shape[1])
+
+            Q1_rand, Q2_rand = self.critic(obs_repeat,
+                                           random_actions)  # TODO might want to double check the logic here see if the repeat is correct
+            Q1_rand = Q1_rand.view(obs.shape[0], self.cql_n_random)
+            Q2_rand = Q2_rand.view(obs.shape[0], self.cql_n_random)
+
+            Q1_curr, Q2_curr = self.critic(obs, current_actions)
+            Q1_curr_next, Q2_curr_next = self.critic(obs, next_current_actions)
+
+            # now concat all these Q values together
+            Q1_cat = torch.cat([Q1_rand, Q1, Q1_curr, Q1_curr_next], 1)
+            Q2_cat = torch.cat([Q2_rand, Q2, Q2_curr, Q2_curr_next], 1)
+
+            cql_min_q1_loss = torch.logsumexp(Q1_cat / self.cql_temp,
+                                              dim=1, ).mean() * conservative_loss_weight * self.cql_temp
+            cql_min_q2_loss = torch.logsumexp(Q2_cat / self.cql_temp,
+                                              dim=1, ).mean() * conservative_loss_weight * self.cql_temp
+
+            """Subtract the log likelihood of data"""
+            conservative_q_loss = cql_min_q1_loss + cql_min_q2_loss - (Q1.mean() + Q2.mean()) * conservative_loss_weight
+            critic_loss_combined = critic_loss + conservative_q_loss
+        else:
+            critic_loss_combined = critic_loss
+
+        # logging
+        metrics['critic_target_q'] = target_Q.mean().item()
+        metrics['critic_q1'] = Q1.mean().item()
+        metrics['critic_q2'] = Q2.mean().item()
+        metrics['critic_loss'] = critic_loss.item()
+
+        # if needed, also update encoder with critic loss
+        if update_encoder:
+            self.encoder_opt.zero_grad(set_to_none=True)
+        self.critic_opt.zero_grad(set_to_none=True)
+        critic_loss_combined.backward()
+        self.critic_opt.step()
+        if update_encoder:
+            self.encoder_opt.step()
+
+        return metrics
+
+    def update_actor_vrl3(self, obs, action, stddev, bc_weight, pretanh_penalty, pretanh_threshold):
+        metrics = dict()
+
+        """
+        get standard actor loss
+        """
+        dist, pretanh = self.actor.forward_with_pretanh(obs, stddev)
+        current_action = dist.sample(clip=self.stddev_clip)
+        log_prob = dist.log_prob(current_action).sum(-1, keepdim=True)
+        Q1, Q2 = self.critic(obs, current_action)
+        Q = torch.min(Q1, Q2)
+        actor_loss = -Q.mean()
+
+        """
+        add BC loss
+        """
+        if bc_weight > 0:
+            # get mean action with no action noise (though this might not be necessary)
+            stddev_bc = 0
+            dist_bc = self.actor(obs, stddev_bc)
+            current_mean_action = dist_bc.sample(clip=self.stddev_clip)
+            actor_loss_bc = F.mse_loss(current_mean_action, action) * bc_weight
+        else:
+            actor_loss_bc = torch.FloatTensor([0]).to(self.device)
+
+        """
+        add pretanh penalty (might not be necessary for Adroit)
+        """
+        pretanh_loss = 0
+        if pretanh_penalty > 0:
+            pretanh_loss = pretanh.abs() - pretanh_threshold
+            pretanh_loss[pretanh_loss < 0] = 0
+            pretanh_loss = (pretanh_loss ** 2).mean() * pretanh_penalty
+
+        """
+        combine actor losses and optimize
+        """
+        actor_loss_combined = actor_loss + actor_loss_bc + pretanh_loss
+
+        self.actor_opt.zero_grad(set_to_none=True)
+        actor_loss_combined.backward()
+        self.actor_opt.step()
+
+        metrics['actor_loss'] = actor_loss.item()
+        metrics['actor_loss_bc'] = actor_loss_bc.item()
+        metrics['actor_logprob'] = log_prob.mean().item()
+        metrics['actor_ent'] = dist.entropy().sum(dim=-1).mean().item()
+        metrics['abs_pretanh'] = pretanh.abs().mean().item()
+        metrics['max_abs_pretanh'] = pretanh.abs().max().item()
+
+        return metrics
+
+    def to(self, device):
+        self.actor.to(device)
+        self.critic.to(device)
+        self.encoder.to(device)
+        self.device = device

gym-0.21.0/.github/stale.yml

@@ -0,0 +1,62 @@
+# Configuration for probot-stale - https://github.com/probot/stale
+
+# Number of days of inactivity before an Issue or Pull Request becomes stale
+daysUntilStale: 60
+
+# Number of days of inactivity before an Issue or Pull Request with the stale label is closed.
+# Set to false to disable. If disabled, issues still need to be closed manually, but will remain marked as stale.
+daysUntilClose: 14
+
+# Only issues or pull requests with all of these labels are check if stale. Defaults to `[]` (disabled)
+onlyLabels:
+  - more-information-needed
+
+# Issues or Pull Requests with these labels will never be considered stale. Set to `[]` to disable
+exemptLabels:
+  - pinned
+  - security
+  - "[Status] Maybe Later"
+
+# Set to true to ignore issues in a project (defaults to false)
+exemptProjects: true
+
+# Set to true to ignore issues in a milestone (defaults to false)
+exemptMilestones: true
+
+# Set to true to ignore issues with an assignee (defaults to false)
+exemptAssignees: true
+
+# Label to use when marking as stale
+staleLabel: stale
+
+# Comment to post when marking as stale. Set to `false` to disable
+markComment: >
+  This issue has been automatically marked as stale because it has not had
+  recent activity. It will be closed if no further activity occurs. Thank you
+  for your contributions.
+
+# Comment to post when removing the stale label.
+# unmarkComment: >
+#   Your comment here.
+
+# Comment to post when closing a stale Issue or Pull Request.
+# closeComment: >
+#   Your comment here.
+
+# Limit the number of actions per hour, from 1-30. Default is 30
+limitPerRun: 30
+
+# Limit to only `issues` or `pulls`
+only: issues
+
+# Optionally, specify configuration settings that are specific to just 'issues' or 'pulls':
+# pulls:
+#   daysUntilStale: 30
+#   markComment: >
+#     This pull request has been automatically marked as stale because it has not had
+#     recent activity. It will be closed if no further activity occurs. Thank you
+#     for your contributions.
+
+# issues:
+#   exemptLabels:
+#     - confirmed

gym-0.21.0/CONTRIBUTING.md

@@ -0,0 +1,18 @@
+# Gym Contribution Guidelines
+
+At this time we are currently accepting the current forms of contributions:
+
+- Bug reports (keep in mind that changing environment behavior should be minimized as that requires releasing a new version of the environment and makes results hard to compare across versions)
+- Pull requests for bug fixes
+- Documentation improvements
+
+Notably, we are not accepting these forms of contributions:
+
+- New environments
+- New features
+
+This may change in the future.
+If you wish to make a Gym environment, follow the instructions in [Creating Environments](https://github.com/openai/gym/blob/master/docs/creating-environments.md).  When your environment works, you can make a PR to add it to the bottom of the [List of Environments](https://github.com/openai/gym/blob/master/docs/environments.md).
+
+
+Edit July 27, 2021: Please see https://github.com/openai/gym/issues/2259 for new contributing standards

gym-0.21.0/README.md

@@ -0,0 +1,57 @@
+## Gym
+
+Gym is an open source Python library for developing and comparing reinforcement learning algorithms by providing a standard API to communicate between learning algorithms and environments, as well as a standard set of environments compliant with that API. Since its release, Gym's API has become the field standard for doing this.
+
+Gym currently has two pieces of documentation: the [documentation website](http://gym.openai.com) and the [FAQ](https://github.com/openai/gym/wiki/FAQ). A new and more comprehensive documentation website is in the works.
+
+## Installation
+
+To install the base Gym library, use `pip install gym`.
+
+This does not include dependencies for all families of environments (there's a massive number, and some can be problematic to install on certain systems). You can install these dependencies for one family like `pip install gym[atari]` or use `pip install gym[all]` to install all dependencies.
+
+We support Python 3.6, 3.7, 3.8 and 3.9 on Linux and macOS. We will accept PRs related to Windows, but do not officially support it.
+
+## API
+
+The Gym API's API models environments as simple Python `env` classes. Creating environment instances and interacting with them is very simple- here's an example using the "CartPole-v1" environment:
+
+```python
+import gym 
+env = gym.make('CartPole-v1')
+
+# env is created, now we can use it: 
+for episode in range(10): 
+    obs = env.reset()
+    for step in range(50):
+        action = env.action_space.sample()  # or given a custom model, action = policy(observation)
+        nobs, reward, done, info = env.step(action)
+```
+
+## Notable Related Libraries
+
+* [Stable Baselines 3](https://github.com/DLR-RM/stable-baselines3) is a learning library based on the Gym API. It is our recommendation for beginners who want to start learning things quickly.
+* [RL Baselines3 Zoo](https://github.com/DLR-RM/rl-baselines3-zoo) builds upon SB3, containing optimal hyperparameters for Gym environments as well as code to easily find new ones. Such tuning is almost always required.
+* The [Autonomous Learning Library](https://github.com/cpnota/autonomous-learning-library) and [Tianshou](https://github.com/thu-ml/tianshou) are two reinforcement learning libraries I like that are generally geared towards more experienced users.
+* [PettingZoo](https://github.com/PettingZoo-Team/PettingZoo) is like Gym, but for environments with multiple agents.
+
+## Environment Versioning
+
+Gym keeps strict versioning for reproducibility reasons. All environments end in a suffix like "\_v0".  When changes are made to environments that might impact learning results, the number is increased by one to prevent potential confusion.
+
+## Citation
+
+A whitepaper from when OpenAI Gym just came out is available https://arxiv.org/pdf/1606.01540, and can be cited with the following bibtex entry:
+
+```
+@misc{1606.01540,
+  Author = {Greg Brockman and Vicki Cheung and Ludwig Pettersson and Jonas Schneider and John Schulman and Jie Tang and Wojciech Zaremba},
+  Title = {OpenAI Gym},
+  Year = {2016},
+  Eprint = {arXiv:1606.01540},
+}
+```
+
+## Release Notes
+
+There used to be release notes for all the new Gym versions here. New release notes are being moved to [releases page](https://github.com/openai/gym/releases) on GitHub, like most other libraries do. Old notes can be viewed [here](https://github.com/openai/gym/blob/31be35ecd460f670f0c4b653a14c9996b7facc6c/README.rst).

gym-0.21.0/docs/toy_text/blackjack.md

@@ -0,0 +1,60 @@
+Blackjack
+---
+|Title|Action Type|Action Shape|Action Values|Observation Shape|Observation Values|Average Total Reward|Import|
+| ----------- | -----------| ----------- | -----------| ----------- | -----------| ----------- | -----------|
+|Blackjack|Discrete|(1,)|(0,1)|(3,)|[(0,31),(0,10),(0,1)]| |from gym.envs.toy_text import blackjack|
+---
+
+Blackjack is a card game where the goal is to obtain cards that sum to as near as possible to 21 without going over.  They're playing against a fixed dealer.
+
+Card Values:
+
+- Face cards (Jack, Queen, King) have point value 10.
+- Aces can either count as 11 or 1, and it's called 'usable ace' at 11.
+- Numerical cards (2-9) have value of their number.
+
+This game is placed with an infinite deck (or with replacement).
+The game starts with dealer having one face up and one face down card, while player having two face up cards. 
+
+The player can request additional cards (hit, action=1) until they decide to stop
+(stick, action=0) or exceed 21 (bust).
+After the player sticks, the dealer reveals their facedown card, and draws
+until their sum is 17 or greater.  If the dealer goes bust the player wins.
+If neither player nor dealer busts, the outcome (win, lose, draw) is
+decided by whose sum is closer to 21.
+
+The agent take a 1-element vector for actions.
+The action space is `(action)`, where: 
+- `action` is used to decide stick/hit for values (0,1).
+
+The observation of a 3-tuple of: the players current sum,
+the dealer's one showing card (1-10 where 1 is ace), and whether or not the player holds a usable ace (0 or 1).
+
+This environment corresponds to the version of the blackjack problem
+described in Example 5.1 in Reinforcement Learning: An Introduction
+by Sutton and Barto.
+http://incompleteideas.net/book/the-book-2nd.html
+
+**Rewards:**
+
+Reward schedule:
+- win game: +1
+- lose game: -1
+- draw game: 0
+- win game with natural blackjack: 
+
+    +1.5 (if <a href="#nat">natural</a> is True.) 
+    
+    +1 (if <a href="#nat">natural</a> is False.)
+
+### Arguments
+
+```
+gym.make('Blackjack-v0', natural=False)
+```
+
+<a id="nat">`natural`</a>: Whether to give an additional reward for starting with a natural blackjack, i.e. starting with an ace and ten (sum is 21).
+
+### Version History
+
+* v0: Initial versions release (1.0.0)

gym-0.21.0/docs/toy_text/taxi.md

@@ -0,0 +1,92 @@
+Taxi
+---
+|Title|Action Type|Action Shape|Action Values|Observation Shape|Observation Values|Average Total Reward|Import|
+| ----------- | -----------| ----------- | -----------| ----------- | -----------| ----------- | -----------|
+|Taxi|Discrete|(1,)|(0,5)|(1,)|(0,499)| |from gym.envs.toy_text import taxi|
+---
+
+
+The Taxi Problem
+from "Hierarchical Reinforcement Learning with the MAXQ Value Function Decomposition"
+
+by Tom Dietterich
+
+
+
+Description:
+
+There are four designated locations in the grid world indicated by R(ed), G(reen), Y(ellow), and B(lue). When the episode starts, the taxi starts off at a random square and the passenger is at a random location. The taxi drives to the passenger's location, picks up the passenger, drives to the passenger's destination (another one of the four specified locations), and then drops off the passenger. Once the passenger is dropped off, the episode ends.
+
+MAP:
+
+    +---------+
+    |R: | : :G|
+    | : | : : |
+    | : : : : |
+    | | : | : |
+    |Y| : |B: |
+    +---------+
+
+Actions:
+
+There are 6 discrete deterministic actions:
+- 0: move south
+- 1: move north
+- 2: move east
+- 3: move west
+- 4: pickup passenger
+- 5: drop off passenger
+
+Observations:
+
+There are 500 discrete states since there are 25 taxi positions, 5 possible locations of the passenger (including the case when the passenger is in the taxi), and 4 destination locations.
+
+Note that there are 400 states that can actually be reached during an episode. The missing states correspond to situations in which the passenger is at the same location as their destination, as this typically signals the end of an episode.
+Four additional states can be observed right after a successful episodes, when both the passenger and the taxi are at the destination.
+This gives a total of 404 reachable discrete states.
+
+Passenger locations:
+- 0: R(ed)
+- 1: G(reen)
+- 2: Y(ellow)
+- 3: B(lue)
+- 4: in taxi
+
+Destinations:
+- 0: R(ed)
+- 1: G(reen)
+- 2: Y(ellow)
+- 3: B(lue)
+
+
+
+**Rewards:**
+
+- -1 per step reward unless other reward is triggered.
+- +20 delivering passenger.
+- -10  executing "pickup" and "drop-off" actions illegally.
+
+
+Rendering:
+- blue: passenger
+- magenta: destination
+- yellow: empty taxi
+- green: full taxi
+- other letters (R, G, Y and B): locations for passengers and destinations
+state space is represented by:
+(taxi_row, taxi_col, passenger_location, destination)
+
+### Arguments
+
+```
+gym.make('Taxi-v3')
+```
+
+
+
+### Version History
+
+* v3: Map Correction + Cleaner Domain Description
+* v2: Disallow Taxi start location = goal location, Update Taxi observations in the rollout, Update Taxi reward threshold.
+* v1: Remove (3,2) from locs, add passidx<4 check
+* v0: Initial versions release

gym-0.21.0/scripts/generate_json.py

@@ -0,0 +1,119 @@
+from gym import envs, logger
+import json
+import os
+import sys
+import argparse
+
+from tests.envs.spec_list import should_skip_env_spec_for_tests
+from tests import generate_rollout_hash
+
+DATA_DIR = os.path.join(os.path.dirname(__file__), os.pardir, "gym", "envs", "tests")
+ROLLOUT_STEPS = 100
+episodes = ROLLOUT_STEPS
+steps = ROLLOUT_STEPS
+
+ROLLOUT_FILE = os.path.join(DATA_DIR, "rollout.json")
+
+if not os.path.isfile(ROLLOUT_FILE):
+    logger.info(
+        "No rollout file found. Writing empty json file to {}".format(ROLLOUT_FILE)
+    )
+    with open(ROLLOUT_FILE, "w") as outfile:
+        json.dump({}, outfile, indent=2)
+
+
+def update_rollout_dict(spec, rollout_dict):
+    """
+    Takes as input the environment spec for which the rollout is to be generated,
+    and the existing dictionary of rollouts. Returns True iff the dictionary was
+    modified.
+    """
+    # Skip platform-dependent
+    if should_skip_env_spec_for_tests(spec):
+        logger.info("Skipping tests for {}".format(spec.id))
+        return False
+
+    # Skip environments that are nondeterministic
+    if spec.nondeterministic:
+        logger.info("Skipping tests for nondeterministic env {}".format(spec.id))
+        return False
+
+    logger.info("Generating rollout for {}".format(spec.id))
+
+    try:
+        (
+            observations_hash,
+            actions_hash,
+            rewards_hash,
+            dones_hash,
+        ) = generate_rollout_hash(spec)
+    except:
+        # If running the env generates an exception, don't write to the rollout file
+        logger.warn(
+            "Exception {} thrown while generating rollout for {}. Rollout not added.".format(
+                sys.exc_info()[0], spec.id
+            )
+        )
+        return False
+
+    rollout = {}
+    rollout["observations"] = observations_hash
+    rollout["actions"] = actions_hash
+    rollout["rewards"] = rewards_hash
+    rollout["dones"] = dones_hash
+
+    existing = rollout_dict.get(spec.id)
+    if existing:
+        differs = False
+        for key, new_hash in rollout.items():
+            differs = differs or existing[key] != new_hash
+        if not differs:
+            logger.debug("Hashes match with existing for {}".format(spec.id))
+            return False
+        else:
+            logger.warn("Got new hash for {}. Overwriting.".format(spec.id))
+
+    rollout_dict[spec.id] = rollout
+    return True
+
+
+def add_new_rollouts(spec_ids, overwrite):
+    environments = [
+        spec for spec in envs.registry.all() if spec.entry_point is not None
+    ]
+    if spec_ids:
+        environments = [spec for spec in environments if spec.id in spec_ids]
+        assert len(environments) == len(spec_ids), "Some specs not found"
+    with open(ROLLOUT_FILE) as data_file:
+        rollout_dict = json.load(data_file)
+    modified = False
+    for spec in environments:
+        if not overwrite and spec.id in rollout_dict:
+            logger.debug("Rollout already exists for {}. Skipping.".format(spec.id))
+        else:
+            modified = update_rollout_dict(spec, rollout_dict) or modified
+
+    if modified:
+        logger.info("Writing new rollout file to {}".format(ROLLOUT_FILE))
+        with open(ROLLOUT_FILE, "w") as outfile:
+            json.dump(rollout_dict, outfile, indent=2, sort_keys=True)
+    else:
+        logger.info("No modifications needed.")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "-f",
+        "--force",
+        action="store_true",
+        help="Overwrite " + "existing rollouts if hashes differ.",
+    )
+    parser.add_argument("-v", "--verbose", action="store_true")
+    parser.add_argument(
+        "specs", nargs="*", help="ids of env specs to check (default: all)"
+    )
+    args = parser.parse_args()
+    if args.verbose:
+        logger.set_level(logger.INFO)
+    add_new_rollouts(args.specs, args.force)

gym-0.21.0/setup.py

@@ -0,0 +1,76 @@
+import os.path
+import sys
+import itertools
+
+from setuptools import find_packages, setup
+
+# Don't import gym module here, since deps may not be installed
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "gym"))
+from version import VERSION
+
+# Environment-specific dependencies.
+extras = {
+    "atari": ["ale-py~=0.7.1"],
+    "accept-rom-license": ["autorom[accept-rom-license]~=0.4.2"],
+    "box2d": ["box2d-py==2.3.5", "pyglet>=1.4.0"],
+    "classic_control": ["pyglet>=1.4.0"],
+    "mujoco": ["mujoco_py>=1.50, <2.0"],
+    "robotics": ["mujoco_py>=1.50, <2.0"],
+    "toy_text": ["scipy>=1.4.1"],
+    "other": ["lz4>=3.1.0", "opencv-python>=3.0"],
+}
+
+# Meta dependency groups.
+nomujoco_blacklist = set(["mujoco", "robotics", "accept-rom-license"])
+nomujoco_groups = set(extras.keys()) - nomujoco_blacklist
+
+extras["nomujoco"] = list(
+    itertools.chain.from_iterable(map(lambda group: extras[group], nomujoco_groups))
+)
+
+
+all_blacklist = set(["accept-rom-license"])
+all_groups = set(extras.keys()) - all_blacklist
+
+extras["all"] = list(
+    itertools.chain.from_iterable(map(lambda group: extras[group], all_groups))
+)
+
+setup(
+    name="gym",
+    version=VERSION,
+    description="Gym: A universal API for reinforcement learning environments.",
+    url="https://github.com/openai/gym",
+    author="OpenAI",
+    author_email="jkterry@umd.edu",
+    license="",
+    packages=[package for package in find_packages() if package.startswith("gym")],
+    zip_safe=False,
+    install_requires=[
+        "numpy>=1.18.0",
+        "cloudpickle>=1.2.0",
+        "importlib_metadata>=4.8.1; python_version < '3.8'",
+    ],
+    extras_require=extras,
+    package_data={
+        "gym": [
+            "envs/mujoco/assets/*.xml",
+            "envs/classic_control/assets/*.png",
+            "envs/robotics/assets/LICENSE.md",
+            "envs/robotics/assets/fetch/*.xml",
+            "envs/robotics/assets/hand/*.xml",
+            "envs/robotics/assets/stls/fetch/*.stl",
+            "envs/robotics/assets/stls/hand/*.stl",
+            "envs/robotics/assets/textures/*.png",
+        ]
+    },
+    tests_require=["pytest", "mock"],
+    python_requires=">=3.6",
+    classifiers=[
+        "Programming Language :: Python :: 3",
+        "Programming Language :: Python :: 3.6",
+        "Programming Language :: Python :: 3.7",
+        "Programming Language :: Python :: 3.8",
+        "Programming Language :: Python :: 3.9",
+    ],
+)

mujoco-py-2.1.2.14/.gitignore

@@ -0,0 +1,55 @@
+mujoco-py-**
+mjkey.txt
+mujoco_py/generated/cymj*
+_pyxbld*
+dist
+cache
+.idea/*
+*~
+.*~
+*#*#
+*.o
+*.dat
+*.prof
+*.lprof
+*.local
+.realsync
+.DS_Store
+**/*.egg-info
+.cache
+*.ckpt
+*.log
+.ipynb_checkpoints
+venv/
+.vimrc
+*.settings
+*.svn
+.project
+.pydevproject
+tags
+*sublime-project
+*sublime-workspace
+# Intermediate outputs
+__pycache__
+**/__pycache__
+*.pb.*
+*.pyc
+*.swp
+*.swo
+# generated data
+*.rdb
+*.db
+*.avi
+# mujoco outputs
+MUJOCO_LOG.TXT
+model.txt
+.window_data
+.idea/*.xml
+outputfile
+tmp*
+cymj.c
+**/.git
+.eggs/
+*.so
+.python-version
+/build

mujoco-py-2.1.2.14/mujoco_py.egg-info/SOURCES.txt

@@ -0,0 +1,67 @@
+LICENSE.md
+MANIFEST.in
+README.md
+pyproject.toml
+requirements.dev.txt
+requirements.txt
+setup.py
+mujoco_py/__init__.py
+mujoco_py/builder.py
+mujoco_py/cymj.pyx
+mujoco_py/mjbatchrenderer.pyx
+mujoco_py/mjpid.pyx
+mujoco_py/mjrendercontext.pyx
+mujoco_py/mjrenderpool.py
+mujoco_py/mjsim.pyx
+mujoco_py/mjsimstate.pyx
+mujoco_py/mjviewer.py
+mujoco_py/modder.py
+mujoco_py/opengl_context.pyx
+mujoco_py/utils.py
+mujoco_py/version.py
+mujoco_py.egg-info/PKG-INFO
+mujoco_py.egg-info/SOURCES.txt
+mujoco_py.egg-info/dependency_links.txt
+mujoco_py.egg-info/requires.txt
+mujoco_py.egg-info/top_level.txt
+mujoco_py/generated/__init__.py
+mujoco_py/generated/const.py
+mujoco_py/generated/wrappers.pxi
+mujoco_py/gl/__init__.py
+mujoco_py/gl/dummyshim.c
+mujoco_py/gl/egl.h
+mujoco_py/gl/eglext.h
+mujoco_py/gl/eglplatform.h
+mujoco_py/gl/eglshim.c
+mujoco_py/gl/glshim.h
+mujoco_py/gl/khrplatform.h
+mujoco_py/gl/osmesashim.c
+mujoco_py/pxd/__init__.py
+mujoco_py/pxd/mjdata.pxd
+mujoco_py/pxd/mjmodel.pxd
+mujoco_py/pxd/mjrender.pxd
+mujoco_py/pxd/mjui.pxd
+mujoco_py/pxd/mjvisualize.pxd
+mujoco_py/pxd/mujoco.pxd
+mujoco_py/tests/__init__.py
+mujoco_py/tests/include.xml
+mujoco_py/tests/test.xml
+mujoco_py/tests/test_composite.py
+mujoco_py/tests/test_cymj.py
+mujoco_py/tests/test_examples.py
+mujoco_py/tests/test_gen_wrappers.py
+mujoco_py/tests/test_modder.py
+mujoco_py/tests/test_opengl_context.py
+mujoco_py/tests/test_pid.py
+mujoco_py/tests/test_render_pool.py
+mujoco_py/tests/test_substep.py
+mujoco_py/tests/test_vfs.py
+mujoco_py/tests/test_viewer.py
+mujoco_py/tests/utils.py
+xmls/claw.xml
+xmls/door.xml
+xmls/juggler.xml
+xmls/key.xml
+xmls/shelf.xml
+xmls/slider.xml
+xmls/tosser.xml

mujoco-py-2.1.2.14/mujoco_py/builder.py

@@ -0,0 +1,518 @@
+import distutils
+import glob
+import os
+import shutil
+import subprocess
+import sys
+from distutils.core import Extension
+from distutils.dist import Distribution
+from distutils.sysconfig import customize_compiler
+from importlib.machinery import ExtensionFileLoader
+from os.path import abspath, dirname, exists, join, getmtime
+from random import choice
+from shutil import move
+from string import ascii_lowercase
+
+import fasteners
+import numpy as np
+from Cython.Build import cythonize
+from Cython.Distutils.old_build_ext import old_build_ext as build_ext
+from cffi import FFI
+
+from mujoco_py.utils import discover_mujoco
+from mujoco_py.version import get_version
+
+
+def get_nvidia_lib_dir():
+    exists_nvidia_smi = subprocess.call("type nvidia-smi", shell=True,
+                                        stdout=subprocess.PIPE, stderr=subprocess.PIPE) == 0
+    if not exists_nvidia_smi:
+        return None
+    docker_path = '/usr/local/nvidia/lib64'
+    if exists(docker_path):
+        return docker_path
+
+    nvidia_path = '/usr/lib/nvidia'
+    if exists(nvidia_path):
+        return nvidia_path
+
+    paths = glob.glob('/usr/lib/nvidia-[0-9][0-9][0-9]')
+    paths = sorted(paths)
+    if len(paths) == 0:
+        return None
+    if len(paths) > 1:
+        print("Choosing the latest nvidia driver: %s, among %s" % (paths[-1], str(paths)))
+
+    return paths[-1]
+
+
+def load_cython_ext(mujoco_path):
+    """
+    Loads the cymj Cython extension. This is safe to be called from
+    multiple processes running on the same machine.
+
+    Cython only gives us back the raw path, regardless of whether
+    it found a cached version or actually compiled. Since we do
+    non-idempotent postprocessing of the DLL, be extra careful
+    to only do that once and then atomically move to the final
+    location.
+    """
+    if ('glfw' in sys.modules and
+            'mujoco' in abspath(sys.modules["glfw"].__file__)):
+        print('''
+WARNING: Existing glfw python module detected!
+
+MuJoCo comes with its own version of GLFW, so it's preferable to use that one.
+
+The easy solution is to `import mujoco_py` _before_ `import glfw`.
+''')
+
+    lib_path = os.path.join(mujoco_path, "bin")
+    if sys.platform == 'darwin':
+        Builder = MacExtensionBuilder
+    elif sys.platform == 'linux':
+        _ensure_set_env_var("LD_LIBRARY_PATH", lib_path)
+        if os.getenv('MUJOCO_PY_FORCE_CPU') is None and get_nvidia_lib_dir() is not None:
+            _ensure_set_env_var("LD_LIBRARY_PATH", get_nvidia_lib_dir())
+            Builder = LinuxGPUExtensionBuilder
+        else:
+            Builder = LinuxGPUExtensionBuilder
+    elif sys.platform.startswith("win"):
+        var = "PATH"
+        if var not in os.environ or lib_path not in os.environ[var].split(";"):
+            raise Exception("Please add mujoco library to your PATH:\n"
+                            "set %s=%s;%%%s%%" % (var, lib_path, var))
+        Builder = WindowsExtensionBuilder
+    else:
+        raise RuntimeError("Unsupported platform %s" % sys.platform)
+
+    builder = Builder(mujoco_path)
+    cext_so_path = builder.get_so_file_path()
+
+    lockpath = os.path.join(os.path.dirname(cext_so_path), 'mujocopy-buildlock')
+
+    with fasteners.InterProcessLock(lockpath):
+        mod = None
+        force_rebuild = os.environ.get('MUJOCO_PY_FORCE_REBUILD')
+        if force_rebuild:
+            # Try to remove the old file, ignore errors if it doesn't exist
+            print("Removing old mujoco_py cext", cext_so_path)
+            try:
+                os.remove(cext_so_path)
+            except OSError:
+                pass
+        if exists(cext_so_path):
+            try:
+                mod = load_dynamic_ext('cymj', cext_so_path)
+            except ImportError:
+                print("Import error. Trying to rebuild mujoco_py.")
+        if mod is None:
+            cext_so_path = builder.build()
+            mod = load_dynamic_ext('cymj', cext_so_path)
+
+    return mod
+
+
+def _ensure_set_env_var(var_name, lib_path):
+    paths = os.environ.get(var_name, "").split(":")
+    paths = [os.path.abspath(path) for path in paths]
+    if lib_path not in paths:
+        raise Exception("\nMissing path to your environment variable. \n"
+                        "Current values %s=%s\n"
+                        "Please add following line to .bashrc:\n"
+                        "export %s=$%s:%s" % (var_name, os.environ.get(var_name, ""),
+                                              var_name, var_name, lib_path))
+
+
+def load_dynamic_ext(name, path):
+    """ Load compiled shared object and return as python module. """
+    loader = ExtensionFileLoader(name, path)
+    return loader.load_module()
+
+
+class custom_build_ext(build_ext):
+    """
+    Custom build_ext to suppress the "-Wstrict-prototypes" warning.
+    It arises from the fact that we're using C++. This seems to be
+    the cleanest way to get rid of the extra flag.
+
+    See http://stackoverflow.com/a/36293331/248400
+    """
+
+    def build_extensions(self):
+        customize_compiler(self.compiler)
+
+        try:
+            self.compiler.compiler_so.remove("-Wstrict-prototypes")
+        except (AttributeError, ValueError):
+            pass
+        build_ext.build_extensions(self)
+
+
+def fix_shared_library(so_file, name, library_path):
+    """ Used to fixup shared libraries on Linux """
+    subprocess.check_call(['patchelf', '--remove-rpath', so_file])
+    ldd_output = subprocess.check_output(['ldd', so_file]).decode('utf-8')
+
+    if name in ldd_output:
+        subprocess.check_call(['patchelf', '--remove-needed', name, so_file])
+    subprocess.check_call(['patchelf', '--add-needed', library_path, so_file])
+
+
+def manually_link_libraries(mujoco_path, raw_cext_dll_path):
+    """ Used to fix mujoco library linking on Mac """
+    root, ext = os.path.splitext(raw_cext_dll_path)
+    final_cext_dll_path = root + '_final' + ext
+
+    # If someone else already built the final DLL, don't bother
+    # recreating it here, even though this should still be idempotent.
+    if (exists(final_cext_dll_path) and
+            getmtime(final_cext_dll_path) >= getmtime(raw_cext_dll_path)):
+        return final_cext_dll_path
+
+    tmp_final_cext_dll_path = final_cext_dll_path + '~'
+    shutil.copyfile(raw_cext_dll_path, tmp_final_cext_dll_path)
+
+    mj_bin_path = join(mujoco_path, 'bin')
+
+    # Fix the rpath of the generated library -- i lost the Stackoverflow
+    # reference here
+    from_mujoco_path = '@executable_path/libmujoco210.dylib'
+    to_mujoco_path = '%s/libmujoco210.dylib' % mj_bin_path
+    subprocess.check_call(['install_name_tool',
+                           '-change',
+                           from_mujoco_path,
+                           to_mujoco_path,
+                           tmp_final_cext_dll_path])
+
+    from_glfw_path = 'libglfw.3.dylib'
+    to_glfw_path = os.path.join(mj_bin_path, 'libglfw.3.dylib')
+    subprocess.check_call(['install_name_tool',
+                           '-change',
+                           from_glfw_path,
+                           to_glfw_path,
+                           tmp_final_cext_dll_path])
+
+    os.rename(tmp_final_cext_dll_path, final_cext_dll_path)
+    return final_cext_dll_path
+
+
+class MujocoExtensionBuilder():
+
+    CYMJ_DIR_PATH = abspath(dirname(__file__))
+
+    def __init__(self, mujoco_path):
+        self.mujoco_path = mujoco_path
+        python_version = str(sys.version_info.major) + str(sys.version_info.minor)
+        self.version = '%s_%s_%s' % (get_version(), python_version, self.build_base())
+        self.extension = Extension(
+            'mujoco_py.cymj',
+            sources=[join(self.CYMJ_DIR_PATH, "cymj.pyx")],
+            include_dirs=[
+                self.CYMJ_DIR_PATH,
+                join(mujoco_path, 'include'),
+                np.get_include(),
+            ],
+            libraries=['mujoco210'],
+            library_dirs=[join(mujoco_path, 'bin')],
+            extra_compile_args=[
+                '-fopenmp',  # needed for OpenMP
+                '-w',  # suppress numpy compilation warnings
+            ],
+            extra_link_args=['-fopenmp'],
+            language='c')
+
+    def build(self):
+        built_so_file_path = self._build_impl()
+        new_so_file_path = self.get_so_file_path()
+        move(built_so_file_path, new_so_file_path)
+        return new_so_file_path
+
+    def build_base(self):
+        return self.__class__.__name__.lower()
+
+    def _build_impl(self):
+        dist = Distribution({
+            "script_name": None,
+            "script_args": ["build_ext"]
+        })
+        dist.ext_modules = cythonize([self.extension])
+        dist.include_dirs = []
+        dist.cmdclass = {'build_ext': custom_build_ext}
+        build = dist.get_command_obj('build')
+        # following the convention of cython's pyxbuild and naming
+        # base directory "_pyxbld"
+        build.build_base = join(self.CYMJ_DIR_PATH, 'generated',
+                                '_pyxbld_%s' % (self.version))
+        dist.parse_command_line()
+        obj_build_ext = dist.get_command_obj("build_ext")
+        dist.run_commands()
+        built_so_file_path, = obj_build_ext.get_outputs()
+        return built_so_file_path
+
+    def get_so_file_path(self):
+        dir_path = abspath(dirname(__file__))
+        python_version = str(sys.version_info.major) + str(sys.version_info.minor)
+        return join(dir_path, "generated", "cymj_{}_{}.so".format(self.version, python_version))
+
+
+class WindowsExtensionBuilder(MujocoExtensionBuilder):
+
+    def __init__(self, mujoco_path):
+        super().__init__(mujoco_path)
+        os.environ["PATH"] += ";" + join(mujoco_path, "bin")
+        self.extension.sources.append(self.CYMJ_DIR_PATH + "/gl/dummyshim.c")
+
+
+class LinuxCPUExtensionBuilder(MujocoExtensionBuilder):
+
+    def __init__(self, mujoco_path):
+        super().__init__(mujoco_path)
+
+        self.extension.sources.append(
+            join(self.CYMJ_DIR_PATH, "gl", "osmesashim.c"))
+        self.extension.libraries.extend(['glewosmesa', 'OSMesa', 'GL'])
+        self.extension.runtime_library_dirs = [join(mujoco_path, 'bin')]
+
+    def _build_impl(self):
+        so_file_path = super()._build_impl()
+        # Removes absolute paths to libraries. Allows for dynamic loading.
+        fix_shared_library(so_file_path, 'libmujoco210.so', 'libmujoco210.so')
+        fix_shared_library(so_file_path, 'libglewosmesa.so', 'libglewosmesa.so')
+        return so_file_path
+
+
+class LinuxGPUExtensionBuilder(MujocoExtensionBuilder):
+
+    def __init__(self, mujoco_path):
+        super().__init__(mujoco_path)
+
+        self.extension.sources.append(self.CYMJ_DIR_PATH + "/gl/eglshim.c")
+        self.extension.include_dirs.append(self.CYMJ_DIR_PATH + '/vendor/egl')
+        self.extension.libraries.extend(['glewegl'])
+        self.extension.runtime_library_dirs = [join(mujoco_path, 'bin')]
+
+    def _build_impl(self):
+        so_file_path = super()._build_impl()
+        fix_shared_library(so_file_path, 'libOpenGL.so', 'libOpenGL.so.0')
+        fix_shared_library(so_file_path, 'libEGL.so', 'libEGL.so.1')
+        fix_shared_library(so_file_path, 'libmujoco210.so', 'libmujoco210.so')
+        fix_shared_library(so_file_path, 'libglewegl.so', 'libglewegl.so')
+        return so_file_path
+
+
+class MacExtensionBuilder(MujocoExtensionBuilder):
+
+    def __init__(self, mujoco_path):
+        super().__init__(mujoco_path)
+
+        self.extension.sources.append(self.CYMJ_DIR_PATH + "/gl/dummyshim.c")
+        self.extension.libraries.extend(['glfw.3'])
+        self.extension.define_macros = [('ONMAC', None)]
+        self.extension.runtime_library_dirs = [join(mujoco_path, 'bin')]
+
+    def _build_impl(self):
+        if not os.environ.get('CC'):
+            # Known-working versions of GCC on mac (prefer latest one)
+            c_compilers = [
+                '/usr/local/bin/gcc-9',
+                '/usr/local/bin/gcc-8',
+                '/usr/local/bin/gcc-7',
+                '/usr/local/bin/gcc-6',
+                '/opt/local/bin/gcc-mp-9',
+                '/opt/local/bin/gcc-mp-8',
+                '/opt/local/bin/gcc-mp-7',
+                '/opt/local/bin/gcc-mp-6',
+            ]
+            available_c_compiler = None
+            for c_compiler in c_compilers:
+                if distutils.spawn.find_executable(c_compiler) is not None:
+                    available_c_compiler = c_compiler
+                    break
+            if available_c_compiler is None:
+                raise RuntimeError(
+                    'Could not find supported GCC executable.\n\n'
+                    'HINT: On OS X, install GCC 9.x with '
+                    '`brew install gcc@9`. or '
+                    '`port install gcc9`.')
+            os.environ['CC'] = available_c_compiler
+
+            so_file_path = super()._build_impl()
+            del os.environ['CC']
+        else:  # User-directed c compiler
+            so_file_path = super()._build_impl()
+        return manually_link_libraries(self.mujoco_path, so_file_path)
+
+
+class MujocoException(Exception):
+    pass
+
+
+def user_warning_raise_exception(warn_bytes):
+    '''
+    User-defined warning callback, which is called by mujoco on warnings.
+    Here we have two primary jobs:
+        - Detect known warnings and suggest fixes (with code)
+        - Decide whether to raise an Exception and raise if needed
+    More cases should be added as we find new failures.
+    '''
+    # TODO: look through test output to see MuJoCo warnings to catch
+    # and recommend. Also fix those tests
+    warn = warn_bytes.decode()  # Convert bytes to string
+    if 'Pre-allocated constraint buffer is full' in warn:
+        raise MujocoException(warn + 'Increase njmax in mujoco XML')
+    if 'Pre-allocated contact buffer is full' in warn:
+        raise MujocoException(warn + 'Increase njconmax in mujoco XML')
+    # This unhelpfully-named warning is what you get if you feed MuJoCo NaNs
+    if 'Unknown warning type' in warn:
+        raise MujocoException(warn + 'Check for NaN in simulation.')
+    raise MujocoException('Got MuJoCo Warning: {}'.format(warn))
+
+
+def user_warning_ignore_exception(warn_bytes):
+    pass
+
+
+class ignore_mujoco_warnings:
+    """
+    Class to turn off mujoco warning exceptions within a scope. Useful for
+    large, vectorized rollouts.
+    """
+
+    def __enter__(self):
+        self.prev_user_warning = cymj.get_warning_callback()
+        cymj.set_warning_callback(user_warning_ignore_exception)
+        return self
+
+    def __exit__(self, type, value, traceback):
+        cymj.set_warning_callback(self.prev_user_warning)
+
+
+def build_fn_cleanup(name):
+    '''
+    Cleanup files generated by building callback.
+    Set the MUJOCO_PY_DEBUG_FN_BUILDER environment variable to disable cleanup.
+    '''
+    if not os.environ.get('MUJOCO_PY_DEBUG_FN_BUILDER', False):
+        for f in glob.glob(name + '*'):
+            try:
+                os.remove(f)
+            except PermissionError as e:
+                # This happens trying to remove libraries on appveyor
+                print('Error removing {}, continuing anyway: {}'.format(f, e))
+
+
+def build_callback_fn(function_string, userdata_names=[]):
+    '''
+    Builds a C callback function and returns a function pointer int.
+
+        function_string : str
+            This is a string of the C function to be compiled
+        userdata_names : list or tuple
+            This is an optional list to defince convenience names
+
+    We compile and link and load the function, and return a function pointer.
+    See `MjSim.set_substep_callback()` for an example use of these callbacks.
+
+    The callback function should match the signature:
+        void fun(const mjModel *m, mjData *d);
+
+    Here's an example function_string:
+        ```
+        """
+        #include <stdio.h>
+        void fun(const mjModel* m, mjData* d) {
+            printf("hello");
+        }
+        """
+        ```
+
+    Input and output for the function pass through userdata in the data struct:
+        ```
+        """
+        void fun(const mjModel* m, mjData* d) {
+            d->userdata[0] += 1;
+        }
+        """
+        ```
+
+    `userdata_names` is expected to match the model where the callback is used.
+    These can bet set on a model with:
+        `model.set_userdata_names([...])`
+
+    If `userdata_names` is supplied, convenience `#define`s are added for each.
+    For example:
+        `userdata_names = ['my_sum']`
+    Will get gerenerated into the extra line:
+        `#define my_sum d->userdata[0]`
+    And prepended to the top of the function before compilation.
+    Here's an example that takes advantage of this:
+        ```
+        """
+        void fun(const mjModel* m, mjData* d) {
+            for (int i = 0; i < m->nu; i++) {
+                my_sum += d->ctrl[i];
+            }
+        }
+        """
+        ```
+    Note these are just C `#define`s and are limited in how they can be used.
+
+    After compilation, the built library containing the function is loaded
+    into memory and all of the files (including the library) are deleted.
+    To retain these for debugging set the `MUJOCO_PY_DEBUG_FN_BUILDER` envvar.
+
+    To save time compiling, these function pointers may be re-used by many
+    different consumers.  They are thread-safe and don't acquire the GIL.
+
+    See the file `tests/test_substep.py` for additional examples,
+    including an example which iterates over contacts to compute penetrations.
+    '''
+    assert isinstance(userdata_names, (list, tuple)), \
+        'invalid userdata_names: {}'.format(userdata_names)
+    ffibuilder = FFI()
+    ffibuilder.cdef('extern uintptr_t __fun;')
+    name = '_fn_' + ''.join(choice(ascii_lowercase) for _ in range(15))
+    source_string = '#include <mujoco.h>\n'
+    # Add defines for each userdata to make setting them easier
+    for i, data_name in enumerate(userdata_names):
+        source_string += '#define {} d->userdata[{}]\n'.format(data_name, i)
+    source_string += function_string
+    source_string += '\nuintptr_t __fun = (uintptr_t) fun;'
+    # Link against mujoco so we can call mujoco functions from within callback
+    ffibuilder.set_source(name, source_string,
+                          include_dirs=[join(mujoco_path, 'include')],
+                          library_dirs=[join(mujoco_path, 'bin')],
+                          libraries=['mujoco210'])
+    # Catch compilation exceptions so we can cleanup partial files in that case
+    try:
+        library_path = ffibuilder.compile(verbose=True)
+    except Exception as e:
+        build_fn_cleanup(name)
+        raise e
+    # On Mac the MuJoCo library is linked strangely, so we have to fix it here
+    if sys.platform == 'darwin':
+        fixed_library_path = manually_link_libraries(mujoco_path, library_path)
+        move(fixed_library_path, library_path)  # Overwrite with fixed library
+    module = load_dynamic_ext(name, library_path)
+    # Now that the module is loaded into memory, we can actually delete it
+    build_fn_cleanup(name)
+    return module.lib.__fun
+
+
+mujoco_path = discover_mujoco()
+cymj = load_cython_ext(mujoco_path)
+
+
+# Trick to expose all mj* functions from mujoco in mujoco_py.*
+class dict2(object):
+    pass
+
+
+functions = dict2()
+for func_name in dir(cymj):
+    if func_name.startswith("_mj"):
+        setattr(functions, func_name[1:], getattr(cymj, func_name))
+
+# Set user-defined callbacks that raise assertion with message
+cymj.set_warning_callback(user_warning_raise_exception)

mujoco-py-2.1.2.14/mujoco_py/gl/eglplatform.h

@@ -0,0 +1,125 @@
+#ifndef __eglplatform_h_
+#define __eglplatform_h_
+
+/*
+** Copyright (c) 2007-2013 The Khronos Group Inc.
+**
+** Permission is hereby granted, free of charge, to any person obtaining a
+** copy of this software and/or associated documentation files (the
+** "Materials"), to deal in the Materials without restriction, including
+** without limitation the rights to use, copy, modify, merge, publish,
+** distribute, sublicense, and/or sell copies of the Materials, and to
+** permit persons to whom the Materials are 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 Materials.
+**
+** THE MATERIALS ARE 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
+** MATERIALS OR THE USE OR OTHER DEALINGS IN THE MATERIALS.
+*/
+
+/* Platform-specific types and definitions for egl.h
+ * $Revision: 30994 $ on $Date: 2015-04-30 13:36:48 -0700 (Thu, 30 Apr 2015) $
+ *
+ * Adopters may modify khrplatform.h and this file to suit their platform.
+ * You are encouraged to submit all modifications to the Khronos group so that
+ * they can be included in future versions of this file.  Please submit changes
+ * by sending them to the public Khronos Bugzilla (http://khronos.org/bugzilla)
+ * by filing a bug against product "EGL" component "Registry".
+ */
+
+#include "khrplatform.h"
+
+/* Macros used in EGL function prototype declarations.
+ *
+ * EGL functions should be prototyped as:
+ *
+ * EGLAPI return-type EGLAPIENTRY eglFunction(arguments);
+ * typedef return-type (EXPAPIENTRYP PFNEGLFUNCTIONPROC) (arguments);
+ *
+ * KHRONOS_APICALL and KHRONOS_APIENTRY are defined in KHR/khrplatform.h
+ */
+
+#ifndef EGLAPI
+#define EGLAPI KHRONOS_APICALL
+#endif
+
+#ifndef EGLAPIENTRY
+#define EGLAPIENTRY  KHRONOS_APIENTRY
+#endif
+#define EGLAPIENTRYP EGLAPIENTRY*
+
+/* The types NativeDisplayType, NativeWindowType, and NativePixmapType
+ * are aliases of window-system-dependent types, such as X Display * or
+ * Windows Device Context. They must be defined in platform-specific
+ * code below. The EGL-prefixed versions of Native*Type are the same
+ * types, renamed in EGL 1.3 so all types in the API start with "EGL".
+ *
+ * Khronos STRONGLY RECOMMENDS that you use the default definitions
+ * provided below, since these changes affect both binary and source
+ * portability of applications using EGL running on different EGL
+ * implementations.
+ */
+
+#if defined(_WIN32) || defined(__VC32__) && !defined(__CYGWIN__) && !defined(__SCITECH_SNAP__) /* Win32 and WinCE */
+#ifndef WIN32_LEAN_AND_MEAN
+#define WIN32_LEAN_AND_MEAN 1
+#endif
+#include <windows.h>
+
+typedef HDC     EGLNativeDisplayType;
+typedef HBITMAP EGLNativePixmapType;
+typedef HWND    EGLNativeWindowType;
+
+#elif defined(__APPLE__) || defined(__WINSCW__) || defined(__SYMBIAN32__)  /* Symbian */
+
+typedef int   EGLNativeDisplayType;
+typedef void *EGLNativeWindowType;
+typedef void *EGLNativePixmapType;
+
+#elif defined(__ANDROID__) || defined(ANDROID)
+
+#include <android/native_window.h>
+
+struct egl_native_pixmap_t;
+
+typedef struct ANativeWindow*           EGLNativeWindowType;
+typedef struct egl_native_pixmap_t*     EGLNativePixmapType;
+typedef void*                           EGLNativeDisplayType;
+
+#elif defined(__unix__)
+
+/* X11 (tentative)  */
+#include <X11/Xlib.h>
+#include <X11/Xutil.h>
+
+typedef Display *EGLNativeDisplayType;
+typedef Pixmap   EGLNativePixmapType;
+typedef Window   EGLNativeWindowType;
+
+#else
+#error "Platform not recognized"
+#endif
+
+/* EGL 1.2 types, renamed for consistency in EGL 1.3 */
+typedef EGLNativeDisplayType NativeDisplayType;
+typedef EGLNativePixmapType  NativePixmapType;
+typedef EGLNativeWindowType  NativeWindowType;
+
+
+/* Define EGLint. This must be a signed integral type large enough to contain
+ * all legal attribute names and values passed into and out of EGL, whether
+ * their type is boolean, bitmask, enumerant (symbolic constant), integer,
+ * handle, or other.  While in general a 32-bit integer will suffice, if
+ * handles are 64 bit types, then EGLint should be defined as a signed 64-bit
+ * integer type.
+ */
+typedef khronos_int32_t EGLint;
+
+#endif /* __eglplatform_h */

mujoco-py-2.1.2.14/mujoco_py/gl/glshim.h

@@ -0,0 +1,30 @@
+#ifndef __GLSHIM_H__
+#define __GLSHIM_H__
+
+#include "mujoco.h"
+#include "mjrender.h"
+
+#ifdef __cplusplus
+extern "C" {
+#endif
+
+int usingEGL();
+int initOpenGL(int device_id);
+void closeOpenGL();
+int makeOpenGLContextCurrent(int device_id);
+int setOpenGLBufferSize(int device_id, int width, int height);
+
+unsigned int createPBO(int width, int height, int batchSize, int use_short);
+void freePBO(unsigned int pixelBuffer);
+void copyFBOToPBO(mjrContext* con,
+                  unsigned int pbo_rgb, unsigned int pbo_depth,
+                  mjrRect viewport, int bufferOffset);
+void readPBO(unsigned char *buffer_rgb, unsigned short *buffer_depth,
+             unsigned int pbo_rgb, unsigned int pbo_depth,
+             int width, int height, int batchSize);
+
+#ifdef __cplusplus
+}  // extern "C"
+#endif
+
+#endif

mujoco-py-2.1.2.14/mujoco_py/gl/khrplatform.h

@@ -0,0 +1,285 @@
+#ifndef __khrplatform_h_
+#define __khrplatform_h_
+
+/*
+** Copyright (c) 2008-2009 The Khronos Group Inc.
+**
+** Permission is hereby granted, free of charge, to any person obtaining a
+** copy of this software and/or associated documentation files (the
+** "Materials"), to deal in the Materials without restriction, including
+** without limitation the rights to use, copy, modify, merge, publish,
+** distribute, sublicense, and/or sell copies of the Materials, and to
+** permit persons to whom the Materials are 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 Materials.
+**
+** THE MATERIALS ARE 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
+** MATERIALS OR THE USE OR OTHER DEALINGS IN THE MATERIALS.
+*/
+
+/* Khronos platform-specific types and definitions.
+ *
+ * $Revision: 32517 $ on $Date: 2016-03-11 02:41:19 -0800 (Fri, 11 Mar 2016) $
+ *
+ * Adopters may modify this file to suit their platform. Adopters are
+ * encouraged to submit platform specific modifications to the Khronos
+ * group so that they can be included in future versions of this file.
+ * Please submit changes by sending them to the public Khronos Bugzilla
+ * (http://khronos.org/bugzilla) by filing a bug against product
+ * "Khronos (general)" component "Registry".
+ *
+ * A predefined template which fills in some of the bug fields can be
+ * reached using http://tinyurl.com/khrplatform-h-bugreport, but you
+ * must create a Bugzilla login first.
+ *
+ *
+ * See the Implementer's Guidelines for information about where this file
+ * should be located on your system and for more details of its use:
+ *    http://www.khronos.org/registry/implementers_guide.pdf
+ *
+ * This file should be included as
+ *        #include <KHR/khrplatform.h>
+ * by Khronos client API header files that use its types and defines.
+ *
+ * The types in khrplatform.h should only be used to define API-specific types.
+ *
+ * Types defined in khrplatform.h:
+ *    khronos_int8_t              signed   8  bit
+ *    khronos_uint8_t             unsigned 8  bit
+ *    khronos_int16_t             signed   16 bit
+ *    khronos_uint16_t            unsigned 16 bit
+ *    khronos_int32_t             signed   32 bit
+ *    khronos_uint32_t            unsigned 32 bit
+ *    khronos_int64_t             signed   64 bit
+ *    khronos_uint64_t            unsigned 64 bit
+ *    khronos_intptr_t            signed   same number of bits as a pointer
+ *    khronos_uintptr_t           unsigned same number of bits as a pointer
+ *    khronos_ssize_t             signed   size
+ *    khronos_usize_t             unsigned size
+ *    khronos_float_t             signed   32 bit floating point
+ *    khronos_time_ns_t           unsigned 64 bit time in nanoseconds
+ *    khronos_utime_nanoseconds_t unsigned time interval or absolute time in
+ *                                         nanoseconds
+ *    khronos_stime_nanoseconds_t signed time interval in nanoseconds
+ *    khronos_boolean_enum_t      enumerated boolean type. This should
+ *      only be used as a base type when a client API's boolean type is
+ *      an enum. Client APIs which use an integer or other type for
+ *      booleans cannot use this as the base type for their boolean.
+ *
+ * Tokens defined in khrplatform.h:
+ *
+ *    KHRONOS_FALSE, KHRONOS_TRUE Enumerated boolean false/true values.
+ *
+ *    KHRONOS_SUPPORT_INT64 is 1 if 64 bit integers are supported; otherwise 0.
+ *    KHRONOS_SUPPORT_FLOAT is 1 if floats are supported; otherwise 0.
+ *
+ * Calling convention macros defined in this file:
+ *    KHRONOS_APICALL
+ *    KHRONOS_APIENTRY
+ *    KHRONOS_APIATTRIBUTES
+ *
+ * These may be used in function prototypes as:
+ *
+ *      KHRONOS_APICALL void KHRONOS_APIENTRY funcname(
+ *                                  int arg1,
+ *                                  int arg2) KHRONOS_APIATTRIBUTES;
+ */
+
+/*-------------------------------------------------------------------------
+ * Definition of KHRONOS_APICALL
+ *-------------------------------------------------------------------------
+ * This precedes the return type of the function in the function prototype.
+ */
+#if defined(_WIN32) && !defined(__SCITECH_SNAP__)
+#   define KHRONOS_APICALL __declspec(dllimport)
+#elif defined (__SYMBIAN32__)
+#   define KHRONOS_APICALL IMPORT_C
+#elif defined(__ANDROID__)
+#   include <sys/cdefs.h>
+#   define KHRONOS_APICALL __attribute__((visibility("default"))) __NDK_FPABI__
+#else
+#   define KHRONOS_APICALL
+#endif
+
+/*-------------------------------------------------------------------------
+ * Definition of KHRONOS_APIENTRY
+ *-------------------------------------------------------------------------
+ * This follows the return type of the function  and precedes the function
+ * name in the function prototype.
+ */
+#if defined(_WIN32) && !defined(_WIN32_WCE) && !defined(__SCITECH_SNAP__)
+    /* Win32 but not WinCE */
+#   define KHRONOS_APIENTRY __stdcall
+#else
+#   define KHRONOS_APIENTRY
+#endif
+
+/*-------------------------------------------------------------------------
+ * Definition of KHRONOS_APIATTRIBUTES
+ *-------------------------------------------------------------------------
+ * This follows the closing parenthesis of the function prototype arguments.
+ */
+#if defined (__ARMCC_2__)
+#define KHRONOS_APIATTRIBUTES __softfp
+#else
+#define KHRONOS_APIATTRIBUTES
+#endif
+
+/*-------------------------------------------------------------------------
+ * basic type definitions
+ *-----------------------------------------------------------------------*/
+#if (defined(__STDC_VERSION__) && __STDC_VERSION__ >= 199901L) || defined(__GNUC__) || defined(__SCO__) || defined(__USLC__)
+
+
+/*
+ * Using <stdint.h>
+ */
+#include <stdint.h>
+typedef int32_t                 khronos_int32_t;
+typedef uint32_t                khronos_uint32_t;
+typedef int64_t                 khronos_int64_t;
+typedef uint64_t                khronos_uint64_t;
+#define KHRONOS_SUPPORT_INT64   1
+#define KHRONOS_SUPPORT_FLOAT   1
+
+#elif defined(__VMS ) || defined(__sgi)
+
+/*
+ * Using <inttypes.h>
+ */
+#include <inttypes.h>
+typedef int32_t                 khronos_int32_t;
+typedef uint32_t                khronos_uint32_t;
+typedef int64_t                 khronos_int64_t;
+typedef uint64_t                khronos_uint64_t;
+#define KHRONOS_SUPPORT_INT64   1
+#define KHRONOS_SUPPORT_FLOAT   1
+
+#elif defined(_WIN32) && !defined(__SCITECH_SNAP__)
+
+/*
+ * Win32
+ */
+typedef __int32                 khronos_int32_t;
+typedef unsigned __int32        khronos_uint32_t;
+typedef __int64                 khronos_int64_t;
+typedef unsigned __int64        khronos_uint64_t;
+#define KHRONOS_SUPPORT_INT64   1
+#define KHRONOS_SUPPORT_FLOAT   1
+
+#elif defined(__sun__) || defined(__digital__)
+
+/*
+ * Sun or Digital
+ */
+typedef int                     khronos_int32_t;
+typedef unsigned int            khronos_uint32_t;
+#if defined(__arch64__) || defined(_LP64)
+typedef long int                khronos_int64_t;
+typedef unsigned long int       khronos_uint64_t;
+#else
+typedef long long int           khronos_int64_t;
+typedef unsigned long long int  khronos_uint64_t;
+#endif /* __arch64__ */
+#define KHRONOS_SUPPORT_INT64   1
+#define KHRONOS_SUPPORT_FLOAT   1
+
+#elif 0
+
+/*
+ * Hypothetical platform with no float or int64 support
+ */
+typedef int                     khronos_int32_t;
+typedef unsigned int            khronos_uint32_t;
+#define KHRONOS_SUPPORT_INT64   0
+#define KHRONOS_SUPPORT_FLOAT   0
+
+#else
+
+/*
+ * Generic fallback
+ */
+#include <stdint.h>
+typedef int32_t                 khronos_int32_t;
+typedef uint32_t                khronos_uint32_t;
+typedef int64_t                 khronos_int64_t;
+typedef uint64_t                khronos_uint64_t;
+#define KHRONOS_SUPPORT_INT64   1
+#define KHRONOS_SUPPORT_FLOAT   1
+
+#endif
+
+
+/*
+ * Types that are (so far) the same on all platforms
+ */
+typedef signed   char          khronos_int8_t;
+typedef unsigned char          khronos_uint8_t;
+typedef signed   short int     khronos_int16_t;
+typedef unsigned short int     khronos_uint16_t;
+
+/*
+ * Types that differ between LLP64 and LP64 architectures - in LLP64,
+ * pointers are 64 bits, but 'long' is still 32 bits. Win64 appears
+ * to be the only LLP64 architecture in current use.
+ */
+#ifdef _WIN64
+typedef signed   long long int khronos_intptr_t;
+typedef unsigned long long int khronos_uintptr_t;
+typedef signed   long long int khronos_ssize_t;
+typedef unsigned long long int khronos_usize_t;
+#else
+typedef signed   long  int     khronos_intptr_t;
+typedef unsigned long  int     khronos_uintptr_t;
+typedef signed   long  int     khronos_ssize_t;
+typedef unsigned long  int     khronos_usize_t;
+#endif
+
+#if KHRONOS_SUPPORT_FLOAT
+/*
+ * Float type
+ */
+typedef          float         khronos_float_t;
+#endif
+
+#if KHRONOS_SUPPORT_INT64
+/* Time types
+ *
+ * These types can be used to represent a time interval in nanoseconds or
+ * an absolute Unadjusted System Time.  Unadjusted System Time is the number
+ * of nanoseconds since some arbitrary system event (e.g. since the last
+ * time the system booted).  The Unadjusted System Time is an unsigned
+ * 64 bit value that wraps back to 0 every 584 years.  Time intervals
+ * may be either signed or unsigned.
+ */
+typedef khronos_uint64_t       khronos_utime_nanoseconds_t;
+typedef khronos_int64_t        khronos_stime_nanoseconds_t;
+#endif
+
+/*
+ * Dummy value used to pad enum types to 32 bits.
+ */
+#ifndef KHRONOS_MAX_ENUM
+#define KHRONOS_MAX_ENUM 0x7FFFFFFF
+#endif
+
+/*
+ * Enumerated boolean type
+ *
+ * Values other than zero should be considered to be true.  Therefore
+ * comparisons should not be made against KHRONOS_TRUE.
+ */
+typedef enum {
+    KHRONOS_FALSE = 0,
+    KHRONOS_TRUE  = 1,
+    KHRONOS_BOOLEAN_ENUM_FORCE_SIZE = KHRONOS_MAX_ENUM
+} khronos_boolean_enum_t;
+
+#endif /* __khrplatform_h_ */

mujoco-py-2.1.2.14/mujoco_py/gl/osmesashim.c

@@ -0,0 +1,75 @@
+#include <GL/osmesa.h>
+#include "glshim.h"
+
+OSMesaContext ctx;
+
+// this size was picked pretty arbitrarily
+int BUFFER_WIDTH = 1024;
+int BUFFER_HEIGHT = 1024;
+// 4 channels for RGBA
+unsigned char buffer[1024 * 1024 * 4];
+
+int is_initialized = 0;
+
+int usingEGL() {
+    return 0;
+}
+
+int initOpenGL(int device_id) {
+    if (is_initialized)
+        return 1;
+
+    // note: device id not used here
+    ctx = OSMesaCreateContextExt(GL_RGBA, 24, 8, 8, 0);
+    if( !ctx ) {
+        printf("OSMesa context creation failed\n");
+        return -1;
+    }
+
+    if( !OSMesaMakeCurrent(ctx, buffer, GL_UNSIGNED_BYTE, BUFFER_WIDTH, BUFFER_HEIGHT) ) {
+        printf("OSMesa make current failed\n");
+        return -1;
+    }
+
+    is_initialized = 1;
+
+    return 1;
+}
+
+int makeOpenGLContextCurrent(int device_id) {
+    // Don't need to make context current here, causes issues with large tests
+    return 1;
+}
+
+int setOpenGLBufferSize(int device_id, int width, int height) {
+    if (width > BUFFER_WIDTH || height > BUFFER_HEIGHT) {
+        printf("Buffer size too big\n");
+        return -1;
+    }
+    // Noop since we don't support changing the actual buffer
+    return 1;
+}
+
+void closeOpenGL() {
+    if (is_initialized) {
+        OSMesaDestroyContext(ctx);
+        is_initialized = 0;
+    }
+}
+
+unsigned int createPBO(int width, int height, int batchSize, int use_short) {
+    return 0;
+}
+
+void freePBO(unsigned int pixelBuffer) {
+}
+
+void copyFBOToPBO(mjrContext* con,
+                  unsigned int pbo_rgb, unsigned int pbo_depth,
+                  mjrRect viewport, int bufferOffset) {
+}
+
+void readPBO(unsigned char *buffer_rgb, unsigned short *buffer_depth,
+             unsigned int pbo_rgb, unsigned int pbo_depth,
+             int width, int height, int batchSize) {
+}

mujoco-py-2.1.2.14/mujoco_py/mjbatchrenderer.pyx

@@ -0,0 +1,301 @@
+try:
+    import pycuda.driver as drv
+except ImportError:
+    drv = None
+
+
+class MjBatchRendererException(Exception):
+    pass
+
+
+class MjBatchRendererNotSupported(MjBatchRendererException):
+    pass
+
+
+class CudaNotEnabledError(MjBatchRendererException):
+    pass
+
+
+class CudaBufferNotMappedError(MjBatchRendererException):
+    pass
+
+
+class CudaBufferMappedError(MjBatchRendererException):
+    pass
+
+
+class MjBatchRenderer(object):
+    """
+    Utility class for rendering into OpenGL Pixel Buffer Objects (PBOs),
+    which allows for accessing multiple rendered images in batch.
+
+    If used with CUDA (i.e. initialized with use_cuda=True), you need
+    to call map/unmap when accessing CUDA buffer pointer. This is to
+    ensure that all OpenGL instructions have completed:
+
+        renderer = MjBatchRenderer(100, 100, use_cuda=True)
+        renderer.render(sim)
+
+        renderer.map()
+        image = renderer.read()
+        renderer.unmap()
+    """
+
+    def __init__(self, width, height, batch_size=1, device_id=0,
+                 depth=False, use_cuda=False):
+        """
+        Args:
+        - width (int): Image width.
+        - height (int): Image height.
+        - batch_size (int): Size of batch to render into. Memory is
+            allocated once upon initialization of object.
+        - device_id (int): Device to use for storing the batch.
+        - depth (bool): if True, render depth in addition to RGB.
+        - use_cuda (bool): if True, use OpenGL-CUDA interop to map
+            the PBO onto a CUDA buffer.
+        """
+        # Early initialization to prevent failure in __del__
+        self._use_cuda = False
+        self.pbo_depth, self.pbo_depth = 0, 0
+
+        if not usingEGL():
+            raise MjBatchRendererNotSupported(
+                "MjBatchRenderer currently only supported with EGL-backed"
+                "rendering context.")
+
+        # Make sure OpenGL Context is available before creating PBOs
+        initOpenGL(device_id)
+        makeOpenGLContextCurrent(device_id)
+
+        self.pbo_rgb = createPBO(width, height, batch_size, 0)
+        self.pbo_depth = createPBO(width, height, batch_size, 1) if depth else 0
+
+        self._depth = depth
+        self._device_id = device_id
+        self._width = width
+        self._height = height
+        self._batch_size = batch_size
+        self._current_batch_offset = 0
+
+        self._use_cuda = use_cuda
+        self._cuda_buffers_are_mapped = False
+        self._cuda_rgb_ptr, self._cuda_depth_ptr = None, None
+        if use_cuda:
+            self._init_cuda()
+
+    def _init_cuda(self):
+        if drv is None:
+            raise ImportError("Failed to import pycuda.")
+        # Use local imports so that we don't have to make pycuda
+        # opengl interop a requirement
+        from pycuda.gl import RegisteredBuffer
+
+        drv.init()
+        device = drv.Device(self._device_id)
+        self._cuda_context = device.make_context()
+        self._cuda_context.push()
+
+        self._cuda_rgb_pbo = RegisteredBuffer(self.pbo_rgb)
+        if self._depth:
+            self._cuda_depth_pbo = RegisteredBuffer(self.pbo_depth)
+
+    def map(self):
+        """ Map OpenGL buffer to CUDA for reading. """
+        if not self._use_cuda:
+            raise CudaNotEnabledError()
+        elif self._cuda_buffers_are_mapped:
+            return  # just make it a no-op
+
+        self._cuda_context.push()
+        self._cuda_rgb_mapping = self._cuda_rgb_pbo.map()
+        ptr, self._cuda_rgb_buf_size = (
+            self._cuda_rgb_mapping.device_ptr_and_size())
+        assert ptr is not None and self._cuda_rgb_buf_size > 0
+        if self._cuda_rgb_ptr is None:
+            self._cuda_rgb_ptr = ptr
+
+        # There doesn't seem to be a guarantee from the API that the
+        # pointer will be the same between mappings, but empirically
+        # this has been true. If this isn't true, we need to modify
+        # the interface to MjBatchRenderer to make this clearer to user.
+        # So, hopefully we won't hit this assert.
+        assert self._cuda_rgb_ptr == ptr, (
+            "Mapped CUDA rgb buffer pointer %d doesn't match old pointer %d" %
+            (ptr, self._cuda_rgb_ptr))
+
+        if self._depth:
+            self._cuda_depth_mapping = self._cuda_depth_pbo.map()
+            ptr, self._cuda_depth_buf_size = (
+                self._cuda_depth_mapping.device_ptr_and_size())
+            assert ptr is not None and self._cuda_depth_buf_size > 0
+            if self._cuda_depth_ptr is None:
+                self._cuda_depth_ptr = ptr
+            assert self._cuda_depth_ptr == ptr, (
+                "Mapped CUDA depth buffer pointer %d doesn't match old pointer %d" %
+                (ptr, self._cuda_depth_ptr))
+
+        self._cuda_buffers_are_mapped = True
+
+    def unmap(self):
+        """ Unmap OpenGL buffer from CUDA so that it can be rendered into. """
+        if not self._use_cuda:
+            raise CudaNotEnabledError()
+        elif not self._cuda_buffers_are_mapped:
+            return  # just make it a no-op
+
+        self._cuda_context.push()
+        self._cuda_rgb_mapping.unmap()
+        self._cuda_rgb_mapping = None
+        self._cuda_rgb_ptr = None
+        if self._depth:
+            self._cuda_depth_mapping.unmap()
+            self._cuda_depth_mapping = None
+            self._cuda_depth_ptr = None
+
+        self._cuda_buffers_are_mapped = False
+
+    def prepare_render_context(self, sim):
+        """
+        Set up the rendering context for an MjSim. Also happens automatically
+        on `.render()`.
+        """
+        for c in sim.render_contexts:
+            if (c.offscreen and
+                    isinstance(c.opengl_context, OffscreenOpenGLContext) and
+                    c.opengl_context.device_id == self._device_id):
+                return c
+
+        return MjRenderContext(sim, device_id=self._device_id)
+
+    def render(self, sim, camera_id=None, batch_offset=None):
+        """
+        Render current scene from the MjSim into the buffer. By
+        default the batch offset is automatically incremented with
+        each call. It can be reset with the batch_offset parameter.
+
+        This method doesn't return anything. Use the `.read` method
+        to read the buffer, or access the buffer pointer directly with
+        e.g. `.cuda_rgb_buffer_pointer` accessor.
+
+        Args:
+        - sim (MjSim): The simulator to use for rendering.
+        - camera_id (int): MuJoCo id for the camera, from
+            `sim.model.camera_name2id()`.
+        - batch_offset (int): offset in batch to render to.
+        """
+        if self._use_cuda and self._cuda_buffers_are_mapped:
+            raise CudaBufferMappedError(
+                "CUDA buffers must be unmapped before calling render.")
+
+        if batch_offset is not None:
+            if batch_offset < 0 or batch_offset >= self._batch_size:
+                raise ValueError("batch_offset out of range")
+            self._current_batch_offset = batch_offset
+
+        # Ensure the correct device context is used (this takes ~1 µs)
+        makeOpenGLContextCurrent(self._device_id)
+
+        render_context = self.prepare_render_context(sim)
+        render_context.update_offscreen_size(self._width, self._height)
+        render_context.render(self._width, self._height, camera_id=camera_id)
+
+        cdef mjrRect viewport
+        viewport.left = 0
+        viewport.bottom = 0
+        viewport.width = self._width
+        viewport.height = self._height
+
+        cdef PyMjrContext con = <PyMjrContext> render_context.con
+        copyFBOToPBO(con.ptr, self.pbo_rgb, self.pbo_depth,
+                     viewport, self._current_batch_offset)
+
+        self._current_batch_offset = (self._current_batch_offset + 1) % self._batch_size
+
+    def read(self):
+        """
+        Transfer a copy of the buffer from the GPU to the CPU as a numpy array.
+
+        Returns:
+        - rgb_batch (numpy array): batch of rgb images in uint8 NHWC format
+        - depth_batch (numpy array): batch of depth images in uint16 NHWC format
+        """
+        if self._use_cuda:
+            return self._read_cuda()
+        else:
+            return self._read_nocuda()
+
+    def _read_cuda(self):
+        if not self._cuda_buffers_are_mapped:
+            raise CudaBufferNotMappedError(
+                "CUDA buffers must be mapped before reading")
+
+        rgb_arr = drv.from_device(
+            self._cuda_rgb_ptr,
+            shape=(self._batch_size, self._height, self._width, 3),
+            dtype=np.uint8)
+
+        if self._depth:
+            depth_arr = drv.from_device(
+                self._cuda_depth_ptr,
+                shape=(self._batch_size, self._height, self._width),
+                dtype=np.uint16)
+        else:
+            depth_arr = None
+
+        return rgb_arr, depth_arr
+
+    def _read_nocuda(self):
+        rgb_arr = np.zeros(3 * self._width * self._height * self._batch_size, dtype=np.uint8)
+        cdef unsigned char[::view.contiguous] rgb_view = rgb_arr
+        depth_arr = np.zeros(self._width * self._height * self._batch_size, dtype=np.uint16)
+        cdef unsigned short[::view.contiguous] depth_view = depth_arr
+
+        if self._depth:
+            readPBO(&rgb_view[0], &depth_view[0], self.pbo_rgb, self.pbo_depth,
+                    self._width, self._height, self._batch_size)
+            depth_arr = depth_arr.reshape(self._batch_size, self._height, self._width)
+        else:
+            readPBO(&rgb_view[0], NULL, self.pbo_rgb, 0,
+                    self._width, self._height, self._batch_size)
+            # Fine to throw aray depth_arr above since malloc/free is cheap
+            depth_arr = None
+
+        rgb_arr = rgb_arr.reshape(self._batch_size, self._height, self._width, 3)
+        return rgb_arr, depth_arr
+
+    @property
+    def cuda_rgb_buffer_pointer(self):
+        """ Pointer to CUDA buffer for RGB batch. """
+        if not self._use_cuda:
+            raise CudaNotEnabledError()
+        elif not self._cuda_buffers_are_mapped:
+            raise CudaBufferNotMappedError()
+        return self._cuda_rgb_ptr
+
+    @property
+    def cuda_depth_buffer_pointer(self):
+        """ Pointer to CUDA buffer for depth batch. """
+        if not self._use_cuda:
+            raise CudaNotEnabledError()
+        elif not self._cuda_buffers_are_mapped:
+            raise CudaBufferNotMappedError()
+        if not self._depth:
+            raise RuntimeError("Depth not enabled. Use depth=True on initialization.")
+        return self._cuda_depth_ptr
+
+    def __del__(self):
+        if self._use_cuda:
+            self._cuda_context.push()
+            self.unmap()
+            self._cuda_rgb_pbo.unregister()
+            if self._depth:
+                self._cuda_depth_pbo.unregister()
+
+            # Clean up context
+            drv.Context.pop()
+            self._cuda_context.detach()
+
+        if self.pbo_depth:
+            freePBO(self.pbo_rgb)
+        if self.pbo_depth:
+            freePBO(self.pbo_depth)

mujoco-py-2.1.2.14/mujoco_py/mjrendercontext.pyx

@@ -0,0 +1,329 @@
+from threading import Lock
+from mujoco_py.generated import const
+import numpy as np
+cimport numpy as np
+
+cdef class MjRenderContext(object):
+    """
+    Class that encapsulates rendering functionality for a
+    MuJoCo simulation.
+    """
+
+    cdef mjModel *_model_ptr
+    cdef mjData *_data_ptr
+
+    cdef mjvScene _scn
+    cdef mjvCamera _cam
+    cdef mjvOption _vopt
+    cdef mjvPerturb _pert
+    cdef mjrContext _con
+
+    # Public wrappers
+    cdef readonly PyMjvScene scn
+    cdef readonly PyMjvCamera cam
+    cdef readonly PyMjvOption vopt
+    cdef readonly PyMjvPerturb pert
+    cdef readonly PyMjrContext con
+
+    cdef readonly object opengl_context
+    cdef readonly int _visible
+    cdef readonly list _markers
+    cdef readonly dict _overlay
+
+    cdef readonly bint offscreen
+    cdef public object sim
+
+    def __cinit__(self):
+        maxgeom = 1000
+        mjv_makeScene(self._model_ptr, &self._scn, maxgeom)
+        mjv_defaultCamera(&self._cam)
+        mjv_defaultPerturb(&self._pert)
+        mjv_defaultOption(&self._vopt)
+        mjr_defaultContext(&self._con)
+
+    def __init__(self, MjSim sim, bint offscreen=True, int device_id=-1, opengl_backend=None, quiet=False):
+        self.sim = sim
+        self._setup_opengl_context(offscreen, device_id, opengl_backend, quiet=quiet)
+        self.offscreen = offscreen
+
+        # Ensure the model data has been updated so that there
+        # is something to render
+        sim.forward()
+
+        sim.add_render_context(self)
+
+        self._model_ptr = sim.model.ptr
+        self._data_ptr = sim.data.ptr
+        self.scn = WrapMjvScene(&self._scn)
+        self.cam = WrapMjvCamera(&self._cam)
+        self.vopt = WrapMjvOption(&self._vopt)
+        self.con = WrapMjrContext(&self._con)
+
+        self._pert.active = 0
+        self._pert.select = 0
+        self._pert.skinselect = -1
+
+        self.pert = WrapMjvPerturb(&self._pert)
+
+        self._markers = []
+        self._overlay = {}
+
+        self._init_camera(sim)
+        self._set_mujoco_buffers()
+
+    def update_sim(self, MjSim new_sim):
+        if new_sim == self.sim:
+            return
+        self._model_ptr = new_sim.model.ptr
+        self._data_ptr = new_sim.data.ptr
+        self._set_mujoco_buffers()
+        for render_context in self.sim.render_contexts:
+            new_sim.add_render_context(render_context)
+        self.sim = new_sim
+
+    def _set_mujoco_buffers(self):
+        mjr_makeContext(self._model_ptr, &self._con, mjFONTSCALE_150)
+        if self.offscreen:
+            mjr_setBuffer(mjFB_OFFSCREEN, &self._con);
+            if self._con.currentBuffer != mjFB_OFFSCREEN:
+                raise RuntimeError('Offscreen rendering not supported')
+        else:
+            mjr_setBuffer(mjFB_WINDOW, &self._con);
+            if self._con.currentBuffer != mjFB_WINDOW:
+                raise RuntimeError('Window rendering not supported')
+        self.con = WrapMjrContext(&self._con)
+
+    def _setup_opengl_context(self, offscreen, device_id, opengl_backend, quiet=False):
+        if opengl_backend is None and (not offscreen or sys.platform == 'darwin'):
+            # default to glfw for onscreen viewing or mac (both offscreen/onscreen)
+            opengl_backend = 'glfw'
+
+        if opengl_backend == 'glfw':
+            self.opengl_context = GlfwContext(offscreen=offscreen, quiet=quiet)
+        else:
+            if device_id < 0:
+                if "GPUS" in os.environ:
+                    device_id = os.environ["GPUS"]
+                else:
+                    device_id = os.getenv('CUDA_VISIBLE_DEVICES', '')
+                if len(device_id) > 0:
+                    device_id = int(device_id.split(',')[0])
+                else:
+                    # Sometimes env variable is an empty string.
+                    device_id = 0
+            self.opengl_context = OffscreenOpenGLContext(device_id)
+
+    def _init_camera(self, sim):
+        # Make the free camera look at the scene
+        self.cam.type = const.CAMERA_FREE
+        self.cam.fixedcamid = -1
+        for i in range(3):
+            self.cam.lookat[i] = np.median(sim.data.geom_xpos[:, i])
+        self.cam.distance = sim.model.stat.extent
+
+    def update_offscreen_size(self, width, height):
+        if width != self._con.offWidth or height != self._con.offHeight:
+            self._model_ptr.vis.global_.offwidth = width
+            self._model_ptr.vis.global_.offheight = height
+            mjr_freeContext(&self._con)
+            self._set_mujoco_buffers()
+
+    def render(self, width, height, camera_id=None, segmentation=False):
+        cdef mjrRect rect
+        rect.left = 0
+        rect.bottom = 0
+        rect.width = width
+        rect.height = height
+
+        if self.sim.render_callback is not None:
+            self.sim.render_callback(self.sim, self)
+
+        # Sometimes buffers are too small.
+        if width > self._con.offWidth or height > self._con.offHeight:
+            new_width = max(width, self._model_ptr.vis.global_.offwidth)
+            new_height = max(height, self._model_ptr.vis.global_.offheight)
+            self.update_offscreen_size(new_width, new_height)
+
+        if camera_id is not None:
+            if camera_id == -1:
+                self.cam.type = const.CAMERA_FREE
+            else:
+                self.cam.type = const.CAMERA_FIXED
+            self.cam.fixedcamid = camera_id
+
+        # This doesn't really do anything else rather than checking for the size of buffer
+        # need to investigate further whi is that a no-op
+        # self.opengl_context.set_buffer_size(width, height)
+
+        mjv_updateScene(self._model_ptr, self._data_ptr, &self._vopt,
+                        &self._pert, &self._cam, mjCAT_ALL, &self._scn)
+
+        if segmentation:
+            self._scn.flags[const.RND_SEGMENT] = 1
+            self._scn.flags[const.RND_IDCOLOR] = 1
+
+        for marker_params in self._markers:
+            self._add_marker_to_scene(marker_params)
+        
+        mjr_render(rect, &self._scn, &self._con)
+        for gridpos, (text1, text2) in self._overlay.items():
+            mjr_overlay(const.FONTSCALE_150, gridpos, rect, text1.encode(), text2.encode(), &self._con)
+
+        if segmentation:
+            self._scn.flags[const.RND_SEGMENT] = 0
+            self._scn.flags[const.RND_IDCOLOR] = 0
+
+    def read_pixels(self, width, height, depth=True, segmentation=False):
+        cdef mjrRect rect
+        rect.left = 0
+        rect.bottom = 0
+        rect.width = width
+        rect.height = height
+
+        rgb_arr = np.zeros(3 * rect.width * rect.height, dtype=np.uint8)
+        depth_arr = np.zeros(rect.width * rect.height, dtype=np.float32)
+        
+        cdef unsigned char[::view.contiguous] rgb_view = rgb_arr
+        cdef float[::view.contiguous] depth_view = depth_arr
+        mjr_readPixels(&rgb_view[0], &depth_view[0], rect, &self._con)
+        rgb_img = rgb_arr.reshape(rect.height, rect.width, 3)
+        cdef np.ndarray[np.npy_uint32, ndim=2] seg_img
+        cdef np.ndarray[np.npy_int32, ndim=2] seg_ids
+
+        ret_img = rgb_img
+        if segmentation:
+            seg_img = (rgb_img[:, :, 0] + rgb_img[:, :, 1] * (2**8) + rgb_img[:, :, 2] * (2 ** 16))
+            seg_img[seg_img >= (self._scn.ngeom + 1)] = 0
+            seg_ids = np.full((self._scn.ngeom + 1, 2), fill_value=-1, dtype=np.int32)
+
+            for i in range(self._scn.ngeom):
+                geom = self._scn.geoms[i]
+                if geom.segid != -1:
+                    seg_ids[geom.segid + 1, 0] = geom.objtype
+                    seg_ids[geom.segid + 1, 1] = geom.objid
+            ret_img = seg_ids[seg_img]
+
+        if depth:
+            depth_img = depth_arr.reshape(rect.height, rect.width)
+            return (ret_img, depth_img)
+        else:
+            return ret_img
+
+    def read_pixels_depth(self, np.ndarray[np.float32_t, mode="c", ndim=2] buffer):
+            ''' Read depth pixels into a preallocated buffer '''
+            cdef mjrRect rect
+            rect.left = 0
+            rect.bottom = 0
+            rect.width = buffer.shape[1]
+            rect.height = buffer.shape[0]
+
+            cdef float[::view.contiguous] buffer_view = buffer.ravel()
+            mjr_readPixels(NULL, &buffer_view[0], rect, &self._con)
+
+    def upload_texture(self, int tex_id):
+        """ Uploads given texture to the GPU. """
+        self.opengl_context.make_context_current()
+        mjr_uploadTexture(self._model_ptr, &self._con, tex_id)
+
+    def draw_pixels(self, np.ndarray[np.uint8_t, ndim=3] image, int left, int bottom):
+        """Draw an image into the OpenGL buffer."""
+        cdef unsigned char[::view.contiguous] image_view = image.ravel()
+        cdef mjrRect viewport
+        viewport.left = left
+        viewport.bottom = bottom
+        viewport.width = image.shape[1]
+        viewport.height = image.shape[0]
+        mjr_drawPixels(&image_view[0], NULL, viewport, &self._con)
+
+    def move_camera(self, int action, double reldx, double reldy):
+        """ Moves the camera based on mouse movements. Action is one of mjMOUSE_*. """
+        mjv_moveCamera(self._model_ptr, action, reldx, reldy, &self._scn, &self._cam)
+
+    def add_overlay(self, int gridpos, str text1, str text2):
+        """ Overlays text on the scene. """
+        if gridpos not in self._overlay:
+            self._overlay[gridpos] = ["", ""]
+        self._overlay[gridpos][0] += text1 + "\n"
+        self._overlay[gridpos][1] += text2 + "\n"
+
+    def add_marker(self, **marker_params):
+        self._markers.append(marker_params)
+
+    def _add_marker_to_scene(self, marker_params):
+        """ Adds marker to scene, and returns the corresponding object. """
+        if self._scn.ngeom >= self._scn.maxgeom:
+            raise RuntimeError('Ran out of geoms. maxgeom: %d' % self._scn.maxgeom)
+
+        cdef mjvGeom *g = self._scn.geoms + self._scn.ngeom
+
+        # default values.
+        g.dataid = -1
+        g.objtype = const.OBJ_UNKNOWN
+        g.objid = -1
+        g.category = const.CAT_DECOR
+        g.texid = -1
+        g.texuniform = 0
+        g.texrepeat[0] = 1
+        g.texrepeat[1] = 1
+        g.emission = 0
+        g.specular = 0.5
+        g.shininess = 0.5
+        g.reflectance = 0
+        g.type = const.GEOM_BOX
+        g.size[:] = np.ones(3) * 0.1
+        g.mat[:] = np.eye(3).flatten()
+        g.rgba[:] = np.ones(4)
+        wrapped = WrapMjvGeom(g)
+
+        for key, value in marker_params.items():
+            if isinstance(value, (int, float)):
+                setattr(wrapped, key, value)
+            elif isinstance(value, (tuple, list, np.ndarray)):
+                attr = getattr(wrapped, key)
+                attr[:] = np.asarray(value).reshape(attr.shape)
+            elif isinstance(value, str):
+                assert key == "label", "Only label is a string in mjvGeom."
+                if value == None:
+                    g.label[0] = 0
+                else:
+                    strncpy(g.label, value.encode(), 100)
+            elif hasattr(wrapped, key):
+                raise ValueError("mjvGeom has attr {} but type {} is invalid".format(key, type(value)))
+            else:
+                raise ValueError("mjvGeom doesn't have field %s" % key)
+
+        self._scn.ngeom += 1
+
+
+    def __dealloc__(self):
+        mjr_freeContext(&self._con)
+        mjv_freeScene(&self._scn)
+
+
+class MjRenderContextOffscreen(MjRenderContext):
+
+    def __cinit__(self, MjSim sim, int device_id):
+        super().__init__(sim, offscreen=True, device_id=device_id)
+
+class MjRenderContextWindow(MjRenderContext):
+
+    def __init__(self, MjSim sim):
+        super().__init__(sim, offscreen=False)
+        self.render_swap_callback = None
+
+        assert isinstance(self.opengl_context, GlfwContext), (
+            "Only GlfwContext supported for windowed rendering")
+
+    @property
+    def window(self):
+        return self.opengl_context.window
+
+    def render(self):
+        if self.window is None or glfw.window_should_close(self.window):
+            return
+
+        glfw.make_context_current(self.window)
+        super().render(*glfw.get_framebuffer_size(self.window))
+        if self.render_swap_callback is not None:
+            self.render_swap_callback()
+        glfw.swap_buffers(self.window)

mujoco-py-2.1.2.14/mujoco_py/mjrenderpool.py

@@ -0,0 +1,241 @@
+import ctypes
+import inspect
+
+from multiprocessing import Array, get_start_method, Pool, Value
+
+import numpy as np
+
+
+class RenderPoolStorage:
+    """
+    Helper object used for storing global data for worker processes.
+    """
+
+    __slots__ = ['shared_rgbs_array',
+                 'shared_depths_array',
+                 'device_id',
+                 'sim',
+                 'modder']
+
+
+class MjRenderPool:
+    """
+    Utilizes a process pool to render a MuJoCo simulation across
+    multiple GPU devices. This can scale the throughput linearly
+    with the number of available GPUs. Throughput can also be
+    slightly increased by using more than one worker per GPU.
+    """
+
+    DEFAULT_MAX_IMAGE_SIZE = 512 * 512  # in pixels
+
+    def __init__(self, model, device_ids=1, n_workers=None,
+                 max_batch_size=None, max_image_size=DEFAULT_MAX_IMAGE_SIZE,
+                 modder=None):
+        """
+        Args:
+        - model (PyMjModel): MuJoCo model to use for rendering
+        - device_ids (int/list): list of device ids to use for rendering.
+            One or more workers will be assigned to each device, depending
+            on how many workers are requested.
+        - n_workers (int): number of parallel processes in the pool. Defaults
+            to the number of device ids.
+        - max_batch_size (int): maximum number of states that can be rendered
+            in batch using .render(). Defaults to the number of workers.
+        - max_image_size (int): maximum number pixels in images requested
+            by .render()
+        - modder (Modder): modder to use for domain randomization.
+        """
+        self._closed, self.pool = False, None
+
+        if not (modder is None or inspect.isclass(modder)):
+            raise ValueError("modder must be a class")
+
+        if isinstance(device_ids, int):
+            device_ids = list(range(device_ids))
+        else:
+            assert isinstance(device_ids, list), (
+                "device_ids must be list of integer")
+
+        n_workers = n_workers or 1
+        self._max_batch_size = max_batch_size or (len(device_ids) * n_workers)
+        self._max_image_size = max_image_size
+
+        array_size = self._max_image_size * self._max_batch_size
+
+        self._shared_rgbs = Array(ctypes.c_uint8, array_size * 3)
+        self._shared_depths = Array(ctypes.c_float, array_size)
+
+        self._shared_rgbs_array = np.frombuffer(
+            self._shared_rgbs.get_obj(), dtype=ctypes.c_uint8)
+        assert self._shared_rgbs_array.size == (array_size * 3), (
+            "Array size is %d, expected %d" % (
+                self._shared_rgbs_array.size, array_size * 3))
+        self._shared_depths_array = np.frombuffer(
+            self._shared_depths.get_obj(), dtype=ctypes.c_float)
+        assert self._shared_depths_array.size == array_size, (
+            "Array size is %d, expected %d" % (
+                self._shared_depths_array.size, array_size))
+
+        worker_id = Value(ctypes.c_int)
+        worker_id.value = 0
+
+        if get_start_method() != "spawn":
+            raise RuntimeError(
+                "Start method must be set to 'spawn' for the "
+                "render pool to work. That is, you must add the "
+                "following to the _TOP_ of your main script, "
+                "before any other imports (since they might be "
+                "setting it otherwise):\n"
+                "  import multiprocessing as mp\n"
+                "  if __name__ == '__main__':\n"
+                "    mp.set_start_method('spawn')\n")
+
+        self.pool = Pool(
+            processes=len(device_ids) * n_workers,
+            initializer=MjRenderPool._worker_init,
+            initargs=(
+                model.get_mjb(),
+                worker_id,
+                device_ids,
+                self._shared_rgbs,
+                self._shared_depths,
+                modder))
+
+    @staticmethod
+    def _worker_init(mjb_bytes, worker_id, device_ids,
+                     shared_rgbs, shared_depths, modder):
+        """
+        Initializes the global state for the workers.
+        """
+        s = RenderPoolStorage()
+
+        with worker_id.get_lock():
+            proc_worker_id = worker_id.value
+            worker_id.value += 1
+        s.device_id = device_ids[proc_worker_id % len(device_ids)]
+
+        s.shared_rgbs_array = np.frombuffer(
+            shared_rgbs.get_obj(), dtype=ctypes.c_uint8)
+        s.shared_depths_array = np.frombuffer(
+            shared_depths.get_obj(), dtype=ctypes.c_float)
+
+        # avoid a circular import
+        from mujoco_py import load_model_from_mjb, MjRenderContext, MjSim
+        s.sim = MjSim(load_model_from_mjb(mjb_bytes))
+        # attach a render context to the sim (needs to happen before
+        # modder is called, since it might need to upload textures
+        # to the GPU).
+        MjRenderContext(s.sim, device_id=s.device_id)
+
+        if modder is not None:
+            s.modder = modder(s.sim, random_state=proc_worker_id)
+            s.modder.whiten_materials()
+        else:
+            s.modder = None
+
+        global _render_pool_storage
+        _render_pool_storage = s
+
+    @staticmethod
+    def _worker_render(worker_id, state, width, height,
+                       camera_name, randomize):
+        """
+        Main target function for the workers.
+        """
+        s = _render_pool_storage
+
+        forward = False
+        if state is not None:
+            s.sim.set_state(state)
+            forward = True
+        if randomize and s.modder is not None:
+            s.modder.randomize()
+            forward = True
+        if forward:
+            s.sim.forward()
+
+        rgb_block = width * height * 3
+        rgb_offset = rgb_block * worker_id
+        rgb = s.shared_rgbs_array[rgb_offset:rgb_offset + rgb_block]
+        rgb = rgb.reshape(height, width, 3)
+
+        depth_block = width * height
+        depth_offset = depth_block * worker_id
+        depth = s.shared_depths_array[depth_offset:depth_offset + depth_block]
+        depth = depth.reshape(height, width)
+
+        rgb[:], depth[:] = s.sim.render(
+            width, height, camera_name=camera_name, depth=True,
+            device_id=s.device_id)
+
+    def render(self, width, height, states=None, camera_name=None,
+               depth=False, randomize=False, copy=True):
+        """
+        Renders the simulations in batch. If no states are provided,
+        the max_batch_size will be used.
+
+        Args:
+        - width (int): width of image to render.
+        - height (int): height of image to render.
+        - states (list): list of MjSimStates; updates the states before
+            rendering. Batch size will be number of states supplied.
+        - camera_name (str): name of camera to render from.
+        - depth (bool): if True, also return depth.
+        - randomize (bool): calls modder.rand_all() before rendering.
+        - copy (bool): return a copy rather than a reference
+
+        Returns:
+        - rgbs: NxHxWx3 numpy array of N images in batch of width W
+            and height H.
+        - depth: NxHxW numpy array of N images in batch of width W
+            and height H. Only returned if depth=True.
+        """
+        if self._closed:
+            raise RuntimeError("The pool has been closed.")
+
+        if (width * height) > self._max_image_size:
+            raise ValueError(
+                "Requested image larger than maximum image size. Create "
+                "a new RenderPool with a larger maximum image size.")
+        if states is None:
+            batch_size = self._max_batch_size
+            states = [None] * batch_size
+        else:
+            batch_size = len(states)
+
+        if batch_size > self._max_batch_size:
+            raise ValueError(
+                "Requested batch size larger than max batch size. Create "
+                "a new RenderPool with a larger max batch size.")
+
+        self.pool.starmap(
+            MjRenderPool._worker_render,
+            [(i, state, width, height, camera_name, randomize)
+             for i, state in enumerate(states)])
+
+        rgbs = self._shared_rgbs_array[:width * height * 3 * batch_size]
+        rgbs = rgbs.reshape(batch_size, height, width, 3)
+        if copy:
+            rgbs = rgbs.copy()
+
+        if depth:
+            depths = self._shared_depths_array[:width * height * batch_size]
+            depths = depths.reshape(batch_size, height, width).copy()
+            if copy:
+                depths = depths.copy()
+            return rgbs, depths
+        else:
+            return rgbs
+
+    def close(self):
+        """
+        Closes the pool and terminates child processes.
+        """
+        if not self._closed:
+            if self.pool is not None:
+                self.pool.close()
+                self.pool.join()
+            self._closed = True
+
+    def __del__(self):
+        self.close()

mujoco-py-2.1.2.14/mujoco_py/mjsim.pyx

@@ -0,0 +1,439 @@
+from xml.dom import minidom
+from mujoco_py.utils import remove_empty_lines
+from mujoco_py.builder import build_callback_fn
+from threading import Lock
+
+_MjSim_render_lock = Lock()
+
+ctypedef void (*substep_udd_t)(const mjModel* m, mjData* d)
+
+
+cdef class MjSim(object):
+    """MjSim represents a running simulation including its state.
+
+    Similar to Gym's ``MujocoEnv``, it internally wraps a :class:`.PyMjModel`
+    and a :class:`.PyMjData`.
+
+    Parameters
+    ----------
+    model : :class:`.PyMjModel`
+        The model to simulate.
+    data : :class:`.PyMjData`
+        Optional container for the simulation state. Will be created if ``None``.
+    nsubsteps : int
+        Optional number of MuJoCo steps to run for every call to :meth:`.step`.
+        Buffers will be swapped only once per step.
+    udd_callback : fn(:class:`.MjSim`) -> dict
+        Optional callback for user-defined dynamics. At every call to
+        :meth:`.step`, it receives an MjSim object ``sim`` containing the
+        current user-defined dynamics state in ``sim.udd_state``, and returns the
+        next ``udd_state`` after applying the user-defined dynamics. This is
+        useful e.g. for reward functions that operate over functions of historical
+        state.
+    substep_callback : str or int or None
+        This uses a compiled C function as user-defined dynamics in substeps.
+        If given as a string, it's compiled as a C function and set as pointer.
+        If given as int, it's interpreted as a function pointer.
+        See :meth:`.set_substep_callback` for detailed info.
+    userdata_names : list of strings or None
+        This is a convenience parameter which is just set on the model.
+        Equivalent to calling ``model.set_userdata_names``
+    render_callback : callback for rendering.
+    """
+    # MjRenderContext for rendering camera views.
+    cdef readonly list render_contexts
+    cdef readonly object _render_context_window
+    cdef readonly object _render_context_offscreen
+
+    # MuJoCo model
+    cdef readonly PyMjModel model
+    # MuJoCo data
+    """
+    DATAZ
+    """
+    cdef readonly PyMjData data
+    # Number of substeps when calling .step
+    cdef public int nsubsteps
+    # User defined state.
+    cdef public dict udd_state
+    # User defined dynamics callback
+    cdef readonly object _udd_callback
+    # Allows to store extra information in MjSim.
+    cdef readonly dict extras
+    # Function pointer for substep callback, stored as uintptr
+    cdef readonly uintptr_t substep_callback_ptr
+    # Callback executed before rendering.
+    cdef public object render_callback
+
+    def __cinit__(self, PyMjModel model, PyMjData data=None, int nsubsteps=1,
+                  udd_callback=None, substep_callback=None, userdata_names=None,
+                  render_callback=None):
+        self.nsubsteps = nsubsteps
+        self.model = model
+        if data is None:
+            with wrap_mujoco_warning():
+                _data = mj_makeData(self.model.ptr)
+            if _data == NULL:
+                raise Exception('mj_makeData failed!')
+            self.data = WrapMjData(_data, self.model)
+        else:
+            self.data = data
+
+        self.render_contexts = []
+        self._render_context_offscreen = None
+        self._render_context_window = None
+        self.udd_state = None
+        self.udd_callback = udd_callback
+        self.render_callback = render_callback
+        self.extras = {}
+        self.set_substep_callback(substep_callback, userdata_names)
+
+    def reset(self):
+        """
+        Resets the simulation data and clears buffers.
+        """
+        with wrap_mujoco_warning():
+            mj_resetData(self.model.ptr, self.data.ptr)
+
+        self.udd_state = None
+        self.step_udd()
+
+    def forward(self):
+        """
+        Computes the forward kinematics. Calls ``mj_forward`` internally.
+        """
+        with wrap_mujoco_warning():
+            mj_forward(self.model.ptr, self.data.ptr)
+
+    def set_constants(self):
+        """
+        Set constant fields of mjModel, corresponding to qpos0 configuration.
+        """
+        with wrap_mujoco_warning():
+            mj_setConst(self.model.ptr, self.data.ptr)
+
+    def step(self, with_udd=True):
+        """
+        Advances the simulation by calling ``mj_step``.
+
+        If ``qpos`` or ``qvel`` have been modified directly, the user is required to call
+        :meth:`.forward` before :meth:`.step` if their ``udd_callback`` requires access to MuJoCo state
+        set during the forward dynamics.
+        """
+        if with_udd:
+            self.step_udd()
+
+        with wrap_mujoco_warning():
+            for _ in range(self.nsubsteps):
+                self.substep_callback()
+                mj_step(self.model.ptr, self.data.ptr)
+
+    def render(self, width=None, height=None, *, camera_name=None, depth=False,
+               mode='offscreen', device_id=-1, segmentation=False):
+        """
+        Renders view from a camera and returns image as an `numpy.ndarray`.
+
+        Args:
+        - width (int): desired image width.
+        - height (int): desired image height.
+        - camera_name (str): name of camera in model. If None, the free
+            camera will be used.
+        - depth (bool): if True, also return depth buffer
+        - device (int): device to use for rendering (only for GPU-backed
+            rendering).
+
+        Returns:
+        - rgb (uint8 array): image buffer from camera
+        - depth (float array): depth buffer from camera (only returned
+            if depth=True)
+        """
+        if camera_name is None:
+            camera_id = None
+        else:
+            camera_id = self.model.camera_name2id(camera_name)
+
+        if mode == 'offscreen':
+            with _MjSim_render_lock:
+                if self._render_context_offscreen is None:
+                    render_context = MjRenderContextOffscreen(
+                        self, device_id=device_id)
+                else:
+                    render_context = self._render_context_offscreen
+
+                render_context.render(
+                    width=width, height=height, camera_id=camera_id, segmentation=segmentation)
+                return render_context.read_pixels(
+                    width, height, depth=depth, segmentation=segmentation)
+        elif mode == 'window':
+            if self._render_context_window is None:
+                from mujoco_py.mjviewer import MjViewer
+                render_context = MjViewer(self)
+            else:
+                render_context = self._render_context_window
+
+            render_context.render()
+
+        else:
+            raise ValueError("Mode must be either 'window' or 'offscreen'.")
+
+    def add_render_context(self, render_context):
+        self.render_contexts.append(render_context)
+        if render_context.offscreen and self._render_context_offscreen is None:
+            self._render_context_offscreen = render_context
+        elif not render_context.offscreen and self._render_context_window is None:
+            self._render_context_window = render_context
+
+    @property
+    def udd_callback(self):
+        return self._udd_callback
+
+    @udd_callback.setter
+    def udd_callback(self, value):
+        self._udd_callback = value
+        self.udd_state = None
+        self.step_udd()
+
+    cpdef substep_callback(self):
+        if self.substep_callback_ptr:
+            (<mjfGeneric>self.substep_callback_ptr)(self.model.ptr, self.data.ptr)
+
+    def set_substep_callback(self, substep_callback, userdata_names=None):
+        '''
+        Set a substep callback function.
+
+        Parameters :
+            substep_callback : str or int or None
+                If `substep_callback` is a string, compile to function pointer and set.
+                    See `builder.build_callback_fn()` for documentation.
+                If `substep_callback` is an int, we interpret it as a function pointer.
+                If `substep_callback` is None, we disable substep_callbacks.
+            userdata_names : list of strings or None
+                This is a convenience parameter, if not None, this is passed
+                onto ``model.set_userdata_names()``.
+        '''
+        if userdata_names is not None:
+            self.model.set_userdata_names(userdata_names)
+        if substep_callback is None:
+            self.substep_callback_ptr = 0
+        elif isinstance(substep_callback, int):
+            self.substep_callback_ptr = substep_callback
+        elif isinstance(substep_callback, str):
+            self.substep_callback_ptr = build_callback_fn(substep_callback,
+                                                          self.model.userdata_names)
+        else:
+            raise TypeError('invalid: {}'.format(type(substep_callback)))
+
+    def step_udd(self):
+        if self._udd_callback is None:
+            self.udd_state = {}
+        else:
+            schema_example = self.udd_state
+            self.udd_state = self._udd_callback(self)
+            # Check to make sure the udd_state has consistent keys and dimension across steps
+            if schema_example is not None:
+                keys = set(schema_example.keys()) | set(self.udd_state.keys())
+                for key in keys:
+                    assert key in schema_example, "Keys cannot be added to udd_state between steps."
+                    assert key in self.udd_state, "Keys cannot be dropped from udd_state between steps."
+                    if isinstance(schema_example[key], Number):
+                        assert isinstance(self.udd_state[key], Number), \
+                            "Every value in udd_state must be either a number or a numpy array"
+                    else:
+                        assert isinstance(self.udd_state[key], np.ndarray), \
+                            "Every value in udd_state must be either a number or a numpy array"
+                        assert self.udd_state[key].shape == schema_example[key].shape, \
+                            "Numpy array values in udd_state must keep the same dimension across steps."
+
+    def get_state(self):
+        """ Returns a copy of the simulator state. """
+        qpos = np.copy(self.data.qpos)
+        qvel = np.copy(self.data.qvel)
+        if self.model.na == 0:
+            act = None
+        else:
+            act = np.copy(self.data.act)
+        udd_state = copy.deepcopy(self.udd_state)
+
+        return MjSimState(self.data.time, qpos, qvel, act, udd_state)
+
+    def set_state(self, value):
+        """
+        Sets the state from an MjSimState.
+        If the MjSimState was previously unflattened from a numpy array, consider
+        set_state_from_flattened, as the defensive copy is a substantial overhead
+        in an inner loop.
+
+        Args:
+        - value (MjSimState): the desired state.
+        - call_forward: optionally call sim.forward(). Called by default if
+            the udd_callback is set.
+        """
+        self.data.time = value.time
+        self.data.qpos[:] = np.copy(value.qpos)
+        self.data.qvel[:] = np.copy(value.qvel)
+        if self.model.na != 0:
+            self.data.act[:] = np.copy(value.act)
+        self.udd_state = copy.deepcopy(value.udd_state)
+
+    def set_state_from_flattened(self, value):
+        """ This helper method sets the state from an array without requiring a defensive copy."""
+        state = MjSimState.from_flattened(value, self)
+
+        self.data.time = state.time
+        self.data.qpos[:] = state.qpos
+        self.data.qvel[:] = state.qvel
+        if self.model.na != 0:
+            self.data.act[:] = state.act
+        self.udd_state = state.udd_state
+
+    def save(self, file, format='xml', keep_inertials=False):
+        """
+        Saves the simulator model and state to a file as either
+        a MuJoCo XML or MJB file. The current state is saved as
+        a keyframe in the model file. This is useful for debugging
+        using MuJoCo's `simulate` utility.
+
+        Note that this doesn't save the UDD-state which is
+        part of MjSimState, since that's not supported natively
+        by MuJoCo. If you want to save the model together with
+        the UDD-state, you should use the `get_xml` or `get_mjb`
+        methods on `MjModel` together with `MjSim.get_state` and
+        save them with e.g. pickle.
+
+        Args:
+        - file (IO stream): stream to write model to.
+        - format: format to use (either 'xml' or 'mjb')
+        - keep_inertials (bool): if False, removes all <inertial>
+          properties derived automatically for geoms by MuJoco. Note
+          that this removes ones that were provided by the user
+          as well.
+        """
+        xml_str = self.model.get_xml()
+        dom = minidom.parseString(xml_str)
+
+        mujoco_node = dom.childNodes[0]
+        assert mujoco_node.tagName == 'mujoco'
+
+        keyframe_el = dom.createElement('keyframe')
+        key_el = dom.createElement('key')
+        keyframe_el.appendChild(key_el)
+        mujoco_node.appendChild(keyframe_el)
+
+        def str_array(arr):
+            return " ".join(map(str, arr))
+
+        key_el.setAttribute('time', str(self.data.time))
+        key_el.setAttribute('qpos', str_array(self.data.qpos))
+        key_el.setAttribute('qvel', str_array(self.data.qvel))
+        if self.data.act is not None:
+            key_el.setAttribute('act', str_array(self.data.act))
+
+        if not keep_inertials:
+            for element in dom.getElementsByTagName('inertial'):
+                element.parentNode.removeChild(element)
+
+        result_xml = remove_empty_lines(dom.toprettyxml(indent=" " * 4))
+
+        if format == 'xml':
+            file.write(result_xml)
+        elif format == 'mjb':
+            new_model = load_model_from_xml(result_xml)
+            file.write(new_model.get_mjb())
+        else:
+            raise ValueError("Unsupported format. Valid ones are 'xml' and 'mjb'")
+
+    def ray(self, pnt, vec, include_static_geoms=True, exclude_body=-1, group_filter=None):
+        """
+        Cast a ray into the scene, and return the first valid geom it intersects.
+            pnt - origin point of the ray in world coordinates (X Y Z)
+            vec - direction of the ray in world coordinates (X Y Z)
+            include_static_geoms - if False, we exclude geoms that are children of worldbody.
+            exclude_body - if this is a body ID, we exclude all children geoms of this body.
+            group_filter - a vector of booleans of length const.NGROUP
+                           which specifies what geom groups (stored in model.geom_group)
+                           to enable or disable.  If none, all groups are used
+        Returns (distance, geomid) where
+            distance - distance along ray until first collision with geom
+            geomid - id of the geom the ray collided with
+        If no collision was found in the scene, return (-1, None)
+
+        NOTE: sometimes self.forward() needs to be called before self.ray().
+
+        See self.ray_fast_group() and self.ray_fast_nogroup() for versions of this call
+        with more stringent type requirements.
+        """
+        cdef mjtNum distance
+        cdef mjtNum[::view.contiguous] pnt_view = pnt
+        cdef mjtNum[::view.contiguous] vec_view = vec
+
+        if group_filter is None:
+            return self.ray_fast_nogroup(
+                np.asarray(pnt, dtype=np.float64),
+                np.asarray(vec, dtype=np.float64),
+                1 if include_static_geoms else 0,
+                exclude_body)
+        else:
+            return self.ray_fast_group(
+                np.asarray(pnt, dtype=np.float64),
+                np.asarray(vec, dtype=np.float64),
+                np.asarray(group_filter, dtype=np.uint8),
+                1 if include_static_geoms else 0,
+                exclude_body)
+
+    def ray_fast_group(self,
+            np.ndarray[np.float64_t, mode="c", ndim=1] pnt,
+            np.ndarray[np.float64_t, mode="c", ndim=1] vec,
+            np.ndarray[np.uint8_t, mode="c", ndim=1] geomgroup,
+            mjtByte flg_static=1,
+            int bodyexclude=-1):
+        """
+        Faster version of sim.ray(), which avoids extra copies,
+        but needs to be given all the correct type arrays.
+
+        See self.ray() for explanation of arguments
+        """
+        cdef int geomid
+        cdef mjtNum distance
+        cdef mjtNum[::view.contiguous] pnt_view = pnt
+        cdef mjtNum[::view.contiguous] vec_view = vec
+        cdef mjtByte[::view.contiguous] geomgroup_view = geomgroup
+
+        distance = mj_ray(self.model.ptr,
+                          self.data.ptr,
+                          &pnt_view[0],
+                          &vec_view[0],
+                          &geomgroup_view[0],
+                          flg_static,
+                          bodyexclude,
+                          &geomid)
+        return (distance, geomid)
+
+
+    def ray_fast_nogroup(self,
+            np.ndarray[np.float64_t, mode="c", ndim=1] pnt,
+            np.ndarray[np.float64_t, mode="c", ndim=1] vec,
+            mjtByte flg_static=1,
+            int bodyexclude=-1):
+        """
+        Faster version of sim.ray(), which avoids extra copies,
+        but needs to be given all the correct type arrays.
+
+        This version hardcodes the geomgroup to NULL.
+        (Can't easily express a signature that is "numpy array of specific type or None")
+
+        See self.ray() for explanation of arguments
+        """
+        cdef int geomid
+        cdef mjtNum distance
+        cdef mjtNum[::view.contiguous] pnt_view = pnt
+        cdef mjtNum[::view.contiguous] vec_view = vec
+
+        distance = mj_ray(self.model.ptr,
+                          self.data.ptr,
+                          &pnt_view[0],
+                          &vec_view[0],
+                          NULL,
+                          flg_static,
+                          bodyexclude,
+                          &geomid)
+        return (distance, geomid)

mujoco-py-2.1.2.14/mujoco_py/pxd/mjdata.pxd

@@ -0,0 +1,312 @@
+cdef extern from "mjdata.h" nogil:
+
+    #---------------------------- primitive types (mjt) ------------------------------------
+
+    ctypedef enum mjtWarning:            # warning types
+        mjWARN_INERTIA      = 0,        # (near) singular inertia matrix
+        mjWARN_CONTACTFULL,             # too many contacts in contact list
+        mjWARN_CNSTRFULL,               # too many constraints
+        mjWARN_VGEOMFULL,               # too many visual geoms
+        mjWARN_BADQPOS,                 # bad number in qpos
+        mjWARN_BADQVEL,                 # bad number in qvel
+        mjWARN_BADQACC,                 # bad number in qacc
+        mjWARN_BADCTRL,                 # bad number in ctrl
+
+    enum: mjNWARNING                      # number of warnings
+
+
+    ctypedef enum mjtTimer:
+        # main api
+        mjTIMER_STEP        = 0,        # step
+        mjTIMER_FORWARD,                # forward
+        mjTIMER_INVERSE,                # inverse
+
+        # breakdown of step/forward
+        mjTIMER_POSITION,               # fwdPosition
+        mjTIMER_VELOCITY,               # fwdVelocity
+        mjTIMER_ACTUATION,              # fwdActuation
+        mjTIMER_ACCELERATION,           # fwdAcceleration
+        mjTIMER_CONSTRAINT,             # fwdConstraint
+
+        # breakdown of fwdPosition
+        mjTIMER_POS_KINEMATICS,         # kinematics, com, tendon, transmission
+        mjTIMER_POS_INERTIA,            # inertia computations
+        mjTIMER_POS_COLLISION,          # collision detection
+        mjTIMER_POS_MAKE,               # make constraints
+        mjTIMER_POS_PROJECT,            # project constraints
+
+    enum: mjNTIMER                        # number of timers
+
+
+    #------------------------------ mjContact ----------------------------------------------
+
+    ctypedef struct mjContact:                   # result of collision detection functions
+        # contact parameters set by geom-specific collision detector
+        mjtNum dist                     # distance between nearest points; neg: penetration
+        mjtNum pos[3]                   # position of contact point: midpoint between geoms
+        mjtNum frame[9]                 # normal is in [0-2]
+
+        # contact parameters set by mj_collideGeoms
+        mjtNum includemargin            # include if dist<includemargin=margin-gap
+        mjtNum friction[5]              # tangent1, 2, spin, roll1, 2
+        mjtNum solref[mjNREF]           # constraint solver reference
+        mjtNum solimp[mjNIMP]           # constraint solver impedance
+
+        # storage used internally by constraint solver
+        mjtNum mu                       # friction of regularized cone
+        mjtNum H[36]                    # cone Hessian, set by mj_updateConstraint
+
+        # contact descriptors set by mj_collideGeoms
+        int dim                         # contact space dimensionality: 1, 3, 4 or 6
+        int geom1                       # id of geom 1
+        int geom2                       # id of geom 2
+
+        # flag set by mj_fuseContact or mj_instantianteEquality
+        int exclude                     # 0: include, 1: in gap, 2: fused, 3: equality
+
+        # address computed by mj_instantiateContact
+        int efc_address                 # address in efc; -1: not included, -2-i: distance constraint i ???
+
+    #------------------------------ diagnostics --------------------------------------------
+
+    ctypedef struct mjWarningStat:      # warning statistics
+        int lastinfo                    # info from last warning
+        int number                      # how many times was warning raised
+
+
+    ctypedef struct mjTimerStat:        # timer statistics
+        mjtNum duration                 # cumulative duration
+        int number                      # how many times was timer called
+
+
+    ctypedef struct mjSolverStat:       # per-iteration solver statistics
+        mjtNum improvement              # cost reduction, scaled by 1/trace(M(qpos0))
+        mjtNum gradient                 # gradient norm (primal only, scaled)
+        mjtNum lineslope                # slope in linesearch
+        int nactive                     # number of active constraints
+        int nchange                     # number of constraint state changes
+        int neval                       # number of cost evaluations in line search
+        int nupdate                     # number of Cholesky updates in line search
+
+
+
+    #---------------------------------- mjData ---------------------------------------------
+    ctypedef struct mjData:
+        # constant sizes
+        int nstack                      # number of mjtNums that can fit in stack
+        int nbuffer                     # size of main buffer in bytes
+
+        # stack pointer
+        int pstack                      # first available mjtNum address in stack
+
+        # memory utilization stats
+        int maxuse_stack                # maximum stack allocation
+        int maxuse_con                  # maximum number of contacts
+        int maxuse_efc                  # maximum number of scalar constraints
+
+        # diagnostics
+        mjWarningStat warning[mjNWARNING] # warning statistics
+        mjTimerStat timer[mjNTIMER]       # timer statistics
+        mjSolverStat solver[mjNSOLVER]    # solver statistics per iteration
+        int solver_iter                   # number of solver iterations
+        int solver_nnz                  # number of non-zeros in Hessian or efc_AR
+        mjtNum solver_fwdinv[2]         # forward-inverse comparison: qfrc, efc
+
+        # variable sizes
+        int ne                          # number of equality constraints
+        int nf                          # number of friction constraints
+        int nefc                        # number of constraints
+        int ncon                        # number of detected contacts
+
+        # global properties
+        mjtNum time                     # simulation time
+        mjtNum energy[2]                # potential, kinetic energy
+
+        #-------------------------------- end of info header
+
+        # buffers
+        void*     buffer                # main buffer; all pointers point in it    (nbuffer bytes)
+        mjtNum*   stack                 # stack buffer                             (nstack mjtNums)
+
+        #-------------------------------- main inputs and outputs of the computation
+
+        # state
+        mjtNum*   qpos                  # position                                 (nq x 1)
+        mjtNum*   qvel                  # velocity                                 (nv x 1)
+        mjtNum*   act                   # actuator activation                      (na x 1)
+        mjtNum*   qacc_warmstart        # acceleration used for warmstart          (nv x 1)
+
+        # control
+        mjtNum*   ctrl                  # control                                  (nu x 1)
+        mjtNum*   qfrc_applied          # applied generalized force                (nv x 1)
+        mjtNum*   xfrc_applied          # applied Cartesian force/torque           (nbody x 6)
+
+        # dynamics
+        mjtNum*   qacc                  # acceleration                             (nv x 1)
+        mjtNum*   act_dot               # time-derivative of actuator activation   (na x 1)
+
+        # mocap data
+        mjtNum*  mocap_pos              # positions of mocap bodies                (nmocap x 3)
+        mjtNum*  mocap_quat             # orientations of mocap bodies             (nmocap x 4)
+
+        # user data
+        mjtNum*  userdata               # user data, not touched by engine         (nuserdata x 1)
+
+        # sensors
+        mjtNum*  sensordata             # sensor data array                        (nsensordata x 1)
+
+        #-------------------------------- POSITION dependent
+
+        # computed by mj_fwdPosition/mj_kinematics
+        mjtNum*   xpos                  # Cartesian position of body frame         (nbody x 3)
+        mjtNum*   xquat                 # Cartesian orientation of body frame      (nbody x 4)
+        mjtNum*   xmat                  # Cartesian orientation of body frame      (nbody x 9)
+        mjtNum*   xipos                 # Cartesian position of body com           (nbody x 3)
+        mjtNum*   ximat                 # Cartesian orientation of body inertia    (nbody x 9)
+        mjtNum*   xanchor               # Cartesian position of joint anchor       (njnt x 3)
+        mjtNum*   xaxis                 # Cartesian joint axis                     (njnt x 3)
+        mjtNum*   geom_xpos             # Cartesian geom position                  (ngeom x 3)
+        mjtNum*   geom_xmat             # Cartesian geom orientation               (ngeom x 9)
+        mjtNum*   site_xpos             # Cartesian site position                  (nsite x 3)
+        mjtNum*   site_xmat             # Cartesian site orientation               (nsite x 9)
+        mjtNum*   cam_xpos              # Cartesian camera position                (ncam x 3)
+        mjtNum*   cam_xmat              # Cartesian camera orientation             (ncam x 9)
+        mjtNum*   light_xpos            # Cartesian light position                 (nlight x 3)
+        mjtNum*   light_xdir            # Cartesian light direction                (nlight x 3)
+
+        # computed by mj_fwdPosition/mj_comPos
+        mjtNum*   subtree_com           # center of mass of each subtree           (nbody x 3)
+        mjtNum*   cdof                  # com-based motion axis of each dof        (nv x 6)
+        mjtNum*   cinert                # com-based body inertia and mass          (nbody x 10)
+
+        # computed by mj_fwdPosition/mj_tendon
+        int*      ten_wrapadr           # start address of tendon's path           (ntendon x 1)
+        int*      ten_wrapnum           # number of wrap points in path            (ntendon x 1)
+        int*      ten_J_rownnz          # number of non-zeros in Jacobian row      (ntendon x 1)
+        int*      ten_J_rowadr          # row start address in colind array        (ntendon x 1)
+        int*      ten_J_colind          # column indices in sparse Jacobian        (ntendon x nv)
+        mjtNum*   ten_length            # tendon lengths                           (ntendon x 1)
+        mjtNum*   ten_J                 # tendon Jacobian                          (ntendon x nv)
+        int*      wrap_obj              # geom id; -1: site; -2: pulley            (nwrap*2 x 1)
+        mjtNum*   wrap_xpos             # Cartesian 3D points in all path          (nwrap*2 x 3)
+
+        # computed by mj_fwdPosition/mj_transmission
+        mjtNum*   actuator_length       # actuator lengths                         (nu x 1)
+        mjtNum*   actuator_moment       # actuator moment arms                     (nu x nv)
+
+        # computed by mj_fwdPosition/mj_crb
+        mjtNum*   crb                   # com-based composite inertia and mass     (nbody x 10)
+        mjtNum*   qM                    # total inertia                            (nM x 1)
+
+        # computed by mj_fwdPosition/mj_factorM
+        mjtNum*   qLD                   # L'*D*L factorization of M                (nM x 1)
+        mjtNum*   qLDiagInv             # 1/diag(D)                                (nv x 1)
+        mjtNum*   qLDiagSqrtInv         # 1/sqrt(diag(D))                          (nv x 1)
+
+        # computed by mj_fwdPosition/mj_collision
+        mjContact* contact              # list of all detected contacts            (nconmax x 1)
+
+        # computed by mj_fwdPosition/mj_makeConstraint
+        int*      efc_type              # constraint type (mjtConstraint)          (njmax x 1)
+        int*      efc_id                # id of object of specified type           (njmax x 1)
+        int*      efc_J_rownnz          # number of non-zeros in Jacobian row      (njmax x 1)
+        int*      efc_J_rowadr          # row start address in colind array        (njmax x 1)
+        int*      efc_J_rowsuper        # number of subsequent rows in supernode   (njmax x 1)
+        int*      efc_J_colind          # column indices in sparse Jacobian        (njmax x nv)
+        int*      efc_JT_rownnz         # number of non-zeros in Jacobian row  T   (nv x 1)
+        int*      efc_JT_rowadr         # row start address in colind array    T   (nv x 1)
+        int*      efc_JT_rowsuper       # number of subsequent rows in supernode T (nv x 1)
+        int*      efc_JT_colind         # column indices in sparse Jacobian    T   (nv x njmax)
+        mjtNum*   efc_solref            # constraint solver reference              (njmax x mjNREF)
+        mjtNum*   efc_solimp            # constraint solver impedance              (njmax x mjNIMP)
+        mjtNum*   efc_J                 # constraint Jacobian                      (njmax x nv)
+        mjtNum*   efc_JT                # sparse constraint Jacobian transposed    (nv x njmax)
+        mjtNum*   efc_pos               # constraint position (equality, contact)  (njmax x 1)
+        mjtNum*   efc_margin            # inclusion margin (contact)               (njmax x 1)
+        mjtNum*   efc_frictionloss      # frictionloss (friction)                  (njmax x 1)
+        mjtNum*   efc_diagApprox        # approximation to diagonal of A           (njmax x 1)
+        mjtNum*   efc_KBIP              # stiffness, damping, impedance, imp'      (njmax x 4)
+        mjtNum*   efc_D                 # constraint mass                          (njmax x 1)
+        mjtNum*   efc_R                 # inverse constraint mass                  (njmax x 1)
+
+        # computed by mj_fwdPosition/mj_projectConstraint
+        int*      efc_AR_rownnz         # number of non-zeros in AR                (njmax x 1)
+        int*      efc_AR_rowadr         # row start address in colind array        (njmax x 1)
+        int*      efc_AR_colind         # column indices in sparse AR              (njmax x njmax)
+        mjtNum*   efc_AR                # J*inv(M)*J' + R                          (njmax x njmax)
+
+        #-------------------------------- POSITION, VELOCITY dependent
+
+        # computed by mj_fwdVelocity
+        mjtNum*   ten_velocity          # tendon velocities                        (ntendon x 1)
+        mjtNum*   actuator_velocity     # actuator velocities                      (nu x 1)
+
+        # computed by mj_fwdVelocity/mj_comVel
+        mjtNum*   cvel                  # com-based velocity [3D rot; 3D tran]     (nbody x 6)
+        mjtNum*   cdof_dot              # time-derivative of cdof                  (nv x 6)
+
+        # computed by mj_fwdVelocity/mj_rne (without acceleration)
+        mjtNum*   qfrc_bias             # C(qpos,qvel)                             (nv x 1)
+
+        # computed by mj_fwdVelocity/mj_passive
+        mjtNum*   qfrc_passive          # passive force                            (nv x 1)
+
+        # computed by mj_fwdVelocity/mj_referenceConstraint
+        mjtNum*   efc_vel               # velocity in constraint space: J*qvel     (njmax x 1)
+        mjtNum*   efc_aref              # reference pseudo-acceleration            (njmax x 1)
+
+        # computed by mj_sensorVel
+        mjtNum*   subtree_linvel        # linear velocity of subtree com           (nbody x 3)
+        mjtNum*   subtree_angmom        # angular momentum about subtree com       (nbody x 3)
+
+        #-------------------------------- POSITION, VELOCITY, CONTROL/ACCELERATION dependent
+
+        # computed by mj_fwdActuation
+        mjtNum*   actuator_force        # actuator force in actuation space        (nu x 1)
+        mjtNum*   qfrc_actuator         # actuator force                           (nv x 1)
+
+        # computed by mj_fwdAcceleration
+        mjtNum*   qfrc_unc              # net unconstrained force                  (nv x 1)
+        mjtNum*   qacc_unc              # unconstrained acceleration               (nv x 1)
+
+        # computed by mj_fwdConstraint/mj_inverse
+        mjtNum*   efc_b                 # linear cost term: J*qacc_unc - aref      (njmax x 1)
+        mjtNum*   efc_force             # constraint force in constraint space     (njmax x 1)
+        int*      efc_state             # constraint state (mjtConstraintState)    (njmax x 1)
+        mjtNum*   qfrc_constraint       # constraint force                         (nv x 1)
+
+        # computed by mj_inverse
+        mjtNum*   qfrc_inverse          # net external force; should equal:        (nv x 1)
+                                        #  qfrc_applied + J'*xfrc_applied + qfrc_actuator
+
+        # computed by mj_sensorAcc/mj_rnePostConstraint; rotation:translation format
+        mjtNum*   cacc                  # com-based acceleration                   (nbody x 6)
+        mjtNum*   cfrc_int              # com-based interaction force with parent  (nbody x 6)
+        mjtNum*   cfrc_ext              # com-based external force on body         (nbody x 6)
+
+
+#---------------------------------- callback function types ----------------------------
+
+# generic MuJoCo function
+ctypedef void (*mjfGeneric)(const mjModel* m, mjData* d)
+
+# sensor simulation
+ctypedef void (*mjfSensor)(const mjModel* m, mjData* d, int stage)
+
+# timer
+ctypedef long long int (*mjfTime)();
+
+# actuator dynamics, gain, bias
+ctypedef mjtNum (*mjfAct)(const mjModel* m, const mjData* d, int id);
+
+# solver impedance
+ctypedef mjtNum (*mjfSolImp)(const mjModel* m, const mjData* d, int id,
+                             mjtNum distance, mjtNum* constimp);
+
+# solver reference
+ctypedef void (*mjfSolRef)(const mjModel* m, const mjData* d, int id,
+                           mjtNum constimp, mjtNum imp, int dim, mjtNum* ref);
+
+# collision detection
+ctypedef int (*mjfCollision)(const mjModel* m, const mjData* d,
+                             mjContact* con, int g1, int g2, mjtNum margin);

mujoco-py-2.1.2.14/mujoco_py/pxd/mjmodel.pxd

@@ -0,0 +1,834 @@
+cdef struct mjVisual_global_:   # global parameters
+    float fovy                  # y-field of view (deg) for free camera
+    float ipd                   # inter-pupilary distance for free camera
+    float linewidth             # line width for wireframe rendering
+    float glow                  # glow coefficient for selected body
+    int offwidth                # width of offscreen buffer
+    int offheight               # height of offscreen buffer
+
+cdef struct mjVisual_quality:   # rendering quality
+    int   shadowsize            # size of shadowmap texture
+    int   offsamples            # number of multisamples for offscreen rendering
+    int   numslices             # number of slices for Glu drawing
+    int   numstacks             # number of stacks for Glu drawing
+    int   numquads              # number of quads for box rendering
+
+cdef struct mjVisual_headlight: # head light
+    float ambient[3]            # ambient rgb (alpha=1)
+    float diffuse[3]            # diffuse rgb (alpha=1)
+    float specular[3]           # specular rgb (alpha=1)
+    int   active                # is headlight active
+
+cdef struct mjVisual_map:       # mapping
+    float stiffness             # mouse perturbation stiffness (space->force)
+    float stiffnessrot          # mouse perturbation stiffness (space->torque)
+    float force                 # from force units to space units
+    float torque                # from torque units to space units
+    float alpha                 # scale geom alphas when transparency is enabled
+    float fogstart              # OpenGL fog starts at fogstart * mjModel.stat.extent
+    float fogend                # OpenGL fog ends at fogend * mjModel.stat.extent
+    float znear                 # near clipping plane = znear * mjModel.stat.extent
+    float zfar                  # far clipping plane = zfar * mjModel.stat.extent
+    float haze                  # haze ratio
+    float shadowclip            # directional light: shadowclip * mjModel.stat.extent
+    float shadowscale           # spot light: shadowscale * light.cutoff
+    float actuatortendon        # scale tendon width
+
+cdef struct mjVisual_scale:     # scale of decor elements relative to mean body size
+    float forcewidth            # width of force arrow
+    float contactwidth          # contact width
+    float contactheight         # contact height
+    float connect               # autoconnect capsule width
+    float com                   # com radius
+    float camera                # camera object
+    float light                 # light object
+    float selectpoint           # selection point
+    float jointlength           # joint length
+    float jointwidth            # joint width
+    float actuatorlength        # actuator length
+    float actuatorwidth         # actuator width
+    float framelength           # bodyframe axis length
+    float framewidth            # bodyframe axis width
+    float constraint            # constraint width
+    float slidercrank           # slidercrank width
+
+cdef struct mjVisual_rgba:      # color of decor elements
+    float fog[4]                # external force
+    float haze[4]               # haze
+    float force[4]              # external force
+    float inertia[4]            # inertia box
+    float joint[4]              # joint
+    float actuator[4]           # actuator
+    float actuatornegative[4]   # actuator, negative limit
+    float actuatorpositive[4]   # actuator, positive limit
+    float com[4]                # center of mass
+    float camera[4]             # camera object
+    float light[4]              # light object
+    float selectpoint[4]        # selection point
+    float connect[4]            # auto connect
+    float contactpoint[4]       # contact point
+    float contactforce[4]       # contact force
+    float contactfriction[4]    # contact friction force
+    float contacttorque[4]      # contact torque
+    float contactgap[4]         # contact point in gap
+    float rangefinder[4]        # rangefinder ray
+    float constraint[4]         # constraint
+    float slidercrank[4]        # slidercrank
+    float crankbroken[4]        # used when crank must be stretched/broken
+
+cdef extern from "mjmodel.h" nogil:
+    # ---------------------------- floating-point definitions -------------------------------
+    ctypedef double mjtNum
+
+    # global constants
+    enum: mjPI
+    enum: mjMAXVAL
+    enum: mjMINMU
+    enum: mjMINIMP
+    enum: mjMAXIMP
+    enum: mjMAXCONPAIR
+    enum: mjMAXVFS
+    enum: mjMAXVFSNAME
+
+    # ---------------------------- sizes ----------------------------------------------------
+    enum: mjNEQDATA
+    enum: mjNDYN
+    enum: mjNGAIN
+    enum: mjNBIAS
+    enum: mjNREF
+    enum: mjNIMP
+    enum: mjNSOLVER
+
+    # ---------------------------- primitive types (mjt) ------------------------------------
+    ctypedef unsigned char mjtByte       # used for true/false
+
+
+    ctypedef enum mjtDisableBit:        # disable default feature bitflags
+        mjDSBL_CONSTRAINT               # entire constraint solver
+        mjDSBL_EQUALITY                 # equality constraints
+        mjDSBL_FRICTIONLOSS             # joint and tendon frictionloss constraints
+        mjDSBL_LIMIT                    # joint and tendon limit constraints
+        mjDSBL_CONTACT                  # contact constraints
+        mjDSBL_PASSIVE                  # passive forces
+        mjDSBL_GRAVITY                  # gravitational forces
+        mjDSBL_CLAMPCTRL                # clamp control to specified range
+        mjDSBL_WARMSTART                # warmstart constraint solver
+        mjDSBL_FILTERPARENT             # remove collisions with parent body
+        mjDSBL_ACTUATION                # apply actuation forces
+        mjDSBL_REFSAFE                  # integrator safety: make ref[0]>=2*timestep
+    enum: mjNDISABLE                      # number of disable flags
+
+    ctypedef enum mjtEnableBit:         # enable optional feature bitflags
+        mjENBL_OVERRIDE                 # override contact parameters
+        mjENBL_ENERGY                   # energy computation
+        mjENBL_FWDINV		        # record solver statistics
+        mjENBL_SENSORNOISE              # add noise to sensor data
+    enum: mjNENABLE                     # number of enable flags
+
+    ctypedef enum mjtJoint:             # type of degree of freedom
+        mjJNT_FREE          = 0,        # global position and orientation (quat)       (7)
+        mjJNT_BALL,                     # orientation (quat) relative to parent        (4)
+        mjJNT_SLIDE,                    # sliding distance along body-fixed axis       (1)
+        mjJNT_HINGE                     # rotation angle (rad) around body-fixed axis  (1)
+
+    ctypedef enum mjtGeom:               # type of geometric shape
+        # regular geom types
+        mjGEOM_PLANE        = 0,        # plane
+        mjGEOM_HFIELD,                  # height field
+        mjGEOM_SPHERE,                  # sphere
+        mjGEOM_CAPSULE,                 # capsule
+        mjGEOM_ELLIPSOID,               # ellipsoid
+        mjGEOM_CYLINDER,                # cylinder
+        mjGEOM_BOX,                     # box
+        mjGEOM_MESH,                    # mesh
+
+        mjNGEOMTYPES,                   # number of regular geom types
+
+        # rendering-only geom types: not used in mjModel, not counted in mjNGEOMTYPES
+        mjGEOM_ARROW        = 100,      # arrow
+        mjGEOM_ARROW1,                  # arrow without wedges
+        mjGEOM_ARROW2,                  # arrow in both directions
+        mjGEOM_LABEL,                   # text label
+
+        mjGEOM_NONE         = 1001      # missing geom type
+
+
+    ctypedef enum mjtCamLight:           # tracking mode for camera and light
+        mjCAMLIGHT_FIXED    = 0,        # pos and rot fixed in body
+        mjCAMLIGHT_TRACK,               # pos tracks body, rot fixed in global
+        mjCAMLIGHT_TRACKCOM,            # pos tracks subtree com, rot fixed in body
+        mjCAMLIGHT_TARGETBODY,          # pos fixed in body, rot tracks target body
+        mjCAMLIGHT_TARGETBODYCOM        # pos fixed in body, rot tracks target subtree com
+
+
+    ctypedef enum mjtTexture:            # type of texture
+        mjTEXTURE_2D        = 0,        # 2d texture, suitable for planes and hfields
+        mjTEXTURE_CUBE,                 # cube texture, suitable for all other geom types
+        mjTEXTURE_SKYBOX                # cube texture used as skybox
+
+
+    ctypedef enum mjtIntegrator:         # integrator mode
+        mjINT_EULER         = 0,        # semi-implicit Euler
+        mjINT_RK4                       # 4th-order Runge Kutta
+
+
+    ctypedef enum mjtCollision:         # collision mode for selecting geom pairs
+        mjCOL_ALL           = 0,        # test precomputed and dynamic pairs
+        mjCOL_PAIR,                     # test predefined pairs only
+        mjCOL_DYNAMIC                   # test dynamic pairs only
+
+
+    ctypedef enum mjtCone:              # type of friction cone
+        mjCONE_PYRAMIDAL     = 0,       # pyramidal
+        mjCONE_ELLIPTIC                 # elliptic
+
+
+    ctypedef enum mjtJacobian:          # type of constraint Jacobian
+        mjJAC_DENSE          = 0,       # dense
+        mjJAC_SPARSE,                   # sparse
+        mjJAC_AUTO                      # dense if nv<60, sparse otherwise
+
+
+    ctypedef enum mjtSolver:            # constraint solver algorithm
+        mjSOL_PGS            = 0,       # PGS    (dual)
+        mjSOL_CG,                       # CG     (primal)
+        mjSOL_NEWTON                    # Newton (primal)
+
+
+    ctypedef enum mjtImp:                # how to interpret solimp parameters
+        mjIMP_CONSTANT      = 0,        # constant solimp[1]
+        mjIMP_SIGMOID,                  # sigmoid from solimp[0] to solimp[1], width solimp[2]
+        mjIMP_LINEAR,                   # piece-wise linear sigmoid
+        mjIMP_USER                      # impedance computed by callback
+
+
+    ctypedef enum mjtRef:                # how to interpret solref parameters
+        mjREF_SPRING        = 0,        # spring-damper: timeconst=solref[0], dampratio=solref[1]
+        mjREF_USER                      # reference computed by callback
+
+
+    ctypedef enum mjtEq:                 # type of equality constraint
+        mjEQ_CONNECT        = 0,        # connect two bodies at a point (ball joint)
+        mjEQ_WELD,                      # fix relative position and orientation of two bodies
+        mjEQ_JOINT,                     # couple the values of two scalar joints with cubic
+        mjEQ_TENDON,                    # couple the lengths of two tendons with cubic
+        mjEQ_DISTANCE                   # fix the contact distance betweent two geoms
+
+
+    ctypedef enum mjtWrap:               # type of tendon wrap object
+        mjWRAP_NONE         = 0,        # null object
+        mjWRAP_JOINT,                   # constant moment arm
+        mjWRAP_PULLEY,                  # pulley used to split tendon
+        mjWRAP_SITE,                    # pass through site
+        mjWRAP_SPHERE,                  # wrap around sphere
+        mjWRAP_CYLINDER                 # wrap around (infinite) cylinder
+
+
+    ctypedef enum mjtTrn:                # type of actuator transmission
+        mjTRN_JOINT         = 0,        # force on joint
+        mjTRN_JOINTINPARENT,            # force on joint, expressed in parent frame
+        mjTRN_SLIDERCRANK,              # force via slider-crank linkage
+        mjTRN_TENDON,                   # force on tendon
+        mjTRN_SITE,                     # force on site
+
+        mjTRN_UNDEFINED     = 1000      # undefined transmission type
+
+
+    ctypedef enum mjtDyn:                # type of actuator dynamics
+        mjDYN_NONE          = 0,        # no internal dynamics; ctrl specifies force
+        mjDYN_INTEGRATOR,               # integrator: da/dt = u
+        mjDYN_FILTER,                   # linear filter: da/dt = (u-a) / tau
+        mjDYN_USER                      # user-defined dynamics type
+
+
+    ctypedef enum mjtGain:               # type of actuator gain
+        mjGAIN_FIXED        = 0,        # fixed gain
+        mjGAIN_USER                     # user-defined gain type
+
+
+    ctypedef enum mjtBias:               # type of actuator bias
+        mjBIAS_NONE         = 0,        # no bias
+        mjBIAS_AFFINE,                  # const + kp*length + kv*velocity
+        mjBIAS_USER                     # user-defined bias type
+
+
+    ctypedef enum mjtObj:                # type of MujoCo object
+        mjOBJ_UNKNOWN       = 0,        # unknown object type
+        mjOBJ_BODY,                     # body
+        mjOBJ_XBODY,                    # body, used to access regular frame instead of i-frame
+        mjOBJ_JOINT,                    # joint
+        mjOBJ_DOF,                      # dof
+        mjOBJ_GEOM,                     # geom
+        mjOBJ_SITE,                     # site
+        mjOBJ_CAMERA,                   # camera
+        mjOBJ_LIGHT,                    # light
+        mjOBJ_MESH,                     # mesh
+        mjOBJ_HFIELD,                   # heightfield
+        mjOBJ_TEXTURE,                  # texture
+        mjOBJ_MATERIAL,                 # material for rendering
+        mjOBJ_PAIR,                     # geom pair to include
+        mjOBJ_EXCLUDE,                  # body pair to exclude
+        mjOBJ_EQUALITY,                 # equality constraint
+        mjOBJ_TENDON,                   # tendon
+        mjOBJ_ACTUATOR,                 # actuator
+        mjOBJ_SENSOR,                   # sensor
+        mjOBJ_NUMERIC,                  # numeric
+        mjOBJ_TEXT,                     # text
+        mjOBJ_TUPLE,                    # tuple
+        mjOBJ_KEY                       # keyframe
+
+
+    ctypedef enum mjtConstraint:         # type of constraint
+        mjCNSTR_EQUALITY    = 0,        # equality constraint
+        mjCNSTR_FRICTION_DOF,           # dof friction
+        mjCNSTR_FRICTION_TENDON,        # tendon friction
+        mjCNSTR_LIMIT_JOINT,            # joint limit
+        mjCNSTR_LIMIT_TENDON,           # tendon limit
+        mjCNSTR_CONTACT_FRICTIONLESS,   # frictionless contact
+        mjCNSTR_CONTACT_PYRAMIDAL,      # frictional contact, pyramidal friction cone
+        mjCNSTR_CONTACT_ELLIPTIC        # frictional contact, elliptic friction cone
+
+
+    ctypedef enum mjtConstraintState:   # constraint state
+        mjCNSTRSTATE_SATISFIED = 0,     # constraint satisfied, zero cost (limit, contact)
+        mjCNSTRSTATE_QUADRATIC,         # quadratic cost (equality, friction, limit, contact)
+        mjCNSTRSTATE_LINEARNEG,         # linear cost, negative side (friction)
+        mjCNSTRSTATE_LINEARPOS,         # linear cost, positive side (friction)
+        mjCNSTRSTATE_CONE               # squared distance to cone cost (elliptic contact)
+
+
+
+    ctypedef enum mjtSensor:            # type of sensor
+        # common robotic sensors, attached to a site
+        mjSENS_TOUCH        = 0,        # scalar contact normal forces summed over sensor zone
+        mjSENS_ACCELEROMETER,           # 3D linear acceleration, in local frame
+        mjSENS_VELOCIMETER,             # 3D linear velocity, in local frame
+        mjSENS_GYRO,                    # 3D angular velocity, in local frame
+        mjSENS_FORCE,                   # 3D force between site's body and its parent body
+        mjSENS_TORQUE,                  # 3D torque between site's body and its parent body
+        mjSENS_MAGNETOMETER,            # 3D magnetometer
+        mjSENS_RANGEFINDER,             # scalar distance to nearest geom or site along z-axis
+
+        # sensors related to scalar joints, tendons, actuators
+        mjSENS_JOINTPOS,                # scalar joint position (hinge and slide only)
+        mjSENS_JOINTVEL,                # scalar joint velocity (hinge and slide only)
+        mjSENS_TENDONPOS,               # scalar tendon position
+        mjSENS_TENDONVEL,               # scalar tendon velocity
+        mjSENS_ACTUATORPOS,             # scalar actuator position
+        mjSENS_ACTUATORVEL,             # scalar actuator velocity
+        mjSENS_ACTUATORFRC,             # scalar actuator force
+
+        # sensors related to ball joints
+        mjSENS_BALLQUAT,                # 4D ball joint quaterion
+        mjSENS_BALLANGVEL,              # 3D ball joint angular velocity
+
+        # sensors attached to an object with spatial frame: (x)body, geom, site, camera
+        mjSENS_FRAMEPOS,                # 3D position
+        mjSENS_FRAMEQUAT,               # 4D unit quaternion orientation
+        mjSENS_FRAMEXAXIS,              # 3D unit vector: x-axis of object's frame
+        mjSENS_FRAMEYAXIS,              # 3D unit vector: y-axis of object's frame
+        mjSENS_FRAMEZAXIS,              # 3D unit vector: z-axis of object's frame
+        mjSENS_FRAMELINVEL,             # 3D linear velocity
+        mjSENS_FRAMEANGVEL,             # 3D angular velocity
+        mjSENS_FRAMELINACC,             # 3D linear acceleration
+        mjSENS_FRAMEANGACC,             # 3D angular acceleration
+
+        # sensors related to kinematic subtrees; attached to a body (which is the subtree root)
+        mjSENS_SUBTREECOM,              # 3D center of mass of subtree
+        mjSENS_SUBTREELINVEL,           # 3D linear velocity of subtree
+        mjSENS_SUBTREEANGMOM,           # 3D angular momentum of subtree
+
+        # user-defined sensor
+        mjSENS_USER                     # sensor data provided by mjcb_sensor callback
+
+
+    ctypedef enum mjtStage:             # computation stage
+        mjSTAGE_NONE        = 0,        # no computations
+        mjSTAGE_POS,                    # position-dependent computations
+        mjSTAGE_VEL,                    # velocity-dependent computations
+        mjSTAGE_ACC                     # acceleration/force-dependent computations
+
+
+    ctypedef enum mjtDataType:           # data type for sensors
+        mjDATATYPE_REAL     = 0,        # real values, no constraints
+        mjDATATYPE_POSITIVE,            # positive values; 0 or negative: inactive
+        mjDATATYPE_AXIS,                # 3D unit vector
+        mjDATATYPE_QUAT                 # unit quaternion
+
+    #------------------------------ mjVFS --------------------------------------------------
+
+    ctypedef struct mjVFS:              # virtual file system for loading from memory
+        int   nfile                     # number of files present
+        char  filename[mjMAXVFS][mjMAXVFSNAME] # file name without path
+        int   filesize[mjMAXVFS]        # file size in bytes
+        void* filedata[mjMAXVFS]        # buffer with file data
+
+
+
+    #------------------------------ mjOption -----------------------------------------------
+
+    ctypedef struct mjOption:                    # physics options
+        # timing parameters
+        mjtNum timestep                 # timestep
+        mjtNum apirate                  # update rate for remote API (Hz)
+
+        # solver parameters
+        mjtNum impratio                 # ratio of friction-to-normal contact impedance
+        mjtNum tolerance                # solver convergence tolerance
+        mjtNum noslip_tolerance         # noslip solver tolerance
+        mjtNum mpr_tolerance            # MPR solver tolerance
+
+        # physical constants
+        mjtNum gravity[3]               # gravitational acceleration
+        mjtNum wind[3]                  # wind (for lift, drag and viscosity)
+        mjtNum magnetic[3]              # global magnetic flux
+        mjtNum density                  # density of medium
+        mjtNum viscosity                # viscosity of medium
+
+        # override contact solver parameters (if enabled)
+        mjtNum o_margin                 # margin
+        mjtNum o_solref[mjNREF]         # solref
+        mjtNum o_solimp[mjNIMP]         # solimp
+
+        # discrete settings
+        int integrator                  # integration mode (mjtIntegrator)
+        int collision                   # collision mode (mjtCollision)
+        int cone                        # type of friction cone (mjtCone)
+        int jacobian                    # type of Jacobian (mjtJacobian)
+        int solver                      # solver mode (mjtSolver)
+        int iterations                  # maximum number of solver iterations
+        int noslip_iterations           # maximum number of noslip solver iterations
+        int mpr_iterations              # maximum number of MPR solver iterations
+        int disableflags                # bit flags for disabling standard features
+        int enableflags                 # bit flags for enabling optional features
+
+    #------------------------------ mjLROpt ------------------------------------------------
+
+    ctypedef struct mjLROpt:
+        # flags
+        int mode                        # which actuators to process (mjtLRMode)
+        int useexisting                 # use existing length range if available
+        int uselimit                    # use joint and tendon limits if available
+
+        # algorithm parameters
+        mjtNum accel                    # target acceleration used to compute force
+        mjtNum maxforce                 # maximum force; 0: no limit
+        mjtNum timeconst                # time constant for velocity reduction; min 0.01
+        mjtNum timestep                 # simulation timestep; 0: use mjOption.timestep
+        mjtNum inttotal                 # total simulation time interval
+        mjtNum inteval                  # evaluation time interval (at the end)
+        mjtNum tolrange                 # convergence tolerance (relative to range)
+
+    #------------------------------ mjVisual -----------------------------------------------
+
+
+    ctypedef struct mjVisual:
+        mjVisual_global_ global_ "global"
+        mjVisual_quality quality
+        mjVisual_headlight headlight
+        mjVisual_map map
+        mjVisual_scale scale
+        mjVisual_rgba rgba
+
+    #------------------------------ mjStatistic --------------------------------------------
+
+    ctypedef struct mjStatistic:                 # model statistics (in qpos0)
+        mjtNum meaninertia              # mean diagonal inertia
+        mjtNum meanmass                 # mean body mass
+        mjtNum meansize                 # mean body size
+        mjtNum extent                   # spatial extent
+        mjtNum center[3]                # center of model
+
+
+    # ---------------------------------- mjModel --------------------------------------------
+    ctypedef struct mjModel:
+        # ------------------------------- sizes
+
+        # sizes needed at mjModel construction
+        int nq                          # number of generalized coordinates = dim(qpos)
+        int nv                          # number of degrees of freedom = dim(qvel)
+        int nu                          # number of actuators/controls = dim(ctrl)
+        int na                          # number of activation states = dim(act)
+        int nbody                       # number of bodies
+        int njnt                        # number of joints
+        int ngeom                       # number of geoms
+        int nsite                       # number of sites
+        int ncam                        # number of cameras
+        int nlight                      # number of lights
+        int nmesh                       # number of meshes
+        int nmeshvert                   # number of vertices in all meshes
+        int nmeshtexvert;               # number of vertices with texcoords in all meshes
+        int nmeshface                   # number of triangular faces in all meshes
+        int nmeshgraph                  # number of ints in mesh auxiliary data
+        int nskin                       # number of skins
+        int nskinvert                   # number of vertices in all skins
+        int nskintexvert                # number of vertiex with texcoords in all skins
+        int nskinface                   # number of triangular faces in all skins
+        int nskinbone                   # number of bones in all skins
+        int nskinbonevert               # number of vertices in all skin bones
+        int nhfield                     # number of heightfields
+        int nhfielddata                 # number of data points in all heightfields
+        int ntex                        # number of textures
+        int ntexdata                    # number of bytes in texture rgb data
+        int nmat                        # number of materials
+        int npair                       # number of predefined geom pairs
+        int nexclude                    # number of excluded geom pairs
+        int neq                         # number of equality constraints
+        int ntendon                     # number of tendons
+        int nwrap                       # number of wrap objects in all tendon paths
+        int nsensor                     # number of sensors
+        int nnumeric                    # number of numeric custom fields
+        int nnumericdata                # number of mjtNums in all numeric fields
+        int ntext                       # number of text custom fields
+        int ntextdata                   # number of mjtBytes in all text fields
+        int ntuple                      # number of tuple custom fields
+        int ntupledata                  # number of objects in all tuple fields
+        int nkey                        # number of keyframes
+        int nmocap                      # number of mocap bodies
+        int nuser_body                  # number of mjtNums in body_user
+        int nuser_jnt                   # number of mjtNums in jnt_user
+        int nuser_geom                  # number of mjtNums in geom_user
+        int nuser_site                  # number of mjtNums in site_user
+        int nuser_cam                   # number of mjtNums in cam_user
+        int nuser_tendon                # number of mjtNums in tendon_user
+        int nuser_actuator              # number of mjtNums in actuator_user
+        int nuser_sensor                # number of mjtNums in sensor_user
+        int nnames                      # number of chars in all names
+
+        # sizes set after jModel construction (only affect mjData)
+        int nM                          # number of non-zeros in sparse inertia matrix
+        int nemax                       # number of potential equality-constraint rows
+        int njmax                       # number of available rows in constraint Jacobian
+        int nconmax                     # number of potential contacts in contact list
+        int nstack                      # number of fields in mjData stack
+        int nuserdata                   # number of extra fields in mjData
+        int nsensordata                 # number of fields in sensor data vector
+
+        int nbuffer                     # number of bytes in buffer
+
+        # ------------------------------- options and statistics
+
+        mjOption opt                    # physics options
+        mjVisual vis                    # visualization options
+        mjStatistic stat                # model statistics
+
+        # ------------------------------- buffers
+
+        # main buffer
+        void*     buffer                # main buffer; all pointers point in it    (nbuffer)
+
+        # default generalized coordinates
+        mjtNum*   qpos0                 # qpos values at default pose              (nq x 1)
+        mjtNum*   qpos_spring           # reference pose for springs               (nq x 1)
+
+        # bodies
+        int*      body_parentid         # id of body's parent                      (nbody x 1)
+        int*      body_rootid           # id of root above body                    (nbody x 1)
+        int*      body_weldid           # id of body that this body is welded to   (nbody x 1)
+        int*      body_mocapid          # id of mocap data; -1: none               (nbody x 1)
+        int*      body_jntnum           # number of joints for this body           (nbody x 1)
+        int*      body_jntadr           # start addr of joints; -1: no joints      (nbody x 1)
+        int*      body_dofnum           # number of motion degrees of freedom      (nbody x 1)
+        int*      body_dofadr           # start addr of dofs; -1: no dofs          (nbody x 1)
+        int*      body_geomnum          # number of geoms                          (nbody x 1)
+        int*      body_geomadr          # start addr of geoms; -1: no geoms        (nbody x 1)
+        mjtByte*  body_simple           # body is simple (has diagonal M)          (nbody x 1)
+        mjtByte*  body_sameframe        # inertial frame is same as body frame     (nbody x 1)
+        mjtNum*   body_pos              # position offset rel. to parent body      (nbody x 3)
+        mjtNum*   body_quat             # orientation offset rel. to parent body   (nbody x 4)
+        mjtNum*   body_ipos             # local position of center of mass         (nbody x 3)
+        mjtNum*   body_iquat            # local orientation of inertia ellipsoid   (nbody x 4)
+        mjtNum*   body_mass             # mass                                     (nbody x 1)
+        mjtNum*   body_subtreemass      # mass of subtree starting at this body    (nbody x 1)
+        mjtNum*   body_inertia          # diagonal inertia in ipos/iquat frame     (nbody x 3)
+        mjtNum*   body_invweight0       # mean inv inert in qpos0 (trn, rot)       (nbody x 2)
+        mjtNum*   body_user             # user data                                (nbody x nuser_body)
+
+        # joints
+        int*      jnt_type              # type of joint (mjtJoint)                 (njnt x 1)
+        int*      jnt_qposadr           # start addr in 'qpos' for joint's data    (njnt x 1)
+        int*      jnt_dofadr            # start addr in 'qvel' for joint's data    (njnt x 1)
+        int*      jnt_bodyid            # id of joint's body                       (njnt x 1)
+        int*      jnt_group             # group for visibility                     (njnt x 1)
+        mjtByte*  jnt_limited           # does joint have limits                   (njnt x 1)
+        mjtNum*   jnt_solref            # constraint solver reference: limit       (njnt x mjNREF)
+        mjtNum*   jnt_solimp            # constraint solver impedance: limit       (njnt x mjNIMP)
+        mjtNum*   jnt_pos               # local anchor position                    (njnt x 3)
+        mjtNum*   jnt_axis              # local joint axis                         (njnt x 3)
+        mjtNum*   jnt_stiffness         # stiffness coefficient                    (njnt x 1)
+        mjtNum*   jnt_range             # joint limits                             (njnt x 2)
+        mjtNum*   jnt_margin            # min distance for limit detection         (njnt x 1)
+        mjtNum*   jnt_user              # user data                                (njnt x nuser_jnt)
+
+        # dofs
+        int*      dof_bodyid            # id of dof's body                         (nv x 1)
+        int*      dof_jntid             # id of dof's joint                        (nv x 1)
+        int*      dof_parentid          # id of dof's parent; -1: none             (nv x 1)
+        int*      dof_Madr              # dof address in M-diagonal                (nv x 1)
+        int*      dof_simplenum         # number of consecutive simple dofs        (nv x 1)
+        mjtNum*   dof_solref            # constraint solver reference:frictionloss (nv x mjNREF)
+        mjtNum*   dof_solimp            # constraint solver impedance:frictionloss (nv x mjNIMP)
+        mjtNum*   dof_frictionloss      # dof friction loss                        (nv x 1)
+        mjtNum*   dof_armature          # dof armature inertia/mass                (nv x 1)
+        mjtNum*   dof_damping           # damping coefficient                      (nv x 1)
+        mjtNum*   dof_invweight0        # inv. diag. inertia in qpos0              (nv x 1)
+        mjtNum*   dof_M0                # diag. inertia in qpos0                   (nv x 1)
+
+        # geoms
+        int*      geom_type             # geometric type (mjtGeom)                 (ngeom x 1)
+        int*      geom_contype          # geom contact type                        (ngeom x 1)
+        int*      geom_conaffinity      # geom contact affinity                    (ngeom x 1)
+        int*      geom_condim           # contact dimensionality (1, 3, 4, 6)      (ngeom x 1)
+        int*      geom_bodyid           # id of geom's body                        (ngeom x 1)
+        int*      geom_dataid           # id of geom's mesh/hfield (-1: none)      (ngeom x 1)
+        int*      geom_matid            # material id for rendering                (ngeom x 1)
+        int*      geom_group            # group for visibility                     (ngeom x 1)
+        int*      geom_priority         # geom contact priority                    (ngeom x 1)
+        mjtByte*  geom_sameframe        # same as body frame (1) or iframe (2)     (ngeom x 1)
+        mjtNum*   geom_solmix           # mixing coef for solref/imp in geom pair  (ngeom x 1)
+        mjtNum*   geom_solref           # constraint solver reference: contact     (ngeom x mjNREF)
+        mjtNum*   geom_solimp           # constraint solver impedance: contact     (ngeom x mjNIMP)
+        mjtNum*   geom_size             # geom-specific size parameters            (ngeom x 3)
+        mjtNum*   geom_rbound           # radius of bounding sphere                (ngeom x 1)
+        mjtNum*   geom_pos              # local position offset rel. to body       (ngeom x 3)
+        mjtNum*   geom_quat             # local orientation offset rel. to body    (ngeom x 4)
+        mjtNum*   geom_friction         # friction for (slide, spin, roll)         (ngeom x 3)
+        mjtNum*   geom_margin           # detect contact if dist<margin            (ngeom x 1)
+        mjtNum*   geom_gap              # include in solver if dist<margin-gap     (ngeom x 1)
+        mjtNum*   geom_user             # user data                                (ngeom x nuser_geom)
+        float*    geom_rgba             # rgba when material is omitted            (ngeom x 4)
+
+        # sites
+        int*      site_type             # geom type for rendering (mjtGeom)        (nsite x 1)
+        int*      site_bodyid           # id of site's body                        (nsite x 1)
+        int*      site_matid            # material id for rendering                (nsite x 1)
+        int*      site_group            # group for visibility                     (nsite x 1)
+        mjtByte*  site_sameframe        # same as body frame (1) or iframe (2)     (nsite x 1)
+        mjtNum*   site_size             # geom size for rendering                  (nsite x 3)
+        mjtNum*   site_pos              # local position offset rel. to body       (nsite x 3)
+        mjtNum*   site_quat             # local orientation offset rel. to body    (nsite x 4)
+        mjtNum*   site_user             # user data                                (nsite x nuser_site)
+        float*    site_rgba             # rgba when material is omitted            (nsite x 4)
+
+        # cameras
+        int*      cam_mode              # camera tracking mode (mjtCamLight)       (ncam x 1)
+        int*      cam_bodyid            # id of camera's body                      (ncam x 1)
+        int*      cam_targetbodyid      # id of targeted body; -1: none            (ncam x 1)
+        mjtNum*   cam_pos               # position rel. to body frame              (ncam x 3)
+        mjtNum*   cam_quat              # orientation rel. to body frame           (ncam x 4)
+        mjtNum*   cam_poscom0           # global position rel. to sub-com in qpos0 (ncam x 3)
+        mjtNum*   cam_pos0              # global position rel. to body in qpos0    (ncam x 3)
+        mjtNum*   cam_mat0              # global orientation in qpos0              (ncam x 9)
+        mjtNum*   cam_fovy              # y-field of view (deg)                    (ncam x 1)
+        mjtNum*   cam_ipd               # inter-pupilary distance                  (ncam x 1)
+        mjtNum*   cam_user              # user data                                (ncam x nuser_cam)
+
+        # lights
+        int*      light_mode            # light tracking mode (mjtCamLight)        (nlight x 1)
+        int*      light_bodyid          # id of light's body                       (nlight x 1)
+        int*      light_targetbodyid    # id of targeted body; -1: none            (nlight x 1)
+        mjtByte*  light_directional     # directional light                        (nlight x 1)
+        mjtByte*  light_castshadow      # does light cast shadows                  (nlight x 1)
+        mjtByte*  light_active          # is light on                              (nlight x 1)
+        mjtNum*   light_pos             # position rel. to body frame              (nlight x 3)
+        mjtNum*   light_dir             # direction rel. to body frame             (nlight x 3)
+        mjtNum*   light_poscom0         # global position rel. to sub-com in qpos0 (nlight x 3)
+        mjtNum*   light_pos0            # global position rel. to body in qpos0    (nlight x 3)
+        mjtNum*   light_dir0            # global direction in qpos0                (nlight x 3)
+        float*    light_attenuation     # OpenGL attenuation (quadratic model)     (nlight x 3)
+        float*    light_cutoff          # OpenGL cutoff                            (nlight x 1)
+        float*    light_exponent        # OpenGL exponent                          (nlight x 1)
+        float*    light_ambient         # ambient rgb (alpha=1)                    (nlight x 3)
+        float*    light_diffuse         # diffuse rgb (alpha=1)                    (nlight x 3)
+        float*    light_specular        # specular rgb (alpha=1)                   (nlight x 3)
+
+        # meshes
+        int*      mesh_vertadr          # first vertex address                     (nmesh x 1)
+        int*      mesh_vertnum          # number of vertices                       (nmesh x 1)
+        int*      mesh_texcoordadr      # texcoord data address; -1: no texcoord   (nmesh x 1)
+        int*      mesh_faceadr          # first face address                       (nmesh x 1)
+        int*      mesh_facenum          # number of faces                          (nmesh x 1)
+        int*      mesh_graphadr         # graph data address; -1: no graph         (nmesh x 1)
+        float*    mesh_vert             # vertex data for all meshes               (nmeshvert x 3)
+        float*    mesh_normal           # vertex normal data for all meshes        (nmeshvert x 3)
+        float*    mesh_texcoord         # vertex texcoords for all meshes          (nmeshtexvert x 2)
+        int*      mesh_face             # triangle face data                       (nmeshface x 3)
+        int*      mesh_graph            # convex graph data                        (nmeshgraph x 1)
+
+        # skins
+        int*      skin_matid            # skin material id; -1: none               (nskin x 1)
+        float*    skin_rgba             # skin rgba                                (nskin x 4)
+        float*    skin_inflate          # inflate skin in normal direction         (nskin x 1)
+        int*      skin_vertadr          # first vertex address                     (nskin x 1)
+        int*      skin_vertnum          # number of vertices                       (nskin x 1)
+        int*      skin_texcoordadr      # texcoord data address; -1: no texcoord   (nskin x 1)
+        int*      skin_faceadr          # first face address                       (nskin x 1)
+        int*      skin_facenum          # number of faces                          (nskin x 1)
+        int*      skin_boneadr          # first bone in skin                       (nskin x 1)
+        int*      skin_bonenum          # number of bones in skin                  (nskin x 1)
+        float*    skin_vert             # vertex positions for all skin meshes     (nskinvert x 3)
+        float*    skin_texcoord         # vertex texcoords for all skin meshes     (nskintexvert x 2)
+        int*      skin_face             # triangle faces for all skin meshes       (nskinface x 3)
+        int*      skin_bonevertadr      # first vertex in each bone                (nskinbone x 1)
+        int*      skin_bonevertnum      # number of vertices in each bone          (nskinbone x 1)
+        float*    skin_bonebindpos      # bind pos of each bone                    (nskinbone x 3)
+        float*    skin_bonebindquat     # bind quat of each bone                   (nskinbone x 4)
+        int*      skin_bonebodyid       # body id of each bone                     (nskinbone x 1)
+        int*      skin_bonevertid       # mesh ids of vertices in each bone        (nskinbonevert x 1)
+        float*    skin_bonevertweight   # weights of vertices in each bone         (nskinbonevert x 1)
+
+        # height fields
+        mjtNum*   hfield_size           # (x, y, z_top, z_bottom)                  (nhfield x 4)
+        int*      hfield_nrow           # number of rows in grid                   (nhfield x 1)
+        int*      hfield_ncol           # number of columns in grid                (nhfield x 1)
+        int*      hfield_adr            # address in hfield_data                   (nhfield x 1)
+        float*    hfield_data           # elevation data                           (nhfielddata x 1)
+
+        # textures
+        int*      tex_type              # texture type (mjtTexture)                (ntex x 1)
+        int*      tex_height            # number of rows in texture image          (ntex x 1)
+        int*      tex_width             # number of columns in texture image       (ntex x 1)
+        int*      tex_adr               # address in rgb                           (ntex x 1)
+        mjtByte*  tex_rgb               # rgb (alpha = 1)                          (ntexdata x 1)
+
+        # materials
+        int*      mat_texid             # texture id; -1: none                     (nmat x 1)
+        mjtByte*  mat_texuniform        # make texture cube uniform                (nmat x 1)
+        float*    mat_texrepeat         # texture repetition for 2d mapping        (nmat x 2)
+        float*    mat_emission          # emission (x rgb)                         (nmat x 1)
+        float*    mat_specular          # specular (x white)                       (nmat x 1)
+        float*    mat_shininess         # shininess coef                           (nmat x 1)
+        float*    mat_reflectance       # reflectance (0: disable)                 (nmat x 1)
+        float*    mat_rgba              # rgba                                     (nmat x 4)
+
+        # predefined geom pairs for collision detection; has precedence over exclude
+        int*      pair_dim              # contact dimensionality                   (npair x 1)
+        int*      pair_geom1            # id of geom1                              (npair x 1)
+        int*      pair_geom2            # id of geom2                              (npair x 1)
+        int*      pair_signature        # (body1+1)<<16 + body2+1                  (npair x 1)
+        mjtNum*   pair_solref           # constraint solver reference: contact     (npair x mjNREF)
+        mjtNum*   pair_solimp           # constraint solver impedance: contact     (npair x mjNIMP)
+        mjtNum*   pair_margin           # detect contact if dist<margin            (npair x 1)
+        mjtNum*   pair_gap              # include in solver if dist<margin-gap     (npair x 1)
+        mjtNum*   pair_friction         # tangent1, 2, spin, roll1, 2              (npair x 5)
+
+        # excluded body pairs for collision detection
+        int*      exclude_signature     # (body1+1)<<16 + body2+1                  (nexclude x 1)
+
+        # equality constraints
+        int*      eq_type               # constraint type (mjtEq)                  (neq x 1)
+        int*      eq_obj1id             # id of object 1                           (neq x 1)
+        int*      eq_obj2id             # id of object 2                           (neq x 1)
+        mjtByte*  eq_active             # enable/disable constraint                (neq x 1)
+        mjtNum*   eq_solref             # constraint solver reference              (neq x mjNREF)
+        mjtNum*   eq_solimp             # constraint solver impedance              (neq x mjNIMP)
+        mjtNum*   eq_data               # numeric data for constraint              (neq x mjNEQDATA)
+
+        # tendons
+        int*      tendon_adr            # address of first object in tendon's path (ntendon x 1)
+        int*      tendon_num            # number of objects in tendon's path       (ntendon x 1)
+        int*      tendon_matid          # material id for rendering                (ntendon x 1)
+        int*      tendon_group          # group for visibility                     (ntendon x 1)
+        mjtByte*  tendon_limited        # does tendon have length limits           (ntendon x 1)
+        mjtNum*   tendon_width          # width for rendering                      (ntendon x 1)
+        mjtNum*   tendon_solref_lim     # constraint solver reference: limit       (ntendon x mjNREF)
+        mjtNum*   tendon_solimp_lim     # constraint solver impedance: limit       (ntendon x mjNIMP)
+        mjtNum*   tendon_solref_fri     # constraint solver reference: friction    (ntendon x mjNREF)
+        mjtNum*   tendon_solimp_fri     # constraint solver impedance: friction    (ntendon x mjNIMP)
+        mjtNum*   tendon_range          # tendon length limits                     (ntendon x 2)
+        mjtNum*   tendon_margin         # min distance for limit detection         (ntendon x 1)
+        mjtNum*   tendon_stiffness      # stiffness coefficient                    (ntendon x 1)
+        mjtNum*   tendon_damping        # damping coefficient                      (ntendon x 1)
+        mjtNum*   tendon_frictionloss;  # loss due to friction                     (ntendon x 1)
+        mjtNum*   tendon_lengthspring;  # tendon length in qpos_spring             (ntendon x 1)
+        mjtNum*   tendon_length0        # tendon length in qpos0                   (ntendon x 1)
+        mjtNum*   tendon_invweight0     # inv. weight in qpos0                     (ntendon x 1)
+        mjtNum*   tendon_user           # user data                                (ntendon x nuser_tendon)
+        float*    tendon_rgba           # rgba when material is omitted            (ntendon x 4)
+
+        # list of all wrap objects in tendon paths
+        int*      wrap_type             # wrap object type (mjtWrap)               (nwrap x 1)
+        int*      wrap_objid            # object id: geom, site, joint             (nwrap x 1)
+        mjtNum*   wrap_prm              # divisor, joint coef, or site id          (nwrap x 1)
+
+        # actuators
+        int*      actuator_trntype      # transmission type (mjtTrn)               (nu x 1)
+        int*      actuator_dyntype      # dynamics type (mjtDyn)                   (nu x 1)
+        int*      actuator_gaintype     # gain type (mjtGain)                      (nu x 1)
+        int*      actuator_biastype     # bias type (mjtBias)                      (nu x 1)
+        int*      actuator_trnid        # transmission id: joint, tendon, site     (nu x 2)
+        int*      actuator_group        # group for visibility                     (nu x 1)
+        mjtByte*  actuator_ctrllimited; # is control limited                       (nu x 1)
+        mjtByte*  actuator_forcelimited;# is force limited                         (nu x 1)
+        mjtNum*   actuator_dynprm       # dynamics parameters                      (nu x mjNDYN)
+        mjtNum*   actuator_gainprm      # gain parameters                          (nu x mjNGAIN)
+        mjtNum*   actuator_biasprm      # bias parameters                          (nu x mjNBIAS)
+        mjtNum*   actuator_ctrlrange    # range of controls                        (nu x 2)
+        mjtNum*   actuator_forcerange;  # range of forces                          (nu x 2)
+        mjtNum*   actuator_gear         # scale length and transmitted force       (nu x 6)
+        mjtNum*   actuator_cranklength; # crank length for slider-crank            (nu x 1)
+        mjtNum*   actuator_acc0         # acceleration from unit force in qpos0    (nu x 1)
+        mjtNum*   actuator_length0      # actuator length in qpos0                 (nu x 1)
+        mjtNum*   actuator_lengthrange  # ... not yet implemented ???              (nu x 2)
+        mjtNum*   actuator_user         # user data                                (nu x nuser_actuator)
+
+        # sensors
+        int*      sensor_type           # sensor type (mjtSensor)                  (nsensor x 1)
+        int*      sensor_datatype       # numeric data type (mjtDataType)          (nsensor x 1)
+        int*      sensor_needstage      # required compute stage (mjtStage)        (nsensor x 1)
+        int*      sensor_objtype        # type of sensorized object (mjtObj)       (nsensor x 1)
+        int*      sensor_objid          # id of sensorized object                  (nsensor x 1)
+        int*      sensor_dim            # number of scalar outputs                 (nsensor x 1)
+        int*      sensor_adr            # address in sensor array                  (nsensor x 1)
+        mjtNum*   sensor_cutoff         # cutoff for real and positive; 0: ignore  (nsensor x 1)
+        mjtNum*   sensor_noise          # noise standard deviation                 (nsensor x 1)
+        mjtNum*   sensor_user           # user data                                (nsensor x nuser_sensor)
+
+        # custom numeric fields
+        int*      numeric_adr           # address of field in numeric_data         (nnumeric x 1)
+        int*      numeric_size          # size of numeric field                    (nnumeric x 1)
+        mjtNum*   numeric_data          # array of all numeric fields              (nnumericdata x 1)
+
+        # custom text fields
+        int*      text_adr              # address of text in text_data             (ntext x 1)
+        int*      text_size             # size of text field (strlen+1)            (ntext x 1)
+        char*     text_data             # array of all text fields (0-terminated)  (ntextdata x 1)
+
+        # custom tuple fields
+        int*      tuple_adr             # address of text in text_data             (ntuple x 1)
+        int*      tuple_size            # number of objects in tuple               (ntuple x 1)
+        int*      tuple_objtype         # array of object types in all tuples      (ntupledata x 1)
+        int*      tuple_objid           # array of object ids in all tuples        (ntupledata x 1)
+        mjtNum*   tuple_objprm          # array of object params in all tuples     (ntupledata x 1)
+
+        # keyframes
+        mjtNum*   key_time              # key time                                 (nkey x 1)
+        mjtNum*   key_qpos              # key position                             (nkey x nq)
+        mjtNum*   key_qvel              # key velocity                             (nkey x nv)
+        mjtNum*   key_act               # key activation                           (nkey x na)
+        mjtNum*   key_mpos              # key mocap position                       (nkey x 3*nmocap)
+        mjtNum*   key_mquat             # key mocap quaternion                     (nkey x 4*nmocap)
+
+        # names
+        int*      name_bodyadr          # body name pointers                       (nbody x 1)
+        int*      name_jntadr           # joint name pointers                      (njnt x 1)
+        int*      name_geomadr          # geom name pointers                       (ngeom x 1)
+        int*      name_siteadr          # site name pointers                       (nsite x 1)
+        int*      name_camadr           # camera name pointers                     (ncam x 1)
+        int*      name_lightadr         # light name pointers                      (nlight x 1)
+        int*      name_meshadr          # mesh name pointers                       (nmesh x 1)
+        int*      name_skinadr          # skin name pointers                       (nskin x 1)
+        int*      name_hfieldadr        # hfield name pointers                     (nhfield x 1)
+        int*      name_texadr           # texture name pointers                    (ntex x 1)
+        int*      name_matadr           # material name pointers                   (nmat x 1)
+        int*      name_pairadr          # geom pair name pointers                  (npair x 1)
+        int*      name_excludeadr       # exclude name pointers                    (nexclude x 1)
+        int*      name_eqadr            # equality constraint name pointers        (neq x 1)
+        int*      name_tendonadr        # tendon name pointers                     (ntendon x 1)
+        int*      name_actuatoradr      # actuator name pointers                   (nu x 1)
+        int*      name_sensoradr        # sensor name pointers                     (nsensor x 1)
+        int*      name_numericadr       # numeric name pointers                    (nnumeric x 1)
+        int*      name_textadr          # text name pointers                       (ntext x 1)
+        int*      name_tupleadr         # tuple name pointers                      (ntuple x 1)
+        int*      name_keyadr           # keyframe name pointers                   (nkey x 1)
+        char*     names                 # names of all objects, 0-terminated       (nnames x 1)

mujoco-py-2.1.2.14/mujoco_py/pxd/mjrender.pxd

@@ -0,0 +1,115 @@
+cdef extern from "mjrender.h" nogil:
+    # Global constants
+    enum: mjNAUX
+    enum: mjMAXTEXTURE
+
+    ctypedef enum mjtGridPos:            # grid position for overlay
+        mjGRID_TOPLEFT      = 0,        # top left
+        mjGRID_TOPRIGHT,                # top right
+        mjGRID_BOTTOMLEFT,              # bottom left
+        mjGRID_BOTTOMRIGHT              # bottom right
+
+
+    ctypedef enum mjtFramebuffer:        # OpenGL framebuffer option
+        mjFB_WINDOW         = 0,        # default/window buffer
+        mjFB_OFFSCREEN                  # offscreen buffer
+
+
+    ctypedef enum mjtFontScale:          # font scale, used at context creation
+        mjFONTSCALE_100     = 100,      # normal scale, suitable in the absence of DPI scaling
+        mjFONTSCALE_150     = 150,      # 150% scale
+        mjFONTSCALE_200     = 200       # 200% scale
+
+
+    ctypedef enum mjtFont:               # font type, used at each text operation
+        mjFONT_NORMAL       = 0,        # normal font
+        mjFONT_SHADOW,                  # normal font with shadow (for higher contrast)
+        mjFONT_BIG                      # big font (for user alerts)
+
+    ctypedef struct mjrRect:                     # OpenGL rectangle
+        int left                        # left (usually 0)
+        int bottom                      # bottom (usually 0)
+        int width                       # width (usually buffer width)
+        int height                      # height (usually buffer height)
+
+
+    ctypedef struct mjrContext:                  # custom OpenGL context
+        # parameters copied from mjVisual
+        float lineWidth                 # line width for wireframe rendering
+        float shadowClip                # clipping radius for directional lights
+        float shadowScale               # fraction of light cutoff for spot lights
+        float fogStart                  # fog start = stat.extent * vis.map.fogstart
+        float fogEnd                    # fog end = stat.extent * vis.map.fogend
+        float fogRGBA[4]                # fog rgba
+        int shadowSize                  # size of shadow map texture
+        int offWidth                    # width of offscreen buffer
+        int offHeight                   # height of offscreen buffer
+        int offSamples                  # number of offscreen buffer multisamples
+
+        # parameters specified at creation
+        int fontScale;                  # font scale
+        int auxWidth[mjNAUX]            # auxiliary buffer width
+        int auxHeight[mjNAUX]           # auxiliary buffer height
+        int auxSamples[mjNAUX]          # auxiliary buffer multisamples
+
+        # offscreen rendering objects
+        unsigned int offFBO             # offscreen framebuffer object
+        unsigned int offFBO_r           # offscreen framebuffer for resolving multisamples
+        unsigned int offColor           # offscreen color buffer
+        unsigned int offColor_r         # offscreen color buffer for resolving multisamples
+        unsigned int offDepthStencil    # offscreen depth and stencil buffer
+        unsigned int offDepthStencil_r  # offscreen depth and stencil buffer for resolving multisamples
+
+        # shadow rendering objects
+        unsigned int shadowFBO          # shadow map framebuffer object
+        unsigned int shadowTex          # shadow map texture
+
+        # auxiliary buffers
+        unsigned int auxFBO[mjNAUX]     # auxiliary framebuffer object
+        unsigned int auxFBO_r[mjNAUX]   # auxiliary framebuffer object for resolving
+        unsigned int auxColor[mjNAUX]   # auxiliary color buffer
+        unsigned int auxColor_r[mjNAUX] # auxiliary color buffer for resolving
+
+        # texture objects and info
+        int ntexture                    # number of allocated textures
+        int textureType[100]            # type of texture (mjtTexture)
+        unsigned int texture[100]       # texture names
+
+        # displaylist starting positions
+        unsigned int basePlane          # all planes from model
+        unsigned int baseMesh           # all meshes from model
+        unsigned int baseHField         # all hfields from model
+        unsigned int baseBuiltin        # all buildin geoms, with quality from model
+        unsigned int baseFontNormal     # normal font
+        unsigned int baseFontShadow     # shadow font
+        unsigned int baseFontBig        # big font
+
+        # displaylist ranges
+        int     rangePlane              # all planes from model
+        int     rangeMesh               # all meshes from model
+        int     rangeHField             # all hfields from model
+        int     rangeBuiltin            # all builtin geoms, with quality from model
+        int     rangeFont               # all characters in font
+
+        # skin VBOs
+        int      nskin                  # number of skins
+        unsigned int* skinvertVBO       # skin vertex position VBOs
+        unsigned int* skinnormalVBO     # skin vertex normal VBOs
+        unsigned int* skintexcoordVBO   # skin vertex texture coordinate VBOs
+        unsigned int* skinfaceVBO       # skin face index VBOs
+
+        # character info
+        int     charWidth[127]          # character widths: normal and shadow
+        int     charWidthBig[127]       # chacarter widths: big
+        int     charHeight              # character heights: normal and shadow
+        int     charHeightBig           # character heights: big
+
+        # capabilities
+        int     glewInitialized         # is glew initialized
+        int     windowAvailable         # is default/window framebuffer available
+        int     windowSamples           # number of samples for default/window framebuffer
+        int     windowStereo            # is stereo available for default/window framebuffer
+        int     windowDoublebuffer      # is default/window framebuffer double buffered
+
+        # only field that changes after mjr_makeContext
+        int     currentBuffer           # currently active framebuffer: mjFB_WINDOW or mjFB_OFFSCREEN

mujoco-py-2.1.2.14/mujoco_py/pxd/mujoco.pxd

@@ -0,0 +1,1083 @@
+include "mjmodel.pxd"
+include "mjdata.pxd"
+include "mjrender.pxd"
+include "mjui.pxd"
+include "mjvisualize.pxd"
+
+
+cdef extern from "mujoco.h" nogil:
+    # macros
+    #define mjMARKSTACK   int _mark = d->pstack;
+    #define mjFREESTACK   d->pstack = _mark;
+    #define mjDISABLED(x) (m->opt.disableflags & (x))
+    #define mjENABLED(x)  (m->opt.enableflags & (x))
+
+
+    # user error and memory handlers
+    void  (*mju_user_error)(const char*);
+    void  (*mju_user_warning)(const char*);
+    void* (*mju_user_malloc)(size_t);
+    void  (*mju_user_free)(void*);
+
+
+    # # callbacks extending computation pipeline
+    # mjfGeneric  mjcb_passive;
+    # mjfGeneric  mjcb_control;
+    # mjfSensor   mjcb_sensor;
+    # mjfTime     mjcb_time;
+    # mjfAct      mjcb_act_dyn;
+    mjfAct      mjcb_act_gain;
+    mjfAct      mjcb_act_bias;
+    # mjfSolImp   mjcb_sol_imp;
+    # mjfSolRef   mjcb_sol_ref;
+    #
+    #
+    # # collision function table
+    # mjfCollision mjCOLLISIONFUNC[mjNGEOMTYPES][mjNGEOMTYPES];
+    #
+    #
+    # # string names
+    const char* mjDISABLESTRING[mjNDISABLE];
+    const char* mjENABLESTRING[mjNENABLE];
+    const char* mjTIMERSTRING[mjNTIMER];
+    const char* mjLABELSTRING[mjNLABEL];
+    const char* mjFRAMESTRING[mjNFRAME];
+    const char* mjVISSTRING[mjNVISFLAG][3];
+    const char* mjRNDSTRING[mjNRNDFLAG][3];
+
+
+    #---------------------- Activation -----------------------------------------------------
+
+    # activate license, call mju_error on failure; return 1 if ok, 0 if failure
+    int mj_activate(const char* filename);
+
+    # deactivate license, free memory
+    void mj_deactivate();
+
+    #---------------------- Virtual file system --------------------------------------------
+
+    # Initialize VFS to empty (no deallocation).
+    void mj_defaultVFS(mjVFS* vfs);
+
+    # Add file to VFS, return 0: success, 1: full, 2: repeated name, -1: not found on disk.
+    int mj_addFileVFS(mjVFS* vfs, const char* directory, const char* filename);
+
+    # Make empty file in VFS, return 0: success, 1: full, 2: repeated name.
+    int mj_makeEmptyFileVFS(mjVFS* vfs, const char* filename, int filesize);
+
+    # Return file index in VFS, or -1 if not found in VFS.
+    int mj_findFileVFS(const mjVFS* vfs, const char* filename);
+
+    # Delete file from VFS, return 0: success, -1: not found in VFS.
+    int mj_deleteFileVFS(mjVFS* vfs, const char* filename);
+
+    # Delete all files from VFS.
+    void mj_deleteVFS(mjVFS* vfs);
+
+    #--------------------- Parse and compile ----------------------------------------------
+
+    # Parse XML file in MJCF or URDF format, compile it, return low-level model.
+    # If vfs is not NULL, look up files in vfs before reading from disk.
+    # If error is not NULL, it must have size error_sz.
+    mjModel* mj_loadXML(const char* filename, const mjVFS* vfs,
+                              char* error, int error_sz);
+
+    # Update XML data structures with info from low-level model, save as MJCF.
+    # If error is not NULL, it must have size error_sz.
+    int mj_saveLastXML(const char* filename, const mjModel* m,
+                             char* error, int error_sz);
+
+    # Free last XML model if loaded. Called internally at each load.
+    void mj_freeLastXML();
+
+    # Print internal XML schema as plain text or HTML, with style-padding or  .
+    int mj_printSchema(const char* filename, char* buffer, int buffer_sz,
+                             int flg_html, int flg_pad);
+
+
+    #--------------------- Main simulation ------------------------------------------------
+
+    # Advance simulation, use control callback to obtain external force and control.
+    void mj_step(const mjModel* m, mjData* d);
+
+    # Advance simulation in two steps: before external force and control is set by user.
+    void mj_step1(const mjModel* m, mjData* d);
+
+    # Advance simulation in two steps: after external force and control is set by user.
+    void mj_step2(const mjModel* m, mjData* d);
+
+    # Forward dynamics: same as mj_step but do not integrate in time.
+    void mj_forward(const mjModel* m, mjData* d);
+
+    # Inverse dynamics: qacc must be set before calling.
+    void mj_inverse(const mjModel* m, mjData* d);
+
+    # Forward dynamics with skip; skipstage is mjtStage.
+    void mj_forwardSkip(const mjModel* m, mjData* d,
+                              int skipstage, int skipsensorenergy);
+
+    # Inverse dynamics with skip; skipstage is mjtStage.
+    void mj_inverseSkip(const mjModel* m, mjData* d,
+                        int skipstage, int skipsensorenergy);
+
+    # Forward dynamics with skip; skipstage is mjtStage.
+    void mj_forwardSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor);
+
+    # Inverse dynamics with skip; skipstage is mjtStage.
+    void mj_inverseSkip(const mjModel* m, mjData* d, int skipstage, int skipsensor);
+
+
+    #--------------------- Initialization -------------------------------------------------
+
+    # Set default options for length range computation.
+    void mj_defaultLROpt(mjLROpt* opt);
+
+    # Set solver parameters to default values.
+    void mj_defaultSolRefImp(mjtNum* solref, mjtNum* solimp);
+
+    # Set physics options to default values.
+    void mj_defaultOption(mjOption* opt);
+
+    # Set visual options to default values.
+    void mj_defaultVisual(mjVisual* vis);
+
+    # Copy mjModel, allocate new if dest is NULL.
+    mjModel* mj_copyModel(mjModel* dest, const mjModel* src);
+
+    # Save model to binary MJB file or memory buffer; buffer has precedence when given.
+    void mj_saveModel(const mjModel* m, const char* filename, void* buffer, int buffer_sz);
+
+    # Load model from binary MJB file.
+    # If vfs is not NULL, look up file in vfs before reading from disk.
+    mjModel* mj_loadModel(const char* filename, mjVFS* vfs);
+
+    # Free memory allocation in model.
+    void mj_deleteModel(mjModel* m);
+
+    # Return size of buffer needed to hold model.
+    int mj_sizeModel(const mjModel* m);
+
+    # Allocate mjData correponding to given model.
+    mjData* mj_makeData(const mjModel* m);
+
+    # Copy mjData.
+    mjData* mj_copyData(mjData* dest, const mjModel* m, const mjData* src);
+
+    # Reset data to defaults.
+    void mj_resetData(const mjModel* m, mjData* d);
+
+    # Reset data to defaults, fill everything else with debug_value.
+    void mj_resetDataDebug(const mjModel* m, mjData* d, unsigned char debug_value);
+
+    # Reset data, set fields from specified keyframe.
+    void mj_resetDataKeyframe(const mjModel* m, mjData* d, int key);
+
+    # Allocate array of specified size on mjData stack. Call mju_error on stack overflow.
+    mjtNum* mj_stackAlloc(mjData* d, int size);
+
+    # Free memory allocation in mjData.
+    void mj_deleteData(mjData* d);
+
+    # Reset all callbacks to NULL pointers (NULL is the default).
+    void mj_resetCallbacks();
+
+    # Set constant fields of mjModel, corresponding to qpos0 configuration.
+    void mj_setConst(mjModel* m, mjData* d);
+
+    # Set actuator_lengthrange for specified actuator; return 1 if ok, 0 if error.
+    int mj_setLengthRange(mjModel* m, mjData* d, int index,
+                          const mjLROpt* opt, char* error, int error_sz);
+
+    #--------------------- Printing -------------------------------------------------------
+
+    # Print model to text file.
+    void mj_printModel(const mjModel* m, const char* filename);
+
+    # Print data to text file.
+    void mj_printData(const mjModel* m, mjData* d, const char* filename);
+
+    # Print matrix to screen.
+    void mju_printMat(const mjtNum* mat, int nr, int nc);
+
+    # Print sparse matrix to screen.
+    void mju_printMatSparse(const mjtNum* mat, int nr,
+                                  const int* rownnz, const int* rowadr,
+                                  const int* colind);
+
+
+    #--------------------- Components -----------------------------------------------------
+
+    # Run position-dependent computations.
+    void mj_fwdPosition(const mjModel* m, mjData* d);
+
+    # Run velocity-dependent computations.
+    void mj_fwdVelocity(const mjModel* m, mjData* d);
+
+    # Compute actuator force qfrc_actuation.
+    void mj_fwdActuation(const mjModel* m, mjData* d);
+
+    # Add up all non-constraint forces, compute qacc_unc.
+    void mj_fwdAcceleration(const mjModel* m, mjData* d);
+
+    # Run selected constraint solver.
+    void mj_fwdConstraint(const mjModel* m, mjData* d);
+
+    # Euler integrator, semi-implicit in velocity.
+    void mj_Euler(const mjModel* m, mjData* d);
+
+    # Runge-Kutta explicit order-N integrator.
+    void mj_RungeKutta(const mjModel* m, mjData* d, int N);
+
+    # Run position-dependent computations in inverse dynamics.
+    void mj_invPosition(const mjModel* m, mjData* d);
+
+    # Run velocity-dependent computations in inverse dynamics.
+    void mj_invVelocity(const mjModel* m, mjData* d);
+
+    # Apply the analytical formula for inverse constraint dynamics.
+    void mj_invConstraint(const mjModel* m, mjData* d);
+
+    # Compare forward and inverse dynamics, save results in fwdinv.
+    void mj_compareFwdInv(const mjModel* m, mjData* d);
+
+
+    #--------------------- Sub components -------------------------------------------------
+
+    # Evaluate position-dependent sensors.
+    void mj_sensorPos(const mjModel* m, mjData* d);
+
+    # Evaluate velocity-dependent sensors.
+    void mj_sensorVel(const mjModel* m, mjData* d);
+
+    # Evaluate acceleration and force-dependent sensors.
+    void mj_sensorAcc(const mjModel* m, mjData* d);
+
+    # Evaluate position-dependent energy (potential).
+    void mj_energyPos(const mjModel* m, mjData* d);
+
+    # Evaluate velocity-dependent energy (kinetic).
+    void mj_energyVel(const mjModel* m, mjData* d);
+
+    # Check qpos, reset if any element is too big or nan.
+    void mj_checkPos(const mjModel* m, mjData* d);
+
+    # Check qvel, reset if any element is too big or nan.
+    void mj_checkVel(const mjModel* m, mjData* d);
+
+    # Check qacc, reset if any element is too big or nan.
+    void mj_checkAcc(const mjModel* m, mjData* d);
+
+    # Run forward kinematics.
+    void mj_kinematics(const mjModel* m, mjData* d);
+
+    # Map inertias and motion dofs to global frame centered at CoM.
+    void mj_comPos(const mjModel* m, mjData* d);
+
+    # Compute camera and light positions and orientations.
+    void mj_camlight(const mjModel* m, mjData* d);
+
+    # Compute tendon lengths, velocities and moment arms.
+    void mj_tendon(const mjModel* m, mjData* d);
+
+    # Compute actuator transmission lengths and moments.
+    void mj_transmission(const mjModel* m, mjData* d);
+
+    # Run composite rigid body inertia algorithm (CRB).
+    void mj_crb(const mjModel* m, mjData* d);
+
+    # Compute sparse L'*D*L factorizaton of inertia matrix.
+    void mj_factorM(const mjModel* m, mjData* d);
+
+    # Solve linear system M * x = y using factorization:  x = inv(L'*D*L)*y
+    void mj_solveM(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
+
+    # Half of linear solve:  x = sqrt(inv(D))*inv(L')*y
+    void mj_solveM2(const mjModel* m, mjData* d, mjtNum* x, const mjtNum* y, int n);
+
+    # Compute cvel, cdof_dot.
+    void mj_comVel(const mjModel* m, mjData* d);
+
+    # Compute qfrc_passive from spring-dampers, viscosity and density.
+    void mj_passive(const mjModel* m, mjData* d);
+
+    # subtree linear velocity and angular momentum
+    void mj_subtreeVel(const mjModel* m, mjData* d);
+
+    # RNE: compute M(qpos)*qacc + C(qpos,qvel); flg_acc=0 removes inertial term.
+    void mj_rne(const mjModel* m, mjData* d, int flg_acc, mjtNum* result);
+
+    # RNE with complete data: compute cacc, cfrc_ext, cfrc_int.
+    void mj_rnePostConstraint(const mjModel* m, mjData* d);
+
+    # Run collision detection.
+    void mj_collision(const mjModel* m, mjData* d);
+
+    # Construct constraints.
+    void mj_makeConstraint(const mjModel* m, mjData* d);
+
+    # Compute inverse constaint inertia efc_AR.
+    void mj_projectConstraint(const mjModel* m, mjData* d);
+
+    # Compute efc_vel, efc_aref.
+    void mj_referenceConstraint(const mjModel* m, mjData* d);
+
+    # Compute efc_state, efc_force, qfrc_constraint, and (optionally) cone Hessians.
+    # If cost is not NULL, set *cost = s(jar) where jar = Jac*qacc-aref.
+    void mj_constraintUpdate(const mjModel* m, mjData* d, const mjtNum* jar,
+                                   mjtNum* cost, int flg_coneHessian);
+
+
+    #--------------------- Support --------------------------------------------------------
+
+    # Add contact to d->contact list; return 0 if success; 1 if buffer full.
+    int mj_addContact(const mjModel* m, mjData* d, const mjContact* con);
+
+    # Determine type of friction cone.
+    int mj_isPyramidal(const mjModel* m);
+
+    # Determine type of constraint Jacobian.
+    int mj_isSparse(const mjModel* m);
+
+    # Determine type of solver (PGS is dual, CG and Newton are primal).
+    int mj_isDual(const mjModel* m);
+
+    # Multiply dense or sparse constraint Jacobian by vector.
+    void mj_mulJacVec(const mjModel* m, mjData* d,
+                            mjtNum* res, const mjtNum* vec);
+
+    # Multiply dense or sparse constraint Jacobian transpose by vector.
+    void mj_mulJacTVec(const mjModel* m, mjData* d, mjtNum* res, const mjtNum* vec);
+
+    # Compute 3/6-by-nv end-effector Jacobian of global point attached to given body.
+    void mj_jac(const mjModel* m, const mjData* d,
+                      mjtNum* jacp, mjtNum* jacr, const mjtNum point[3], int body);
+
+    # Compute body frame end-effector Jacobian.
+    void mj_jacBody(const mjModel* m, const mjData* d,
+                          mjtNum* jacp, mjtNum* jacr, int body);
+
+    # Compute body center-of-mass end-effector Jacobian.
+    void mj_jacBodyCom(const mjModel* m, const mjData* d,
+                             mjtNum* jacp, mjtNum* jacr, int body);
+
+    # Compute geom end-effector Jacobian.
+    void mj_jacGeom(const mjModel* m, const mjData* d,
+                          mjtNum* jacp, mjtNum* jacr, int geom);
+
+    # Compute site end-effector Jacobian.
+    void mj_jacSite(const mjModel* m, const mjData* d,
+                          mjtNum* jacp, mjtNum* jacr, int site);
+
+    # Compute translation end-effector Jacobian of point, and rotation Jacobian of axis.
+    void mj_jacPointAxis(const mjModel* m, mjData* d,
+                               mjtNum* jacPoint, mjtNum* jacAxis,
+                               const mjtNum point[3], const mjtNum axis[3], int body);
+
+    # Get id of object with specified name, return -1 if not found; type is mjtObj.
+    int mj_name2id(const mjModel* m, int type, const char* name);
+
+    # Get name of object with specified id, return 0 if invalid type or id; type is mjtObj.
+    const char* mj_id2name(const mjModel* m, int type, int id);
+
+    # Convert sparse inertia matrix M into full (i.e. dense) matrix.
+    void mj_fullM(const mjModel* m, mjtNum* dst, const mjtNum* M);
+
+    # Multiply vector by inertia matrix.
+    void mj_mulM(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
+
+    # Multiply vector by (inertia matrix)^(1/2).
+    void mj_mulM2(const mjModel* m, const mjData* d, mjtNum* res, const mjtNum* vec);
+
+    # Add inertia matrix to destination matrix.
+    # Destination can be sparse uncompressed, or dense when all int* are NULL
+    void mj_addM(const mjModel* m, mjData* d, mjtNum* dst,
+                       int* rownnz, int* rowadr, int* colind);
+
+    # Apply cartesian force and torque (outside xfrc_applied mechanism).
+    void mj_applyFT(const mjModel* m, mjData* d,
+                          const mjtNum* force, const mjtNum* torque,
+                          const mjtNum* point, int body, mjtNum* qfrc_target);
+
+    # Compute object 6D velocity in object-centered frame, world/local orientation.
+    void mj_objectVelocity(const mjModel* m, const mjData* d,
+                                 int objtype, int objid, mjtNum* res, int flg_local);
+
+    # Compute object 6D acceleration in object-centered frame, world/local orientation.
+    void mj_objectAcceleration(const mjModel* m, const mjData* d,
+                                     int objtype, int objid, mjtNum* res, int flg_local);
+
+    # Extract 6D force:torque for one contact, in contact frame.
+    void mj_contactForce(const mjModel* m, const mjData* d, int id, mjtNum* result);
+
+    # Compute velocity by finite-differencing two positions.
+    void mj_differentiatePos(const mjModel* m, mjtNum* qvel, mjtNum dt,
+                             const mjtNum* qpos1, const mjtNum* qpos2);
+
+    # Integrate position with given velocity.
+    void mj_integratePos(const mjModel* m, mjtNum* qpos, const mjtNum* qvel, mjtNum dt);
+
+    # Normalize all quaterions in qpos-type vector.
+    void mj_normalizeQuat(const mjModel* m, mjtNum* qpos);
+
+    # Map from body local to global Cartesian coordinates.
+    void mj_local2Global(mjData* d, mjtNum* xpos, mjtNum* xmat, const mjtNum* pos, const mjtNum* quat,
+                         int body, mjtByte sameframe);
+
+    # Sum all body masses.
+    mjtNum mj_getTotalmass(const mjModel* m);
+
+    # Scale body masses and inertias to achieve specified total mass.
+    void mj_setTotalmass(mjModel* m, mjtNum newmass);
+
+    # Return version number: 1.0.2 is encoded as 102.
+    int mj_version();
+
+
+    #--------------------- Ray collisions -------------------------------------------------
+
+    # Intersect ray (pnt+x*vec, x>=0) with visible geoms, except geoms in bodyexclude.
+    # Return geomid and distance (x) to nearest surface, or -1 if no intersection.
+    # geomgroup, flg_static are as in mjvOption; geomgroup==NULL skips group exclusion.
+    mjtNum mj_ray(const mjModel* m, const mjData* d, const mjtNum* pnt, const mjtNum* vec,
+                        const mjtByte* geomgroup, mjtByte flg_static, int bodyexclude,
+                        int* geomid);
+
+    # Interect ray with hfield, return nearest distance or -1 if no intersection.
+    mjtNum mj_rayHfield(const mjModel* m, const mjData* d, int geomid,
+                              const mjtNum* pnt, const mjtNum* vec);
+
+    # Interect ray with mesh, return nearest distance or -1 if no intersection.
+    mjtNum mj_rayMesh(const mjModel* m, const mjData* d, int geomid,
+                            const mjtNum* pnt, const mjtNum* vec);
+
+    # Interect ray with pure geom, return nearest distance or -1 if no intersection.
+    mjtNum mju_rayGeom(const mjtNum* pos, const mjtNum* mat, const mjtNum* size,
+                             const mjtNum* pnt, const mjtNum* vec, int geomtype);
+
+
+    #--------------------- Interaction ----------------------------------------------------
+
+    # Set default camera.
+    void mjv_defaultCamera(mjvCamera* cam);
+
+    # Set default perturbation.
+    void mjv_defaultPerturb(mjvPerturb* pert);
+
+    # Transform pose from room to model space.
+    void mjv_room2model(mjtNum* modelpos, mjtNum* modelquat, const mjtNum* roompos,
+                              const mjtNum* roomquat, const mjvScene* scn);
+
+    # Transform pose from model to room space.
+    void mjv_model2room(mjtNum* roompos, mjtNum* roomquat, const mjtNum* modelpos,
+                              const mjtNum* modelquat, const mjvScene* scn);
+
+    # Get camera info in model space; average left and right OpenGL cameras.
+    void mjv_cameraInModel(mjtNum* headpos, mjtNum* forward, mjtNum* up,
+                                 const mjvScene* scn);
+
+    # Get camera info in room space; average left and right OpenGL cameras.
+    void mjv_cameraInRoom(mjtNum* headpos, mjtNum* forward, mjtNum* up,
+                                const mjvScene* scn);
+
+    # Get frustum height at unit distance from camera; average left and right OpenGL cameras.
+    mjtNum mjv_frustumHeight(const mjvScene* scn);
+
+    # Rotate 3D vec in horizontal plane by angle between (0,1) and (forward_x,forward_y).
+    void mjv_alignToCamera(mjtNum* res, const mjtNum* vec, const mjtNum* forward);
+
+    # Move camera with mouse; action is mjtMouse.
+    void mjv_moveCamera(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
+                              const mjvScene* scn, mjvCamera* cam);
+
+    # Move perturb object with mouse; action is mjtMouse.
+    void mjv_movePerturb(const mjModel* m, const mjData* d, int action, mjtNum reldx,
+                               mjtNum reldy, const mjvScene* scn, mjvPerturb* pert);
+
+    # Move model with mouse; action is mjtMouse.
+    void mjv_moveModel(const mjModel* m, int action, mjtNum reldx, mjtNum reldy,
+                             const mjtNum* roomup, mjvScene* scn);
+
+    # Copy perturb pos,quat from selected body; set scale for perturbation.
+    void mjv_initPerturb(const mjModel* m, const mjData* d,
+                               const mjvScene* scn, mjvPerturb* pert);
+
+    # Set perturb pos,quat in d->mocap when selected body is mocap, and in d->qpos otherwise.
+    # Write d->qpos only if flg_paused and subtree root for selected body has free joint.
+    void mjv_applyPerturbPose(const mjModel* m, mjData* d, const mjvPerturb* pert,
+                                    int flg_paused);
+
+    # Set perturb force,torque in d->xfrc_applied, if selected body is dynamic.
+    void mjv_applyPerturbForce(const mjModel* m, mjData* d, const mjvPerturb* pert);
+
+    # Return the average of two OpenGL cameras.
+    mjvGLCamera mjv_averageCamera(const mjvGLCamera* cam1, const mjvGLCamera* cam2);
+
+    # Select geom or skin with mouse, return bodyid; -1: none selected.
+    int mjv_select(const mjModel* m, const mjData* d, const mjvOption* vopt,
+                   mjtNum aspectratio, mjtNum relx, mjtNum rely,
+                   const mjvScene* scn, mjtNum* selpnt, int* geomid, int* skinid);
+
+    #--------------------- Visualization --------------------------------------------------
+
+    # Set default visualization options.
+    void mjv_defaultOption(mjvOption* opt);
+
+    # Set default figure.
+    void mjv_defaultFigure(mjvFigure* fig);
+
+    # Initialize given geom fields when not NULL, set the rest to their default values.
+    void mjv_initGeom(mjvGeom* geom, int type, const mjtNum* size,
+                            const mjtNum* pos, const mjtNum* mat, const float* rgba);
+
+    # Set (type, size, pos, mat) for connector-type geom between given points.
+    # Assume that mjv_initGeom was already called to set all other properties.
+    void mjv_makeConnector(mjvGeom* geom, int type, mjtNum width,
+                                 mjtNum a0, mjtNum a1, mjtNum a2,
+                                 mjtNum b0, mjtNum b1, mjtNum b2);
+
+    # Set default abstract scene.
+    void mjv_defaultScene(mjvScene* scn);
+
+    # Allocate resources in abstract scene.
+    void mjv_makeScene(const mjModel* m, mjvScene* scn, int maxgeom);
+
+    # Free abstract scene.
+    void mjv_freeScene(mjvScene* scn);
+
+    # Update entire scene given model state.
+    void mjv_updateScene(const mjModel* m, mjData* d, const mjvOption* opt,
+                               const mjvPerturb* pert, mjvCamera* cam, int catmask, mjvScene* scn);
+
+    # Add geoms from selected categories to existing scene.
+    void mjv_addGeoms(const mjModel* m, mjData* d, const mjvOption* opt,
+                            const mjvPerturb* pert, int catmask, mjvScene* scn);
+
+    # Make list of lights.
+    void mjv_makeLights(const mjModel* m, mjData* d, mjvScene* scn);
+
+    # Update camera only.
+    void mjv_updateCamera(const mjModel* m, mjData* d, mjvCamera* cam, mjvScene* scn);
+
+    # Update skins.
+    void mjv_updateSkin(const mjModel* m, mjData* d, mjvScene* scn);
+
+    #--------------------- OpenGL rendering -----------------------------------------------
+
+    # Set default mjrContext.
+    void mjr_defaultContext(mjrContext* con);
+
+    # Allocate resources in custom OpenGL context; fontscale is mjtFontScale.
+    void mjr_makeContext(const mjModel* m, mjrContext* con, int fontscale);
+
+    # Change font of existing context.
+    void mjr_changeFont(int fontscale, mjrContext* con);
+
+    # Add Aux buffer with given index to context; free previous Aux buffer.
+    void mjr_addAux(int index, int width, int height, int samples, mjrContext* con);
+
+    # Free resources in custom OpenGL context, set to default.
+    void mjr_freeContext(mjrContext* con);
+
+    # Upload texture to GPU, overwriting previous upload if any.
+    void mjr_uploadTexture(const mjModel* m, const mjrContext* con, int texid);
+
+    # Upload mesh to GPU, overwriting previous upload if any.
+    void mjr_uploadMesh(const mjModel* m, const mjrContext* con, int meshid);
+
+    # Upload height field to GPU, overwriting previous upload if any.
+    void mjr_uploadHField(const mjModel* m, const mjrContext* con, int hfieldid);
+
+    # Make con->currentBuffer current again.
+    void mjr_restoreBuffer(const mjrContext* con);
+
+    # Set OpenGL framebuffer for rendering: mjFB_WINDOW or mjFB_OFFSCREEN.
+    # If only one buffer is available, set that buffer and ignore framebuffer argument.
+    void mjr_setBuffer(int framebuffer, mjrContext* con);
+
+    # Read pixels from current OpenGL framebuffer to client buffer.
+    # Viewport is in OpenGL framebuffer; client buffer starts at (0,0).
+    void mjr_readPixels(unsigned char* rgb, float* depth,
+                              mjrRect viewport, const mjrContext* con);
+
+    # Draw pixels from client buffer to current OpenGL framebuffer.
+    # Viewport is in OpenGL framebuffer; client buffer starts at (0,0).
+    void mjr_drawPixels(const unsigned char* rgb, const float* depth,
+                              mjrRect viewport, const mjrContext* con);
+
+    # Blit from src viewpoint in current framebuffer to dst viewport in other framebuffer.
+    # If src, dst have different size and flg_depth==0, color is interpolated with GL_LINEAR.
+    void mjr_blitBuffer(mjrRect src, mjrRect dst, int flg_color, int flg_depth, const mjrContext* con);
+
+    # Set Aux buffer for custom OpenGL rendering (call restoreBuffer when done).
+    void mjr_setAux(int index, const mjrContext* con);
+
+    # Blit from Aux buffer to con->currentBuffer.
+    void mjr_blitAux(int index, mjrRect src, int left, int bottom, const mjrContext* con);
+
+    # Draw text at (x,y) in relative coordinates; font is mjtFont.
+    void mjr_text(int font, const char* txt, const mjrContext* con,
+                        float x, float y, float r, float g, float b);
+
+    # Draw text overlay; font is mjtFont; gridpos is mjtGridPos.
+    void mjr_overlay(int font, int gridpos, mjrRect viewport,
+                           const char* overlay, const char* overlay2, const mjrContext* con);
+
+    # Get maximum viewport for active buffer.
+    mjrRect mjr_maxViewport(const mjrContext* con);
+
+    # Draw rectangle.
+    void mjr_rectangle(mjrRect viewport, float r, float g, float b, float a);
+
+    # Draw rectangle with centered text.
+    void mjr_label(mjrRect viewport, int font, const char* txt,
+                   float r, float g, float b, float a, float rt, float gt, float bt,
+				           const mjrContext* con);
+
+    # Draw 2D figure.
+    void mjr_figure(mjrRect viewport, const mjvFigure* fig, const mjrContext* con);
+
+    # Render 3D scene.
+    void mjr_render(mjrRect viewport, mjvScene* scn, const mjrContext* con);
+
+    # Call glFinish.
+    void mjr_finish();
+
+    # Call glGetError and return result.
+    int mjr_getError();
+
+    # Find first rectangle containing mouse, -1: not found.
+    int mjr_findRect(int x, int y, int nrect, const mjrRect* rect);
+
+    #---------------------- UI framework ---------------------------------------------------
+
+    # Add definitions to UI.
+    void mjui_add(mjUI* ui, const mjuiDef* _def);
+
+    # Add definitions to UI section.
+    void mjui_addToSection(mjUI* ui, int sect, const mjuiDef* _def);
+
+
+    # Compute UI sizes.
+    void mjui_resize(mjUI* ui, const mjrContext* con);
+
+    # Update specific section/item; -1: update all.
+    void mjui_update(int section, int item, const mjUI* ui, const mjuiState* state, const mjrContext* con);
+
+    # Handle UI event, return pointer to changed item, NULL if no change.
+    mjuiItem* mjui_event(mjUI* ui, mjuiState* state, const mjrContext* con);
+
+    # Copy UI image to current buffer.
+    void mjui_render(mjUI* ui, const mjuiState* state, const mjrContext* con);
+
+
+    #--------------------- Error and memory -----------------------------------------------
+
+    # Main error function; does not return to caller.
+    void mju_error(const char* msg);
+
+    # Error function with int argument; msg is a printf format string.
+    void mju_error_i(const char* msg, int i);
+
+    # Error function with string argument.
+    void mju_error_s(const char* msg, const char* text);
+
+    # Main warning function; returns to caller.
+    void mju_warning(const char* msg);
+
+    # Warning function with int argument.
+    void mju_warning_i(const char* msg, int i);
+
+    # Warning function with string argument.
+    void mju_warning_s(const char* msg, const char* text);
+
+    # Clear user error and memory handlers.
+    void mju_clearHandlers();
+
+    # Allocate memory; byte-align on 8; pad size to multiple of 8.
+    void* mju_malloc(size_t size);
+
+    # Free memory, using free() by default.
+    void mju_free(void* ptr);
+
+    # High-level warning function: count warnings in mjData, print only the first.
+    void mj_warning(mjData* d, int warning, int info);
+
+    # Write [datetime, type: message] to MUJOCO_LOG.TXT.
+    void mju_writeLog(const char* type, const char* msg);
+
+
+    #--------------------- Standard math --------------------------------------------------
+
+#define mjMAX(a,b) (((a) > (b)) ? (a) : (b))
+#define mjMIN(a,b) (((a) < (b)) ? (a) : (b))
+
+#ifdef mjUSEDOUBLE
+        #define mju_sqrt    sqrt
+        #define mju_exp     exp
+        #define mju_sin     sin
+        #define mju_cos     cos
+        #define mju_tan     tan
+        #define mju_asin    asin
+        #define mju_acos    acos
+        #define mju_atan2   atan2
+        #define mju_tanh    tanh
+        #define mju_pow     pow
+        #define mju_abs     fabs
+        #define mju_log     log
+        #define mju_log10   log10
+        #define mju_floor   floor
+        #define mju_ceil    ceil
+
+#else
+        #define mju_sqrt    sqrtf
+        #define mju_exp     expf
+        #define mju_sin     sinf
+        #define mju_cos     cosf
+        #define mju_tan     tanf
+        #define mju_asin    asinf
+        #define mju_acos    acosf
+        #define mju_atan2   atan2f
+        #define mju_tanh    tanhf
+        #define mju_pow     powf
+        #define mju_abs     fabsf
+        #define mju_log     logf
+        #define mju_log10   log10f
+        #define mju_floor   floorf
+        #define mju_ceil    ceilf
+#endif
+
+
+    #----------------------------- Vector math --------------------------------------------
+
+    # Set res = 0.
+    void mju_zero3(mjtNum res[3]);
+
+    # Set res = vec.
+    void mju_copy3(mjtNum res[3], const mjtNum data[3]);
+
+    # Set res = vec*scl.
+    void mju_scl3(mjtNum res[3], const mjtNum vec[3], mjtNum scl);
+
+    # Set res = vec1 + vec2.
+    void mju_add3(mjtNum res[3], const mjtNum vec1[3], const mjtNum vec2[3]);
+
+    # Set res = vec1 - vec2.
+    void mju_sub3(mjtNum res[3], const mjtNum vec1[3], const mjtNum vec2[3]);
+
+    # Set res = res + vec.
+    void mju_addTo3(mjtNum res[3], const mjtNum vec[3]);
+
+    # Set res = res - vec.
+    void mju_subFrom3(mjtNum res[3], const mjtNum vec[3]);
+
+    # Set res = res + vec*scl.
+    void mju_addToScl3(mjtNum res[3], const mjtNum vec[3], mjtNum scl);
+
+    # Set res = vec1 + vec2*scl.
+    void mju_addScl3(mjtNum res[3], const mjtNum vec1[3], const mjtNum vec2[3], mjtNum scl);
+
+    # Normalize vector, return length before normalization.
+    mjtNum mju_normalize3(mjtNum res[3]);
+
+    # Return vector length (without normalizing the vector).
+    mjtNum mju_norm3(const mjtNum vec[3]);
+
+    # Return dot-product of vec1 and vec2.
+    mjtNum mju_dot3(const mjtNum vec1[3], const mjtNum vec2[3]);
+
+    # Return Cartesian distance between 3D vectors pos1 and pos2.
+    mjtNum mju_dist3(const mjtNum pos1[3], const mjtNum pos2[3]);
+
+    # Multiply vector by 3D rotation matrix: res = mat * vec.
+    void mju_rotVecMat(mjtNum res[3], const mjtNum vec[3], const mjtNum mat[9]);
+
+    # Multiply vector by transposed 3D rotation matrix: res = mat' * vec.
+    void mju_rotVecMatT(mjtNum res[3], const mjtNum vec[3], const mjtNum mat[9]);
+
+    # Compute cross-product: res = cross(a, b).
+    void mju_cross(mjtNum res[3], const mjtNum a[3], const mjtNum b[3]);
+
+    # Set res = 0.
+    void mju_zero4(mjtNum res[4]);
+
+    # Set res = (1,0,0,0).
+    void mju_unit4(mjtNum res[4]);
+
+    # Set res = vec.
+    void mju_copy4(mjtNum res[4], const mjtNum data[4]);
+
+    # Normalize vector, return length before normalization.
+    mjtNum mju_normalize4(mjtNum res[4]);
+
+    # Set res = 0.
+    void mju_zero(mjtNum* res, int n);
+
+    # Set res = vec.
+    void mju_copy(mjtNum* res, const mjtNum* data, int n);
+
+    # Return sum(vec).
+    mjtNum mju_sum(const mjtNum* vec, int n);
+
+    # Return L1 norm: sum(abs(vec)).
+    mjtNum mju_L1(const mjtNum* vec, int n);
+
+    # Set res = vec*scl.
+    void mju_scl(mjtNum* res, const mjtNum* vec, mjtNum scl, int n);
+
+    # Set res = vec1 + vec2.
+    void mju_add(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, int n);
+
+    # Set res = vec1 - vec2.
+    void mju_sub(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, int n);
+
+    # Set res = res + vec.
+    void mju_addTo(mjtNum* res, const mjtNum* vec, int n);
+
+    # Set res = res - vec.
+    void mju_subFrom(mjtNum* res, const mjtNum* vec, int n);
+
+    # Set res = res + vec*scl.
+    void mju_addToScl(mjtNum* res, const mjtNum* vec, mjtNum scl, int n);
+
+    # Set res = vec1 + vec2*scl.
+    void mju_addScl(mjtNum* res, const mjtNum* vec1, const mjtNum* vec2, mjtNum scl, int n);
+
+    # Normalize vector, return length before normalization.
+    mjtNum mju_normalize(mjtNum* res, int n);
+
+    # Return vector length (without normalizing vector).
+    mjtNum mju_norm(const mjtNum* res, int n);
+
+    # Return dot-product of vec1 and vec2.
+    mjtNum mju_dot(const mjtNum* vec1, const mjtNum* vec2, const int n);
+
+    # Multiply matrix and vector: res = mat * vec.
+    void mju_mulMatVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
+                             int nr, int nc);
+
+    # Multiply transposed matrix and vector: res = mat' * vec.
+    void mju_mulMatTVec(mjtNum* res, const mjtNum* mat, const mjtNum* vec,
+                              int nr, int nc);
+
+    # Transpose matrix: res = mat'.
+    void mju_transpose(mjtNum* res, const mjtNum* mat, int nr, int nc);
+
+    # Multiply matrices: res = mat1 * mat2.
+    void mju_mulMatMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
+                             int r1, int c1, int c2);
+
+    # Multiply matrices, second argument transposed: res = mat1 * mat2'.
+    void mju_mulMatMatT(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
+                              int r1, int c1, int r2);
+
+    # Multiply matrices, first argument transposed: res = mat1' * mat2.
+    void mju_mulMatTMat(mjtNum* res, const mjtNum* mat1, const mjtNum* mat2,
+                              int r1, int c1, int c2);
+
+    # Set res = mat' * diag * mat if diag is not NULL, and res = mat' * mat otherwise.
+    void mju_sqrMatTD(mjtNum* res, const mjtNum* mat, const mjtNum* diag, int nr, int nc);
+
+    # Coordinate transform of 6D motion or force vector in rotation:translation format.
+    # rotnew2old is 3-by-3, NULL means no rotation; flg_force specifies force or motion type.
+    void mju_transformSpatial(mjtNum res[6], const mjtNum vec[6], int flg_force,
+                                    const mjtNum newpos[3], const mjtNum oldpos[3],
+                                    const mjtNum rotnew2old[9]);
+
+
+    #--------------------- Sparse math ----------------------------------------------------
+
+    # Return dot-product of vec1 and vec2, where vec1 is sparse.
+    mjtNum mju_dotSparse(const mjtNum* vec1, const mjtNum* vec2,
+                               const int nnz1, const int* ind1);
+
+    # Return dot-product of vec1 and vec2, where both vectors are sparse.
+    mjtNum mju_dotSparse2(const mjtNum* vec1, const mjtNum* vec2,
+                                const int nnz1, const int* ind1,
+                                const int nnz2, const int* ind2);
+
+    # Convert matrix from dense to sparse format.
+    void mju_dense2sparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
+                                int* rownnz, int* rowadr, int* colind);
+
+    # Convert matrix from sparse to dense format.
+    void mju_sparse2dense(mjtNum* res, const mjtNum* mat, int nr, int nc,
+                                const int* rownnz, const int* rowadr, const int* colind);
+
+    # Multiply sparse matrix and dense vector:  res = mat * vec.
+    void mju_mulMatVecSparse(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int nr,
+                                   const int* rownnz, const int* rowadr, const int* colind);
+
+    # Compress layout of sparse matrix.
+    void mju_compressSparse(mjtNum* mat, int nr, int nc,
+                                  int* rownnz, int* rowadr, int* colind);
+
+    # Set dst = a*dst + b*src, return nnz of result, modify dst sparsity pattern as needed.
+    # Both vectors are sparse. The required scratch space is 2*n.
+    int mju_combineSparse(mjtNum* dst, const mjtNum* src, int n, mjtNum a, mjtNum b,
+                                int dst_nnz, int src_nnz, int* dst_ind, const int* src_ind,
+                                mjtNum* scratch, int nscratch);
+
+    # Set res = matT * diag * mat if diag is not NULL, and res = matT * mat otherwise.
+    # The required scratch space is 3*nc. The result has uncompressed layout.
+    void mju_sqrMatTDSparse(mjtNum* res, const mjtNum* mat, const mjtNum* matT,
+                                  const mjtNum* diag, int nr, int nc,
+                                  int* res_rownnz, int* res_rowadr, int* res_colind,
+                                  const int* rownnz, const int* rowadr, const int* colind,
+                                  const int* rownnzT, const int* rowadrT, const int* colindT,
+                                  mjtNum* scratch, int nscratch);
+
+    # Transpose sparse matrix.
+    void mju_transposeSparse(mjtNum* res, const mjtNum* mat, int nr, int nc,
+                                  int* res_rownnz, int* res_rowadr, int* res_colind,
+                                  const int* rownnz, const int* rowadr, const int* colind);
+
+
+    #--------------------- Quaternions ----------------------------------------------------
+
+    # Rotate vector by quaternion.
+    void mju_rotVecQuat(mjtNum res[3], const mjtNum vec[3], const mjtNum quat[4]);
+
+    # Negate quaternion.
+    void mju_negQuat(mjtNum res[4], const mjtNum quat[4]);
+
+    # Muiltiply quaternions.
+    void mju_mulQuat(mjtNum res[4], const mjtNum quat1[4], const mjtNum quat2[4]);
+
+    # Muiltiply quaternion and axis.
+    void mju_mulQuatAxis(mjtNum res[4], const mjtNum quat[4], const mjtNum axis[3]);
+
+    # Convert axisAngle to quaternion.
+    void mju_axisAngle2Quat(mjtNum res[4], const mjtNum axis[3], mjtNum angle);
+
+    # Convert quaternion (corresponding to orientation difference) to 3D velocity.
+    void mju_quat2Vel(mjtNum res[3], const mjtNum quat[4], mjtNum dt);
+
+    # Subtract quaternions, express as 3D velocity: qb*quat(res) = qa.
+    void mju_subQuat(mjtNum res[3], const mjtNum qa[4], const mjtNum qb[4]);
+
+    # Convert quaternion to 3D rotation matrix.
+    void mju_quat2Mat(mjtNum res[9], const mjtNum quat[4]);
+
+    # Convert 3D rotation matrix to quaterion.
+    void mju_mat2Quat(mjtNum quat[4], const mjtNum mat[9]);
+
+    # Compute time-derivative of quaternion, given 3D rotational velocity.
+    void mju_derivQuat(mjtNum res[4], const mjtNum quat[4], const mjtNum vel[3]);
+
+    # Integrate quaterion given 3D angular velocity.
+    void mju_quatIntegrate(mjtNum quat[4], const mjtNum vel[3], mjtNum scale);
+
+    # Construct quaternion performing rotation from z-axis to given vector.
+    void mju_quatZ2Vec(mjtNum quat[4], const mjtNum vec[3]);
+
+
+    #--------------------- Poses ----------------------------------------------------------
+
+    # Multiply two poses.
+    void mju_mulPose(mjtNum posres[3], mjtNum quatres[4],
+                           const mjtNum pos1[3], const mjtNum quat1[4],
+                           const mjtNum pos2[3], const mjtNum quat2[4]);
+
+    # Negate pose.
+    void mju_negPose(mjtNum posres[3], mjtNum quatres[4],
+                           const mjtNum pos[3], const mjtNum quat[4]);
+
+    # Transform vector by pose.
+    void mju_trnVecPose(mjtNum res[3], const mjtNum pos[3], const mjtNum quat[4],
+                              const mjtNum vec[3]);
+
+
+    #--------------------- Decompositions --------------------------------------------------
+
+    # Cholesky decomposition: mat = L*L'; return rank.
+    int mju_cholFactor(mjtNum* mat, int n, mjtNum mindiag);
+
+    # Solve mat * res = vec, where mat is Cholesky-factorized
+    void mju_cholSolve(mjtNum* res, const mjtNum* mat, const mjtNum* vec, int n);
+
+    # Cholesky rank-one update: L*L' +/- x*x'; return rank.
+    int mju_cholUpdate(mjtNum* mat, mjtNum* x, int n, int flg_plus);
+
+    # Eigenvalue decomposition of symmetric 3x3 matrix.
+    int mju_eig3(mjtNum* eigval, mjtNum* eigvec, mjtNum* quat, const mjtNum* mat);
+
+
+    #--------------------- Miscellaneous --------------------------------------------------
+
+    # Muscle active force, prm = (range[2], force, scale, lmin, lmax, vmax, fpmax, fvmax).
+    mjtNum mju_muscleGain(mjtNum len, mjtNum vel, const mjtNum lengthrange[2],
+                          mjtNum acc0, const mjtNum prm[9]);
+
+    # Muscle passive force, prm = (range[2], force, scale, lmin, lmax, vmax, fpmax, fvmax).
+    mjtNum mju_muscleBias(mjtNum len, const mjtNum lengthrange[2],
+                          mjtNum acc0, const mjtNum prm[9]);
+
+    # Muscle activation dynamics, prm = (tau_act, tau_deact).
+    mjtNum mju_muscleDynamics(mjtNum ctrl, mjtNum act, const mjtNum prm[2]);
+
+
+    # Convert contact force to pyramid representation.
+    void mju_encodePyramid(mjtNum* pyramid, const mjtNum* force,
+                                 const mjtNum* mu, int dim);
+
+    # Convert pyramid representation to contact force.
+    void mju_decodePyramid(mjtNum* force, const mjtNum* pyramid,
+                                 const mjtNum* mu, int dim);
+
+    # Integrate spring-damper analytically, return pos(dt).
+    mjtNum mju_springDamper(mjtNum pos0, mjtNum vel0, mjtNum Kp, mjtNum Kv, mjtNum dt);
+
+    # Return min(a,b) with single evaluation of a and b.
+    mjtNum mju_min(mjtNum a, mjtNum b);
+
+    # Return max(a,b) with single evaluation of a and b.
+    mjtNum mju_max(mjtNum a, mjtNum b);
+
+    # Return sign of x: +1, -1 or 0.
+    mjtNum mju_sign(mjtNum x);
+
+    # Round x to nearest integer.
+    int mju_round(mjtNum x);
+
+    # Convert type id (mjtObj) to type name.
+    const char* mju_type2Str(int type);
+
+    # Convert type name to type id (mjtObj).
+    int mju_str2Type(const char* str);
+
+    # Construct a warning message given the warning type and info.
+    const char* mju_warningText(int warning, int info);
+
+    # Return 1 if nan or abs(x)>mjMAXVAL, 0 otherwise. Used by check functions.
+    int mju_isBad(mjtNum x);
+
+    # Return 1 if all elements are 0.
+    int mju_isZero(mjtNum* vec, int n);
+
+    # Standard normal random number generator (optional second number).
+    mjtNum mju_standardNormal(mjtNum* num2);
+
+    # Convert from float to mjtNum.
+    void mju_f2n(mjtNum* res, const float* vec, int n);
+
+    # Convert from mjtNum to float.
+    void mju_n2f(float* res, const mjtNum* vec, int n);
+
+    # Convert from double to mjtNum.
+    void mju_d2n(mjtNum* res, const double* vec, int n);
+
+    # Convert from mjtNum to double.
+    void mju_n2d(double* res, const mjtNum* vec, int n);
+
+    # Insertion sort, resulting list is in increasing order.
+    void mju_insertionSort(mjtNum* list, int n);
+
+    # Integer insertion sort, resulting list is in increasing order.
+    void mju_insertionSortInt(int* list, int n);
+
+    # Generate Halton sequence.
+    mjtNum mju_Halton(int index, int base);
+
+    # Sigmoid function over 0<=x<=1 constructed from half-quadratics.
+    mjtNum mju_sigmoid(mjtNum x);
+