Add files using upload-large-folder tool

applications/avoiding/avoiding_baselines/DiRecT/.gitignore

@@ -0,0 +1,28 @@
+# Caches
+__pycache__/
+*.pyc
+*.pyo
+
+# Python packaging
+*.egg-info/
+*.so
+*.o
+
+# Experiments, logs, and run outputs
+logs/
+slurm_logs/
+
+# Mujoco
+MUJOCO_LOG.TXT
+panda_tmp_*.xml
+
+# Notebooks
+.ipynb_checkpoints/
+*.ipynb
+
+# Slurm
+*.out
+*.err
+
+# Wandb
+wandb/

applications/avoiding/avoiding_baselines/DiRecT/README.md

@@ -0,0 +1,155 @@
+# Safe robotic manipulation in D3IL avoiding
+
+This folder presents the benchmarks for _DiRecT_ and relevant baselines on the [D3IL](https://github.com/ALRhub/d3il) `avoiding-v0` navigation task with dynamics constraints. The agent must navigate upward through 6 circular pillar obstacles to reach a goal region while satisfying linearized dynamics and obstacle-avoidance constraints.
+![D3IL avoiding-v0 visualization](assets/d3il-grid.gif)
+
+**Table of Contents**:
+- [Installation](#installation)
+- [Fit the linearized dynamics](#fit-the-linearized-dynamics)
+- [Train the model](#train-the-model)
+- [Evaluation](#evaluation)
+- [Results](#results)
+- [Code structure](#code-structure)
+- [References](#references)
+
+## Installation
+> [!NOTE]
+> The `avoiding-v0` environment uses MuJoCo 2.3.x via the `mujoco` Python bindings — no separate system-wide MuJoCo install is required. CasADi / IPOPT and Pinocchio are pulled in through conda-forge.
+
+1. Change the CUDA and torch versions in `environment.yml` to ones compatible with your system requirements (default is CUDA 12.6 + torch 2.11).
+2. Create and activate the conda environment. The provided `setup_env.sh` script handles environment creation, the `qpth` install (which needs `--no-build-isolation` due to a broken numpy specifier), and the `LD_LIBRARY_PATH` activation hook:
+```bash
+bash setup_env.sh
+conda activate direct-d3il
+```
+
+## Fit the linearized dynamics
+Fit a linear dynamics model `s_{t+1} = A·s_t + B·a_t + c` from the D3IL `avoiding-v0` expert trajectories (96 episodes). The fitted matrices are later imposed as linear constraints during planning.
+
+```bash
+python fit_dynamics.py \
+    --env avoiding-v0                # gym env id (defaults to avoiding-v0)
+```
+Outputs (written to `logs/avoiding-v0/dynamics/`):
+- `linear_model.npz` — fitted `A`, `B`, `c`, metadata, and normalizer (consumed by `eval.py`).
+- `linear_model_predictions.png` — per-dimension true-vs-predicted scatter plot.
+- `dynamics_verification.png` — single-step comparison against the real environment.
+- `all_trajectories.png` — overlay of the expert dataset.
+
+## Train the model
+Train the unconstrained diffusion planner (Temporal U-Net, x_0-prediction parameterization with a cosine noise schedule) on the same D3IL dataset. Weights & Biases is used for logging (project `direct-d3il`).
+
+```bash
+python train.py \
+    --exp-name x0pred \              # subfolder under <log-folder>/<env>/diffusion/
+    --seed 0 \
+    --batch-size 32 \                # overrides default 32
+    --n-train-steps 1000001          # default 1M steps
+```
+
+Resume from the latest checkpoint with `--resume`. Checkpoints (raw + EMA weights), sample previews, and the resolved config are written to `logs/avoiding-v0/diffusion/<exp-name>/`:
+- `model_*.pth`, `model_ema_*.pth` — saved every `--save-freq` steps (default 50k).
+- `samples/step_*.png` — periodic sample previews from the EMA model.
+- `tensorboard_logs/` — TensorBoard event files.
+- `config.yaml` — snapshot of the fully-resolved config.
+
+To submit as a SLURM job (edit the `#SBATCH` directives and conda env name first):
+```bash
+sbatch run.sbatch
+```
+
+## Evaluation
+Evaluate any of the planning policies on the constrained `avoiding-v0` task. The trained diffusion weights and fitted dynamics are both required.
+
+```bash
+python eval.py \
+    --diffusion-exp-name x0pred \    # which trained model to load
+    --diffusion-cp 10 \              # checkpoint index (model_ema_<cp>.pth)
+    --guidance-method direct \       # one of the methods listed below
+    --constraints \                  # enable the obstacle constraint set
+    --random-repeat 100 \            # number of evaluation rollouts
+    --exp-name direct \              # subfolder under <log-folder>/<env>/eval/
+    --seed 0 \
+    --device cuda:0
+```
+
+Available planning policies (selected via `--guidance-method`): `no` (unconditional), `gradient`, `classifier`, `projection`, `safediffuser` (with `--cbf-algorithm {RoS,ReS,TVS}`), `primal-dual`, `augmented-lagrangian`, `direct`. The DiRecT method additionally supports `--pred-objective distance` to add a goal-distance objective in clean x_0 space.
+
+To run the full sweep over every policy (dynamics constraint always on, 100 episodes each), use the provided SLURM script (edit the `#SBATCH` directives and conda env name first):
+```bash
+sbatch eval_all.sbatch
+```
+
+## Results
+For each evaluation, results are written under `logs/avoiding-v0/eval/<exp-name>/`:
+- `config.yaml` — resolved config used for the run.
+- `runs.csv` — one row per rollout with per-run metrics (success, safety, steps, total constraint violations, total reward, average computation time per replan).
+- `summary.json` — aggregated statistics (success rate, safety rate, average steps / violations / rewards / compute time).
+- `all_real.png` — overlay of every closed-loop rollout on the obstacle field.
+- `predictions/ep_*.png` — per-episode plot overlaying the diffused plans and the executed rollout.
+- `trajectories.npz` — raw rollout state arrays (suppressed by `--no-trajectories`).
+
+## Code structure
+```
+d3il/
+├── d3il/                       # ported D3IL `avoiding-v0` environment + expert data (96 episodes)
+│   └── environments/
+├── datasets/
+│   ├── sequence.py             # `SequenceDataset` (H=16 windows, LimitsNormalizer)
+│   ├── buffer.py               # replay buffer over expert episodes
+│   ├── normalization.py        # observation / action normalizers
+│   └── preprocessing.py
+├── policy/                     # planning policies (one file per method)
+│   ├── unconditional.py        # baseline DDIM sampling
+│   ├── gradient_guidance.py    # DPS-style posterior sampling
+│   ├── classifier_guidance.py  # classifier guidance on noisy trajectory
+│   ├── projected_diffusion.py  # IPOPT projection per denoising step (noisy space)
+│   ├── safediffuser.py         # CBF-guided QP (RoS / ReS / TVS)
+│   ├── augmented_lagrangian.py # primal-dual + augmented Lagrangian
+│   ├── direct.py               # DiRecT: IPOPT projection in clean x_0 space
+│   └── value_model.py          # differentiable proxy value function (obstacles + dynamics)
+├── model.py                    # `TemporalUnet` + `CosineScheduleDiffusion` (x_0 parameterization)
+├── projector.py                # `IpoptProjector` — CasADi/IPOPT constraint projection
+├── obstacles.py                # circular & planar obstacles, constraint sets, violation checks
+├── fit_dynamics.py             # fit `s_{t+1} = A s_t + B a_t + c` from offline data
+├── train.py                    # train the diffusion planner
+├── eval.py                     # rollout + metric aggregation for a chosen policy
+├── utils/                      # rendering, array helpers, training utilities
+├── run.sbatch                  # SLURM script for training
+├── eval_all.sbatch             # SLURM script sweeping all 11 policies
+├── setup_env.sh                # conda env bootstrap (incl. qpth + LD_LIBRARY_PATH hook)
+└── environment.yml             # conda environment spec
+```
+
+## ✏️ Citation
+If you find our code or paper useful for your research, please consider citing our work:
+```
+@misc{giaretta2026directsafediffusionbasedplanning,
+      title={DiRecT: Safe Diffusion-Based Planning via Receding-Horizon Denoising}, 
+      author={Paolo Giaretta and Zeyang Li and Navid Azizan},
+      year={2026},
+      eprint={2606.15359},
+      archivePrefix={arXiv},
+      primaryClass={cs.LG},
+      url={https://arxiv.org/abs/2606.15359}, 
+}
+```
+
+The `avoiding-v0` environment and expert data come from the D3IL benchmark:
+```bibtex
+@inproceedings{jia2024towards,
+  title={Towards Diverse Behaviors: A Benchmark for Imitation Learning with Human Demonstrations},
+  author={Xiaogang Jia and Denis Blessing and Xinkai Jiang and Moritz Reuss and Atalay Donat and Rudolf Lioutikov and Gerhard Neumann},
+  booktitle={The Twelfth International Conference on Learning Representations},
+  year={2024},
+  url={https://openreview.net/forum?id=6pPYRXKPpw}
+}
+```
+
+## References
+- Michael Janner, Yilun Du, Joshua B. Tenenbaum, and Sergey Levine. Planning with diffusion for flexible behavior synthesis. In Proceedings of the 39th International Conference on Machine Learning, volume 162 of Proceedings of Machine Learning Research, pages 9902–9915, 2022.
+- Prafulla Dhariwal and Alexander Nichol. Diffusion models beat gans on image synthesis. In Advances in Neural Information Processing Systems, volume 34, pages 8780–8794, 2021.
+- Jacob K Christopher, Stephen Baek, and Ferdinando Fioretto. Constrained synthesis with projected diffusion models. Advances in Neural Information Processing Systems, 37:89307–89333, 2024.
+- Wei Xiao, Tsun-Hsuan Wang, Chuang Gan, Ramin Hasani, Mathias Lechner, and Daniela Rus. Safediffuser: Safe planning with diffusion probabilistic models. In The Thirteenth International Conference on Learning Representations, 2025. URL https://openreview.net/forum?id=ig2wk7kK9J
+- Jichen Zhang, Liqun Zhao, Antonis Papachristodoulou, and Jack Umenberger. Constrained diffusers for safe planning and control, 2025. URL https://arxiv.org/abs/2506.12544.
+- Xiaogang Jia, Denis Blessing, Xinkai Jiang, Moritz Reuss, Atalay Donat, Rudolf Lioutikov, and Gerhard Neumann. Towards diverse behaviors: A benchmark for imitation learning with human demonstrations. In The Twelfth International Conference on Learning Representations, 2024.

applications/avoiding/avoiding_baselines/DiRecT/d3il/__init__.py

@@ -0,0 +1 @@
+from .environments.d3il.envs.gym_avoiding_env.gym_avoiding import envs

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/controllers/Config/fictive_robot_controller_config.gin

@@ -0,0 +1,19 @@
+# Joint PD Controller
+environments.d3il.d3il_sim.controllers.GainsInterface.JointPDGains.pgain = [1200.0, 1200.0, 1200.0, 1200.0, 1200.0, 300.0, 100.0]
+environments.d3il.d3il_sim.controllers.GainsInterface.JointPDGains.dgain = [25.0, 25.0, 25.0, 25.0, 15.0, 12.5, 7.5]
+
+# Damping Controller
+environments.d3il.d3il_sim.controllers.GainsInterface.DampingGains.dgain = [5. , 5. , 5. , 5. , 3. , 2.5, 1.5]
+
+# Cartesian Position Impedance Controller
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosPDGains.pgain = [1000.0, 1000.0, 1000.0]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosPDGains.dgain = [31.622776601683793, 31.622776601683793, 31.622776601683793]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosPDGains.pgain_null = [1.2, 1.2, 1.2, 1.2, 0.5, 0.3, 0.1]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosPDGains.dgain_null = [0.1 , 0.1 , 0.1 , 0.1 , 0.06, 0.05, 0.03]
+
+# Cartesian Position and Orientation Impedance Controller
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatPDGains.pgain = [33.13189603604459, 187.9403713446798, 31.0064033017648, 27.69370238555632, 33.98706171459314, 180.9185531893281]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatPDGains.dgain = [16.329377526884908,  7.672464154315535,  19.307738014214987,  19.369027815549774,  14.709666893271407,  13.529185054450448]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatPDGains.pgain_null = [7.675519770796831, 2.676935478437176, 8.539040163444975, 1.270446361314313, 8.87752182480855, 2.186782233762969, 4.414432577659688]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatPDGains.dgain_null = [0.35158682223052035,0.03532768494155088,0.9812676349486628,0.791470915462905,0.00066062504055376,0.9520533614689963,0.04623941774486632]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatPDGains.dgain_velcontroller = [25.0, 25.0, 25.0, 25.0, 15.0, 12.5, 7.5]

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/controllers/Config/mujoco_controller_config.gin

@@ -0,0 +1,48 @@
+# Joint PD Controller
+# ##these gains seem to be too high ...
+#environments.d3il.d3il_sim.controllers.GainsInterface.JointPDGains.pgain = [600.0, 600.0, 600.0, 600.0, 250.0, 150.0, 50.0]
+#environments.d3il.d3il_sim.controllers.GainsInterface.JointPDGains.dgain = [50.0, 50.0, 50.0, 50.0, 30.0, 25.0, 15.0]
+
+environments.d3il.d3il_sim.controllers.GainsInterface.JointPDGains.pgain = [120.0, 120.0, 120.0, 120.0, 50.0, 30.0, 10.0]
+environments.d3il.d3il_sim.controllers.GainsInterface.JointPDGains.dgain = [10.0, 10.0, 10.0, 10.0, 6.0, 5.0, 3.0]
+
+# Damping Controller
+environments.d3il.d3il_sim.controllers.GainsInterface.DampingGains.dgain = [5. , 5. , 5. , 5. , 3. , 2.5, 1.5]
+
+# Cartesian Position Impedance Controller
+#environments.d3il.d3il_sim.controllers.GainsInterface.CartPosPDGains.pgain = [80.0, 80.0, 80.0]
+#environments.d3il.d3il_sim.controllers.GainsInterface.CartPosPDGains.dgain = [18.973665961010276, 18.973665961010276, 18.973665961010276]
+#environments.d3il.d3il_sim.controllers.GainsInterface.CartPosPDGains.pgain_null = [1.2, 1.2, 1.2, 1.2, 0.5, 0.3, 0.1]
+#environments.d3il.d3il_sim.controllers.GainsInterface.CartPosPDGains.dgain_null = [0.1 , 0.1 , 0.1 , 0.1 , 0.06, 0.05, 0.03]
+
+# Cartesian Position Impedance Controller
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosControllerConfig.pgain_pos = [160.0, 160.0, 160.0]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosControllerConfig.pgain_null = [1.2, 1.2, 1.2, 1.2, 0.5, 0.3, 0.1]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosControllerConfig.J_reg = 1e-6
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosControllerConfig.W = [1, 1, 1, 1, 1, 1, 1]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosControllerConfig.rest_posture = [0, 0.174, 0, -0.872, 0, 1.222, 0.785]
+
+# Cartesian Position Impedance Controller And Orientation
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.pgain_pos = [200.0, 200.0, 800.0]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.pgain_null = [40, 40, 40, 40, 40, 40, 40]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.J_reg = 1e-12
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.W = [1, 1, 1, 1, 1, 1, 1]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.rest_posture = [0, 0.174, 0, -0.872, 0, 1.222, 0.785]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.pgain_quat = [30.0, 30.0, 30.0]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.ddgain = [0.4, 0.4, 0.4, 0.4 ,0.4, 0.4, 0.4]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.joint_filter_coefficient = 1.0
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.min_svd_values = 1e-2
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.max_svd_values = 1e2
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.num_iter = 3
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.learningRate = 0.001
+
+
+# Cartesian Position Jacobian Transpose Impedance Controller And Orientation
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatJacTransposeControllerConfig.pgain_pos = [450.0, 450.0, 450.0]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatJacTransposeControllerConfig.dgain = [25.0, 25.0, 25.0, 25.0, 25.0, 25.0]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatJacTransposeControllerConfig.pgain_null = [1.0, 1.0, 1.0, 1.0, 1.0, 1.0, 1.0]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatJacTransposeControllerConfig.dgain_null = [0.25 , 0.25 , 0.25 , 0.25 , 0.25, 0.25, 0.25]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatJacTransposeControllerConfig.J_reg = 1e-3
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatJacTransposeControllerConfig.W = [1, 1, 1, 10, 1, 1, 1]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatJacTransposeControllerConfig.rest_posture = [0, 0.174, 0, -0.872, 0, 1.222, 0.785]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatJacTransposeControllerConfig.pgain_quat = [250.0, 250.0, 250.0]

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/controllers/Config/pybullet_controller_config.gin

@@ -0,0 +1,34 @@
+# Joint PD Controller
+
+environments.d3il.d3il_sim.controllers.GainsInterface.JointPDGains.pgain = [1200.0, 1200.0, 1200.0, 1200.0, 1200.0, 300.0, 100.0]
+environments.d3il.d3il_sim.controllers.GainsInterface.JointPDGains.dgain = [25.0, 25.0, 25.0, 25.0, 25.0, 15.0, 10.0]
+
+# Damping Controller
+environments.d3il.d3il_sim.controllers.GainsInterface.DampingGains.dgain = [5. , 5. , 5. , 5. , 3. , 2.5, 1.5]
+
+# Cartesian Position Impedance Controller
+#environments.d3il.d3il_sim.controllers.GainsInterface.CartPosPDGains.pgain = [80.0, 80.0, 80.0]
+#environments.d3il.d3il_sim.controllers.GainsInterface.CartPosPDGains.dgain = [18.973665961010276, 18.973665961010276, 18.973665961010276]
+#environments.d3il.d3il_sim.controllers.GainsInterface.CartPosPDGains.pgain_null = [1.2, 1.2, 1.2, 1.2, 0.5, 0.3, 0.1]
+#environments.d3il.d3il_sim.controllers.GainsInterface.CartPosPDGains.dgain_null = [0.1 , 0.1 , 0.1 , 0.1 , 0.06, 0.05, 0.03]
+
+# Cartesian Position Impedance Controller
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosControllerConfig.pgain_pos = [160.0, 160.0, 160.0]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosControllerConfig.pgain_null = [1.2, 1.2, 1.2, 1.2, 0.5, 0.3, 0.1]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosControllerConfig.J_reg = 1e-6
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosControllerConfig.W = [1, 1, 1, 1, 1, 1, 1]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosControllerConfig.rest_posture = [0, 0.174, 0, -0.872, 0, 1.222, 0.785]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.ddgain = [0.4, 0.4, 0.4, 0.4 ,0.4, 0.4, 0.4]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.joint_filter_coefficient = 0.99
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.min_svd_values = 1e-2
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.max_svd_values = 1e2
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.num_iter = 3
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.learningRate = 0.001
+
+# Cartesian Position Impedance Controller And Orientation
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.pgain_pos = [160.0, 160.0, 160.0]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.pgain_null = [1.2, 1.2, 1.2, 1.2, 0.5, 0.3, 0.1]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.J_reg = 1e-6
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.W = [1, 1, 1, 1, 1, 1, 1]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.rest_posture = [0, 0.174, 0, -0.872, 0, 1.222, 0.785]
+environments.d3il.d3il_sim.controllers.GainsInterface.CartPosQuatControllerConfig.pgain_quat = [800.0, 800.0, 800.0]

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/controllers/Controller.py

@@ -0,0 +1,253 @@
+import threading
+from abc import abstractmethod
+
+import numpy as np
+
+import d3il.environments.d3il.d3il_sim.controllers.GainsInterface as gains
+
+
+class ControllerBase:
+
+    def __init__(self):
+        self.paramsLock = threading.Lock()
+        self.last_control_timestamp = np.nan
+        self._max_duration = None
+        self._max_timesteps = None
+        self._controller_timer = None
+
+    def isFinished(self, robot):
+
+        if self._max_duration is not None:
+            return robot.time_stamp - self._controller_timer >= self._max_duration
+
+        if self._max_timesteps is not None:
+            return robot.step_count - self._controller_timer >= self._max_timesteps
+        return False
+
+    def initController(self, robot, maxDuration):
+        return
+
+    def getControl(self, robot):
+        self.last_control_timestamp = robot.time_stamp
+        return 0
+
+    def is_used(self, robot):
+        return (
+            not np.isnan(self.last_control_timestamp)
+            and robot.time_stamp - self.last_control_timestamp < 0.03
+        )
+
+    def setAction(self, action):
+        return 0
+
+    def run(self, robot, log=True):
+
+        while not self.isFinished(robot):
+            robot.nextStep(log)
+
+    def executeController(self, robot, maxDuration=10, block=True, log=True):
+
+        self._max_duration = maxDuration
+        self._max_timesteps = None
+        self._controller_timer = robot.time_stamp
+
+        self.initController(robot, maxDuration)
+        robot.activeController = self
+
+        if block:
+            self.run(robot, log=log)
+
+    def executeControllerTimeSteps(self, robot, timeSteps=10, block=True, log=True):
+
+        self._max_duration = None
+        self._max_timesteps = timeSteps
+        self._controller_timer = robot.step_count
+
+        self.initController(robot, timeSteps * robot.dt)
+        robot.activeController = self
+
+        if block:
+            self.run(robot, log)
+
+    @abstractmethod
+    def reset(self):
+        pass
+
+
+class TorqueController(ControllerBase):
+
+    def __init__(self):
+        ControllerBase.__init__(self)
+        self.reset()
+
+    def getControl(self, robot):
+        super(TorqueController, self).getControl(robot)
+        return self.torque
+
+    def setAction(self, action):
+        self.torque = action.copy()
+
+    def reset(self):
+        self.torque = []
+
+
+class TrackingController(ControllerBase):
+
+    def __init__(self, dimSetPoint):
+        ControllerBase.__init__(self)
+        self.dimSetPoint = dimSetPoint
+        self.tracking_error = False
+
+    def setSetPoint(self, desired_pos, desired_vel=None, desired_acc=None):
+        pass
+
+    def getCurrentPos(self, robot):
+        pass
+
+    def getDesiredPos(self, robot):
+        pass
+
+    @abstractmethod
+    def reset(self):
+        pass
+
+
+class JointPDController(TrackingController, gains.JointPDGains):
+
+    def __init__(self):
+        TrackingController.__init__(self, dimSetPoint=7)
+        gains.JointPDGains.__init__(self)
+
+        self.reset()
+
+    def reset(self):
+        self.desired_joint_pos = np.array([0, 0, 0, -1.562, 0, 1.914, 0])
+        self.desired_joint_vel = np.zeros((7,))
+        self.desired_joint_acc = np.zeros((7,))
+
+    def getControl(self, robot):
+
+        super(JointPDController, self).getControl(robot)
+        self.paramsLock.acquire()
+        qd_d = self.desired_joint_pos - robot.current_j_pos
+        vd_d = self.desired_joint_vel - robot.current_j_vel
+
+        target_j_acc = self.pgain * qd_d + self.dgain * vd_d
+
+        robot.des_joint_pos = self.desired_joint_pos.copy()
+        robot.des_joint_vel = self.desired_joint_vel.copy()
+        robot.des_joint_acc = self.desired_joint_acc.copy()
+
+        self.paramsLock.release()
+        return target_j_acc
+
+    def setSetPoint(self, desired_pos, desired_vel=None, desired_acc=None):
+
+        self.paramsLock.acquire()
+        self.desired_joint_pos = desired_pos
+        if desired_vel is not None:
+            self.desired_joint_vel = desired_vel
+        if desired_acc is not None:
+            self.desired_joint_acc = desired_acc
+        self.paramsLock.release()
+
+    def getCurrentPos(self, robot):
+
+        return robot.current_j_pos
+
+    def getDesiredPos(self, robot):
+
+        return robot.des_joint_pos
+
+
+class ModelBasedFeedforwardController(JointPDController):
+
+    def __init__(self):
+        JointPDController.__init__(self)
+
+    def getControl(self, robot):
+
+        super(ModelBasedFeedforwardController, self).getControl(robot)
+        self.paramsLock.acquire()
+        qd_d = self.desired_joint_pos - robot.current_j_pos
+        vd_d = self.desired_joint_vel - robot.current_j_vel
+
+        target_j_acc = self.pgain * qd_d + self.dgain * vd_d
+        uff = robot.get_mass_matrix(self.desired_joint_pos).dot(
+            self.desired_joint_acc
+        ) + robot.get_coriolis(self.desired_joint_pos, self.desired_joint_vel)
+
+        robot.des_joint_pos = self.desired_joint_pos.copy()
+        robot.des_joint_vel = self.desired_joint_vel.copy()
+        robot.des_joint_acc = self.desired_joint_acc.copy()
+
+        self.paramsLock.release()
+        return target_j_acc + uff
+
+
+class ModelBasedFeedbackController(JointPDController):
+
+    def __init__(self):
+        JointPDController.__init__(self)
+
+    def getControl(self, robot):
+
+        super(ModelBasedFeedbackController, self).getControl(robot)
+        self.paramsLock.acquire()
+        qd_d = self.desired_joint_pos - robot.current_j_pos
+        vd_d = self.desired_joint_vel - robot.current_j_vel
+
+        target_j_acc = self.pgain * qd_d + self.dgain * vd_d + self.desired_joint_acc
+        uff = robot.get_mass_matrix(robot.current_j_pos).dot(
+            target_j_acc
+        ) + robot.get_coriolis(robot.current_j_pos, robot.current_j_vel)
+
+        robot.des_joint_pos = self.desired_joint_pos.copy()
+        robot.des_joint_vel = self.desired_joint_vel.copy()
+        robot.des_joint_acc = self.desired_joint_acc.copy()
+
+        self.paramsLock.release()
+        return uff
+
+
+class JointPositionController(JointPDController):
+    def setAction(self, action):
+        self.desired_joint_pos = action
+
+
+class JointVelocityController(JointPDController):
+    def __init__(self):
+        JointPDController.__init__(self)
+        self.pgain = np.zeros((self.dimSetPoint,))
+
+    def setAction(self, action):
+        self.desired_joint_vel = action
+
+
+class ZeroTorqueController(TrackingController):
+
+    def __init__(self, dimSetPoint=7):
+        TrackingController.__init__(self, dimSetPoint=dimSetPoint)
+
+    def getControl(self, robot):
+        super().getControl(robot)
+        target_j_acc = np.zeros((self.dimSetPoint,))
+        return target_j_acc
+
+    def reset(self):
+        pass
+
+
+class DampingController(ControllerBase, gains.DampingGains):
+
+    def __init__(self):
+        ControllerBase.__init__(self)
+        gains.DampingGains.__init__(self)
+
+    def getControl(self, robot):
+
+        super(DampingController, self).getControl(robot)
+        self.paramsLock.acquire()
+        target_j_acc = -self.dgain * robot.current_j_vel
+        self.paramsLock.release()
+        return target_j_acc

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/controllers/GainsInterface.py

@@ -0,0 +1,90 @@
+import gin
+
+
+class PDGains:
+    def __init__(self, pgain, dgain):
+        self.pgain = pgain
+        self.dgain = dgain
+
+
+@gin.configurable
+class JointPDGains(PDGains):
+    def __init__(self, pgain, dgain):
+        super(JointPDGains, self).__init__(pgain, dgain)
+
+
+@gin.configurable
+class DampingGains(PDGains):
+    def __init__(self, dgain):
+        super(DampingGains, self).__init__(None, dgain)
+
+
+@gin.configurable
+class CartPosControllerConfig:
+    def __init__(self, pgain_pos, J_reg, W, pgain_null, rest_posture):
+        self.pgain_pos = pgain_pos
+        self.J_reg = J_reg
+        self.W = W
+
+        self.rest_posture = rest_posture
+        self.pgain_null = pgain_null
+
+
+@gin.configurable
+class CartPosQuatControllerConfig:
+    def __init__(
+        self,
+        pgain_pos,
+        pgain_quat,
+        J_reg,
+        W,
+        pgain_null,
+        rest_posture,
+        ddgain,
+        joint_filter_coefficient,
+        min_svd_values,
+        max_svd_values,
+        num_iter,
+        learningRate,
+    ):
+
+        self.pgain_pos = pgain_pos
+        self.pgain_quat = pgain_quat
+        self.J_reg = J_reg
+        self.W = W
+
+        self.rest_posture = rest_posture
+        self.pgain_null = pgain_null
+
+        self.ddgain = ddgain
+        self.joint_filter_coefficient = joint_filter_coefficient
+        self.min_svd_values = min_svd_values
+        self.max_svd_values = max_svd_values
+
+        self.num_iter = num_iter
+        self.learningRate = learningRate
+
+
+@gin.configurable
+class CartPosQuatJacTransposeControllerConfig:
+    def __init__(
+        self,
+        pgain_pos,
+        dgain,
+        pgain_quat,
+        J_reg,
+        W,
+        pgain_null,
+        dgain_null,
+        rest_posture,
+    ):
+
+        self.pgain_pos = pgain_pos
+        self.dgain = dgain
+        self.pgain_quat = pgain_quat
+        self.dgain_null = dgain_null
+        self.J_reg = J_reg
+        self.W = W
+
+        self.rest_posture = rest_posture
+        self.pgain_null = pgain_null

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/controllers/IKControllers.py

@@ -0,0 +1,386 @@
+from abc import abstractmethod
+import logging
+
+import numpy as np
+
+import d3il.environments.d3il.d3il_sim.controllers.GainsInterface as gains
+import d3il.environments.d3il.d3il_sim.utils as utils
+from d3il.environments.d3il.d3il_sim.controllers.Controller import TrackingController
+
+
+class CartPosImpedenceController(TrackingController, gains.CartPosControllerConfig):
+
+    def __init__(self):
+
+        TrackingController.__init__(self, dimSetPoint=3)
+        gains.CartPosControllerConfig.__init__(self)
+
+        self.J_reg = 1e-6
+        self.W = np.diag([1, 1, 1, 1, 1, 1, 1])
+
+        self.target_th_null = np.array(
+            [
+                3.57795216e-09,
+                1.74532920e-01,
+                3.30500960e-08,
+                -8.72664630e-01,
+                -1.14096181e-07,
+                1.22173047e00,
+                7.85398126e-01,
+            ]
+        )
+
+        self.reset()
+
+    def reset(self):
+        self.desired_c_pos = np.array([0.624, 0, 0.55])
+        self.desired_c_vel = np.zeros((3,))
+        self.desired_c_acc = np.zeros((3,))
+
+    def getControl(self, robot):
+
+        self.paramsLock.acquire()
+        xd_d = self.desired_c_pos - robot.current_c_pos
+        target_c_acc = self.pgain * xd_d
+
+        J = robot.getJacobian()
+        J = J[:3, :]
+        Jw = J.dot(self.W)
+
+        JwJ_reg = Jw.dot(J.T) + self.J_reg * np.eye(3)
+
+        qd_null = self.pgain_null * (self.target_th_null - robot.current_j_pos)
+
+        qd_d = np.linalg.solve(JwJ_reg, target_c_acc - J.dot(qd_null))
+        qd_d = self.W.dot(J.transpose()).dot(qd_d) + qd_null
+
+        robot.des_c_pos = self.desired_c_pos
+        robot.des_c_vel = self.desired_c_vel
+
+        robot.jointTrackingController.setSetPoint(robot.current_j_pos, qd_d)
+        self.paramsLock.release()
+
+        return robot.jointTrackingController.getControl(robot)
+
+    def setGains(self, pGain, dGain):
+
+        self.paramsLock.acquire()
+        self.pgain = pGain
+        self.dgain = dGain
+        self.paramsLock.release()
+
+    def setSetPoint(self, desired_pos, desired_vel=None, desired_acc=None):
+
+        self.paramsLock.acquire()
+        self.desired_c_pos = desired_pos
+        if desired_vel is not None:
+            self.desired_c_vel = desired_vel
+        if desired_acc is not None:
+            self.desired_c_acc = desired_acc
+        self.paramsLock.release()
+
+    def getCurrentPos(self, robot):
+
+        return robot.current_c_pos
+
+    def getDesiredPos(self, robot):
+        return robot.des_c_pos
+
+
+class CartesianPositionController(CartPosImpedenceController):
+    def setAction(self, action):
+        self.desired_c_pos = action
+
+    @abstractmethod
+    def reset(self):
+        pass
+
+
+class CartPosQuatImpedenceController(
+    TrackingController, gains.CartPosQuatControllerConfig
+):
+
+    def __init__(self):
+
+        TrackingController.__init__(self, dimSetPoint=7)
+        gains.CartPosQuatControllerConfig.__init__(self)
+        self.W = np.diag(self.W)
+        self.dgain_null = np.sqrt(self.pgain_null) * 2
+
+        self.reset()
+        self.neglect_dynamics = False
+
+    def reset(self):
+        self.desired_c_pos = np.array([0.624, 0, 0.55])
+        self.desired_quat = np.array([0.0, 0.984, 0, 0.177])
+        self.desired_quat = self.desired_quat / np.linalg.norm(self.desired_quat)
+        self.desired_c_vel = np.zeros((3,))
+        self.desired_quat_vel = np.zeros((4,))
+
+        self.old_des_joint_vel = np.zeros((7,))
+        self.old_q = np.zeros((7,))
+        self.old_q[:] = np.nan
+
+    def getControl(self, robot):
+
+        self.paramsLock.acquire()
+        super(CartPosQuatImpedenceController, self).getControl(robot)
+
+        if any(np.isnan(self.old_q)):
+            self.old_q = robot.current_j_pos.copy()
+
+        q = self.old_q.copy()
+
+        q = (
+            self.joint_filter_coefficient * q
+            + (1 - self.joint_filter_coefficient) * robot.current_j_pos
+        )
+
+        qd_dsum = np.zeros(q.shape)
+
+        oldErrorNorm = np.inf
+        qd_d = np.zeros(q.shape)
+
+        [current_c_pos, current_c_quat] = robot.getForwardKinematics(q)
+        target_cpos_acc = self.desired_c_pos - current_c_pos
+
+        curr_quat = current_c_quat
+
+        if np.linalg.norm(curr_quat - self.desired_quat) > np.linalg.norm(
+            curr_quat + self.desired_quat
+        ):
+            curr_quat = -curr_quat
+        oldErrorNorm = np.sum(target_cpos_acc**2) + np.sum(
+            (curr_quat - self.desired_quat) ** 2
+        )
+
+        des_quat = self.desired_quat
+        for i in range(self.num_iter):
+
+            [current_c_pos, current_c_quat] = robot.getForwardKinematics(q)
+            target_cpos_acc = self.desired_c_pos - current_c_pos
+
+            curr_quat = current_c_quat
+
+            if np.linalg.norm(curr_quat - des_quat) > np.linalg.norm(
+                curr_quat + des_quat
+            ):
+                des_quat = -des_quat
+            errNorm = np.sum(target_cpos_acc**2) + np.sum((curr_quat - des_quat) ** 2)
+
+            target_cquat = utils.get_quaternion_error(curr_quat, des_quat)
+
+            target_cpos_acc = np.clip(target_cpos_acc, -0.01, 0.01)
+            target_cquat = np.clip(target_cquat, -0.1, 0.1)
+
+            target_c_acc = np.hstack(
+                (self.pgain_pos * target_cpos_acc, self.pgain_quat * target_cquat)
+            )
+
+            J = robot.getJacobian(q)
+
+            Jw = J.dot(self.W)
+
+            condNumber = np.linalg.cond(Jw.dot(J.T))
+            JwJ_reg = Jw.dot(J.T) + self.J_reg * np.eye(J.shape[0])
+
+            u, s, v = np.linalg.svd(JwJ_reg, full_matrices=False)
+            s_orig = s
+            s = np.clip(s, self.min_svd_values, self.max_svd_values)
+            JwJ_reg = u @ np.diag(s) @ v
+            condNumber2 = np.linalg.cond(JwJ_reg)
+            largestSV = np.max(s_orig)
+
+            qdev_rest = np.clip(self.rest_posture - q, -0.2, 0.2)
+
+            qd_null = np.array(self.pgain_null * (qdev_rest))
+
+            margin_to_limit = 0.01
+            pgain_limit = 20
+
+            qd_null_limit = np.zeros(qd_null.shape)
+            qd_null_limit_max = pgain_limit * (
+                robot.joint_pos_max - margin_to_limit - q
+            )
+            qd_null_limit_min = pgain_limit * (
+                robot.joint_pos_min + margin_to_limit - q
+            )
+            qd_null_limit[
+                q > robot.joint_pos_max - margin_to_limit
+            ] += qd_null_limit_max[q > robot.joint_pos_max - margin_to_limit]
+            qd_null_limit[
+                q < robot.joint_pos_min + margin_to_limit
+            ] += qd_null_limit_min[q < robot.joint_pos_min + margin_to_limit]
+
+            qd_d = np.linalg.solve(JwJ_reg, target_c_acc - J.dot(qd_null))
+            qd_d = self.W.dot(J.transpose()).dot(qd_d) + qd_null
+
+            if np.linalg.norm(qd_d) > 3:
+                qd_d = qd_d * 3 / np.linalg.norm(qd_d)
+
+            qd_dsum = qd_dsum + qd_d
+
+            q = q + self.learningRate * qd_d
+            q = np.clip(q, robot.joint_pos_min, robot.joint_pos_max)
+
+        self.tracking_error = np.sum(np.abs(self.desired_c_pos - current_c_pos)) > 0.01
+
+        qd_dsum = (q - self.old_q) / robot.dt
+        des_acc = self.ddgain * (qd_dsum - self.old_des_joint_vel) / robot.dt
+
+        if np.sum(np.abs(self.desired_c_pos - current_c_pos)) > 0.1:
+            target_cquat = utils.get_quaternion_error(curr_quat, self.desired_quat)
+
+            logging.getLogger(__name__).debug(
+                "i: %d, Time: %f, Pos_error: %f, Quat_error: %f,  qd_d: %f"
+                % (
+                    i,
+                    robot.time_stamp,
+                    np.linalg.norm(self.desired_c_pos - current_c_pos),
+                    np.linalg.norm(target_cquat),
+                    np.linalg.norm(qd_d),
+                ),
+                qd_d,
+                self.desired_c_pos,
+                des_acc,
+            )
+
+        if np.linalg.norm(des_acc) > 10000:
+            des_acc = des_acc * 10000 / np.linalg.norm(des_acc)
+
+        robot.jointTrackingController.setSetPoint(q, qd_dsum, des_acc)
+
+        self.old_q = q.copy()
+        self.old_des_joint_vel = qd_dsum
+
+        robot.des_c_pos = self.desired_c_pos
+        robot.des_c_vel = self.desired_c_vel
+        robot.des_quat = self.desired_quat
+        robot.des_quat_vel = self.desired_quat_vel
+        robot.misc_data = np.array([errNorm, condNumber, condNumber2, largestSV])
+
+        self.paramsLock.release()
+
+        if self.neglect_dynamics:
+            robot.beam_to_joint_pos(q, resetDesired=False)
+            return np.zeros((7,))
+        else:
+            control = robot.jointTrackingController.getControl(robot)
+
+            return control
+
+    def setSetPoint(self, desired_pos, desired_vel=None, desired_acc=None):
+
+        self.paramsLock.acquire()
+        self.desired_c_pos = desired_pos[:3].copy()
+        self.desired_quat = desired_pos[3:] / np.linalg.norm(desired_pos[3:])
+        if desired_vel is not None:
+            self.desired_c_vel = desired_vel[:3]
+            self.desired_quat_vel = desired_vel[3:]
+
+        self.paramsLock.release()
+
+    def getCurrentPos(self, robot):
+
+        return np.hstack((robot.current_c_pos, robot.current_c_quat))
+
+    def getDesiredPos(self, robot):
+        return np.concatenate((robot.des_c_pos, robot.des_quat))
+
+
+class CartVelocityImpedenceController(CartPosQuatImpedenceController):
+
+    def __init__(self, fixed_orientation=None, max_cart_vel=0.5):
+
+        super(CartVelocityImpedenceController, self).__init__()
+        self.fixed_orientation = fixed_orientation
+        self.max_cart_vel = max_cart_vel
+
+        self.max_cart_pos = np.array([0.7, 0.4, 0.75])
+        self.min_cart_pos = np.array([0.2, -0.4, 0.0])
+
+    def getControl(self, robot):
+
+        self.desired_c_pos = np.array(robot.current_c_pos) + robot.dt * np.array(
+            self.desired_c_vel
+        )
+        self.desired_c_pos = np.clip(
+            self.desired_c_pos, self.min_cart_pos, self.max_cart_pos
+        )
+
+        if self.fixed_orientation is None:
+            self.desired_quat = self.fixed_orientation
+        else:
+            self.desired_quat = self.fixed_orientation
+
+        return super(CartVelocityImpedenceController, self).getControl(robot)
+
+    def setSetPoint(self, desired_pos, desired_vel, desired_acc=None):
+
+        self.paramsLock.acquire()
+        self.desired_c_vel = desired_vel[:3]
+        if self.fixed_orientation is None:
+            self.desired_quat_vel = desired_vel[3:]
+
+        self.paramsLock.release()
+
+
+class CartPosQuatImpedenceJacTransposeController(
+    CartPosQuatImpedenceController, gains.CartPosQuatJacTransposeControllerConfig
+):
+    def __init__(self):
+        super(CartPosQuatImpedenceJacTransposeController, self).__init__()
+        gains.CartPosQuatJacTransposeControllerConfig.__init__(self)
+
+    def getControl(self, robot):
+        self.paramsLock.acquire()
+        qd_d = self.desired_c_pos - robot.current_c_pos
+        target_cpos_acc = self.pgain_pos * qd_d
+        curr_quat = (
+            -robot.current_c_quat
+            if (robot.current_c_quat @ self.desired_quat) < 0
+            else robot.current_c_quat
+        )
+        target_cquat = self.pgain_quat * utils.get_quaternion_error(
+            robot.current_c_quat, self.desired_quat
+        )
+        J = robot.getJacobian()
+
+        qd_null = self.pgain_null * (
+            self.rest_posture - robot.current_j_pos
+        ) + self.dgain_null * (-robot.current_j_vel)
+        target_acc = np.concatenate((target_cpos_acc, target_cquat))
+        qdd = J.T @ (target_acc - self.dgain * (J @ robot.current_j_vel)) + qd_null
+        robot.des_c_pos = self.desired_c_pos
+        robot.des_c_vel = self.desired_c_vel
+        robot.des_quat = self.desired_quat
+        robot.des_quat_vel = self.desired_quat_vel
+        self.paramsLock.release()
+        return qdd
+
+
+class CartPosQuatCartesianRobotController(TrackingController):
+    def __init__(self):
+        TrackingController.__init__(self, dimSetPoint=7)
+        self.reset()
+
+    def reset(self):
+        self.desired_pos = np.zeros((self.dimSetPoint,))
+
+    def setAction(self, action):
+        self.desired_pos = action
+
+    def setSetPoint(self, desired_pos, desired_vel=None, desired_acc=None):
+        self.desired_pos = desired_pos
+
+    def getControl(self, robot):
+        robot.des_c_pos = self.desired_pos[:3]
+        robot.des_c_vel = np.zeros((3,))
+        if self.desired_pos.shape[0] > 3:
+            robot.des_quat = self.desired_pos[3:]
+        robot.des_quat_vel = np.zeros((4,))
+
+        return self.desired_pos
+
+    def getCurrentPos(self, robot):
+        return np.hstack((robot.current_c_pos, robot.current_c_quat))

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/controllers/TrajectoryTracking.py

@@ -0,0 +1,416 @@
+import logging
+
+import numpy as np
+from scipy.interpolate import make_interp_spline
+
+import d3il.environments.d3il.d3il_sim.utils as utils
+from d3il.environments.d3il.d3il_sim.controllers import ControllerBase
+
+
+class TrajectoryGenerator:
+    def __init__(self):
+        pass
+
+    def generate_trajectory(self, trackingController, robot, duration):
+        pass
+
+    def setDesiredPos(self, desiredPos):
+        pass
+
+
+class TrajectoryTracker(ControllerBase):
+
+    def __init__(self, dt, trajectory_generator=None):
+        ControllerBase.__init__(self)
+
+        self.trajectory_generator = trajectory_generator
+        self.startingTime = None
+        self.trajectory = None
+        self.trajectoryVel = None
+        self.trajectoryAcc = None
+        self.dt = dt
+        self.additionalDuration = 0
+        self.delay_trajectory = 0
+        self.old_time_stamp = None
+        self.delay_trajectory_on_error = False
+
+    def initController(self, robot, maxDuration):
+
+        self.get_tracking_controller(robot).initController(robot, maxDuration)
+        self.startingTime = robot.time_stamp
+        self.duration = maxDuration
+        self.delay_trajectory = 0
+        self.old_time_stamp = robot.time_stamp
+        if self.trajectory_generator is not None:
+            self.setTrajectory(
+                self.trajectory_generator.generate_trajectory(
+                    self.get_tracking_controller(robot), robot, maxDuration
+                )
+            )
+
+        self.set_point_for_time_step(robot, 0)
+
+    def get_tracking_controller(self, robot):
+        pass
+
+    def setAction(self, action):
+        if action.ndim == 1:
+            self.setDesiredPos(action)
+        else:
+            self.setTrajectory(action)
+
+    def setDesiredPos(self, desiredPosition):
+        self.trajectory_generator.setDesiredPos(desiredPosition)
+
+    def set_point_for_time_step(self, robot, timeStep):
+        desired_pos = self.trajectory[timeStep, :]
+
+        if timeStep < self.trajectory.shape[0] - 1:
+            desired_vel = self.trajectoryVel[timeStep, :]
+        else:
+            desired_vel = np.zeros((self.trajectory.shape[1],))
+
+        if timeStep < self.trajectory.shape[0] - 2:
+            desired_acc = self.trajectoryAcc[timeStep, :]
+        else:
+            desired_acc = np.zeros((self.trajectory.shape[1],))
+
+        self.get_tracking_controller(robot).setSetPoint(
+            desired_pos, desired_vel, desired_acc
+        )
+
+    def getControl(self, robot):
+        if self.trajectory is None:
+            logging.getLogger(__name__).warning("Error: Trajectory is empty")
+
+        self.paramsLock.acquire()
+
+        if (
+            self.delay_trajectory_on_error
+            and self.get_tracking_controller(robot).tracking_error
+        ):
+            self.delay_trajectory += robot.time_stamp - self.old_time_stamp
+        self.old_time_stamp = robot.time_stamp
+
+        timeStep = np.round(
+            (robot.time_stamp - self.startingTime - self.delay_trajectory) / self.dt
+        )
+
+        timeStep = int(np.min([timeStep, self.trajectory.shape[0] - 1]))
+
+        self.set_point_for_time_step(robot, timeStep)
+
+        self.paramsLock.release()
+
+        return self.get_tracking_controller(robot).getControl(robot)
+
+    def setTrajectory(self, trajectory):
+
+        self.paramsLock.acquire()
+
+        self.trajectory = trajectory
+        self.trajectoryVel = np.diff(trajectory, 1, axis=0) / self.dt
+        self.trajectoryAcc = np.diff(trajectory, 2, axis=0) / (self.dt**2)
+
+        self.paramsLock.release()
+
+    def resetTrajectory(self):
+
+        self.trajectory = None
+
+
+class JointTrajectoryTracker(TrajectoryTracker):
+
+    def __init__(self, dt, trajectory_generator=None):
+        TrajectoryTracker.__init__(self, dt, trajectory_generator)
+
+    def get_tracking_controller(self, robot):
+        return robot.jointTrackingController
+
+
+class CartPosTrajectoryTracker(TrajectoryTracker):
+
+    def __init__(self, dt, trajectory_generator=None):
+        TrajectoryTracker.__init__(self, dt, trajectory_generator)
+
+    def get_tracking_controller(self, robot):
+        return robot.cartesianPosTrackingController
+
+
+class CartPosQuatTrajectoryTracker(TrajectoryTracker):
+
+    def __init__(self, dt, trajectory_generator=None):
+        TrajectoryTracker.__init__(self, dt, trajectory_generator)
+
+    def get_tracking_controller(self, robot):
+        return robot.cartesianPosQuatTrackingController
+
+
+class GotoTrajectoryBase(TrajectoryGenerator):
+    def __init__(self):
+
+        TrajectoryGenerator.__init__(self)
+
+        self.desiredPosition = np.array([0, 0, 0, -1.562, 0, 1.914, 0])
+
+    def get_init_pos(self, trackingController, robot):
+
+        cur_state = trackingController.getCurrentPos(robot)
+        des_state = trackingController.getDesiredPos(robot)
+        if (
+            robot.smooth_spline
+            and not any(np.isnan(des_state))
+            and not trackingController.tracking_error
+            and trackingController.is_used(robot)
+        ):
+            cur_state = des_state
+        return cur_state
+
+    def setDesiredPos(self, desiredPosition):
+
+        self.desiredPosition = desiredPosition
+
+
+class GotoTrajectorySpline(GotoTrajectoryBase):
+
+    def __init__(self):
+
+        GotoTrajectoryBase.__init__(self)
+
+    def generate_trajectory(self, trackingController, robot, duration):
+
+        super().generate_trajectory(trackingController, robot, duration)
+
+        cur_state = self.get_init_pos(trackingController, robot)
+
+        time = np.linspace(0, duration, int(duration / robot.dt) + 1)
+        trajectory = np.zeros((time.shape[0], trackingController.dimSetPoint))
+
+        for i in range(trackingController.dimSetPoint):
+            try:
+                l, r = [(1, 0.0), (2, 0.0)], [(1, 0.0), (2, 0.0)]
+                bsplinef = make_interp_spline(
+                    x=[0, duration],
+                    y=[cur_state[i], self.desiredPosition[i]],
+                    bc_type=(l, r),
+                    k=5,
+                )
+                trajectory[:, i] = bsplinef(time)
+            except ValueError:
+                raise ValueError(
+                    "Robot might be already be at this configuration.\nDesired Pos: {}\nCurrent Pos: {}".format(
+                        self.desiredPosition[i], cur_state[i]
+                    )
+                )
+        return trajectory
+
+
+class GotoTrajectoryLinear(GotoTrajectoryBase):
+
+    def __init__(self):
+
+        GotoTrajectoryBase.__init__(self)
+
+    def generate_trajectory(self, trackingController, robot, duration):
+
+        super().generate_trajectory(trackingController, robot, duration)
+
+        cur_state = self.get_init_pos(trackingController, robot)
+
+        time = np.linspace(0, duration, int(duration / robot.dt) + 1)
+        trajectory = np.zeros((time.shape[0], trackingController.dimSetPoint))
+
+        for i in range(trackingController.dimSetPoint):
+            trajectory[:, i] = np.linspace(
+                cur_state[i], self.desiredPosition[i], int(duration / robot.dt) + 1
+            )
+
+        trajectory = trajectory[1:, :]
+        return trajectory
+
+
+class OfflineIKTrajectoryGenerator(TrajectoryGenerator):
+    def __init__(self, interpolation_trajectory=None):
+        TrajectoryGenerator.__init__(self)
+
+        self.desiredTaskPosition = np.zeros(
+            7,
+        )
+        self.J_reg = 1e-6
+        self.W = np.diag([1, 1, 1, 1, 1, 1, 1])
+        self.target_th_null = np.array(
+            [
+                3.57795216e-09,
+                1.74532920e-01,
+                3.30500960e-08,
+                -8.72664630e-01,
+                -1.14096181e-07,
+                1.22173047e00,
+                7.85398126e-01,
+            ]
+        )
+        if interpolation_trajectory is None:
+            self.interpolation_trajectory = GotoTrajectorySpline()
+        else:
+            self.interpolation_trajectory = interpolation_trajectory
+
+    def generate_trajectory(self, trackingController, robot, duration):
+
+        eps = 1e-5
+        IT_MAX = 1000
+        DT = 1e-3
+
+        i = 0
+        self.pgain = [
+            33.9403713446798,
+            30.9403713446798,
+            33.9403713446798,
+            27.69370238555632,
+            33.98706171459314,
+            30.9185531893281,
+        ]
+        self.pgain_null = 5 * np.array(
+            [
+                7.675519770796831,
+                2.676935478437176,
+                8.539040163444975,
+                1.270446361314313,
+                8.87752182480855,
+                2.186782233762969,
+                4.414432577659688,
+            ]
+        )
+        self.pgain_limit = 20
+
+        q = robot.current_j_pos.copy()
+        qd_d = np.zeros(q.shape)
+        oldErrNorm = np.inf
+        while True:
+            oldQ = q
+            q = q + DT * qd_d
+
+            q = np.clip(q, robot.joint_pos_min, robot.joint_pos_max)
+
+            cartPos, orient = robot.getForwardKinematics(q)
+            cpos_err = self.desiredTaskPosition[:3] - cartPos
+
+            if np.linalg.norm(orient - self.desiredTaskPosition[3:]) > np.linalg.norm(
+                orient + self.desiredTaskPosition[3:]
+            ):
+                orient = -orient
+
+            cpos_err = np.clip(cpos_err, -0.1, 0.1)
+            cquat_err = np.clip(
+                utils.get_quaternion_error(orient, self.desiredTaskPosition[3:]),
+                -0.5,
+                0.5,
+            )
+            err = np.hstack((cpos_err, cquat_err))
+
+            errNorm = np.sum(cpos_err**2) + np.sum(
+                (orient - self.desiredTaskPosition[3:]) ** 2
+            )
+            if errNorm > oldErrNorm:
+                q = oldQ
+                DT = DT * 0.7
+                continue
+            else:
+                DT = DT * 1.025
+
+            if errNorm < eps:
+                success = True
+                break
+            if i >= IT_MAX:
+                success = False
+                break
+
+            oldErrNorm = errNorm
+
+            J = robot.getJacobian(q)
+
+            Jw = J.dot(self.W)
+
+            JwJ_reg = Jw.dot(J.T) + self.J_reg * np.eye(J.shape[0])
+
+            qd_null = self.pgain_null * (self.target_th_null - q)
+
+            margin_to_limit = 0.1
+            qd_null_limit = np.zeros(qd_null.shape)
+            qd_null_limit_max = self.pgain_limit * (
+                robot.joint_pos_max - margin_to_limit - q
+            )
+            qd_null_limit_min = self.pgain_limit * (
+                robot.joint_pos_min + margin_to_limit - q
+            )
+            qd_null_limit[
+                q > robot.joint_pos_max - margin_to_limit
+            ] += qd_null_limit_max[q > robot.joint_pos_max - margin_to_limit]
+            qd_null_limit[
+                q < robot.joint_pos_min + margin_to_limit
+            ] += qd_null_limit_min[q < robot.joint_pos_min + margin_to_limit]
+
+            qd_null += qd_null_limit
+            qd_d = np.linalg.solve(JwJ_reg, self.pgain * err - J.dot(qd_null))
+            qd_d = self.W.dot(J.transpose()).dot(qd_d) + qd_null
+
+            i += 1
+
+        print("Final IK error (%d iterations):  %s" % (i, errNorm))
+
+        logging.getLogger(__name__).debug(
+            "Final IK error (%d iterations):  %s" % (i, errNorm)
+        )
+        self.interpolation_trajectory.setDesiredPos(q)
+        return self.interpolation_trajectory.generate_trajectory(
+            trackingController, robot, duration
+        )
+
+    def setDesiredPos(self, desiredTaskPosition):
+
+        self.desiredTaskPosition = desiredTaskPosition
+
+
+class GotoCartPosQuatOfflineIKController(JointTrajectoryTracker):
+    def __init__(self, dt):
+        super(JointTrajectoryTracker, self).__init__(
+            dt, trajectory_generator=OfflineIKTrajectoryGenerator()
+        )
+
+
+class GotoJointController(JointTrajectoryTracker):
+    def __init__(self, dt):
+        super(JointTrajectoryTracker, self).__init__(
+            dt, trajectory_generator=GotoTrajectorySpline()
+        )
+
+
+class GotoCartPosImpedanceController(CartPosTrajectoryTracker):
+    def __init__(self, dt):
+        super(CartPosTrajectoryTracker, self).__init__(
+            dt, trajectory_generator=GotoTrajectorySpline()
+        )
+
+
+class GotoCartPosQuatImpedanceController(CartPosQuatTrajectoryTracker):
+    def __init__(self, dt):
+        super(CartPosQuatTrajectoryTracker, self).__init__(
+            dt, trajectory_generator=GotoTrajectorySpline()
+        )
+
+
+class GotoCartPosCartesianRobotController(JointTrajectoryTracker):
+
+    def __init__(self, dt):
+        assert False
+        super(JointTrajectoryTracker, self).__init__(
+            dt, trajectory_generator=GotoTrajectory()
+        )
+
+
+class GotoCartPosQuatCartesianRobotController(JointTrajectoryTracker):
+
+    def __init__(self, dt):
+        assert False
+        super(JointTrajectoryTracker, self).__init__(
+            dt, trajectory_generator=GotoTrajectory()
+        )

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/controllers/__init__.py

@@ -0,0 +1,3 @@
+from .Controller import *
+from .IKControllers import *
+from .TrajectoryTracking import *

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/controllers/executor.py

@@ -0,0 +1,48 @@
+import numpy as np
+
+
+def execute_controller(
+    robots: list,
+    controllers: list,
+    max_duration: float = 4.0,
+):
+    start_times = [rb.time_stamp for rb in robots]
+    robot_ctrl_matches = list(zip(robots, controllers, start_times))
+
+    for rb, ctrl, _ in robot_ctrl_matches:
+        ctrl.initController(rb, max_duration)
+
+    while (
+        not np.all([ctrl.isFinished(rb) for rb, ctrl, _ in robot_ctrl_matches])
+    ) and (
+        np.all([rb.time_stamp - st < max_duration for rb, _, st in robot_ctrl_matches])
+    ):
+        for rb, ctrl, _ in robot_ctrl_matches:
+            ctrl_action = ctrl.getControl(rb)
+            rb.command = ctrl_action
+
+        robots[0].nextStep()
+
+    for rb in robots:
+        rb.hold_joint_position()
+
+
+def execute_controller_timesteps(
+    robots: list,
+    controllers: list,
+    time_steps: int = 10,
+):
+    robot_ctrl_matches = list(zip(robots, controllers))
+
+    for rb, ctrl in robot_ctrl_matches:
+        ctrl.initController(rb, time_steps * rb.dt)
+
+    for i in range(time_steps):
+        for rb, ctrl in robot_ctrl_matches:
+            ctrl_action = ctrl.getControl(rb)
+            rb.command = ctrl_action
+
+        robots[0].nextStep()
+
+    for rb in robots:
+        rb.hold_joint_position()

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/core/Camera.py

@@ -0,0 +1,205 @@
+from abc import ABC, abstractmethod
+from typing import Optional, Tuple
+
+import cv2
+import numpy as np
+
+from d3il.environments.d3il.d3il_sim.core.sim_object.sim_object import (
+    IntelligentSimObject,
+)
+
+
+class Camera(IntelligentSimObject, ABC):
+
+    def __init__(
+        self,
+        name: str,
+        width: int = 1000,
+        height: int = 1000,
+        init_pos=None,
+        init_quat=None,
+        near: float = 0.01,
+        far: float = 10,
+        fovy: int = 45,
+        *args,
+        **kwargs
+    ):
+
+        if init_pos is None:
+            init_pos = [0, 0, 0]
+        if init_quat is None:
+            init_quat = [0, 1, 0, 0]
+
+        super(Camera, self).__init__(name, init_pos, init_quat)
+        self.width = width
+        self.height = height
+
+        self.near = near
+        self.far = far
+        self.fovy = fovy
+        self.fovx = (
+            2
+            * np.arctan(
+                self.width
+                * 0.5
+                / (self.height * 0.5 / np.tan(self.fovy * np.pi / 360 / 2))
+            )
+            / np.pi
+            * 360
+        )
+
+        self.fx = (self.width / 2) / (np.tan(self.fovx * np.pi / 180 / 2))
+        self.fy = (self.height / 2) / (np.tan(self.fovy * np.pi / 180 / 2))
+        self.cx = self.width / 2
+        self.cy = self.height / 2
+
+        self.intrinsics = np.array(
+            [[self.fx, 0, self.cx], [0, self.fy, self.cy], [0, 0, 1]]
+        )
+
+    def set_cam_params(
+        self,
+        width: Optional[int] = None,
+        height: Optional[int] = None,
+        near: Optional[float] = None,
+        far: Optional[float] = None,
+        fovy: Optional[int] = None,
+    ):
+
+        self.width = width or self.width
+        self.height = height or self.height
+
+        self.near = near or self.near
+        self.far = far or self.far
+        self.fovy = fovy or self.fovy
+        self.fovx = (
+            2
+            * np.arctan(
+                self.width
+                * 0.5
+                / (self.height * 0.5 / np.tan(self.fovy * np.pi / 360 / 2))
+            )
+            / np.pi
+            * 360
+        )
+
+        self.fx = (self.width / 2) / (np.tan(self.fovx * np.pi / 180 / 2))
+        self.fy = (self.height / 2) / (np.tan(self.fovy * np.pi / 180 / 2))
+        self.cx = self.width / 2
+        self.cy = self.height / 2
+
+        self.intrinsics = np.array(
+            [[self.fx, 0, self.cx], [0, self.fy, self.cy], [0, 0, 1]]
+        )
+
+    def get_segmentation(
+        self, width: int = None, height: int = None, depth: bool = True
+    ) -> np.ndarray:
+
+        return self._get_img_data(
+            width=width, height=height, depth=depth, segmentation=True
+        )
+
+    def get_image(
+        self,
+        width: int = None,
+        height: int = None,
+        depth: bool = True,
+        denormalize_depth: bool = True,
+    ) -> np.ndarray:
+
+        return self._get_img_data(width, height, depth, denormalize_depth, False)
+
+    def calc_point_cloud(
+        self, width: int = None, height: int = None, denormalize: bool = True
+    ) -> Tuple[np.ndarray, np.ndarray]:
+
+        rgb_img, depth_img = self.get_image(
+            width, height, denormalize_depth=denormalize
+        )
+
+        return self.calc_point_cloud_from_images(rgb_img=rgb_img, depth_img=depth_img)
+
+    def calc_point_cloud_from_images(
+        self, rgb_img: np.ndarray, depth_img: np.ndarray
+    ) -> Tuple[np.ndarray, np.ndarray]:
+
+        true_width = rgb_img.shape[1]
+        true_height = rgb_img.shape[0]
+        if self.height == true_height and self.width == true_width:
+            fx = self.fx
+            fy = self.fy
+            cx = self.cx
+            cy = self.cy
+        else:
+            fx = (true_width / 2) / (np.tan(self.fovx * np.pi / 180 / 2))
+            fy = (true_height / 2) / (np.tan(self.fovy * np.pi / 180 / 2))
+            cx = true_width / 2
+            cy = true_height / 2
+
+        z = depth_img
+        u = np.arange(true_width) - cx
+        v = np.arange(true_height) - cy
+
+        x = (z * u) / fx
+        y = (z.T * v).T / fy
+
+        points = np.stack((x, y, z), axis=-1).reshape((true_width * true_height, 3))
+        colors = rgb_img.reshape((true_width * true_height, 3)) / 255.0
+
+        valid_points = ~np.isnan(points).any(axis=1)
+        points = points[valid_points]
+        colors = colors[valid_points]
+
+        return points, colors
+
+    def denormalize_depth(self, depth_img: np.ndarray) -> np.ndarray:
+
+        z = self.near / (1 - depth_img * (1 - self.near / self.far))
+        return z
+
+    def apply_noise(self, depth_img: np.ndarray) -> np.ndarray:
+
+        z = self.denormalize_depth(depth_img)
+
+        depth_img = depth_img + (
+            0.0001 * np.power(z - 0.5, 2) + 0.0004
+        ) * np.random.rand(self.height, self.width)
+
+        depth_img = ((40000 * depth_img).astype(int) / 40000.0).astype(np.float32)
+
+        z = (
+            2
+            * self.far
+            * self.near
+            / (self.far + self.near - (self.far - self.near) * (2 * depth_img - 1))
+        )
+
+        z = cv2.bilateralFilter(z, 5, 0.1, 5)
+
+        return z
+
+    def get_poi(self) -> list:
+
+        return [self.name]
+
+    @abstractmethod
+    def _get_img_data(
+        self,
+        width: int = None,
+        height: int = None,
+        depth: bool = True,
+        denormalize_depth: bool = True,
+        segmentation: bool = False,
+    ) -> np.ndarray:
+
+        pass
+
+    @abstractmethod
+    def get_cart_pos_quat(self) -> Tuple[np.ndarray, np.ndarray]:
+
+        pass
+
+    @property
+    def fov(self):
+        return self.fovy

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/core/Model.py

@@ -0,0 +1,91 @@
+import numpy as np
+import pinocchio
+
+from d3il.environments.d3il.d3il_sim.utils.sim_path import d3il_path
+
+
+class RobotKinematicsInterface:
+    def getForwardKinematics(self, q=None):
+        pass
+
+    def getJacobian(self, q=None):
+        pass
+
+
+class RobotDynamicsInterface:
+    def get_gravity(self, q=None):
+        pass
+
+    def get_coriolis(self, q=None, qd=None):
+        pass
+
+    def get_mass_matrix(self, q=None):
+        pass
+
+
+class RobotModelFromPinochio(RobotKinematicsInterface, RobotDynamicsInterface):
+    def __init__(self, obj_urdf="./models/common/robots/panda_arm_hand_pinocchio.urdf"):
+        obj_urdf = d3il_path(obj_urdf)
+        self.pin_model = pinocchio.buildModelFromUrdf(obj_urdf)
+        self.pin_data = self.pin_model.createData()
+
+        self.pin_end_effector_frame_id = self.pin_model.getFrameId("panda_grasptarget")
+
+        self.pin_q = np.zeros(self.pin_model.nv)
+        self.pin_qd = np.zeros(self.pin_model.nv)
+
+    def getForwardKinematics(self, q):
+
+        self.pin_q[:7] = q
+        pinocchio.framesForwardKinematics(self.pin_model, self.pin_data, self.pin_q)
+
+        current_c_pos = np.array(
+            self.pin_data.oMf[self.pin_end_effector_frame_id].translation
+        )
+
+        quat_pin = pinocchio.Quaternion(
+            self.pin_data.oMf[self.pin_end_effector_frame_id].rotation
+        ).coeffs()
+        current_c_quat = np.zeros(4)
+        current_c_quat[1:] = quat_pin[:3]
+        current_c_quat[0] = quat_pin[-1]
+
+        return current_c_pos, current_c_quat
+
+    def getJacobian(self, q):
+        self.pin_q[:7] = q
+
+        pinocchio.computeJointJacobians(self.pin_model, self.pin_data, self.pin_q)
+        pinocchio.framesForwardKinematics(self.pin_model, self.pin_data, self.pin_q)
+        return pinocchio.getFrameJacobian(
+            self.pin_model,
+            self.pin_data,
+            self.pin_end_effector_frame_id,
+            pinocchio.LOCAL_WORLD_ALIGNED,
+        )[:, :7]
+
+    def get_gravity(self, q=None):
+        self.pin_q[:7] = q
+
+        vq = np.zeros(self.pin_q.shape)
+        aq0 = np.zeros(self.pin_q.shape)
+
+        b = pinocchio.rnea(self.pin_model, self.pin_data, self.pin_q, vq, aq0)
+        return b[:7]
+
+    def get_coriolis(self, q, qd):
+        self.pin_q[:7] = q
+        self.pin_qd[:7] = qd
+
+        aq0 = np.zeros(self.pin_q.shape)
+
+        b = pinocchio.rnea(self.pin_model, self.pin_data, self.pin_q, self.pin_qd, aq0)
+        return b[:7] - self.get_gravity(q)
+
+    def get_mass_matrix(self, q):
+        if q.shape[0] == 7:
+            q_tmp = q.copy()
+            q = np.zeros(9)
+            q[:7] = q_tmp
+        M = pinocchio.crba(self.pin_model, self.pin_data, q)
+        return M[:7, :7]

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/core/Robots.py

@@ -0,0 +1,491 @@
+from enum import Enum, auto
+
+import numpy as np
+
+import d3il.environments.d3il.d3il_sim.controllers as ctrl
+import d3il.environments.d3il.d3il_sim.core.logger as logger
+import d3il.environments.d3il.d3il_sim.core.Model as model
+import d3il.environments.d3il.d3il_sim.core.time_keeper as time_keeper
+import d3il.environments.d3il.d3il_sim.utils.geometric_transformation as geom_trans
+
+
+class RobotControlInterface(Enum):
+    TorqueInterface = (auto(),)
+    CartesianInterface = auto()
+
+
+class RobotBase:
+
+    def __init__(self, scene, dt, num_DoF=7, base_position=None, base_orientation=None):
+
+        self.scene = scene
+        self.dt = dt
+        self.num_DoF = num_DoF
+
+        self.inhand_cam = None
+
+        if base_position is None:
+            base_position = [0.0, 0.0, 0.0]
+        self.base_position = np.array(base_position)
+
+        if base_orientation is None:
+            base_orientation = [1.0, 0.0, 0.0, 0.0]
+        self.base_orientation = np.array(base_orientation)
+
+        self.misc_data = None
+
+        self.use_inv_dyn = False
+        self.gravity_comp = True
+        self.clip_rate = False
+
+        self.kinModel = model.RobotModelFromPinochio()
+        self.dynModel = self.kinModel
+
+        self.smooth_spline = True
+
+        self.torque_limit = np.array([80, 80, 80, 80, 10, 10, 10], dtype=np.float64)
+
+        self.joint_vel_limit = np.array([2.00, 2.00, 2.00, 2.00, 2.50, 2.50, 2.50])
+        self.joint_pos_min = np.array(
+            [-2.8973, -1.7628, -2.8973, -3.0718, -2.8973, -0.0175, -2.8973]
+        )
+        self.joint_pos_max = np.array(
+            [2.8973, 1.7628, 2.0, -0.0698, 2.8973, 3.7525, 2.8973]
+        )
+
+        self.clip_actions = True
+        self.rate_limit = 0.8
+        self.end_effector = None
+
+        self.jointTrackingController = ctrl.JointPDController()
+        self.cartesianPosTrackingController = ctrl.CartPosImpedenceController()
+        self.cartesianPosQuatTrackingController = ctrl.CartPosQuatImpedenceController()
+        self.activeController = self.jointTrackingController
+
+        self.jointPosController = ctrl.JointPositionController()
+        self.jointVelController = ctrl.JointVelocityController()
+        self.torqueController = ctrl.TorqueController()
+
+        self.gotoJointController = ctrl.GotoJointController(self.dt)
+
+        self.gotoCartPosController = ctrl.GotoCartPosQuatOfflineIKController(self.dt)
+        self.gotoCartPosQuatController = ctrl.GotoCartPosQuatOfflineIKController(
+            self.dt
+        )
+
+        self.gotoCartPosImpedanceController = ctrl.GotoCartPosImpedanceController(
+            self.dt
+        )
+        self.gotoCartPosQuatImpedanceController = (
+            ctrl.GotoCartPosQuatImpedanceController(self.dt)
+        )
+
+        self.jointTrajectoryTracker = ctrl.JointTrajectoryTracker(self.dt)
+        self.cartPosQuatTrajectoryTracker = ctrl.CartPosQuatTrajectoryTracker(self.dt)
+
+        self.controlInterface = RobotControlInterface.TorqueInterface
+
+        self.robot_logger = logger.RobotLogger(self)
+        self.time_keeper = time_keeper.TimeKeeper(self.dt)
+        self.additional_loggers = []
+
+        self.scene.add_robot(self)
+        self.reset()
+
+    def reset(self):
+
+        self.current_j_pos = np.zeros((self.num_DoF,)) * np.nan
+        self.current_j_vel = np.zeros((self.num_DoF,)) * np.nan
+        self.des_joint_pos = np.zeros((self.num_DoF,)) * np.nan
+        self.des_joint_vel = np.zeros((self.num_DoF,)) * np.nan
+        self.des_joint_acc = np.zeros((self.num_DoF,)) * np.nan
+
+        self.current_fing_pos = np.zeros((2,)) * np.nan
+        self.current_fing_vel = np.zeros((2,)) * np.nan
+        self.des_fing_pos = np.zeros(2) * np.nan
+        self.gripper_width = 0
+        self.set_gripper_width = 0.001
+        self.finger_commands = np.zeros((2,))
+        self.grasp_flag = False
+
+        self.current_c_pos = np.zeros((3,)) * np.nan
+        self.current_c_vel = np.zeros((3,)) * np.nan
+        self.current_c_quat = np.zeros((4,)) * np.nan
+        self.current_c_quat_vel = np.zeros((4,)) * np.nan
+        self.des_c_pos = np.zeros((3,)) * np.nan
+        self.des_c_vel = np.zeros((3,)) * np.nan
+        self.des_quat = np.zeros((4,)) * np.nan
+        self.des_quat_vel = np.zeros((4,)) * np.nan
+
+        self.current_c_pos_global = np.zeros((3,)) * np.nan
+        self.current_c_vel_global = np.zeros((3,)) * np.nan
+        self.current_c_quat_global = np.zeros((4,)) * np.nan
+        self.current_c_quat_vel_global = np.zeros((4,)) * np.nan
+        self.des_c_pos_global = np.zeros((3,)) * np.nan
+        self.des_c_vel_global = np.zeros((3,)) * np.nan
+        self.des_quat_global = np.zeros((4,)) * np.nan
+        self.des_quat_vel_global = np.zeros((4,)) * np.nan
+
+        self.uff = np.zeros((self.num_DoF + 2,))
+        self.uff_last = np.zeros((self.num_DoF + 2,))
+        self.last_cmd = np.zeros((self.num_DoF,))
+        self.command = np.zeros((self.num_DoF,))
+        self.grav_terms = np.zeros(9) * np.nan
+        self.current_load = np.zeros((self.num_DoF,))
+
+        self.robot_logger = logger.RobotLogger(self)
+        self.time_keeper = time_keeper.TimeKeeper(self.dt)
+
+        if self.gotoJointController is not None:
+            self.gotoJointController.resetTrajectory()
+
+        if self.gotoCartPosController is not None:
+            self.gotoCartPosController.resetTrajectory()
+        if self.gotoCartPosQuatController is not None:
+            self.gotoCartPosQuatController.resetTrajectory()
+        if self.gotoCartPosImpedanceController is not None:
+            self.gotoCartPosImpedanceController.resetTrajectory()
+        if self.gotoCartPosQuatImpedanceController is not None:
+            self.gotoCartPosQuatImpedanceController.resetTrajectory()
+
+        self.jointTrackingController = ctrl.JointPDController()
+
+        self.cartesianPosTrackingController.reset()
+        self.cartesianPosQuatTrackingController.reset()
+        self.activeController.reset()
+
+    def getForwardKinematics(self, q=None):
+        if q is None:
+            q = self.current_j_pos
+
+        return self.kinModel.getForwardKinematics(q)
+
+    def getJacobian(self, q=None):
+        if q is None:
+            q = self.current_j_pos
+
+        return self.kinModel.getJacobian(q)
+
+    def get_gravity(self, q=None):
+        if q is None:
+            q = self.current_j_pos
+
+        return self.dynModel.get_gravity(q)
+
+    def get_coriolis(self, q=None, qd=None):
+        if q is None:
+            q = self.current_j_pos
+
+        if qd is None:
+            qd = self.current_j_vel
+
+        return self.dynModel.get_coriolis(q, qd)
+
+    def get_mass_matrix(self, q=None):
+        if q is None:
+            q = self.current_j_pos
+        return self.dynModel.get_mass_matrix(q)
+
+    def start_logging(self, duration: float = 300.0, **kwargs):
+        self.robot_logger.start_logging(duration, **kwargs)
+
+    def stop_logging(self):
+        self.robot_logger.stop_logging()
+
+    def log_data(self):
+        self.robot_logger.log_data()
+
+    def hold_joint_position(self):
+        self.des_joint_pos = self.current_j_pos
+        self.command = None
+        self.activeController = self.jointTrackingController
+        self.jointTrackingController.setSetPoint(self.current_j_pos)
+
+    def gotoJointPosition(self, desiredPos, duration=4.0, block=True):
+
+        self.gotoJointController.setDesiredPos(desiredPos)
+        self.gotoJointController.executeController(self, duration, block=block)
+
+    def executeJointPosCtrlTimeStep(self, action, timeSteps=1, block=True):
+        self.jointPosController.setAction(action)
+        self.jointPosController.executeControllerTimeSteps(self, timeSteps, block=block)
+
+    def executeJointVelCtrlTimeStep(self, action, timeSteps=1, block=True):
+        self.jointVelController.setAction(action)
+        self.jointVelController.executeControllerTimeSteps(self, timeSteps, block=block)
+
+    def executeTorqueCtrlTimeStep(self, action, timeSteps=1, block=True):
+        self.torqueController.setAction(action)
+        self.torqueController.executeControllerTimeSteps(self, timeSteps, block=block)
+
+    def follow_JointTraj(self, desiredTraj, goto_start=True, block=True):
+
+        if goto_start:
+            if np.linalg.norm(self.current_j_pos - desiredTraj[0, :]) > 0.05:
+                self.gotoJointPosition(desiredTraj[0, :], 4.0)
+
+        self.jointTrajectoryTracker.setTrajectory(trajectory=desiredTraj)
+        self.jointTrajectoryTracker.executeController(
+            self, maxDuration=desiredTraj.shape[0] * self.dt, block=block
+        )
+
+    def gotoCartPosition_ImpedanceCtrl(
+        self, desiredPos, duration=4.0, global_coord=True, block=True
+    ):
+        if global_coord:
+            desiredPos = self._localize_cart_pos(desiredPos)
+
+        self.gotoCartPosImpedanceController.setDesiredPos(desiredPos)
+        self.gotoCartPosImpedanceController.executeController(
+            self, duration, block=block
+        )
+
+    def gotoCartPositionAndQuat_ImpedanceCtrl(
+        self, desiredPos, desiredQuat, duration=4.0, global_coord=True, block=True
+    ):
+        if global_coord:
+            desiredPos = self._localize_cart_pos(desiredPos)
+            desiredQuat = self._localize_cart_quat(desiredQuat)
+
+        self.gotoCartPosQuatImpedanceController.setDesiredPos(
+            np.hstack((desiredPos, desiredQuat))
+        )
+        self.gotoCartPosQuatImpedanceController.executeController(
+            self, duration, block=block
+        )
+
+    def gotoCartPosition(self, desiredPos, duration=4.0, global_coord=True, block=True):
+
+        if global_coord:
+            desiredPos = self._localize_cart_pos(desiredPos)
+        self.gotoCartPosController.setDesiredPos(desiredPos)
+        self.gotoCartPosController.executeController(self, duration, block=block)
+
+    def gotoCartPositionAndQuat(
+        self,
+        desiredPos,
+        desiredQuat,
+        duration=4.0,
+        global_coord=True,
+        block=True,
+        log=True,
+    ):
+
+        if global_coord:
+            desiredPos = self._localize_cart_pos(desiredPos)
+            desiredQuat = self._localize_cart_quat(desiredQuat)
+
+        self.gotoCartPosQuatController.setDesiredPos(
+            np.hstack((desiredPos, desiredQuat))
+        )
+        self.gotoCartPosQuatController.executeController(
+            self, duration, block=block, log=log
+        )
+
+    def follow_CartPositionAndQuatTraj(
+        self, desiredPos, desiredQuat, goto_start=True, global_coord=True, block=True
+    ):
+
+        if global_coord:
+            desiredPos = np.array([self._localize_cart_pos(p) for p in desiredPos])
+            desiredQuat = np.array([self._localize_cart_quat(q) for q in desiredQuat])
+
+        desiredTraj = np.hstack((desiredPos, desiredQuat))
+        duration = desiredTraj.shape[0] * self.dt
+
+        if goto_start:
+            is_position_far = (
+                np.linalg.norm(self.current_c_pos - desiredPos[0, :]) > 0.005
+            )
+            is_rotation_far = (
+                np.linalg.norm(self.current_c_quat - desiredQuat[0, :]) > 0.005
+            ) and (np.linalg.norm(self.current_c_quat + desiredQuat[0, :]) > 0.005)
+
+            if is_position_far or is_rotation_far:
+                self.gotoCartPositionAndQuat(
+                    desiredPos=desiredPos[0, :],
+                    desiredQuat=desiredQuat[0, :],
+                    duration=4.0,
+                )
+
+        self.cartPosQuatTrajectoryTracker.setTrajectory(trajectory=desiredTraj)
+        self.cartPosQuatTrajectoryTracker.executeController(
+            self, maxDuration=duration, block=block
+        )
+
+    def _localize_cart_pos(self, pos):
+        if not np.any(self.base_position) and np.array_equal(
+            self.base_orientation, [1, 0, 0, 0]
+        ):
+            return pos
+
+        pos_vec = np.array(pos)
+
+        local_pos = pos_vec - self.base_position
+        local_pos = geom_trans.quat_rot_vec(
+            geom_trans.quat_conjugate(self.base_orientation), local_pos
+        )
+        return local_pos
+
+    def _localize_cart_quat(self, quat):
+        if np.array_equal(self.base_orientation, [1, 0, 0, 0]):
+            return quat
+
+        local_quat = quat
+        if local_quat is not None:
+            quat_vec = np.array(local_quat)
+            local_quat = geom_trans.quat_mul(
+                geom_trans.quat_conjugate(self.base_orientation), quat_vec
+            )
+        return local_quat
+
+    def _localize_cart_coords(self, pos, quat=None):
+        if quat is not None:
+            quat = self._localize_cart_quat(quat)
+        return self._localize_cart_pos(pos), quat
+
+    def wait(self, duration=1.0, block=True):
+        self.activeController.executeController(self, duration, block=block)
+
+    def open_fingers(self):
+        self.set_desired_gripper_width(0.04)
+
+    def close_fingers(self, duration=0.5):
+        self.set_gripper_width = 0.0
+        self.grasp_flag = False
+        if duration > 0:
+            self.wait(duration)
+        self.grasp_flag = True
+
+    def set_desired_gripper_width(self, desired_width):
+        self.set_gripper_width = desired_width
+        self.grasp_flag = False
+
+    def fing_ctrl_step(self):
+
+        pgain = 1 * np.array([500, 500], dtype=np.float64)
+        dgain = 1 * np.array([10, 10], dtype=np.float64)
+
+        gripper_width = self.current_fing_pos
+        gripper_width_vel = self.current_fing_vel
+
+        des_fing_pos = np.array(
+            [self.set_gripper_width, self.set_gripper_width], dtype=np.float64
+        )
+
+        mean_finger_pos = np.mean(gripper_width)
+        force = pgain * (mean_finger_pos - gripper_width)
+
+        if (np.mean(gripper_width) - self.set_gripper_width) > 0.005:
+            if self.grasp_flag:
+                force2 = np.array([-20, -20])
+            else:
+                force2 = dgain * (np.array([-0.2, -0.2]) - gripper_width_vel)
+        else:
+            force2 = pgain * (des_fing_pos - gripper_width) - dgain * gripper_width_vel
+            force2 = np.clip(force2, -5, 5)
+
+        force = force + force2
+
+        return force
+
+    def receiveState(self):
+
+        raise NotImplementedError
+
+    def tick(self):
+        self.time_keeper.tick()
+
+    def nextStep(self, log=True):
+
+        self.scene.next_step(log)
+
+    def prepare_step(self):
+
+        self.des_joint_pos[:] = np.nan
+        self.des_joint_vel[:] = np.nan
+        self.des_joint_acc[:] = np.nan
+
+        self.des_c_pos[:] = np.nan
+        self.des_c_vel[:] = np.nan
+        self.des_quat[:] = np.nan
+        self.des_quat_vel[:] = np.nan
+
+    def get_command_from_inverse_dynamics(
+        self, target_j_acc, mj_calc_inv=False, robot_id=None, client_id=None
+    ):
+        raise NotImplementedError
+
+    def preprocessCommand(self, target_j_acc):
+        if len(target_j_acc.shape) == 2:
+            target_j_acc = target_j_acc.reshape((-1))
+        if target_j_acc.shape[0] != self.num_DoF:
+            raise ValueError(
+                "Specified motor command vector needs to be of size {} is {}\n".format(
+                    self.num_DoF, target_j_acc.shape[0]
+                )
+            )
+
+        self.finger_commands = self.fing_ctrl_step()
+        if self.use_inv_dyn:
+            target_j_acc = np.append(target_j_acc, self.finger_commands[0])
+            target_j_acc = np.append(target_j_acc, self.finger_commands[1])
+            self.des_joint_acc = target_j_acc.copy()
+            self.uff = self.get_command_from_inverse_dynamics(
+                target_j_acc, mj_calc_inv=True
+            )
+            self.uff[7:9] = self.finger_commands
+
+        else:
+            target_j_acc = np.append(target_j_acc, self.finger_commands[0])
+            target_j_acc = np.append(target_j_acc, self.finger_commands[1])
+            self.uff = target_j_acc
+            if self.gravity_comp:
+                comp_forces = self.get_command_from_inverse_dynamics(
+                    target_j_acc=np.zeros(9), mj_calc_inv=False
+                )
+                self.grav_terms = comp_forces.copy()
+                self.uff += comp_forces
+
+        if self.step_count == 0:
+            self.uff_last = 0
+
+        if self.clip_rate:
+            uff_diff = self.uff - self.uff_last
+            uff_diff = np.clip(uff_diff, -self.rate_limit, self.rate_limit)
+            self.uff = self.uff_last + uff_diff
+
+        if self.clip_actions:
+            self.uff[:7] = np.clip(self.uff[:7], -self.torque_limit, self.torque_limit)
+
+    def beam_to_cart_pos_and_quat(self, desiredPos, desiredQuat):
+        q = self.gotoCartPosQuatController.trajectory_generator.findJointPosition(
+            self, desiredPos=desiredPos, desiredQuat=desiredQuat
+        )
+        self.beam_to_joint_pos(q)
+
+    def beam_to_joint_pos(self, desiredJoints, run=True, log=False):
+        self.set_q(desiredJoints)
+        self.receiveState()
+        self.gotoJointController.resetTrajectory()
+        self.gotoCartPosQuatController.resetTrajectory()
+        self.gotoCartPosImpedanceController.resetTrajectory()
+        self.jointTrackingController.setSetPoint(desiredJoints)
+        if run:
+            self.jointTrackingController.executeControllerTimeSteps(self, 1, log)
+
+    def set_q(self, joints, robot_id=None, physicsClientId=None):
+
+        raise NotImplementedError
+
+    @property
+    def step_count(self):
+        return self.time_keeper.step_count
+
+    @property
+    def time_stamp(self):
+        return self.time_keeper.time_stamp
+
+    @property
+    def step_count(self):
+        return self.time_keeper.step_count

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/core/Scene.py

@@ -0,0 +1,319 @@
+import enum
+from abc import ABC, abstractmethod
+from typing import List
+
+import numpy as np
+
+import d3il.environments.d3il.d3il_sim.core.sim_object as sim_object
+import d3il.environments.d3il.d3il_sim.core.time_keeper as time_keeper
+import d3il.environments.d3il.d3il_sim.utils.geometric_transformation as geom_trans
+from d3il.environments.d3il.d3il_sim.core import Camera, Robots
+
+
+class Scene(ABC):
+
+    class RenderMode(enum.Enum):
+        BLIND = "blind"
+        OFFSCREEN = "offscreen"
+        HUMAN = "human"
+
+    def __init__(
+        self,
+        object_list=None,
+        dt=0.001,
+        render: RenderMode = RenderMode.HUMAN,
+        *args,
+        **kwargs
+    ):
+        if object_list is None:
+            object_list = []
+
+        self.obj_repo = sim_object.SimObjectRepository(object_list)
+
+        self.dt = dt
+        self.render_mode = render
+
+        self.setup_done = False
+
+        self._robots: List[Robots.RobotBase] = []
+        self.inhand_cam = None
+        self.cage_cam = None
+
+        self.step_callbacks = []
+        self.additional_loggers = []
+        self.time_keeper = time_keeper.TimeKeeper(self.dt)
+
+    def add_robot(self, robot):
+        self._robots.append(robot)
+
+    @property
+    def robots(self):
+        return self._robots
+
+    def register_callback(self, fn, **kwargs):
+
+        self.step_callbacks.append((fn, kwargs))
+
+    @property
+    @abstractmethod
+    def sim_name(self) -> str:
+
+        return "NONE"
+
+    @abstractmethod
+    def _setup_scene(self):
+
+        self._setup_objects(self.object_list)
+
+    @abstractmethod
+    def reset(self, obj_pos=None):
+
+        raise NotImplementedError
+
+    def start(self, robot_init_qpos: np.ndarray = None):
+
+        self._setup_scene()
+        self.load_robot_to_scene(robot_init_qpos=robot_init_qpos)
+        self.setup_done = True
+        self.reset()
+
+    @abstractmethod
+    def render(self):
+
+        raise NotImplementedError
+
+    def next_step(self, log=True):
+        for rb in self.robots:
+            rb.prepare_step()
+
+        self._sim_step()
+
+        self.time_keeper.tick()
+
+        for rb in self.robots:
+            rb.tick()
+            rb.receiveState()
+
+        for call_back, kwargs in self.step_callbacks:
+            call_back(**kwargs)
+        if log:
+            self.log_data()
+
+        self.render()
+
+    @abstractmethod
+    def _sim_step(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def load_robot_to_scene(self, robot_init_qpos: np.ndarray = None):
+
+        raise NotImplementedError
+
+    @abstractmethod
+    def _setup_objects(self, sim_objs: List[sim_object.SimObject]):
+
+        raise NotImplementedError
+
+    @abstractmethod
+    def _rt_add_object(self, sim_obj: sim_object.SimObject):
+
+        raise NotImplementedError
+
+    def add_object(self, sim_obj: sim_object.SimObject):
+
+        if not self.setup_done:
+            self.obj_repo.add_object(sim_obj)
+        else:
+            self._rt_add_object(sim_obj)
+
+    def list_objects(self):
+        return list(self._objects.keys())
+
+    def get_object(self, name: str = None, obj_id: int = None) -> sim_object.SimObject:
+
+        return self.obj_repo.get_object(name, obj_id)
+
+    def _query_pois(self, sim_obj: sim_object.SimObject, fn) -> np.ndarray:
+
+        if len(sim_obj.get_poi()) > 1:
+            return np.array([fn(poi, sim_obj) for poi in sim_obj.get_poi()])
+        else:
+            return fn(sim_obj.get_poi()[0], sim_obj)
+
+    @abstractmethod
+    def _get_obj_seg_id(self, obj_name: str) -> int:
+
+        raise NotImplementedError
+
+    def get_obj_seg_id(
+        self,
+        sim_obj: sim_object.SimObject = None,
+        obj_name: str = None,
+        obj_id: int = None,
+    ) -> int:
+
+        if obj_name is None:
+            if sim_obj is None:
+                sim_obj = self.get_object(obj_id=obj_id)
+            obj_name = sim_obj.name
+
+        return self._get_obj_seg_id(obj_name=obj_name)
+
+    @abstractmethod
+    def _get_obj_pos(self, poi, sim_obj: sim_object.SimObject) -> np.ndarray:
+
+        raise NotImplementedError
+
+    def get_obj_pos(
+        self,
+        sim_obj: sim_object.SimObject = None,
+        obj_name: str = None,
+        obj_id: int = None,
+    ) -> np.ndarray:
+
+        if sim_obj is None:
+            sim_obj = self.get_object(obj_name, obj_id)
+        return self._query_pois(sim_obj, self._get_obj_pos)
+
+    @abstractmethod
+    def _get_obj_quat(self, poi, sim_obj: sim_object.SimObject) -> np.ndarray:
+
+        raise NotImplementedError
+
+    def get_obj_quat(
+        self,
+        sim_obj: sim_object.SimObject = None,
+        obj_name: str = None,
+        obj_id: int = None,
+    ) -> np.ndarray:
+
+        if sim_obj is None:
+            sim_obj = self.get_object(obj_name, obj_id)
+        return self._query_pois(sim_obj, self._get_obj_quat)
+
+    def _get_obj_rot_mat(self, poi, sim_obj: sim_object.SimObject) -> np.ndarray:
+
+        quat = self._get_obj_quat(poi, sim_obj)
+        return geom_trans.mat2quat(quat)
+
+    def get_obj_rot_mat(
+        self,
+        sim_obj: sim_object.SimObject = None,
+        obj_name: str = None,
+        obj_id: int = None,
+    ) -> np.ndarray:
+
+        if sim_obj is None:
+            sim_obj = self.get_object(obj_name, obj_id)
+        return self._query_pois(sim_obj, self._get_obj_rot_mat)
+
+    @abstractmethod
+    def _set_obj_pos(self, new_pos, sim_obj: sim_object.SimObject):
+
+        raise NotImplementedError
+
+    def set_obj_pos(
+        self,
+        new_pos,
+        sim_obj: sim_object.SimObject = None,
+        obj_name: str = None,
+        obj_id: int = None,
+    ):
+
+        if sim_obj is None:
+            sim_obj = self.get_object(obj_name, obj_id)
+
+        self._set_obj_pos(new_pos, sim_obj)
+
+    @abstractmethod
+    def _set_obj_quat(self, new_quat, sim_obj: sim_object.SimObject) -> np.ndarray:
+
+        raise NotImplementedError
+
+    def set_obj_quat(
+        self,
+        new_quat,
+        sim_obj: sim_object.SimObject = None,
+        obj_name: str = None,
+        obj_id: int = None,
+    ):
+
+        if sim_obj is None:
+            sim_obj = self.get_object(obj_name, obj_id)
+
+        self._set_obj_quat(new_quat, sim_obj)
+
+    @abstractmethod
+    def _set_obj_pos_and_quat(
+        self, new_pos, new_quat, sim_obj: sim_object.SimObject
+    ) -> np.ndarray:
+
+        raise NotImplementedError
+
+    def set_obj_pos_and_quat(
+        self,
+        new_pos,
+        new_quat,
+        sim_obj: sim_object.SimObject = None,
+        obj_name: str = None,
+        obj_id: int = None,
+    ):
+
+        if sim_obj is None:
+            sim_obj = self.get_object(obj_name, obj_id)
+
+        self._set_obj_pos_and_quat(new_pos, new_quat, sim_obj)
+
+    @abstractmethod
+    def _remove_object(self, sim_obj: sim_object.SimObject):
+        raise NotImplementedError
+
+    def remove_object(
+        self,
+        sim_obj: sim_object.SimObject = None,
+        obj_name: str = None,
+        obj_id: int = None,
+    ):
+        if sim_obj is None:
+            sim_obj = self.get_object(obj_name, obj_id)
+
+        self._remove_object(sim_obj)
+        self.obj_repo.remove_object(sim_obj)
+
+    def get_cage_cam(self) -> Camera.Camera:
+        return self.cage_cam
+
+    def start_logging(self, duration: float = 300.0, **kwargs):
+
+        for rb in self.robots:
+            rb.start_logging(duration, **kwargs)
+
+        for logger in self.additional_loggers:
+            logger.start_logging(duration, **kwargs)
+
+    def stop_logging(self):
+
+        for rb in self.robots:
+            rb.stop_logging()
+
+        for logger in self.additional_loggers:
+            logger.stop_logging()
+
+    def add_logger(self, logger):
+
+        self.additional_loggers.append(logger)
+
+    def log_data(self):
+        for rb in self.robots:
+            rb.log_data()
+        for logger in self.additional_loggers:
+            logger.log_data()
+
+    @property
+    def step_count(self):
+        return self.time_keeper.step_count
+
+    @property
+    def time_stamp(self):
+        return self.time_keeper.time_stamp

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/core/__init__.py

@@ -0,0 +1,4 @@
+from .logger import *
+from .Model import *
+from .Robots import *
+from .Scene import *

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/core/logger.py

@@ -0,0 +1,867 @@
+import abc
+import datetime
+import os
+from enum import Flag, auto
+
+import numpy as np
+
+try:
+    import wandb
+except ImportError:
+    pass
+
+
+def reset_wandb_env():
+    exclude = {
+        "WANDB_PROJECT",
+        "WANDB_ENTITY",
+        "WANDB_API_KEY",
+        "WANDB_START_METHOD",
+    }
+    for k, v in os.environ.items():
+        if k.startswith("WANDB_") and k not in exclude:
+            del os.environ[k]
+
+
+class RobotPlotFlags(Flag):
+    JOINT_POS = auto()
+    JOINT_VEL = auto()
+    JOINT_ACC = auto()
+
+    CART_POS = auto()
+    CART_VEL = auto()
+    ORIENTATION = auto()
+    ORIENTATION_VEL = auto()
+
+    CART_POS_GLOBAL = auto()
+    CART_VEL_GLOBAL = auto()
+    ORIENTATION_GLOBAL = auto()
+    ORIENTATION_VEL_GLOBAL = auto()
+
+    COMMAND = auto()
+    TORQUES = auto()
+    GRAVITY = auto()
+    LOADS = auto()
+    TIME_STAMPS = auto()
+
+    GRIPPER_POS = auto()
+    GRIPPER_WIDTH = auto()
+    GRIPPER_FORCE = auto()
+    GRIPPER_VEL = auto()
+
+    MISC = auto()
+
+    GRIPPER = GRIPPER_WIDTH | GRIPPER_POS | GRIPPER_FORCE | GRIPPER_VEL
+    JOINTS = JOINT_POS | JOINT_ACC | JOINT_VEL
+    END_EFFECTOR = CART_POS | CART_VEL | ORIENTATION | ORIENTATION_VEL
+    END_EFFECTOR_GLOBAL = (
+        CART_POS_GLOBAL | CART_VEL_GLOBAL | ORIENTATION_GLOBAL | ORIENTATION_VEL_GLOBAL
+    )
+
+    FULL = GRIPPER | JOINTS | END_EFFECTOR | END_EFFECTOR_GLOBAL
+
+
+class ObjectPlotFlags(Flag):
+    POSITION = auto()
+    ORIENTATION = auto()
+    FULL = POSITION | ORIENTATION
+
+
+class CamPlotFlags(Flag):
+    IMAGE = auto()
+    FULL = IMAGE
+
+
+class LoggerBase(abc.ABC):
+    def __init__(self):
+
+        self.logged_time = 0
+        self.log_interval = None
+        self.log_duration = None
+        self.logged_time_steps = 0
+        self.max_time_steps = None
+        self.last_log_time = None
+        self.is_logging = False
+
+    def start_logging(self, duration: float = 300.0, **kwargs):
+
+        self.logged_time = 0
+        self.logged_time_steps = 0
+        self.log_duration = duration
+        self._start(duration, **kwargs)
+        self.is_logging = True
+
+    def log_data(self):
+
+        if self.is_logging and self._check_log_interval():
+            self._log()
+            self.logged_time_steps += 1
+            if (
+                self.max_time_steps is not None
+                and self.logged_time_steps >= self.max_time_steps
+            ) or self.logged_time >= self.log_duration:
+                self.stop_logging()
+
+    def stop_logging(self):
+
+        if self.is_logging:
+            self.is_logging = False
+            if self.logged_time_steps > 0:
+                self._stop()
+
+    @abc.abstractmethod
+    def _start(self, duration=600, **kwargs):
+        pass
+
+    @abc.abstractmethod
+    def _log(self):
+        pass
+
+    @abc.abstractmethod
+    def _stop(self):
+        pass
+
+    @abc.abstractmethod
+    def _check_log_interval(self) -> bool:
+        return True
+
+    @abc.abstractmethod
+    def reset(self):
+        pass
+
+
+class RobotLogger(LoggerBase):
+    def __init__(self, robot):
+        super().__init__()
+
+        self.robot = robot
+        self.plot_selection = RobotPlotFlags.FULL
+        self.use_wandb = False
+
+        self.reset()
+
+    def reset(self):
+        self.log_dict_full = None
+
+    @property
+    def log_dict(self):
+        log_dict = dict()
+
+        if RobotPlotFlags.JOINTS & self.plot_selection:
+            log_dict.update(
+                {
+                    "joint_pos": self.robot.current_j_pos,
+                    "joint_vel": self.robot.current_j_vel,
+                    "des_joint_pos": self.robot.des_joint_pos,
+                    "des_joint_vel": self.robot.des_joint_vel,
+                    "des_joint_acc": self.robot.des_joint_acc,
+                }
+            )
+
+        if RobotPlotFlags.GRIPPER & self.plot_selection:
+            log_dict.update(
+                {
+                    "finger_pos": self.robot.current_fing_pos,
+                    "finger_vel": self.robot.current_fing_vel,
+                    "des_finger_pos": self.robot.set_gripper_width,
+                    "gripper_width": self.robot.gripper_width,
+                    "uff": self.robot.uff,
+                }
+            )
+
+        if RobotPlotFlags.END_EFFECTOR & self.plot_selection:
+            log_dict.update(
+                {
+                    "cart_pos": self.robot.current_c_pos,
+                    "cart_vel": self.robot.current_c_vel,
+                    "cart_quat": self.robot.current_c_quat,
+                    "cart_quat_vel": self.robot.current_c_quat_vel,
+                    "des_c_pos": self.robot.des_c_pos,
+                    "des_c_vel": self.robot.des_c_vel,
+                    "des_quat": self.robot.des_quat,
+                    "des_quat_vel": self.robot.des_quat_vel,
+                }
+            )
+
+        if RobotPlotFlags.END_EFFECTOR_GLOBAL & self.plot_selection:
+            log_dict.update(
+                {
+                    "cart_pos_global": self.robot.current_c_pos_global,
+                    "cart_vel_global": self.robot.current_c_vel_global,
+                    "cart_quat_global": self.robot.current_c_quat_global,
+                    "cart_quat_vel_global": self.robot.current_c_quat_vel_global,
+                }
+            )
+
+        if RobotPlotFlags.COMMAND & self.plot_selection:
+            log_dict.update(
+                {
+                    "uff": self.robot.uff,
+                    "last_cmd": self.robot.last_cmd,
+                    "command": self.robot.command,
+                    "grav_terms": self.robot.grav_terms,
+                    "load": self.robot.current_load,
+                }
+            )
+
+        if RobotPlotFlags.MISC & self.plot_selection:
+            log_dict.update({"misc": self.robot.misc_data})
+
+        log_dict.update(
+            {
+                "time_stamp": self.robot.time_stamp,
+                "step": self.robot.step_count,
+                "wall_clock": self.robot.time_keeper.wall_clock,
+            }
+        )
+
+        return log_dict
+
+    def _start(
+        self,
+        duration=600,
+        log_interval=None,
+        plot_selection: RobotPlotFlags = RobotPlotFlags.FULL,
+        use_wandb=False,
+        **kwargs
+    ):
+
+        self.plot_selection = plot_selection
+        self.use_wandb = use_wandb
+
+        if log_interval is None:
+            self.log_interval = 1
+        else:
+            self.log_interval = log_interval
+
+        assert self.log_interval >= 1
+
+        self.max_time_steps = int(duration / self.log_interval)
+
+        if self.max_time_steps > 1800000:
+            self.max_time_steps = 1800000
+
+        self.log_dict_full = {
+            k: [None] * self.max_time_steps for k, v in self.log_dict.items()
+        }
+
+        self.initial_time_stamp = self.robot.step_count
+
+        if self.use_wandb:
+            reset_wandb_env()
+            now = datetime.datetime.now().strftime("%d_%m_%Y_%H_%M_%S")
+            job_name = "robot_log_" + now
+            self.run = wandb.init(
+                project="sf",
+                group="test_log",
+                name=job_name,
+                settings=wandb.Settings(_disable_stats=True),
+            )
+
+    def _log(self):
+        log_properties = self.log_dict
+
+        for k, v in log_properties.items():
+            self.log_dict_full[k][self.logged_time_steps] = v
+
+        self.last_log_time = log_properties["step"]
+
+        self.logged_time = self.last_log_time - self.initial_time_stamp
+
+        if self.use_wandb:
+            flat_log_dict = {
+                k + "_" + str(i): vv
+                for k, v in log_properties.items()
+                for i, vv in enumerate(np.atleast_1d(v))
+            }
+            self.run.log(flat_log_dict, step=self.logged_time_steps)
+
+    def _stop(self):
+
+        if self.logged_time_steps < self.max_time_steps:
+            for v in self.log_dict_full.values():
+                del v[self.logged_time_steps :]
+
+        if self.use_wandb:
+            if self.run is not None:
+                self.run.finish()
+
+        self.__dict__.update({k: np.array(v) for k, v in self.log_dict_full.items()})
+
+    def _check_log_interval(self) -> bool:
+
+        if (
+            self.last_log_time is None
+            or self.robot.step_count >= self.last_log_time + self.log_interval
+        ):
+            return True
+        return False
+
+    def plot(self, plot_selection=RobotPlotFlags.FULL, block=True):
+
+        import matplotlib.pyplot as plt
+
+        valid_plot_selections = plot_selection & self.plot_selection
+
+        if self.is_logging:
+            self.stop_logging()
+
+        self.time_stamp = self.time_stamp - self.time_stamp[0]
+
+        j1_limit_lower = -2.9671
+        j1_limit_upper = 2.9671
+        j2_limit_lower = -1.8326
+        j2_limit_upper = 1.8326
+        j3_limit_lower = -2.9671
+        j3_limit_upper = 2.9671
+        j4_limit_lower = -3.1416
+        j4_limit_upper = 0.0873
+        j5_limit_lower = -2.9671
+        j5_limit_upper = 2.9671
+        j6_limit_lower = -0.0873
+        j6_limit_upper = 3.8223
+        j7_limit_lower = -2.9671
+        j7_limit_upper = 2.9671
+
+        j1_limit_lower_safety = -2.8973
+        j1_limit_upper_safety = 2.8973
+        j2_limit_lower_safety = -1.7628
+        j2_limit_upper_safety = 1.7628
+        j3_limit_lower_safety = -2.8973
+        j3_limit_upper_safety = 2.8973
+        j4_limit_lower_safety = -3.0718
+        j4_limit_upper_safety = 0.0175
+        j5_limit_lower_safety = -2.8973
+        j5_limit_upper_safety = 2.8973
+        j6_limit_lower_safety = -0.0175
+        j6_limit_upper_safety = 3.7525
+        j7_limit_lower_safety = -2.8973
+        j7_limit_upper_safety = 2.8973
+
+        j1_vel_limit_lower = -2.3925
+        j1_vel_limit_upper = 2.3925
+        j2_vel_limit_lower = -2.3925
+        j2_vel_limit_upper = 2.3925
+        j3_vel_limit_lower = -2.3925
+        j3_vel_limit_upper = 2.3925
+        j4_vel_limit_lower = -2.3925
+        j4_vel_limit_upper = 2.3925
+        j5_vel_limit_lower = -2.8710
+        j5_vel_limit_upper = 2.8710
+        j6_vel_limit_lower = -2.8710
+        j6_vel_limit_upper = 2.8710
+        j7_vel_limit_lower = -2.8710
+        j7_vel_limit_upper = 2.8710
+
+        upper_limits = [
+            j1_limit_upper,
+            j2_limit_upper,
+            j3_limit_upper,
+            j4_limit_upper,
+            j5_limit_upper,
+            j6_limit_upper,
+            j7_limit_upper,
+        ]
+        lower_limits = [
+            j1_limit_lower,
+            j2_limit_lower,
+            j3_limit_lower,
+            j4_limit_lower,
+            j5_limit_lower,
+            j6_limit_lower,
+            j7_limit_lower,
+        ]
+
+        soft_upper_limits = [
+            j1_limit_upper_safety,
+            j2_limit_upper_safety,
+            j3_limit_upper_safety,
+            j4_limit_upper_safety,
+            j5_limit_upper_safety,
+            j6_limit_upper_safety,
+            j7_limit_upper_safety,
+        ]
+
+        soft_lower_limits = [
+            j1_limit_lower_safety,
+            j2_limit_lower_safety,
+            j3_limit_lower_safety,
+            j4_limit_lower_safety,
+            j5_limit_lower_safety,
+            j6_limit_lower_safety,
+            j7_limit_lower_safety,
+        ]
+
+        upper_limits_vel = [
+            j1_vel_limit_upper,
+            j2_vel_limit_upper,
+            j3_vel_limit_upper,
+            j4_vel_limit_upper,
+            j5_vel_limit_upper,
+            j6_vel_limit_upper,
+            j7_vel_limit_upper,
+        ]
+
+        lower_limits_vel = [
+            j1_vel_limit_lower,
+            j2_vel_limit_lower,
+            j3_vel_limit_lower,
+            j4_vel_limit_lower,
+            j5_vel_limit_lower,
+            j6_vel_limit_lower,
+            j7_vel_limit_lower,
+        ]
+
+        if RobotPlotFlags.JOINT_POS in valid_plot_selections:
+            joint_pos_fig = plt.figure()
+            for k in range(7):
+                plt.figure(joint_pos_fig.number)
+                plt.subplot(7, 1, k + 1)
+                plt.plot(self.time_stamp, self.joint_pos[:, k])
+                plt.plot(self.time_stamp, self.des_joint_pos[:, k], "r")
+            plt.subplot(7, 1, 1)
+            plt.title("joint positions ")
+
+        if RobotPlotFlags.JOINT_VEL in valid_plot_selections:
+            joint_vel_fig = plt.figure()
+            for k in range(7):
+                plt.figure(joint_vel_fig.number)
+                plt.subplot(7, 1, k + 1)
+                plt.plot(self.time_stamp, self.joint_vel[:, k])
+                plt.plot(self.time_stamp, self.des_joint_vel[:, k], "r")
+            plt.subplot(7, 1, 1)
+            plt.title("joint velocities ")
+
+        if RobotPlotFlags.JOINT_ACC in valid_plot_selections:
+            robot_acceleration = np.diff(self.joint_vel, 1, axis=0) / self.robot.dt
+            acceleration_fig = plt.figure()
+            for k in range(7):
+                plt.figure(acceleration_fig.number)
+                plt.subplot(7, 1, k + 1)
+                plt.plot(self.time_stamp[1:], robot_acceleration[:, k])
+                plt.plot(self.time_stamp, self.des_joint_acc[:, k], "r")
+            plt.subplot(7, 1, 1)
+            plt.title("joint accelerations ")
+
+        if RobotPlotFlags.COMMAND in valid_plot_selections:
+            ctrl_torques_fig = plt.figure()
+            for k in range(7):
+                plt.figure(ctrl_torques_fig.number)
+                plt.subplot(7, 1, k + 1)
+                plt.plot(self.time_stamp, self.command[:, k])
+            plt.subplot(7, 1, 1)
+            plt.title(" control torques (without gravity comp.)")
+
+        if RobotPlotFlags.TORQUES in valid_plot_selections:
+            clipped_torques_fig = plt.figure()
+            for k in range(7):
+                plt.figure(clipped_torques_fig.number)
+                plt.subplot(7, 1, k + 1)
+                plt.plot(self.time_stamp, self.uff[:, k])
+            plt.subplot(7, 1, 1)
+            plt.title(" clipped torques (with gravity comp.)")
+
+        if RobotPlotFlags.GRAVITY in valid_plot_selections:
+            grav_terms_fig = plt.figure()
+            for k in range(7):
+                plt.figure(grav_terms_fig.number)
+                plt.subplot(9, 1, k + 1)
+                plt.plot(self.time_stamp, self.grav_terms[:, k])
+            plt.subplot(7, 1, 1)
+            plt.title(" gravity (+coriolis) ")
+
+        if RobotPlotFlags.LOADS in valid_plot_selections:
+            loads_fig = plt.figure()
+            for k in range(7):
+                plt.figure(loads_fig.number)
+                plt.subplot(7, 1, k + 1)
+                plt.plot(self.time_stamp, self.load[:, k])
+            plt.subplot(7, 1, 1)
+            plt.title(" loads - forces in joints")
+
+        if RobotPlotFlags.GRIPPER_WIDTH in valid_plot_selections:
+            gripper_fig = plt.figure()
+            plt.figure(gripper_fig.number)
+            plt.plot(self.time_stamp, self.gripper_width)
+            plt.title(" gripper width ")
+
+        if RobotPlotFlags.GRIPPER_FORCE in valid_plot_selections:
+            gripper_force_fig = plt.figure()
+            plt.figure(gripper_force_fig.number)
+            plt.plot(self.time_stamp, self.uff[:, -2:])
+
+            plt.title(" gripper force (without contact forces)")
+
+        if RobotPlotFlags.GRIPPER_POS in valid_plot_selections:
+            gripper_pos_fig = plt.figure()
+            plt.figure(gripper_pos_fig.number)
+            plt.plot(self.time_stamp, self.finger_pos[:, -2:], label="actual")
+            plt.plot(self.time_stamp, self.des_finger_pos, label="desired")
+            plt.legend()
+            plt.title(" gripper pos ")
+
+        if RobotPlotFlags.GRIPPER_VEL in valid_plot_selections:
+            gripper_vel_fig = plt.figure()
+            plt.figure(gripper_vel_fig.number)
+            plt.plot(self.time_stamp, self.finger_vel[:, -2:])
+            plt.title(" gripper vel ")
+
+        if RobotPlotFlags.CART_POS in valid_plot_selections:
+            cart_pos_fig = plt.figure()
+            for j in range(3):
+                plt.figure(cart_pos_fig.number)
+                plt.subplot(3, 1, j + 1)
+                plt.plot(self.time_stamp, self.cart_pos[:, j])
+                plt.plot(self.time_stamp, self.des_c_pos[:, j], "r")
+            plt.subplot(3, 1, 1)
+            plt.title(" endeffector pos ")
+
+        if RobotPlotFlags.CART_VEL in valid_plot_selections:
+            cart_vel_fig = plt.figure()
+            for j in range(3):
+                plt.figure(cart_vel_fig.number)
+                plt.subplot(3, 1, j + 1)
+                plt.plot(self.time_stamp, self.cart_vel[:, j])
+                plt.plot(self.time_stamp, self.des_c_vel[:, j], "r")
+            plt.subplot(3, 1, 1)
+            plt.title(" endeffector vel ")
+
+        if RobotPlotFlags.ORIENTATION in valid_plot_selections:
+            orientation_fig = plt.figure()
+            for j in range(4):
+                plt.figure(orientation_fig.number)
+                plt.subplot(4, 1, j + 1)
+                plt.plot(self.time_stamp, self.cart_quat[:, j])
+                plt.plot(self.time_stamp, self.des_quat[:, j], "r")
+            plt.subplot(4, 1, 1)
+            plt.title(" endeffector orientation ")
+
+        if RobotPlotFlags.ORIENTATION_VEL in valid_plot_selections:
+            orientation_vel_fig = plt.figure()
+            for j in range(4):
+                plt.figure(orientation_vel_fig.number)
+                plt.subplot(4, 1, j + 1)
+                plt.plot(self.time_stamp, self.cart_quat_vel[:, j])
+                plt.plot(self.time_stamp, self.des_quat_vel[:, j], "r")
+            plt.subplot(4, 1, 1)
+            plt.title(" endeffector angular velocity ")
+
+        if RobotPlotFlags.CART_POS_GLOBAL in valid_plot_selections:
+            cart_pos_fig = plt.figure()
+            for j in range(3):
+                plt.figure(cart_pos_fig.number)
+                plt.subplot(3, 1, j + 1)
+                plt.plot(self.time_stamp, self.cart_pos_global[:, j])
+            plt.subplot(3, 1, 1)
+            plt.title(" endeffector pos global")
+
+        if RobotPlotFlags.CART_VEL_GLOBAL in valid_plot_selections:
+            cart_vel_fig = plt.figure()
+            for j in range(3):
+                plt.figure(cart_vel_fig.number)
+                plt.subplot(3, 1, j + 1)
+                plt.plot(self.time_stamp, self.cart_vel_global[:, j])
+            plt.subplot(3, 1, 1)
+            plt.title(" endeffector vel global")
+
+        if RobotPlotFlags.ORIENTATION_GLOBAL in valid_plot_selections:
+            orientation_fig = plt.figure()
+            for j in range(4):
+                plt.figure(orientation_fig.number)
+                plt.subplot(4, 1, j + 1)
+                plt.plot(self.time_stamp, self.cart_quat_global[:, j])
+            plt.subplot(4, 1, 1)
+            plt.title(" endeffector orientation global")
+
+        if RobotPlotFlags.ORIENTATION_VEL_GLOBAL in valid_plot_selections:
+            orientation_vel_fig = plt.figure()
+            for j in range(4):
+                plt.figure(orientation_vel_fig.number)
+                plt.subplot(4, 1, j + 1)
+                plt.plot(self.time_stamp, self.cart_quat_vel_global[:, j])
+            plt.subplot(4, 1, 1)
+            plt.title(" endeffector angular velocity global")
+
+        if RobotPlotFlags.TIME_STAMPS in valid_plot_selections:
+            timeStamp_fig = plt.figure()
+            plt.figure(timeStamp_fig.number)
+            plt.plot(np.diff(self.time_stamp))
+            plt.title("time stamp difference ")
+
+        if RobotPlotFlags.MISC in valid_plot_selections:
+            misc_fig = plt.figure()
+            numplots = self.misc.shape[1]
+            for j in range(numplots):
+                plt.figure(misc_fig.number)
+                plt.subplot(numplots, 1, j + 1)
+                plt.plot(self.misc[:, j])
+
+            plt.subplot(numplots, 1, 1)
+            plt.title("misc debug data ")
+        plt.show(block=block)
+
+
+class ObjectLogger(LoggerBase):
+
+    def __init__(self, scene, sim_object=None, obj_name=None):
+        super().__init__()
+        self.scene = scene
+
+        if obj_name is not None:
+            sim_object = scene.get_object(name=obj_name)
+
+        if sim_object is None:
+            raise ValueError("No Simobject defined.")
+        self.object = sim_object
+        self.plot_selection = ObjectPlotFlags.FULL
+
+    def reset(self):
+        self.log_dict_full = None
+        self.last_log_time = None
+
+    @property
+    def log_dict(self):
+        log_dict = dict()
+
+        if ObjectPlotFlags.POSITION in self.plot_selection:
+            log_dict.update({"pos": self.scene.get_obj_pos(self.object).flatten()})
+
+        if ObjectPlotFlags.ORIENTATION in self.plot_selection:
+            log_dict.update({"orientation": self.scene.get_obj_quat(self.object)})
+
+        log_dict.update({"time_stamp": self.scene.robots[0].time_stamp})
+        log_dict.update({"step": self.scene.robots[0].step_count})
+
+        return log_dict
+
+    def _start(
+        self,
+        duration=600,
+        log_interval=None,
+        obj_plot_selection: ObjectPlotFlags = ObjectPlotFlags.FULL,
+        use_wandb=False,
+        **kwargs
+    ):
+
+        self.plot_selection = obj_plot_selection
+        self.use_wandb = use_wandb
+
+        if log_interval is None:
+            self.log_interval = 1
+        else:
+            self.log_interval = log_interval
+
+        assert self.log_interval >= 1
+
+        self.max_time_steps = int(duration / self.log_interval)
+
+        if self.max_time_steps > 1800000:
+            self.max_time_steps = 1800000
+
+        self.log_dict_full = {
+            k: [None] * self.max_time_steps for k, v in self.log_dict.items()
+        }
+
+        self.initial_time_stamp = self.scene.robots[0].step_count
+
+        if self.use_wandb:
+            reset_wandb_env()
+            now = datetime.datetime.now().strftime("%d_%m_%Y_%H_%M_%S")
+            job_name = "robot_log_" + now
+            self.run = wandb.init(
+                project="sf",
+                group="test_log",
+                name=job_name,
+                settings=wandb.Settings(_disable_stats=True),
+            )
+
+    def _log(self):
+        log_properties = self.log_dict
+
+        for k, v in log_properties.items():
+            self.log_dict_full[k][self.logged_time_steps] = v
+
+        self.last_log_time = log_properties["step"]
+
+        self.logged_time = self.last_log_time - self.initial_time_stamp
+
+        if self.use_wandb:
+            flat_log_dict = {
+                k + "_" + str(i): vv
+                for k, v in log_properties.items()
+                for i, vv in enumerate(np.atleast_1d(v))
+            }
+            self.run.log(flat_log_dict, step=self.logged_time_steps)
+
+    def _stop(self):
+
+        if self.logged_time_steps < self.max_time_steps:
+            for v in self.log_dict_full.values():
+                del v[self.logged_time_steps :]
+
+        if self.use_wandb:
+            if self.run is not None:
+                self.run.finish()
+
+        self.__dict__.update({k: np.array(v) for k, v in self.log_dict_full.items()})
+
+    def _check_log_interval(self) -> bool:
+        if (
+            self.last_log_time is None
+            or self.scene.robots[0].step_count >= self.last_log_time + self.log_interval
+        ):
+            return True
+        return False
+
+    def plot(self, plot_selection=ObjectPlotFlags.FULL, block=True):
+        valid_plot_selections = plot_selection & self.plot_selection
+
+        import matplotlib.pyplot as plt
+
+        self.time_stamp = self.time_stamp - self.time_stamp[0]
+
+        if ObjectPlotFlags.POSITION in valid_plot_selections:
+            pos_fig = plt.figure()
+            for k in range(3):
+                plt.figure(pos_fig.number)
+                plt.subplot(3, 1, k + 1)
+                plt.plot(self.time_stamp, self.pos[:, k])
+            plt.subplot(3, 1, 1)
+            plt.title("object positions ")
+
+        if ObjectPlotFlags.ORIENTATION in valid_plot_selections:
+            orientation_fig = plt.figure()
+            for k in range(4):
+                plt.figure(orientation_fig.number)
+                plt.subplot(4, 1, k + 1)
+                plt.plot(self.time_stamp, self.orientation[:, k])
+            plt.title(" object orientation ")
+        plt.show(block=block)
+
+
+class CamLogger(LoggerBase):
+
+    def __init__(self, scene, camera):
+        super().__init__()
+
+        self.scene = scene
+        self.cam = camera
+
+        self.plot_selection = CamPlotFlags.FULL
+
+    def reset(self):
+        self.log_dict_full = None
+        self.last_log_time = None
+        self.cam.get_image(depth=False)
+
+    @property
+    def log_dict(self):
+        log_dict = dict()
+
+        if CamPlotFlags.IMAGE in self.plot_selection:
+
+            log_dict.update({"color_image": self.cam.get_image(depth=False)})
+
+        log_dict.update({"time_stamp": self.scene.robots[0].step_count})
+        log_dict.update({"step": self.scene.robots[0].step_count})
+
+        return log_dict
+
+    def _start(
+        self,
+        duration=600,
+        log_interval=None,
+        obj_plot_selection: CamPlotFlags = CamPlotFlags.FULL,
+        use_wandb=False,
+        **kwargs
+    ):
+
+        self.plot_selection = obj_plot_selection
+        self.use_wandb = use_wandb
+
+        if log_interval is None:
+            self.log_interval = 1
+        else:
+            self.log_interval = log_interval
+
+        assert self.log_interval >= 1
+
+        self.max_time_steps = int(duration / self.log_interval)
+
+        if self.max_time_steps > 1800000:
+            self.max_time_steps = 1800000
+
+        self.log_dict_full = {
+            k: [None] * self.max_time_steps for k, v in self.log_dict.items()
+        }
+
+        self.initial_time_stamp = self.scene.robots[0].step_count
+
+        if self.use_wandb:
+            reset_wandb_env()
+            now = datetime.datetime.now().strftime("%d_%m_%Y_%H_%M_%S")
+            job_name = "robot_log_" + now
+            self.run = wandb.init(
+                project="sf",
+                group="test_log",
+                name=job_name,
+                settings=wandb.Settings(_disable_stats=True),
+            )
+
+    def _log(self):
+
+        log_properties = self.log_dict
+
+        for k, v in log_properties.items():
+            self.log_dict_full[k][self.logged_time_steps] = v
+
+        self.last_log_time = log_properties["step"]
+
+        self.logged_time = self.last_log_time - self.initial_time_stamp
+
+        if self.use_wandb:
+            flat_log_dict = {
+                k + "_" + str(i): vv
+                for k, v in log_properties.items()
+                for i, vv in enumerate(np.atleast_1d(v))
+            }
+            self.run.log(flat_log_dict, step=self.logged_time_steps)
+
+    def _stop(self):
+
+        if self.logged_time_steps < self.max_time_steps:
+            for v in self.log_dict_full.values():
+                del v[self.logged_time_steps :]
+
+        if self.use_wandb:
+            if self.run is not None:
+                self.run.finish()
+
+        self.__dict__.update({k: np.array(v) for k, v in self.log_dict_full.items()})
+
+    def _check_log_interval(self) -> bool:
+
+        if (
+            self.last_log_time is None
+            or self.scene.robots[0].step_count >= self.last_log_time + self.log_interval
+        ):
+            return True
+        return False
+
+    def plot(self, plot_selection=CamPlotFlags.FULL, block=True):
+        valid_plot_selections = plot_selection & self.plot_selection
+
+        import matplotlib.pyplot as plt
+
+        self.time_stamp = self.time_stamp - self.time_stamp[0]
+
+        if CamPlotFlags.IMAGE in valid_plot_selections:
+            pos_fig = plt.figure()
+            for k in range(3):
+                plt.figure(pos_fig.number)
+                plt.subplot(3, 1, k + 1)
+                plt.plot(self.time_stamp, self.pos[:, k])
+            plt.subplot(3, 1, 1)
+            plt.title("camera images ")
+
+        plt.show(block=block)

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/core/sim_object/__init__.py

@@ -0,0 +1,2 @@
+from .sim_obj_repo import SimObjectRepository
+from .sim_object import IntelligentSimObject, SimObject

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/core/sim_object/sim_obj_repo.py

@@ -0,0 +1,46 @@
+from typing import List
+
+from d3il.environments.d3il.d3il_sim.utils.unique_dict import UniqueDict
+
+from .sim_object import SimObject
+
+
+class SimObjectRepository:
+
+    def __init__(self, obj_list=None) -> None:
+        self._name2id_map = UniqueDict(err_msg="Duplicate object name:")
+        self._id2name_map = UniqueDict(err_msg="Duplicate object id:")
+        self._objects = UniqueDict(err_msg="Duplicate object name:")
+
+        if obj_list is not None:
+            for obj in obj_list:
+                self.add_object(obj)
+
+    def add_object(self, sim_obj: SimObject) -> None:
+        self._objects[sim_obj.name] = sim_obj
+
+    def remove_object(self, sim_obj: SimObject) -> None:
+        del self._name2id_map[sim_obj.name]
+
+        del self._objects[sim_obj.name]
+
+    def register_obj_id(self, sim_obj: SimObject, obj_id: int):
+        sim_obj.obj_id = obj_id
+        self._name2id_map[sim_obj.name] = obj_id
+
+    def get_obj_list(self) -> List[SimObject]:
+        return list(self._objects.values())
+
+    def get_object(self, name: str = None, obj_id: int = None) -> SimObject:
+
+        if obj_id is not None:
+            name = self.get_name_from_id(obj_id)
+        return self._objects[name]
+
+    def get_id_from_name(self, name: str) -> int:
+
+        return self._name2id_map[name]
+
+    def get_name_from_id(self, obj_id: int) -> str:
+
+        return self._id2name_map[obj_id]

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/core/sim_object/sim_object.py

@@ -0,0 +1,49 @@
+import abc
+from typing import List
+
+
+class SimObject(abc.ABC):
+
+    GLOBAL_NAME_COUNTER = 0
+
+    def __init__(self, name: str = None, init_pos=None, init_quat=None):
+
+        if name is None:
+            name = "SIM_OBJ_{}".format(SimObject.GLOBAL_NAME_COUNTER)
+
+        self.name = name
+
+        self.init_pos = init_pos
+        self.init_quat = init_quat
+
+        self.obj_id = None
+
+        SimObject.GLOBAL_NAME_COUNTER += 1
+
+    @abc.abstractmethod
+    def get_poi(self) -> list:
+
+        return [self.name]
+
+
+class IntelligentSimObject(SimObject, abc.ABC):
+
+    def __init__(self, name: str = None, init_pos=None, init_quat=None):
+
+        super(IntelligentSimObject, self).__init__(name, init_pos, init_quat)
+
+        self.sim = None
+        self.sim_name = None
+
+    def register_sim(self, sim, sim_name):
+
+        self.sim = sim
+        self.sim_name = sim_name
+
+
+class DummyObject(SimObject):
+    def __init__(self, pois: List[str]):
+        self.pois = pois
+
+    def get_poi(self) -> list:
+        return self.pois

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/core/time_keeper.py

@@ -0,0 +1,13 @@
+import time
+
+
+class TimeKeeper:
+    def __init__(self, dt) -> None:
+        self.dt = dt
+        self.step_count = 0
+        self.time_stamp = 0
+
+    def tick(self):
+        self.step_count += 1
+        self.time_stamp += self.dt
+        self.wall_clock = time.time_ns()

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/gyms/gym_controllers.py

@@ -0,0 +1,162 @@
+import abc
+
+import gym
+import numpy as np
+
+import d3il.environments.d3il.d3il_sim.controllers.Controller as ctrl
+from d3il.environments.d3il.d3il_sim.controllers import CartPosQuatImpedenceController
+
+
+class GymController(abc.ABC):
+    def __init__(self, robot, controller):
+        self.controller = controller
+        self.robot = robot
+
+    @abc.abstractmethod
+    def action_space(self):
+        raise NotImplementedError
+
+    def set_action(self, action):
+        pass
+
+    def reset(self):
+        self.controller.reset()
+
+    def execute_action(self, n_time_steps):
+        self.controller.executeControllerTimeSteps(self.robot, timeSteps=n_time_steps)
+
+
+class GymJointPosController(GymController):
+    def __init__(self, robot):
+
+        super().__init__(robot, ctrl.JointPositionController())
+
+    def action_space(self):
+        low = [-2.8973, -1.7628, -2.8973, -3.0718, -2.8973, -0.0175, -2.8973]
+        high = [2.8973, 1.7628, 2.8973, -0.0698, 2.8973, 3.7525, 2.8973]
+
+        action_space = gym.spaces.Box(low=np.array(low), high=np.array(high))
+
+        return action_space
+
+    def set_action(self, action):
+        self.controller.setSetPoint(desired_pos=action)
+
+
+class GymJointVelController(GymController):
+    def __init__(self, robot):
+
+        super().__init__(robot, ctrl.JointVelocityController())
+
+    def action_space(self):
+        low = [-2.175, -2.1750, -2.1750, -2.1750, -2.6100, -2.6100, -2.6100]
+        high = [2.175, 2.1750, 2.1750, 2.1750, 2.6100, 2.6100, 2.6100]
+
+        action_space = gym.spaces.Box(low=np.array(low), high=np.array(high))
+
+        return action_space
+
+    def set_action(self, action):
+        self.controller.setSetPoint(desired_vel=action)
+
+
+class GymTorqueController(GymController):
+    def __init__(self, robot):
+        super().__init__(robot, ctrl.TorqueController())
+
+    def action_space(self):
+        low = [-2.175, -2.1750, -2.1750, -2.1750, -2.6100, -2.6100, -2.6100]
+        high = [2.175, 2.1750, 2.1750, 2.1750, 2.6100, 2.6100, 2.6100]
+
+        action_space = gym.spaces.Box(low=np.array(low), high=np.array(high))
+
+        return action_space
+
+    def set_action(self, action):
+        self.controller.setAction(action)
+
+
+class GymCartesianVelController(GymController):
+    def __init__(self, robot, fixed_orientation=None, max_cart_vel=0.05):
+        controller = CartPosQuatImpedenceController()
+        super(GymCartesianVelController, self).__init__(robot, controller)
+        self.fixed_orientation = fixed_orientation
+        self.max_cart_vel = max_cart_vel
+        self.max_cart_pos = np.array([0.7, 0.4, 0.6])
+        self.min_cart_pos = np.array([0.2, -0.4, 0.0])
+        self.desired_c_pos = robot.des_c_pos
+
+    def action_space(self):
+        low_xyz_gripper = np.array([-1.0, -1.0, -1.0]) * self.max_cart_vel
+        high_xyz_gripper = np.array([1.0, 1.0, 1.0]) * self.max_cart_vel
+
+        quat_low_wxyz = np.array([-1, -1, -1, -1])
+        quat_high_wxyz = np.array([1, 1, 1, 1])
+
+        low_limits = low_xyz_gripper
+        high_limits = high_xyz_gripper
+
+        if self.fixed_orientation is None:
+            low_limits.append(quat_low_wxyz)
+            high_limits.append(quat_high_wxyz)
+
+        action_space = gym.spaces.Box(low=low_limits, high=high_limits)
+        return action_space
+
+    def set_action(self, action):
+        self.desired_c_pos = np.array(action[:3])
+
+        if self.fixed_orientation is None:
+            desired_quat = np.array(action[3:])
+        else:
+            desired_quat = self.fixed_orientation
+
+        self.controller.setDesiredPos(
+            np.concatenate([self.desired_c_pos, desired_quat])
+        )
+
+
+class GymXYVelController(GymController):
+    def __init__(
+        self, robot, fixed_orientation=None, fixed_z_pos=0.4, max_cart_vel=0.05
+    ):
+        controller = CartPosQuatImpedenceController()
+        super(GymCartesianVelController, self).__init__(robot, controller)
+        self.fixed_orientation = fixed_orientation
+        self.fixed_z_pos = fixed_z_pos
+        self.max_cart_vel = max_cart_vel
+        self.max_cart_pos = np.array([0.7, 0.4])
+        self.min_cart_pos = np.array([0.2, -0.4])
+        self.desired_c_pos = robot.des_c_pos
+
+    def action_space(self):
+        low_xyz_gripper = np.array([-1.0, -1.0]) * self.max_cart_vel
+        high_xyz_gripper = np.array([1.0, 1.0]) * self.max_cart_vel
+
+        quat_low_wxyz = np.array([-1, -1, -1, -1])
+        quat_high_wxyz = np.array([1, 1, 1, 1])
+
+        low_limits = low_xyz_gripper
+        high_limits = high_xyz_gripper
+
+        if self.fixed_orientation is None:
+            low_limits.append(quat_low_wxyz)
+            high_limits.append(quat_high_wxyz)
+
+        action_space = gym.spaces.Box(low=low_limits, high=high_limits)
+        return action_space
+
+    def set_action(self, action):
+        self.desired_c_pos = np.array(self.robot.des_c_pos) + np.array(action[:2])
+        self.desired_c_pos = np.clip(
+            self.desired_c_pos, self.min_cart_pos, self.max_cart_pos
+        )
+
+        if self.fixed_orientation is None:
+            desired_quat = np.array(action[2:])
+        else:
+            desired_quat = self.fixed_orientation
+
+        self.controller.setDesiredPos(
+            np.concatenate([self.desired_c_pos, desired_quat])
+        )

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/gyms/gym_env_wrapper.py

@@ -0,0 +1,119 @@
+from abc import ABC, abstractmethod
+
+import gym
+from gym.utils import seeding
+import numpy as np
+
+from d3il.environments.d3il.d3il_sim.controllers.Controller import ControllerBase
+from d3il.environments.d3il.d3il_sim.core import Scene
+
+
+class GymEnvWrapper(gym.Env, ABC):
+
+    metadata = {"render.modes": ["human", "rgb_array"], "video.frames_per_second": 50}
+
+    def __init__(
+        self,
+        scene: Scene,
+        controller: ControllerBase,
+        max_steps_per_episode,
+        n_substeps,
+        debug: bool = False,
+    ):
+        self.scene = scene
+        self.robot = scene.robots[0]
+
+        self.controller = controller
+
+        self.max_steps_per_episode = max_steps_per_episode
+        self.n_substeps = n_substeps
+        self.env_step_counter = 0
+
+        self.episode = 0
+        self.terminated = False
+        self.debug = debug
+
+    def start(self):
+        self.scene.start()
+        self.controller.reset()
+
+    def seed(self, seed=None):
+        self.np_random, seed = seeding.np_random(seed)
+        return [seed]
+
+    def step(self, action, gripper_width=None, desired_vel=None, desired_acc=None):
+
+        if gripper_width is not None:
+            self.robot.set_gripper_width = gripper_width
+
+        self.robot.open_fingers()
+
+        self.controller.setSetPoint(action)
+        self.controller.executeControllerTimeSteps(
+            self.robot, self.n_substeps, block=False
+        )
+
+        for i in range(self.n_substeps):
+            self.scene.next_step()
+
+        observation = self.get_observation()
+        reward = self.get_reward()
+        done = self.is_finished()
+
+        debug_info = {}
+        if self.debug:
+            debug_info = self.debug_msg()
+
+        self.env_step_counter += 1
+        return observation, reward, done, debug_info
+
+    @abstractmethod
+    def get_observation(self) -> np.ndarray:
+        raise NotImplementedError
+
+    @abstractmethod
+    def get_reward(self):
+        raise NotImplementedError
+
+    @abstractmethod
+    def _check_early_termination(self) -> bool:
+        return self.terminated
+
+    def is_finished(self):
+        if (
+            self.terminated
+            or self._check_early_termination()
+            or self.env_step_counter >= self.max_steps_per_episode - 1
+        ):
+            return True
+        return False
+
+    def debug_msg(self) -> dict:
+        return {}
+
+    @abstractmethod
+    def _reset_env(self):
+        raise NotImplementedError
+
+    def reset(self):
+        self.terminated = False
+        self.env_step_counter = 0
+        self.episode += 1
+        obs = self._reset_env()
+
+        return obs
+
+    def robot_state(self):
+        self.robot.receiveState()
+
+        joint_pos = self.robot.current_j_pos
+        joint_vel = self.robot.current_j_vel
+
+        gripper_vel = self.robot.current_fing_vel
+        gripper_width = [self.robot.gripper_width]
+
+        tcp_pos = self.robot.current_c_pos
+        tcp_vel = self.robot.current_c_vel
+        tcp_quad = self.robot.current_c_quat
+
+        return tcp_pos

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/gyms/gym_utils/helpers.py

@@ -0,0 +1,7 @@
+import numpy as np
+
+
+def obj_distance(obj1_pos, obj2_pos):
+    dist = np.linalg.norm(obj1_pos - obj2_pos)
+    rel_dist = (obj1_pos - obj2_pos) / dist + 1e-8
+    return dist, rel_dist

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/gyms/gym_utils/rewards.py

@@ -0,0 +1,13 @@
+import numpy as np
+
+
+class Distance:
+    @staticmethod
+    def linear(start, goal):
+        return np.linalg.norm(start - goal, axis=-1)
+
+    @staticmethod
+    def quadratic(start, goal): ...
+
+    @staticmethod
+    def tanh(start, goal): ...

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/sims/SimFactory.py

@@ -0,0 +1,81 @@
+import abc
+from functools import wraps
+
+from d3il.environments.d3il.d3il_sim.core import Camera, RobotBase, Scene
+from d3il.environments.d3il.d3il_sim.sims.universal_sim.PrimitiveObjects import (
+    PrimitiveObject,
+)
+
+
+def _register_function_on_class(cls, fn):
+
+    @wraps(fn)
+    def self_wrapper(self, *args, **kwargs):
+        return fn(self, *args, **kwargs)
+
+    setattr(cls, fn.__name__, self_wrapper)
+
+
+class SimFactory(abc.ABC):
+
+    def __init__(self) -> None:
+        super().__init__()
+        fn = self.prim_loading()
+        if fn is not None:
+            _register_function_on_class(PrimitiveObject, fn)
+
+    @abc.abstractmethod
+    def create_scene(
+        self,
+        gin_config=None,
+        object_list: list = None,
+        dt: float = 0.001,
+        render: Scene.RenderMode = Scene.RenderMode.HUMAN,
+        *args,
+        **kwargs
+    ) -> Scene:
+        pass
+
+    @abc.abstractmethod
+    def create_robot(self, scene, *args, **kwargs) -> RobotBase:
+        pass
+
+    @abc.abstractmethod
+    def create_camera(
+        self,
+        name: str,
+        width: int = 1000,
+        height: int = 1000,
+        init_pos=None,
+        init_quat=None,
+        *args,
+        **kwargs
+    ) -> Camera.Camera:
+        pass
+
+    @abc.abstractmethod
+    def prim_loading(self):
+        return None
+
+    @property
+    def RenderMode(self):
+        return Scene.RenderMode
+
+
+class SimRepository:
+
+    _repository = {}
+
+    @classmethod
+    def register(cls, factory: SimFactory, sim_name: str):
+        if sim_name in cls._repository:
+            return
+        cls._repository[sim_name] = factory
+
+    @classmethod
+    def get_factory(cls, sim_name: str) -> SimFactory:
+        return cls._repository[sim_name]
+
+    @classmethod
+    def list_all_sims(cls):
+        return cls._repository.keys()

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/sims/__init__.py

@@ -0,0 +1,6 @@
+from .mj_beta import *
+
+try:
+    from .mj_beta import MjCamera
+except ImportError:
+    pass

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/sims/mj_beta/MjCamera.py

@@ -0,0 +1,196 @@
+from typing import Tuple
+from xml.etree import ElementTree as Et
+
+import numpy as np
+from mujoco import MjData, MjModel, mj_name2id, mjtObj
+
+from d3il.environments.d3il.d3il_sim.core.Camera import Camera
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.mj_utils.mj_helper import IncludeType
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.mj_utils.mj_render_singleton import (
+    render,
+)
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.MjLoadable import MjLoadable
+from d3il.environments.d3il.d3il_sim.utils.geometric_transformation import (
+    mat2posQuat,
+    posRotMat2TFMat,
+    quat2mat,
+)
+
+
+class MjCamera(Camera, MjLoadable):
+    def __init__(
+        self,
+        name: str,
+        width: int = 1000,
+        height: int = 1000,
+        init_pos=None,
+        init_quat=None,
+        near=0.01,
+        far=10,
+        *args,
+        **kwargs,
+    ):
+        super(MjCamera, self).__init__(
+            name, width, height, init_pos, init_quat, near=near, far=far
+        )
+        self.viewer = None
+        self.data = None
+        self.model = None
+
+    def register_sim(self, sim: Tuple[MjModel, MjData], sim_name: str):
+        self.model = sim[0]
+        self.data = sim[1]
+        super().register_sim(sim, sim_name)
+
+        extent = self.model.stat.extent
+        self.near = self.model.vis.map.znear * extent
+        self.far = self.model.vis.map.zfar * extent
+
+        cam_id = mj_name2id(m=self.model, name=self.name, type=mjtObj.mjOBJ_CAMERA)
+        self.fovy = self.model.cam_fovy[cam_id]
+        self.fovx = (
+            2
+            * np.arctan(
+                self.width
+                * 0.5
+                / (self.height * 0.5 / np.tan(self.fovy * np.pi / 360 / 2))
+            )
+            / np.pi
+            * 360
+        )
+
+        self.fx = (self.width / 2) / (np.tan(self.fovx * np.pi / 180 / 2))
+        self.fy = (self.height / 2) / (np.tan(self.fovy * np.pi / 180 / 2))
+        self.cx = self.width / 2
+        self.cy = self.height / 2
+
+        self.intrinsics = np.array(
+            [[self.fx, 0, self.cx], [0, self.fy, self.cy], [0, 0, 1]]
+        )
+
+        return
+
+    def _get_img_data(
+        self,
+        width: int = None,
+        height: int = None,
+        depth: bool = True,
+        denormalize_depth: bool = True,
+        segmentation: bool = False,
+    ):
+        if width is None:
+            width = self.width
+        if height is None:
+            height = self.height
+
+        rendered_images = render(
+            cam_name=self.name,
+            width=width,
+            height=height,
+            depth=depth,
+            segmentation=segmentation,
+            model=self.model,
+            data=self.data,
+        )
+
+        if depth:
+            color_img, depth_img = rendered_images
+            color_img = self.vertical_flip(color_img)
+            depth_img = self.vertical_flip(depth_img)
+            if denormalize_depth:
+                depth_img = self.denormalize_depth(depth_img)
+            return color_img, depth_img
+        elif segmentation:
+            seg_img = rendered_images
+            seg_img = self.vertical_flip(seg_img)
+            return seg_img
+        else:
+            color_img = rendered_images
+            color_img = self.vertical_flip(color_img)
+            return color_img
+
+    def get_cart_pos_quat(self):
+        cam_id = mj_name2id(m=self.model, name=self.name, type=mjtObj.mjOBJ_CAMERA)
+
+        cam_pos = self.data.cam_xpos[cam_id]
+        cam_mat = self.data.cam_xmat[cam_id]
+        c2b_r = cam_mat.reshape((3, 3))
+        b2w_r = quat2mat([0, 1, 0, 0])
+        c2w_r = np.matmul(c2b_r, b2w_r)
+        c2w = posRotMat2TFMat(cam_pos, c2w_r)
+        pos, quat = mat2posQuat(c2w)
+
+        return pos, quat
+
+    def mj_load(self):
+        pos_str = " ".join(map(str, self.init_pos))
+        quat_str = " ".join(map(str, self.init_quat))
+
+        object_body = Et.Element("body")
+        object_body.set("pos", pos_str)
+        object_body.set("quat", quat_str)
+        object_body.set("name", self.name)
+
+        cam = Et.SubElement(object_body, "camera")
+        cam.set("name", self.name)
+        cam.set("fovy", str(self.fovy))
+
+        return object_body, {}, IncludeType.WORLD_BODY
+
+    def get_segmentation(self, width=None, height=None, depth=True):
+        data = super(MjCamera, self).get_segmentation(width, height, depth)
+
+        if depth:
+            data, d = data
+
+        obj_types = data[..., 0]
+        obj_ids = self.model.geom_bodyid[data[..., 1]]
+
+        if depth:
+            return obj_ids, d
+        return obj_ids
+
+    def vertical_flip(self, img):
+        return np.flipud(img)
+
+
+class MjInhandCamera(MjCamera):
+
+    def __init__(
+        self,
+        robot_cam_id: str,
+        width: int = 96,
+        height: int = 96,
+        *args,
+        **kwargs,
+    ):
+        super().__init__(
+            robot_cam_id,
+            width,
+            height,
+            *args,
+            **kwargs,
+        )
+
+    def mj_load(self):
+
+        return None, None, IncludeType.VIRTUAL_INCLUDE
+
+
+class MjCageCam(MjCamera):
+
+    def __init__(self, width: int = 1000, height: int = 1000, *args, **kwargs):
+        super().__init__(
+            "rgbd_cage",
+            width,
+            height,
+            init_pos=[0.7, 0.1, 0.9],
+            init_quat=[
+                0.6830127,
+                0.1830127,
+                0.1830127,
+                0.683012,
+            ],
+            *args,
+            **kwargs,
+        )

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/sims/mj_beta/MjFactory.py

@@ -0,0 +1,50 @@
+import gin
+
+import d3il.environments.d3il.d3il_sim.sims.SimFactory as Sims
+from d3il.environments.d3il.d3il_sim.core import Camera, RobotBase, Scene
+from d3il.environments.d3il.d3il_sim.utils.sim_path import d3il_path
+
+from .MjCamera import MjCamera
+from .MjPrimLoader import mj_load
+from .MjRobot import MjRobot
+from .MjScene import MjScene
+
+
+class MjFactory(Sims.SimFactory):
+    def create_scene(
+        self,
+        gin_path=None,
+        object_list: list = None,
+        dt: float = 0.001,
+        render: Scene.RenderMode = Scene.RenderMode.HUMAN,
+        *args,
+        **kwargs
+    ) -> Scene:
+
+        if gin_path is None:
+            gin_path = d3il_path(
+                "d3il_sim/controllers/Config/mujoco_controller_config.gin"
+            )
+        gin.parse_config_file(gin_path)
+        return MjScene(object_list, dt, render, *args, **kwargs)
+
+    def create_robot(self, scene, *args, **kwargs) -> RobotBase:
+        return MjRobot(scene, *args, **kwargs)
+
+    def create_camera(
+        self,
+        name: str,
+        width: int = 1000,
+        height: int = 1000,
+        init_pos=None,
+        init_quat=None,
+        *args,
+        **kwargs
+    ) -> Camera:
+        return MjCamera(name, width, height, init_pos, init_quat)
+
+    def prim_loading(self):
+        return mj_load
+
+
+Sims.SimRepository.register(MjFactory(), "mj_beta")

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/sims/mj_beta/MjLoadable.py

@@ -0,0 +1,86 @@
+import abc
+import os
+import xml.etree.ElementTree as Et
+from typing import Tuple
+
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.mj_utils.mj_helper import IncludeType
+
+
+class MjLoadable(abc.ABC):
+
+    @abc.abstractmethod
+    def mj_load(self) -> Tuple[Et.Element, list, IncludeType]:
+        raise NotImplementedError
+
+
+class MjXmlLoadable(MjLoadable):
+
+    def __init__(self, full_xml_path, asset_path=None) -> None:
+        super().__init__()
+        self._full_xml_path = full_xml_path
+        self._asset_path = asset_path
+
+    @property
+    def loadable_dir(self):
+        return os.path.split(self._full_xml_path)[0]
+
+    @property
+    def asset_path(self):
+        if self._asset_path is not None:
+            return self._asset_path
+        return os.path.join(self.loadable_dir, "assets")
+
+    @property
+    def file_name(self):
+        return os.path.split(self._full_xml_path)[1]
+
+    def mj_load(self) -> Tuple[Et.Element, list, IncludeType]:
+        et_include = Et.Element("include")
+        et_include.set("file", self.file_name)
+        assets = {}
+
+        if os.path.isdir(self.asset_path):
+            for dirpath, dirnames, files in os.walk(self.asset_path):
+                for f in files:
+                    with open(os.path.join(dirpath, f), "rb") as file:
+                        assets[f] = file.read()
+
+        with open(os.path.join(self._full_xml_path), "rb") as file:
+            assets[self.file_name] = file.read()
+        return et_include, assets, IncludeType.FILE_INCLUDE
+
+
+class MjIncludeTemplate(MjXmlLoadable):
+    def __init__(self, full_xml_path, asset_path=None) -> None:
+        super().__init__(full_xml_path, asset_path)
+        self._tmp_filled_xml = None
+
+    @abc.abstractmethod
+    def modify_template(self, et: Et.ElementTree) -> Et.ElementTree:
+        pass
+
+    def mj_load(self) -> Tuple[Et.Element, list, IncludeType]:
+        inc, assets, include_type = super().mj_load()
+
+        obj = Et.parse(self._full_xml_path)
+        new_xml = self.modify_template(obj)
+
+        self._tmp_filled_xml = new_xml
+
+        et_include = Et.Element("include")
+        et_include.set("file", new_xml)
+        return et_include, assets, include_type
+
+    def cleanup(self):
+        if self._tmp_filled_xml is not None:
+            os.remove(self._tmp_filled_xml)
+
+
+class MjFreezable:
+    @abc.abstractmethod
+    def freeze(self, data, model):
+        pass
+
+    @abc.abstractmethod
+    def unfreeze(self, data, model):
+        pass

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/sims/mj_beta/MjPrimLoader.py

@@ -0,0 +1,38 @@
+import xml.etree.ElementTree as Et
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.mj_utils.mj_helper import IncludeType
+
+
+def mj_load(prim_obj, *args, **kwargs):
+    pos_str = " ".join(map(str, prim_obj.init_pos))
+    orientation_str = " ".join(map(str, prim_obj.init_quat))
+    mass_str = str(prim_obj.mass)
+    size_str = " ".join(map(str, prim_obj.size))
+    rgba_str = " ".join(map(str, prim_obj.rgba))
+
+    object_body = Et.Element("body")
+    object_body.set("name", prim_obj.name)
+    object_body.set("pos", pos_str)
+    object_body.set("quat", orientation_str)
+
+    geom = Et.SubElement(object_body, "geom")
+    geom.set("type", prim_obj.type.value)
+    geom.set("name", "{}:geom".format(prim_obj.name))
+    if prim_obj.mass:
+        geom.set("mass", mass_str)
+    geom.set("size", size_str)
+    geom.set("rgba", rgba_str)
+
+    if prim_obj.visual_only:
+        geom.set("contype", "0")
+        geom.set("conaffinity", "0")
+
+    if not prim_obj.static:
+        Et.SubElement(object_body, "freejoint")
+
+    if prim_obj.solimp is not None:
+        geom.set("solimp", " ".join(prim_obj.solimp))
+    if prim_obj.solref is not None:
+        geom.set("solref", " ".join(prim_obj.solref))
+
+    inc_type = IncludeType.WORLD_BODY
+    return object_body, {}, inc_type

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/sims/mj_beta/MjRobot.py

@@ -0,0 +1,287 @@
+import xml.etree.ElementTree as Et
+
+import mujoco
+import numpy as np
+
+from d3il.environments.d3il.d3il_sim.controllers.Controller import (
+    ModelBasedFeedforwardController,
+)
+from d3il.environments.d3il.d3il_sim.core import RobotBase
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.mj_utils.mj_helper import (
+    IncludeType,
+    get_body_xvelp,
+    get_body_xvelr,
+)
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.MjCamera import MjInhandCamera
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.MjLoadable import (
+    MjIncludeTemplate,
+    MjXmlLoadable,
+)
+from d3il.environments.d3il.d3il_sim.utils import sim_path
+
+
+class MjRobot(RobotBase, MjIncludeTemplate):
+    GLOBAL_MJ_ROBOT_COUNTER = 0
+
+    def __init__(
+        self,
+        scene,
+        dt=1e-3,
+        num_DoF=7,
+        base_position=None,
+        base_orientation=None,
+        gravity_comp=True,
+        clip_actions=False,
+        root=sim_path.D3IL_DIR,
+        xml_path=None,
+    ):
+        RobotBase.__init__(self, scene, dt, num_DoF, base_position, base_orientation)
+
+        if xml_path is None:
+            xml_path = sim_path.d3il_path("models/mj/robot/panda.xml")
+        MjXmlLoadable.__init__(self, xml_path)
+
+        self.clip_actions = clip_actions
+        self.gravity_comp = gravity_comp
+
+        self.joint_names = None
+        self.joint_indices = None
+        self.joint_act_indices = None
+
+        self.gripper_names = None
+        self.gripper_indices = None
+        self.gripper_act_indices = None
+
+        self.jointTrackingController = ModelBasedFeedforwardController()
+
+        self._mj_robot_id = MjRobot.GLOBAL_MJ_ROBOT_COUNTER
+        MjRobot.GLOBAL_MJ_ROBOT_COUNTER += 1
+        self.root = root
+
+        self.inhand_cam = MjInhandCamera(self.add_id2model_key("rgbd"))
+
+        self.qpos_freeze = None
+        self.qvel_freeze = None
+        self.ctrl_freeze = None
+
+    def _getJacobian_internal(self, q=None):
+        jac = np.zeros((6, 7))
+        if q is not None:
+
+            cur_sim_state = self.scene.sim.get_state()
+            qpos_idx = self.joint_indices
+            self.scene.sim.data.qpos[qpos_idx] = q
+            mujoco.mj_kinematics(self.scene.model, self.scene.sim.data)
+            mujoco.mj_comPos(self.scene.model, self.scene.sim.data)
+
+        tcp_id = self.add_id2model_key("tcp")
+        jac[:3, :] = self.scene.sim.data.get_body_jacp(tcp_id).reshape((3, -1))[
+            :, -9:-2
+        ]
+        jac[3:, :] = self.scene.sim.data.get_body_jacr(tcp_id).reshape((3, -1))[
+            :, -9:-2
+        ]
+
+        if q is not None:
+            self.scene.sim.set_state(cur_sim_state)
+        return jac
+
+    def _getForwardKinematics_internal(self, q=None):
+        if q is not None:
+            cur_sim_state = self.scene.sim.get_state()
+            qpos_idx = self.joint_indices
+            self.scene.data.qpos[qpos_idx] = q
+            mujoco.mj_kinematics(self.scene.model, self.scene.data)
+
+        tcp_id = self.add_id2model_key("tcp")
+        cart_pos = self.scene.data.get_body_xpos(tcp_id)
+        cart_or = self.scene.data.get_body_xquat(tcp_id)
+        if q is not None:
+            self.scene.sim.set_state(cur_sim_state)
+        return cart_pos, cart_or
+
+    def prepare_step(self):
+        self.command = self.activeController.getControl(self)
+
+        self.preprocessCommand(self.command)
+        self.scene.data.ctrl[self.joint_act_indices] = self.uff.copy()[: self.num_DoF]
+        self.scene.data.ctrl[self.gripper_act_indices] = self.finger_commands
+        self.receiveState()
+
+    def receiveState(self):
+        if self.joint_indices is None:
+            self._init_jnt_indices()
+
+        tcp_name = self.add_id2model_key("tcp")
+
+        self.current_j_pos = np.array(
+            [self.scene.data.joint(name).qpos.copy() for name in self.joint_names]
+        ).squeeze()
+        self.current_j_vel = np.array(
+            [self.scene.data.joint(name).qvel.copy() for name in self.joint_names]
+        ).squeeze()
+
+        test = self.scene.data.body(tcp_name)
+        self.current_c_pos_global = self.scene.data.body(tcp_name).xpos.copy()
+        self.current_c_vel_global = get_body_xvelp(
+            self.scene.model,
+            self.scene.data,
+            tcp_name,
+        )
+        self.current_c_quat_global = self.scene.data.body(tcp_name).xquat.copy()
+        self.current_c_quat_vel_global = np.zeros(4)
+
+        self.current_c_quat_vel_global[1:] = get_body_xvelr(
+            self.scene.model, self.scene.data, tcp_name
+        )
+        self.current_c_quat_vel_global *= 0.5 * self.current_c_quat_global
+
+        self.current_c_pos, self.current_c_quat = self._localize_cart_coords(
+            self.current_c_pos_global, self.current_c_quat_global
+        )
+        self.current_c_vel, _ = self._localize_cart_coords(
+            self.current_c_vel_global + self.base_position
+        )
+        _, self.current_c_quat_vel = self._localize_cart_coords(
+            self.base_position, self.current_c_quat_vel_global
+        )
+
+        self.current_fing_pos = np.array(
+            [self.scene.data.joint(name).qpos.copy() for name in self.gripper_names]
+        ).squeeze()
+        self.current_fing_vel = np.array(
+            [self.scene.data.joint(name).qvel.copy() for name in self.gripper_names]
+        ).squeeze()
+        self.gripper_width = self.current_fing_pos[-2] + self.current_fing_pos[-1]
+
+    def get_command_from_inverse_dynamics(self, target_j_acc, mj_calc_inv=False):
+        if mj_calc_inv:
+            self.scene.data.qacc[self.joint_vel_indices + self.gripper_vel_indices] = (
+                target_j_acc
+            )
+            mujoco.mj_inverse(self.scene.model, self.scene.sim.data)
+            return self.scene.data.qfrc_inverse[
+                self.joint_vel_indices + self.gripper_vel_indices
+            ]
+        else:
+            return self.scene.data.qfrc_bias[
+                self.joint_vel_indices + self.gripper_vel_indices
+            ]
+
+    def get_init_qpos(self):
+        return np.array(
+            [
+                3.57795216e-09,
+                1.74532920e-01,
+                3.30500960e-08,
+                -8.72664630e-01,
+                -1.14096181e-07,
+                1.22173047e00,
+                7.85398126e-01,
+            ]
+        )
+
+    def _init_jnt_indices(self):
+        n_joints = len(self.joint_names)
+        assert (
+            n_joints == self.num_DoF
+        ), "Error, found {} joints, but expected {}".format(n_joints, self.num_DoF)
+        assert (
+            len(self.gripper_names) == 2
+        ), "Error, found more gripper joints than expected."
+
+        self.joint_indices = []
+        self.joint_vel_indices = []
+        self.gripper_indices = []
+        self.gripper_vel_indices = []
+        for jnt_name in self.joint_names:
+            jnt_id = mujoco.mj_name2id(
+                self.scene.model, type=mujoco.mjtObj.mjOBJ_JOINT, name=jnt_name
+            )
+            self.joint_indices.append(self.scene.model.jnt_qposadr[jnt_id])
+            self.joint_vel_indices.append(self.scene.model.jnt_dofadr[jnt_id])
+        for grp_name in self.gripper_names:
+            grp_id = mujoco.mj_name2id(
+                self.scene.model, type=mujoco.mjtObj.mjOBJ_JOINT, name=grp_name
+            )
+            self.gripper_indices.append(self.scene.model.jnt_qposadr[grp_id])
+            self.gripper_vel_indices.append(self.scene.model.jnt_dofadr[grp_id])
+        self.joint_act_indices = [
+            mujoco.mj_name2id(
+                self.scene.model, type=mujoco.mjtObj.mjOBJ_ACTUATOR, name=name + "_act"
+            )
+            for name in self.joint_names
+        ]
+        self.gripper_act_indices = [
+            mujoco.mj_name2id(
+                self.scene.model, type=mujoco.mjtObj.mjOBJ_ACTUATOR, name=name + "_act"
+            )
+            for name in self.gripper_names
+        ]
+
+    def set_q(self, joint_pos):
+        if self.joint_indices is None:
+            self._init_jnt_indices()
+
+        qpos = self.scene.data.qpos.copy()
+        qpos[self.joint_indices] = joint_pos
+
+        qvel = self.scene.data.qvel.copy()
+        qvel[self.joint_vel_indices] = np.zeros(len(self.joint_vel_indices))
+        self.scene.set_state(qpos, qvel)
+
+    def add_id2model_key(self, model_key_id: str) -> str:
+        attrib_split = model_key_id.split("_")
+        attrib_split.insert(1, "rb{}".format(self._mj_robot_id))
+        attrib_id = "_".join(attrib_split)
+        return attrib_id
+
+    def modify_template(self, et: Et.ElementTree) -> str:
+        self.joint_names = None
+        self.joint_indices = None
+        self.joint_act_indices = None
+
+        self.gripper_names = None
+        self.gripper_indices = None
+        self.gripper_act_indices = None
+
+        for node in et.iter():
+            for attrib in [
+                "body1",
+                "body2",
+                "name",
+                "class",
+                "childclass",
+                "mesh",
+                "site",
+                "joint",
+            ]:
+                if attrib in node.attrib:
+                    model_key_id = node.get(attrib)
+                    attrib_id = self.add_id2model_key(model_key_id)
+                    node.set(attrib, attrib_id)
+
+        wb = et.find("worldbody")
+        body_root = wb.find("body")
+        body_root.set("pos", " ".join(map(str, self.base_position)))
+        body_root.set("quat", " ".join(map(str, self.base_orientation)))
+
+        self.joint_names = []
+        self.gripper_names = []
+        for jnt in wb.iter("joint"):
+            if jnt.get("name") is not None and jnt.get("name").startswith(
+                self.add_id2model_key("panda_joint")
+            ):
+                self.joint_names.append(jnt.get("name"))
+            if jnt.get("name") is not None and jnt.get("name").startswith(
+                self.add_id2model_key("panda_finger_joint")
+            ):
+                self.gripper_names.append(jnt.get("name"))
+
+        import uuid
+
+        new_path = sim_path.d3il_path(
+            f"models/mj/robot/panda_tmp_rb{self._mj_robot_id}_{uuid.uuid1()}.xml"
+        )
+        et.write(new_path)
+        return new_path

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/sims/mj_beta/MjScene.py

@@ -0,0 +1,282 @@
+import itertools
+import logging
+import os
+from typing import List, Tuple
+
+import mujoco
+import numpy as np
+from mujoco import MjData, MjModel, mj_name2id, mjtObj
+
+from d3il.environments.d3il.d3il_sim.core.Scene import Scene
+from d3il.environments.d3il.d3il_sim.core.sim_object.sim_object import (
+    IntelligentSimObject,
+    SimObject,
+)
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.mj_utils.mj_render_singleton import (
+    reset_singleton as reset_render_singleton,
+)
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.mj_utils.mj_renderer import Viewer
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.mj_utils.mj_scene_object import (
+    MjSurrounding,
+)
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.mj_utils.mj_scene_parser import (
+    MjSceneParser,
+)
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.MjCamera import MjCageCam, MjCamera
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.MjLoadable import (
+    MjFreezable,
+    MjXmlLoadable,
+)
+from d3il.environments.d3il.d3il_sim.utils import sim_path
+
+
+class MjScene(Scene):
+    def __init__(
+        self,
+        object_list=None,
+        dt=0.001,
+        render=Scene.RenderMode.HUMAN,
+        surrounding=None,
+        random_env=False,
+        main_xml_path=None,
+    ):
+
+        super(MjScene, self).__init__(object_list=object_list, dt=dt, render=render)
+
+        self.data: MjData = None
+        self.model: MjModel = None
+        self.viewer: Viewer = None
+        self.rt_obj = []
+
+        self.init_qpos, self.init_qvel = None, None
+        self.random_env = random_env
+
+        if surrounding is None:
+            surrounding = sim_path.d3il_path(
+                "models/mujoco/surroundings/lab_surrounding.xml"
+            )
+
+        self.surrounding = MjSurrounding(
+            surrounding_name=os.path.split(os.path.splitext(surrounding)[0])[-1],
+            root=os.path.dirname(surrounding),
+        )
+
+        self.mj_scene_parser = MjSceneParser(main_xml_path=main_xml_path)
+
+        self.cage_cam = MjCageCam()
+        self.add_object(self.cage_cam)
+
+    @property
+    def sim_name(self) -> str:
+        return "mj"
+
+    def _setup_scene(self):
+        for rb in self.robots:
+            self.add_object(rb.inhand_cam)
+
+        self.model, self.data = self.mj_scene_parser.create_scene(
+            self.robots, self.surrounding, self.obj_repo.get_obj_list(), self.dt
+        )
+
+        self.viewer = None
+        if self.render_mode == self.RenderMode.HUMAN:
+            self.viewer = Viewer(self.model, self.data)
+        self.init_qpos = self.data.qpos.copy()
+        self.init_qvel = np.zeros(self.data.qvel.shape)
+
+        self._setup_objects(self.obj_repo.get_obj_list())
+
+    def set_views(self, views):
+        if self.render_mode == self.RenderMode.HUMAN:
+            cam_ids = []
+            for v in views:
+                cam_ids.append(
+                    mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_CAMERA, v.name)
+                )
+            self.viewer.set_cameras(cam_ids)
+
+    def load_robot_to_scene(self, robot_init_qpos: np.ndarray = None):
+        if robot_init_qpos is None:
+            robot_init_qpos = np.stack([robot.get_init_qpos() for robot in self.robots])
+        else:
+            robot_init_qpos = np.asarray(robot_init_qpos)
+        if robot_init_qpos.ndim == 1:
+            robot_init_qpos = np.expand_dims(robot_init_qpos, 0)
+
+        for i in range(len(self.robots)):
+            self.robots[i].beam_to_joint_pos(robot_init_qpos[i], run=False)
+        self.init_qpos = self.data.qpos.copy()
+
+    def _sim_step(self):
+        mujoco.mj_step(self.model, self.data)
+
+    def sim_steps(self, n_steps):
+        mujoco.mj_step(self.model, self.data, n_steps)
+
+    def render(self):
+        if self.render_mode == self.RenderMode.HUMAN:
+            self.viewer.render()
+
+    def reset(self, obj_pos=None):
+        if obj_pos is None:
+            obj_pos = []
+
+        for rb in self.robots:
+            rb.reset()
+
+        mujoco.mj_resetData(self.model, self.data)
+
+        qpos = self.data.qpos.copy()
+        qpos[:] = self.init_qpos
+
+        qvel = np.zeros(self.data.qvel.shape)
+        self.set_state(qpos, qvel)
+
+        for obj, new_pos in obj_pos:
+            self.set_obj_pos(new_pos, obj)
+
+        mujoco.mj_forward(self.model, self.data)
+
+        for rb in self.robots:
+            rb.receiveState()
+
+    def add_object_rt(self, new_obj: MjXmlLoadable):
+        for obj in itertools.chain(self.obj_repo.get_obj_list(), self.robots):
+            if isinstance(obj, MjFreezable):
+                obj.freeze(self.model, self.data)
+
+        self.obj_repo.add_object(new_obj)
+        self.obj_repo.register_obj_id(new_obj, len(self.obj_repo.get_obj_list()))
+        if isinstance(new_obj, IntelligentSimObject):
+            new_obj.register_sim((self.model, self.data), self.sim_name)
+
+        self.mj_scene_parser = MjSceneParser()
+        self.model, self.data = self.mj_scene_parser.create_scene(
+            self.robots, self.surrounding, self.obj_repo.get_obj_list(), self.dt
+        )
+
+        reset_render_singleton()
+        if self.render_mode == self.RenderMode.HUMAN:
+            self.viewer.reinit(self.model, self.data)
+
+        self.init_qpos = self.data.qpos.copy()
+        self.init_qvel = np.zeros(self.data.qvel.shape)
+        self.load_robot_to_scene()
+
+        for obj in itertools.chain(self.obj_repo.get_obj_list(), self.robots):
+            if isinstance(obj, MjFreezable):
+                obj.unfreeze(self.model, self.data)
+        self.rt_obj.append(new_obj)
+
+    def obj_reset(self):
+        for obj in self.rt_obj:
+            self.obj_repo.remove_object(obj)
+
+        for obj in itertools.chain(self.obj_repo.get_obj_list(), self.robots):
+            if isinstance(obj, MjFreezable):
+                obj.freeze(self.model, self.data)
+
+        self.mj_scene_parser = MjSceneParser()
+        self.model, self.data = self.mj_scene_parser.create_scene(
+            self.robots, self.surrounding, self.obj_repo.get_obj_list(), self.dt
+        )
+        if self.render_mode == self.RenderMode.HUMAN:
+            self.viewer.reinit(self.model, self.data)
+
+        self.init_qpos = self.data.qpos.copy()
+        self.init_qvel = np.zeros(self.data.qvel.shape)
+        self.load_robot_to_scene()
+
+        for obj in itertools.chain(self.obj_repo.get_obj_list(), self.robots):
+            if isinstance(obj, MjFreezable):
+                obj.unfreeze(self.model, self.data)
+        self.rt_objects = []
+
+    def start_recording(self, nframes):
+        if self.render_mode == self.RenderMode.HUMAN:
+            self.viewer.start_recording(nframes=nframes)
+
+    def _setup_objects(self, sim_objs: List[SimObject]):
+        for i, obj in enumerate(sim_objs):
+            self.obj_repo.register_obj_id(obj, i)
+
+            if isinstance(obj, IntelligentSimObject):
+                obj.register_sim((self.model, self.data), self.sim_name)
+
+    def _rt_add_object(self, sim_obj: SimObject):
+        raise RuntimeError(
+            "Adding objects in MuJoCo only possible prior to scene setup."
+        )
+
+    def _get_obj_seg_id(self, obj_name: str):
+        return mj_name2id(self.model, 1, obj_name)
+
+    def _get_obj_pos(self, poi, sim_obj: SimObject):
+        pos, quat = self._get_obj_pos_and_quat(poi=poi, sim_obj=sim_obj)
+        return pos
+
+    def _get_obj_quat(self, poi, sim_obj: SimObject):
+        pos, quat = self._get_obj_pos_and_quat(poi=poi, sim_obj=sim_obj)
+        return quat
+
+    def _get_obj_pos_and_quat(self, poi, sim_obj: SimObject):
+        body_id = mj_name2id(self.model, mjtObj.mjOBJ_BODY, sim_obj.name)
+        body_jnt_addr = self.model.body_jntadr[body_id]
+        qposadr = self.model.jnt_qposadr[body_jnt_addr]
+
+        if body_jnt_addr == -1:
+            pos = self.model.body_pos[body_id]
+            quat = self.model.body_quat[body_id]
+        else:
+            pos = self.data.qpos[qposadr : qposadr + 3]
+            quat = self.data.qpos[qposadr + 3 : qposadr + 7]
+
+        return pos.copy(), quat.copy()
+
+    def _set_obj_pos(self, new_pos, sim_obj: SimObject):
+        self._set_obj_pos_and_quat(new_pos, None, sim_obj=sim_obj)
+
+    def _set_obj_quat(self, new_quat, sim_obj: SimObject):
+        self._set_obj_pos_and_quat(None, new_quat, sim_obj=sim_obj)
+
+    def _set_obj_pos_and_quat(self, new_pos, new_quat, sim_obj: SimObject):
+        if new_pos is None and new_quat is None:
+            logging.getLogger(__name__).warning(
+                "Expected at least either a new position or quaternion for set_obj_pos_and_quat"
+            )
+            return
+
+        body_id = mj_name2id(self.model, mjtObj.mjOBJ_BODY, sim_obj.name)
+        body_jnt_addr = self.model.body_jntadr[body_id]
+        qposadr = self.model.jnt_qposadr[body_jnt_addr]
+
+        if new_pos is not None:
+            assert len(new_pos) == 3, print(
+                f"Expected a positions list of 3 values, got {len(new_pos)}"
+            )
+
+            if body_jnt_addr == -1:
+                self.model.body_pos[body_id] = new_pos
+            else:
+                self.data.qpos[qposadr : qposadr + 3] = new_pos
+
+        if new_quat is not None:
+            assert len(new_quat) == 4, print(
+                f"Expected a quaternions list of 4 values, got {len(new_quat)}"
+            )
+            if body_jnt_addr == -1:
+                self.model.body_quat[body_id] = new_quat
+            else:
+                self.data.qpos[qposadr + 3 : qposadr + 7] = new_quat
+
+    def _remove_object(self, sim_obj: SimObject):
+        raise RuntimeError(
+            "Removing objects in MuJoCo only possible prior to scene setup."
+        )
+
+    def set_state(self, qpos, qvel):
+        self.data.qpos[:] = np.copy(qpos)
+        self.data.qvel[:] = np.copy(qvel)
+        if self.model.na == 0:
+            self.data.act[:] = None
+        mujoco.mj_forward(self.model, self.data)

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/sims/mj_beta/__init__.py

@@ -0,0 +1,14 @@
+import logging
+
+try:
+    import mujoco
+
+    from .MjFactory import *
+    from .MjCamera import MjCamera
+except ImportError as e:
+    logging.getLogger(__name__).info(e)
+    logging.getLogger(__name__).info(
+        "No MuJoCo bindings are(new bindings from DeepMind) installed.",
+        "Install MuJoCo >= 2.1.5 and load bindings via pip, conda or mamba",
+    )
+    pass

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/sims/mj_beta/mj_utils/mj_helper.py

@@ -0,0 +1,75 @@
+from enum import Enum
+
+import mujoco
+import numpy as np
+
+
+class IncludeType(Enum):
+    FILE_INCLUDE = 0
+    MJ_INCLUDE = 1
+    WORLD_BODY = 2
+    VIRTUAL_INCLUDE = 3
+
+
+def change_domain(value, in_low, in_high, out_low, out_high):
+    return (out_high - out_low) * ((value - in_low) / (in_high - in_low)) + out_low
+
+
+def obj_position(data, obj_id: str):
+    return data.get_body_xpos(obj_id)
+
+
+def obj_distance(data, obj1: str, obj2: str):
+    obj1_pos = obj_position(data, obj1)
+    obj2_pos = obj_position(data, obj2)
+    dist = np.linalg.norm(obj1_pos - obj2_pos)
+    rel_dist = (obj1_pos - obj2_pos) / dist + 1e-8
+    return dist, rel_dist
+
+
+def reset_mocap2body_xpos(model, data):
+
+    if model.eq_type is None or model.eq_obj1id is None or model.eq_obj2id is None:
+        return
+    for eq_type, obj1_id, obj2_id in zip(
+        model.eq_type, model.eq_obj1id, model.eq_obj2id
+    ):
+        if eq_type != mujoco.const.EQ_WELD:
+            continue
+
+        mocap_id = model.body_mocapid[obj1_id]
+        if mocap_id != -1:
+            body_idx = obj2_id
+        else:
+            mocap_id = model.body_mocapid[obj2_id]
+            body_idx = obj1_id
+
+        assert mocap_id != -1
+        data.mocap_pos[mocap_id][:] = data.body_xpos[body_idx]
+        data.mocap_quat[mocap_id][:] = data.body_xquat[body_idx]
+
+
+def get_body_jacr(model, data, name):
+    id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, name)
+    jacr = np.zeros((3, model.nv))
+    mujoco.mj_jacBody(model, data, None, jacr, body=id)
+    return jacr
+
+
+def get_body_jacp(model, data, name):
+    id = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, name)
+    jacp = np.zeros((3, model.nv))
+    mujoco.mj_jacBody(model, data, jacp, None, body=id)
+    return jacp
+
+
+def get_body_xvelp(model, data, name):
+    jacp = get_body_jacp(model, data, name).reshape((3, model.nv))
+    xvelp = np.dot(jacp, data.qvel.copy())
+    return xvelp
+
+
+def get_body_xvelr(model, data, name):
+    jacr = get_body_jacr(model, data, name).reshape((3, model.nv))
+    xvelr = np.dot(jacr, data.qvel.copy())
+    return xvelr

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/sims/mj_beta/mj_utils/mj_render_singleton.py

@@ -0,0 +1,44 @@
+from mujoco import MjData, MjModel, mj_name2id, mjtObj
+
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.mj_utils.mj_renderer import (
+    RenderContext,
+    RenderContextOffscreen,
+)
+
+__RENDER_CTX_MAP = {}
+
+
+def reset_singleton():
+    global __RENDER_CTX_MAP
+    keys = list(__RENDER_CTX_MAP.keys())
+    for key in keys:
+        del __RENDER_CTX_MAP[key]
+    __RENDER_CTX_MAP = {}
+
+
+def get_renderer(
+    name: str, width: int, height: int, model: MjModel, data: MjData
+) -> RenderContextOffscreen:
+    global __RENDER_CTX_MAP
+    if name not in __RENDER_CTX_MAP:
+        ctx = RenderContextOffscreen(width, height, model, data)
+        __RENDER_CTX_MAP[name] = ctx
+    ctx = __RENDER_CTX_MAP[name]
+    return ctx
+
+
+def render(
+    cam_name: str,
+    width: int,
+    height: int,
+    model: MjModel,
+    data: MjData,
+    depth: bool,
+    segmentation: bool,
+):
+
+    ctx = get_renderer(cam_name, width, height, model, data)
+    ctx.opengl_context.make_current()
+    cam_id = mj_name2id(m=model, name=cam_name, type=mjtObj.mjOBJ_CAMERA)
+    ctx.render(width, height, cam_id, segmentation=segmentation)
+    return ctx.read_pixels(width, height, depth=depth, segmentation=segmentation)

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/sims/mj_beta/mj_utils/mj_renderer.py

@@ -0,0 +1,593 @@
+import collections
+import os
+import time
+from threading import Lock
+from typing import List
+
+import glfw
+import imageio
+import mujoco
+import numpy as np
+
+
+def _import_egl(width, height):
+    from mujoco.egl import GLContext
+
+    return GLContext(width, height)
+
+
+def _import_glfw(width, height):
+    from mujoco.glfw import GLContext
+
+    return GLContext(width, height)
+
+
+def _import_osmesa(width, height):
+    from mujoco.osmesa import GLContext
+
+    return GLContext(width, height)
+
+
+_ALL_RENDERERS = collections.OrderedDict(
+    [
+        ("glfw", _import_glfw),
+        ("egl", _import_egl),
+        ("osmesa", _import_osmesa),
+    ]
+)
+
+
+class RenderContext:
+
+    def __init__(self, model, data, offscreen=True):
+
+        self.model = model
+        self.data = data
+        self.offscreen = offscreen
+        max_geom = 1000
+
+        mujoco.mj_forward(self.model, self.data)
+
+        self.scn = mujoco.MjvScene(self.model, max_geom)
+        self.cam = mujoco.MjvCamera()
+        self.vopt = mujoco.MjvOption()
+        self.pert = mujoco.MjvPerturb()
+        self.con = mujoco.MjrContext(self.model, mujoco.mjtFontScale.mjFONTSCALE_150)
+
+        self._markers = []
+        self._overlays = {}
+
+        self._init_camera()
+        self._set_mujoco_buffers()
+
+    def _set_mujoco_buffers(self):
+        if self.offscreen:
+            mujoco.mjr_setBuffer(mujoco.mjtFramebuffer.mjFB_OFFSCREEN, self.con)
+            if self.con.currentBuffer != mujoco.mjtFramebuffer.mjFB_OFFSCREEN:
+                raise RuntimeError("Offscreen rendering not supported")
+        else:
+            mujoco.mjr_setBuffer(mujoco.mjtFramebuffer.mjFB_WINDOW, self.con)
+            if self.con.currentBuffer != mujoco.mjtFramebuffer.mjFB_WINDOW:
+                raise RuntimeError("Window rendering not supported")
+
+    def update_offscreen_size(self, width, height):
+        if width != self.con.offWidth or height != self.con.offHeight:
+            self.model.vis.global_.offwidth = width
+            self.model.vis.global_.offheight = height
+            self.con.free()
+            self.con = mujoco.MjrContext(
+                self.model, mujoco.mjtFontScale.mjFONTSCALE_150
+            )
+            self._set_mujoco_buffers()
+
+    def render(self, width, height, camera_id=None, segmentation=False):
+        rect = mujoco.MjrRect(left=0, bottom=0, width=width, height=height)
+
+        if width != self.con.offWidth or height != self.con.offHeight:
+            new_width = max(width, self.model.vis.global_.offwidth)
+            new_height = max(height, self.model.vis.global_.offheight)
+            self.update_offscreen_size(new_width, new_height)
+
+        if camera_id is not None:
+            if camera_id == -1:
+                self.cam.type = mujoco.mjtCamera.mjCAMERA_FREE
+            else:
+                self.cam.type = mujoco.mjtCamera.mjCAMERA_FIXED
+            self.cam.fixedcamid = camera_id
+
+        mujoco.mjv_updateScene(
+            self.model,
+            self.data,
+            self.vopt,
+            self.pert,
+            self.cam,
+            mujoco.mjtCatBit.mjCAT_ALL,
+            self.scn,
+        )
+
+        if segmentation:
+            self.scn.flags[mujoco.mjtRndFlag.mjRND_SEGMENT] = 1
+            self.scn.flags[mujoco.mjtRndFlag.mjRND_IDCOLOR] = 1
+
+        for marker_params in self._markers:
+            self._add_marker_to_scene(marker_params)
+
+        mujoco.mjr_render(rect, self.scn, self.con)
+
+        for gridpos, (text1, text2) in self._overlays.items():
+            mujoco.mjr_overlay(
+                mujoco.mjtFontScale.mjFONTSCALE_150,
+                gridpos,
+                rect,
+                text1.encode(),
+                text2.encode(),
+                self.con,
+            )
+
+        if segmentation:
+            self.scn.flags[mujoco.mjtRndFlag.mjRND_SEGMENT] = 0
+            self.scn.flags[mujoco.mjtRndFlag.mjRND_IDCOLOR] = 0
+
+    def read_pixels(self, width, height, depth=True, segmentation=False):
+        rect = mujoco.MjrRect(left=0, bottom=0, width=width, 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)
+
+        mujoco.mjr_readPixels(rgb_arr, depth_arr, rect, self.con)
+        rgb_img = rgb_arr.reshape(rect.height, rect.width, 3)
+
+        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 _init_camera(self):
+        self.cam.type = mujoco.mjtCamera.mjCAMERA_FREE
+        self.cam.fixedcamid = -1
+        for i in range(3):
+            self.cam.lookat[i] = np.median(self.data.geom_xpos[:, i])
+
+        self.cam.azimuth = 180
+        self.cam.elevation = -30
+        self.cam.distance = self.model.stat.extent * 0.75
+
+    def add_overlay(self, gridpos: int, text1: str, text2: str):
+        if gridpos not in self._overlays:
+            self._overlays[gridpos] = ["", ""]
+        self._overlays[gridpos][0] += text1 + "\n"
+        self._overlays[gridpos][1] += text2 + "\n"
+
+    def add_marker(self, **marker_params):
+        self._markers.append(marker_params)
+
+    def _add_marker_to_scene(self, marker):
+        if self.scn.ngeom >= self.scn.maxgeom:
+            raise RuntimeError("Ran out of geoms. maxgeom: %d" % self.scn.maxgeom)
+
+        g = self.scn.geoms[self.scn.ngeom]
+        g.dataid = -1
+        g.objtype = mujoco.mjtObj.mjOBJ_UNKNOWN
+        g.objid = -1
+        g.category = mujoco.mjtCatBit.mjCAT_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 = mujoco.mjtGeom.mjGEOM_BOX
+        g.size[:] = np.ones(3) * 0.1
+        g.mat[:] = np.eye(3)
+        g.rgba[:] = np.ones(4)
+
+        for key, value in marker.items():
+            if isinstance(value, (int, float, mujoco._enums.mjtGeom)):
+                setattr(g, key, value)
+            elif isinstance(value, (tuple, list, np.ndarray)):
+                attr = getattr(g, key)
+                attr[:] = np.asarray(value).reshape(attr.shape)
+            elif isinstance(value, str):
+                assert key == "label", "Only label is a string in mjtGeom."
+                if value is None:
+                    g.label[0] = 0
+                else:
+                    g.label = value
+            elif hasattr(g, key):
+                raise ValueError(
+                    "mjtGeom has attr {} but type {} is invalid".format(
+                        key, type(value)
+                    )
+                )
+            else:
+                raise ValueError("mjtGeom doesn't have field %s" % key)
+
+        self.scn.ngeom += 1
+
+    def close(self):
+        pass
+
+
+class RenderContextOffscreen(RenderContext):
+
+    def __init__(self, width, height, model, data):
+
+        self._get_opengl_backend(width, height)
+        self.opengl_context.make_current()
+
+        super().__init__(model, data, offscreen=True)
+
+    def _get_opengl_backend(self, width, height):
+
+        backend = os.environ.get("MUJOCO_GL")
+        if backend is not None:
+            try:
+                self.opengl_context = _ALL_RENDERERS[backend](width, height)
+            except KeyError:
+                raise RuntimeError(
+                    "Environment variable {} must be one of {!r}: got {!r}.".format(
+                        "MUJOCO_GL", _ALL_RENDERERS.keys(), backend
+                    )
+                )
+
+        else:
+            for name, _ in _ALL_RENDERERS.items():
+                try:
+                    self.opengl_context = _ALL_RENDERERS[name](width, height)
+                    backend = name
+                    break
+                except:
+                    pass
+            if backend is None:
+                raise RuntimeError(
+                    "No OpenGL backend could be imported. Attempting to create a "
+                    "rendering context will result in a RuntimeError."
+                )
+
+
+class Viewer(RenderContext):
+
+    def __init__(self, model, data):
+        self._gui_lock = Lock()
+        self._button_left_pressed = False
+        self._button_right_pressed = False
+        self._last_mouse_x = 0
+        self._last_mouse_y = 0
+        self._paused = False
+        self._transparent = False
+        self._contacts = False
+        self._render_every_frame = False
+        self._image_idx = 0
+        self._image_path = "/tmp/frame_%07d.png"
+        self._run_speed = 1.0
+        self.last_time_rendered = 0.0
+        self._advance_by_one_step = False
+        self._hide_menu = False
+        self.windows = []
+        self.cams = [-1]
+
+        glfw.init()
+        self.width, self.height = glfw.get_video_mode(glfw.get_primary_monitor()).size
+        self.windows.append(
+            glfw.create_window(self.width // 2, self.height // 2, "mujoco", None, None)
+        )
+        glfw.make_context_current(self.windows[0])
+        glfw.swap_interval(1)
+
+        framebuffer_width, framebuffer_height = glfw.get_framebuffer_size(
+            self.windows[0]
+        )
+        window_width, _ = glfw.get_window_size(self.windows[0])
+        self._scale = framebuffer_width * 1.0 / window_width
+
+        glfw.set_cursor_pos_callback(self.windows[0], self._cursor_pos_callback)
+        glfw.set_mouse_button_callback(self.windows[0], self._mouse_button_callback)
+        glfw.set_scroll_callback(self.windows[0], self._scroll_callback)
+        glfw.set_key_callback(self.windows[0], self._key_callback)
+
+        self.viewport = mujoco.MjrRect(0, 0, framebuffer_width, framebuffer_height)
+        super().__init__(model, data, offscreen=False)
+
+    def set_cameras(self, cameras: List[int]):
+        for cam in cameras:
+            self.windows.append(
+                glfw.create_window(
+                    self.width // 2,
+                    self.height // 2,
+                    "Camera: {}".format(str(cam)),
+                    None,
+                    self.windows[0],
+                )
+            )
+            glfw.make_context_current(self.windows[-1])
+            self.cams.append(cam)
+
+    def _key_callback(self, window, key, scancode, action, mods):
+        if action != glfw.RELEASE:
+            return
+        elif key == glfw.KEY_TAB:
+            self.cams[0] += 1
+            if self.cams[0] >= self.model.ncam:
+                self.cams[0] = -1
+        elif key == glfw.KEY_A:
+            self.windows.append(
+                glfw.create_window(
+                    self.width // 2,
+                    self.height // 2,
+                    "Camera: {}".format(str(self.cams[0])),
+                    None,
+                    self.windows[0],
+                )
+            )
+            glfw.make_context_current(self.windows[-1])
+            self.cams.append(self.cams[0])
+        elif key == glfw.KEY_P:
+            print("lookat: ", self.cam.lookat)
+            print("elevation: ", self.cam.elevation)
+            print("distance: ", self.cam.distance)
+            print("azimuth: ", self.cam.azimuth)
+        elif key == glfw.KEY_SPACE and self._paused is not None:
+            self._paused = not self._paused
+        elif key == glfw.KEY_RIGHT and self._paused is not None:
+            self._advance_by_one_step = True
+            self._paused = True
+        elif key == glfw.KEY_S:
+            self._run_speed -= 0.25
+            if self._run_speed <= 0.0:
+                self._run_speed = 0.25
+        elif key == glfw.KEY_F:
+            self._run_speed += 0.25
+        elif key == glfw.KEY_D:
+            self._render_every_frame = not self._render_every_frame
+        elif key == glfw.KEY_T:
+            img = np.zeros(
+                (
+                    glfw.get_framebuffer_size(window)[1],
+                    glfw.get_framebuffer_size(window)[0],
+                    3,
+                ),
+                dtype=np.uint8,
+            )
+            mujoco.mjr_readPixels(img, None, self.viewport, self.con)
+            imageio.imwrite(self._image_path % self._image_idx, np.flipud(img))
+            self._image_idx += 1
+        elif key == glfw.KEY_C:
+            self._contacts = not self._contacts
+            self.vopt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTPOINT] = self._contacts
+            self.vopt.flags[mujoco.mjtVisFlag.mjVIS_CONTACTFORCE] = self._contacts
+        elif key == glfw.KEY_E:
+            self.vopt.frame = 1 - self.vopt.frame
+        elif key == glfw.KEY_H:
+            self._hide_menu = not self._hide_menu
+        elif key == glfw.KEY_R:
+            self._transparent = not self._transparent
+            if self._transparent:
+                self.model.geom_rgba[:, 3] /= 5.0
+            else:
+                self.model.geom_rgba[:, 3] *= 5.0
+        elif key in (glfw.KEY_0, glfw.KEY_1, glfw.KEY_2, glfw.KEY_3, glfw.KEY_4):
+            self.vopt.geomgroup[key - glfw.KEY_0] ^= 1
+        if key == glfw.KEY_ESCAPE:
+            print("Pressed ESC")
+            print("Quitting.")
+            glfw.destroy_window(window)
+            glfw.terminate()
+
+    def _cursor_pos_callback(self, window, xpos, ypos):
+        if not (self._button_left_pressed or self._button_right_pressed):
+            return
+
+        mod_shift = (
+            glfw.get_key(window, glfw.KEY_LEFT_SHIFT) == glfw.PRESS
+            or glfw.get_key(window, glfw.KEY_RIGHT_SHIFT) == glfw.PRESS
+        )
+        if self._button_right_pressed:
+            action = (
+                mujoco.mjtMouse.mjMOUSE_MOVE_H
+                if mod_shift
+                else mujoco.mjtMouse.mjMOUSE_MOVE_V
+            )
+        elif self._button_left_pressed:
+            action = (
+                mujoco.mjtMouse.mjMOUSE_ROTATE_H
+                if mod_shift
+                else mujoco.mjtMouse.mjMOUSE_ROTATE_V
+            )
+        else:
+            action = mujoco.mjtMouse.mjMOUSE_ZOOM
+
+        dx = int(self._scale * xpos) - self._last_mouse_x
+        dy = int(self._scale * ypos) - self._last_mouse_y
+        width, height = glfw.get_framebuffer_size(window)
+
+        with self._gui_lock:
+            mujoco.mjv_moveCamera(
+                self.model, action, dx / height, dy / height, self.scn, self.cam
+            )
+
+        self._last_mouse_x = int(self._scale * xpos)
+        self._last_mouse_y = int(self._scale * ypos)
+
+    def _mouse_button_callback(self, window, button, act, mods):
+        self._button_left_pressed = (
+            glfw.get_mouse_button(window, glfw.MOUSE_BUTTON_LEFT) == glfw.PRESS
+        )
+        self._button_right_pressed = (
+            glfw.get_mouse_button(window, glfw.MOUSE_BUTTON_RIGHT) == glfw.PRESS
+        )
+
+        x, y = glfw.get_cursor_pos(window)
+        self._last_mouse_x = int(self._scale * x)
+        self._last_mouse_y = int(self._scale * y)
+
+    def _scroll_callback(self, window, x_offset, y_offset):
+        with self._gui_lock:
+            mujoco.mjv_moveCamera(
+                self.model,
+                mujoco.mjtMouse.mjMOUSE_ZOOM,
+                0,
+                -0.05 * y_offset,
+                self.scn,
+                self.cam,
+            )
+
+    def _create_overlay(self):
+        topleft = mujoco.mjtGridPos.mjGRID_TOPLEFT
+        bottomleft = mujoco.mjtGridPos.mjGRID_BOTTOMLEFT
+
+        if self._render_every_frame:
+            self.add_overlay(topleft, "", "")
+        else:
+            self.add_overlay(
+                topleft,
+                "Run speed = %.3f x real time" % self._run_speed,
+                "[S]lower, [F]aster",
+            )
+        self.add_overlay(
+            topleft, "Ren[d]er every frame", "On" if self._render_every_frame else "Off"
+        )
+        self.add_overlay(
+            topleft,
+            "Switch camera (#cams = %d)" % (self.model.ncam + 1),
+            "[Tab] (camera ID = %d)" % self.cam.fixedcamid,
+        )
+        self.add_overlay(topleft, "[C]ontact forces", "On" if self._contacts else "Off")
+        self.add_overlay(topleft, "T[r]ansparent", "On" if self._transparent else "Off")
+        self.add_overlay(topleft, "Camera detach", "[A]")
+        self.add_overlay(topleft, "Print camera position", "[p]")
+        if self._paused is not None:
+            if not self._paused:
+                self.add_overlay(topleft, "Stop", "[Space]")
+            else:
+                self.add_overlay(topleft, "Start", "[Space]")
+                self.add_overlay(
+                    topleft, "Advance simulation by one step", "[right arrow]"
+                )
+        self.add_overlay(
+            topleft, "Referenc[e] frames", "On" if self.vopt.frame == 1 else "Off"
+        )
+        self.add_overlay(topleft, "[H]ide Menu", "")
+        if self._image_idx > 0:
+            fname = self._image_path % (self._image_idx - 1)
+            self.add_overlay(topleft, "Cap[t]ure frame", "Saved as %s" % fname)
+        else:
+            self.add_overlay(topleft, "Cap[t]ure frame", "")
+        self.add_overlay(topleft, "Toggle geomgroup visibility", "0-4")
+
+        self.add_overlay(
+            bottomleft, "Solver iterations", str(self.data.solver_iter + 1)
+        )
+        self.add_overlay(
+            bottomleft, "Step", str(round(self.data.time / self.model.opt.timestep))
+        )
+        self.add_overlay(bottomleft, "timestep", "%.5f" % self.model.opt.timestep)
+
+    def render(self):
+        def get_scaled_time():
+            return (1.0 / self._run_speed) * self.data.time
+
+        def time_for_next_frame(next_frame_after=1.0 / 30.0):
+            if (
+                get_scaled_time() < self.last_time_rendered
+                or self.last_time_rendered == 0.0
+            ):
+                return True
+            return (get_scaled_time() - self.last_time_rendered) >= next_frame_after
+
+        def update():
+            for i, wc in enumerate(zip(self.windows, self.cams)):
+                window, cam_id = wc
+                if i == 0:
+                    self._create_overlay()
+                self.cam.fixedcamid = cam_id
+                self.cam.type = mujoco.mjtCamera.mjCAMERA_FIXED
+                if self.cam.fixedcamid == -1:
+                    self.cam.type = mujoco.mjtCamera.mjCAMERA_FREE
+
+                if window is None:
+                    continue
+                elif glfw.window_should_close(window):
+                    if i == 0:
+                        glfw.terminate()
+                    glfw.destroy_window(window)
+                    self.windows.remove(window)
+                    self.cams.pop(i)
+                    continue
+                self.viewport.width, self.viewport.height = glfw.get_framebuffer_size(
+                    window
+                )
+                with self._gui_lock:
+                    mujoco.mjv_updateScene(
+                        self.model,
+                        self.data,
+                        self.vopt,
+                        mujoco.MjvPerturb(),
+                        self.cam,
+                        mujoco.mjtCatBit.mjCAT_ALL.value,
+                        self.scn,
+                    )
+                    for marker in self._markers:
+                        self._add_marker_to_scene(marker)
+                    glfw.make_context_current(window)
+                    mujoco.mjr_render(self.viewport, self.scn, self.con)
+                    if not self._hide_menu:
+                        for gridpos, [t1, t2] in self._overlays.items():
+                            mujoco.mjr_overlay(
+                                mujoco.mjtFontScale.mjFONTSCALE_150,
+                                gridpos,
+                                self.viewport,
+                                t1,
+                                t2,
+                                self.con,
+                            )
+                    glfw.swap_buffers(window)
+                glfw.poll_events()
+                if i == 0:
+                    self._overlays.clear()
+
+        if self._paused:
+            while self._paused:
+                update()
+                if self._advance_by_one_step:
+                    self._advance_by_one_step = False
+                    break
+        else:
+            if self._render_every_frame or time_for_next_frame():
+                update()
+                self.last_time_rendered = get_scaled_time()
+
+        self._markers[:] = []
+
+    def close(self):
+        for window in self.windows:
+            glfw.destroy_window(window)
+        glfw.terminate()
+
+    def reinit(self, model, data):
+        self.model = model
+        self.data = data
+        max_geom = 1000
+        mujoco.mj_forward(self.model, self.data)
+        self.scn = mujoco.MjvScene(self.model, max_geom)
+        glfw.make_context_current(self.windows[-1])
+        self.con = mujoco.MjrContext(self.model, mujoco.mjtFontScale.mjFONTSCALE_150)

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/sims/mj_beta/mj_utils/mj_scene_object.py

@@ -0,0 +1,262 @@
+import os
+import xml.etree.ElementTree as Et
+from typing import Tuple
+
+import numpy as np
+import yaml
+
+import d3il.environments.d3il.d3il_sim.sims.mj_beta.MjLoadable as mj_load
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.mj_utils.mj_helper import IncludeType
+from d3il.environments.d3il.d3il_sim.utils import sim_path
+
+
+class MujocoObject(mj_load.MjXmlLoadable):
+    def __init__(self, object_name, pos, quat, root=sim_path.D3IL_DIR):
+        if pos is None:
+            pos = [0, 0, 0]
+        else:
+            assert len(pos) == 3, "Error, parameter pos has to be three dimensional."
+
+        if quat is None:
+            quat = [0, 0, 0, 0]
+        else:
+            assert len(quat) == 4, "Error, parameter quat has to be four dimensional."
+
+        self.obj_dir_path = "models/mj/common-objects/" + object_name
+        self.obj_file_name = object_name + ".xml"
+        self.root = root
+        self.pos = pos
+        self.quat = quat
+        self.name = object_name
+        super().__init__(os.path.join(root, self.obj_dir_path, self.obj_file_name))
+
+    def mj_load(self) -> Tuple[Et.Element, list, IncludeType]:
+        include_et, xml, is_include = super(MujocoObject, self).mj_load()
+        xml_file = xml[0]
+        obj = Et.ElementTree(Et.fromstring(xml_file))
+        worldbody = obj.find("worldbody")
+        body = worldbody.find("body")
+
+        obj_pos_str = " ".join(map(str, self.pos))
+        obj_quat_str = " ".join(map(str, self.quat))
+
+        body.set("pos", obj_pos_str)
+        body.set("quat", obj_quat_str)
+
+        return include_et, Et.tostring(obj), IncludeType.FILE_INCLUDE
+
+    def get_poi(self) -> list:
+        return [self.name]
+
+
+class CompositeMujocoObject(mj_load.MjLoadable, mj_load.MjFreezable):
+    def __init__(self, object_folder, object_name, pos, quat):
+        if pos is None:
+            pos = [0, 0, 0]
+        else:
+            assert len(pos) == 3, "Error, parameter pos has to be three dimensional."
+
+        if quat is None:
+            quat = [0, 0, 0, 0]
+        else:
+            assert len(quat) == 4, "Error, parameter quat has to be four dimensional."
+
+        self.object_folder = object_folder
+        self.pos = pos
+        self.quat = quat
+        self.name = object_name
+        self.file_name = f"{self.name}.xml"
+
+    def mj_load(self) -> Tuple[Et.Element, list, bool]:
+        et_include = Et.Element("include")
+        et_include.set("file", self.file_name)
+
+        xml_string = self.generate_xml()
+        return (
+            et_include,
+            {os.path.join(self.object_folder, self.file_name): xml_string},
+            IncludeType.FILE_INCLUDE,
+        )
+
+    def generate_xml(self) -> str:
+        info_file = os.path.join(self.object_folder, "info.yml")
+        assert os.path.isfile(
+            info_file
+        ), f"The file {info_file} was not found. Did you specify the path to the object folder correctly?"
+
+        with open(info_file, "r") as f:
+            info_dict = yaml.safe_load(f)
+
+        original_file = info_dict["original_file"]
+        submesh_files = info_dict["submesh_files"]
+        submesh_props = info_dict["submesh_props"]
+        weight = info_dict["weight"]
+        material_map = info_dict["material_map"]
+
+        root = Et.Element("mujoco", attrib={"model": self.name})
+
+        assets = Et.SubElement(root, "asset")
+        worldbody = Et.SubElement(root, "worldbody")
+        body_attributes = {
+            "name": f"{self.name}",
+            "pos": " ".join(map(str, self.pos)),
+            "quat": " ".join(map(str, self.quat)),
+        }
+        body = Et.SubElement(worldbody, "body", attrib=body_attributes)
+
+        texture_attributes = {
+            "type": "2d",
+            "name": f"{self.name}_tex",
+            "file": os.path.join(self.object_folder, material_map),
+        }
+        texture = Et.SubElement(assets, "texture", attrib=texture_attributes)
+
+        material_attributes = {
+            "name": f"{self.name}_mat",
+            "texture": texture_attributes["name"],
+            "specular": "0.5",
+            "shininess": "0.5",
+        }
+        material = Et.SubElement(assets, "material", attrib=material_attributes)
+
+        orig_mesh_attributes = {
+            "name": f"{self.name}_orig",
+            "file": os.path.join(self.object_folder, original_file),
+        }
+        orig_mesh = Et.SubElement(assets, "mesh", attrib=orig_mesh_attributes)
+
+        orig_geom_attributes = {
+            "material": material_attributes["name"],
+            "mesh": orig_mesh_attributes["name"],
+            "group": "2",
+            "type": "mesh",
+            "contype": "0",
+            "conaffinity": "0",
+        }
+        orig_geom = Et.SubElement(body, "geom", attrib=orig_geom_attributes)
+
+        for i, (submesh_file, submesh_prop) in enumerate(
+            zip(submesh_files, submesh_props)
+        ):
+            collision_mesh_attributes = {
+                "name": f"{self.name}_coll_{i}",
+                "file": os.path.join(self.object_folder, submesh_file),
+            }
+            collision_mesh = Et.SubElement(
+                assets, "mesh", attrib=collision_mesh_attributes
+            )
+            collision_geom_attributes = {
+                "mesh": collision_mesh_attributes["name"],
+                "mass": str(weight * submesh_prop),
+                "group": "3",
+                "type": "mesh",
+                "conaffinity": "1",
+                "contype": "1",
+                "condim": "4",
+                "friction": "0.95 0.3 0.1",
+                "rgba": "1 1 1 1",
+                "solimp": "0.998 0.998 0.001",
+                "solref": "0.001 1",
+            }
+            collision_geom = Et.SubElement(
+                body, "geom", attrib=collision_geom_attributes
+            )
+
+        joint_attributes = {
+            "damping": "0.0001",
+            "name": f"{self.name}:joint",
+            "type": "free",
+        }
+        joint = Et.SubElement(body, "joint", attrib=joint_attributes)
+
+        return Et.tostring(root)
+
+    def get_poi(self) -> list:
+        return [self.name]
+
+    def freeze(self, model, data):
+        pos: np.ndarray = data.jnt(self.name + ":joint").qpos
+        self.pos = pos[0:3].tolist()
+        self.quat = pos[3:].tolist()
+
+    def unfreeze(self, data, model):
+        pass
+
+
+class YCBMujocoObject(CompositeMujocoObject):
+    def __init__(self, ycb_base_folder, object_id, object_name, pos, quat):
+        self.ycb_base_folder = ycb_base_folder
+        self.ycb_object_folder = os.path.join(ycb_base_folder, object_id)
+
+        super().__init__(self.ycb_object_folder, object_name, pos, quat)
+
+
+class CustomMujocoObject(mj_load.MjXmlLoadable):
+    def __init__(self, object_name, object_dir_path, pos, quat, root=sim_path.D3IL_DIR):
+        if pos is None:
+            pos = [0, 0, 0]
+        else:
+            assert len(pos) == 3, "Error, parameter pos has to be three dimensional."
+
+        if quat is None:
+            quat = [0, 0, 0, 0]
+        else:
+            assert len(quat) == 4, "Error, parameter quat has to be four dimensional."
+
+        self.obj_dir_path = object_dir_path
+        self.obj_file_name = object_name + ".xml"
+        self.root = root
+        self.pos = pos
+        self.quat = quat
+        self.name = object_name
+
+        super().__init__(os.path.join(root, self.obj_dir_path, self.obj_file_name))
+
+    def mj_load(self) -> Tuple[Et.Element, list, bool]:
+        include_et, xml, is_include = super(CustomMujocoObject, self).mj_load()
+        xml_file = list(xml.values())[0]
+        obj = Et.ElementTree(Et.fromstring(xml_file))
+        worldbody = obj.find("worldbody")
+        body = worldbody.find("body")
+
+        obj_pos_str = " ".join(map(str, self.pos))
+        obj_quat_str = " ".join(map(str, self.quat))
+
+        body.set("pos", obj_pos_str)
+        body.set("quat", obj_quat_str)
+
+        return (
+            include_et,
+            {
+                os.path.join(self.loadable_dir, self.file_name): Et.tostring(
+                    obj.getroot()
+                )
+            },
+            is_include,
+        )
+
+    def get_poi(self) -> list:
+        return [self.name]
+
+
+class MjWorkspace(mj_load.MjXmlLoadable):
+    def __init__(self, size: str, root=sim_path.D3IL_DIR):
+        if size not in ["small", "medium", "large"]:
+            raise ValueError(
+                "Error, please choose a size between <small>, <medium> or <large> ."
+            )
+        self.size = size
+        self.root = root
+        super().__init__(
+            os.path.join(root, "models/mj/workspace/{}".format(self.size), "ws.xml")
+        )
+
+
+class MjSurrounding(mj_load.MjXmlLoadable):
+    def __init__(self, surrounding_name, root=None):
+        self.surrounding_name = surrounding_name
+        if root is None:
+            root = os.path.join(sim_path.D3IL_DIR, "models/mj/surroundings")
+        self.root = root
+
+        super().__init__(os.path.join(root, "{}.xml".format(self.surrounding_name)))

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/sims/mj_beta/mj_utils/mj_scene_parser.py

@@ -0,0 +1,94 @@
+import xml.etree.ElementTree as Et
+from itertools import chain
+from xml.etree.ElementTree import Element
+
+import mujoco
+
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.mj_utils.mj_helper import IncludeType
+from d3il.environments.d3il.d3il_sim.sims.mj_beta.MjLoadable import (
+    MjLoadable,
+    MjIncludeTemplate,
+)
+from d3il.environments.d3il.d3il_sim.utils.sim_path import d3il_path
+
+
+class MjSceneParser:
+
+    def __init__(self, main_xml_path=None):
+        if main_xml_path is None:
+            main_xml_path = "./models/mj/surroundings/base.xml"
+
+        with open(d3il_path(main_xml_path), "r") as file:
+            self.scene_xml = file.read()
+
+        self.assets = {}
+        self.mj_xml_string = ""
+
+        self._tree = Et.ElementTree(Et.fromstring(self.scene_xml))
+        self._root = self._tree.getroot()
+        self._worldbody = self._root.find("worldbody")
+        assert self._worldbody, "Error, xml file does contain a world body."
+
+    def create_scene(self, mj_robots, mj_surrounding, object_list: list, dt):
+        if object_list is None:
+            object_list = []
+        if not isinstance(object_list, list):
+            object_list = [object_list]
+
+        all_objects = chain([mj_surrounding], object_list, mj_robots)
+
+        for obj in all_objects:
+            self.load_mj_loadable(obj)
+
+        self.set_dt(dt)
+
+        model = mujoco.MjModel.from_xml_string(self.mj_xml_string, self.assets)
+        data = mujoco.MjData(model)
+
+        self.cleanup(all_objects)
+        return model, data
+
+    def load_mj_loadable(self, mj_loadable: MjLoadable):
+        xml_element, assets, include_type = mj_loadable.mj_load()
+        if include_type == IncludeType.FILE_INCLUDE:
+            self._root.append(xml_element)
+        elif include_type == IncludeType.WORLD_BODY:
+            self._worldbody.append(xml_element)
+        elif include_type == IncludeType.MJ_INCLUDE:
+            for child in xml_element:
+                self._root.append(child)
+        elif include_type == IncludeType.VIRTUAL_INCLUDE:
+            return
+        else:
+            raise ValueError("Unknown IncludeType")
+
+        self.indent(self._root)
+
+        self.mj_xml_string = Et.tostring(self._root, encoding="unicode", method="xml")
+        self.assets.update(assets)
+
+    def set_dt(self, dt=0.001):
+        options = self._root.find("option")
+        options.set("timestep", str(dt))
+        self.indent(self._root)
+        self.mj_xml_string = Et.tostring(self._root, encoding="unicode", method="xml")
+
+    def indent(self, elem: Element, level=0):
+        i = "\n" + level * "  "
+        if len(elem):
+            if not elem.text or not elem.text.strip():
+                elem.text = i + "  "
+            if not elem.tail or not elem.tail.strip():
+                elem.tail = i
+            for elem in elem:
+                self.indent(elem, level + 1)
+            if not elem.tail or not elem.tail.strip():
+                elem.tail = i
+        else:
+            if level and (not elem.tail or not elem.tail.strip()):
+                elem.tail = i
+
+    def cleanup(self, objects: list):
+        for obj in objects:
+            if isinstance(obj, MjIncludeTemplate):
+                obj.cleanup()

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/sims/universal_sim/PrimitiveObjects.py

@@ -0,0 +1,144 @@
+from enum import Enum
+from typing import List, Tuple
+
+from d3il.environments.d3il.d3il_sim.core.sim_object.sim_object import SimObject
+
+
+class PrimitiveObject(SimObject):
+    class Shape(Enum):
+        BOX = "box"
+        SPHERE = "sphere"
+        CYLINDER = "cylinder"
+
+    def __init__(
+        self,
+        name: str,
+        init_pos,
+        init_quat,
+        mass: float,
+        size,
+        rgba,
+        static: bool,
+        visual_only: bool,
+        obj_type: Shape,
+        solimp: List[float] = None,
+        solref: List[float] = None,
+    ):
+        super(PrimitiveObject, self).__init__(name, init_pos, init_quat)
+        self.type = obj_type
+        self.mass = mass
+        self.size = size
+        self.static = static
+        self.visual_only = visual_only
+
+        if rgba is None:
+            rgba = [0.5, 0.5, 0.5, 1.0]
+        self.rgba = rgba
+
+        self.solimp = solimp
+        self.solref = solref
+
+    def get_poi(self) -> list:
+        return [self.name]
+
+
+class Box(PrimitiveObject):
+    def __init__(
+        self,
+        name: str,
+        init_pos,
+        init_quat,
+        mass: float = 0.1,
+        size=None,
+        rgba=None,
+        static: bool = False,
+        visual_only: bool = False,
+        solimp: List[float] = None,
+        solref: List[float] = None,
+    ):
+        if size is None:
+            size = [0.02, 0.02, 0.02]
+
+        if len(size) != 3:
+            raise ValueError("Expected list of size three for size attribute.")
+        super(Box, self).__init__(
+            name,
+            init_pos,
+            init_quat,
+            mass,
+            size,
+            rgba,
+            static,
+            visual_only,
+            self.Shape.BOX,
+            solimp,
+            solref,
+        )
+
+
+class Sphere(PrimitiveObject):
+    def __init__(
+        self,
+        name: str,
+        init_pos,
+        init_quat,
+        mass: float = 0.1,
+        size=None,
+        rgba=None,
+        static: bool = False,
+        visual_only: bool = False,
+        solimp: List[float] = None,
+        solref: List[float] = None,
+    ):
+        if size is None:
+            size = [0.02]
+
+        if len(size) != 1:
+            raise ValueError("Expected list of size one for size attribute")
+        super(Sphere, self).__init__(
+            name,
+            init_pos,
+            init_quat,
+            mass,
+            size,
+            rgba,
+            static,
+            visual_only,
+            self.Shape.SPHERE,
+            solimp,
+            solref,
+        )
+
+
+class Cylinder(PrimitiveObject):
+    def __init__(
+        self,
+        name: str,
+        init_pos,
+        init_quat,
+        mass: float = 0.1,
+        size=None,
+        rgba=None,
+        static: bool = False,
+        visual_only: bool = False,
+        solimp: List[float] = None,
+        solref: List[float] = None,
+    ):
+        if size is None:
+            size = [0.02, 0.02]
+
+        if len(size) != 2:
+            raise ValueError("Expected list of size two for size attribute.")
+        super(Cylinder, self).__init__(
+            name,
+            init_pos,
+            init_quat,
+            mass,
+            size,
+            rgba,
+            static,
+            visual_only,
+            self.Shape.CYLINDER,
+            solimp,
+            solref,
+        )

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/utils/__init__.py

@@ -0,0 +1 @@
+from .geometric_transformation import *

applications/avoiding/avoiding_baselines/DiRecT/d3il/environments/d3il/d3il_sim/utils/geometric_transformation.py

@@ -0,0 +1,319 @@
+import itertools
+
+import numpy as np
+
+_FLOAT_EPS = np.finfo(np.float64).eps
+_EPS4 = _FLOAT_EPS * 4.0
+
+
+def get_quaternion_error(curr_quat, des_quat):
+
+    quatError = np.zeros((3,))
+
+    quatError[0] = (
+        curr_quat[0] * des_quat[1]
+        - des_quat[0] * curr_quat[1]
+        - curr_quat[3] * des_quat[2]
+        + curr_quat[2] * des_quat[3]
+    )
+
+    quatError[1] = (
+        curr_quat[0] * des_quat[2]
+        - des_quat[0] * curr_quat[2]
+        + curr_quat[3] * des_quat[1]
+        - curr_quat[1] * des_quat[3]
+    )
+
+    quatError[2] = (
+        curr_quat[0] * des_quat[3]
+        - des_quat[0] * curr_quat[3]
+        - curr_quat[2] * des_quat[1]
+        + curr_quat[1] * des_quat[2]
+    )
+
+    return quatError
+
+
+def euler2mat(euler):
+    euler = np.asarray(euler, dtype=np.float64)
+    assert euler.shape[-1] == 3, "Invalid shaped euler {}".format(euler)
+
+    ai, aj, ak = -euler[..., 2], -euler[..., 1], -euler[..., 0]
+    si, sj, sk = np.sin(ai), np.sin(aj), np.sin(ak)
+    ci, cj, ck = np.cos(ai), np.cos(aj), np.cos(ak)
+    cc, cs = ci * ck, ci * sk
+    sc, ss = si * ck, si * sk
+
+    mat = np.empty(euler.shape[:-1] + (3, 3), dtype=np.float64)
+    mat[..., 2, 2] = cj * ck
+    mat[..., 2, 1] = sj * sc - cs
+    mat[..., 2, 0] = sj * cc + ss
+    mat[..., 1, 2] = cj * sk
+    mat[..., 1, 1] = sj * ss + cc
+    mat[..., 1, 0] = sj * cs - sc
+    mat[..., 0, 2] = -sj
+    mat[..., 0, 1] = cj * si
+    mat[..., 0, 0] = cj * ci
+    return mat
+
+
+def euler2quat(euler):
+    euler = np.asarray(euler, dtype=np.float64)
+    assert euler.shape[-1] == 3, "Invalid shape euler {}".format(euler)
+
+    ai, aj, ak = euler[..., 2] / 2, -euler[..., 1] / 2, euler[..., 0] / 2
+    si, sj, sk = np.sin(ai), np.sin(aj), np.sin(ak)
+    ci, cj, ck = np.cos(ai), np.cos(aj), np.cos(ak)
+    cc, cs = ci * ck, ci * sk
+    sc, ss = si * ck, si * sk
+
+    quat = np.empty(euler.shape[:-1] + (4,), dtype=np.float64)
+    quat[..., 0] = cj * cc + sj * ss
+    quat[..., 3] = cj * sc - sj * cs
+    quat[..., 2] = -(cj * ss + sj * cc)
+    quat[..., 1] = cj * cs - sj * sc
+    return quat
+
+
+def mat2euler(mat):
+    mat = np.asarray(mat, dtype=np.float64)
+    assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat)
+
+    cy = np.sqrt(mat[..., 2, 2] * mat[..., 2, 2] + mat[..., 1, 2] * mat[..., 1, 2])
+    condition = cy > _EPS4
+    euler = np.empty(mat.shape[:-1], dtype=np.float64)
+    euler[..., 2] = np.where(
+        condition,
+        -np.arctan2(mat[..., 0, 1], mat[..., 0, 0]),
+        -np.arctan2(-mat[..., 1, 0], mat[..., 1, 1]),
+    )
+    euler[..., 1] = np.where(
+        condition, -np.arctan2(-mat[..., 0, 2], cy), -np.arctan2(-mat[..., 0, 2], cy)
+    )
+    euler[..., 0] = np.where(
+        condition, -np.arctan2(mat[..., 1, 2], mat[..., 2, 2]), 0.0
+    )
+    return euler
+
+
+def mat2quat(mat):
+    mat = np.asarray(mat, dtype=np.float64)
+    assert mat.shape[-2:] == (3, 3), "Invalid shape matrix {}".format(mat)
+
+    Qxx, Qyx, Qzx = mat[..., 0, 0], mat[..., 0, 1], mat[..., 0, 2]
+    Qxy, Qyy, Qzy = mat[..., 1, 0], mat[..., 1, 1], mat[..., 1, 2]
+    Qxz, Qyz, Qzz = mat[..., 2, 0], mat[..., 2, 1], mat[..., 2, 2]
+    K = np.zeros(mat.shape[:-2] + (4, 4), dtype=np.float64)
+    K[..., 0, 0] = Qxx - Qyy - Qzz
+    K[..., 1, 0] = Qyx + Qxy
+    K[..., 1, 1] = Qyy - Qxx - Qzz
+    K[..., 2, 0] = Qzx + Qxz
+    K[..., 2, 1] = Qzy + Qyz
+    K[..., 2, 2] = Qzz - Qxx - Qyy
+    K[..., 3, 0] = Qyz - Qzy
+    K[..., 3, 1] = Qzx - Qxz
+    K[..., 3, 2] = Qxy - Qyx
+    K[..., 3, 3] = Qxx + Qyy + Qzz
+    K /= 3.0
+    q = np.empty(K.shape[:-2] + (4,))
+    it = np.nditer(q[..., 0], flags=["multi_index"])
+    while not it.finished:
+        vals, vecs = np.linalg.eigh(K[it.multi_index])
+        q[it.multi_index] = vecs[[3, 0, 1, 2], np.argmax(vals)]
+        if q[it.multi_index][0] < 0:
+            q[it.multi_index] *= -1
+        it.iternext()
+    return q
+
+
+def quat2euler(quat):
+    return mat2euler(quat2mat(quat))
+
+
+def subtract_euler(e1, e2):
+    assert e1.shape == e2.shape
+    assert e1.shape[-1] == 3
+    q1 = euler2quat(e1)
+    q2 = euler2quat(e2)
+    q_diff = quat_mul(q1, quat_conjugate(q2))
+    return quat2euler(q_diff)
+
+
+def quat2mat(quat):
+    quat = np.asarray(quat, dtype=np.float64)
+    assert quat.shape[-1] == 4, "Invalid shape quat {}".format(quat)
+
+    w, x, y, z = quat[..., 0], quat[..., 1], quat[..., 2], quat[..., 3]
+    Nq = np.sum(quat * quat, axis=-1)
+    s = 2.0 / Nq
+    X, Y, Z = x * s, y * s, z * s
+    wX, wY, wZ = w * X, w * Y, w * Z
+    xX, xY, xZ = x * X, x * Y, x * Z
+    yY, yZ, zZ = y * Y, y * Z, z * Z
+
+    mat = np.empty(quat.shape[:-1] + (3, 3), dtype=np.float64)
+    mat[..., 0, 0] = 1.0 - (yY + zZ)
+    mat[..., 0, 1] = xY - wZ
+    mat[..., 0, 2] = xZ + wY
+    mat[..., 1, 0] = xY + wZ
+    mat[..., 1, 1] = 1.0 - (xX + zZ)
+    mat[..., 1, 2] = yZ - wX
+    mat[..., 2, 0] = xZ - wY
+    mat[..., 2, 1] = yZ + wX
+    mat[..., 2, 2] = 1.0 - (xX + yY)
+    return np.where((Nq > _FLOAT_EPS)[..., np.newaxis, np.newaxis], mat, np.eye(3))
+
+
+def quat_conjugate(q):
+    inv_q = -q
+    inv_q[..., 0] *= -1
+    return inv_q
+
+
+def quat_mul(q0, q1):
+    assert q0.shape == q1.shape
+    assert q0.shape[-1] == 4
+    assert q1.shape[-1] == 4
+
+    w0 = q0[..., 0]
+    x0 = q0[..., 1]
+    y0 = q0[..., 2]
+    z0 = q0[..., 3]
+
+    w1 = q1[..., 0]
+    x1 = q1[..., 1]
+    y1 = q1[..., 2]
+    z1 = q1[..., 3]
+
+    w = w0 * w1 - x0 * x1 - y0 * y1 - z0 * z1
+    x = w0 * x1 + x0 * w1 + y0 * z1 - z0 * y1
+    y = w0 * y1 + y0 * w1 + z0 * x1 - x0 * z1
+    z = w0 * z1 + z0 * w1 + x0 * y1 - y0 * x1
+    q = np.array([w, x, y, z])
+    if q.ndim == 2:
+        q = q.swapaxes(0, 1)
+    assert q.shape == q0.shape
+    return q
+
+
+def quat_rot_vec(q, v0):
+    q_v0 = np.array([0, v0[0], v0[1], v0[2]])
+    q_v = quat_mul(q, quat_mul(q_v0, quat_conjugate(q)))
+    v = q_v[1:]
+    return v
+
+
+def quat_identity():
+    return np.array([1, 0, 0, 0])
+
+
+def quat2axisangle(quat):
+    theta = 0
+    axis = np.array([0, 0, 1])
+    sin_theta = np.linalg.norm(quat[1:])
+
+    if sin_theta > 0.0001:
+        theta = 2 * np.arcsin(sin_theta)
+        theta *= 1 if quat[0] >= 0 else -1
+        axis = quat[1:] / sin_theta
+
+    return axis, theta
+
+
+def euler2point_euler(euler):
+    _euler = euler.copy()
+    if len(_euler.shape) < 2:
+        _euler = np.expand_dims(_euler, 0)
+    assert _euler.shape[1] == 3
+    _euler_sin = np.sin(_euler)
+    _euler_cos = np.cos(_euler)
+    return np.concatenate([_euler_sin, _euler_cos], axis=-1)
+
+
+def point_euler2euler(euler):
+    _euler = euler.copy()
+    if len(_euler.shape) < 2:
+        _euler = np.expand_dims(_euler, 0)
+    assert _euler.shape[1] == 6
+    angle = np.arctan(_euler[..., :3] / _euler[..., 3:])
+    angle[_euler[..., 3:] < 0] += np.pi
+    return angle
+
+
+def quat2point_quat(quat):
+    _quat = quat.copy()
+    if len(_quat.shape) < 2:
+        _quat = np.expand_dims(_quat, 0)
+    assert _quat.shape[1] == 4
+    angle = np.arccos(_quat[:, [0]]) * 2
+    xyz = _quat[:, 1:]
+    xyz[np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5] = (xyz / np.sin(angle / 2))[
+        np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5
+    ]
+    return np.concatenate([np.sin(angle), np.cos(angle), xyz], axis=-1)
+
+
+def point_quat2quat(quat):
+    _quat = quat.copy()
+    if len(_quat.shape) < 2:
+        _quat = np.expand_dims(_quat, 0)
+    assert _quat.shape[1] == 5
+    angle = np.arctan(_quat[:, [0]] / _quat[:, [1]])
+    qw = np.cos(angle / 2)
+
+    qxyz = _quat[:, 2:]
+    qxyz[np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5] = (qxyz * np.sin(angle / 2))[
+        np.squeeze(np.abs(np.sin(angle / 2))) >= 1e-5
+    ]
+    return np.concatenate([qw, qxyz], axis=-1)
+
+
+def normalize_angles(angles):
+    angles = angles.copy()
+    if angles.size > 0:
+        angles = (angles + np.pi) % (2 * np.pi) - np.pi
+        assert -np.pi - 1e-6 <= angles.min() and angles.max() <= np.pi + 1e-6
+    return angles
+
+
+def round_to_straight_angles(angles):
+    angles = np.round(angles / (np.pi / 2)) * (np.pi / 2)
+    return normalize_angles(angles)
+
+
+def get_parallel_rotations():
+    mult90 = [0, np.pi / 2, -np.pi / 2, np.pi]
+    parallel_rotations = []
+    for euler in itertools.product(mult90, repeat=3):
+        canonical = mat2euler(euler2mat(euler))
+        canonical = np.round(canonical / (np.pi / 2))
+        if canonical[0] == -2:
+            canonical[0] = 2
+        if canonical[2] == -2:
+            canonical[2] = 2
+        canonical *= np.pi / 2
+        if all([(canonical != rot).any() for rot in parallel_rotations]):
+            parallel_rotations += [canonical]
+    assert len(parallel_rotations) == 24
+    return parallel_rotations
+
+
+def posRotMat2TFMat(pos, rot_mat):
+    t_mat = np.eye(4)
+    t_mat[:3, :3] = rot_mat
+    t_mat[:3, 3] = np.array(pos)
+    return t_mat
+
+
+def mat2posQuat(mat):
+    pos = mat[:3, 3]
+    quat = mat2quat(mat[:3, :3])
+    return pos, quat
+
+
+def wxyz_to_xyzw(quat):
+    return np.roll(quat, -1)
+
+
+def xyzw_to_wxyz(quat):
+    return np.roll(quat, 1)