first commit

Open_Duck_Playground/.gitignore

@@ -0,0 +1,11 @@
+Open_Duck_Playground.egg-info/*
+uv.lock
+__pycache__/
+playground/sigmaban2019/
+playground/open_duck_mini_v2_skate/
+robot_saved_obs.pkl
+mujoco_saved_obs.pkl
+checkpoints/
+.tmp/
+*.onnx
+*.TXT

Open_Duck_Playground/README.md

@@ -0,0 +1,98 @@
+# Open Duck Playground
+
+# Installation 
+
+Install uv
+
+```bash
+curl -LsSf https://astral.sh/uv/install.sh | sh
+```
+
+# Training
+
+If you want to use the [imitation reward](https://la.disneyresearch.com/wp-content/uploads/BD_X_paper.pdf), you can generate reference motion with [this repo](https://github.com/apirrone/Open_Duck_reference_motion_generator)
+
+Then copy `polynomial_coefficients.pkl` in `playground/<robot>/data/`
+
+You'll also have to set `USE_IMITATION_REWARD=True` in it's `joystick.py` file
+
+Run: 
+
+```bash
+uv run playground/<robot>/runner.py 
+```
+
+## Tensorboard
+
+```bash
+uv run tensorboard --logdir=<yourlogdir>
+```
+
+# Inference 
+
+Infer mujoco
+
+(for now this is specific to open_duck_mini_v2)
+
+```bash
+uv run playground/open_duck_mini_v2/mujoco_infer.py -o <path_to_.onnx>
+```
+
+# Documentation
+
+## Project structure : 
+
+```
+.
+├── pyproject.toml
+├── README.md
+├── playground
+│   ├── common
+│   │   ├── export_onnx.py
+│   │   ├── onnx_infer.py
+│   │   ├── poly_reference_motion.py
+│   │   ├── randomize.py
+│   │   ├── rewards.py
+│   │   └── runner.py
+│   ├── open_duck_mini_v2
+│   │   ├── base.py
+│   │   ├── data
+│   │   │   └── polynomial_coefficients.pkl
+│   │   ├── joystick.py
+│   │   ├── mujoco_infer.py
+│   │   ├── constants.py
+│   │   ├── runner.py
+│   │   └── xmls
+│   │       ├── assets
+│   │       ├── open_duck_mini_v2_no_head.xml
+│   │       ├── open_duck_mini_v2.xml
+│   │       ├── scene_mjx_flat_terrain.xml
+│   │       ├── scene_mjx_rough_terrain.xml
+│   │       └── scene.xml
+```
+
+## Adding a new robot
+
+Create a new directory in `playground` named after `<your robot>`. You can copy the `open_duck_mini_v2` directory as a starting point.
+
+You will need to:
+- Edit `base.py`: Mainly renaming stuff to match you robot's name
+- Edit `constants.py`: specify the names of some important geoms, sensors etc
+  - In your `mjcf`, you'll probably have to add some sites, name some bodies/geoms and add the sensors. Look at how we did it for `open_duck_mini_v2`
+- Add your `mjcf` assets in `xmls`. 
+- Edit `joystick.py` : to choose the rewards you are interested in
+  - Note: for now there is still some hard coded values etc. We'll improve things on the way
+- Edit `runner.py`
+
+
+
+# Notes
+
+Inspired from https://github.com/kscalelabs/mujoco_playground
+
+
+## Current win
+
+```bash
+uv run playground/open_duck_mini_v2/runner.py --task flat_terrain_backlash --num_timesteps 300000000
+```

Open_Duck_Playground/playground/common/export_onnx.py

@@ -0,0 +1,189 @@
+
+import tensorflow as tf
+from tensorflow.keras import layers
+import tf2onnx
+import numpy as np
+
+def export_onnx(
+    params, act_size, ppo_params, obs_size, output_path="ONNX.onnx"
+):
+    print(" === EXPORT ONNX === ")
+
+    # inference_fn = make_inference_fn(params, deterministic=True)
+
+    class MLP(tf.keras.Model):
+        def __init__(
+            self,
+            layer_sizes,
+            activation=tf.nn.relu,
+            kernel_init="lecun_uniform",
+            activate_final=False,
+            bias=True,
+            layer_norm=False,
+            mean_std=None,
+        ):
+            super().__init__()
+
+            self.layer_sizes = layer_sizes
+            self.activation = activation
+            self.kernel_init = kernel_init
+            self.activate_final = activate_final
+            self.bias = bias
+            self.layer_norm = layer_norm
+
+            if mean_std is not None:
+                self.mean = tf.Variable(mean_std[0], trainable=False, dtype=tf.float32)
+                self.std = tf.Variable(mean_std[1], trainable=False, dtype=tf.float32)
+            else:
+                self.mean = None
+                self.std = None
+
+            self.mlp_block = tf.keras.Sequential(name="MLP_0")
+            for i, size in enumerate(self.layer_sizes):
+                dense_layer = layers.Dense(
+                    size,
+                    activation=self.activation,
+                    kernel_initializer=self.kernel_init,
+                    name=f"hidden_{i}",
+                    use_bias=self.bias,
+                )
+                self.mlp_block.add(dense_layer)
+                if self.layer_norm:
+                    self.mlp_block.add(
+                        layers.LayerNormalization(name=f"layer_norm_{i}")
+                    )
+            if not self.activate_final and self.mlp_block.layers:
+                if (
+                    hasattr(self.mlp_block.layers[-1], "activation")
+                    and self.mlp_block.layers[-1].activation is not None
+                ):
+                    self.mlp_block.layers[-1].activation = None
+
+            self.submodules = [self.mlp_block]
+
+        def call(self, inputs):
+            if isinstance(inputs, list):
+                inputs = inputs[0]
+            if self.mean is not None and self.std is not None:
+                print(self.mean.shape, self.std.shape)
+                inputs = (inputs - self.mean) / self.std
+            logits = self.mlp_block(inputs)
+            loc, _ = tf.split(logits, 2, axis=-1)
+            return tf.tanh(loc)
+
+    def make_policy_network(
+        param_size,
+        mean_std,
+        hidden_layer_sizes=[256, 256],
+        activation=tf.nn.relu,
+        kernel_init="lecun_uniform",
+        layer_norm=False,
+    ):
+        policy_network = MLP(
+            layer_sizes=list(hidden_layer_sizes) + [param_size],
+            activation=activation,
+            kernel_init=kernel_init,
+            layer_norm=layer_norm,
+            mean_std=mean_std,
+        )
+        return policy_network
+
+    mean = params[0].mean["state"]
+    std = params[0].std["state"]
+
+    # Convert mean/std jax arrays to tf tensors.
+    mean_std = (tf.convert_to_tensor(mean), tf.convert_to_tensor(std))
+
+    tf_policy_network = make_policy_network(
+        param_size=act_size * 2,
+        mean_std=mean_std,
+        hidden_layer_sizes=ppo_params.network_factory.policy_hidden_layer_sizes,
+        activation=tf.nn.swish,
+    )
+
+    example_input = tf.zeros((1, obs_size))
+    example_output = tf_policy_network(example_input)
+    print(example_output.shape)
+
+    def transfer_weights(jax_params, tf_model):
+        """
+        Transfer weights from a JAX parameter dictionary to the TensorFlow model.
+
+        Parameters:
+        - jax_params: dict
+        Nested dictionary with structure {block_name: {layer_name: {params}}}.
+        For example:
+        {
+            'CNN_0': {
+            'Conv_0': {'kernel': np.ndarray},
+            'Conv_1': {'kernel': np.ndarray},
+            'Conv_2': {'kernel': np.ndarray},
+            },
+            'MLP_0': {
+            'hidden_0': {'kernel': np.ndarray, 'bias': np.ndarray},
+            'hidden_1': {'kernel': np.ndarray, 'bias': np.ndarray},
+            'hidden_2': {'kernel': np.ndarray, 'bias': np.ndarray},
+            }
+        }
+
+        - tf_model: tf.keras.Model
+        An instance of the adapted VisionMLP model containing named submodules and layers.
+        """
+        for layer_name, layer_params in jax_params.items():
+            try:
+                tf_layer = tf_model.get_layer("MLP_0").get_layer(name=layer_name)
+            except ValueError:
+                print(f"Layer {layer_name} not found in TensorFlow model.")
+                continue
+            if isinstance(tf_layer, tf.keras.layers.Dense):
+                kernel = np.array(layer_params["kernel"])
+                bias = np.array(layer_params["bias"])
+                print(
+                    f"Transferring Dense layer {layer_name}, kernel shape {kernel.shape}, bias shape {bias.shape}"
+                )
+                tf_layer.set_weights([kernel, bias])
+            else:
+                print(f"Unhandled layer type in {layer_name}: {type(tf_layer)}")
+
+        print("Weights transferred successfully.")
+
+    net_params = params[1]
+    # try attribute‐style .policy
+    if hasattr(net_params, "policy"):
+        policy_tree = net_params.policy
+    # try dict‐style ["policy"]
+    elif isinstance(net_params, dict) and "policy" in net_params:
+        policy_tree = net_params["policy"]
+    # fallback to Flax FrozenDict { "params": … }
+    elif isinstance(net_params, dict) and "params" in net_params:
+        policy_tree = net_params["params"]
+    else:
+        raise KeyError(f"Cannot locate policy params in {type(net_params)}; keys = {list(net_params.keys())}")
+
+    # policy_tree is now the dict of weight arrays
+    transfer_weights(policy_tree, tf_policy_network)
+
+    # Example inputs for the model
+    test_input = [np.ones((1, obs_size), dtype=np.float32)]
+
+    # Define the TensorFlow input signature
+    spec = [
+        tf.TensorSpec(shape=(1, obs_size), dtype=tf.float32, name="obs")
+    ]
+
+    tensorflow_pred = tf_policy_network(test_input)[0]
+    # Build the model by calling it with example data
+    print(f"Tensorflow prediction: {tensorflow_pred}")
+
+    tf_policy_network.output_names = ["continuous_actions"]
+
+    # opset 11 matches isaac lab.
+    model_proto, _ = tf2onnx.convert.from_keras(
+        tf_policy_network, input_signature=spec, opset=11, output_path=output_path
+    )
+
+    # For Antoine :)
+    model_proto, _ = tf2onnx.convert.from_keras(
+        tf_policy_network, input_signature=spec, opset=11, output_path="ONNX.onnx"
+    )
+    return

Open_Duck_Playground/playground/common/onnx_infer.py

@@ -0,0 +1,46 @@
+import onnxruntime
+
+
+class OnnxInfer:
+    def __init__(self, onnx_model_path, input_name="obs", awd=False):
+        self.onnx_model_path = onnx_model_path
+        self.ort_session = onnxruntime.InferenceSession(
+            self.onnx_model_path, providers=["CPUExecutionProvider"]
+        )
+        self.input_name = input_name
+        self.awd = awd
+
+    def infer(self, inputs):
+        if self.awd:
+            outputs = self.ort_session.run(None, {self.input_name: [inputs]})
+            return outputs[0][0]
+        else:
+            outputs = self.ort_session.run(
+                None, {self.input_name: inputs.astype("float32")}
+            )
+            return outputs[0]
+
+
+if __name__ == "__main__":
+    import argparse
+    import numpy as np
+    import time
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-o", "--onnx_model_path", type=str, required=True)
+    args = parser.parse_args()
+
+    obs_size = 46
+
+    oi = OnnxInfer(args.onnx_model_path, awd=True)
+    times = []
+    for i in range(1000):
+        inputs = np.random.uniform(size=obs_size).astype(np.float32)
+        # inputs = np.arange(obs_size).astype(np.float32)
+        # print(inputs)
+        start = time.time()
+        print(oi.infer(inputs))
+        times.append(time.time() - start)
+
+    print("Average time: ", sum(times) / len(times))
+    print("Average fps: ", 1 / (sum(times) / len(times)))

Open_Duck_Playground/playground/common/plot_saved_obs.py

@@ -0,0 +1,222 @@
+import pickle
+
+import matplotlib.pyplot as plt
+import numpy as np
+from scipy.spatial.transform import Rotation as R
+import argparse
+
+parser = argparse.ArgumentParser()
+parser.add_argument(
+    "-d", "--data", type=str, required=False, default="mujoco_saved_obs.pkl"
+)
+args = parser.parse_args()
+
+
+init_pos = np.array(
+    [
+        0.002,
+        0.053,
+        -0.63,
+        1.368,
+        -0.784,
+        0.0,
+        0,
+        0,
+        0,
+        # 0,
+        # 0,
+        -0.003,
+        -0.065,
+        0.635,
+        1.379,
+        -0.796,
+    ]
+)
+
+joints_order = [
+    "left_hip_yaw",
+    "left_hip_roll",
+    "left_hip_pitch",
+    "left_knee",
+    "left_ankle",
+    "neck_pitch",
+    "head_pitch",
+    "head_yaw",
+    "head_roll",
+    "right_hip_yaw",
+    "right_hip_roll",
+    "right_hip_pitch",
+    "right_knee",
+    "right_ankle",
+]
+
+obses = pickle.load(open(args.data, "rb"))
+
+num_dofs = 14
+dof_poses = []  # (dof, num_obs)
+actions = []  # (dof, num_obs)
+
+for i in range(num_dofs):
+    print(i)
+    dof_poses.append([])
+    actions.append([])
+    for obs in obses:
+        dof_poses[i].append(obs[13 : 13 + num_dofs][i])
+        actions[i].append(obs[26 : 26 + num_dofs][i])
+
+# plot action vs dof pos
+
+nb_dofs = len(dof_poses)
+nb_rows = int(np.sqrt(nb_dofs))
+nb_cols = int(np.ceil(nb_dofs / nb_rows))
+
+fig, axs = plt.subplots(nb_rows, nb_cols, sharex=True, sharey=True)
+
+for i in range(nb_rows):
+    for j in range(nb_cols):
+        if i * nb_cols + j >= nb_dofs:
+            break
+        axs[i, j].plot(actions[i * nb_cols + j], label="action")
+        axs[i, j].plot(dof_poses[i * nb_cols + j], label="dof_pos")
+        axs[i, j].legend()
+        axs[i, j].set_title(f"{joints_order[i * nb_cols + j]}")
+
+fig.suptitle(f"{args.data}")
+plt.show()
+
+obses_names = [
+    "gyro x",
+    "gyro y",
+    "gyro z",
+    "accelo x",
+    "accelo y",
+    "accelo z",
+    # commands
+    "command 0",
+    "command 1",
+    "command 2",
+    "command 3",
+    "command 4",
+    "command 5",
+    "command 6",
+    # dof pos
+    "pos_" + str(joints_order[0]),
+    "pos_" + str(joints_order[1]),
+    "pos_" + str(joints_order[2]),
+    "pos_" + str(joints_order[3]),
+    "pos_" + str(joints_order[4]),
+    "pos_" + str(joints_order[5]),
+    "pos_" + str(joints_order[6]),
+    "pos_" + str(joints_order[7]),
+    "pos_" + str(joints_order[8]),
+    "pos_" + str(joints_order[9]),
+    "pos_" + str(joints_order[10]),
+    "pos_" + str(joints_order[11]),
+    "pos_" + str(joints_order[12]),
+    "pos_" + str(joints_order[13]),
+    # dof vel
+    "vel_" + str(joints_order[0]),
+    "vel_" + str(joints_order[1]),
+    "vel_" + str(joints_order[2]),
+    "vel_" + str(joints_order[3]),
+    "vel_" + str(joints_order[4]),
+    "vel_" + str(joints_order[5]),
+    "vel_" + str(joints_order[6]),
+    "vel_" + str(joints_order[7]),
+    "vel_" + str(joints_order[8]),
+    "vel_" + str(joints_order[9]),
+    "vel_" + str(joints_order[10]),
+    "vel_" + str(joints_order[11]),
+    "vel_" + str(joints_order[12]),
+    "vel_" + str(joints_order[13]),
+    # action
+    "last_action_" + str(joints_order[0]),
+    "last_action_" + str(joints_order[1]),
+    "last_action_" + str(joints_order[2]),
+    "last_action_" + str(joints_order[3]),
+    "last_action_" + str(joints_order[4]),
+    "last_action_" + str(joints_order[5]),
+    "last_action_" + str(joints_order[6]),
+    "last_action_" + str(joints_order[7]),
+    "last_action_" + str(joints_order[8]),
+    "last_action_" + str(joints_order[9]),
+    "last_action_" + str(joints_order[10]),
+    "last_action_" + str(joints_order[11]),
+    "last_action_" + str(joints_order[12]),
+    "last_action_" + str(joints_order[13]),
+    "last_last_action_" + str(joints_order[0]),
+    "last_last_action_" + str(joints_order[1]),
+    "last_last_action_" + str(joints_order[2]),
+    "last_last_action_" + str(joints_order[3]),
+    "last_last_action_" + str(joints_order[4]),
+    "last_last_action_" + str(joints_order[5]),
+    "last_last_action_" + str(joints_order[6]),
+    "last_last_action_" + str(joints_order[7]),
+    "last_last_action_" + str(joints_order[8]),
+    "last_last_action_" + str(joints_order[9]),
+    "last_last_action_" + str(joints_order[10]),
+    "last_last_action_" + str(joints_order[11]),
+    "last_last_action_" + str(joints_order[12]),
+    "last_last_action_" + str(joints_order[13]),
+    "last_last_last_action_" + str(joints_order[0]),
+    "last_last_last_action_" + str(joints_order[1]),
+    "last_last_last_action_" + str(joints_order[2]),
+    "last_last_last_action_" + str(joints_order[3]),
+    "last_last_last_action_" + str(joints_order[4]),
+    "last_last_last_action_" + str(joints_order[5]),
+    "last_last_last_action_" + str(joints_order[6]),
+    "last_last_last_action_" + str(joints_order[7]),
+    "last_last_last_action_" + str(joints_order[8]),
+    "last_last_last_action_" + str(joints_order[9]),
+    "last_last_last_action_" + str(joints_order[10]),
+    "last_last_last_action_" + str(joints_order[11]),
+    "last_last_last_action_" + str(joints_order[12]),
+    "last_last_last_action_" + str(joints_order[13]),
+    "motor_targets_" + str(joints_order[0]),
+    "motor_targets_" + str(joints_order[1]),
+    "motor_targets_" + str(joints_order[2]),
+    "motor_targets_" + str(joints_order[3]),
+    "motor_targets_" + str(joints_order[4]),
+    "motor_targets_" + str(joints_order[5]),
+    "motor_targets_" + str(joints_order[6]),
+    "motor_targets_" + str(joints_order[7]),
+    "motor_targets_" + str(joints_order[8]),
+    "motor_targets_" + str(joints_order[9]),
+    "motor_targets_" + str(joints_order[10]),
+    "motor_targets_" + str(joints_order[11]),
+    "motor_targets_" + str(joints_order[12]),
+    "motor_targets_" + str(joints_order[13]),
+    "contact left",
+    "contact right",
+    "imitation_phase 1",
+    "imitation_phase 2"
+    # ref (ignored)
+]
+# print(len(obses_names))
+# exit()
+
+
+# obses = [[56 obs at time 0], [56 obs at time 1], ...]
+
+nb_obs = len(obses[0])
+print(nb_obs)
+nb_rows = int(np.sqrt(nb_obs))
+nb_cols = int(np.ceil(nb_obs / nb_rows))
+
+fig, axs = plt.subplots(nb_rows, nb_cols, sharex=True, sharey=True)
+
+for i in range(nb_rows):
+    for j in range(nb_cols):
+        if i * nb_cols + j >= nb_obs:
+            break
+        axs[i, j].plot([obs[i * nb_cols + j] for obs in obses])
+        axs[i, j].set_title(obses_names[i * nb_cols + j])
+
+# set ylim between -5 and 5
+
+for ax in axs.flat:
+    ax.set_ylim([-5, 5])
+
+
+fig.suptitle(f"{args.data}")
+plt.show()

Open_Duck_Playground/playground/common/poly_reference_motion.py

@@ -0,0 +1,187 @@
+import jax.numpy as jp
+from jax import vmap
+import pickle
+
+
+# dimensions_names = [
+#     0  "pos left_hip_yaw",
+#     1  "pos left_hip_roll",
+#     2  "pos left_hip_pitch",
+#     3  "pos left_knee",
+#     4  "pos left_ankle",
+#     5  "pos neck_pitch",
+#     6  "pos head_pitch",
+#     7  "pos head_yaw",
+#     8  "pos head_roll",
+#     9  "pos left_antenna",
+#     10 "pos right_antenna",
+#     11 "pos right_hip_yaw",
+#     12 "pos right_hip_roll",
+#     13 "pos right_hip_pitch",
+#     14 "pos right_knee",
+#     15 "pos right_ankle",
+
+#     16 "vel left_hip_yaw",
+#     17 "vel left_hip_roll",
+#     18 "vel left_hip_pitch",
+#     19 "vel left_knee",
+#     20 "vel left_ankle",
+#     21 "vel neck_pitch",
+#     22 "vel head_pitch",
+#     23 "vel head_yaw",
+#     24 "vel head_roll",
+#     25 "vel left_antenna",
+#     26 "vel right_antenna",
+#     27 "vel right_hip_yaw",
+#     28 "vel right_hip_roll",
+#     29 "vel right_hip_pitch",
+#     30 "vel right_knee",
+#     31 "vel right_ankle",
+
+#     32 "foot_contacts left",
+#     33 "foot_contacts right",
+
+#     34 "base_linear_vel x",
+#     35 "base_linear_vel y",
+#     36 "base_linear_vel z",
+
+#     37 "base_angular_vel x",
+#     38 "base_angular_vel y",
+#     39 "base_angular_vel z",
+# ]
+
+
+class PolyReferenceMotion:
+    def __init__(self, polynomial_coefficients: str):
+        data = pickle.load(open(polynomial_coefficients, "rb"))
+        # data = json.load(open(polynomial_coefficients))
+        self.dx_range = [0, 0]
+        self.dy_range = [0, 0]
+        self.dtheta_range = [0, 0]
+        self.dxs = []
+        self.dys = []
+        self.dthetas = []
+        self.data_array = []
+        self.period = None
+        self.fps = None
+        self.frame_offsets = None
+        self.startend_double_support_ratio = None
+        self.start_offset = None
+        self.nb_steps_in_period = None
+
+        self.process(data)
+
+    def process(self, data):
+        print("[Poly ref data] Processing ...")
+        _data = {}
+        for name in data.keys():
+            split = name.split("_")
+            dx = float(split[0])
+            dy = float(split[1])
+            dtheta = float(split[2])
+
+            if self.period is None:
+                self.period = data[name]["period"]
+                self.fps = data[name]["fps"]
+                self.frame_offsets = data[name]["frame_offsets"]
+                self.startend_double_support_ratio = data[name][
+                    "startend_double_support_ratio"
+                ]
+                self.start_offset = int(self.startend_double_support_ratio * self.fps)
+                self.nb_steps_in_period = int(self.period * self.fps)
+
+            if dx not in self.dxs:
+                self.dxs.append(dx)
+
+            if dy not in self.dys:
+                self.dys.append(dy)
+
+            if dtheta not in self.dthetas:
+                self.dthetas.append(dtheta)
+
+            self.dx_range = [min(dx, self.dx_range[0]), max(dx, self.dx_range[1])]
+            self.dy_range = [min(dy, self.dy_range[0]), max(dy, self.dy_range[1])]
+            self.dtheta_range = [
+                min(dtheta, self.dtheta_range[0]),
+                max(dtheta, self.dtheta_range[1]),
+            ]
+
+            if dx not in _data:
+                _data[dx] = {}
+
+            if dy not in _data[dx]:
+                _data[dx][dy] = {}
+
+            if dtheta not in _data[dx][dy]:
+                _data[dx][dy][dtheta] = data[name]
+
+            _coeffs = data[name]["coefficients"]
+
+            coeffs = []
+            for k, v in _coeffs.items():
+                coeffs.append(jp.flip(jp.array(v)))
+            _data[dx][dy][dtheta] = coeffs
+
+        # print(self.dtheta_range)
+        # exit()
+
+        self.dxs = sorted(self.dxs)
+        self.dys = sorted(self.dys)
+        self.dthetas = sorted(self.dthetas)
+
+        nb_dx = len(self.dxs)
+        nb_dy = len(self.dys)
+        nb_dtheta = len(self.dthetas)
+
+        self.data_array = nb_dx * [None]
+        for x, dx in enumerate(self.dxs):
+            self.data_array[x] = nb_dy * [None]
+            for y, dy in enumerate(self.dys):
+                self.data_array[x][y] = nb_dtheta * [None]
+                for th, dtheta in enumerate(self.dthetas):
+                    self.data_array[x][y][th] = jp.array(_data[dx][dy][dtheta])
+
+        self.data_array = jp.array(self.data_array)
+
+        print("[Poly ref data] Done processing")
+
+    def vel_to_index(self, dx, dy, dtheta):
+
+        dx = jp.clip(dx, self.dx_range[0], self.dx_range[1])
+        dy = jp.clip(dy, self.dy_range[0], self.dy_range[1])
+        dtheta = jp.clip(dtheta, self.dtheta_range[0], self.dtheta_range[1])
+
+        ix = jp.argmin(jp.abs(jp.array(self.dxs) - dx))
+        iy = jp.argmin(jp.abs(jp.array(self.dys) - dy))
+        itheta = jp.argmin(jp.abs(jp.array(self.dthetas) - dtheta))
+
+        return ix, iy, itheta
+
+    def sample_polynomial(self, t, coeffs):
+        return vmap(lambda c: jp.polyval(c, t))(coeffs)
+
+    def get_reference_motion(self, dx, dy, dtheta, i):
+        ix, iy, itheta = self.vel_to_index(dx, dy, dtheta)
+        t = i % self.nb_steps_in_period / self.nb_steps_in_period
+        t = jp.clip(t, 0.0, 1.0)  # safeguard
+        ret = self.sample_polynomial(t, self.data_array[ix][iy][itheta])
+        return ret
+
+
+if __name__ == "__main__":
+
+    PRM = PolyReferenceMotion(
+        "playground/open_duck_mini_v2/data/polynomial_coefficients.pkl"
+    )
+    vals = []
+    select_dim = -1
+    for i in range(PRM.nb_steps_in_period):
+        vals.append(PRM.get_reference_motion(0.0, -0.05, -0.1, i)[select_dim])
+
+    # plot
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    ts = np.arange(0, PRM.nb_steps_in_period)
+    plt.plot(ts, vals)
+    plt.show()

Open_Duck_Playground/playground/common/poly_reference_motion_numpy.py

@@ -0,0 +1,138 @@
+import numpy as np
+import pickle
+
+
+class PolyReferenceMotion:
+    def __init__(self, polynomial_coefficients: str):
+        data = pickle.load(open(polynomial_coefficients, "rb"))
+        self.dx_range = [0, 0]
+        self.dy_range = [0, 0]
+        self.dtheta_range = [0, 0]
+        self.dxs = []
+        self.dys = []
+        self.dthetas = []
+        self.data_array = []
+        self.period = None
+        self.fps = None
+        self.frame_offsets = None
+        self.startend_double_support_ratio = None
+        self.start_offset = None
+        self.nb_steps_in_period = None
+
+        self.process(data)
+
+    def process(self, data):
+        print("[Poly ref data] Processing ...")
+        _data = {}
+        for name in data.keys():
+            split = name.split("_")
+            dx = float(split[0])
+            dy = float(split[1])
+            dtheta = float(split[2])
+
+            if self.period is None:
+                self.period = data[name]["period"]
+                self.fps = data[name]["fps"]
+                self.frame_offsets = data[name]["frame_offsets"]
+                self.startend_double_support_ratio = data[name][
+                    "startend_double_support_ratio"
+                ]
+                self.start_offset = int(self.startend_double_support_ratio * self.fps)
+                self.nb_steps_in_period = int(self.period * self.fps)
+
+            if dx not in self.dxs:
+                self.dxs.append(dx)
+
+            if dy not in self.dys:
+                self.dys.append(dy)
+
+            if dtheta not in self.dthetas:
+                self.dthetas.append(dtheta)
+
+            self.dx_range = [min(dx, self.dx_range[0]), max(dx, self.dx_range[1])]
+            self.dy_range = [min(dy, self.dy_range[0]), max(dy, self.dy_range[1])]
+            self.dtheta_range = [
+                min(dtheta, self.dtheta_range[0]),
+                max(dtheta, self.dtheta_range[1]),
+            ]
+
+            if dx not in _data:
+                _data[dx] = {}
+
+            if dy not in _data[dx]:
+                _data[dx][dy] = {}
+
+            if dtheta not in _data[dx][dy]:
+                _data[dx][dy][dtheta] = data[name]
+
+            _coeffs = data[name]["coefficients"]
+
+            coeffs = []
+            for k, v in _coeffs.items():
+                coeffs.append(v)
+            _data[dx][dy][dtheta] = coeffs
+
+        self.dxs = sorted(self.dxs)
+        self.dys = sorted(self.dys)
+        self.dthetas = sorted(self.dthetas)
+
+        nb_dx = len(self.dxs)
+        nb_dy = len(self.dys)
+        nb_dtheta = len(self.dthetas)
+
+        self.data_array = nb_dx * [None]
+        for x, dx in enumerate(self.dxs):
+            self.data_array[x] = nb_dy * [None]
+            for y, dy in enumerate(self.dys):
+                self.data_array[x][y] = nb_dtheta * [None]
+                for th, dtheta in enumerate(self.dthetas):
+                    self.data_array[x][y][th] = _data[dx][dy][dtheta]
+
+        self.data_array = self.data_array
+
+        print("[Poly ref data] Done processing")
+
+    def vel_to_index(self, dx, dy, dtheta):
+
+        dx = np.clip(dx, self.dx_range[0], self.dx_range[1])
+        dy = np.clip(dy, self.dy_range[0], self.dy_range[1])
+        dtheta = np.clip(dtheta, self.dtheta_range[0], self.dtheta_range[1])
+
+        ix = np.argmin(np.abs(np.array(self.dxs) - dx))
+        iy = np.argmin(np.abs(np.array(self.dys) - dy))
+        itheta = np.argmin(np.abs(np.array(self.dthetas) - dtheta))
+
+        return int(ix), int(iy), int(itheta)
+
+    def sample_polynomial(self, t, coeffs):
+        ret = []
+        for c in coeffs:
+            ret.append(np.polyval(np.flip(c), t))
+
+        return ret
+
+    def get_reference_motion(self, dx, dy, dtheta, i):
+        ix, iy, itheta = self.vel_to_index(dx, dy, dtheta)
+        t = i % self.nb_steps_in_period / self.nb_steps_in_period
+        t = np.clip(t, 0.0, 1.0)  # safeguard
+        ret = self.sample_polynomial(t, self.data_array[ix][iy][itheta])
+        return ret
+
+
+if __name__ == "__main__":
+
+    PRM = PolyReferenceMotion(
+        "playground/open_duck_mini_v2/data/polynomial_coefficients.pkl"
+    )
+    vals = []
+    select_dim = -1
+    for i in range(PRM.nb_steps_in_period):
+        vals.append(PRM.get_reference_motion(0.0, -0.05, -0.1, i)[select_dim])
+
+    # plot
+    import matplotlib.pyplot as plt
+    import numpy as np
+
+    ts = np.arange(0, PRM.nb_steps_in_period)
+    plt.plot(ts, vals)
+    plt.show()

Open_Duck_Playground/playground/common/randomize.py

@@ -0,0 +1,146 @@
+# Copyright 2025 DeepMind Technologies Limited
+# Copyright 2025 Antoine Pirrone - Steve Nguyen
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Domain randomization for the Open Duck Mini V2 environment. (based on Berkeley Humanoid)"""
+
+import jax
+from mujoco import mjx
+import jax.numpy as jp
+
+FLOOR_GEOM_ID = 0
+TORSO_BODY_ID = 1
+
+
+def domain_randomize(model: mjx.Model, rng: jax.Array):
+
+    # _dof_addr=jp.array([6,8,10,12,14,16,18,20,22,24])
+    # _joint_addr=jp.array([7,9,11,13,15,17,19,21,23,25])
+
+    dof_id = jp.array(
+        [idx for idx, fr in enumerate(model.dof_hasfrictionloss) if fr == True]
+    )  # for backlash joint we disable frictionloss
+    jnt_id = model.dof_jntid[dof_id]
+
+    dof_addr = jp.array([jadd for jadd in model.jnt_dofadr if jadd in dof_id])
+    joint_addr = model.jnt_qposadr[jnt_id]
+
+    @jax.vmap
+    def rand_dynamics(rng):
+        # Floor friction: =U(0.4, 1.0).
+        rng, key = jax.random.split(rng)
+        geom_friction = model.geom_friction.at[FLOOR_GEOM_ID, 0].set(
+            jax.random.uniform(key, minval=0.5, maxval=1.0)  # was 0.4, 1.0
+        )
+
+        # Scale static friction: *U(0.9, 1.1).
+        rng, key = jax.random.split(rng)
+        frictionloss = model.dof_frictionloss[dof_addr] * jax.random.uniform(
+            key, shape=(model.nu,), minval=0.9, maxval=1.1
+        )
+        dof_frictionloss = model.dof_frictionloss.at[dof_addr].set(frictionloss)
+
+        # Scale armature: *U(1.0, 1.05).
+        rng, key = jax.random.split(rng)
+        armature = model.dof_armature[dof_addr] * jax.random.uniform(
+            key, shape=(model.nu,), minval=1.0, maxval=1.05
+        )
+        dof_armature = model.dof_armature.at[dof_addr].set(armature)
+
+        # Jitter center of mass positiion: +U(-0.05, 0.05).
+        rng, key = jax.random.split(rng)
+        dpos = jax.random.uniform(key, (3,), minval=-0.05, maxval=0.05)
+        body_ipos = model.body_ipos.at[TORSO_BODY_ID].set(
+            model.body_ipos[TORSO_BODY_ID] + dpos
+        )
+
+        # Scale all link masses: *U(0.9, 1.1).
+        rng, key = jax.random.split(rng)
+        dmass = jax.random.uniform(key, shape=(model.nbody,), minval=0.9, maxval=1.1)
+        body_mass = model.body_mass.at[:].set(model.body_mass * dmass)
+
+        # Add mass to torso: +U(-0.2, 0.2).
+        rng, key = jax.random.split(rng)
+        dmass = jax.random.uniform(key, minval=-0.1, maxval=0.1)  # was -0.2, 0.2
+        body_mass = body_mass.at[TORSO_BODY_ID].set(body_mass[TORSO_BODY_ID] + dmass)
+
+        # Jitter qpos0: +U(-0.05, 0.05).
+        rng, key = jax.random.split(rng)
+        qpos0 = model.qpos0
+        qpos0 = qpos0.at[joint_addr].set(
+            qpos0[joint_addr]
+            + jax.random.uniform(
+                key, shape=(model.nu,), minval=-0.03, maxval=0.03
+            )  # was -0.05 0.05
+        )
+
+        # # Randomize KP
+        rng, key = jax.random.split(rng)
+        factor = jax.random.uniform(
+            key, shape=(model.nu,), minval=0.9, maxval=1.1
+        )  # was 0.8, 1.2
+        current_kp = model.actuator_gainprm[:, 0]
+        actuator_gainprm = model.actuator_gainprm.at[:, 0].set(current_kp * factor)
+        actuator_biasprm = model.actuator_biasprm.at[:, 1].set(-current_kp * factor)
+
+        return (
+            geom_friction,
+            body_ipos,
+            dof_frictionloss,
+            dof_armature,
+            body_mass,
+            qpos0,
+            actuator_gainprm,
+            actuator_biasprm,
+        )
+
+    (
+        friction,
+        body_ipos,
+        frictionloss,
+        armature,
+        body_mass,
+        qpos0,
+        actuator_gainprm,
+        actuator_biasprm,
+    ) = rand_dynamics(rng)
+
+    in_axes = jax.tree_util.tree_map(lambda x: None, model)
+    in_axes = in_axes.tree_replace(
+        {
+            "geom_friction": 0,
+            "body_ipos": 0,
+            "dof_frictionloss": 0,
+            "dof_armature": 0,
+            "body_mass": 0,
+            "qpos0": 0,
+            "actuator_gainprm": 0,
+            "actuator_biasprm": 0,
+        }
+    )
+
+    model = model.tree_replace(
+        {
+            "geom_friction": friction,
+            "body_ipos": body_ipos,
+            "dof_frictionloss": frictionloss,
+            "dof_armature": armature,
+            "body_mass": body_mass,
+            "qpos0": qpos0,
+            "actuator_gainprm": actuator_gainprm,
+            "actuator_biasprm": actuator_biasprm,
+        }
+    )
+
+    return model, in_axes

Open_Duck_Playground/playground/common/rewards.py

@@ -0,0 +1,241 @@
+"""
+Set of commonly used rewards
+For examples on how to use some rewards, look at https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/_src/locomotion/berkeley_humanoid/joystick.py
+"""
+
+import jax
+import jax.numpy as jp
+
+
+# Tracking rewards.
+def reward_tracking_lin_vel(
+    commands: jax.Array,
+    local_vel: jax.Array,
+    tracking_sigma: float,
+) -> jax.Array:
+    # lin_vel_error = jp.sum(jp.square(commands[:2] - local_vel[:2]))
+    # return jp.nan_to_num(jp.exp(-lin_vel_error / self._config.reward_config.tracking_sigma))
+    y_tol = 0.1
+    error_x = jp.square(commands[0] - local_vel[0])
+    error_y = jp.clip(jp.abs(local_vel[1] - commands[1]) - y_tol, 0.0, None)
+    lin_vel_error = error_x + jp.square(error_y)
+    return jp.nan_to_num(jp.exp(-lin_vel_error / tracking_sigma))
+
+
+def reward_tracking_ang_vel(
+    commands: jax.Array,
+    ang_vel: jax.Array,
+    tracking_sigma: float,
+) -> jax.Array:
+    ang_vel_error = jp.square(commands[2] - ang_vel[2])
+    return jp.nan_to_num(jp.exp(-ang_vel_error / tracking_sigma))
+
+
+# Base-related rewards.
+
+
+def cost_lin_vel_z(global_linvel) -> jax.Array:
+    return jp.nan_to_num(jp.square(global_linvel[2]))
+
+
+def cost_ang_vel_xy(global_angvel) -> jax.Array:
+    return jp.nan_to_num(jp.sum(jp.square(global_angvel[:2])))
+
+
+def cost_orientation(torso_zaxis: jax.Array) -> jax.Array:
+    return jp.nan_to_num(jp.sum(jp.square(torso_zaxis[:2])))
+
+
+def cost_base_height(base_height: jax.Array, base_height_target: float) -> jax.Array:
+    return jp.nan_to_num(jp.square(base_height - base_height_target))
+
+
+def reward_base_y_swing(
+    base_y_speed: jax.Array,
+    freq: float,
+    amplitude: float,
+    t: float,
+    tracking_sigma: float,
+) -> jax.Array:
+    target_y_speed = amplitude * jp.sin(2 * jp.pi * freq * t)
+    y_speed_error = jp.square(target_y_speed - base_y_speed)
+    return jp.nan_to_num(jp.exp(-y_speed_error / tracking_sigma))
+
+
+# Energy related rewards.
+
+
+def cost_torques(torques: jax.Array) -> jax.Array:
+    return jp.nan_to_num(jp.sum(jp.square(torques)))
+    # return jp.nan_to_num(jp.sum(jp.abs(torques)))
+
+
+def cost_energy(qvel: jax.Array, qfrc_actuator: jax.Array) -> jax.Array:
+    return jp.nan_to_num(jp.sum(jp.abs(qvel) * jp.abs(qfrc_actuator)))
+
+
+def cost_action_rate(act: jax.Array, last_act: jax.Array) -> jax.Array:
+    c1 = jp.nan_to_num(jp.sum(jp.square(act - last_act)))
+    return c1
+
+
+# Other rewards.
+
+
+def cost_joint_pos_limits(
+    qpos: jax.Array, soft_lowers: float, soft_uppers: float
+) -> jax.Array:
+    out_of_limits = -jp.clip(qpos - soft_lowers, None, 0.0)
+    out_of_limits += jp.clip(qpos - soft_uppers, 0.0, None)
+    return jp.nan_to_num(jp.sum(out_of_limits))
+
+
+def cost_stand_still(
+    commands: jax.Array,
+    qpos: jax.Array,
+    qvel: jax.Array,
+    default_pose: jax.Array,
+    ignore_head: bool = False,
+) -> jax.Array:
+    # TODO no hard coded slices
+    cmd_norm = jp.linalg.norm(commands[:3])
+    if not ignore_head:
+        pose_cost = jp.sum(jp.abs(qpos - default_pose))
+        vel_cost = jp.sum(jp.abs(qvel))
+    else:
+        left_leg_pos = qpos[:5]
+        right_leg_pos = qpos[9:]
+        left_leg_vel = qvel[:5]
+        right_leg_vel = qvel[9:]
+        left_leg_default = default_pose[:5]
+        right_leg_default = default_pose[9:]
+        pose_cost = jp.sum(jp.abs(left_leg_pos - left_leg_default)) + jp.sum(
+            jp.abs(right_leg_pos - right_leg_default)
+        )
+        vel_cost = jp.sum(jp.abs(left_leg_vel)) + jp.sum(jp.abs(right_leg_vel))
+
+    return jp.nan_to_num(pose_cost + vel_cost) * (cmd_norm < 0.01)
+
+
+def cost_termination(done: jax.Array) -> jax.Array:
+    return done
+
+
+def reward_alive() -> jax.Array:
+    return jp.array(1.0)
+
+
+# Pose-related rewards.
+
+
+def cost_head_pos(
+    joints_qpos: jax.Array,
+    joints_qvel: jax.Array,
+    cmd: jax.Array,
+) -> jax.Array:
+    move_cmd_norm = jp.linalg.norm(cmd[:3])
+    head_cmd = cmd[3:]
+    head_pos = joints_qpos[5:9]
+    # head_vel = joints_qvel[5:9]
+
+    # target_head_qvel = jp.zeros_like(head_cmd)
+
+    head_pos_error = jp.sum(jp.square(head_pos - head_cmd))
+
+    # head_vel_error = jp.sum(jp.square(head_vel - target_head_qvel))
+
+    return jp.nan_to_num(head_pos_error) * (move_cmd_norm > 0.01)
+    # return jp.nan_to_num(head_pos_error + head_vel_error)
+
+
+# FIXME
+def cost_joint_deviation_hip(
+    qpos: jax.Array, cmd: jax.Array, hip_indices: jax.Array, default_pose: jax.Array
+) -> jax.Array:
+    cost = jp.sum(jp.abs(qpos[hip_indices] - default_pose[hip_indices]))
+    cost *= jp.abs(cmd[1]) > 0.1
+    return jp.nan_to_num(cost)
+
+
+# FIXME
+def cost_joint_deviation_knee(
+    qpos: jax.Array, knee_indices: jax.Array, default_pose: jax.Array
+) -> jax.Array:
+    return jp.nan_to_num(
+        jp.sum(jp.abs(qpos[knee_indices] - default_pose[knee_indices]))
+    )
+
+
+# FIXME
+def cost_pose(
+    qpos: jax.Array, default_pose: jax.Array, weights: jax.Array
+) -> jax.Array:
+    return jp.nan_to_num(jp.sum(jp.square(qpos - default_pose) * weights))
+
+
+# Feet related rewards.
+
+
+# FIXME
+def cost_feet_slip(contact: jax.Array, global_linvel: jax.Array) -> jax.Array:
+    body_vel = global_linvel[:2]
+    reward = jp.sum(jp.linalg.norm(body_vel, axis=-1) * contact)
+    return jp.nan_to_num(reward)
+
+
+# FIXME
+def cost_feet_clearance(
+    feet_vel: jax.Array,
+    foot_pos: jax.Array,
+    max_foot_height: float,
+) -> jax.Array:
+    # feet_vel = data.sensordata[self._foot_linvel_sensor_adr]
+    vel_xy = feet_vel[..., :2]
+    vel_norm = jp.sqrt(jp.linalg.norm(vel_xy, axis=-1))
+    # foot_pos = data.site_xpos[self._feet_site_id]
+    foot_z = foot_pos[..., -1]
+    delta = jp.abs(foot_z - max_foot_height)
+    return jp.nan_to_num(jp.sum(delta * vel_norm))
+
+
+# FIXME
+def cost_feet_height(
+    swing_peak: jax.Array,
+    first_contact: jax.Array,
+    max_foot_height: float,
+) -> jax.Array:
+    error = swing_peak / max_foot_height - 1.0
+    return jp.nan_to_num(jp.sum(jp.square(error) * first_contact))
+
+
+# FIXME
+def reward_feet_air_time(
+    air_time: jax.Array,
+    first_contact: jax.Array,
+    commands: jax.Array,
+    threshold_min: float = 0.1,  # 0.2
+    threshold_max: float = 0.5,
+) -> jax.Array:
+    cmd_norm = jp.linalg.norm(commands[:3])
+    air_time = (air_time - threshold_min) * first_contact
+    air_time = jp.clip(air_time, max=threshold_max - threshold_min)
+    reward = jp.sum(air_time)
+    reward *= cmd_norm > 0.01  # No reward for zero commands.
+    return jp.nan_to_num(reward)
+
+
+# FIXME
+def reward_feet_phase(
+    foot_pos: jax.Array,
+    rz: jax.Array,
+) -> jax.Array:
+    # Reward for tracking the desired foot height.
+    # foot_pos = data.site_xpos[self._feet_site_id]
+    foot_z = foot_pos[..., -1]
+    # rz = gait.get_rz(phase, swing_height=foot_height)
+    error = jp.sum(jp.square(foot_z - rz))
+    reward = jp.exp(-error / 0.01)
+    # TODO(kevin): Ensure no movement at 0 command.
+    # cmd_norm = jp.linalg.norm(commands)
+    # reward *= cmd_norm > 0.1  # No reward for zero commands.
+    return jp.nan_to_num(reward)

Open_Duck_Playground/playground/common/rewards_numpy.py

@@ -0,0 +1,196 @@
+"""
+Set of commonly used rewards
+For examples on how to use some rewards, look at https://github.com/google-deepmind/mujoco_playground/blob/main/mujoco_playground/_src/locomotion/berkeley_humanoid/joystick.py
+"""
+
+# import jax
+# import jax.numpy as np
+
+import numpy as np
+
+
+# Tracking rewards.
+def reward_tracking_lin_vel(commands, local_vel, tracking_sigma):
+    # lin_vel_error = np.sum(np.square(commands[:2] - local_vel[:2]))
+    # return np.nan_to_num(np.exp(-lin_vel_error / self._config.reward_config.tracking_sigma))
+    y_tol = 0.1
+    error_x = np.square(commands[0] - local_vel[0])
+    error_y = np.clip(np.abs(local_vel[1] - commands[1]) - y_tol, 0.0, None)
+    lin_vel_error = error_x + np.square(error_y)
+    return np.nan_to_num(np.exp(-lin_vel_error / tracking_sigma))
+
+
+def reward_tracking_ang_vel(commands, ang_vel, tracking_sigma):
+    ang_vel_error = np.square(commands[2] - ang_vel[2])
+    return np.nan_to_num(np.exp(-ang_vel_error / tracking_sigma))
+
+
+# Base-related rewards.
+
+
+def cost_lin_vel_z(global_linvel):
+    return np.nan_to_num(np.square(global_linvel[2]))
+
+
+def cost_ang_vel_xy(global_angvel):
+    return np.nan_to_num(np.sum(np.square(global_angvel[:2])))
+
+
+def cost_orientation(torso_zaxis):
+    return np.nan_to_num(np.sum(np.square(torso_zaxis[:2])))
+
+
+def cost_base_height(base_height, base_height_target):
+    return np.nan_to_num(np.square(base_height - base_height_target))
+
+
+def reward_base_y_swing(base_y_speed, freq, amplitude, t, tracking_sigma):
+    target_y_speed = amplitude * np.sin(2 * np.pi * freq * t)
+    y_speed_error = np.square(target_y_speed - base_y_speed)
+    return np.nan_to_num(np.exp(-y_speed_error / tracking_sigma))
+
+
+# Energy related rewards.
+
+
+def cost_torques(torques):
+    return np.nan_to_num(np.sum(np.square(torques)))
+    # return np.nan_to_num(np.sum(np.abs(torques)))
+
+
+def cost_energy(qvel, qfrc_actuator):
+    return np.nan_to_num(np.sum(np.abs(qvel) * np.abs(qfrc_actuator)))
+
+
+def cost_action_rate(act, last_act):
+    c1 = np.nan_to_num(np.sum(np.square(act - last_act)))
+    return c1
+
+
+# Other rewards.
+
+
+def cost_joint_pos_limits(qpos, soft_lowers, soft_uppers):
+    out_of_limits = -np.clip(qpos - soft_lowers, None, 0.0)
+    out_of_limits += np.clip(qpos - soft_uppers, 0.0, None)
+    return np.nan_to_num(np.sum(out_of_limits))
+
+
+def cost_stand_still(commands, qpos, qvel, default_pose, ignore_head=False):
+    # TODO no hard coded slices
+    cmd_norm = np.linalg.norm(commands[:3])
+    if not ignore_head:
+        pose_cost = np.sum(np.abs(qpos - default_pose))
+        vel_cost = np.sum(np.abs(qvel))
+    else:
+        left_leg_pos = qpos[:5]
+        right_leg_pos = qpos[9:]
+        left_leg_vel = qvel[:5]
+        right_leg_vel = qvel[9:]
+        left_leg_default = default_pose[:5]
+        right_leg_default = default_pose[9:]
+        pose_cost = np.sum(np.abs(left_leg_pos - left_leg_default)) + np.sum(
+            np.abs(right_leg_pos - right_leg_default)
+        )
+        vel_cost = np.sum(np.abs(left_leg_vel)) + np.sum(np.abs(right_leg_vel))
+
+    return np.nan_to_num(pose_cost + vel_cost) * (cmd_norm < 0.01)
+
+
+def cost_termination(done):
+    return done
+
+
+def reward_alive():
+    return np.array(1.0)
+
+
+# Pose-related rewards.
+
+
+def cost_head_pos(joints_qpos, joints_qvel, cmd):
+    move_cmd_norm = np.linalg.norm(cmd[:3])
+    head_cmd = cmd[3:]
+    head_pos = joints_qpos[5:9]
+    # head_vel = joints_qvel[5:9]
+
+    # target_head_qvel = np.zeros_like(head_cmd)
+
+    head_pos_error = np.sum(np.square(head_pos - head_cmd))
+
+    # head_vel_error = np.sum(np.square(head_vel - target_head_qvel))
+
+    return np.nan_to_num(head_pos_error) * (move_cmd_norm > 0.01)
+    # return np.nan_to_num(head_pos_error + head_vel_error)
+
+
+# FIXME
+def cost_joint_deviation_hip(qpos, cmd, hip_indices, default_pose):
+    cost = np.sum(np.abs(qpos[hip_indices] - default_pose[hip_indices]))
+    cost *= np.abs(cmd[1]) > 0.1
+    return np.nan_to_num(cost)
+
+
+# FIXME
+def cost_joint_deviation_knee(qpos, knee_indices, default_pose):
+    return np.nan_to_num(
+        np.sum(np.abs(qpos[knee_indices] - default_pose[knee_indices]))
+    )
+
+
+# FIXME
+def cost_pose(qpos, default_pose, weights):
+    return np.nan_to_num(np.sum(np.square(qpos - default_pose) * weights))
+
+
+# Feet related rewards.
+
+
+# FIXME
+def cost_feet_slip(contact, global_linvel):
+    body_vel = global_linvel[:2]
+    reward = np.sum(np.linalg.norm(body_vel, axis=-1) * contact)
+    return np.nan_to_num(reward)
+
+
+# FIXME
+def cost_feet_clearance(feet_vel, foot_pos, max_foot_height):
+    # feet_vel = data.sensordata[self._foot_linvel_sensor_adr]
+    vel_xy = feet_vel[..., :2]
+    vel_norm = np.sqrt(np.linalg.norm(vel_xy, axis=-1))
+    # foot_pos = data.site_xpos[self._feet_site_id]
+    foot_z = foot_pos[..., -1]
+    delta = np.abs(foot_z - max_foot_height)
+    return np.nan_to_num(np.sum(delta * vel_norm))
+
+
+# FIXME
+def cost_feet_height(swing_peak, first_contact, max_foot_height):
+    error = swing_peak / max_foot_height - 1.0
+    return np.nan_to_num(np.sum(np.square(error) * first_contact))
+
+
+# FIXME
+def reward_feet_air_time(
+    air_time, first_contact, commands, threshold_min=0.1, threshold_max=0.5  # 0.2
+):
+    cmd_norm = np.linalg.norm(commands[:3])
+    air_time = (air_time - threshold_min) * first_contact
+    air_time = np.clip(air_time, max=threshold_max - threshold_min)
+    reward = np.sum(air_time)
+    reward *= cmd_norm > 0.01  # No reward for zero commands.
+    return np.nan_to_num(reward)
+
+
+# FIXME
+def reward_feet_phase(foot_pos, rz):
+    # Reward for tracking the desired foot height.
+    # foot_pos = data.site_xpos[self._feet_site_id]
+    foot_z = foot_pos[..., -1]
+    # rz = gait.get_rz(phase, swing_height=foot_height)
+    error = np.sum(np.square(foot_z - rz))
+    reward = np.exp(-error / 0.01)
+    # TODO(kevin): Ensure no movement at 0 command.
+    # cmd_norm = np.linalg.norm(commands)
+    # reward *= cmd_norm > 0.1  # No reward for zero commands.
+    return np.nan_to_num(reward)

Open_Duck_Playground/playground/common/runner.py

@@ -0,0 +1,118 @@
+"""
+Defines a common runner between the different robots.
+Inspired from https://github.com/kscalelabs/mujoco_playground/blob/master/playground/common/runner.py
+"""
+
+from pathlib import Path
+from abc import ABC
+import argparse
+import functools
+from datetime import datetime
+from flax.training import orbax_utils
+from tensorboardX import SummaryWriter
+
+import os
+from brax.training.agents.ppo import networks as ppo_networks, train as ppo
+from mujoco_playground import wrapper
+from mujoco_playground.config import locomotion_params
+from orbax import checkpoint as ocp
+import jax
+
+from playground.common.export_onnx import export_onnx
+
+
+class BaseRunner(ABC):
+    def __init__(self, args: argparse.Namespace) -> None:
+        """Initialize the Runner class.
+
+        Args:
+            args (argparse.Namespace): Command line arguments.
+        """
+        self.args = args
+        self.output_dir = args.output_dir
+        self.output_dir = Path.cwd() / Path(self.output_dir)
+
+        self.env_config = None
+        self.env = None
+        self.eval_env = None
+        self.randomizer = None
+        self.writer = SummaryWriter(log_dir=self.output_dir)
+        self.action_size = None
+        self.obs_size = None
+        self.num_timesteps = args.num_timesteps
+        self.restore_checkpoint_path = None
+        
+        # CACHE STUFF
+        os.makedirs(".tmp", exist_ok=True)
+        jax.config.update("jax_compilation_cache_dir", ".tmp/jax_cache")
+        jax.config.update("jax_persistent_cache_min_entry_size_bytes", -1)
+        jax.config.update("jax_persistent_cache_min_compile_time_secs", 0)
+        jax.config.update(
+            "jax_persistent_cache_enable_xla_caches",
+            "xla_gpu_per_fusion_autotune_cache_dir",
+        )
+        os.environ["JAX_COMPILATION_CACHE_DIR"] = ".tmp/jax_cache"
+
+    def progress_callback(self, num_steps: int, metrics: dict) -> None:
+
+        for metric_name, metric_value in metrics.items():
+            # Convert to float, but watch out for 0-dim JAX arrays
+            self.writer.add_scalar(metric_name, metric_value, num_steps)
+
+        print("-----------")
+        print(
+            f'STEP: {num_steps} reward: {metrics["eval/episode_reward"]} reward_std: {metrics["eval/episode_reward_std"]}'
+        )
+        print("-----------")
+
+    def policy_params_fn(self, current_step, make_policy, params):
+        # save checkpoints
+
+        orbax_checkpointer = ocp.PyTreeCheckpointer()
+        save_args = orbax_utils.save_args_from_target(params)
+        d = datetime.now().strftime("%Y_%m_%d_%H%M%S")
+        path = f"{self.output_dir}/{d}_{current_step}"
+        print(f"Saving checkpoint (step: {current_step}): {path}")
+        orbax_checkpointer.save(path, params, force=True, save_args=save_args)
+        onnx_export_path = f"{self.output_dir}/{d}_{current_step}.onnx"
+        export_onnx(
+            params,
+            self.action_size,
+            self.ppo_params,
+            self.obs_size,  # may not work
+            output_path=onnx_export_path
+        )
+
+    def train(self) -> None:
+        self.ppo_params = locomotion_params.brax_ppo_config(
+            "BerkeleyHumanoidJoystickFlatTerrain"
+        )  # TODO
+        self.ppo_training_params = dict(self.ppo_params)
+        # self.ppo_training_params["num_timesteps"] = 150000000 * 20
+        
+
+        if "network_factory" in self.ppo_params:
+            network_factory = functools.partial(
+                ppo_networks.make_ppo_networks, **self.ppo_params.network_factory
+            )
+            del self.ppo_training_params["network_factory"]
+        else:
+            network_factory = ppo_networks.make_ppo_networks
+        self.ppo_training_params["num_timesteps"] = self.num_timesteps
+        print(f"PPO params: {self.ppo_training_params}")
+
+        train_fn = functools.partial(
+            ppo.train,
+            **self.ppo_training_params,
+            network_factory=network_factory,
+            randomization_fn=self.randomizer,
+            progress_fn=self.progress_callback,
+            policy_params_fn=self.policy_params_fn,
+            restore_checkpoint_path=self.restore_checkpoint_path,
+        )
+
+        _, params, _ = train_fn(
+            environment=self.env,
+            eval_env=self.eval_env,
+            wrap_env_fn=wrapper.wrap_for_brax_training,
+        )

Open_Duck_Playground/playground/common/utils.py

@@ -0,0 +1,25 @@
+import jax.numpy as jp
+import jax
+
+
+class LowPassActionFilter:
+    def __init__(self, control_freq, cutoff_frequency=30.0):
+        self.last_action = 0
+        self.current_action = 0
+        self.control_freq = float(control_freq)
+        self.cutoff_frequency = float(cutoff_frequency)
+        self.alpha = self.compute_alpha()
+
+    def compute_alpha(self):
+        return (1.0 / self.cutoff_frequency) / (
+            1.0 / self.control_freq + 1.0 / self.cutoff_frequency
+        )
+
+    def push(self, action: jax.Array) -> None:
+        self.current_action = jp.array(action)
+
+    def get_filtered_action(self) -> jax.Array:
+        self.last_action = (
+            self.alpha * self.last_action + (1 - self.alpha) * self.current_action
+        )
+        return self.last_action

Open_Duck_Playground/playground/open_duck_mini_v2/base.py

@@ -0,0 +1,291 @@
+# Copyright 2025 DeepMind Technologies Limited
+# Copyright 2025 Antoine Pirrone - Steve Nguyen
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Base classes for Open Duck Mini V2. (based on Berkeley Humanoid)"""
+
+from typing import Any, Dict, Optional, Union
+
+from etils import epath
+import jax
+import jax.numpy as jp
+from ml_collections import config_dict
+import mujoco
+from mujoco import mjx
+
+from mujoco_playground._src import mjx_env
+from . import constants
+
+
+def get_assets() -> Dict[str, bytes]:
+    assets = {}
+    mjx_env.update_assets(assets, constants.ROOT_PATH / "xmls", "*.xml")
+    mjx_env.update_assets(assets, constants.ROOT_PATH / "xmls" / "assets")
+    path = constants.ROOT_PATH
+    mjx_env.update_assets(assets, path, "*.xml")
+    mjx_env.update_assets(assets, path / "assets")
+    return assets
+
+
+class OpenDuckMiniV2Env(mjx_env.MjxEnv):
+    """Base class for Open Duck Mini V2 environments."""
+
+    def __init__(
+        self,
+        xml_path: str,
+        config: config_dict.ConfigDict,
+        config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None,
+    ) -> None:
+        super().__init__(config, config_overrides)
+
+        print(f"xml: {xml_path}")
+        self._mj_model = mujoco.MjModel.from_xml_string(
+            epath.Path(xml_path).read_text(), assets=get_assets()
+        )
+        self._mj_model.opt.timestep = self.sim_dt
+
+        self._mj_model.vis.global_.offwidth = 3840
+        self._mj_model.vis.global_.offheight = 2160
+
+        self._mjx_model = mjx.put_model(self._mj_model)
+        self._xml_path = xml_path
+        self.floating_base_name= [self._mj_model.jnt(k).name for k in range(0, self._mj_model.njnt) if self._mj_model.jnt(k).type == 0][0] #assuming only one floating object!
+        self.actuator_names = [
+            self._mj_model.actuator(k).name for k in range(0, self._mj_model.nu)
+        ]  # will be useful to get only the actuators we care about
+        self.joint_names = [ #njnt = all joints (including floating base, actuators and backlash joints)
+            self._mj_model.jnt(k).name for k in range(0, self._mj_model.njnt)
+        ]  # all the joint (including the backlash joints)
+        self.backlash_joint_names = [
+            j for j in self.joint_names if j not in self.actuator_names and j not in self.floating_base_name
+        ]  # only the dummy backlash joint
+        self.all_joint_ids = [self.get_joint_id_from_name(n) for n in self.joint_names]
+        self.all_joint_qpos_addr = [self.get_joint_addr_from_name(n) for n in self.joint_names]
+
+        self.actuator_joint_ids = [
+            self.get_joint_id_from_name(n) for n in self.actuator_names
+        ]
+        self.actuator_joint_qpos_addr = [
+            self.get_joint_addr_from_name(n) for n in self.actuator_names
+        ]
+
+        self.backlash_joint_ids=[
+            self.get_joint_id_from_name(n) for n in self.backlash_joint_names
+        ]
+
+        self.backlash_joint_qpos_addr=[
+            self.get_joint_addr_from_name(n) for n in self.backlash_joint_names
+        ]
+
+        self.all_qvel_addr=jp.array([self._mj_model.jnt_dofadr[jad] for jad in self.all_joint_ids])
+        self.actuator_qvel_addr=jp.array([self._mj_model.jnt_dofadr[jad] for jad in self.actuator_joint_ids])
+
+        self.actuator_joint_dict = {
+            n: self.get_joint_id_from_name(n) for n in self.actuator_names
+        }
+
+        self._floating_base_qpos_addr = self._mj_model.jnt_qposadr[
+            jp.where(self._mj_model.jnt_type == 0)
+        ][
+            0
+        ]  # Assuming there is only one floating base! the jnt_type==0 is a floating joint. 3 is a hinge
+
+        self._floating_base_qvel_addr = self._mj_model.jnt_dofadr[
+            jp.where(self._mj_model.jnt_type == 0)
+        ][
+            0
+        ]  # Assuming there is only one floating base! the jnt_type==0 is a floating joint. 3 is a hinge
+
+        self._floating_base_id = self._mj_model.joint(self.floating_base_name).id
+
+        # self.all_joint_no_backlash_ids=jp.zeros(7+self._mj_model.nu)
+        # all_idx=self.actuator_joint_ids+list(range(self._floating_base_qpos_addr,self._floating_base_qpos_addr+7))
+        # all_idx=jp.array(all_idx).sort()
+        all_idx=self.actuator_joint_ids+list([self.get_joint_id_from_name("trunk_assembly_freejoint")])
+        all_idx=jp.array(all_idx).sort()
+        # self.all_joint_no_backlash_ids=[idx for idx in self.all_joint_ids if idx not in self.backlash_joint_ids]+list(range(self._floating_base_add,self._floating_base_add+7))
+        self.all_joint_no_backlash_ids=[idx for idx in all_idx]
+        # print(f"ALL: {self.all_joint_no_backlash_ids} back_id: {self.backlash_joint_ids} base_id: {list(range(self._floating_base_qpos_addr,self._floating_base_qpos_addr+7))}")
+
+        self.backlash_idx_to_add = []
+
+        for i, actuator_name in enumerate(self.actuator_names):
+            if actuator_name + "_backlash" not in self.backlash_joint_names:
+                self.backlash_idx_to_add.append(i)
+
+        print(f"actuators: {self.actuator_names}")
+        print(f"joints: {self.joint_names}")
+        print(f"backlash joints: {self.backlash_joint_names}")
+        print(f"actuator joints ids: {self.actuator_joint_ids}")
+        print(f"actuator joints dict: {self.actuator_joint_dict}")
+        print(f"floating qpos addr: {self._floating_base_qpos_addr} qvel addr: {self._floating_base_qvel_addr}")
+
+
+
+    def get_actuator_id_from_name(self, name: str) -> int:
+        """Return the id of a specified actuator"""
+        return mujoco.mj_name2id(self._mj_model, mujoco.mjtObj.mjOBJ_ACTUATOR, name)
+
+    def get_joint_id_from_name(self, name: str) -> int:
+        """Return the id of a specified joint"""
+        return mujoco.mj_name2id(self._mj_model, mujoco.mjtObj.mjOBJ_JOINT, name)
+
+
+    def get_joint_addr_from_name(self, name: str) -> int:
+        """Return the address of a specified joint"""
+        return self._mj_model.joint(name).qposadr
+
+    def get_dof_id_from_name(self, name: str) -> int:
+        """Return the id of a specified dof"""
+        return mujoco.mj_name2id(self._mj_model, mujoco.mjtObj.mjOBJ_DOF, name)
+
+
+    def get_actuator_joint_qpos_from_name(self, data: jax.Array, name: str) -> jax.Array:
+        """Return the qpos of a given actual joint"""
+        addr = self._mj_model.jnt_qposadr[self.actuator_joint_dict[name]]
+        return data[addr]
+
+    def get_actuator_joints_qpos_addr(self) -> jax.Array:
+        """Return the all the idx of actual joints"""
+        addr = jp.array(
+            [self._mj_model.jnt_qposadr[idx] for idx in self.actuator_joint_ids]
+        )
+        return addr
+
+    def get_floating_base_qpos(self, data:jax.Array) -> jax.Array:
+        return data[self._floating_base_qpos_addr:self._floating_base_qvel_addr+7]
+
+    def get_floating_base_qvel(self, data:jax.Array) -> jax.Array:
+        return data[self._floating_base_qvel_addr:self._floating_base_qvel_addr+6]
+
+
+    def set_floating_base_qpos(self, new_qpos:jax.Array, qpos:jax.Array) -> jax.Array:
+        return qpos.at[self._floating_base_qpos_addr:self._floating_base_qpos_addr+7].set(new_qpos)
+
+    def set_floating_base_qvel(self, new_qvel:jax.Array, qvel:jax.Array) -> jax.Array:
+        return qvel.at[self._floating_base_qvel_addr:self._floating_base_qvel_addr+6].set(new_qvel)
+
+
+    def exclude_backlash_joints_addr(self) -> jax.Array:
+        """Return the all the idx of actual joints and floating base"""
+        addr = jp.array(
+            [self._mj_model.jnt_qposadr[idx] for idx in self.all_joint_no_backlash_ids]
+        )
+        return addr
+
+
+    def get_all_joints_addr(self) -> jax.Array:
+        """Return the all the idx of all joints"""
+        addr = jp.array([self._mj_model.jnt_qposadr[idx] for idx in self.all_joint_ids])
+        return addr
+
+    def get_actuator_joints_qpos(self, data: jax.Array) -> jax.Array:
+        """Return the all the qpos of actual joints"""
+        return data[self.get_actuator_joints_qpos_addr()]
+
+    def set_actuator_joints_qpos(self, new_qpos: jax.Array, qpos: jax.Array) -> jax.Array:
+        """Set the qpos only for the actual joints (omit the backlash joint)"""
+        return qpos.at[self.get_actuator_joints_qpos_addr()].set(new_qpos)
+
+    def get_actuator_backlash_qpos(self, data: jax.Array) -> jax.Array:
+        """Return the all the qpos of backlash joints"""
+        if self.backlash_joint_qpos_addr == []:
+            return jp.array([])
+        return data[jp.array(self.backlash_joint_qpos_addr)]
+
+
+    def get_actuator_joints_qvel(self, data: jax.Array) -> jax.Array:
+        """Return the all the qvel of actual joints"""
+        return data[self.actuator_qvel_addr]
+
+    def set_actuator_joints_qvel(self, new_qvel: jax.Array, qvel: jax.Array) -> jax.Array:
+        """Set the qvel only for the actual joints (omit the backlash joint)"""
+        return qvel.at[self.actuator_qvel_addr].set(new_qvel)
+
+    def get_all_joints_qpos(self, data: jax.Array) -> jax.Array:
+        """Return the all the qpos of all joints"""
+        return data[self.get_all_joints_addr()]
+
+    def get_all_joints_qvel(self, data: jax.Array) -> jax.Array:
+        """Return the all the qvel of all joints"""
+        return data[self.all_qvel_addr]
+
+    def get_joints_nobacklash_qpos(self, data: jax.Array) -> jax.Array:
+        """Return the all the qpos of actual joints with the floating base"""
+        return data[self.exclude_backlash_joints_addr()]
+
+    def set_complete_qpos_from_joints(self, no_backlash_qpos: jax.Array, full_qpos: jax.Array) -> jax.Array:
+        """In the case of backlash joints, we want to ignore them (remove them) but we still need to set the complete state incuding them"""
+        full_qpos.at[self.exclude_backlash_joints_addr()].set(no_backlash_qpos)
+        return jp.array(full_qpos)
+
+    # Sensor readings.
+    def get_gravity(self, data: mjx.Data) -> jax.Array:
+        """Return the gravity vector in the world frame."""
+        return mjx_env.get_sensor_data(self.mj_model, data, constants.GRAVITY_SENSOR)
+
+    def get_global_linvel(self, data: mjx.Data) -> jax.Array:
+        """Return the linear velocity of the robot in the world frame."""
+        return mjx_env.get_sensor_data(
+            self.mj_model, data, constants.GLOBAL_LINVEL_SENSOR
+        )
+
+    def get_global_angvel(self, data: mjx.Data) -> jax.Array:
+        """Return the angular velocity of the robot in the world frame."""
+        return mjx_env.get_sensor_data(
+            self.mj_model, data, constants.GLOBAL_ANGVEL_SENSOR
+        )
+
+    def get_local_linvel(self, data: mjx.Data) -> jax.Array:
+        """Return the linear velocity of the robot in the local frame."""
+        return mjx_env.get_sensor_data(
+            self.mj_model, data, constants.LOCAL_LINVEL_SENSOR
+        )
+
+    def get_accelerometer(self, data: mjx.Data) -> jax.Array:
+        """Return the accelerometer readings in the local frame."""
+        return mjx_env.get_sensor_data(
+            self.mj_model, data, constants.ACCELEROMETER_SENSOR
+        )
+
+    def get_gyro(self, data: mjx.Data) -> jax.Array:
+        """Return the gyroscope readings in the local frame."""
+        return mjx_env.get_sensor_data(self.mj_model, data, constants.GYRO_SENSOR)
+
+    def get_feet_pos(self, data: mjx.Data) -> jax.Array:
+        """Return the position of the feet in the world frame."""
+        return jp.vstack(
+            [
+                mjx_env.get_sensor_data(self.mj_model, data, sensor_name)
+                for sensor_name in constants.FEET_POS_SENSOR
+            ]
+        )
+
+    # Accessors.
+
+    @property
+    def xml_path(self) -> str:
+        return self._xml_path
+
+    @property
+    def action_size(self) -> int:
+        return self._mjx_model.nu
+
+    @property
+    def mj_model(self) -> mujoco.MjModel:
+        return self._mj_model
+
+    @property
+    def mjx_model(self) -> mjx.Model:
+        return self._mjx_model

Open_Duck_Playground/playground/open_duck_mini_v2/constants.py

@@ -0,0 +1,89 @@
+# Copyright 2025 DeepMind Technologies Limited
+# Copyright 2025 Antoine Pirrone - Steve Nguyen
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Constants for Open Duck Mini V2. (based on Berkeley Humanoid)"""
+
+from etils import epath
+
+
+ROOT_PATH = epath.Path(__file__).parent
+FLAT_TERRAIN_XML = ROOT_PATH / "xmls" / "scene_flat_terrain.xml"
+ROUGH_TERRAIN_XML = ROOT_PATH / "xmls" / "scene_rough_terrain.xml"
+FLAT_TERRAIN_BACKLASH_XML = ROOT_PATH / "xmls" / "scene_flat_terrain_backlash.xml"
+ROUGH_TERRAIN_BACKLASH_XML = ROOT_PATH / "xmls" / "scene_rough_terrain_backlash.xml"
+
+
+def task_to_xml(task_name: str) -> epath.Path:
+    return {
+        "flat_terrain": FLAT_TERRAIN_XML,
+        "rough_terrain": ROUGH_TERRAIN_XML,
+        "flat_terrain_backlash": FLAT_TERRAIN_BACKLASH_XML,
+        "rough_terrain_backlash": ROUGH_TERRAIN_BACKLASH_XML,
+    }[task_name]
+
+
+FEET_SITES = [
+    "left_foot",
+    "right_foot",
+]
+
+LEFT_FEET_GEOMS = [
+    "left_foot_bottom_tpu",
+]
+
+RIGHT_FEET_GEOMS = [
+    "right_foot_bottom_tpu",
+]
+
+HIP_JOINT_NAMES = [
+    "left_hip_yaw",
+    "left_hip_roll",
+    "left_hip_pitch",
+    "right_hip_yaw",
+    "right_hip_roll",
+    "right_hip_pitch",
+]
+
+KNEE_JOINT_NAMES = [
+    "left_knee",
+    "right_knee",
+]
+
+# There should be a way to get that from the mjModel...
+JOINTS_ORDER_NO_HEAD = [
+    "left_hip_yaw",
+    "left_hip_roll",
+    "left_hip_pitch",
+    "left_knee",
+    "left_ankle",
+    "right_hip_yaw",
+    "right_hip_roll",
+    "right_hip_pitch",
+    "right_knee",
+    "right_ankle",
+]
+
+FEET_GEOMS = LEFT_FEET_GEOMS + RIGHT_FEET_GEOMS
+
+FEET_POS_SENSOR = [f"{site}_pos" for site in FEET_SITES]
+
+ROOT_BODY = "trunk_assembly"
+
+GRAVITY_SENSOR = "upvector"
+GLOBAL_LINVEL_SENSOR = "global_linvel"
+GLOBAL_ANGVEL_SENSOR = "global_angvel"
+LOCAL_LINVEL_SENSOR = "local_linvel"
+ACCELEROMETER_SENSOR = "accelerometer"
+GYRO_SENSOR = "gyro"

Open_Duck_Playground/playground/open_duck_mini_v2/custom_rewards.py

@@ -0,0 +1,149 @@
+import jax
+import jax.numpy as jp
+
+def reward_imitation(
+    base_qpos: jax.Array,
+    base_qvel: jax.Array,
+    joints_qpos: jax.Array,
+    joints_qvel: jax.Array,
+    contacts: jax.Array,
+    reference_frame: jax.Array,
+    cmd: jax.Array,
+    use_imitation_reward: bool = False,
+) -> jax.Array:
+    if not use_imitation_reward:
+        return jp.nan_to_num(0.0)
+
+    # TODO don't reward for moving when the command is zero.
+    cmd_norm = jp.linalg.norm(cmd[:3])
+
+    w_torso_pos = 1.0
+    w_torso_orientation = 1.0
+    w_lin_vel_xy = 1.0
+    w_lin_vel_z = 1.0
+    w_ang_vel_xy = 0.5
+    w_ang_vel_z = 0.5
+    w_joint_pos = 15.0
+    w_joint_vel = 1.0e-3
+    w_contact = 1.0
+
+    #  TODO : double check if the slices are correct
+    linear_vel_slice_start = 34
+    linear_vel_slice_end = 37
+
+    angular_vel_slice_start = 37
+    angular_vel_slice_end = 40
+
+    joint_pos_slice_start = 0
+    joint_pos_slice_end = 16
+
+    joint_vels_slice_start = 16
+    joint_vels_slice_end = 32
+
+    # root_pos_slice_start = 0
+    # root_pos_slice_end = 3
+
+    root_quat_slice_start = 3
+    root_quat_slice_end = 7
+
+    # left_toe_pos_slice_start = 23
+    # left_toe_pos_slice_end = 26
+
+    # right_toe_pos_slice_start = 26
+    # right_toe_pos_slice_end = 29
+
+    foot_contacts_slice_start = 32
+    foot_contacts_slice_end = 34
+
+    # ref_base_pos = reference_frame[root_pos_slice_start:root_pos_slice_end]
+    # base_pos = qpos[:3]
+
+    ref_base_orientation_quat = reference_frame[
+        root_quat_slice_start:root_quat_slice_end
+    ]
+    ref_base_orientation_quat = ref_base_orientation_quat / jp.linalg.norm(
+        ref_base_orientation_quat
+    )  # normalize the quat
+    base_orientation = base_qpos[3:7]
+    base_orientation = base_orientation / jp.linalg.norm(
+        base_orientation
+    )  # normalize the quat
+
+    ref_base_lin_vel = reference_frame[linear_vel_slice_start:linear_vel_slice_end]
+    base_lin_vel = base_qvel[:3]
+
+    ref_base_ang_vel = reference_frame[angular_vel_slice_start:angular_vel_slice_end]
+    base_ang_vel = base_qvel[3:6]
+
+    ref_joint_pos = reference_frame[joint_pos_slice_start:joint_pos_slice_end]
+    # remove neck head and antennas
+    ref_joint_pos = jp.concatenate([ref_joint_pos[:5], ref_joint_pos[11:]])
+    # joint_pos = joints_qpos
+    joint_pos = jp.concatenate([joints_qpos[:5], joints_qpos[9:]])
+
+    ref_joint_vels = reference_frame[joint_vels_slice_start:joint_vels_slice_end]
+    # remove neck head and antennas
+    ref_joint_vels = jp.concatenate([ref_joint_vels[:5], ref_joint_vels[11:]])
+    # joint_vel = joints_qvel
+    joint_vel = jp.concatenate([joints_qvel[:5], joints_qvel[9:]])
+
+    # ref_left_toe_pos = reference_frame[left_toe_pos_slice_start:left_toe_pos_slice_end]
+    # ref_right_toe_pos = reference_frame[right_toe_pos_slice_start:right_toe_pos_slice_end]
+
+    ref_foot_contacts = reference_frame[
+        foot_contacts_slice_start:foot_contacts_slice_end
+    ]
+
+    # reward
+    # torso_pos_rew = jp.exp(-200.0 * jp.sum(jp.square(base_pos[:2] - ref_base_pos[:2]))) * w_torso_pos
+
+    # real quaternion angle doesn't have the expected  effect, switching back for now
+    # torso_orientation_rew = jp.exp(-20 * self.quaternion_angle(base_orientation, ref_base_orientation_quat)) * w_torso_orientation
+
+    # TODO ignore yaw here, we just want xy orientation
+    torso_orientation_rew = (
+        jp.exp(-20.0 * jp.sum(jp.square(base_orientation - ref_base_orientation_quat)))
+        * w_torso_orientation
+    )
+
+    lin_vel_xy_rew = (
+        jp.exp(-8.0 * jp.sum(jp.square(base_lin_vel[:2] - ref_base_lin_vel[:2])))
+        * w_lin_vel_xy
+    )
+    lin_vel_z_rew = (
+        jp.exp(-8.0 * jp.sum(jp.square(base_lin_vel[2] - ref_base_lin_vel[2])))
+        * w_lin_vel_z
+    )
+
+    ang_vel_xy_rew = (
+        jp.exp(-2.0 * jp.sum(jp.square(base_ang_vel[:2] - ref_base_ang_vel[:2])))
+        * w_ang_vel_xy
+    )
+    ang_vel_z_rew = (
+        jp.exp(-2.0 * jp.sum(jp.square(base_ang_vel[2] - ref_base_ang_vel[2])))
+        * w_ang_vel_z
+    )
+
+    joint_pos_rew = -jp.sum(jp.square(joint_pos - ref_joint_pos)) * w_joint_pos
+    joint_vel_rew = -jp.sum(jp.square(joint_vel - ref_joint_vels)) * w_joint_vel
+
+    ref_foot_contacts = jp.where(
+        ref_foot_contacts > 0.5,
+        jp.ones_like(ref_foot_contacts),
+        jp.zeros_like(ref_foot_contacts),
+    )
+    contact_rew = jp.sum(contacts == ref_foot_contacts) * w_contact
+
+    reward = (
+        lin_vel_xy_rew
+        + lin_vel_z_rew
+        + ang_vel_xy_rew
+        + ang_vel_z_rew
+        + joint_pos_rew
+        + joint_vel_rew
+        + contact_rew
+        # + torso_orientation_rew
+    )
+
+    reward *= cmd_norm > 0.01  # No reward for zero commands.
+    return jp.nan_to_num(reward)

Open_Duck_Playground/playground/open_duck_mini_v2/custom_rewards_numpy.py

@@ -0,0 +1,151 @@
+# import jax
+# import jax.numpy as np
+import numpy as np
+
+
+def reward_imitation(
+    base_qpos,
+    base_qvel,
+    joints_qpos,
+    joints_qvel,
+    contacts,
+    reference_frame,
+    cmd,
+    use_imitation_reward=False,
+):
+    if not use_imitation_reward:
+        return np.nan_to_num(0.0)
+
+    # TODO don't reward for moving when the command is zero.
+    cmd_norm = np.linalg.norm(cmd[:3])
+
+    w_torso_pos = 1.0
+    w_torso_orientation = 1.0
+    w_lin_vel_xy = 1.0
+    w_lin_vel_z = 1.0
+    w_ang_vel_xy = 0.5
+    w_ang_vel_z = 0.5
+    w_joint_pos = 15.0
+    w_joint_vel = 1.0e-3
+    w_contact = 1.0
+
+    #  TODO : double check if the slices are correct
+    linear_vel_slice_start = 34
+    linear_vel_slice_end = 37
+
+    angular_vel_slice_start = 37
+    angular_vel_slice_end = 40
+
+    joint_pos_slice_start = 0
+    joint_pos_slice_end = 16
+
+    joint_vels_slice_start = 16
+    joint_vels_slice_end = 32
+
+    # root_pos_slice_start = 0
+    # root_pos_slice_end = 3
+
+    root_quat_slice_start = 3
+    root_quat_slice_end = 7
+
+    # left_toe_pos_slice_start = 23
+    # left_toe_pos_slice_end = 26
+
+    # right_toe_pos_slice_start = 26
+    # right_toe_pos_slice_end = 29
+
+    foot_contacts_slice_start = 32
+    foot_contacts_slice_end = 34
+
+    # ref_base_pos = reference_frame[root_pos_slice_start:root_pos_slice_end]
+    # base_pos = qpos[:3]
+
+    ref_base_orientation_quat = reference_frame[
+        root_quat_slice_start:root_quat_slice_end
+    ]
+    ref_base_orientation_quat = ref_base_orientation_quat / np.linalg.norm(
+        ref_base_orientation_quat
+    )  # normalize the quat
+    base_orientation = base_qpos[3:7]
+    base_orientation = base_orientation / np.linalg.norm(
+        base_orientation
+    )  # normalize the quat
+
+    ref_base_lin_vel = reference_frame[linear_vel_slice_start:linear_vel_slice_end]
+    base_lin_vel = base_qvel[:3]
+
+    ref_base_ang_vel = reference_frame[angular_vel_slice_start:angular_vel_slice_end]
+    base_ang_vel = base_qvel[3:6]
+
+    ref_joint_pos = reference_frame[joint_pos_slice_start:joint_pos_slice_end]
+    # remove neck head and antennas
+    ref_joint_pos = np.concatenate([ref_joint_pos[:5], ref_joint_pos[11:]])
+    # joint_pos = joints_qpos
+    joint_pos = np.concatenate([joints_qpos[:5], joints_qpos[9:]])
+
+    ref_joint_vels = reference_frame[joint_vels_slice_start:joint_vels_slice_end]
+    # remove neck head and antennas
+    ref_joint_vels = np.concatenate([ref_joint_vels[:5], ref_joint_vels[11:]])
+    # joint_vel = joints_qvel
+    joint_vel = np.concatenate([joints_qvel[:5], joints_qvel[9:]])
+
+    # ref_left_toe_pos = reference_frame[left_toe_pos_slice_start:left_toe_pos_slice_end]
+    # ref_right_toe_pos = reference_frame[right_toe_pos_slice_start:right_toe_pos_slice_end]
+
+    ref_foot_contacts = reference_frame[
+        foot_contacts_slice_start:foot_contacts_slice_end
+    ]
+
+    # reward
+    # torso_pos_rew = np.exp(-200.0 * np.sum(np.square(base_pos[:2] - ref_base_pos[:2]))) * w_torso_pos
+
+    # real quaternion angle doesn't have the expected  effect, switching back for now
+    # torso_orientation_rew = np.exp(-20 * self.quaternion_angle(base_orientation, ref_base_orientation_quat)) * w_torso_orientation
+
+    # TODO ignore yaw here, we just want xy orientation
+    torso_orientation_rew = (
+        np.exp(-20.0 * np.sum(np.square(base_orientation - ref_base_orientation_quat)))
+        * w_torso_orientation
+    )
+
+    lin_vel_xy_rew = (
+        np.exp(-8.0 * np.sum(np.square(base_lin_vel[:2] - ref_base_lin_vel[:2])))
+        * w_lin_vel_xy
+    )
+    lin_vel_z_rew = (
+        np.exp(-8.0 * np.sum(np.square(base_lin_vel[2] - ref_base_lin_vel[2])))
+        * w_lin_vel_z
+    )
+
+    ang_vel_xy_rew = (
+        np.exp(-2.0 * np.sum(np.square(base_ang_vel[:2] - ref_base_ang_vel[:2])))
+        * w_ang_vel_xy
+    )
+    ang_vel_z_rew = (
+        np.exp(-2.0 * np.sum(np.square(base_ang_vel[2] - ref_base_ang_vel[2])))
+        * w_ang_vel_z
+    )
+
+    joint_pos_rew = -np.sum(np.square(joint_pos - ref_joint_pos)) * w_joint_pos
+    joint_vel_rew = -np.sum(np.square(joint_vel - ref_joint_vels)) * w_joint_vel
+
+    ref_foot_contacts = np.where(
+        np.array(ref_foot_contacts) > 0.5,
+        np.ones_like(ref_foot_contacts),
+        np.zeros_like(ref_foot_contacts),
+    )
+    contact_rew = np.sum(contacts == ref_foot_contacts) * w_contact
+
+    reward = (
+        lin_vel_xy_rew
+        + lin_vel_z_rew
+        + ang_vel_xy_rew
+        + ang_vel_z_rew
+        + joint_pos_rew
+        + joint_vel_rew
+        + contact_rew
+        # + torso_orientation_rew
+    )
+
+    reward *= cmd_norm > 0.01  # No reward for zero commands.
+    return np.nan_to_num(reward)

Open_Duck_Playground/playground/open_duck_mini_v2/joystick.py

@@ -0,0 +1,725 @@
+# Copyright 2025 DeepMind Technologies Limited
+# Copyright 2025 Antoine Pirrone - Steve Nguyen
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Joystick task for Open Duck Mini V2. (based on Berkeley Humanoid)"""
+
+from typing import Any, Dict, Optional, Union
+import jax
+import jax.numpy as jp
+from ml_collections import config_dict
+from mujoco import mjx
+from mujoco.mjx._src import math
+import numpy as np
+
+from mujoco_playground._src import mjx_env
+from mujoco_playground._src.collision import geoms_colliding
+
+from . import constants
+from . import base as open_duck_mini_v2_base
+
+# from playground.common.utils import LowPassActionFilter
+from playground.common.poly_reference_motion import PolyReferenceMotion
+from playground.common.rewards import (
+    reward_tracking_lin_vel,
+    reward_tracking_ang_vel,
+    cost_torques,
+    cost_action_rate,
+    cost_stand_still,
+    reward_alive,
+)
+from playground.open_duck_mini_v2.custom_rewards import reward_imitation
+
+# if set to false, won't require the reference data to be present and won't compute the reference motions polynoms for nothing
+USE_IMITATION_REWARD = True
+USE_MOTOR_SPEED_LIMITS = True
+
+
+def default_config() -> config_dict.ConfigDict:
+    return config_dict.create(
+        ctrl_dt=0.02,
+        sim_dt=0.002,
+        episode_length=1000,
+        action_repeat=1,
+        action_scale=0.25,
+        dof_vel_scale=0.05,
+        history_len=0,
+        soft_joint_pos_limit_factor=0.95,
+        max_motor_velocity=5.24,  # rad/s
+        noise_config=config_dict.create(
+            level=1.0,  # Set to 0.0 to disable noise.
+            action_min_delay=0,  # env steps
+            action_max_delay=3,  # env steps
+            imu_min_delay=0,  # env steps
+            imu_max_delay=3,  # env steps
+            scales=config_dict.create(
+                hip_pos=0.03,  # rad, for each hip joint
+                knee_pos=0.05,  # rad, for each knee joint
+                ankle_pos=0.08,  # rad, for each ankle joint
+                joint_vel=2.5,  # rad/s # Was 1.5
+                gravity=0.1,
+                linvel=0.1,
+                gyro=0.1,
+                accelerometer=0.05,
+            ),
+        ),
+        reward_config=config_dict.create(
+            scales=config_dict.create(
+                tracking_lin_vel=2.5,
+                tracking_ang_vel=6.0,
+                torques=-1.0e-3,
+                action_rate=-0.5,  # was -1.5
+                stand_still=-0.2,  # was -1.0 TODO try to relax this a bit ?
+                alive=20.0,
+                imitation=1.0,
+            ),
+            tracking_sigma=0.01,  # was working at 0.01
+        ),
+        push_config=config_dict.create(
+            enable=True,
+            interval_range=[5.0, 10.0],
+            magnitude_range=[0.1, 1.0],
+        ),
+        lin_vel_x=[-0.15, 0.15],
+        lin_vel_y=[-0.2, 0.2],
+        ang_vel_yaw=[-1.0, 1.0],  # [-1.0, 1.0]
+        neck_pitch_range=[-0.34, 1.1],
+        head_pitch_range=[-0.78, 0.78],
+        head_yaw_range=[-1.5, 1.5],
+        head_roll_range=[-0.5, 0.5],
+        head_range_factor=1.0,  # to make it easier
+    )
+
+
+class Joystick(open_duck_mini_v2_base.OpenDuckMiniV2Env):
+    """Track a joystick command."""
+
+    def __init__(
+        self,
+        task: str = "flat_terrain",
+        config: config_dict.ConfigDict = default_config(),
+        config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None,
+    ):
+        super().__init__(
+            xml_path=constants.task_to_xml(task).as_posix(),
+            config=config,
+            config_overrides=config_overrides,
+        )
+        self._post_init()
+
+    def _post_init(self) -> None:
+
+        self._init_q = jp.array(self._mj_model.keyframe("home").qpos)
+        self._default_actuator = self._mj_model.keyframe(
+            "home"
+        ).ctrl  # ctrl of all the actual joints (no floating base and no backlash)
+
+        if USE_IMITATION_REWARD:
+            self.PRM = PolyReferenceMotion(
+                "playground/open_duck_mini_v2/data/polynomial_coefficients.pkl"
+            )
+
+        # Note: First joint is freejoint.
+        # get the range of the joints
+        self._lowers, self._uppers = self.mj_model.jnt_range[1:].T
+        c = (self._lowers + self._uppers) / 2
+        r = self._uppers - self._lowers
+        self._soft_lowers = c - 0.5 * r * self._config.soft_joint_pos_limit_factor
+        self._soft_uppers = c + 0.5 * r * self._config.soft_joint_pos_limit_factor
+
+        # weights for computing the cost of each joints compared to a reference pose
+        self._weights = jp.array(
+            [
+                1.0,
+                1.0,
+                0.01,
+                0.01,
+                1.0,  # left leg.
+                # 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, #head
+                1.0,
+                1.0,
+                0.01,
+                0.01,
+                1.0,  # right leg.
+            ]
+        )
+
+        self._njoints = self._mj_model.njnt  # number of joints
+        self._actuators = self._mj_model.nu  # number of actuators
+
+        self._torso_body_id = self._mj_model.body(constants.ROOT_BODY).id
+        self._torso_mass = self._mj_model.body_subtreemass[self._torso_body_id]
+        self._site_id = self._mj_model.site("imu").id
+
+        self._feet_site_id = np.array(
+            [self._mj_model.site(name).id for name in constants.FEET_SITES]
+        )
+        self._floor_geom_id = self._mj_model.geom("floor").id
+        self._feet_geom_id = np.array(
+            [self._mj_model.geom(name).id for name in constants.FEET_GEOMS]
+        )
+
+        foot_linvel_sensor_adr = []
+        for site in constants.FEET_SITES:
+            sensor_id = self._mj_model.sensor(f"{site}_global_linvel").id
+            sensor_adr = self._mj_model.sensor_adr[sensor_id]
+            sensor_dim = self._mj_model.sensor_dim[sensor_id]
+            foot_linvel_sensor_adr.append(
+                list(range(sensor_adr, sensor_adr + sensor_dim))
+            )
+        self._foot_linvel_sensor_adr = jp.array(foot_linvel_sensor_adr)
+
+        # # noise in the simu?
+        qpos_noise_scale = np.zeros(self._actuators)
+
+        hip_ids = [
+            idx for idx, j in enumerate(constants.JOINTS_ORDER_NO_HEAD) if "_hip" in j
+        ]
+        knee_ids = [
+            idx for idx, j in enumerate(constants.JOINTS_ORDER_NO_HEAD) if "_knee" in j
+        ]
+        ankle_ids = [
+            idx for idx, j in enumerate(constants.JOINTS_ORDER_NO_HEAD) if "_ankle" in j
+        ]
+
+        qpos_noise_scale[hip_ids] = self._config.noise_config.scales.hip_pos
+        qpos_noise_scale[knee_ids] = self._config.noise_config.scales.knee_pos
+        qpos_noise_scale[ankle_ids] = self._config.noise_config.scales.ankle_pos
+        # qpos_noise_scale[faa_ids] = self._config.noise_config.scales.faa_pos
+        self._qpos_noise_scale = jp.array(qpos_noise_scale)
+
+        # self.action_filter = LowPassActionFilter(
+        #     1 / self._config.ctrl_dt, cutoff_frequency=37.5
+        # )
+
+    def reset(self, rng: jax.Array) -> mjx_env.State:
+        qpos = self._init_q  # the complete qpos
+        # print(f'DEBUG0 init qpos: {qpos}')
+        qvel = jp.zeros(self.mjx_model.nv)
+
+        # init position/orientation in environment
+        # x=+U(-0.05, 0.05), y=+U(-0.05, 0.05), yaw=U(-3.14, 3.14).
+        rng, key = jax.random.split(rng)
+        dxy = jax.random.uniform(key, (2,), minval=-0.05, maxval=0.05)
+
+        # floating base
+        base_qpos = self.get_floating_base_qpos(qpos)
+        base_qpos = base_qpos.at[0:2].set(
+            qpos[self._floating_base_qpos_addr : self._floating_base_qpos_addr + 2]
+            + dxy
+        )  # x y noise
+
+        rng, key = jax.random.split(rng)
+        yaw = jax.random.uniform(key, (1,), minval=-3.14, maxval=3.14)
+        quat = math.axis_angle_to_quat(jp.array([0, 0, 1]), yaw)
+        new_quat = math.quat_mul(
+            qpos[self._floating_base_qpos_addr + 3 : self._floating_base_qpos_addr + 7],
+            quat,
+        )  # yaw noise
+
+        base_qpos = base_qpos.at[3:7].set(new_quat)
+
+        qpos = self.set_floating_base_qpos(base_qpos, qpos)
+        # print(f'DEBUG1 base qpos: {qpos}')
+        # init joint position
+        # qpos[7:]=*U(0.0, 0.1)
+        rng, key = jax.random.split(rng)
+
+        # multiply actual joints with noise (excluding floating base and backlash)
+        qpos_j = self.get_actuator_joints_qpos(qpos) * jax.random.uniform(
+            key, (self._actuators,), minval=0.5, maxval=1.5
+        )
+        qpos = self.set_actuator_joints_qpos(qpos_j, qpos)
+        # print(f'DEBUG2 joint qpos: {qpos}')
+        # init joint vel
+        # d(xyzrpy)=U(-0.05, 0.05)
+        rng, key = jax.random.split(rng)
+        # qvel = qvel.at[self._floating_base_qvel_addr : self._floating_base_qvel_addr + 6].set(
+        #     jax.random.uniform(key, (6,), minval=-0.5, maxval=0.5)
+        # )
+
+        qvel = self.set_floating_base_qvel(
+            jax.random.uniform(key, (6,), minval=-0.05, maxval=0.05), qvel
+        )
+        # print(f'DEBUG3 base qvel: {qvel}')
+        ctrl = self.get_actuator_joints_qpos(qpos)
+        # print(f'DEBUG4 ctrl: {ctrl}')
+        data = mjx_env.init(self.mjx_model, qpos=qpos, qvel=qvel, ctrl=ctrl)
+        rng, cmd_rng = jax.random.split(rng)
+        cmd = self.sample_command(cmd_rng)
+
+        # Sample push interval.
+        rng, push_rng = jax.random.split(rng)
+        push_interval = jax.random.uniform(
+            push_rng,
+            minval=self._config.push_config.interval_range[0],
+            maxval=self._config.push_config.interval_range[1],
+        )
+        push_interval_steps = jp.round(push_interval / self.dt).astype(jp.int32)
+
+        if USE_IMITATION_REWARD:
+            current_reference_motion = self.PRM.get_reference_motion(
+                cmd[0], cmd[1], cmd[2], 0
+            )
+        else:
+            current_reference_motion = jp.zeros(0)
+
+        info = {
+            "rng": rng,
+            "step": 0,
+            "command": cmd,
+            "last_act": jp.zeros(self.mjx_model.nu),
+            "last_last_act": jp.zeros(self.mjx_model.nu),
+            "last_last_last_act": jp.zeros(self.mjx_model.nu),
+            "motor_targets": self._default_actuator,
+            "feet_air_time": jp.zeros(2),
+            "last_contact": jp.zeros(2, dtype=bool),
+            "swing_peak": jp.zeros(2),
+            # Push related.
+            "push": jp.array([0.0, 0.0]),
+            "push_step": 0,
+            "push_interval_steps": push_interval_steps,
+            # History related.
+            "action_history": jp.zeros(
+                self._config.noise_config.action_max_delay * self._actuators
+            ),
+            "imu_history": jp.zeros(self._config.noise_config.imu_max_delay * 3),
+            # imitation related
+            "imitation_i": 0,
+            "current_reference_motion": current_reference_motion,
+            "imitation_phase": jp.zeros(2),
+        }
+
+        metrics = {}
+        for k, v in self._config.reward_config.scales.items():
+            if v != 0:
+                if v > 0:
+                    metrics[f"reward/{k}"] = jp.zeros(())
+                else:
+                    metrics[f"cost/{k}"] = jp.zeros(())
+        metrics["swing_peak"] = jp.zeros(())
+
+        contact = jp.array(
+            [
+                geoms_colliding(data, geom_id, self._floor_geom_id)
+                for geom_id in self._feet_geom_id
+            ]
+        )
+        obs = self._get_obs(data, info, contact)
+        reward, done = jp.zeros(2)
+        return mjx_env.State(data, obs, reward, done, metrics, info)
+
+    def step(self, state: mjx_env.State, action: jax.Array) -> mjx_env.State:
+
+        if USE_IMITATION_REWARD:
+            state.info["imitation_i"] += 1
+            state.info["imitation_i"] = (
+                state.info["imitation_i"] % self.PRM.nb_steps_in_period
+            )  # not critical, is already moduloed in get_reference_motion
+            state.info["imitation_phase"] = jp.array(
+                [
+                    jp.cos(
+                        (state.info["imitation_i"] / self.PRM.nb_steps_in_period)
+                        * 2
+                        * jp.pi
+                    ),
+                    jp.sin(
+                        (state.info["imitation_i"] / self.PRM.nb_steps_in_period)
+                        * 2
+                        * jp.pi
+                    ),
+                ]
+            )
+        else:
+            state.info["imitation_i"] = 0
+
+        if USE_IMITATION_REWARD:
+            state.info["current_reference_motion"] = self.PRM.get_reference_motion(
+                state.info["command"][0],
+                state.info["command"][1],
+                state.info["command"][2],
+                state.info["imitation_i"],
+            )
+        else:
+            state.info["current_reference_motion"] = jp.zeros(0)
+
+        state.info["rng"], push1_rng, push2_rng, action_delay_rng = jax.random.split(
+            state.info["rng"], 4
+        )
+
+        # Handle action delay
+        action_history = (
+            jp.roll(state.info["action_history"], self._actuators)
+            .at[: self._actuators]
+            .set(action)
+        )
+        state.info["action_history"] = action_history
+        action_idx = jax.random.randint(
+            action_delay_rng,
+            (1,),
+            minval=self._config.noise_config.action_min_delay,
+            maxval=self._config.noise_config.action_max_delay,
+        )
+        action_w_delay = action_history.reshape((-1, self._actuators))[
+            action_idx[0]
+        ]  # action with delay
+
+        # self.action_filter.push(action_w_delay)
+        # action_w_delay = self.action_filter.get_filtered_action()
+
+        push_theta = jax.random.uniform(push1_rng, maxval=2 * jp.pi)
+        push_magnitude = jax.random.uniform(
+            push2_rng,
+            minval=self._config.push_config.magnitude_range[0],
+            maxval=self._config.push_config.magnitude_range[1],
+        )
+        push = jp.array([jp.cos(push_theta), jp.sin(push_theta)])
+        push *= (
+            jp.mod(state.info["push_step"] + 1, state.info["push_interval_steps"]) == 0
+        )
+        push *= self._config.push_config.enable
+        qvel = state.data.qvel
+        qvel = qvel.at[
+            self._floating_base_qvel_addr : self._floating_base_qvel_addr + 2
+        ].set(
+            push * push_magnitude
+            + qvel[self._floating_base_qvel_addr : self._floating_base_qvel_addr + 2]
+        )  # floating base x,y
+        data = state.data.replace(qvel=qvel)
+        state = state.replace(data=data)
+
+        ####
+
+        motor_targets = (
+            self._default_actuator + action_w_delay * self._config.action_scale
+        )
+
+        if USE_MOTOR_SPEED_LIMITS:
+            prev_motor_targets = state.info["motor_targets"]
+
+            motor_targets = jp.clip(
+                motor_targets,
+                prev_motor_targets
+                - self._config.max_motor_velocity * self.dt,  # control dt
+                prev_motor_targets
+                + self._config.max_motor_velocity * self.dt,  # control dt
+            )
+
+        # motor_targets.at[5:9].set(state.info["command"][3:])  # head joints
+        data = mjx_env.step(self.mjx_model, state.data, motor_targets, self.n_substeps)
+
+        state.info["motor_targets"] = motor_targets
+
+        contact = jp.array(
+            [
+                geoms_colliding(data, geom_id, self._floor_geom_id)
+                for geom_id in self._feet_geom_id
+            ]
+        )
+        contact_filt = contact | state.info["last_contact"]
+        first_contact = (state.info["feet_air_time"] > 0.0) * contact_filt
+        state.info["feet_air_time"] += self.dt
+        p_f = data.site_xpos[self._feet_site_id]
+        p_fz = p_f[..., -1]
+        state.info["swing_peak"] = jp.maximum(state.info["swing_peak"], p_fz)
+
+        obs = self._get_obs(data, state.info, contact)
+        done = self._get_termination(data)
+
+        rewards = self._get_reward(
+            data, action, state.info, state.metrics, done, first_contact, contact
+        )
+        # FIXME
+        rewards = {
+            k: v * self._config.reward_config.scales[k] for k, v in rewards.items()
+        }
+        reward = jp.clip(sum(rewards.values()) * self.dt, 0.0, 10000.0)
+        # jax.debug.print('STEP REWARD: {}',reward)
+        state.info["push"] = push
+        state.info["step"] += 1
+        state.info["push_step"] += 1
+        state.info["last_last_last_act"] = state.info["last_last_act"]
+        state.info["last_last_act"] = state.info["last_act"]
+        state.info["last_act"] = action  # was
+        # state.info["last_act"] = motor_targets  # became
+        state.info["rng"], cmd_rng = jax.random.split(state.info["rng"])
+        state.info["command"] = jp.where(
+            state.info["step"] > 500,
+            self.sample_command(cmd_rng),
+            state.info["command"],
+        )
+        state.info["step"] = jp.where(
+            done | (state.info["step"] > 500),
+            0,
+            state.info["step"],
+        )
+        state.info["feet_air_time"] *= ~contact
+        state.info["last_contact"] = contact
+        state.info["swing_peak"] *= ~contact
+        for k, v in rewards.items():
+            rew_scale = self._config.reward_config.scales[k]
+            if rew_scale != 0:
+                if rew_scale > 0:
+                    state.metrics[f"reward/{k}"] = v
+                else:
+                    state.metrics[f"cost/{k}"] = -v
+        state.metrics["swing_peak"] = jp.mean(state.info["swing_peak"])
+
+        done = done.astype(reward.dtype)
+        state = state.replace(data=data, obs=obs, reward=reward, done=done)
+        return state
+
+    def _get_termination(self, data: mjx.Data) -> jax.Array:
+        fall_termination = self.get_gravity(data)[-1] < 0.0
+        return fall_termination | jp.isnan(data.qpos).any() | jp.isnan(data.qvel).any()
+
+    def _get_obs(
+        self, data: mjx.Data, info: dict[str, Any], contact: jax.Array
+    ) -> mjx_env.Observation:
+
+        gyro = self.get_gyro(data)
+        info["rng"], noise_rng = jax.random.split(info["rng"])
+        noisy_gyro = (
+            gyro
+            + (2 * jax.random.uniform(noise_rng, shape=gyro.shape) - 1)
+            * self._config.noise_config.level
+            * self._config.noise_config.scales.gyro
+        )
+
+        accelerometer = self.get_accelerometer(data)
+        # accelerometer[0] += 1.3 # TODO testing
+        accelerometer.at[0].set(accelerometer[0] + 1.3)
+
+        info["rng"], noise_rng = jax.random.split(info["rng"])
+        noisy_accelerometer = (
+            accelerometer
+            + (2 * jax.random.uniform(noise_rng, shape=accelerometer.shape) - 1)
+            * self._config.noise_config.level
+            * self._config.noise_config.scales.accelerometer
+        )
+
+        gravity = data.site_xmat[self._site_id].T @ jp.array([0, 0, -1])
+        info["rng"], noise_rng = jax.random.split(info["rng"])
+        noisy_gravity = (
+            gravity
+            + (2 * jax.random.uniform(noise_rng, shape=gravity.shape) - 1)
+            * self._config.noise_config.level
+            * self._config.noise_config.scales.gravity
+        )
+
+        # Handle IMU delay
+        imu_history = jp.roll(info["imu_history"], 3).at[:3].set(noisy_gravity)
+        info["imu_history"] = imu_history
+        imu_idx = jax.random.randint(
+            noise_rng,
+            (1,),
+            minval=self._config.noise_config.imu_min_delay,
+            maxval=self._config.noise_config.imu_max_delay,
+        )
+        noisy_gravity = imu_history.reshape((-1, 3))[imu_idx[0]]
+
+        # joint_angles = data.qpos[7:]
+
+        # Handling backlash
+        joint_angles = self.get_actuator_joints_qpos(data.qpos)
+        joint_backlash = self.get_actuator_backlash_qpos(data.qpos)
+
+        for i in self.backlash_idx_to_add:
+            joint_backlash = jp.insert(joint_backlash, i, 0)
+
+        joint_angles = joint_angles + joint_backlash
+
+        info["rng"], noise_rng = jax.random.split(info["rng"])
+        noisy_joint_angles = (
+            joint_angles
+            + (2.0 * jax.random.uniform(noise_rng, shape=joint_angles.shape) - 1.0)
+            * self._config.noise_config.level
+            * self._qpos_noise_scale
+        )
+
+        # joint_vel = data.qvel[6:]
+        joint_vel = self.get_actuator_joints_qvel(data.qvel)
+        info["rng"], noise_rng = jax.random.split(info["rng"])
+        noisy_joint_vel = (
+            joint_vel
+            + (2.0 * jax.random.uniform(noise_rng, shape=joint_vel.shape) - 1.0)
+            * self._config.noise_config.level
+            * self._config.noise_config.scales.joint_vel
+        )
+
+        linvel = self.get_local_linvel(data)
+        # info["rng"], noise_rng = jax.random.split(info["rng"])
+        # noisy_linvel = (
+        #     linvel
+        #     + (2 * jax.random.uniform(noise_rng, shape=linvel.shape) - 1)
+        #     * self._config.noise_config.level
+        #     * self._config.noise_config.scales.linvel
+        # )
+
+        state = jp.hstack(
+            [
+                # noisy_linvel,  # 3
+                # noisy_gyro,  # 3
+                # noisy_gravity,  # 3
+                noisy_gyro,  # 3
+                noisy_accelerometer,  # 3
+                info["command"],  # 3
+                noisy_joint_angles - self._default_actuator,  # 10
+                noisy_joint_vel * self._config.dof_vel_scale,  # 10
+                info["last_act"],  # 10
+                info["last_last_act"],  # 10
+                info["last_last_last_act"],  # 10
+                info["motor_targets"],  # 10
+                contact,  # 2
+                # info["current_reference_motion"],
+                # info["imitation_i"],
+                info["imitation_phase"],
+            ]
+        )
+
+        accelerometer = self.get_accelerometer(data)
+        global_angvel = self.get_global_angvel(data)
+        feet_vel = data.sensordata[self._foot_linvel_sensor_adr].ravel()
+        root_height = data.qpos[self._floating_base_qpos_addr + 2]
+
+        privileged_state = jp.hstack(
+            [
+                state,
+                gyro,  # 3
+                accelerometer,  # 3
+                gravity,  # 3
+                linvel,  # 3
+                global_angvel,  # 3
+                joint_angles - self._default_actuator,
+                joint_vel,
+                root_height,  # 1
+                data.actuator_force,  # 10
+                contact,  # 2
+                feet_vel,  # 4*3
+                info["feet_air_time"],  # 2
+                info["current_reference_motion"],
+                info["imitation_i"],
+                info["imitation_phase"],
+            ]
+        )
+
+        return {
+            "state": state,
+            "privileged_state": privileged_state,
+        }
+
+    def _get_reward(
+        self,
+        data: mjx.Data,
+        action: jax.Array,
+        info: dict[str, Any],
+        metrics: dict[str, Any],
+        done: jax.Array,
+        first_contact: jax.Array,
+        contact: jax.Array,
+    ) -> dict[str, jax.Array]:
+        del metrics  # Unused.
+
+        ret = {
+            "tracking_lin_vel": reward_tracking_lin_vel(
+                info["command"],
+                self.get_local_linvel(data),
+                self._config.reward_config.tracking_sigma,
+            ),
+            "tracking_ang_vel": reward_tracking_ang_vel(
+                info["command"],
+                self.get_gyro(data),
+                self._config.reward_config.tracking_sigma,
+            ),
+            # "orientation": cost_orientation(self.get_gravity(data)),
+            "torques": cost_torques(data.actuator_force),
+            "action_rate": cost_action_rate(action, info["last_act"]),
+            "alive": reward_alive(),
+            "imitation": reward_imitation(  # FIXME, this reward is so adhoc...
+                self.get_floating_base_qpos(data.qpos),  # floating base qpos
+                self.get_floating_base_qvel(data.qvel),  # floating base qvel
+                self.get_actuator_joints_qpos(data.qpos),
+                self.get_actuator_joints_qvel(data.qvel),
+                contact,
+                info["current_reference_motion"],
+                info["command"],
+                USE_IMITATION_REWARD,
+            ),
+            "stand_still": cost_stand_still(
+                # info["command"], data.qpos[7:], data.qvel[6:], self._default_pose
+                info["command"],
+                self.get_actuator_joints_qpos(data.qpos),
+                self.get_actuator_joints_qvel(data.qvel),
+                self._default_actuator,
+                ignore_head=False,
+            ),
+        }
+
+        return ret
+
+    def sample_command(self, rng: jax.Array) -> jax.Array:
+        rng1, rng2, rng3, rng4, rng5, rng6, rng7, rng8 = jax.random.split(rng, 8)
+
+        lin_vel_x = jax.random.uniform(
+            rng1, minval=self._config.lin_vel_x[0], maxval=self._config.lin_vel_x[1]
+        )
+        lin_vel_y = jax.random.uniform(
+            rng2, minval=self._config.lin_vel_y[0], maxval=self._config.lin_vel_y[1]
+        )
+        ang_vel_yaw = jax.random.uniform(
+            rng3,
+            minval=self._config.ang_vel_yaw[0],
+            maxval=self._config.ang_vel_yaw[1],
+        )
+
+        neck_pitch = jax.random.uniform(
+            rng5,
+            minval=self._config.neck_pitch_range[0] * self._config.head_range_factor,
+            maxval=self._config.neck_pitch_range[1] * self._config.head_range_factor,
+        )
+
+        head_pitch = jax.random.uniform(
+            rng6,
+            minval=self._config.head_pitch_range[0] * self._config.head_range_factor,
+            maxval=self._config.head_pitch_range[1] * self._config.head_range_factor,
+        )
+
+        head_yaw = jax.random.uniform(
+            rng7,
+            minval=self._config.head_yaw_range[0] * self._config.head_range_factor,
+            maxval=self._config.head_yaw_range[1] * self._config.head_range_factor,
+        )
+
+        head_roll = jax.random.uniform(
+            rng8,
+            minval=self._config.head_roll_range[0] * self._config.head_range_factor,
+            maxval=self._config.head_roll_range[1] * self._config.head_range_factor,
+        )
+
+        # With 10% chance, set everything to zero.
+        return jp.where(
+            jax.random.bernoulli(rng4, p=0.1),
+            jp.zeros(7),
+            jp.hstack(
+                [
+                    lin_vel_x,
+                    lin_vel_y,
+                    ang_vel_yaw,
+                    neck_pitch,
+                    head_pitch,
+                    head_yaw,
+                    head_roll,
+                ]
+            ),
+        )

Open_Duck_Playground/playground/open_duck_mini_v2/mujoco_infer.py

@@ -0,0 +1,266 @@
+import mujoco
+import pickle
+import numpy as np
+import mujoco
+import mujoco.viewer
+import time
+import argparse
+from playground.common.onnx_infer import OnnxInfer
+from playground.common.poly_reference_motion_numpy import PolyReferenceMotion
+from playground.common.utils import LowPassActionFilter
+
+from playground.open_duck_mini_v2.mujoco_infer_base import MJInferBase
+
+USE_MOTOR_SPEED_LIMITS = True
+
+
+class MjInfer(MJInferBase):
+    def __init__(
+        self, model_path: str, reference_data: str, onnx_model_path: str, standing: bool
+    ):
+        super().__init__(model_path)
+
+        self.standing = standing
+        self.head_control_mode = self.standing
+
+        # Params
+        self.linearVelocityScale = 1.0
+        self.angularVelocityScale = 1.0
+        self.dof_pos_scale = 1.0
+        self.dof_vel_scale = 0.05
+        self.action_scale = 0.25
+
+        self.action_filter = LowPassActionFilter(50, cutoff_frequency=37.5)
+
+        if not self.standing:
+            self.PRM = PolyReferenceMotion(reference_data)
+
+        self.policy = OnnxInfer(onnx_model_path, awd=True)
+
+        self.COMMANDS_RANGE_X = [-0.15, 0.15]
+        self.COMMANDS_RANGE_Y = [-0.2, 0.2]
+        self.COMMANDS_RANGE_THETA = [-1.0, 1.0]  # [-1.0, 1.0]
+
+        self.NECK_PITCH_RANGE = [-0.34, 1.1]
+        self.HEAD_PITCH_RANGE = [-0.78, 0.78]
+        self.HEAD_YAW_RANGE = [-1.5, 1.5]
+        self.HEAD_ROLL_RANGE = [-0.5, 0.5]
+
+        self.last_action = np.zeros(self.num_dofs)
+        self.last_last_action = np.zeros(self.num_dofs)
+        self.last_last_last_action = np.zeros(self.num_dofs)
+        self.commands = [0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 0.0]
+
+        self.imitation_i = 0
+        self.imitation_phase = np.array([0, 0])
+        self.saved_obs = []
+
+        self.max_motor_velocity = 5.24  # rad/s
+
+        self.phase_frequency_factor = 1.0
+
+        print(f"joint names: {self.joint_names}")
+        print(f"actuator names: {self.actuator_names}")
+        print(f"backlash joint names: {self.backlash_joint_names}")
+        # print(f"actual joints idx: {self.get_actual_joints_idx()}")
+
+    def get_obs(
+        self,
+        data,
+        command,  # , qvel_history, qpos_error_history, gravity_history
+    ):
+        gyro = self.get_gyro(data)
+        accelerometer = self.get_accelerometer(data)
+        accelerometer[0] += 1.3
+
+        joint_angles = self.get_actuator_joints_qpos(data.qpos)
+        joint_vel = self.get_actuator_joints_qvel(data.qvel)
+
+        contacts = self.get_feet_contacts(data)
+
+        # if not self.standing:
+        # ref = self.PRM.get_reference_motion(*command[:3], self.imitation_i)
+
+        obs = np.concatenate(
+            [
+                gyro,
+                accelerometer,
+                # gravity,
+                command,
+                joint_angles - self.default_actuator,
+                joint_vel * self.dof_vel_scale,
+                self.last_action,
+                self.last_last_action,
+                self.last_last_last_action,
+                self.motor_targets,
+                contacts,
+                # ref if not self.standing else np.array([]),
+                # [self.imitation_i]
+                self.imitation_phase,
+            ]
+        )
+
+        return obs
+
+    def key_callback(self, keycode):
+        print(f"key: {keycode}")
+        if keycode == 72:  # h
+            self.head_control_mode = not self.head_control_mode
+        lin_vel_x = 0
+        lin_vel_y = 0
+        ang_vel = 0
+        if not self.head_control_mode:
+            if keycode == 265:  # arrow up
+                lin_vel_x = self.COMMANDS_RANGE_X[1]
+            if keycode == 264:  # arrow down
+                lin_vel_x = self.COMMANDS_RANGE_X[0]
+            if keycode == 263:  # arrow left
+                lin_vel_y = self.COMMANDS_RANGE_Y[1]
+            if keycode == 262:  # arrow right
+                lin_vel_y = self.COMMANDS_RANGE_Y[0]
+            if keycode == 81:  # a
+                ang_vel = self.COMMANDS_RANGE_THETA[1]
+            if keycode == 69:  # e
+                ang_vel = self.COMMANDS_RANGE_THETA[0]
+            if keycode == 80:  # p
+                self.phase_frequency_factor += 0.1
+            if keycode == 59:  # m
+                self.phase_frequency_factor -= 0.1
+        else:
+            neck_pitch = 0
+            head_pitch = 0
+            head_yaw = 0
+            head_roll = 0
+            if keycode == 265:  # arrow up
+                head_pitch = self.NECK_PITCH_RANGE[1]
+            if keycode == 264:  # arrow down
+                head_pitch = self.NECK_PITCH_RANGE[0]
+            if keycode == 263:  # arrow left
+                head_yaw = self.HEAD_YAW_RANGE[1]
+            if keycode == 262:  # arrow right
+                head_yaw = self.HEAD_YAW_RANGE[0]
+            if keycode == 81:  # a
+                head_roll = self.HEAD_ROLL_RANGE[1]
+            if keycode == 69:  # e
+                head_roll = self.HEAD_ROLL_RANGE[0]
+
+            self.commands[3] = neck_pitch
+            self.commands[4] = head_pitch
+            self.commands[5] = head_yaw
+            self.commands[6] = head_roll
+
+        self.commands[0] = lin_vel_x
+        self.commands[1] = lin_vel_y
+        self.commands[2] = ang_vel
+
+    def run(self):
+        try:
+            with mujoco.viewer.launch_passive(
+                self.model,
+                self.data,
+                show_left_ui=False,
+                show_right_ui=False,
+                key_callback=self.key_callback,
+            ) as viewer:
+                counter = 0
+                while True:
+
+                    step_start = time.time()
+
+                    mujoco.mj_step(self.model, self.data)
+
+                    counter += 1
+
+                    if counter % self.decimation == 0:
+                        if not self.standing:
+                            self.imitation_i += 1.0 * self.phase_frequency_factor
+                            self.imitation_i = (
+                                self.imitation_i % self.PRM.nb_steps_in_period
+                            )
+                            # print(self.PRM.nb_steps_in_period)
+                            # exit()
+                            self.imitation_phase = np.array(
+                                [
+                                    np.cos(
+                                        self.imitation_i
+                                        / self.PRM.nb_steps_in_period
+                                        * 2
+                                        * np.pi
+                                    ),
+                                    np.sin(
+                                        self.imitation_i
+                                        / self.PRM.nb_steps_in_period
+                                        * 2
+                                        * np.pi
+                                    ),
+                                ]
+                            )
+                        obs = self.get_obs(
+                            self.data,
+                            self.commands,
+                        )
+                        self.saved_obs.append(obs)
+                        action = self.policy.infer(obs)
+
+                        # self.action_filter.push(action)
+                        # action = self.action_filter.get_filtered_action()
+
+                        self.last_last_last_action = self.last_last_action.copy()
+                        self.last_last_action = self.last_action.copy()
+                        self.last_action = action.copy()
+
+                        self.motor_targets = (
+                            self.default_actuator + action * self.action_scale
+                        )
+
+                        if USE_MOTOR_SPEED_LIMITS:
+                            self.motor_targets = np.clip(
+                                self.motor_targets,
+                                self.prev_motor_targets
+                                - self.max_motor_velocity
+                                * (self.sim_dt * self.decimation),
+                                self.prev_motor_targets
+                                + self.max_motor_velocity
+                                * (self.sim_dt * self.decimation),
+                            )
+
+                            self.prev_motor_targets = self.motor_targets.copy()
+
+                        # head_targets = self.commands[3:]
+                        # self.motor_targets[5:9] = head_targets
+                        self.data.ctrl = self.motor_targets.copy()
+
+                    viewer.sync()
+
+                    time_until_next_step = self.model.opt.timestep - (
+                        time.time() - step_start
+                    )
+                    if time_until_next_step > 0:
+                        time.sleep(time_until_next_step)
+        except KeyboardInterrupt:
+            pickle.dump(self.saved_obs, open("mujoco_saved_obs.pkl", "wb"))
+
+
+if __name__ == "__main__":
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("-o", "--onnx_model_path", type=str, required=True)
+    # parser.add_argument("-k", action="store_true", default=False)
+    parser.add_argument(
+        "--reference_data",
+        type=str,
+        default="playground/open_duck_mini_v2/data/polynomial_coefficients.pkl",
+    )
+    parser.add_argument(
+        "--model_path",
+        type=str,
+        default="playground/open_duck_mini_v2/xmls/scene_flat_terrain.xml",
+    )
+    parser.add_argument("--standing", action="store_true", default=False)
+
+    args = parser.parse_args()
+
+    mjinfer = MjInfer(
+        args.model_path, args.reference_data, args.onnx_model_path, args.standing
+    )
+    mjinfer.run()

Open_Duck_Playground/playground/open_duck_mini_v2/mujoco_infer_base.py

@@ -0,0 +1,283 @@
+import mujoco
+import numpy as np
+from etils import epath
+from playground.open_duck_mini_v2 import base
+
+
+class MJInferBase:
+    def __init__(self, model_path):
+
+        self.model = mujoco.MjModel.from_xml_string(
+            epath.Path(model_path).read_text(), assets=base.get_assets()
+        )
+        print(model_path)
+
+        self.sim_dt = 0.002
+        self.decimation = 10
+        self.model.opt.timestep = self.sim_dt
+        self.data = mujoco.MjData(self.model)
+        mujoco.mj_step(self.model, self.data)
+
+        self.num_dofs = self.model.nu
+        self.floating_base_name = [
+            self.model.jnt(k).name
+            for k in range(0, self.model.njnt)
+            if self.model.jnt(k).type == 0
+        ][
+            0
+        ]  # assuming only one floating object!
+        self.actuator_names = [
+            self.model.actuator(k).name for k in range(0, self.model.nu)
+        ]  # will be useful to get only the actuators we care about
+        self.joint_names = [  # njnt = all joints (including floating base, actuators and backlash joints)
+            self.model.jnt(k).name for k in range(0, self.model.njnt)
+        ]  # all the joint (including the backlash joints)
+        self.backlash_joint_names = [
+            j
+            for j in self.joint_names
+            if j not in self.actuator_names and j not in self.floating_base_name
+        ]  # only the dummy backlash joint
+        self.all_joint_ids = [self.get_joint_id_from_name(n) for n in self.joint_names]
+        self.all_joint_qpos_addr = [
+            self.get_joint_addr_from_name(n) for n in self.joint_names
+        ]
+
+        self.actuator_joint_ids = [
+            self.get_joint_id_from_name(n) for n in self.actuator_names
+        ]
+        self.actuator_joint_qpos_addr = [
+            self.get_joint_addr_from_name(n) for n in self.actuator_names
+        ]
+
+        self.backlash_joint_ids = [
+            self.get_joint_id_from_name(n) for n in self.backlash_joint_names
+        ]
+
+        self.backlash_joint_qpos_addr = [
+            self.get_joint_addr_from_name(n) for n in self.backlash_joint_names
+        ]
+
+        self.all_qvel_addr = np.array(
+            [self.model.jnt_dofadr[jad] for jad in self.all_joint_ids]
+        )
+        self.actuator_qvel_addr = np.array(
+            [self.model.jnt_dofadr[jad] for jad in self.actuator_joint_ids]
+        )
+
+        self.actuator_joint_dict = {
+            n: self.get_joint_id_from_name(n) for n in self.actuator_names
+        }
+
+        self._floating_base_qpos_addr = self.model.jnt_qposadr[
+            np.where(self.model.jnt_type == 0)
+        ][
+            0
+        ]  # Assuming there is only one floating base! the jnt_type==0 is a floating joint. 3 is a hinge
+
+        self._floating_base_qvel_addr = self.model.jnt_dofadr[
+            np.where(self.model.jnt_type == 0)
+        ][
+            0
+        ]  # Assuming there is only one floating base! the jnt_type==0 is a floating joint. 3 is a hinge
+
+        self._floating_base_id = self.model.joint(self.floating_base_name).id
+
+        # self.all_joint_no_backlash_ids=np.zeros(7+self.model.nu)
+        all_idx = self.backlash_joint_ids + list(
+            range(self._floating_base_qpos_addr, self._floating_base_qpos_addr + 7)
+        )
+        all_idx.sort()
+
+        # self.all_joint_no_backlash_ids=[idx for idx in self.all_joint_ids if idx not in self.backlash_joint_ids]+list(range(self._floating_base_add,self._floating_base_add+7))
+        self.all_joint_no_backlash_ids = [idx for idx in all_idx]
+
+        self.gyro_id = mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_SENSOR, "gyro")
+        self.gyro_addr = self.model.sensor_adr[self.gyro_id]
+        self.gyro_dimensions = 3
+
+        self.accelerometer_id = mujoco.mj_name2id(
+            self.model, mujoco.mjtObj.mjOBJ_SENSOR, "accelerometer"
+        )
+        self.accelerometer_dimensions = 3
+        self.accelerometer_addr = self.model.sensor_adr[self.accelerometer_id]
+
+        self.linvel_id = mujoco.mj_name2id(
+            self.model, mujoco.mjtObj.mjOBJ_SENSOR, "local_linvel"
+        )
+        self.linvel_dimensions = 3
+
+        self.imu_site_id = mujoco.mj_name2id(
+            self.model, mujoco.mjtObj.mjOBJ_SITE, "imu"
+        )
+
+        self.gravity_id = mujoco.mj_name2id(
+            self.model, mujoco.mjtObj.mjOBJ_SENSOR, "upvector"
+        )
+        self.gravity_dimensions = 3
+
+        self.init_pos = np.array(
+            self.get_all_joints_qpos(self.model.keyframe("home").qpos)
+        )  # pose of all the joints (no floating base)
+        self.default_actuator = self.model.keyframe(
+            "home"
+        ).ctrl  # ctrl of all the actual joints (no floating base and no backlash)
+        self.motor_targets = self.default_actuator
+        self.prev_motor_targets = self.default_actuator
+
+        self.data.qpos[:] = self.model.keyframe("home").qpos
+        self.data.ctrl[:] = self.default_actuator
+
+    def get_actuator_id_from_name(self, name: str) -> int:
+        """Return the id of a specified actuator"""
+        return mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_ACTUATOR, name)
+
+    def get_joint_id_from_name(self, name: str) -> int:
+        """Return the id of a specified joint"""
+        return mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_JOINT, name)
+
+    def get_joint_addr_from_name(self, name: str) -> int:
+        """Return the address of a specified joint"""
+        return self.model.joint(name).qposadr
+
+    def get_dof_id_from_name(self, name: str) -> int:
+        """Return the id of a specified dof"""
+        return mujoco.mj_name2id(self.model, mujoco.mjtObj.mjOBJ_DOF, name)
+
+    def get_actuator_joint_qpos_from_name(
+        self, data: np.ndarray, name: str
+    ) -> np.ndarray:
+        """Return the qpos of a given actual joint"""
+        addr = self.model.jnt_qposadr[self.actuator_joint_dict[name]]
+        return data[addr]
+
+    def get_actuator_joints_addr(self) -> np.ndarray:
+        """Return the all the idx of actual joints"""
+        addr = np.array(
+            [self.model.jnt_qposadr[idx] for idx in self.actuator_joint_ids]
+        )
+        return addr
+
+    def get_floating_base_qpos(self, data: np.ndarray) -> np.ndarray:
+        return data[self._floating_base_qpos_addr : self._floating_base_qvel_addr + 7]
+
+    def get_floating_base_qvel(self, data: np.ndarray) -> np.ndarray:
+        return data[self._floating_base_qvel_addr : self._floating_base_qvel_addr + 6]
+
+    def set_floating_base_qpos(
+        self, new_qpos: np.ndarray, qpos: np.ndarray
+    ) -> np.ndarray:
+        qpos[self._floating_base_qpos_addr : self._floating_base_qpos_addr + 7] = (
+            new_qpos
+        )
+        return qpos
+
+    def set_floating_base_qvel(
+        self, new_qvel: np.ndarray, qvel: np.ndarray
+    ) -> np.ndarray:
+        qvel[self._floating_base_qvel_addr : self._floating_base_qvel_addr + 6] = (
+            new_qvel
+        )
+        return qvel
+
+    def exclude_backlash_joints_addr(self) -> np.ndarray:
+        """Return the all the idx of actual joints and floating base"""
+        addr = np.array(
+            [self.model.jnt_qposadr[idx] for idx in self.all_joint_no_backlash_ids]
+        )
+        return addr
+
+    def get_all_joints_addr(self) -> np.ndarray:
+        """Return the all the idx of all joints"""
+        addr = np.array([self.model.jnt_qposadr[idx] for idx in self.all_joint_ids])
+        return addr
+
+    def get_actuator_joints_qpos(self, data: np.ndarray) -> np.ndarray:
+        """Return the all the qpos of actual joints"""
+        return data[self.get_actuator_joints_addr()]
+
+    def set_actuator_joints_qpos(
+        self, new_qpos: np.ndarray, qpos: np.ndarray
+    ) -> np.ndarray:
+        """Set the qpos only for the actual joints (omit the backlash joint)"""
+        qpos[self.get_actuator_joints_addr()] = new_qpos
+        return qpos
+
+    def get_actuator_joints_qvel(self, data: np.ndarray) -> np.ndarray:
+        """Return the all the qvel of actual joints"""
+        return data[self.actuator_qvel_addr]
+
+    def set_actuator_joints_qvel(
+        self, new_qvel: np.ndarray, qvel: np.ndarray
+    ) -> np.ndarray:
+        """Set the qvel only for the actual joints (omit the backlash joint)"""
+        qvel[self.actuator_qvel_addr] = new_qvel
+        return qvel
+
+    def get_all_joints_qpos(self, data: np.ndarray) -> np.ndarray:
+        """Return the all the qpos of all joints"""
+        return data[self.get_all_joints_addr()]
+
+    def get_all_joints_qvel(self, data: np.ndarray) -> np.ndarray:
+        """Return the all the qvel of all joints"""
+        return data[self.all_qvel_addr]
+
+    def get_joints_nobacklash_qpos(self, data: np.ndarray) -> np.ndarray:
+        """Return the all the qpos of actual joints with the floating base"""
+        return data[self.exclude_backlash_joints_addr()]
+
+    def set_complete_qpos_from_joints(
+        self, no_backlash_qpos: np.ndarray, full_qpos: np.ndarray
+    ) -> np.ndarray:
+        """In the case of backlash joints, we want to ignore them (remove them) but we still need to set the complete state incuding them"""
+        full_qpos[self.exclude_backlash_joints_addr()] = no_backlash_qpos
+        return np.array(full_qpos)
+
+    def get_sensor(self, data, name, dimensions):
+        i = self.model.sensor_name2id(name)
+        return data.sensordata[i : i + dimensions]
+
+    def get_gyro(self, data):
+        return data.sensordata[self.gyro_addr : self.gyro_addr + self.gyro_dimensions]
+
+    def get_accelerometer(self, data):
+        return data.sensordata[
+            self.accelerometer_addr : self.accelerometer_addr
+            + self.accelerometer_dimensions
+        ]
+
+    def get_linvel(self, data):
+        return data.sensordata[self.linvel_id : self.linvel_id + self.linvel_dimensions]
+
+    # def get_gravity(self, data):
+    #     return data.site_xmat[self.imu_site_id].reshape((3, 3)).T @ np.array([0, 0, -1])
+
+    def get_gravity(self, data):
+        return data.sensordata[
+            self.gravity_id : self.gravity_id + self.gravity_dimensions
+        ]
+
+    def check_contact(self, data, body1_name, body2_name):
+        body1_id = data.body(body1_name).id
+        body2_id = data.body(body2_name).id
+
+        for i in range(data.ncon):
+            try:
+                contact = data.contact[i]
+            except Exception as e:
+                return False
+
+            if (
+                self.model.geom_bodyid[contact.geom1] == body1_id
+                and self.model.geom_bodyid[contact.geom2] == body2_id
+            ) or (
+                self.model.geom_bodyid[contact.geom1] == body2_id
+                and self.model.geom_bodyid[contact.geom2] == body1_id
+            ):
+                return True
+
+        return False
+
+    def get_feet_contacts(self, data):
+        left_contact = self.check_contact(data, "foot_assembly", "floor")
+        right_contact = self.check_contact(data, "foot_assembly_2", "floor")
+        return left_contact, right_contact

Open_Duck_Playground/playground/open_duck_mini_v2/ref_motion_viewer.py

@@ -0,0 +1,214 @@
+import mujoco
+import numpy as np
+import time
+import argparse
+import os, sys
+import pygame
+from etils import epath
+import mujoco.viewer
+
+# Import the reference motion class.
+from playground.open_duck_mini_v2 import base
+
+from playground.common.poly_reference_motion_numpy import PolyReferenceMotion
+
+SCRIPT_PATH = os.path.dirname(os.path.abspath(__file__))
+SCENE_PATH = f"{SCRIPT_PATH}/xmls"
+
+COMMANDS_RANGE_X = [-0.15, 0.15]
+COMMANDS_RANGE_Y = [-0.2, 0.2]
+COMMANDS_RANGE_THETA = [-1.0, 1.0]  # [-1.0, 1.0]
+
+available_scenes = []
+if os.path.isdir(SCENE_PATH):
+    for name in os.listdir(SCENE_PATH):
+        if (
+            name.startswith("scene_")
+            and name.endswith(".xml")
+            and os.path.isfile(os.path.join(SCENE_PATH, name))
+        ):
+            scene_name = name[len("scene_") : -len(".xml")]
+            available_scenes.append(scene_name)
+if len(available_scenes) == 0:
+    print(f"No scenes found in: {SCENE_PATH}")
+    sys.exit(1)
+
+# Parse command-line arguments.
+parser = argparse.ArgumentParser(description="Reference Motion Viewer")
+parser.add_argument(
+    "--reference_data",
+    type=str,
+    default="playground/go_bdx/data/polynomial_coefficients.pkl",
+    help="Path to the polynomial coefficients pickle file.",
+)
+parser.add_argument(
+    "-joystick", action="store_true", default=False, help="Use joystick control"
+)
+# Command parameters: dx, dy, dtheta
+parser.add_argument(
+    "--command",
+    nargs=3,
+    type=float,
+    default=[0.0, -0.05, -0.1],
+    help="Reference motion command as three floats: dx, dy, dtheta.",
+)
+parser.add_argument(
+    "--scene",
+    type=str,
+    choices=available_scenes,
+    default="flat_terrain",
+)
+args = parser.parse_args()
+# model_path = f"playground/go_bdx/xmls/scene_mjx_{args.scene}.xml"
+model_path = f"playground/open_duck_mini_v2/xmls/scene_{args.scene}.xml"
+
+command = args.command
+
+joystick1 = None
+joystick2 = None
+if args.joystick:
+    pygame.init()
+    pygame.joystick.init()
+    if pygame.joystick.get_count() > 0:
+        joystick1 = pygame.joystick.Joystick(0)
+        joystick1.init()
+        command = [0.0, 0.0, 0.0]
+        print("Joystick initialized:", joystick1.get_name())
+        if pygame.joystick.get_count() > 1:
+            joystick2 = pygame.joystick.Joystick(1)
+            joystick2.init()
+            print("Joystick 2 (theta) initialized:", joystick2.get_name())
+        else:
+            print(
+                "Only one joystick detected; theta via second joystick will be disabled."
+            )
+    else:
+        print("No joystick found!")
+
+# Load the Mujoco model XML.
+model = mujoco.MjModel.from_xml_string(
+    epath.Path(model_path).read_text(), assets=base.get_assets()
+)
+data = mujoco.MjData(model)
+model.opt.timestep = 0.002
+
+# Step the simulation once to initialize.
+mujoco.mj_step(model, data)
+
+# Load the polynomial reference motion.
+PRM = PolyReferenceMotion(args.reference_data)
+
+# Get the "home" keyframe to use as a default pose.
+home_frame = model.keyframe("home")
+default_qpos = np.array(home_frame.qpos)
+default_ctrl = np.array(home_frame.ctrl)
+default_qpos[2] += 0.2  # Increase the base height by 0.2 meters
+
+# Set initial state.
+data.qpos[:] = default_qpos.copy()
+data.ctrl[:] = default_ctrl.copy()
+
+decimation = 10  # 50 hz
+
+
+def key_callback(keycode):
+    if joystick1 is not None:
+        return
+
+    print(f"key: {keycode}")
+    lin_vel_x = 0
+    lin_vel_y = 0
+    ang_vel = 0
+    if keycode == 265:  # arrow up
+        lin_vel_x = COMMANDS_RANGE_X[1]
+    if keycode == 264:  # arrow down
+        lin_vel_x = COMMANDS_RANGE_X[0]
+    if keycode == 263:  # arrow left
+        lin_vel_y = COMMANDS_RANGE_Y[1]
+    if keycode == 262:  # arrow right
+        lin_vel_y = COMMANDS_RANGE_Y[0]
+    if keycode == 81:  # a
+        ang_vel = COMMANDS_RANGE_THETA[1]
+    if keycode == 69:  # e
+        ang_vel = COMMANDS_RANGE_THETA[0]
+
+    command[0] = lin_vel_x
+    command[1] = lin_vel_y
+    command[2] = ang_vel
+    print(f"command: {command}")
+
+
+def handle_joystick():
+    if joystick1 is None:
+        return
+
+    joy_z = 0
+    pygame.event.pump()
+    joy_y = joystick1.get_axis(1)
+    joy_x = joystick1.get_axis(0)
+    if joystick2 is not None:
+        joy_z = joystick2.get_axis(0)
+    if joy_y < 0:
+        lin_vel_x = (-joy_y) * COMMANDS_RANGE_X[1]
+    else:
+        lin_vel_x = -joy_y * abs(COMMANDS_RANGE_X[0])
+    lin_vel_y = -joy_x * COMMANDS_RANGE_Y[1]
+    lin_vel_z = -joy_z * COMMANDS_RANGE_THETA[1]
+
+    command[0] = lin_vel_x
+    command[1] = lin_vel_y
+    command[2] = lin_vel_z
+    print(f"command: {command}")
+
+
+# Create a Mujoco viewer to display the model.
+with mujoco.viewer.launch_passive(
+    model, data, show_left_ui=False, show_right_ui=False, key_callback=key_callback
+) as viewer:
+    step = 0
+    dt = model.opt.timestep
+    counter = 0
+    new_qpos = default_qpos.copy()
+    while viewer.is_running():
+        step_start = time.time()
+        handle_joystick()
+        counter += 1
+        new_qpos[:7] = default_qpos[:7].copy()
+        if counter % decimation == 0:
+            new_qpos = default_qpos.copy()
+            if not all(val == 0.0 for val in command):
+                imitation_i = step % PRM.nb_steps_in_period
+
+                ref_motion = PRM.get_reference_motion(
+                    command[0], command[1], command[2], imitation_i
+                )
+                ref_motion = np.array(ref_motion)
+
+                if ref_motion.shape[0] == 40:
+                    joints_pos = ref_motion[0:16]
+                    ref_joint_pos = np.concatenate([joints_pos[:9], joints_pos[11:]])
+                else:
+                    print(
+                        "Error: Unexpected reference motion dimension:",
+                        ref_motion.shape,
+                    )
+                    sys.exit(1)
+
+                new_qpos = default_qpos.copy()
+                if new_qpos[7 : 7 + 14].shape[0] == ref_joint_pos.shape[0]:
+                    new_qpos[7 : 7 + 14] = ref_joint_pos
+                else:
+                    print(
+                        "Error: Actuated joint dimension mismatch. Using default pose."
+                    )
+                step += 1
+            else:
+                step = 0
+        data.qpos[:] = new_qpos
+
+        # Step the simulation to update any dependent quantities.
+        mujoco.mj_step(model, data)
+        viewer.sync()
+        time_until_next_step = model.opt.timestep - (time.time() - step_start)
+        if time_until_next_step > 0:
+            time.sleep(time_until_next_step)

Open_Duck_Playground/playground/open_duck_mini_v2/runner.py

@@ -0,0 +1,64 @@
+"""Runs training and evaluation loop for Open Duck Mini V2."""
+
+import argparse
+
+from playground.common import randomize
+from playground.common.runner import BaseRunner
+from playground.open_duck_mini_v2 import joystick, standing
+
+
+class OpenDuckMiniV2Runner(BaseRunner):
+
+    def __init__(self, args):
+        super().__init__(args)
+        available_envs = {
+            "joystick": (joystick, joystick.Joystick),
+            "standing": (standing, standing.Standing),
+        }
+        if args.env not in available_envs:
+            raise ValueError(f"Unknown env {args.env}")
+
+        self.env_file = available_envs[args.env]
+
+        self.env_config = self.env_file[0].default_config()
+        self.env = self.env_file[1](task=args.task)
+        self.eval_env = self.env_file[1](task=args.task)
+        self.randomizer = randomize.domain_randomize
+        self.action_size = self.env.action_size
+        self.obs_size = int(
+            self.env.observation_size["state"][0]
+        )  # 0: state 1: privileged_state
+        self.restore_checkpoint_path = args.restore_checkpoint_path
+        print(f"Observation size: {self.obs_size}")
+
+
+def main() -> None:
+    parser = argparse.ArgumentParser(description="Open Duck Mini Runner Script")
+    parser.add_argument(
+        "--output_dir",
+        type=str,
+        default="checkpoints",
+        help="Where to save the checkpoints",
+    )
+    # parser.add_argument("--num_timesteps", type=int, default=300000000)
+    parser.add_argument("--num_timesteps", type=int, default=150000000)
+    parser.add_argument("--env", type=str, default="joystick", help="env")
+    parser.add_argument("--task", type=str, default="flat_terrain", help="Task to run")
+    parser.add_argument(
+        "--restore_checkpoint_path",
+        type=str,
+        default=None,
+        help="Resume training from this checkpoint",
+    )
+    # parser.add_argument(
+    #     "--debug", action="store_true", help="Run in debug mode with minimal parameters"
+    # )
+    args = parser.parse_args()
+
+    runner = OpenDuckMiniV2Runner(args)
+
+    runner.train()
+
+
+if __name__ == "__main__":
+    main()

Open_Duck_Playground/playground/open_duck_mini_v2/standing.py

@@ -0,0 +1,661 @@
+# Copyright 2025 DeepMind Technologies Limited
+# Copyright 2025 Antoine Pirrone - Steve Nguyen
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+# ==============================================================================
+"""Standing task for Open Duck Mini V2. (based on Berkeley Humanoid)"""
+
+from typing import Any, Dict, Optional, Union
+import jax
+import jax.numpy as jp
+from ml_collections import config_dict
+from mujoco import mjx
+from mujoco.mjx._src import math
+import numpy as np
+
+from mujoco_playground._src import mjx_env
+from mujoco_playground._src.collision import geoms_colliding
+
+from . import constants
+from . import base as open_duck_mini_v2_base
+from playground.common.poly_reference_motion import PolyReferenceMotion
+from playground.common.rewards import (
+    cost_orientation,
+    cost_torques,
+    cost_action_rate,
+    cost_stand_still,
+    reward_alive,
+    cost_head_pos,
+)
+
+# if set to false, won't require the reference data to be present and won't compute the reference motions polynoms for nothing
+USE_IMITATION_REWARD = False
+
+
+def default_config() -> config_dict.ConfigDict:
+    return config_dict.create(
+        ctrl_dt=0.02,
+        sim_dt=0.002,
+        # episode_length=450,
+        episode_length=1000,
+        action_repeat=1,
+        action_scale=0.25,
+        dof_vel_scale=0.05,
+        history_len=0,
+        soft_joint_pos_limit_factor=0.95,
+        noise_config=config_dict.create(
+            level=1.0,  # Set to 0.0 to disable noise.
+            action_min_delay=0,  # env steps
+            action_max_delay=3,  # env steps
+            imu_min_delay=0,  # env steps
+            imu_max_delay=3,  # env steps
+            scales=config_dict.create(
+                hip_pos=0.03,  # rad, for each hip joint
+                knee_pos=0.05,  # rad, for each knee joint
+                ankle_pos=0.08,  # rad, for each ankle joint
+                joint_vel=2.5,  # rad/s # Was 1.5
+                gravity=0.1,
+                linvel=0.1,
+                gyro=0.05,
+                accelerometer=0.005,
+            ),
+        ),
+        reward_config=config_dict.create(
+            scales=config_dict.create(
+                # tracking_lin_vel=2.5,
+                # tracking_ang_vel=4.0,
+                orientation=-0.5,
+                torques=-1.0e-3,
+                action_rate=-0.375,  # was -1.5
+                stand_still=-0.3,  # was -1.0 TODO try to relax this a bit ?
+                alive=20.0,
+                # imitation=1.0,
+                head_pos=-2.0,
+            ),
+            tracking_sigma=0.01,  # was working at 0.01
+        ),
+        push_config=config_dict.create(
+            enable=True,
+            interval_range=[5.0, 10.0],
+            magnitude_range=[0.1, 1.0],
+        ),
+        # lin_vel_x=[-0.1, 0.15],
+        # lin_vel_y=[-0.2, 0.2],
+        # ang_vel_yaw=[-1.0, 1.0],  # [-1.0, 1.0]
+        neck_pitch_range=[-0.34, 1.1],
+        head_pitch_range=[-0.78, 0.78],
+        head_yaw_range=[-2.7, 2.7],
+        head_roll_range=[-0.5, 0.5],
+        head_range_factor=1.0,
+    )
+
+
+class Standing(open_duck_mini_v2_base.OpenDuckMiniV2Env):
+    """Standing policy"""
+
+    def __init__(
+        self,
+        task: str = "flat_terrain",
+        config: config_dict.ConfigDict = default_config(),
+        config_overrides: Optional[Dict[str, Union[str, int, list[Any]]]] = None,
+    ):
+        super().__init__(
+            xml_path=constants.task_to_xml(task).as_posix(),
+            config=config,
+            config_overrides=config_overrides,
+        )
+        self._post_init()
+
+    def _post_init(self) -> None:
+
+        self._init_q = jp.array(self._mj_model.keyframe("home").qpos)
+        self._default_actuator = self._mj_model.keyframe(
+            "home"
+        ).ctrl  # ctrl of all the actual joints (no floating base and no backlash)
+
+        if USE_IMITATION_REWARD:
+            self.PRM = PolyReferenceMotion(
+                "playground/open_duck_mini_v2/data/polynomial_coefficients.pkl"
+            )
+
+        # Note: First joint is freejoint.
+        # get the range of the joints
+        self._lowers, self._uppers = self.mj_model.jnt_range[1:].T
+        c = (self._lowers + self._uppers) / 2
+        r = self._uppers - self._lowers
+        self._soft_lowers = c - 0.5 * r * self._config.soft_joint_pos_limit_factor
+        self._soft_uppers = c + 0.5 * r * self._config.soft_joint_pos_limit_factor
+
+        # weights for computing the cost of each joints compared to a reference pose
+        self._weights = jp.array(
+            [
+                1.0,
+                1.0,
+                0.01,
+                0.01,
+                1.0,  # left leg.
+                # 0.01, 0.01, 0.01, 0.01, 0.01, 0.01, #head
+                1.0,
+                1.0,
+                0.01,
+                0.01,
+                1.0,  # right leg.
+            ]
+        )
+
+        self._njoints = self._mj_model.njnt  # number of joints
+        self._actuators = self._mj_model.nu  # number of actuators
+
+        self._torso_body_id = self._mj_model.body(constants.ROOT_BODY).id
+        self._torso_mass = self._mj_model.body_subtreemass[self._torso_body_id]
+        self._site_id = self._mj_model.site("imu").id
+
+        self._feet_site_id = np.array(
+            [self._mj_model.site(name).id for name in constants.FEET_SITES]
+        )
+        self._floor_geom_id = self._mj_model.geom("floor").id
+        self._feet_geom_id = np.array(
+            [self._mj_model.geom(name).id for name in constants.FEET_GEOMS]
+        )
+
+        foot_linvel_sensor_adr = []
+        for site in constants.FEET_SITES:
+            sensor_id = self._mj_model.sensor(f"{site}_global_linvel").id
+            sensor_adr = self._mj_model.sensor_adr[sensor_id]
+            sensor_dim = self._mj_model.sensor_dim[sensor_id]
+            foot_linvel_sensor_adr.append(
+                list(range(sensor_adr, sensor_adr + sensor_dim))
+            )
+        self._foot_linvel_sensor_adr = jp.array(foot_linvel_sensor_adr)
+
+        # noise in the simu?
+        qpos_noise_scale = np.zeros(self._actuators)
+
+        hip_ids = [
+            idx for idx, j in enumerate(constants.JOINTS_ORDER_NO_HEAD) if "_hip" in j
+        ]
+        knee_ids = [
+            idx for idx, j in enumerate(constants.JOINTS_ORDER_NO_HEAD) if "_knee" in j
+        ]
+        ankle_ids = [
+            idx for idx, j in enumerate(constants.JOINTS_ORDER_NO_HEAD) if "_ankle" in j
+        ]
+
+        qpos_noise_scale[hip_ids] = self._config.noise_config.scales.hip_pos
+        qpos_noise_scale[knee_ids] = self._config.noise_config.scales.knee_pos
+        qpos_noise_scale[ankle_ids] = self._config.noise_config.scales.ankle_pos
+        # qpos_noise_scale[faa_ids] = self._config.noise_config.scales.faa_pos
+        self._qpos_noise_scale = jp.array(qpos_noise_scale)
+
+    def reset(self, rng: jax.Array) -> mjx_env.State:
+        qpos = self._init_q  # the complete qpos
+        # print(f'DEBUG0 init qpos: {qpos}')
+        qvel = jp.zeros(self.mjx_model.nv)
+
+        # init position/orientation in environment
+        # x=+U(-0.05, 0.05), y=+U(-0.05, 0.05), yaw=U(-3.14, 3.14).
+        rng, key = jax.random.split(rng)
+        dxy = jax.random.uniform(key, (2,), minval=-0.05, maxval=0.05)
+
+        # floating base
+        base_qpos = self.get_floating_base_qpos(qpos)
+        base_qpos = base_qpos.at[0:2].set(
+            qpos[self._floating_base_qpos_addr : self._floating_base_qpos_addr + 2]
+            + dxy
+        )  # x y noise
+
+        rng, key = jax.random.split(rng)
+        yaw = jax.random.uniform(key, (1,), minval=-3.14, maxval=3.14)
+        quat = math.axis_angle_to_quat(jp.array([0, 0, 1]), yaw)
+        new_quat = math.quat_mul(
+            qpos[self._floating_base_qpos_addr + 3 : self._floating_base_qpos_addr + 7],
+            quat,
+        )  # yaw noise
+
+        base_qpos = base_qpos.at[3:7].set(new_quat)
+
+        qpos = self.set_floating_base_qpos(base_qpos, qpos)
+        # print(f'DEBUG1 base qpos: {qpos}')
+        # init joint position
+        # qpos[7:]=*U(0.0, 0.1)
+        rng, key = jax.random.split(rng)
+
+        # multiply actual joints with noise (excluding floating base and backlash)
+        qpos_j = self.get_actuator_joints_qpos(qpos) * jax.random.uniform(
+            key, (self._actuators,), minval=0.5, maxval=1.5
+        )
+        qpos = self.set_actuator_joints_qpos(qpos_j, qpos)
+        # print(f'DEBUG2 joint qpos: {qpos}')
+        # init joint vel
+        # d(xyzrpy)=U(-0.05, 0.05)
+        rng, key = jax.random.split(rng)
+        # qvel = qvel.at[self._floating_base_qvel_addr : self._floating_base_qvel_addr + 6].set(
+        #     jax.random.uniform(key, (6,), minval=-0.5, maxval=0.5)
+        # )
+
+        qvel = self.set_floating_base_qvel(
+            jax.random.uniform(key, (6,), minval=-0.5, maxval=0.5), qvel
+        )
+        # print(f'DEBUG3 base qvel: {qvel}')
+        ctrl = self.get_actuator_joints_qpos(qpos)
+        # print(f'DEBUG4 ctrl: {ctrl}')
+        data = mjx_env.init(self.mjx_model, qpos=qpos, qvel=qvel, ctrl=ctrl)
+        rng, cmd_rng = jax.random.split(rng)
+        cmd = self.sample_command(cmd_rng)
+
+        # Sample push interval.
+        rng, push_rng = jax.random.split(rng)
+        push_interval = jax.random.uniform(
+            push_rng,
+            minval=self._config.push_config.interval_range[0],
+            maxval=self._config.push_config.interval_range[1],
+        )
+        push_interval_steps = jp.round(push_interval / self.dt).astype(jp.int32)
+
+        if USE_IMITATION_REWARD:
+            current_reference_motion = self.PRM.get_reference_motion(
+                cmd[0], cmd[1], cmd[2], 0
+            )
+        else:
+            current_reference_motion = jp.zeros(0)
+
+        info = {
+            "rng": rng,
+            "step": 0,
+            "command": cmd,
+            "last_act": jp.zeros(self.mjx_model.nu),
+            "last_last_act": jp.zeros(self.mjx_model.nu),
+            "last_last_last_act": jp.zeros(self.mjx_model.nu),
+            "motor_targets": jp.zeros(self.mjx_model.nu),
+            "feet_air_time": jp.zeros(2),
+            "last_contact": jp.zeros(2, dtype=bool),
+            "swing_peak": jp.zeros(2),
+            # Push related.
+            "push": jp.array([0.0, 0.0]),
+            "push_step": 0,
+            "push_interval_steps": push_interval_steps,
+            # History related.
+            "action_history": jp.zeros(
+                self._config.noise_config.action_max_delay * self._actuators
+            ),
+            "imu_history": jp.zeros(self._config.noise_config.imu_max_delay * 3),
+            # imitation related
+            "imitation_i": 0,
+            "current_reference_motion": current_reference_motion,
+        }
+
+        metrics = {}
+        for k, v in self._config.reward_config.scales.items():
+            if v != 0:
+                if v > 0:
+                    metrics[f"reward/{k}"] = jp.zeros(())
+                else:
+                    metrics[f"cost/{k}"] = jp.zeros(())
+        metrics["swing_peak"] = jp.zeros(())
+
+        contact = jp.array(
+            [
+                geoms_colliding(data, geom_id, self._floor_geom_id)
+                for geom_id in self._feet_geom_id
+            ]
+        )
+        obs = self._get_obs(data, info, contact)
+        reward, done = jp.zeros(2)
+        return mjx_env.State(data, obs, reward, done, metrics, info)
+
+    def step(self, state: mjx_env.State, action: jax.Array) -> mjx_env.State:
+
+        if USE_IMITATION_REWARD:
+            state.info["imitation_i"] += 1
+            state.info["imitation_i"] = (
+                state.info["imitation_i"] % self.PRM.nb_steps_in_period
+            )  # not critical, is already moduloed in get_reference_motion
+        else:
+            state.info["imitation_i"] = 0
+
+        if USE_IMITATION_REWARD:
+            state.info["current_reference_motion"] = self.PRM.get_reference_motion(
+                state.info["command"][0],
+                state.info["command"][1],
+                state.info["command"][2],
+                state.info["imitation_i"],
+            )
+        else:
+            state.info["current_reference_motion"] = jp.zeros(0)
+
+        state.info["rng"], push1_rng, push2_rng, action_delay_rng = jax.random.split(
+            state.info["rng"], 4
+        )
+
+        # Handle action delay
+        action_history = (
+            jp.roll(state.info["action_history"], self._actuators)
+            .at[: self._actuators]
+            .set(action)
+        )
+        state.info["action_history"] = action_history
+        action_idx = jax.random.randint(
+            action_delay_rng,
+            (1,),
+            minval=self._config.noise_config.action_min_delay,
+            maxval=self._config.noise_config.action_max_delay,
+        )
+        action_w_delay = action_history.reshape((-1, self._actuators))[
+            action_idx[0]
+        ]  # action with delay
+
+        push_theta = jax.random.uniform(push1_rng, maxval=2 * jp.pi)
+        push_magnitude = jax.random.uniform(
+            push2_rng,
+            minval=self._config.push_config.magnitude_range[0],
+            maxval=self._config.push_config.magnitude_range[1],
+        )
+        push = jp.array([jp.cos(push_theta), jp.sin(push_theta)])
+        push *= (
+            jp.mod(state.info["push_step"] + 1, state.info["push_interval_steps"]) == 0
+        )
+        push *= self._config.push_config.enable
+        qvel = state.data.qvel
+        qvel = qvel.at[
+            self._floating_base_qvel_addr : self._floating_base_qvel_addr + 2
+        ].set(
+            push * push_magnitude
+            + qvel[self._floating_base_qvel_addr : self._floating_base_qvel_addr + 2]
+        )  # floating base x,y
+        data = state.data.replace(qvel=qvel)
+        state = state.replace(data=data)
+
+        motor_targets = (
+            self._default_actuator + action_w_delay * self._config.action_scale
+        )
+        data = mjx_env.step(self.mjx_model, state.data, motor_targets, self.n_substeps)
+        state.info["motor_targets"] = motor_targets
+
+        contact = jp.array(
+            [
+                geoms_colliding(data, geom_id, self._floor_geom_id)
+                for geom_id in self._feet_geom_id
+            ]
+        )
+        contact_filt = contact | state.info["last_contact"]
+        first_contact = (state.info["feet_air_time"] > 0.0) * contact_filt
+        state.info["feet_air_time"] += self.dt
+        p_f = data.site_xpos[self._feet_site_id]
+        p_fz = p_f[..., -1]
+        state.info["swing_peak"] = jp.maximum(state.info["swing_peak"], p_fz)
+
+        obs = self._get_obs(data, state.info, contact)
+        done = self._get_termination(data)
+
+        rewards = self._get_reward(
+            data, action, state.info, state.metrics, done, first_contact, contact
+        )
+        # FIXME
+        rewards = {
+            k: v * self._config.reward_config.scales[k] for k, v in rewards.items()
+        }
+        reward = jp.clip(sum(rewards.values()) * self.dt, 0.0, 10000.0)
+        # jax.debug.print('STEP REWARD: {}',reward)
+        state.info["push"] = push
+        state.info["step"] += 1
+        state.info["push_step"] += 1
+        state.info["last_last_last_act"] = state.info["last_last_act"]
+        state.info["last_last_act"] = state.info["last_act"]
+        state.info["last_act"] = action
+        state.info["rng"], cmd_rng = jax.random.split(state.info["rng"])
+        state.info["command"] = jp.where(
+            state.info["step"] > 500,
+            self.sample_command(cmd_rng),
+            state.info["command"],
+        )
+        state.info["step"] = jp.where(
+            done | (state.info["step"] > 500),
+            0,
+            state.info["step"],
+        )
+        state.info["feet_air_time"] *= ~contact
+        state.info["last_contact"] = contact
+        state.info["swing_peak"] *= ~contact
+        for k, v in rewards.items():
+            rew_scale = self._config.reward_config.scales[k]
+            if rew_scale != 0:
+                if rew_scale > 0:
+                    state.metrics[f"reward/{k}"] = v
+                else:
+                    state.metrics[f"cost/{k}"] = -v
+        state.metrics["swing_peak"] = jp.mean(state.info["swing_peak"])
+
+        done = done.astype(reward.dtype)
+        state = state.replace(data=data, obs=obs, reward=reward, done=done)
+        return state
+
+    def _get_termination(self, data: mjx.Data) -> jax.Array:
+        fall_termination = self.get_gravity(data)[-1] < 0.0
+        return fall_termination | jp.isnan(data.qpos).any() | jp.isnan(data.qvel).any()
+
+    def _get_obs(
+        self, data: mjx.Data, info: dict[str, Any], contact: jax.Array
+    ) -> mjx_env.Observation:
+
+        gyro = self.get_gyro(data)
+        info["rng"], noise_rng = jax.random.split(info["rng"])
+        noisy_gyro = (
+            gyro
+            + (2 * jax.random.uniform(noise_rng, shape=gyro.shape) - 1)
+            * self._config.noise_config.level
+            * self._config.noise_config.scales.gyro
+        )
+
+        accelerometer = self.get_accelerometer(data)
+        info["rng"], noise_rng = jax.random.split(info["rng"])
+        noisy_accelerometer = (
+            accelerometer
+            + (2 * jax.random.uniform(noise_rng, shape=accelerometer.shape) - 1)
+            * self._config.noise_config.level
+            * self._config.noise_config.scales.accelerometer
+        )
+
+        gravity = data.site_xmat[self._site_id].T @ jp.array([0, 0, -1])
+        info["rng"], noise_rng = jax.random.split(info["rng"])
+        noisy_gravity = (
+            gravity
+            + (2 * jax.random.uniform(noise_rng, shape=gravity.shape) - 1)
+            * self._config.noise_config.level
+            * self._config.noise_config.scales.gravity
+        )
+
+        # Handle IMU delay
+        imu_history = jp.roll(info["imu_history"], 3).at[:3].set(noisy_gravity)
+        info["imu_history"] = imu_history
+        imu_idx = jax.random.randint(
+            noise_rng,
+            (1,),
+            minval=self._config.noise_config.imu_min_delay,
+            maxval=self._config.noise_config.imu_max_delay,
+        )
+        noisy_gravity = imu_history.reshape((-1, 3))[imu_idx[0]]
+
+        # joint_angles = data.qpos[7:]
+
+        # Handling backlash
+        joint_angles = self.get_actuator_joints_qpos(data.qpos)
+        joint_backlash = self.get_actuator_backlash_qpos(data.qpos)
+
+        for i in self.backlash_idx_to_add:
+            joint_backlash = jp.insert(joint_backlash, i, 0)
+
+        joint_angles = joint_angles + joint_backlash
+
+        info["rng"], noise_rng = jax.random.split(info["rng"])
+        noisy_joint_angles = (
+            joint_angles
+            + (2.0 * jax.random.uniform(noise_rng, shape=joint_angles.shape) - 1.0)
+            * self._config.noise_config.level
+            * self._qpos_noise_scale
+        )
+
+        # joint_vel = data.qvel[6:]
+        joint_vel = self.get_actuator_joints_qvel(data.qvel)
+        info["rng"], noise_rng = jax.random.split(info["rng"])
+        noisy_joint_vel = (
+            joint_vel
+            + (2.0 * jax.random.uniform(noise_rng, shape=joint_vel.shape) - 1.0)
+            * self._config.noise_config.level
+            * self._config.noise_config.scales.joint_vel
+        )
+
+        linvel = self.get_local_linvel(data)
+        # info["rng"], noise_rng = jax.random.split(info["rng"])
+        # noisy_linvel = (
+        #     linvel
+        #     + (2 * jax.random.uniform(noise_rng, shape=linvel.shape) - 1)
+        #     * self._config.noise_config.level
+        #     * self._config.noise_config.scales.linvel
+        # )
+
+        state = jp.hstack(
+            [
+                # noisy_linvel,  # 3
+                # noisy_gyro,  # 3
+                # noisy_gravity,  # 3
+                noisy_gyro,  # 3
+                noisy_accelerometer,  # 3
+                info["command"],  # 3
+                noisy_joint_angles - self._default_actuator,  # 10
+                noisy_joint_vel * self._config.dof_vel_scale,  # 10
+                info["last_act"],  # 10
+                info["last_last_act"],  # 10
+                info["last_last_last_act"],  # 10
+                contact,  # 2
+                info["current_reference_motion"],
+            ]
+        )
+
+        accelerometer = self.get_accelerometer(data)
+        global_angvel = self.get_global_angvel(data)
+        feet_vel = data.sensordata[self._foot_linvel_sensor_adr].ravel()
+        root_height = data.qpos[self._floating_base_qpos_addr + 2]
+
+        privileged_state = jp.hstack(
+            [
+                state,
+                gyro,  # 3
+                accelerometer,  # 3
+                gravity,  # 3
+                linvel,  # 3
+                global_angvel,  # 3
+                joint_angles - self._default_actuator,
+                joint_vel,
+                root_height,  # 1
+                data.actuator_force,  # 10
+                contact,  # 2
+                feet_vel,  # 4*3
+                info["feet_air_time"],  # 2
+                info["current_reference_motion"],
+            ]
+        )
+
+        return {
+            "state": state,
+            "privileged_state": privileged_state,
+        }
+
+    def _get_reward(
+        self,
+        data: mjx.Data,
+        action: jax.Array,
+        info: dict[str, Any],
+        metrics: dict[str, Any],
+        done: jax.Array,
+        first_contact: jax.Array,
+        contact: jax.Array,
+    ) -> dict[str, jax.Array]:
+        del metrics  # Unused.
+
+        ret = {
+            "orientation": cost_orientation(self.get_gravity(data)),
+            "torques": cost_torques(data.actuator_force),
+            "action_rate": cost_action_rate(action, info["last_act"]),
+            "alive": reward_alive(),
+            "stand_still": cost_stand_still(
+                # info["command"], data.qpos[7:], data.qvel[6:], self._default_pose
+                info["command"],
+                self.get_actuator_joints_qpos(data.qpos),
+                self.get_actuator_joints_qvel(data.qvel),
+                self._default_actuator,
+                True
+            ),
+            "head_pos": cost_head_pos(
+                self.get_actuator_joints_qpos(data.qpos),
+                self.get_actuator_joints_qvel(data.qvel),
+                info["command"],
+            ),
+        }
+
+        return ret
+
+    def sample_command(self, rng: jax.Array) -> jax.Array:
+        rng1, rng2, rng3, rng4, rng5, rng6, rng7, rng8 = jax.random.split(rng, 8)
+
+        # lin_vel_x = jax.random.uniform(
+        #     rng1, minval=self._config.lin_vel_x[0], maxval=self._config.lin_vel_x[1]
+        # )
+        # lin_vel_y = jax.random.uniform(
+        #     rng2, minval=self._config.lin_vel_y[0], maxval=self._config.lin_vel_y[1]
+        # )
+        # ang_vel_yaw = jax.random.uniform(
+        #     rng3,
+        #     minval=self._config.ang_vel_yaw[0],
+        #     maxval=self._config.ang_vel_yaw[1],
+        # )
+
+        neck_pitch = jax.random.uniform(
+            rng5,
+            minval=self._config.neck_pitch_range[0] * self._config.head_range_factor,
+            maxval=self._config.neck_pitch_range[1] * self._config.head_range_factor,
+        )
+
+        head_pitch = jax.random.uniform(
+            rng6,
+            minval=self._config.head_pitch_range[0] * self._config.head_range_factor,
+            maxval=self._config.head_pitch_range[1] * self._config.head_range_factor,
+        )
+
+        head_yaw = jax.random.uniform(
+            rng7,
+            minval=self._config.head_yaw_range[0] * self._config.head_range_factor,
+            maxval=self._config.head_yaw_range[1] * self._config.head_range_factor,
+        )
+
+        head_roll = jax.random.uniform(
+            rng8,
+            minval=self._config.head_roll_range[0] * self._config.head_range_factor,
+            maxval=self._config.head_roll_range[1] * self._config.head_range_factor,
+        )
+
+        # With 10% chance, set everything to zero.
+        return jp.where(
+            jax.random.bernoulli(rng4, p=0.1),
+            jp.zeros(7),
+            jp.hstack(
+                [
+                    0.0,
+                    0.0,
+                    0.0,
+                    neck_pitch,
+                    head_pitch,
+                    head_yaw,
+                    head_roll,
+                ]
+            ),
+        )

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/antenna.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MmXLbiyIJZ1T9tiX3",
+    "isStandardContent": false,
+    "name": "antenna <1>",
+    "partId": "RGPD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/body_back.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MfCA08JspCUYCoOzy",
+    "isStandardContent": false,
+    "name": "body_back <1>",
+    "partId": "RqKD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/body_front.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MAILAHna+4EMDjwEB",
+    "isStandardContent": false,
+    "name": "body_front <1>",
+    "partId": "RbKD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/body_middle_bottom.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "Mi7/ZBgz7j8FhNN8n",
+    "isStandardContent": false,
+    "name": "body_middle_bottom <1>",
+    "partId": "R/KH",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/body_middle_top.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MwLs2m7WqJdIp5rCM",
+    "isStandardContent": false,
+    "name": "body_middle_top <1>",
+    "partId": "R/KD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/foot_bottom_pla.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MbWaXbNFhwXrvfPr1",
+    "isStandardContent": false,
+    "name": "foot_bottom_pla <1>",
+    "partId": "R7GH",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/foot_bottom_tpu.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MSRJjnMROc2yqGlBV",
+    "isStandardContent": false,
+    "name": "foot_bottom_tpu <1>",
+    "partId": "R7GD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/foot_side.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "M9QvdrmpVd96BQQ3J",
+    "isStandardContent": false,
+    "name": "foot_side <1>",
+    "partId": "RsGD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/foot_top.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MryNg763VjsN4xBWU",
+    "isStandardContent": false,
+    "name": "foot_top <1>",
+    "partId": "RsGH",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/head.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MGPnRRHSCLRuHbf9s",
+    "isStandardContent": false,
+    "name": "head <1>",
+    "partId": "RXMD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/head_bot_sheet.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MB88Zh5nm8uQRQJB5",
+    "isStandardContent": false,
+    "name": "head_bot_sheet <1>",
+    "partId": "RwMD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/head_pitch_to_yaw.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MMYrJ8GU6Tulh9Ezl",
+    "isStandardContent": false,
+    "name": "head_pitch_to_yaw <1>",
+    "partId": "RGCD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/head_yaw_to_roll.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MSSQrz89LHCpx23/c",
+    "isStandardContent": false,
+    "name": "head_yaw_to_roll <1>",
+    "partId": "RyCD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/left_antenna_holder.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MqpShx/iKn+WMc6nz",
+    "isStandardContent": false,
+    "name": "left_antenna_holder <1>",
+    "partId": "R3PL",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/left_cache.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "Mh206kD+LDwOu9+3I",
+    "isStandardContent": false,
+    "name": "left_cache <1>",
+    "partId": "RELD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/left_knee_to_ankle_left_sheet.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MkexR4anxdIDvoNnP",
+    "isStandardContent": false,
+    "name": "left_knee_to_ankle_left_sheet <1>",
+    "partId": "RyDD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/left_knee_to_ankle_right_sheet.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MFSYykuSOyFPl8MEE",
+    "isStandardContent": false,
+    "name": "left_knee_to_ankle_right_sheet <1>",
+    "partId": "RyDH",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/left_roll_to_pitch.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "Mg7gXWiWLFdTkX/WZ",
+    "isStandardContent": false,
+    "name": "left_roll_to_pitch <1>",
+    "partId": "RRFD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/leg_spacer.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "M1AJs3nTP2bD6DgAF",
+    "isStandardContent": false,
+    "name": "leg_spacer <1>",
+    "partId": "RfED",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/neck_left_sheet.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "M8Je0JZqWOimCS7cY",
+    "isStandardContent": false,
+    "name": "neck_left_sheet <1>",
+    "partId": "RgLD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/neck_right_sheet.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "M8fJckSPn4u+ct8pO",
+    "isStandardContent": false,
+    "name": "neck_right_sheet <1>",
+    "partId": "RgLH",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/right_antenna_holder.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "M/p8PGArLA1nrXwe8",
+    "isStandardContent": false,
+    "name": "right_antenna_holder <1>",
+    "partId": "R0OD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/right_cache.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MBtFsLan/52XkruX7",
+    "isStandardContent": false,
+    "name": "right_cache <1>",
+    "partId": "R3PH",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/right_roll_to_pitch.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "Mvz4SlAytL1YUO5L4",
+    "isStandardContent": false,
+    "name": "right_roll_to_pitch <1>",
+    "partId": "R3PD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/roll_motor_bottom.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MjsDdlagt/LopvMjR",
+    "isStandardContent": false,
+    "name": "roll_motor_bottom <1>",
+    "partId": "RGHD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/roll_motor_top.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MoZYOvTkbbGLGW0df",
+    "isStandardContent": false,
+    "name": "roll_motor_top <1>",
+    "partId": "RyHD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/trunk_bottom.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MMYHrJ4hqH24cpa8I",
+    "isStandardContent": false,
+    "name": "trunk_bottom <1>",
+    "partId": "RvJH",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/assets/trunk_top.part

@@ -0,0 +1,13 @@
+{
+    "configuration": "default",
+    "documentId": "64074dfcfa379b37d8a47762",
+    "documentMicroversion": "78f84c802a6a701851609660",
+    "elementId": "560456e30aa4a28c1f2fba59",
+    "fullConfiguration": "default",
+    "id": "MNg8avQSkDpa0UkRX",
+    "isStandardContent": false,
+    "name": "trunk_top <1>",
+    "partId": "RvJD",
+    "suppressed": false,
+    "type": "Part"
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/config.json

@@ -0,0 +1,51 @@
+{
+    // NEED TO WRAP all in a body named "base" and add site 
+    // <site name="imu" pos="-0.08 -0.0 0.05"/>
+    // and comment the default kp frictionloss etc
+    "documentId": "64074dfcfa379b37d8a47762",
+    "outputFormat": "mujoco",
+    "assemblyName": "Open_Duck_Mini_v2",
+    "robot_name": "open_duck_mini_v2",
+    "output_filename": "open_duck_mini_v2",
+    // Simplify STL meshes (default: false)
+    "simplify_stls": true,
+    // Maximum size of the STL files in MB (default: 3)
+    "max_stl_size": 0.005,
+    // "collisions_as_visual": true,
+    "ignore": {
+        "wj*": "visual",
+        "drive_palonier": "visual",
+        "passive_palonier": "visual",
+        "bulb": "visual",
+        "cam": "visual",
+        "sg90": "visual",
+        "head_roll_mount": "visual",
+        "raspberrypizerow": "visual",
+        "speaker_stand": "visual",
+        "speaker_interface": "visual",
+        "holder": "visual",
+        "cell": "visual",
+        "bms": "visual",
+        "usb_c_charger": "visual",
+        "battery_pack_lid": "visual",
+        "right_eye": "visual",
+        "left_eye": "visual",
+        "glass": "visual",
+        "full_speaker": "visual",
+        "roll_bearing": "visual",
+        "power_switch": "visual",
+        "board": "visual",
+        "bno055": "visual",
+        "*": "collision",
+        "!foot_bottom_tpu": "collision"
+    },
+    "joint_properties": {
+        "default": {
+            "class": "sts3215"
+        }
+    },
+    "additional_xml": [
+        "sensors.xml",
+        "joints_properties.xml"
+    ]
+}

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/joints_properties.xml

@@ -0,0 +1,12 @@
+<default>
+    <default class="sts3215">
+        <geom contype="0" conaffinity="0"/>
+        <joint damping="0.60" frictionloss="0.052" armature="0.028"/>
+        <position kp="17.8" kv="0.0" forcerange="-3.35 3.35"/>
+    </default>
+    <default class="backlash">
+        <!-- +/- 0.5° of backlash -->
+        <joint damping="0.01" frictionloss="0" armature="0.01" limited="true"
+            range="-0.008726646259971648 0.008726646259971648"/>
+    </default>
+</default>

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/open_duck_mini_v2.xml

@@ -0,0 +1,503 @@
+<?xml version="1.0"?>
+<!-- Generated using onshape-to-robot -->
+<!-- Onshape document_id: 64074dfcfa379b37d8a47762 -->
+<mujoco model="open_duck_mini_v2">
+
+  <option iterations="1" ls_iterations="5">
+    <flag eulerdamp="disable"/>
+  </option>
+  <compiler angle="radian" meshdir="assets"/>
+
+  <!-- <option noslip_iterations="1"/>
+  <compiler angle="radian" meshdir="assets" autolimits="true"/> -->
+  <default>
+    <default class="open_duck_mini_v2">
+      <joint frictionloss="0.1" armature="0.005"/>
+      <position kp="50" dampratio="1"/>
+      <default class="visual">
+        <geom type="mesh" contype="0" conaffinity="0" group="2"/>
+      </default>
+      <default class="collision">
+        <geom group="3"/>
+      </default>
+    </default>
+  </default>
+  <!-- Additional sensors.xml -->
+  <sensor>
+    <gyro site="imu" name="gyro"/>
+    <velocimeter site="imu" name="local_linvel"/>
+    <accelerometer site="imu" name="accelerometer"/>
+    <framezaxis objtype="site" objname="imu" name="upvector"/>
+    <framexaxis objtype="site" objname="imu" name="forwardvector"/>
+    <framelinvel objtype="site" objname="imu" name="global_linvel"/>
+    <frameangvel objtype="site" objname="imu" name="global_angvel"/>
+    <framepos objtype="site" objname="imu" name="position"/>
+    <framequat objtype="site" objname="imu" name="orientation"/>
+    <framelinvel objtype="site" objname="right_foot" name="right_foot_global_linvel"/>
+    <framelinvel objtype="site" objname="left_foot" name="left_foot_global_linvel"/>
+    <framexaxis objtype="site" objname="left_foot" name="left_foot_upvector"/>
+    <framexaxis objtype="site" objname="right_foot" name="right_foot_upvector"/>
+    <framepos objtype="site" objname="left_foot" name="left_foot_pos"/>
+    <framepos objtype="site" objname="right_foot" name="right_foot_pos"/>
+  </sensor>
+  <!-- Additional joints_properties.xml -->
+  <default>
+    <default class="sts3215">
+      <geom contype="0" conaffinity="0"/>
+      <joint damping="0.56" frictionloss="0.068" armature="0.027"/>
+      <!-- <position kp="17.12" kv="0.0" forcerange="-3.23 3.23"/> -->
+      <position kp="13.37" kv="0.0" forcerange="-3.23 3.23"/>
+    </default>
+    <default class="backlash">
+      <!-- +/- 0.5° of backlash -->
+      <joint damping="0.01" frictionloss="0" armature="0.01" limited="true"
+        range="-0.008726646259971648 0.008726646259971648"/>
+    </default>
+  </default>
+  <worldbody>
+    <body name="base" pos="0 0 0.22">
+      <freejoint name="floating_base"/>
+      <site name="imu" pos="-0.08 -0.0 0.05"/>
+      <!-- Link trunk_assembly -->
+      <body name="trunk_assembly" pos="0 0 0.0" quat="1 0 0 0" childclass="open_duck_mini_v2">
+        <inertial pos="-0.0483259 -9.97823e-05 0.0384971" mass="0.698526"
+          fullinertia="0.00167701 0.00344479 0.00292635 -1.31662e-05 -3.24953e-05 -2.15581e-06"/>
+        <!-- Part body_front -->
+        <geom type="mesh" class="visual" pos="-0.019 0 0.0648909" quat="1 0 0 0" mesh="body_front"
+          material="body_front_material"/>
+        <!-- Part usb_c_charger -->
+        <!-- Part body_back -->
+        <geom type="mesh" class="visual" pos="-0.019 3.46945e-18 0.0648909" quat="1 0 0 0" mesh="body_back"
+          material="body_back_material"/>
+        <!-- Part body_middle_bottom -->
+        <geom type="mesh" class="visual" pos="-0.019 -3.46945e-18 0.0648909" quat="1 0 0 0" mesh="body_middle_bottom"
+          material="body_middle_bottom_material"/>
+        <!-- Part power_switch -->
+        <!-- Part bms -->
+        <!-- Part battery_pack_lid -->
+        <!-- Part body_middle_top -->
+        <geom type="mesh" class="visual" pos="-0.019 0 0.0648909" quat="1 0 0 0" mesh="body_middle_top"
+          material="body_middle_top_material"/>
+        <!-- Part bno055 -->
+        <!-- Part roll_bearing -->
+        <!-- Part trunk_bottom -->
+        <geom type="mesh" class="visual" pos="-0.019 9.1073e-18 0.0648909" quat="1 -0 -0 5.42101e-20"
+          mesh="trunk_bottom" material="trunk_bottom_material"/>
+        <!-- Part trunk_top -->
+        <geom type="mesh" class="visual" pos="-0.019 -2.42428e-16 0.0648909"
+          quat="1 -4.82108e-17 -5.5355e-15 2.77556e-17" mesh="trunk_top" material="trunk_top_material"/>
+        <!-- Part board -->
+        <!-- Part roll_bearing_2 -->
+        <!-- Part cell -->
+        <!-- Part holder -->
+        <!-- Part cell_2 -->
+        <!-- Part wj_wk00_0123middlecase_56 -->
+        <!-- Part drive_palonier -->
+        <!-- Part wj_wk00_0124bottomcase_45 -->
+        <!-- Part wj_wk00_0122topcabinetcase_95 -->
+        <!-- Part passive_palonier -->
+        <!-- Part wj_wk00_0123middlecase_56_2 -->
+        <!-- Part drive_palonier_2 -->
+        <!-- Part wj_wk00_0124bottomcase_45_2 -->
+        <!-- Part wj_wk00_0122topcabinetcase_95_2 -->
+        <!-- Part passive_palonier_2 -->
+        <!-- Part wj_wk00_0123middlecase_56_3 -->
+        <!-- Part drive_palonier_3 -->
+        <!-- Part wj_wk00_0124bottomcase_45_3 -->
+        <!-- Part wj_wk00_0122topcabinetcase_95_3 -->
+        <!-- Part passive_palonier_3 -->
+        <!-- Frame trunk -->
+        <site group="0" name="trunk" pos="-0.024 -2.41127e-16 0.0881909" quat="1 -2.77556e-17 -0 -0"/>
+        <!-- Frame imu -->
+        <!-- <site group="0" name="imu" pos="-0.08711 1.249e-16 0.0422909" quat="1 -7.03105e-17 -2.77556e-17
+        -3.14419e-16"/> -->
+        <!-- Link hip_roll_assembly -->
+        <body name="hip_roll_assembly" pos="-0.019 0.035 0.0459409" quat="1 -1.17663e-16 -1.09357e-14 2.77556e-17">
+          <!-- Joint from trunk_assembly to hip_roll_assembly -->
+          <joint name="left_hip_yaw" type="hinge" range="-0.5235987755982979 0.5235987755982997" class="sts3215"/>
+          <inertial pos="0.000795081 -5.46541e-06 -0.03306" mass="0.06648"
+            fullinertia="2.44411e-05 2.81818e-05 1.42687e-05 2.35658e-09 -3.81059e-07 -7.50036e-09"/>
+          <!-- Part roll_motor_bottom -->
+          <geom type="mesh" class="visual" pos="-3.33067e-16 -0.035 0.01905"
+            quat="1 -4.38499e-16 2.22045e-16 -3.34802e-16" mesh="roll_motor_bottom"
+            material="roll_motor_bottom_material"/>
+          <!-- Part roll_motor_top -->
+          <geom type="mesh" class="visual" pos="0 -0.035 0.01905" quat="1 -0 -1.80411e-16 -4.16334e-17"
+            mesh="roll_motor_top" material="roll_motor_top_material"/>
+          <!-- Part wj_wk00_0123middlecase_56_4 -->
+          <!-- Part drive_palonier_4 -->
+          <!-- Part wj_wk00_0124bottomcase_45_4 -->
+          <!-- Part wj_wk00_0122topcabinetcase_95_4 -->
+          <!-- Part passive_palonier_4 -->
+          <!-- Link left_roll_to_pitch_assembly -->
+          <body name="left_roll_to_pitch_assembly" pos="0.01905 4.61436e-15 -0.046015" quat="0.5 -0.5 -0.5 0.5">
+            <!-- Joint from hip_roll_assembly to left_roll_to_pitch_assembly -->
+            <joint name="left_hip_roll" type="hinge" range="-0.4363323129985815 0.43633231299858327" class="sts3215"/>
+            <inertial pos="0.0508042 -0.00041105 0.0204704" mass="0.07516"
+              fullinertia="2.52335e-05 4.13855e-05 2.8127e-05 1.22133e-07 -5.90499e-07 -1.07544e-07"/>
+            <!-- Part left_roll_to_pitch -->
+            <geom type="mesh" class="visual" pos="-0.035 -0.065 0.01905" quat="0.5 0.5 0.5 -0.5"
+              mesh="left_roll_to_pitch" material="left_roll_to_pitch_material"/>
+            <!-- Part wj_wk00_0123middlecase_56_5 -->
+            <!-- Part drive_palonier_5 -->
+            <!-- Part wj_wk00_0124bottomcase_45_5 -->
+            <!-- Part wj_wk00_0122topcabinetcase_95_5 -->
+            <!-- Part passive_palonier_5 -->
+            <!-- Link knee_and_ankle_assembly -->
+            <body name="knee_and_ankle_assembly" pos="0.07415 -3.25261e-17 0.03511"
+              quat="8.49887e-24 -0.707107 -1.11022e-16 0.707107">
+              <!-- Joint from left_roll_to_pitch_assembly to knee_and_ankle_assembly -->
+              <joint name="left_hip_pitch" type="hinge" range="-1.2217304763960306 0.5235987755982988" class="sts3215"/>
+              <inertial pos="0.00253369 -0.0390636 0.0102776" mass="0.12407"
+                fullinertia="0.000211406 7.83759e-05 0.000225619 -2.68706e-05 3.81906e-06 2.08009e-05"/>
+              <!-- Part left_cache -->
+              <geom type="mesh" class="visual" pos="0.01606 0.065 0.10915"
+                quat="0.707107 -0.707107 3.33067e-16 5.76761e-16" mesh="left_cache" material="left_cache_material"/>
+              <!-- Part leg_spacer -->
+              <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10925"
+                quat="0.707107 -0.707107 3.33067e-16 5.94711e-16" mesh="leg_spacer" material="leg_spacer_material"/>
+              <!-- Part left_knee_to_ankle_right_sheet -->
+              <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10915" quat="0.707107 -0.707107 0 -5.93802e-24"
+                mesh="left_knee_to_ankle_right_sheet" material="left_knee_to_ankle_right_sheet_material"/>
+              <!-- Part left_knee_to_ankle_left_sheet -->
+              <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10915" quat="0.707107 -0.707107 0 5.12923e-30"
+                mesh="left_knee_to_ankle_left_sheet" material="left_knee_to_ankle_left_sheet_material"/>
+              <!-- Part wj_wk00_0123middlecase_56_6 -->
+              <!-- Part drive_palonier_6 -->
+              <!-- Part wj_wk00_0124bottomcase_45_6 -->
+              <!-- Part wj_wk00_0122topcabinetcase_95_6 -->
+              <!-- Part passive_palonier_6 -->
+              <!-- Link knee_and_ankle_assembly_2 -->
+              <body name="knee_and_ankle_assembly_2" pos="0 -0.07865 0.0001"
+                quat="1 1.11022e-16 1.11022e-16 -2.64698e-23">
+                <!-- Joint from knee_and_ankle_assembly to knee_and_ankle_assembly_2 -->
+                <joint name="left_knee" type="hinge" range="-1.5707963267948966 1.5707963267948966" class="sts3215"/>
+                <inertial pos="5.01859e-06 -0.0577465 0.0181136" mass="0.07259"
+                  fullinertia="4.99575e-05 1.87273e-05 4.23932e-05 1.03817e-08 2.25175e-09 8.81049e-08"/>
+                <!-- Part leg_spacer_2 -->
+                <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10925"
+                  quat="0.707107 -0.707107 3.33067e-16 5.72905e-16" mesh="leg_spacer" material="leg_spacer_material"/>
+                <!-- Part left_knee_to_ankle_right_sheet_2 -->
+                <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10915" quat="0.707107 -0.707107 0 1.33119e-24"
+                  mesh="left_knee_to_ankle_right_sheet" material="left_knee_to_ankle_right_sheet_material"/>
+                <!-- Part left_knee_to_ankle_left_sheet_2 -->
+                <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10915" quat="0.707107 -0.707107 0 -5.36743e-30"
+                  mesh="left_knee_to_ankle_left_sheet" material="left_knee_to_ankle_left_sheet_material"/>
+                <!-- Part wj_wk00_0123middlecase_56_7 -->
+                <!-- Part drive_palonier_7 -->
+                <!-- Part wj_wk00_0124bottomcase_45_7 -->
+                <!-- Part wj_wk00_0122topcabinetcase_95_7 -->
+                <!-- Part passive_palonier_7 -->
+                <!-- Link foot_assembly -->
+                <body name="foot_assembly" pos="1.8735e-16 -0.07865 0.0001" quat="1 2.77556e-16 5.55112e-16 1.52656e-16">
+                  <!-- Joint from knee_and_ankle_assembly_2 to foot_assembly -->
+                  <joint name="left_ankle" type="hinge" range="-1.5707963267948957 1.5707963267948974" class="sts3215"/>
+                  <inertial pos="0.0110718 -0.0246608 0.0190626" mass="0.07524"
+                    fullinertia="1.87482e-05 6.74427e-05 6.061e-05 1.58839e-06 -6.52354e-09 5.67391e-08"/>
+                  <!-- Part foot_side -->
+                  <geom type="mesh" class="visual" pos="0.01606 0.2223 0.10905"
+                    quat="0.707107 -0.707107 -0 -4.61924e-30" mesh="foot_side" material="foot_side_material"/>
+                  <!-- Part foot_bottom_tpu -->
+                  <geom type="mesh" class="visual" pos="0.01656 0.2228 0.10955" quat="0.707107 -0.707107 -0 1.72496e-18"
+                    mesh="foot_bottom_tpu" material="foot_bottom_tpu_material"/>
+                  <geom type="mesh" class="collision" pos="0.01656 0.2228 0.10955"
+                    quat="0.707107 -0.707107 -0 1.72496e-18" mesh="foot_bottom_tpu" material="foot_bottom_tpu_material"
+                    name="left_foot_bottom_tpu"/>
+                  <!-- Part foot_bottom_pla -->
+                  <geom type="mesh" class="visual" pos="0.01656 0.2228 0.10955" quat="0.707107 -0.707107 -0 1.72496e-18"
+                    mesh="foot_bottom_pla" material="foot_bottom_pla_material"/>
+                  <!-- Part foot_top -->
+                  <geom type="mesh" class="visual" pos="0.01606 0.2223 0.10905"
+                    quat="0.707107 -0.707107 -0 -4.61924e-30" mesh="foot_top" material="foot_top_material"/>
+                  <!-- Frame left_foot -->
+                  <site group="0" name="left_foot" pos="0.0005 -0.036225 0.01955"
+                    quat="0.707107 -0.707107 -0 1.72494e-18"/>
+                </body>
+              </body>
+            </body>
+          </body>
+        </body>
+        <!-- Link neck_pitch_assembly -->
+        <body name="neck_pitch_assembly" pos="0.001 0.01915 0.0900009" quat="0.707107 0.707107 -7.88258e-15 7.79925e-15">
+          <!-- Joint from trunk_assembly to neck_pitch_assembly -->
+          <joint name="neck_pitch" type="hinge" range="-0.3490658503988437 1.1344640137963364" class="sts3215"/>
+          <inertial pos="-5.63137e-06 0.0492968 0.0181786" mass="0.06618"
+            fullinertia="3.49456e-05 1.70231e-05 2.80043e-05 8.75633e-09 -2.15592e-09 -1.85188e-08"/>
+          <!-- Part neck_left_sheet -->
+          <geom type="mesh" class="visual" pos="-0.02 -0.02511 0.01905" quat="0.707107 -0.707107 0 -5.1423e-31"
+            mesh="neck_left_sheet" material="neck_left_sheet_material"/>
+          <!-- Part neck_right_sheet -->
+          <geom type="mesh" class="visual" pos="-0.02 -0.02511 0.01905" quat="0.707107 -0.707107 -0 3.06492e-17"
+            mesh="neck_right_sheet" material="neck_right_sheet_material"/>
+          <!-- Part wj_wk00_0123middlecase_56_8 -->
+          <!-- Part drive_palonier_8 -->
+          <!-- Part wj_wk00_0124bottomcase_45_8 -->
+          <!-- Part wj_wk00_0122topcabinetcase_95_8 -->
+          <!-- Part passive_palonier_8 -->
+          <!-- Link head_pitch_to_yaw -->
+          <body name="head_pitch_to_yaw" pos="-2.9924e-17 0.066 -6.93889e-18" quat="1 -1.7203e-16 0 1.10553e-17">
+            <!-- Joint from neck_pitch_assembly to head_pitch_to_yaw -->
+            <joint name="head_pitch" type="hinge" range="-0.7853981633974483 0.7853981633974483" class="sts3215"/>
+            <inertial pos="-0.00766247 0.026015 0.0186681" mass="0.0169378"
+              fullinertia="9.43036e-06 4.64763e-06 8.02915e-06 1.35195e-06 3.05899e-08 -8.42802e-08"/>
+            <!-- Part head_pitch_to_yaw -->
+            <geom type="mesh" class="visual" pos="-0.02 -0.09111 0.01905" quat="0.707107 -0.707107 0 1.82143e-39"
+              mesh="head_pitch_to_yaw" material="head_pitch_to_yaw_material"/>
+            <!-- Link neck_yaw_assembly -->
+            <body name="neck_yaw_assembly" pos="-4.33681e-19 0.057 0.01905"
+              quat="0.707107 -0.707107 1.11022e-16 -5.88785e-17">
+              <!-- Joint from head_pitch_to_yaw to neck_yaw_assembly -->
+              <joint name="head_yaw" type="hinge" range="-2.792526803190927 2.792526803190927" class="sts3215"/>
+              <inertial pos="0.00412907 3.95745e-06 -0.0222828" mass="0.0918099"
+                fullinertia="3.11706e-05 6.94935e-05 7.04025e-05 4.29429e-09 -4.41251e-06 -3.82028e-09"/>
+              <!-- Part head_yaw_to_roll -->
+              <geom type="mesh" class="visual" pos="-0.02 1.4985e-16 -0.14821"
+                quat="1 6.93889e-18 -5.55112e-17 7.04614e-17" mesh="head_yaw_to_roll"
+                material="head_yaw_to_roll_material"/>
+              <!-- Part wj_wk00_0123middlecase_56_9 -->
+              <!-- Part drive_palonier_9 -->
+              <!-- Part wj_wk00_0124bottomcase_45_9 -->
+              <!-- Part wj_wk00_0122topcabinetcase_95_9 -->
+              <!-- Part passive_palonier_9 -->
+              <!-- Link head_assembly -->
+              <body name="head_assembly" pos="0.04095 2.1764e-16 -0.01915"
+                quat="0.707107 2.89132e-14 -0.707107 -2.86438e-14">
+                <!-- Joint from neck_yaw_assembly to head_assembly -->
+                <joint name="head_roll" type="hinge" range="-0.523598775598218 0.5235987755983796" class="sts3215"/>
+                <inertial pos="0.00815416 -0.00390754 0.0227726" mass="0.406607"
+                  fullinertia="0.00244627 0.0016864 0.00108002 -2.02404e-06 9.2382e-05 -1.84944e-05"/>
+                <!-- Part right_antenna_holder -->
+                <geom type="mesh" class="visual" pos="-0.12906 -2.4549e-16 0.06095"
+                  quat="0.707107 5.22659e-16 0.707107 3.88578e-15" mesh="right_antenna_holder"
+                  material="right_antenna_holder_material"/>
+                <!-- Part antenna -->
+                <geom type="mesh" class="visual" pos="-0.0830595 0.182212 0.06095"
+                  quat="0.685585 0.173128 0.685585 -0.173128" mesh="antenna" material="antenna_material"/>
+                <!-- Part antenna_2 -->
+                <geom type="mesh" class="visual" pos="-0.12808 0.000247519 0.06095"
+                  quat="0.707107 5.42285e-16 0.707107 4.21885e-15" mesh="antenna" material="antenna_material"/>
+                <!-- Part left_antenna_holder -->
+                <geom type="mesh" class="visual" pos="-0.12906 3.26887e-17 0.06095"
+                  quat="0.707107 2.12052e-15 0.707107 4.996e-15" mesh="left_antenna_holder"
+                  material="left_antenna_holder_material"/>
+                <!-- Part cam -->
+                <!-- Part head_bot_sheet -->
+                <geom type="mesh" class="visual" pos="-0.12906 1.72388e-17 0.06095"
+                  quat="0.707107 2.48509e-16 0.707107 5.27356e-16" mesh="head_bot_sheet"
+                  material="head_bot_sheet_material"/>
+                <!-- Part head_2 -->
+                <geom type="mesh" class="visual" pos="-0.12906 1.92717e-17 0.06095"
+                  quat="0.707107 2.80428e-16 0.707107 5.82867e-16" mesh="head" material="head_material"/>
+                <!-- Part bulb -->
+                <!-- Part glass -->
+                <!-- Part head_roll_mount -->
+                <!-- Part raspberrypizerow -->
+                <!-- Part left_eye -->
+                <!-- Part full_speaker -->
+                <!-- Part roll_bearing_3 -->
+                <!-- Part sg90 -->
+                <!-- Part glass_2 -->
+                <!-- Part sg90_2 -->
+                <!-- Part bulb_2 -->
+                <!-- Part speaker_interface -->
+                <!-- Part speaker_stand -->
+                <!-- Part right_eye -->
+                <!-- Part glass_3 -->
+                <!-- Part glass_4 -->
+                <!-- Part wj_wk00_0123middlecase_56_10 -->
+                <!-- Part drive_palonier_10 -->
+                <!-- Part wj_wk00_0124bottomcase_45_10 -->
+                <!-- Part wj_wk00_0122topcabinetcase_95_10 -->
+                <!-- Part passive_palonier_10 -->
+                <!-- Frame head -->
+                <site group="0" name="head" pos="0.04245 1.88922e-17 0.03595"
+                  quat="0.707107 2.48509e-16 0.707107 5.27356e-16"/>
+              </body>
+            </body>
+          </body>
+        </body>
+        <!-- Link hip_roll_assembly_2 -->
+        <body name="hip_roll_assembly_2" pos="-0.019 -0.035 0.0459409" quat="1 -4.25414e-16 -1.07137e-14 2.77556e-17">
+          <!-- Joint from trunk_assembly to hip_roll_assembly_2 -->
+          <joint name="right_hip_yaw" type="hinge" range="-0.523598775598297 0.5235987755983006" class="sts3215"/>
+          <inertial pos="0.000795081 -5.46541e-06 -0.03306" mass="0.06648"
+            fullinertia="2.44411e-05 2.81818e-05 1.42687e-05 2.35658e-09 -3.81059e-07 -7.50036e-09"/>
+          <!-- Part roll_motor_bottom_2 -->
+          <geom type="mesh" class="visual" pos="-3.64292e-16 -0.035 0.01905" quat="1 -0 -0 -1.52656e-16"
+            mesh="roll_motor_bottom" material="roll_motor_bottom_material"/>
+          <!-- Part roll_motor_top_2 -->
+          <geom type="mesh" class="visual" pos="3.46945e-18 -0.035 0.01905" quat="1 0 2.22045e-16 -1.11022e-16"
+            mesh="roll_motor_top" material="roll_motor_top_material"/>
+          <!-- Part wj_wk00_0123middlecase_56_11 -->
+          <!-- Part drive_palonier_11 -->
+          <!-- Part wj_wk00_0124bottomcase_45_11 -->
+          <!-- Part wj_wk00_0122topcabinetcase_95_11 -->
+          <!-- Part passive_palonier_11 -->
+          <!-- Link right_roll_to_pitch_assembly -->
+          <body name="right_roll_to_pitch_assembly" pos="0.01905 4.05787e-14 -0.046015" quat="0.5 -0.5 -0.5 0.5">
+            <!-- Joint from hip_roll_assembly_2 to right_roll_to_pitch_assembly -->
+            <joint name="right_hip_roll" type="hinge" range="-0.4363323129985797 0.43633231299858505" class="sts3215"/>
+            <inertial pos="-0.0508042 -0.000420742 0.0204704" mass="0.07516"
+              fullinertia="2.52329e-05 4.13855e-05 2.81264e-05 -1.30528e-07 5.90499e-07 -9.13427e-08"/>
+            <!-- Part right_roll_to_pitch -->
+            <geom type="mesh" class="visual" pos="0.035 -0.065 0.01905" quat="0.5 0.5 0.5 -0.5"
+              mesh="right_roll_to_pitch" material="right_roll_to_pitch_material"/>
+            <!-- Part wj_wk00_0123middlecase_56_12 -->
+            <!-- Part drive_palonier_12 -->
+            <!-- Part wj_wk00_0124bottomcase_45_12 -->
+            <!-- Part wj_wk00_0122topcabinetcase_95_12 -->
+            <!-- Part passive_palonier_12 -->
+            <!-- Link knee_and_ankle_assembly_3 -->
+            <body name="knee_and_ankle_assembly_3" pos="-0.07415 2.48437e-15 0.03511"
+              quat="0.707107 6.41892e-22 0.707107 -2.22045e-16">
+              <!-- Joint from right_roll_to_pitch_assembly to knee_and_ankle_assembly_3 -->
+              <joint name="right_hip_pitch" type="hinge" range="-0.5235987755982988 1.2217304763960306" class="sts3215"/>
+              <inertial pos="0.00253369 0.0390636 0.010809" mass="0.12407"
+                fullinertia="0.000212363 7.93324e-05 0.000225619 2.68706e-05 3.9656e-06 -2.17407e-05"/>
+              <!-- Part right_cache -->
+              <geom type="mesh" class="visual" pos="0.01606 -0.065 0.1092"
+                quat="0.707107 0.707107 -3.33067e-16 -3.92523e-16" mesh="right_cache" material="right_cache_material"/>
+              <!-- Part leg_spacer_3 -->
+              <geom type="mesh" class="visual" pos="0.01606 -0.14365 -0.07215"
+                quat="0.707107 0.707107 -3.33067e-16 -3.92523e-16" mesh="leg_spacer" material="leg_spacer_material"/>
+              <!-- Part left_knee_to_ankle_right_sheet_3 -->
+              <geom type="mesh" class="visual" pos="0.01606 -0.14365 -0.07205" quat="0.707107 0.707107 0 -1.74653e-29"
+                mesh="left_knee_to_ankle_right_sheet" material="left_knee_to_ankle_right_sheet_material"/>
+              <!-- Part left_knee_to_ankle_left_sheet_3 -->
+              <geom type="mesh" class="visual" pos="0.01606 -0.14365 -0.07205" quat="0.707107 0.707107 0 -3.2666e-25"
+                mesh="left_knee_to_ankle_left_sheet" material="left_knee_to_ankle_left_sheet_material"/>
+              <!-- Part wj_wk00_0123middlecase_56_13 -->
+              <!-- Part drive_palonier_13 -->
+              <!-- Part wj_wk00_0124bottomcase_45_13 -->
+              <!-- Part wj_wk00_0122topcabinetcase_95_13 -->
+              <!-- Part passive_palonier_13 -->
+              <!-- Link knee_and_ankle_assembly_4 -->
+              <body name="knee_and_ankle_assembly_4" pos="1.38778e-16 0.07865 0.037"
+                quat="1.32068e-16 1 -1.10334e-26 3.55346e-22">
+                <!-- Joint from knee_and_ankle_assembly_3 to knee_and_ankle_assembly_4 -->
+                <joint name="right_knee" type="hinge" range="-1.5707963267948966 1.5707963267948966" class="sts3215"/>
+                <inertial pos="5.01859e-06 -0.0577465 0.0181136" mass="0.07259"
+                  fullinertia="4.99575e-05 1.87273e-05 4.23932e-05 1.03817e-08 2.25175e-09 8.81049e-08"/>
+                <!-- Part leg_spacer_4 -->
+                <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10925"
+                  quat="0.707107 -0.707107 -3.33067e-16 3.92523e-16" mesh="leg_spacer" material="leg_spacer_material"/>
+                <!-- Part left_knee_to_ankle_right_sheet_4 -->
+                <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10915" quat="0.707107 -0.707107 0 -2.51202e-22"
+                  mesh="left_knee_to_ankle_right_sheet" material="left_knee_to_ankle_right_sheet_material"/>
+                <!-- Part left_knee_to_ankle_left_sheet_4 -->
+                <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10915" quat="0.707107 -0.707107 0 5.41603e-30"
+                  mesh="left_knee_to_ankle_left_sheet" material="left_knee_to_ankle_left_sheet_material"/>
+                <!-- Part wj_wk00_0123middlecase_56_14 -->
+                <!-- Part drive_palonier_14 -->
+                <!-- Part wj_wk00_0124bottomcase_45_14 -->
+                <!-- Part wj_wk00_0122topcabinetcase_95_14 -->
+                <!-- Part passive_palonier_14 -->
+                <!-- Link foot_assembly_2 -->
+                <body name="foot_assembly_2" pos="1.32533e-15 -0.07865 0.0001"
+                  quat="1 3.88578e-16 -3.33067e-16 5.55111e-16">
+                  <!-- Joint from knee_and_ankle_assembly_4 to foot_assembly_2 -->
+                  <joint name="right_ankle" type="hinge" range="-1.5707963267948957 1.5707963267948974" class="sts3215"/>
+                  <inertial pos="0.0110718 -0.0246608 0.0190626" mass="0.07524"
+                    fullinertia="1.87482e-05 6.74427e-05 6.061e-05 1.58839e-06 -6.52354e-09 5.67391e-08"/>
+                  <!-- Part foot_side_2 -->
+                  <geom type="mesh" class="visual" pos="0.01606 0.2223 0.10905"
+                    quat="0.707107 -0.707107 -0 -8.61814e-30" mesh="foot_side" material="foot_side_material"/>
+                  <!-- Part foot_bottom_tpu_2 -->
+                  <geom type="mesh" class="visual" pos="0.01656 0.2228 0.10955"
+                    quat="0.707107 -0.707107 -0 -8.61814e-30" mesh="foot_bottom_tpu" material="foot_bottom_tpu_material"/>
+                  <geom type="mesh" class="collision" pos="0.01656 0.2228 0.10955"
+                    quat="0.707107 -0.707107 -0 -8.61814e-30" mesh="foot_bottom_tpu" material="foot_bottom_tpu_material"
+                    name="right_foot_bottom_tpu"/>
+                  <!-- Part foot_bottom_pla_2 -->
+                  <geom type="mesh" class="visual" pos="0.01656 0.2228 0.10955"
+                    quat="0.707107 -0.707107 -0 -8.61814e-30" mesh="foot_bottom_pla" material="foot_bottom_pla_material"/>
+                  <!-- Part foot_top_2 -->
+                  <geom type="mesh" class="visual" pos="0.01606 0.2223 0.10905"
+                    quat="0.707107 -0.707107 -0 -8.61814e-30" mesh="foot_top" material="foot_top_material"/>
+                  <!-- Frame right_foot -->
+                  <site group="0" name="right_foot" pos="0.0005 -0.036225 0.01955"
+                    quat="0.707107 -0.707107 -0 1.26645e-25"/>
+                </body>
+              </body>
+            </body>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+  <asset>
+    <mesh file="body_back.stl"/>
+    <mesh file="foot_bottom_pla.stl"/>
+    <mesh file="neck_left_sheet.stl"/>
+    <mesh file="right_roll_to_pitch.stl"/>
+    <mesh file="trunk_bottom.stl"/>
+    <mesh file="body_middle_bottom.stl"/>
+    <mesh file="antenna.stl"/>
+    <mesh file="foot_side.stl"/>
+    <mesh file="left_roll_to_pitch.stl"/>
+    <mesh file="foot_top.stl"/>
+    <mesh file="head.stl"/>
+    <mesh file="foot_bottom_tpu.stl"/>
+    <mesh file="neck_right_sheet.stl"/>
+    <mesh file="head_yaw_to_roll.stl"/>
+    <mesh file="roll_motor_top.stl"/>
+    <mesh file="left_knee_to_ankle_left_sheet.stl"/>
+    <mesh file="body_middle_top.stl"/>
+    <mesh file="left_antenna_holder.stl"/>
+    <mesh file="head_pitch_to_yaw.stl"/>
+    <mesh file="right_antenna_holder.stl"/>
+    <mesh file="roll_motor_bottom.stl"/>
+    <mesh file="right_cache.stl"/>
+    <mesh file="left_knee_to_ankle_right_sheet.stl"/>
+    <mesh file="body_front.stl"/>
+    <mesh file="leg_spacer.stl"/>
+    <mesh file="trunk_top.stl"/>
+    <mesh file="head_bot_sheet.stl"/>
+    <mesh file="left_cache.stl"/>
+    <material name="body_front_material" rgba="0.917647 0.917647 0.917647 1"/>
+    <material name="body_back_material" rgba="0.917647 0.917647 0.917647 1"/>
+    <material name="body_middle_bottom_material" rgba="0.917647 0.917647 0.917647 1"/>
+    <material name="body_middle_top_material" rgba="0.917647 0.917647 0.917647 1"/>
+    <material name="trunk_bottom_material" rgba="0.180392 0.180392 0.180392 1"/>
+    <material name="trunk_top_material" rgba="0.180392 0.180392 0.180392 1"/>
+    <material name="roll_motor_bottom_material" rgba="0.647059 0.647059 0.647059 1"/>
+    <material name="roll_motor_top_material" rgba="0.901961 0.901961 0.901961 1"/>
+    <material name="left_roll_to_pitch_material" rgba="0.909804 0.572549 0.164706 1"/>
+    <material name="left_cache_material" rgba="0.917647 0.917647 0.917647 1"/>
+    <material name="leg_spacer_material" rgba="0.647059 0.647059 0.647059 1"/>
+    <material name="left_knee_to_ankle_right_sheet_material" rgba="0.223529 0.219608 0.219608 1"/>
+    <material name="left_knee_to_ankle_left_sheet_material" rgba="0.223529 0.219608 0.219608 1"/>
+    <material name="foot_side_material" rgba="0.980392 0.713725 0.00392157 1"/>
+    <material name="foot_bottom_tpu_material" rgba="0.305882 0.298039 0.278431 1"/>
+    <material name="foot_bottom_pla_material" rgba="0.305882 0.298039 0.278431 1"/>
+    <material name="foot_top_material" rgba="0.980392 0.713725 0.00392157 1"/>
+    <material name="neck_left_sheet_material" rgba="0.223529 0.219608 0.219608 1"/>
+    <material name="neck_right_sheet_material" rgba="0.223529 0.219608 0.219608 1"/>
+    <material name="head_pitch_to_yaw_material" rgba="0.4 0.4 0.4 1"/>
+    <material name="head_yaw_to_roll_material" rgba="0.647059 0.647059 0.647059 1"/>
+    <material name="right_antenna_holder_material" rgba="0.980392 0.713725 0.00392157 1"/>
+    <material name="antenna_material" rgba="0.647059 0.647059 0.647059 1"/>
+    <material name="left_antenna_holder_material" rgba="0.980392 0.713725 0.00392157 1"/>
+    <material name="head_bot_sheet_material" rgba="0.411765 0.411765 0.411765 1"/>
+    <material name="head_material" rgba="0.917647 0.917647 0.917647 1"/>
+    <material name="right_roll_to_pitch_material" rgba="0.909804 0.572549 0.164706 1"/>
+    <material name="right_cache_material" rgba="0.917647 0.917647 0.917647 1"/>
+  </asset>
+  <actuator>
+    <position class="sts3215" name="left_hip_yaw" joint="left_hip_yaw" inheritrange="1"/>
+    <position class="sts3215" name="left_hip_roll" joint="left_hip_roll" inheritrange="1"/>
+    <position class="sts3215" name="left_hip_pitch" joint="left_hip_pitch" inheritrange="1"/>
+    <position class="sts3215" name="left_knee" joint="left_knee" inheritrange="1"/>
+    <position class="sts3215" name="left_ankle" joint="left_ankle" inheritrange="1"/>
+    <position class="sts3215" name="neck_pitch" joint="neck_pitch" inheritrange="1"/>
+    <position class="sts3215" name="head_pitch" joint="head_pitch" inheritrange="1"/>
+    <position class="sts3215" name="head_yaw" joint="head_yaw" inheritrange="1"/>
+    <position class="sts3215" name="head_roll" joint="head_roll" inheritrange="1"/>
+    <position class="sts3215" name="right_hip_yaw" joint="right_hip_yaw" inheritrange="1"/>
+    <position class="sts3215" name="right_hip_roll" joint="right_hip_roll" inheritrange="1"/>
+    <position class="sts3215" name="right_hip_pitch" joint="right_hip_pitch" inheritrange="1"/>
+    <position class="sts3215" name="right_knee" joint="right_knee" inheritrange="1"/>
+    <position class="sts3215" name="right_ankle" joint="right_ankle" inheritrange="1"/>
+  </actuator>
+  <equality/>
+</mujoco>

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/open_duck_mini_v2_backlash.xml

@@ -0,0 +1,514 @@
+<?xml version="1.0"?>
+<!-- Generated using onshape-to-robot -->
+<!-- Onshape document_id: 64074dfcfa379b37d8a47762 -->
+<mujoco model="open_duck_mini_v2">
+
+  <option iterations="1" ls_iterations="5">
+    <flag eulerdamp="disable"/>
+  </option>
+  <compiler angle="radian" meshdir="assets"/>
+
+  <!-- <option noslip_iterations="1"/>
+  <compiler angle="radian" meshdir="assets" autolimits="true"/> -->
+  <default>
+    <default class="open_duck_mini_v2">
+      <!-- <joint frictionloss="0.1" armature="0.005"/>
+      <position kp="50" dampratio="1"/> -->
+      <default class="visual">
+        <geom type="mesh" contype="0" conaffinity="0" group="2"/>
+      </default>
+      <default class="collision">
+        <geom group="3"/>
+      </default>
+    </default>
+  </default>
+  <!-- Additional sensors.xml -->
+  <sensor>
+    <gyro site="imu" name="gyro"/>
+    <velocimeter site="imu" name="local_linvel"/>
+    <accelerometer site="imu" name="accelerometer"/>
+    <framezaxis objtype="site" objname="imu" name="upvector"/>
+    <framexaxis objtype="site" objname="imu" name="forwardvector"/>
+    <framelinvel objtype="site" objname="imu" name="global_linvel"/>
+    <frameangvel objtype="site" objname="imu" name="global_angvel"/>
+    <framepos objtype="site" objname="imu" name="position"/>
+    <framequat objtype="site" objname="imu" name="orientation"/>
+    <framelinvel objtype="site" objname="right_foot" name="right_foot_global_linvel"/>
+    <framelinvel objtype="site" objname="left_foot" name="left_foot_global_linvel"/>
+    <framexaxis objtype="site" objname="left_foot" name="left_foot_upvector"/>
+    <framexaxis objtype="site" objname="right_foot" name="right_foot_upvector"/>
+    <framepos objtype="site" objname="left_foot" name="left_foot_pos"/>
+    <framepos objtype="site" objname="right_foot" name="right_foot_pos"/>
+  </sensor>
+  <!-- Additional joints_properties.xml -->
+  <default>
+    <default class="sts3215">
+      <geom contype="0" conaffinity="0"/>
+      <!-- <joint damping="0.60" frictionloss="0.052" armature="0.028"/>
+      <position kp="17.8" kv="0.0" forcerange="-3.35 3.35"/> -->
+      <joint damping="0.56" frictionloss="0.068" armature="0.027"/>
+      <position kp="17.11" kv="0.0" forcerange="-3.23 3.23"/>
+    </default>
+    <default class="backlash">
+      <!-- +/- 0.5° of backlash -->
+      <joint damping="0.01" frictionloss="0" armature="0.01" limited="true"
+        range="-0.008726646259971648 0.008726646259971648"/>
+    </default>
+  </default>
+  <worldbody>
+    <body name="base" pos="0 0 0.22">
+      <!-- <freejoint name="trunk_assembly_freejoint"/> -->
+      <freejoint name="floating_base"/>
+      <site name="imu" pos="-0.08 -0.0 0.05"/>
+      <!-- Link trunk_assembly -->
+      <body name="trunk_assembly" pos="0 0 0.0" quat="1 0 0 0" childclass="open_duck_mini_v2">
+        <inertial pos="-0.0483259 -9.97823e-05 0.0384971" mass="0.698526"
+          fullinertia="0.00167701 0.00344479 0.00292635 -1.31662e-05 -3.24953e-05 -2.15581e-06"/>
+        <!-- Part body_front -->
+        <geom type="mesh" class="visual" pos="-0.019 0 0.0648909" quat="1 0 0 0" mesh="body_front"
+          material="body_front_material"/>
+        <!-- Part usb_c_charger -->
+        <!-- Part body_back -->
+        <geom type="mesh" class="visual" pos="-0.019 3.46945e-18 0.0648909" quat="1 0 0 0" mesh="body_back"
+          material="body_back_material"/>
+        <!-- Part body_middle_bottom -->
+        <geom type="mesh" class="visual" pos="-0.019 -3.46945e-18 0.0648909" quat="1 0 0 0" mesh="body_middle_bottom"
+          material="body_middle_bottom_material"/>
+        <!-- Part power_switch -->
+        <!-- Part bms -->
+        <!-- Part battery_pack_lid -->
+        <!-- Part body_middle_top -->
+        <geom type="mesh" class="visual" pos="-0.019 0 0.0648909" quat="1 0 0 0" mesh="body_middle_top"
+          material="body_middle_top_material"/>
+        <!-- Part bno055 -->
+        <!-- Part roll_bearing -->
+        <!-- Part trunk_bottom -->
+        <geom type="mesh" class="visual" pos="-0.019 9.1073e-18 0.0648909" quat="1 -0 -0 5.42101e-20"
+          mesh="trunk_bottom" material="trunk_bottom_material"/>
+        <!-- Part trunk_top -->
+        <geom type="mesh" class="visual" pos="-0.019 -2.42428e-16 0.0648909"
+          quat="1 -4.82108e-17 -5.5355e-15 2.77556e-17" mesh="trunk_top" material="trunk_top_material"/>
+        <!-- Part board -->
+        <!-- Part roll_bearing_2 -->
+        <!-- Part cell -->
+        <!-- Part holder -->
+        <!-- Part cell_2 -->
+        <!-- Part wj_wk00_0123middlecase_56 -->
+        <!-- Part drive_palonier -->
+        <!-- Part wj_wk00_0124bottomcase_45 -->
+        <!-- Part wj_wk00_0122topcabinetcase_95 -->
+        <!-- Part passive_palonier -->
+        <!-- Part wj_wk00_0123middlecase_56_2 -->
+        <!-- Part drive_palonier_2 -->
+        <!-- Part wj_wk00_0124bottomcase_45_2 -->
+        <!-- Part wj_wk00_0122topcabinetcase_95_2 -->
+        <!-- Part passive_palonier_2 -->
+        <!-- Part wj_wk00_0123middlecase_56_3 -->
+        <!-- Part drive_palonier_3 -->
+        <!-- Part wj_wk00_0124bottomcase_45_3 -->
+        <!-- Part wj_wk00_0122topcabinetcase_95_3 -->
+        <!-- Part passive_palonier_3 -->
+        <!-- Frame trunk -->
+        <site group="0" name="trunk" pos="-0.024 -2.41127e-16 0.0881909" quat="1 -2.77556e-17 -0 -0"/>
+        <!-- Frame imu -->
+        <!-- <site group="0" name="imu" pos="-0.08711 1.249e-16 0.0422909" quat="1 -7.03105e-17 -2.77556e-17 -3.14419e-16"/> -->
+        <!-- Link hip_roll_assembly -->
+        <body name="hip_roll_assembly" pos="-0.019 0.035 0.0459409" quat="1 -1.17663e-16 -1.09357e-14 2.77556e-17">
+          <!-- Joint from trunk_assembly to hip_roll_assembly -->
+          <joint name="left_hip_yaw" type="hinge" range="-0.5235987755982979 0.5235987755982997" class="sts3215"/>
+          <joint name="left_hip_yaw_backlash" pos="0 0 0" axis="0 0 1" class="backlash"/>
+          <inertial pos="0.000795081 -5.46541e-06 -0.03306" mass="0.06648"
+            fullinertia="2.44411e-05 2.81818e-05 1.42687e-05 2.35658e-09 -3.81059e-07 -7.50036e-09"/>
+          <!-- Part roll_motor_bottom -->
+          <geom type="mesh" class="visual" pos="-3.33067e-16 -0.035 0.01905"
+            quat="1 -4.38499e-16 2.22045e-16 -3.34802e-16" mesh="roll_motor_bottom"
+            material="roll_motor_bottom_material"/>
+          <!-- Part roll_motor_top -->
+          <geom type="mesh" class="visual" pos="0 -0.035 0.01905" quat="1 -0 -1.80411e-16 -4.16334e-17"
+            mesh="roll_motor_top" material="roll_motor_top_material"/>
+          <!-- Part wj_wk00_0123middlecase_56_4 -->
+          <!-- Part drive_palonier_4 -->
+          <!-- Part wj_wk00_0124bottomcase_45_4 -->
+          <!-- Part wj_wk00_0122topcabinetcase_95_4 -->
+          <!-- Part passive_palonier_4 -->
+          <!-- Link left_roll_to_pitch_assembly -->
+          <body name="left_roll_to_pitch_assembly" pos="0.01905 4.61436e-15 -0.046015" quat="0.5 -0.5 -0.5 0.5">
+            <!-- Joint from hip_roll_assembly to left_roll_to_pitch_assembly -->
+            <joint name="left_hip_roll" type="hinge" range="-0.4363323129985815 0.43633231299858327" class="sts3215"/>
+            <joint name="left_hip_roll_backlash" pos="0 0 0" axis="0 0 1" class="backlash"/>
+            <inertial pos="0.0508042 -0.00041105 0.0204704" mass="0.07516"
+              fullinertia="2.52335e-05 4.13855e-05 2.8127e-05 1.22133e-07 -5.90499e-07 -1.07544e-07"/>
+            <!-- Part left_roll_to_pitch -->
+            <geom type="mesh" class="visual" pos="-0.035 -0.065 0.01905" quat="0.5 0.5 0.5 -0.5"
+              mesh="left_roll_to_pitch" material="left_roll_to_pitch_material"/>
+            <!-- Part wj_wk00_0123middlecase_56_5 -->
+            <!-- Part drive_palonier_5 -->
+            <!-- Part wj_wk00_0124bottomcase_45_5 -->
+            <!-- Part wj_wk00_0122topcabinetcase_95_5 -->
+            <!-- Part passive_palonier_5 -->
+            <!-- Link knee_and_ankle_assembly -->
+            <body name="knee_and_ankle_assembly" pos="0.07415 -3.25261e-17 0.03511"
+              quat="8.49887e-24 -0.707107 -1.11022e-16 0.707107">
+              <!-- Joint from left_roll_to_pitch_assembly to knee_and_ankle_assembly -->
+              <joint name="left_hip_pitch" type="hinge" range="-1.2217304763960306 0.5235987755982988" class="sts3215"/>
+              <joint name="left_hip_pitch_backlash" pos="0 0 0" axis="0 0 1" class="backlash"/>
+              <inertial pos="0.00253369 -0.0390636 0.0102776" mass="0.12407"
+                fullinertia="0.000211406 7.83759e-05 0.000225619 -2.68706e-05 3.81906e-06 2.08009e-05"/>
+              <!-- Part left_cache -->
+              <geom type="mesh" class="visual" pos="0.01606 0.065 0.10915"
+                quat="0.707107 -0.707107 3.33067e-16 5.76761e-16" mesh="left_cache" material="left_cache_material"/>
+              <!-- Part leg_spacer -->
+              <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10925"
+                quat="0.707107 -0.707107 3.33067e-16 5.94711e-16" mesh="leg_spacer" material="leg_spacer_material"/>
+              <!-- Part left_knee_to_ankle_right_sheet -->
+              <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10915" quat="0.707107 -0.707107 0 -5.93802e-24"
+                mesh="left_knee_to_ankle_right_sheet" material="left_knee_to_ankle_right_sheet_material"/>
+              <!-- Part left_knee_to_ankle_left_sheet -->
+              <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10915" quat="0.707107 -0.707107 0 5.12923e-30"
+                mesh="left_knee_to_ankle_left_sheet" material="left_knee_to_ankle_left_sheet_material"/>
+              <!-- Part wj_wk00_0123middlecase_56_6 -->
+              <!-- Part drive_palonier_6 -->
+              <!-- Part wj_wk00_0124bottomcase_45_6 -->
+              <!-- Part wj_wk00_0122topcabinetcase_95_6 -->
+              <!-- Part passive_palonier_6 -->
+              <!-- Link knee_and_ankle_assembly_2 -->
+              <body name="knee_and_ankle_assembly_2" pos="0 -0.07865 0.0001"
+                quat="1 1.11022e-16 1.11022e-16 -2.64698e-23">
+                <!-- Joint from knee_and_ankle_assembly to knee_and_ankle_assembly_2 -->
+                <joint name="left_knee" type="hinge" range="-1.5707963267948966 1.5707963267948966" class="sts3215"/>
+                <joint name="left_knee_backlash" pos="0 0 0" axis="0 0 1" class="backlash"/>
+                <inertial pos="5.01859e-06 -0.0577465 0.0181136" mass="0.07259"
+                  fullinertia="4.99575e-05 1.87273e-05 4.23932e-05 1.03817e-08 2.25175e-09 8.81049e-08"/>
+                <!-- Part leg_spacer_2 -->
+                <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10925"
+                  quat="0.707107 -0.707107 3.33067e-16 5.72905e-16" mesh="leg_spacer" material="leg_spacer_material"/>
+                <!-- Part left_knee_to_ankle_right_sheet_2 -->
+                <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10915" quat="0.707107 -0.707107 0 1.33119e-24"
+                  mesh="left_knee_to_ankle_right_sheet" material="left_knee_to_ankle_right_sheet_material"/>
+                <!-- Part left_knee_to_ankle_left_sheet_2 -->
+                <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10915" quat="0.707107 -0.707107 0 -5.36743e-30"
+                  mesh="left_knee_to_ankle_left_sheet" material="left_knee_to_ankle_left_sheet_material"/>
+                <!-- Part wj_wk00_0123middlecase_56_7 -->
+                <!-- Part drive_palonier_7 -->
+                <!-- Part wj_wk00_0124bottomcase_45_7 -->
+                <!-- Part wj_wk00_0122topcabinetcase_95_7 -->
+                <!-- Part passive_palonier_7 -->
+                <!-- Link foot_assembly -->
+                <body name="foot_assembly" pos="1.8735e-16 -0.07865 0.0001" quat="1 2.77556e-16 5.55112e-16 1.52656e-16">
+                  <!-- Joint from knee_and_ankle_assembly_2 to foot_assembly -->
+                  <joint name="left_ankle" type="hinge" range="-1.5707963267948957 1.5707963267948974" class="sts3215"/>
+                  <joint name="left_ankle_backlash" pos="0 0 0" axis="0 0 1" class="backlash"/>
+                  <inertial pos="0.0110718 -0.0246608 0.0190626" mass="0.07524"
+                    fullinertia="1.87482e-05 6.74427e-05 6.061e-05 1.58839e-06 -6.52354e-09 5.67391e-08"/>
+                  <!-- Part foot_side -->
+                  <geom type="mesh" class="visual" pos="0.01606 0.2223 0.10905"
+                    quat="0.707107 -0.707107 -0 -4.61924e-30" mesh="foot_side" material="foot_side_material"/>
+                  <!-- Part foot_bottom_tpu -->
+                  <geom type="mesh" class="visual" pos="0.01656 0.2228 0.10955" quat="0.707107 -0.707107 -0 1.72496e-18"
+                    mesh="foot_bottom_tpu" material="foot_bottom_tpu_material"/>
+                  <geom type="mesh" class="collision" pos="0.01656 0.2228 0.10955"
+                    quat="0.707107 -0.707107 -0 1.72496e-18" mesh="foot_bottom_tpu" material="foot_bottom_tpu_material"
+                    name="left_foot_bottom_tpu"/>
+                  <!-- Part foot_bottom_pla -->
+                  <geom type="mesh" class="visual" pos="0.01656 0.2228 0.10955" quat="0.707107 -0.707107 -0 1.72496e-18"
+                    mesh="foot_bottom_pla" material="foot_bottom_pla_material"/>
+                  <!-- Part foot_top -->
+                  <geom type="mesh" class="visual" pos="0.01606 0.2223 0.10905"
+                    quat="0.707107 -0.707107 -0 -4.61924e-30" mesh="foot_top" material="foot_top_material"/>
+                  <!-- Frame left_foot -->
+                  <site group="0" name="left_foot" pos="0.0005 -0.036225 0.01955"
+                    quat="0.707107 -0.707107 -0 1.72494e-18"/>
+                </body>
+              </body>
+            </body>
+          </body>
+        </body>
+        <!-- Link neck_pitch_assembly -->
+        <body name="neck_pitch_assembly" pos="0.001 0.01915 0.0900009" quat="0.707107 0.707107 -7.88258e-15 7.79925e-15">
+          <!-- Joint from trunk_assembly to neck_pitch_assembly -->
+          <joint name="neck_pitch" type="hinge" range="-0.3490658503988437 1.1344640137963364" class="sts3215"/>
+          <inertial pos="-5.63137e-06 0.0492968 0.0181786" mass="0.06618"
+            fullinertia="3.49456e-05 1.70231e-05 2.80043e-05 8.75633e-09 -2.15592e-09 -1.85188e-08"/>
+          <!-- Part neck_left_sheet -->
+          <geom type="mesh" class="visual" pos="-0.02 -0.02511 0.01905" quat="0.707107 -0.707107 0 -5.1423e-31"
+            mesh="neck_left_sheet" material="neck_left_sheet_material"/>
+          <!-- Part neck_right_sheet -->
+          <geom type="mesh" class="visual" pos="-0.02 -0.02511 0.01905" quat="0.707107 -0.707107 -0 3.06492e-17"
+            mesh="neck_right_sheet" material="neck_right_sheet_material"/>
+          <!-- Part wj_wk00_0123middlecase_56_8 -->
+          <!-- Part drive_palonier_8 -->
+          <!-- Part wj_wk00_0124bottomcase_45_8 -->
+          <!-- Part wj_wk00_0122topcabinetcase_95_8 -->
+          <!-- Part passive_palonier_8 -->
+          <!-- Link head_pitch_to_yaw -->
+          <body name="head_pitch_to_yaw" pos="-2.9924e-17 0.066 -6.93889e-18" quat="1 -1.7203e-16 0 1.10553e-17">
+            <!-- Joint from neck_pitch_assembly to head_pitch_to_yaw -->
+            <joint name="head_pitch" type="hinge" range="-0.7853981633974483 0.7853981633974483" class="sts3215"/>
+            <inertial pos="-0.00766247 0.026015 0.0186681" mass="0.0169378"
+              fullinertia="9.43036e-06 4.64763e-06 8.02915e-06 1.35195e-06 3.05899e-08 -8.42802e-08"/>
+            <!-- Part head_pitch_to_yaw -->
+            <geom type="mesh" class="visual" pos="-0.02 -0.09111 0.01905" quat="0.707107 -0.707107 0 1.82143e-39"
+              mesh="head_pitch_to_yaw" material="head_pitch_to_yaw_material"/>
+            <!-- Link neck_yaw_assembly -->
+            <body name="neck_yaw_assembly" pos="-4.33681e-19 0.057 0.01905"
+              quat="0.707107 -0.707107 1.11022e-16 -5.88785e-17">
+              <!-- Joint from head_pitch_to_yaw to neck_yaw_assembly -->
+              <joint name="head_yaw" type="hinge" range="-2.792526803190927 2.792526803190927" class="sts3215"/>
+              <inertial pos="0.00412907 3.95745e-06 -0.0222828" mass="0.0918099"
+                fullinertia="3.11706e-05 6.94935e-05 7.04025e-05 4.29429e-09 -4.41251e-06 -3.82028e-09"/>
+              <!-- Part head_yaw_to_roll -->
+              <geom type="mesh" class="visual" pos="-0.02 1.4985e-16 -0.14821"
+                quat="1 6.93889e-18 -5.55112e-17 7.04614e-17" mesh="head_yaw_to_roll"
+                material="head_yaw_to_roll_material"/>
+              <!-- Part wj_wk00_0123middlecase_56_9 -->
+              <!-- Part drive_palonier_9 -->
+              <!-- Part wj_wk00_0124bottomcase_45_9 -->
+              <!-- Part wj_wk00_0122topcabinetcase_95_9 -->
+              <!-- Part passive_palonier_9 -->
+              <!-- Link head_assembly -->
+              <body name="head_assembly" pos="0.04095 2.1764e-16 -0.01915"
+                quat="0.707107 2.89132e-14 -0.707107 -2.86438e-14">
+                <!-- Joint from neck_yaw_assembly to head_assembly -->
+                <joint name="head_roll" type="hinge" range="-0.523598775598218 0.5235987755983796" class="sts3215"/>
+                <inertial pos="0.00815416 -0.00390754 0.0227726" mass="0.406607"
+                  fullinertia="0.00244627 0.0016864 0.00108002 -2.02404e-06 9.2382e-05 -1.84944e-05"/>
+                <!-- Part right_antenna_holder -->
+                <geom type="mesh" class="visual" pos="-0.12906 -2.4549e-16 0.06095"
+                  quat="0.707107 5.22659e-16 0.707107 3.88578e-15" mesh="right_antenna_holder"
+                  material="right_antenna_holder_material"/>
+                <!-- Part antenna -->
+                <geom type="mesh" class="visual" pos="-0.0830595 0.182212 0.06095"
+                  quat="0.685585 0.173128 0.685585 -0.173128" mesh="antenna" material="antenna_material"/>
+                <!-- Part antenna_2 -->
+                <geom type="mesh" class="visual" pos="-0.12808 0.000247519 0.06095"
+                  quat="0.707107 5.42285e-16 0.707107 4.21885e-15" mesh="antenna" material="antenna_material"/>
+                <!-- Part left_antenna_holder -->
+                <geom type="mesh" class="visual" pos="-0.12906 3.26887e-17 0.06095"
+                  quat="0.707107 2.12052e-15 0.707107 4.996e-15" mesh="left_antenna_holder"
+                  material="left_antenna_holder_material"/>
+                <!-- Part cam -->
+                <!-- Part head_bot_sheet -->
+                <geom type="mesh" class="visual" pos="-0.12906 1.72388e-17 0.06095"
+                  quat="0.707107 2.48509e-16 0.707107 5.27356e-16" mesh="head_bot_sheet"
+                  material="head_bot_sheet_material"/>
+                <!-- Part head_2 -->
+                <geom type="mesh" class="visual" pos="-0.12906 1.92717e-17 0.06095"
+                  quat="0.707107 2.80428e-16 0.707107 5.82867e-16" mesh="head" material="head_material"/>
+                <!-- Part bulb -->
+                <!-- Part glass -->
+                <!-- Part head_roll_mount -->
+                <!-- Part raspberrypizerow -->
+                <!-- Part left_eye -->
+                <!-- Part full_speaker -->
+                <!-- Part roll_bearing_3 -->
+                <!-- Part sg90 -->
+                <!-- Part glass_2 -->
+                <!-- Part sg90_2 -->
+                <!-- Part bulb_2 -->
+                <!-- Part speaker_interface -->
+                <!-- Part speaker_stand -->
+                <!-- Part right_eye -->
+                <!-- Part glass_3 -->
+                <!-- Part glass_4 -->
+                <!-- Part wj_wk00_0123middlecase_56_10 -->
+                <!-- Part drive_palonier_10 -->
+                <!-- Part wj_wk00_0124bottomcase_45_10 -->
+                <!-- Part wj_wk00_0122topcabinetcase_95_10 -->
+                <!-- Part passive_palonier_10 -->
+                <!-- Frame head -->
+                <site group="0" name="head" pos="0.04245 1.88922e-17 0.03595"
+                  quat="0.707107 2.48509e-16 0.707107 5.27356e-16"/>
+              </body>
+            </body>
+          </body>
+        </body>
+        <!-- Link hip_roll_assembly_2 -->
+        <body name="hip_roll_assembly_2" pos="-0.019 -0.035 0.0459409" quat="1 -4.25414e-16 -1.07137e-14 2.77556e-17">
+          <!-- Joint from trunk_assembly to hip_roll_assembly_2 -->
+          <joint name="right_hip_yaw" type="hinge" range="-0.523598775598297 0.5235987755983006" class="sts3215"/>
+          <joint name="right_hip_yaw_backlash" pos="0 0 0" axis="0 0 1" class="backlash"/>
+          <inertial pos="0.000795081 -5.46541e-06 -0.03306" mass="0.06648"
+            fullinertia="2.44411e-05 2.81818e-05 1.42687e-05 2.35658e-09 -3.81059e-07 -7.50036e-09"/>
+          <!-- Part roll_motor_bottom_2 -->
+          <geom type="mesh" class="visual" pos="-3.64292e-16 -0.035 0.01905" quat="1 -0 -0 -1.52656e-16"
+            mesh="roll_motor_bottom" material="roll_motor_bottom_material"/>
+          <!-- Part roll_motor_top_2 -->
+          <geom type="mesh" class="visual" pos="3.46945e-18 -0.035 0.01905" quat="1 0 2.22045e-16 -1.11022e-16"
+            mesh="roll_motor_top" material="roll_motor_top_material"/>
+          <!-- Part wj_wk00_0123middlecase_56_11 -->
+          <!-- Part drive_palonier_11 -->
+          <!-- Part wj_wk00_0124bottomcase_45_11 -->
+          <!-- Part wj_wk00_0122topcabinetcase_95_11 -->
+          <!-- Part passive_palonier_11 -->
+          <!-- Link right_roll_to_pitch_assembly -->
+          <body name="right_roll_to_pitch_assembly" pos="0.01905 4.05787e-14 -0.046015" quat="0.5 -0.5 -0.5 0.5">
+            <!-- Joint from hip_roll_assembly_2 to right_roll_to_pitch_assembly -->
+            <joint name="right_hip_roll" type="hinge" range="-0.4363323129985797 0.43633231299858505" class="sts3215"/>
+            <joint name="right_hip_roll_backlash" pos="0 0 0" axis="0 0 1" class="backlash"/>
+            <inertial pos="-0.0508042 -0.000420742 0.0204704" mass="0.07516"
+              fullinertia="2.52329e-05 4.13855e-05 2.81264e-05 -1.30528e-07 5.90499e-07 -9.13427e-08"/>
+            <!-- Part right_roll_to_pitch -->
+            <geom type="mesh" class="visual" pos="0.035 -0.065 0.01905" quat="0.5 0.5 0.5 -0.5"
+              mesh="right_roll_to_pitch" material="right_roll_to_pitch_material"/>
+            <!-- Part wj_wk00_0123middlecase_56_12 -->
+            <!-- Part drive_palonier_12 -->
+            <!-- Part wj_wk00_0124bottomcase_45_12 -->
+            <!-- Part wj_wk00_0122topcabinetcase_95_12 -->
+            <!-- Part passive_palonier_12 -->
+            <!-- Link knee_and_ankle_assembly_3 -->
+            <body name="knee_and_ankle_assembly_3" pos="-0.07415 2.48437e-15 0.03511"
+              quat="0.707107 6.41892e-22 0.707107 -2.22045e-16">
+              <!-- Joint from right_roll_to_pitch_assembly to knee_and_ankle_assembly_3 -->
+              <joint name="right_hip_pitch" type="hinge" range="-0.5235987755982988 1.2217304763960306" class="sts3215"/>
+              <joint name="right_hip_pitch_backlash" pos="0 0 0" axis="0 0 1" class="backlash"/>
+              <inertial pos="0.00253369 0.0390636 0.010809" mass="0.12407"
+                fullinertia="0.000212363 7.93324e-05 0.000225619 2.68706e-05 3.9656e-06 -2.17407e-05"/>
+              <!-- Part right_cache -->
+              <geom type="mesh" class="visual" pos="0.01606 -0.065 0.1092"
+                quat="0.707107 0.707107 -3.33067e-16 -3.92523e-16" mesh="right_cache" material="right_cache_material"/>
+              <!-- Part leg_spacer_3 -->
+              <geom type="mesh" class="visual" pos="0.01606 -0.14365 -0.07215"
+                quat="0.707107 0.707107 -3.33067e-16 -3.92523e-16" mesh="leg_spacer" material="leg_spacer_material"/>
+              <!-- Part left_knee_to_ankle_right_sheet_3 -->
+              <geom type="mesh" class="visual" pos="0.01606 -0.14365 -0.07205" quat="0.707107 0.707107 0 -1.74653e-29"
+                mesh="left_knee_to_ankle_right_sheet" material="left_knee_to_ankle_right_sheet_material"/>
+              <!-- Part left_knee_to_ankle_left_sheet_3 -->
+              <geom type="mesh" class="visual" pos="0.01606 -0.14365 -0.07205" quat="0.707107 0.707107 0 -3.2666e-25"
+                mesh="left_knee_to_ankle_left_sheet" material="left_knee_to_ankle_left_sheet_material"/>
+              <!-- Part wj_wk00_0123middlecase_56_13 -->
+              <!-- Part drive_palonier_13 -->
+              <!-- Part wj_wk00_0124bottomcase_45_13 -->
+              <!-- Part wj_wk00_0122topcabinetcase_95_13 -->
+              <!-- Part passive_palonier_13 -->
+              <!-- Link knee_and_ankle_assembly_4 -->
+              <body name="knee_and_ankle_assembly_4" pos="1.38778e-16 0.07865 0.037"
+                quat="1.32068e-16 1 -1.10334e-26 3.55346e-22">
+                <!-- Joint from knee_and_ankle_assembly_3 to knee_and_ankle_assembly_4 -->
+                <joint name="right_knee" type="hinge" range="-1.5707963267948966 1.5707963267948966" class="sts3215"/>
+                <joint name="right_knee_backlash" pos="0 0 0" axis="0 0 1" class="backlash"/>
+                <inertial pos="5.01859e-06 -0.0577465 0.0181136" mass="0.07259"
+                  fullinertia="4.99575e-05 1.87273e-05 4.23932e-05 1.03817e-08 2.25175e-09 8.81049e-08"/>
+                <!-- Part leg_spacer_4 -->
+                <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10925"
+                  quat="0.707107 -0.707107 -3.33067e-16 3.92523e-16" mesh="leg_spacer" material="leg_spacer_material"/>
+                <!-- Part left_knee_to_ankle_right_sheet_4 -->
+                <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10915" quat="0.707107 -0.707107 0 -2.51202e-22"
+                  mesh="left_knee_to_ankle_right_sheet" material="left_knee_to_ankle_right_sheet_material"/>
+                <!-- Part left_knee_to_ankle_left_sheet_4 -->
+                <geom type="mesh" class="visual" pos="0.01606 0.14365 0.10915" quat="0.707107 -0.707107 0 5.41603e-30"
+                  mesh="left_knee_to_ankle_left_sheet" material="left_knee_to_ankle_left_sheet_material"/>
+                <!-- Part wj_wk00_0123middlecase_56_14 -->
+                <!-- Part drive_palonier_14 -->
+                <!-- Part wj_wk00_0124bottomcase_45_14 -->
+                <!-- Part wj_wk00_0122topcabinetcase_95_14 -->
+                <!-- Part passive_palonier_14 -->
+                <!-- Link foot_assembly_2 -->
+                <body name="foot_assembly_2" pos="1.32533e-15 -0.07865 0.0001"
+                  quat="1 3.88578e-16 -3.33067e-16 5.55111e-16">
+                  <!-- Joint from knee_and_ankle_assembly_4 to foot_assembly_2 -->
+                  <joint name="right_ankle" type="hinge" range="-1.5707963267948957 1.5707963267948974" class="sts3215"/>
+                  <joint name="right_ankle_backlash" pos="0 0 0" axis="0 0 1" class="backlash"/>
+                  <inertial pos="0.0110718 -0.0246608 0.0190626" mass="0.07524"
+                    fullinertia="1.87482e-05 6.74427e-05 6.061e-05 1.58839e-06 -6.52354e-09 5.67391e-08"/>
+                  <!-- Part foot_side_2 -->
+                  <geom type="mesh" class="visual" pos="0.01606 0.2223 0.10905"
+                    quat="0.707107 -0.707107 -0 -8.61814e-30" mesh="foot_side" material="foot_side_material"/>
+                  <!-- Part foot_bottom_tpu_2 -->
+                  <geom type="mesh" class="visual" pos="0.01656 0.2228 0.10955"
+                    quat="0.707107 -0.707107 -0 -8.61814e-30" mesh="foot_bottom_tpu" material="foot_bottom_tpu_material"/>
+                  <geom type="mesh" class="collision" pos="0.01656 0.2228 0.10955"
+                    quat="0.707107 -0.707107 -0 -8.61814e-30" mesh="foot_bottom_tpu" material="foot_bottom_tpu_material"
+                    name="right_foot_bottom_tpu"/>
+                  <!-- Part foot_bottom_pla_2 -->
+                  <geom type="mesh" class="visual" pos="0.01656 0.2228 0.10955"
+                    quat="0.707107 -0.707107 -0 -8.61814e-30" mesh="foot_bottom_pla" material="foot_bottom_pla_material"/>
+                  <!-- Part foot_top_2 -->
+                  <geom type="mesh" class="visual" pos="0.01606 0.2223 0.10905"
+                    quat="0.707107 -0.707107 -0 -8.61814e-30" mesh="foot_top" material="foot_top_material"/>
+                  <!-- Frame right_foot -->
+                  <site group="0" name="right_foot" pos="0.0005 -0.036225 0.01955"
+                    quat="0.707107 -0.707107 -0 1.26645e-25"/>
+                </body>
+              </body>
+            </body>
+          </body>
+        </body>
+      </body>
+    </body>
+  </worldbody>
+  <asset>
+    <mesh file="body_back.stl"/>
+    <mesh file="foot_bottom_pla.stl"/>
+    <mesh file="neck_left_sheet.stl"/>
+    <mesh file="right_roll_to_pitch.stl"/>
+    <mesh file="trunk_bottom.stl"/>
+    <mesh file="body_middle_bottom.stl"/>
+    <mesh file="antenna.stl"/>
+    <mesh file="foot_side.stl"/>
+    <mesh file="left_roll_to_pitch.stl"/>
+    <mesh file="foot_top.stl"/>
+    <mesh file="head.stl"/>
+    <mesh file="foot_bottom_tpu.stl"/>
+    <mesh file="neck_right_sheet.stl"/>
+    <mesh file="head_yaw_to_roll.stl"/>
+    <mesh file="roll_motor_top.stl"/>
+    <mesh file="left_knee_to_ankle_left_sheet.stl"/>
+    <mesh file="body_middle_top.stl"/>
+    <mesh file="left_antenna_holder.stl"/>
+    <mesh file="head_pitch_to_yaw.stl"/>
+    <mesh file="right_antenna_holder.stl"/>
+    <mesh file="roll_motor_bottom.stl"/>
+    <mesh file="right_cache.stl"/>
+    <mesh file="left_knee_to_ankle_right_sheet.stl"/>
+    <mesh file="body_front.stl"/>
+    <mesh file="leg_spacer.stl"/>
+    <mesh file="trunk_top.stl"/>
+    <mesh file="head_bot_sheet.stl"/>
+    <mesh file="left_cache.stl"/>
+    <material name="body_front_material" rgba="0.917647 0.917647 0.917647 1"/>
+    <material name="body_back_material" rgba="0.917647 0.917647 0.917647 1"/>
+    <material name="body_middle_bottom_material" rgba="0.917647 0.917647 0.917647 1"/>
+    <material name="body_middle_top_material" rgba="0.917647 0.917647 0.917647 1"/>
+    <material name="trunk_bottom_material" rgba="0.180392 0.180392 0.180392 1"/>
+    <material name="trunk_top_material" rgba="0.180392 0.180392 0.180392 1"/>
+    <material name="roll_motor_bottom_material" rgba="0.647059 0.647059 0.647059 1"/>
+    <material name="roll_motor_top_material" rgba="0.901961 0.901961 0.901961 1"/>
+    <material name="left_roll_to_pitch_material" rgba="0.909804 0.572549 0.164706 1"/>
+    <material name="left_cache_material" rgba="0.917647 0.917647 0.917647 1"/>
+    <material name="leg_spacer_material" rgba="0.647059 0.647059 0.647059 1"/>
+    <material name="left_knee_to_ankle_right_sheet_material" rgba="0.223529 0.219608 0.219608 1"/>
+    <material name="left_knee_to_ankle_left_sheet_material" rgba="0.223529 0.219608 0.219608 1"/>
+    <material name="foot_side_material" rgba="0.980392 0.713725 0.00392157 1"/>
+    <material name="foot_bottom_tpu_material" rgba="0.305882 0.298039 0.278431 1"/>
+    <material name="foot_bottom_pla_material" rgba="0.305882 0.298039 0.278431 1"/>
+    <material name="foot_top_material" rgba="0.980392 0.713725 0.00392157 1"/>
+    <material name="neck_left_sheet_material" rgba="0.223529 0.219608 0.219608 1"/>
+    <material name="neck_right_sheet_material" rgba="0.223529 0.219608 0.219608 1"/>
+    <material name="head_pitch_to_yaw_material" rgba="0.4 0.4 0.4 1"/>
+    <material name="head_yaw_to_roll_material" rgba="0.647059 0.647059 0.647059 1"/>
+    <material name="right_antenna_holder_material" rgba="0.980392 0.713725 0.00392157 1"/>
+    <material name="antenna_material" rgba="0.647059 0.647059 0.647059 1"/>
+    <material name="left_antenna_holder_material" rgba="0.980392 0.713725 0.00392157 1"/>
+    <material name="head_bot_sheet_material" rgba="0.411765 0.411765 0.411765 1"/>
+    <material name="head_material" rgba="0.917647 0.917647 0.917647 1"/>
+    <material name="right_roll_to_pitch_material" rgba="0.909804 0.572549 0.164706 1"/>
+    <material name="right_cache_material" rgba="0.917647 0.917647 0.917647 1"/>
+  </asset>
+  <actuator>
+    <position class="sts3215" name="left_hip_yaw" joint="left_hip_yaw" inheritrange="1"/>
+    <position class="sts3215" name="left_hip_roll" joint="left_hip_roll" inheritrange="1"/>
+    <position class="sts3215" name="left_hip_pitch" joint="left_hip_pitch" inheritrange="1"/>
+    <position class="sts3215" name="left_knee" joint="left_knee" inheritrange="1"/>
+    <position class="sts3215" name="left_ankle" joint="left_ankle" inheritrange="1"/>
+    <position class="sts3215" name="neck_pitch" joint="neck_pitch" inheritrange="1"/>
+    <position class="sts3215" name="head_pitch" joint="head_pitch" inheritrange="1"/>
+    <position class="sts3215" name="head_yaw" joint="head_yaw" inheritrange="1"/>
+    <position class="sts3215" name="head_roll" joint="head_roll" inheritrange="1"/>
+    <position class="sts3215" name="right_hip_yaw" joint="right_hip_yaw" inheritrange="1"/>
+    <position class="sts3215" name="right_hip_roll" joint="right_hip_roll" inheritrange="1"/>
+    <position class="sts3215" name="right_hip_pitch" joint="right_hip_pitch" inheritrange="1"/>
+    <position class="sts3215" name="right_knee" joint="right_knee" inheritrange="1"/>
+    <position class="sts3215" name="right_ankle" joint="right_ankle" inheritrange="1"/>
+  </actuator>
+  <equality/>
+</mujoco>

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/scene_flat_terrain.xml

@@ -0,0 +1,76 @@
+<mujoco model="scene">
+    <include file="open_duck_mini_v2.xml"/>
+
+    <visual>
+        <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0"/>
+        <rgba haze="0.15 0.25 0.35 1"/>
+        <global azimuth="160" elevation="-20"/>
+    </visual>
+    <!-- 
+    <asset>
+        <texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512" height="3072"/>
+        <texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3"
+            markrgb="0.8 0.8 0.8" width="300" height="300"/>
+        <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
+    </asset> -->
+
+    <!-- <worldbody>
+        <body name="floor">
+            <light pos="0 0 3.5" dir="0 0 -1" directional="true"/>
+            <geom name="floor" size="0 0 0.05" pos="0 0 0" type="plane" material="groundplane" contype="1"
+                conaffinity="0" priority="1"/>
+        </body>
+    </worldbody> -->
+
+
+    <asset>
+        <texture type="skybox" builtin="gradient" rgb1="1 1 1" rgb2="1 1 1" width="800" height="800"/>
+        <texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="1 1 1" rgb2="1 1 1" markrgb="0 0 0"
+            width="300" height="300"/>
+        <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0"/>
+    </asset>
+
+    <worldbody>
+        <body name="floor">
+            <geom name="floor" size="0 0 0.01" type="plane" material="groundplane" contype="1" conaffinity="0"
+                priority="1" friction="0.6" condim="3"/>
+        </body>
+    </worldbody>
+    <keyframe>
+        <key name="home"
+            qpos="
+    0 0 0.15
+    1 0 0 0
+      0.002
+        0.053
+        -0.63
+        1.368
+        -0.784
+        0
+        0
+        0
+        0
+        -0.003
+        -0.065
+        0.635
+        1.379
+        -0.796
+"
+            ctrl="
+          0.002
+        0.053
+        -0.63
+        1.368
+        -0.784
+        0
+        0
+        0
+        0
+        -0.003
+        -0.065
+        0.635
+        1.379
+        -0.796
+          "/>
+    </keyframe>
+</mujoco>

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/scene_flat_terrain_backlash.xml

@@ -0,0 +1,79 @@
+<mujoco model="scene">
+    <include file="open_duck_mini_v2_backlash.xml"/>
+
+    <visual>
+        <headlight diffuse="0.6 0.6 0.6" ambient="0.3 0.3 0.3" specular="0 0 0"/>
+        <rgba haze="0.15 0.25 0.35 1"/>
+        <global azimuth="160" elevation="-20"/>
+    </visual>
+    <!-- 
+    <asset>
+        <texture type="skybox" builtin="gradient" rgb1="0.3 0.5 0.7" rgb2="0 0 0" width="512" height="3072"/>
+        <texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="0.2 0.3 0.4" rgb2="0.1 0.2 0.3"
+            markrgb="0.8 0.8 0.8" width="300" height="300"/>
+        <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0.2"/>
+    </asset> -->
+
+    <!-- <worldbody>
+        <body name="floor">
+            <light pos="0 0 3.5" dir="0 0 -1" directional="true"/>
+            <geom name="floor" size="0 0 0.05" pos="0 0 0" type="plane" material="groundplane" contype="1"
+                conaffinity="0" priority="1"/>
+        </body>
+    </worldbody> -->
+
+
+    <asset>
+        <texture type="skybox" builtin="gradient" rgb1="1 1 1" rgb2="1 1 1" width="800" height="800"/>
+        <texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="1 1 1" rgb2="1 1 1" markrgb="0 0 0"
+            width="300" height="300"/>
+        <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance="0"/>
+    </asset>
+
+    <worldbody>
+        <body name="floor">
+            <geom name="floor" size="0 0 0.01" type="plane" material="groundplane" contype="1" conaffinity="0"
+                priority="1" friction="0.6" condim="3"/>
+        </body>
+    </worldbody>
+    <keyframe>
+        <key name="home"
+            qpos="
+        0 0 0.15
+        1 0 0 0
+
+        0.002 0
+        0.053 0
+        -0.63 0
+        1.368 0
+        -0.784 0
+
+        0
+        0
+        0
+        0
+
+        -0.003 0
+        -0.065 0
+        0.635 0
+        1.379 0
+        -0.796 0
+"
+            ctrl="
+          0.002
+        0.053
+        -0.63
+        1.368
+        -0.784
+        0
+        0
+        0
+        0
+        -0.003
+        -0.065
+        0.635
+        1.379
+        -0.796
+          "/>
+    </keyframe>
+</mujoco>

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/scene_rough_terrain_backlash.xml

@@ -0,0 +1,75 @@
+<mujoco model="Open Duck Mini V2 rough terrain scene">
+  <!-- <include file="open_duck_mini_v2_no_head.xml"/> -->
+  <include file="open_duck_mini_v2_backlash.xml"/>
+
+  <statistic center="0 0 0.1" extent="0.8" meansize="0.04"/>
+
+  <visual>
+    <rgba force="1 0 0 1"/>
+    <global azimuth="120" elevation="-20"/>
+    <map force="0.01"/>
+    <scale forcewidth="0.3" contactwidth="0.5" contactheight="0.2"/>
+    <quality shadowsize="8192"/>
+  </visual>
+
+  <asset>
+    <!-- https://polyhaven.com/a/rock_face -->
+    <texture type="skybox" builtin="gradient" rgb1="1 1 1" rgb2="1 1 1" width="800" height="800"/>
+        <texture type="2d" name="groundplane" builtin="checker" mark="edge" rgb1="1 1 1" rgb2="1 1 1" markrgb="0 0 0"
+      width="300" height="300"/>
+    <!-- <texture type="2d" name="groundplane" file="assets/rocky_texture.png"/> -->
+    <material name="groundplane" texture="groundplane" texuniform="true" texrepeat="5 5" reflectance=".8"/>
+    <hfield name="hfield" file="assets/hfield.png" size="10 10 .01 0.1"/>
+  </asset>
+
+  <worldbody>
+    <body name="floor">
+      <geom name="floor" type="hfield" hfield="hfield" material="groundplane" contype="1" conaffinity="0" priority="1"
+        friction="1.0"/>
+    </body>
+  </worldbody>
+
+  <keyframe>
+    <key name="home"
+      qpos="
+        0 0 0.15
+        1 0 0 0
+
+        0.002 0
+        0.053 0
+        -0.63 0
+        1.368 0
+        -0.784 0
+
+        0
+        0
+        0
+        0
+
+        -0.003 0
+        -0.065 0
+        0.635 0
+        1.379 0
+        -0.796 0
+
+      "
+      ctrl="
+          0.002
+          0.053
+          -0.63
+          1.368
+          -0.784
+
+          0
+          0
+          0
+          0
+          
+          -0.003
+          -0.065
+          0.635
+          1.379
+          -0.796
+      "/>
+  </keyframe>
+</mujoco>

Open_Duck_Playground/playground/open_duck_mini_v2/xmls/sensors.xml

@@ -0,0 +1,17 @@
+<sensor>
+    <gyro site="imu" name="gyro"/>
+    <velocimeter site="imu" name="local_linvel"/>
+    <accelerometer site="imu" name="accelerometer"/>
+    <framezaxis objtype="site" objname="imu" name="upvector"/>
+    <framexaxis objtype="site" objname="imu" name="forwardvector"/>
+    <framelinvel objtype="site" objname="imu" name="global_linvel"/>
+    <frameangvel objtype="site" objname="imu" name="global_angvel"/>
+    <framepos objtype="site" objname="imu" name="position"/>
+    <framequat objtype="site" objname="imu" name="orientation"/>
+    <framelinvel objtype="site" objname="right_foot" name="right_foot_global_linvel"/>
+    <framelinvel objtype="site" objname="left_foot" name="left_foot_global_linvel"/>
+    <framexaxis objtype="site" objname="left_foot" name="left_foot_upvector"/>
+    <framexaxis objtype="site" objname="right_foot" name="right_foot_upvector"/>
+    <framepos objtype="site" objname="left_foot" name="left_foot_pos"/>
+    <framepos objtype="site" objname="right_foot" name="right_foot_pos"/>
+</sensor>

Open_Duck_Playground/pyproject.toml

@@ -0,0 +1,23 @@
+[project]
+name = "Open_Duck_Playground"
+version = "0.1.0"
+description = "Add your description here"
+readme = "README.md"
+requires-python = ">=3.11"
+dependencies = [
+    "framesviewer>=1.0.2",
+    "jax[cuda12]>=0.5.0",
+    "jaxlib>=0.5.0",
+    "jaxtyping>=0.2.38",
+    "matplotlib>=3.10.0",
+    "mediapy>=1.2.2",
+    "onnxruntime>=1.20.1",
+    "playground>=0.0.3",
+    "pygame>=2.6.1",
+    "tensorflow>=2.18.0",
+    "tf2onnx>=1.16.1",
+]
+[build-system]
+requires = ["setuptools"]
+[tool.setuptools]
+packages = ["playground"]