Upload folder using huggingface_hub

bundle.yaml

@@ -0,0 +1,107 @@
+bundle_format: augmpc_model_bundle_v1
+bundle_name: d2026_03_07_h19_m22_s30-CentauroCloopPartialNoYawUb_FakePosTrackingEnv
+checkpoint_file: d2026_03_07_h19_m22_s30-CentauroCloopPartialNoYawUb_FakePosTrackingEnv_model
+preserved_training_cfgs:
+  - ibrido_run__2026_03_07_19_20_05_ID/training_cfg_centauro_no_yaw_ub_cloop.sh
+framework:
+  repos:
+    AugMPC:
+      commit: f2ff28085ea76d2b548841de6f363f7183480f86
+      branch: ibrido
+      remote: git@github.com:AndrePatri/AugMPC.git
+      dirty: false
+    AugMPCEnvs:
+      commit: 46c5d4161cb4b249b3a2e50c93c6bc2aa087f027
+      branch: ibrido
+      remote: git@github.com:AndrePatri/AugMPCEnvs.git
+      dirty: false
+    AugMPCModels:
+      commit: 8b15c800c0f5684c61fdaf4847156ff71df61ebc
+      branch: main
+      remote: https://huggingface.co/AndrePatri/AugMPCModels
+      dirty: true
+    CentauroHybridMPC:
+      commit: 640889d4c3b9c8d360b5a3ccde6fc2bd8f139891
+      branch: ibrido
+      remote: git@github.com:ADVRHumanoids/CentauroHybridMPC.git
+      dirty: false
+    EigenIPC:
+      commit: 7c1c3ecd08716e61ed164a6e0bb788aa716705ca
+      branch: devel
+      remote: git@github.com:AndrePatri/EigenIPC.git
+      dirty: false
+    IBRIDO:
+      commit: 0ceb5f3e0508b6ecdce12bcc8f0dcbcde8f29a93
+      branch: main
+      remote: git@github.com:AndrePatri/IBRIDO.git
+      dirty: false
+    IsaacLab:
+      commit: a520a883ce996d855cc9d5255d71fd1c1307633f
+      branch: HEAD
+      remote: git@github.com:isaac-sim/IsaacLab.git
+      dirty: true
+    KyonRLStepping:
+      commit: 2bea2b8d70078974869df0549e90fc27ff31f851
+      branch: ibrido
+      remote: git@github.com:ADVRHumanoids/KyonRLStepping.git
+      dirty: false
+    MPCHive:
+      commit: 45b4fc692850cef607020dda2e32fd708e7fff62
+      branch: devel
+      remote: git@github.com:AndrePatri/MPCHive.git
+      dirty: false
+    MPCViz:
+      commit: 806d03efcc9d8ab1fb04991a159c19ba89bfb85b
+      branch: ros2_humble
+      remote: git@github.com:AndrePatri/MPCViz.git
+      dirty: false
+    adarl:
+      commit: f585499e49fa05fdd070a77f3211c0996599b87b
+      branch: ibrido
+      remote: git@github.com:c-rizz/adarl.git
+      dirty: false
+    casadi:
+      commit: 79cebe3b416dee22abc87de0076b80a5b97bd345
+      branch: optional_float
+      remote: git@github.com:AndrePatri/casadi.git
+      dirty: true
+    horizon:
+      commit: 3b084317709ff9b88d4a07ac5436f5988b39eece
+      branch: ibrido
+      remote: git@github.com:AndrePatri/horizon.git
+      dirty: true
+    iit-centauro-ros-pkg:
+      commit: 6069807ebb37a6d7df39430a02685e09ed9b167a
+      branch: ibrido_ros2
+      remote: git@github.com:AndrePatri/iit-centauro-ros-pkg.git
+      dirty: false
+    iit-dagana-ros-pkg:
+      commit: f7ecd6c84a0b7f3320c1b7de6449cbcd4445d2fe
+      branch: ibrido_ros2
+      remote: git@github.com:AndrePatri/iit-dagana-ros-pkg.git
+      dirty: false
+    iit-kyon-description:
+      commit: 50fd7c8909b3ddfd1ebbe67c61d6b775b86df6b1
+      branch: ibrido_ros2
+      remote: git@github.com:ADVRHumanoids/iit-kyon-ros-pkg.git
+      dirty: false
+    iit-kyon-description-mpc:
+      commit: 3a92bee28405172e8f6c628b4498703724d35bf5
+      branch: ibrido_ros2
+      remote: git@github.com:ADVRHumanoids/iit-kyon-ros-pkg.git
+      dirty: false
+    iit-kyon-ros-pkg:
+      commit: 298917efffb63dbca540e0aedbd21b41bf393fbf
+      branch: ibrido_ros2_simple
+      remote: git@github.com:ADVRHumanoids/iit-kyon-ros-pkg.git
+      dirty: false
+    phase_manager:
+      commit: 9925d18c0d7a55d013cbaa4998c61d85a3a8944f
+      branch: ibrido
+      remote: git@github.com:AndrePatri/phase_manager.git
+      dirty: false
+    unitree_ros:
+      commit: c75a622b8782d11dd6aa4c2ebd3b0f9c13a56aae
+      branch: ibrido
+      remote: git@github.com:AndrePatri/unitree_ros.git
+      dirty: true

centauro_big_wheels_no_yaw_ub_2026_03_07_19_20_05_ID.srdf

@@ -0,0 +1,147 @@
+<?xml version="1.0" ?>
+<!-- =================================================================================== -->
+<!-- |    This document was autogenerated by xacro from /root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_srdf/srdf/centauro.srdf.xacro | -->
+<!-- |    EDITING THIS FILE BY HAND IS NOT RECOMMENDED                                 | -->
+<!-- =================================================================================== -->
+<!--This does not replace URDF, and is not an extension of URDF.
+    This is a format for representing semantic information about the robot structure.
+    A URDF file must exist for this robot as well, where the joints and the links that are referenced are defined
+-->
+<robot name="centauro">
+  <!-- none, ball, dagana_fixed, dagana -->
+  <group name="base">
+    <link name="pelvis"/>
+  </group>
+  <group name="imu_sensors">
+    <link name="imu_link"/>
+  </group>
+  <group name="front_left_leg">
+    <chain base_link="pelvis" tip_link="wheel_1"/>
+  </group>
+  <group name="front_right_leg">
+    <chain base_link="pelvis" tip_link="wheel_2"/>
+  </group>
+  <group name="rear_right_leg">
+    <chain base_link="pelvis" tip_link="wheel_4"/>
+  </group>
+  <group name="rear_left_leg">
+    <chain base_link="pelvis" tip_link="wheel_3"/>
+  </group>
+  <group name="left_arm">
+    <chain base_link="torso_2" tip_link="arm1_8"/>
+  </group>
+  <group name="right_arm">
+    <chain base_link="torso_2" tip_link="arm2_8"/>
+  </group>
+  <group name="torso">
+    <chain base_link="pelvis" tip_link="torso_2"/>
+  </group>
+  <group name="velodyne">
+    <chain base_link="torso_2" tip_link="velodyne_motor"/>
+  </group>
+  <group name="d435_head">
+    <chain base_link="torso_2" tip_link="d435_head_motor"/>
+  </group>
+  <group name="chains">
+    <group name="front_left_leg"/>
+    <group name="front_right_leg"/>
+    <group name="rear_right_leg"/>
+    <group name="rear_left_leg"/>
+    <group name="left_arm"/>
+    <group name="right_arm"/>
+    <group name="torso"/>
+    <group name="velodyne"/>
+    <group name="d435_head"/>
+  </group>
+  <!-- Homing posture -->
+  <!-- <group_state name="home" group="chains">
+        <joint name="hip_yaw_1" value="-0.746874"/>
+        <joint name="hip_pitch_1" value="-1.25409"/>
+        <joint name="knee_pitch_1" value="-1.55576"/>
+        <joint name="ankle_pitch_1" value="-0.301666"/>
+        <joint name="ankle_yaw_1" value="0.746874"/>
+        <joint name="hip_yaw_2" value="0.746874"/>
+        <joint name="hip_pitch_2" value="1.25409"/>
+        <joint name="knee_pitch_2" value="1.55576"/>
+        <joint name="ankle_pitch_2" value="0.301666"/>
+        <joint name="ankle_yaw_2" value="-0.746874"/>
+        <joint name="hip_yaw_3" value="0.746874"/>
+        <joint name="hip_pitch_3" value="1.25409"/>
+        <joint name="knee_pitch_3" value="1.55576"/>
+        <joint name="ankle_pitch_3" value="0.301667"/>
+        <joint name="ankle_yaw_3" value="-0.746874"/>
+        <joint name="hip_yaw_4" value="-0.746874"/>
+        <joint name="hip_pitch_4" value="-1.25409"/>
+        <joint name="knee_pitch_4" value="-1.55576"/>
+        <joint name="ankle_pitch_4" value="-0.301667"/>
+        <joint name="ankle_yaw_4" value="0.746874"/>
+        <joint name="torso_yaw" value="3.56617e-13"/>
+        <joint name="velodyne_joint" value="0"/>
+        <joint name="d435_head_joint" value="0"/>
+        <joint name="j_arm1_1" value="0.520149"/>
+        <joint name="j_arm1_2" value="0.320865"/>
+        <joint name="j_arm1_3" value="0.274669"/>
+        <joint name="j_arm1_4" value="-2.23604"/>
+        <joint name="j_arm1_5" value="0.0500815"/>
+        <joint name="j_arm1_6" value="-0.781461"/>
+        <joint name="j_arm2_1" value="0.520149"/>
+        <joint name="j_arm2_2" value="-0.320865"/>
+        <joint name="j_arm2_3" value="-0.274669"/>
+        <joint name="j_arm2_4" value="-2.23604"/>
+        <joint name="j_arm2_5" value="-0.0500815"/>
+        <joint name="j_arm2_6" value="-0.781461"/>
+        <joint name="j_wheel_1" value="0.0"/>
+        <joint name="j_wheel_2" value="0.0"/>
+        <joint name="j_wheel_3" value="0.0"/>
+        <joint name="j_wheel_4" value="0.0"/>
+        <xacro:if value="${END_EFFECTOR_LEFT == 'dagana'}">
+            <joint name="dagana_1_claw_joint" value="0"/>
+        </xacro:if>
+        <xacro:if value="${END_EFFECTOR_RIGHT == 'dagana'}">
+            <joint name="dagana_2_claw_joint" value="0"/>
+        </xacro:if>
+        
+    </group_state> -->
+  <!-- Balanced weight homing -->
+  <group_state group="chains" name="home">
+    <joint name="hip_yaw_1" value="-0.5"/>
+    <joint name="hip_pitch_1" value="-1.37"/>
+    <joint name="knee_pitch_1" value="-1.34"/>
+    <joint name="ankle_pitch_1" value="-0.03"/>
+    <joint name="ankle_yaw_1" value="0.6"/>
+    <joint name="hip_yaw_2" value="0.5"/>
+    <joint name="hip_pitch_2" value="1.37"/>
+    <joint name="knee_pitch_2" value="1.34"/>
+    <joint name="ankle_pitch_2" value="-0.03"/>
+    <joint name="ankle_yaw_2" value="-0.6"/>
+    <joint name="hip_yaw_3" value="1.28"/>
+    <joint name="hip_pitch_3" value="1.29"/>
+    <joint name="knee_pitch_3" value="1.71"/>
+    <joint name="ankle_pitch_3" value="0.43"/>
+    <joint name="ankle_yaw_3" value="-1.3"/>
+    <joint name="hip_yaw_4" value="-1.28"/>
+    <joint name="hip_pitch_4" value="-1.29"/>
+    <joint name="knee_pitch_4" value="-1.71"/>
+    <joint name="ankle_pitch_4" value="-0.43"/>
+    <joint name="ankle_yaw_4" value="1.3"/>
+    <joint name="torso_yaw" value="3.56617e-13"/>
+    <joint name="velodyne_joint" value="0"/>
+    <joint name="d435_head_joint" value="0"/>
+    <joint name="j_arm1_1" value="0.520149"/>
+    <joint name="j_arm1_2" value="0.320865"/>
+    <joint name="j_arm1_3" value="0.274669"/>
+    <joint name="j_arm1_4" value="-2.23604"/>
+    <joint name="j_arm1_5" value="0.0500815"/>
+    <joint name="j_arm1_6" value="-0.781461"/>
+    <joint name="j_arm2_1" value="0.520149"/>
+    <joint name="j_arm2_2" value="-0.320865"/>
+    <joint name="j_arm2_3" value="-0.274669"/>
+    <joint name="j_arm2_4" value="-2.23604"/>
+    <joint name="j_arm2_5" value="-0.0500815"/>
+    <joint name="j_arm2_6" value="-0.781461"/>
+    <joint name="j_wheel_1" value="0.0"/>
+    <joint name="j_wheel_2" value="0.0"/>
+    <joint name="j_wheel_3" value="0.0"/>
+    <joint name="j_wheel_4" value="0.0"/>
+  </group_state>
+</robot>

centauro_big_wheels_no_yaw_ub_2026_03_07_19_20_05_ID.urdf

@@ -0,0 +1,1569 @@
+<?xml version="1.0" ?>
+<!-- =================================================================================== -->
+<!-- |    This document was autogenerated by xacro from /root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/urdf/centauro.urdf.xacro | -->
+<!-- |    EDITING THIS FILE BY HAND IS NOT RECOMMENDED                                 | -->
+<!-- =================================================================================== -->
+<robot name="centauro">
+  <!-- none, ball, dagana_fixed, dagana -->
+  <!-- scaling facto for available peack torque link side, from 0 to 1 -->
+  <material name="black">
+    <color rgba="0.4 0.4 0.4 1"/>
+  </material>
+  <material name="dark">
+    <color rgba="0.1 0.1 0.1 1"/>
+  </material>
+  <material name="dark_grey">
+    <color rgba="0.2 0.2 0.2 1"/>
+  </material>
+  <!-- Note that the "23" suffix denotes the new bigger-size wheel -->
+  <!-- ??? -->
+  <!-- ??? -->
+  <!-- Origin -->
+  <!-- <xacro:property name="hip_yaw_lower" value="${[-2.827,-2.583, -2.583 ,-2.827]}"/>
+<xacro:property name="hip_yaw_upper" value="${[ 2.583, 2.827,  2.827 , 2.583]}"/> -->
+  <link name="base_link"/>
+  <joint name="base_joint" type="fixed">
+    <parent link="base_link"/>
+    <child link="pelvis"/>
+    <origin rpy="0 0 0" xyz="0 0 0"/>
+  </joint>
+  <link name="pelvis">
+    <inertial>
+      <origin xyz="-0.0016320207 0.0026534004000000003 0.079167428"/>
+      <mass value="26.584529"/>
+      <inertia ixx="0.39195477" ixy="0.024781258999999996" ixz="0.017108606999999998" iyy="1.0000708999999999" iyz="0.0052627106000000005" izz="1.0679988999999999"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/pelvis.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/pelvis_reduced.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="imu_link"/>
+  <joint name="imu_joint" type="fixed">
+    <parent link="pelvis"/>
+    <origin rpy="0.0 0.0 0.0" xyz="0.248316 0.0 0.015"/>
+    <child link="imu_link"/>
+  </joint>
+  <!-- torso links -->
+  <link name="torso_2">
+    <inertial>
+      <origin xyz="-0.019497794 0.0045317835 0.13768283"/>
+      <mass value="12.553731"/>
+      <inertia ixx="0.063643392" ixy="8.939e-05" ixz="-0.00086873" iyy="0.02680235" iyz="-4.657e-05" izz="0.04743015"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/link1_no_head.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/link1_no_head.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <!-- /torso links -->
+  <!-- torso joints -->
+  <joint name="torso_yaw" type="revolute">
+    <parent link="pelvis"/>
+    <origin xyz="0.2 0.0 0.256"/>
+    <child link="torso_2"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="147" lower="-2.618" upper="2.618" velocity="5.7"/>
+    <!-- TODO -->
+  </joint>
+  <!-- /torso joints -->
+  <!-- /macro arms -->
+  <!-- LINKS -->
+  <!-- shoulder yaw-roll-->
+  <link name="arm1_1">
+    <inertial>
+      <origin xyz="-0.0074457212 0.03410796 0.00010978102"/>
+      <mass value="1.9628675"/>
+      <inertia ixx="0.0053547717" ixy="0.00036428926" ixz="1.5089568e-05" iyy="0.0033923328" iyz="5.5692312e-05" izz="0.0068921413"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/ShoulderPitch.STL" scale="0.001 -0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/ShoulderPitch.STL" scale="0.001 -0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <!-- shoulder yaw-roll-->
+  <!-- shoulder roll-pitch-->
+  <link name="arm1_2">
+    <inertial>
+      <origin xyz="0.058142302 5.7450803e-05 -0.077477683"/>
+      <mass value="1.8595811"/>
+      <inertia ixx="0.013776643" ixy="3.7788675e-05" ixz="0.0037690171" iyy="0.015677464" iyz="-9.4893549e-06" izz="0.0046317657"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/ShoulderRoll.STL" scale="0.001 -0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/ShoulderRoll.STL" scale="0.001 -0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <!-- shoulder roll-pitch-->
+  <!-- shoulder-elbow -->
+  <link name="arm1_3">
+    <inertial>
+      <origin xyz="0.014625194 0.0008172672 -0.028333545"/>
+      <mass value="1.6678109"/>
+      <inertia ixx="0.0064480435" ixy="-0.00015639093" ixz="0.0012205359" iyy="0.0073372077" iyz="8.9941532e-05" izz="0.0036738448"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/ShoulderYaw.STL" scale="0.001 -0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/ShoulderYaw.STL" scale="0.001 -0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <!-- /shoulder-elbow -->
+  <!-- elbow yaw-pitch -->
+  <link name="arm1_4">
+    <inertial>
+      <origin xyz="-0.0076833067 -0.040302205 -0.043492779"/>
+      <mass value="1.3157289"/>
+      <inertia ixx="0.004330394" ixy="-0.00011737391" ixz="-0.00041923199" iyy="0.0038539919" iyz="-0.00079573038" izz="0.0017594689"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/Elbow.STL" scale="0.001 -0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/Elbow.STL" scale="0.001 -0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <!-- /elbow yaw-pitch -->
+  <!-- elbow-wrist -->
+  <link name="arm1_5">
+    <inertial>
+      <origin xyz="-0.00011079615 0.011590836 -0.07816026"/>
+      <mass value="1.4908547"/>
+      <inertia ixx="0.0085692128" ixy="1.7856252e-05" ixz="1.9379365e-05" iyy="0.0077454159" iyz="-0.00032860094" izz="0.0027441921"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/Forearm.STL" scale="0.001 -0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/Forearm.STL" scale="0.001 -0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <!-- /elbow-wrist -->
+  <!-- wrist yaw-pitch -->
+  <link name="arm1_6">
+    <inertial>
+      <origin xyz="-4.6502396e-06 -0.038014094 -0.069926878"/>
+      <mass value="1.1263612"/>
+      <inertia ixx="0.0051871784" ixy="2.724437e-05" ixz="2.2833496e-06" iyy="0.0048037789" iyz="-0.00072165653" izz="0.0012771388"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/wrist_pitch_right.stl" scale="0.001 -0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/wrist_pitch_right.stl" scale="0.001 -0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <!-- /wrist yaw-pitch -->
+  <!-- <link name="arm${arm_num}_7">
+      <inertial>
+         <origin xyz="${Wrist_3_x} ${rot*Wrist_3_y} ${Wrist_3_z}"/>
+         <mass value="${Wrist_3_mass}"/>
+         <inertia ixx="${Wrist_3_xx}" ixy="${rot*Wrist_3_xy}" ixz="${Wrist_3_xz}" iyy="${Wrist_3_yy}" iyz="${rot*Wrist_3_yz}" izz="${Wrist_3_zz}"/>
+      </inertial>
+
+      <visual>
+         <origin xyz="0.0 0.0 0.0"  rpy="0.0 0.0 0.0"/>
+   <geometry>
+      <mesh filename="${MESH_PATH}/ForearmYaw.STL" scale="0.001 ${-rot*0.001} 0.001" />
+   </geometry>
+         <material name="dark_grey"/>
+      </visual>
+
+      <collision>
+         <origin xyz="0.0 0.0 0.0"  rpy="0.0 0.0 0.0"/>
+   <geometry>
+      <mesh filename="${MESH_PATH}/simple/ForearmYaw.STL" scale="0.001 ${-rot*0.001} 0.001" />
+   </geometry>
+      </collision>
+   </link> -->
+  <!-- /LINKS -->
+  <!-- JOINTS -->
+  <!-- shoulder yaw -->
+  <joint name="j_arm1_1" type="revolute">
+    <parent link="torso_2"/>
+    <child link="arm1_1"/>
+    <origin rpy="0.174533 0.0 -0.5235988" xyz="0.0457475 0.169137 0.116626"/>
+    <axis xyz="0 1 0"/>
+    <limit effort="147.0" lower="-3.312" upper="1.615" velocity="3.86"/>
+  </joint>
+  <!-- /shoulder yaw -->
+  <!-- shoulder roll -->
+  <joint name="j_arm1_2" type="revolute">
+    <parent link="arm1_1"/>
+    <child link="arm1_2"/>
+    <origin rpy="-0.1745327 0.0 0.0" xyz="-0.09015 0.062 0.0"/>
+    <axis xyz="1 0 0"/>
+    <limit effort="147.0" lower="0.02" upper="3.431" velocity="3.86"/>
+  </joint>
+  <!-- /shoulder roll -->
+  <!-- shoulder pitch-->
+  <joint name="j_arm1_3" type="revolute">
+    <parent link="arm1_2"/>
+    <child link="arm1_3"/>
+    <origin xyz="0.09015 0.0 -0.21815"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="147.0" lower="-2.552" upper="2.566" velocity="6.06"/>
+  </joint>
+  <!-- /shoulder pitch -->
+  <!-- elbow yaw -->
+  <joint name="j_arm1_4" type="revolute">
+    <parent link="arm1_3"/>
+    <child link="arm1_4"/>
+    <origin xyz="0.045 0.05515 -0.074"/>
+    <axis xyz="0 1 0"/>
+    <limit effort="147.0" lower="-2.465" upper="0.28" velocity="6.06"/>
+  </joint>
+  <!-- /elbow yaw -->
+  <!-- elbow pitch-->
+  <joint name="j_arm1_5" type="revolute">
+    <parent link="arm1_4"/>
+    <child link="arm1_5"/>
+    <origin xyz="-0.015 -0.05515 -0.095"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="55.0" lower="-2.569" upper="2.562" velocity="11.72"/>
+  </joint>
+  <!-- /elbow pitch-->
+  <!-- wrist yaw-->
+  <joint name="j_arm1_6" type="revolute">
+    <parent link="arm1_5"/>
+    <child link="arm1_6"/>
+    <origin xyz="0.0 0.049 -0.156"/>
+    <axis xyz="0 1 0"/>
+    <limit effort="55.0" lower="-1.529" upper="1.509" velocity="11.72"/>
+  </joint>
+  <!-- /wrist yaw-->
+  <!-- wrist pitch-->
+  <!-- <joint name="j_arm${arm_num}_7" type="fixed">
+      <parent link="arm${arm_num}_6"/>
+      <child link="arm${arm_num}_7"/>
+      <origin xyz="${Wrist_3_Ox} ${rot*Wrist_3_Oy} ${Wrist_3_Oz}"/>
+      <axis xyz="0 0 1"/>
+      <limit lower="${j_arm_7_lower[arm_num-1]}" upper="${j_arm_7_upper[arm_num-1]}" effort="${j_arm_7_torque[arm_num-1]}" velocity="${j_arm_7_velocity[arm_num-1]}"/>
+   </joint> -->
+  <!-- wrist pitch-->
+  <!-- force-troque sensor -->
+  <!-- end-effector -->
+  <!-- wrist pitch-->
+  <!-- <joint name="j_ft_${arm_num}" type="fixed">
+      <parent link="arm${arm_num}_7"/>
+      <child link="ft_arm${arm_num}"/>
+      <origin xyz="${Ft_arm_Ox} ${Ft_arm_Oy} ${Ft_arm_Oz}" rpy="${Ft_arm_roll} ${Ft_arm_pitch} ${Ft_arm_yaw}" />
+      <axis xyz="0 0 1"/>
+      <limit lower="0.0" upper="0.0" effort="0" velocity="0"/>
+   </joint> -->
+  <!-- wrist pitch-->
+  <!-- <link name="ft_arm${arm_num}">
+      <inertial>
+         <origin xyz="${Ft_arm_x} ${rot*Ft_arm_y} ${Ft_arm_z}"/>
+         <mass value="${Ft_arm_mass}"/>
+         <inertia ixx="${Ft_arm_xx}" ixy="${rot*Ft_arm_xy}" ixz="${Ft_arm_xz}" iyy="${Ft_arm_yy}" iyz="${rot*Ft_arm_yz}" izz="${Ft_arm_zz}"/>
+      </inertial> -->
+  <!--visual>
+         <origin xyz="0.0 0.0 0.0"  rpy="${PI} 0.0 0.0"/>
+   <geometry>
+      <mesh filename="${MESH_PATH}/Wrist.STL" scale="0.001 ${-rot*0.001} 0.001" />
+   </geometry>
+      </visual-->
+  <!--collision>
+         <origin xyz="0.0 0.0 0.0"  rpy="${PI} 0.0 0.0"/>
+   <geometry>
+      <mesh filename="${MESH_PATH}/simple/Wrist.STL" scale="0.001 ${-rot*0.001} 0.001" />
+   </geometry>
+      </collision-->
+  <!-- </link> -->
+  <!-- LINKS -->
+  <!-- shoulder yaw-roll-->
+  <link name="arm2_1">
+    <inertial>
+      <origin xyz="-0.0074457212 -0.03410796 0.00010978102"/>
+      <mass value="1.9628675"/>
+      <inertia ixx="0.0053547717" ixy="-0.00036428926" ixz="1.5089568e-05" iyy="0.0033923328" iyz="-5.5692312e-05" izz="0.0068921413"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/ShoulderPitch.STL" scale="0.001 0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/ShoulderPitch.STL" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <!-- shoulder yaw-roll-->
+  <!-- shoulder roll-pitch-->
+  <link name="arm2_2">
+    <inertial>
+      <origin xyz="0.058142302 -5.7450803e-05 -0.077477683"/>
+      <mass value="1.8595811"/>
+      <inertia ixx="0.013776643" ixy="-3.7788675e-05" ixz="0.0037690171" iyy="0.015677464" iyz="9.4893549e-06" izz="0.0046317657"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/ShoulderRoll.STL" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/ShoulderRoll.STL" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <!-- shoulder roll-pitch-->
+  <!-- shoulder-elbow -->
+  <link name="arm2_3">
+    <inertial>
+      <origin xyz="0.014625194 -0.0008172672 -0.028333545"/>
+      <mass value="1.6678109"/>
+      <inertia ixx="0.0064480435" ixy="0.00015639093" ixz="0.0012205359" iyy="0.0073372077" iyz="-8.9941532e-05" izz="0.0036738448"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/ShoulderYaw.STL" scale="0.001 0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/ShoulderYaw.STL" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <!-- /shoulder-elbow -->
+  <!-- elbow yaw-pitch -->
+  <link name="arm2_4">
+    <inertial>
+      <origin xyz="-0.0076833067 0.040302205 -0.043492779"/>
+      <mass value="1.3157289"/>
+      <inertia ixx="0.004330394" ixy="0.00011737391" ixz="-0.00041923199" iyy="0.0038539919" iyz="0.00079573038" izz="0.0017594689"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/Elbow.STL" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/Elbow.STL" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <!-- /elbow yaw-pitch -->
+  <!-- elbow-wrist -->
+  <link name="arm2_5">
+    <inertial>
+      <origin xyz="-0.00011079615 -0.011590836 -0.07816026"/>
+      <mass value="1.4908547"/>
+      <inertia ixx="0.0085692128" ixy="-1.7856252e-05" ixz="1.9379365e-05" iyy="0.0077454159" iyz="0.00032860094" izz="0.0027441921"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/Forearm.STL" scale="0.001 0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/Forearm.STL" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <!-- /elbow-wrist -->
+  <!-- wrist yaw-pitch -->
+  <link name="arm2_6">
+    <inertial>
+      <origin xyz="-4.6502396e-06 0.038014094 -0.069926878"/>
+      <mass value="1.1263612"/>
+      <inertia ixx="0.0051871784" ixy="-2.724437e-05" ixz="2.2833496e-06" iyy="0.0048037789" iyz="0.00072165653" izz="0.0012771388"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/wrist_pitch_right.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/wrist_pitch_right.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <!-- /wrist yaw-pitch -->
+  <!-- <link name="arm${arm_num}_7">
+      <inertial>
+         <origin xyz="${Wrist_3_x} ${rot*Wrist_3_y} ${Wrist_3_z}"/>
+         <mass value="${Wrist_3_mass}"/>
+         <inertia ixx="${Wrist_3_xx}" ixy="${rot*Wrist_3_xy}" ixz="${Wrist_3_xz}" iyy="${Wrist_3_yy}" iyz="${rot*Wrist_3_yz}" izz="${Wrist_3_zz}"/>
+      </inertial>
+
+      <visual>
+         <origin xyz="0.0 0.0 0.0"  rpy="0.0 0.0 0.0"/>
+   <geometry>
+      <mesh filename="${MESH_PATH}/ForearmYaw.STL" scale="0.001 ${-rot*0.001} 0.001" />
+   </geometry>
+         <material name="dark_grey"/>
+      </visual>
+
+      <collision>
+         <origin xyz="0.0 0.0 0.0"  rpy="0.0 0.0 0.0"/>
+   <geometry>
+      <mesh filename="${MESH_PATH}/simple/ForearmYaw.STL" scale="0.001 ${-rot*0.001} 0.001" />
+   </geometry>
+      </collision>
+   </link> -->
+  <!-- /LINKS -->
+  <!-- JOINTS -->
+  <!-- shoulder yaw -->
+  <joint name="j_arm2_1" type="revolute">
+    <parent link="torso_2"/>
+    <child link="arm2_1"/>
+    <origin rpy="-0.174533 0.0 0.5235988" xyz="0.0457475 -0.169137 0.116626"/>
+    <axis xyz="0 1 0"/>
+    <limit effort="147.0" lower="-3.3458" upper="1.6012" velocity="3.86"/>
+  </joint>
+  <!-- /shoulder yaw -->
+  <!-- shoulder roll -->
+  <joint name="j_arm2_2" type="revolute">
+    <parent link="arm2_1"/>
+    <child link="arm2_2"/>
+    <origin rpy="0.1745327 0.0 0.0" xyz="-0.09015 -0.062 0.0"/>
+    <axis xyz="1 0 0"/>
+    <limit effort="147.0" lower="-3.4258" upper="-0.0138" velocity="3.86"/>
+  </joint>
+  <!-- /shoulder roll -->
+  <!-- shoulder pitch-->
+  <joint name="j_arm2_3" type="revolute">
+    <parent link="arm2_2"/>
+    <child link="arm2_3"/>
+    <origin xyz="0.09015 0.0 -0.21815"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="147.0" lower="-2.5614" upper="2.5606" velocity="6.06"/>
+  </joint>
+  <!-- /shoulder pitch -->
+  <!-- elbow yaw -->
+  <joint name="j_arm2_4" type="revolute">
+    <parent link="arm2_3"/>
+    <child link="arm2_4"/>
+    <origin xyz="0.045 -0.05515 -0.074"/>
+    <axis xyz="0 1 0"/>
+    <limit effort="147.0" lower="-2.4794" upper="0.2886" velocity="6.06"/>
+  </joint>
+  <!-- /elbow yaw -->
+  <!-- elbow pitch-->
+  <joint name="j_arm2_5" type="revolute">
+    <parent link="arm2_4"/>
+    <child link="arm2_5"/>
+    <origin xyz="-0.015 0.05515 -0.095"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="55.0" lower="-2.5394" upper="2.5546" velocity="11.72"/>
+  </joint>
+  <!-- /elbow pitch-->
+  <!-- wrist yaw-->
+  <joint name="j_arm2_6" type="revolute">
+    <parent link="arm2_5"/>
+    <child link="arm2_6"/>
+    <origin xyz="0.0 -0.049 -0.156"/>
+    <axis xyz="0 1 0"/>
+    <limit effort="55.0" lower="-1.5154" upper="1.5156" velocity="11.72"/>
+  </joint>
+  <!-- /wrist yaw-->
+  <!-- wrist pitch-->
+  <!-- <joint name="j_arm${arm_num}_7" type="fixed">
+      <parent link="arm${arm_num}_6"/>
+      <child link="arm${arm_num}_7"/>
+      <origin xyz="${Wrist_3_Ox} ${rot*Wrist_3_Oy} ${Wrist_3_Oz}"/>
+      <axis xyz="0 0 1"/>
+      <limit lower="${j_arm_7_lower[arm_num-1]}" upper="${j_arm_7_upper[arm_num-1]}" effort="${j_arm_7_torque[arm_num-1]}" velocity="${j_arm_7_velocity[arm_num-1]}"/>
+   </joint> -->
+  <!-- wrist pitch-->
+  <!-- force-troque sensor -->
+  <!-- end-effector -->
+  <!-- wrist pitch-->
+  <!-- <joint name="j_ft_${arm_num}" type="fixed">
+      <parent link="arm${arm_num}_7"/>
+      <child link="ft_arm${arm_num}"/>
+      <origin xyz="${Ft_arm_Ox} ${Ft_arm_Oy} ${Ft_arm_Oz}" rpy="${Ft_arm_roll} ${Ft_arm_pitch} ${Ft_arm_yaw}" />
+      <axis xyz="0 0 1"/>
+      <limit lower="0.0" upper="0.0" effort="0" velocity="0"/>
+   </joint> -->
+  <!-- wrist pitch-->
+  <!-- <link name="ft_arm${arm_num}">
+      <inertial>
+         <origin xyz="${Ft_arm_x} ${rot*Ft_arm_y} ${Ft_arm_z}"/>
+         <mass value="${Ft_arm_mass}"/>
+         <inertia ixx="${Ft_arm_xx}" ixy="${rot*Ft_arm_xy}" ixz="${Ft_arm_xz}" iyy="${Ft_arm_yy}" iyz="${rot*Ft_arm_yz}" izz="${Ft_arm_zz}"/>
+      </inertial> -->
+  <!--visual>
+         <origin xyz="0.0 0.0 0.0"  rpy="${PI} 0.0 0.0"/>
+   <geometry>
+      <mesh filename="${MESH_PATH}/Wrist.STL" scale="0.001 ${-rot*0.001} 0.001" />
+   </geometry>
+      </visual-->
+  <!--collision>
+         <origin xyz="0.0 0.0 0.0"  rpy="${PI} 0.0 0.0"/>
+   <geometry>
+      <mesh filename="${MESH_PATH}/simple/Wrist.STL" scale="0.001 ${-rot*0.001} 0.001" />
+   </geometry>
+      </collision-->
+  <!-- </link> -->
+  <link name="ball1">
+    <inertial>
+      <origin rpy="0 0 0" xyz="-1.0072498e-05 3.7590658e-05 0.019772332"/>
+      <mass value="1.27"/>
+      <inertia ixx="0.0023061092" ixy="2.9181758e-07" ixz="-2.1246683999999997e-07" iyy="0.0023053155000000002" iyz="7.9290978e-07" izz="0.0030675615999999997"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry name="ball1_visual">
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/BallHand.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry name="ball1_collision">
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/BallHand.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="ball1_tip"/>
+  <joint name="j_arm1_8" type="fixed">
+    <parent link="arm1_6"/>
+    <child link="ball1"/>
+    <origin rpy="0. 0. 0." xyz="0. -0.05 -0.105"/>
+  </joint>
+  <joint name="j_ball1_tip" type="fixed">
+    <parent link="ball1"/>
+    <child link="ball1_tip"/>
+    <origin rpy="0 0 0" xyz="0 0 -0.07"/>
+  </joint>
+  <!-- add armX_8 link for legacy reasons (same as ballX) -->
+  <joint name="j_ball1_fixed" type="fixed">
+    <parent link="ball1"/>
+    <child link="arm1_8"/>
+  </joint>
+  <link name="arm1_8">
+    <inertial>
+      <mass value="0"/>
+      <inertia ixx="0" ixy="0" ixz="0" iyy="0" iyz="0" izz="0"/>
+    </inertial>
+  </link>
+  <link name="ball2">
+    <inertial>
+      <origin rpy="0 0 0" xyz="-1.0072498e-05 3.7590658e-05 0.019772332"/>
+      <mass value="1.27"/>
+      <inertia ixx="0.0023061092" ixy="2.9181758e-07" ixz="-2.1246683999999997e-07" iyy="0.0023053155000000002" iyz="7.9290978e-07" izz="0.0030675615999999997"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry name="ball2_visual">
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/BallHand.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry name="ball2_collision">
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/BallHand.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="ball2_tip"/>
+  <joint name="j_arm2_8" type="fixed">
+    <parent link="arm2_6"/>
+    <child link="ball2"/>
+    <origin rpy="0. 0. 0." xyz="0. 0.05 -0.105"/>
+  </joint>
+  <joint name="j_ball2_tip" type="fixed">
+    <parent link="ball2"/>
+    <child link="ball2_tip"/>
+    <origin rpy="0 0 0" xyz="0 0 -0.07"/>
+  </joint>
+  <!-- add armX_8 link for legacy reasons (same as ballX) -->
+  <joint name="j_ball2_fixed" type="fixed">
+    <parent link="ball2"/>
+    <child link="arm2_8"/>
+  </joint>
+  <link name="arm2_8">
+    <inertial>
+      <mass value="0"/>
+      <inertia ixx="0" ixy="0" ixz="0" iyy="0" iyz="0" izz="0"/>
+    </inertial>
+  </link>
+  <!-- LINKS -->
+  <link name="neck_1">
+    <inertial>
+      <origin xyz="0 0 0"/>
+      <mass value="0.33878686"/>
+      <!-- MASS/INERTIA TO DO -->
+      <inertia ixx="0.0003718568" ixy="-9.42983e-05" ixz="-2.80644e-05" iyy="0.0001886879" iyz="-1.54534e-05" izz="0.0005023443"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/head-base_mesh.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/head-base_mesh.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="velodyne_motor">
+    <inertial>
+      <origin xyz="0 0 0"/>
+      <mass value="0.33878686"/>
+      <!-- MASS/INERTIA TO DO -->
+      <inertia ixx="0.0003718568" ixy="-9.42983e-05" ixz="-2.80644e-05" iyy="0.0001886879" iyz="-1.54534e-05" izz="0.0005023443"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/lidar_mesh.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/lidar_mesh.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="d435_head_motor">
+    <inertial>
+      <origin xyz="0 0 0"/>
+      <mass value="0.33878686"/>
+      <!-- MASS/INERTIA TO DO -->
+      <inertia ixx="0.0003718568" ixy="-9.42983e-05" ixz="-2.80644e-05" iyy="0.0001886879" iyz="-1.54534e-05" izz="0.0005023443"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/neck-pitch_mesh.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/simple/neck-pitch_mesh.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <!--JOINTS -->
+  <joint name="neck_base" type="fixed">
+    <parent link="torso_2"/>
+    <origin rpy="0 0 0" xyz="0.0030637045 0.0061607221 0.158"/>
+    <child link="neck_1"/>
+    <axis xyz="0 0 1"/>
+  </joint>
+  <joint name="velodyne_joint" type="revolute">
+    <parent link="neck_1"/>
+    <origin rpy="0 0 0" xyz="0.0822724 0 0.242"/>
+    <child link="velodyne_motor"/>
+    <axis xyz="0 -1 0"/>
+    <limit effort="35" lower="-1.1" upper="1.1" velocity="5.7"/>
+    <dynamics friction="10.0"/>
+  </joint>
+  <joint name="d435_head_joint" type="revolute">
+    <parent link="neck_1"/>
+    <origin rpy="0 0 0" xyz="0.06 0 0.1375"/>
+    <child link="d435_head_motor"/>
+    <axis xyz="0 -1 0"/>
+    <limit effort="35" lower="-1.1" upper="1.1" velocity="5.7"/>
+    <dynamics friction="10.0"/>
+  </joint>
+  <!-- camera body, with origin at bottom screw mount -->
+  <joint name="D435_head_camera_joint" type="fixed">
+    <origin rpy="0 0 0" xyz="0.03015 0 -0.0125"/>
+    <parent link="d435_head_motor"/>
+    <child link="D435_head_camera_bottom_screw_frame"/>
+  </joint>
+  <link name="D435_head_camera_bottom_screw_frame"/>
+  <joint name="D435_head_camera_link_joint" type="fixed">
+    <origin rpy="0 0 0" xyz="0 0.0175 0.0125"/>
+    <parent link="D435_head_camera_bottom_screw_frame"/>
+    <child link="D435_head_camera_link"/>
+  </joint>
+  <link name="D435_head_camera_link">
+    <visual>
+      <origin rpy="1.570796326795 0 1.570796326795" xyz="0.0149 -0.0175 0"/>
+      <geometry>
+        <box size="0.09 0.025 0.02505"/>
+        <!-- <mesh filename="${MESH_PATH}/realsense/d435.dae"/> -->
+      </geometry>
+      <material name="D435_head_camera_aluminum"/>
+    </visual>
+    <collision>
+      <origin rpy="0 0 0" xyz="0 -0.0175 0"/>
+      <geometry>
+        <box size="0.02505 0.09 0.025"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <!-- The following are not reliable values, and should not be used for modeling -->
+      <mass value="0.564"/>
+      <origin xyz="0 0 0"/>
+      <inertia ixx="0.003881243" ixy="0.0" ixz="0.0" iyy="0.000498940" iyz="0.0" izz="0.003879257"/>
+    </inertial>
+  </link>
+  <!-- camera body, with origin at bottom screw mount -->
+  <joint name="D435i_camera_joint" type="fixed">
+    <origin rpy="0 0.54 0" xyz="0.272 0 0.084"/>
+    <parent link="pelvis"/>
+    <child link="D435i_camera_bottom_screw_frame"/>
+  </joint>
+  <link name="D435i_camera_bottom_screw_frame"/>
+  <joint name="D435i_camera_link_joint" type="fixed">
+    <origin rpy="0 0.0050 -0.008" xyz="0.053 0.0 -0.035"/>
+    <parent link="D435i_camera_bottom_screw_frame"/>
+    <child link="D435i_camera_link"/>
+  </joint>
+  <link name="D435i_camera_link">
+    <visual>
+      <origin rpy="1.570796326795 0 1.570796326795" xyz="0.053 -0.0 0"/>
+      <geometry>
+        <!-- <box size="${d435_cam_width} ${d435_cam_height} ${d435_cam_depth}"/> -->
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/realsense/d435.dae"/>
+      </geometry>
+      <material name="D435i_camera_aluminum"/>
+    </visual>
+    <collision>
+      <origin rpy="0 0 0" xyz="0 -0.0 0"/>
+      <geometry>
+        <box size="0.02505 0.09 0.025"/>
+      </geometry>
+    </collision>
+    <inertial>
+      <!-- The following are not reliable values, and should not be used for modeling -->
+      <mass value="0.564"/>
+      <origin xyz="0 0 0"/>
+      <inertia ixx="0.003881243" ixy="0.0" ixz="0.0" iyy="0.000498940" iyz="0.0" izz="0.003879257"/>
+    </inertial>
+  </link>
+  <!-- /macro centauro_wheel -->
+  <!-- /macro arms -->
+  <!-- /macro legs -->
+  <link name="hip1_1">
+    <inertial>
+      <origin xyz="0.0003595855 0.033338818000000006 -0.059190309"/>
+      <mass value="2.3349153"/>
+      <!-- org 2 -->
+      <inertia ixx="0.0068405374" ixy="1.8572633e-05" ixz="-4.8325972e-05" iyy="0.0026261815" iyz="0.0001564273" izz="0.0066933354"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/hip-yaw.stl" scale="0.001 -0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/hip-yaw.stl" scale="0.001 -0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="hip2_1">
+    <inertial>
+      <origin xyz="-4.6285952e-06 -0.2214019 0.03316685"/>
+      <mass value="3.4572118"/>
+      <!-- org 4.5 -->
+      <inertia ixx="0.050469853" ixy="-7.784938699999999e-06" ixz="-9.220580999999999e-06" iyy="0.008181596999999999" iyz="0.0059796848" izz="0.048213808999999996"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/hip-pitch.stl" scale="0.001 0.001 -0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/hip-pitch.stl" scale="0.001 0.001 -0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="knee_1">
+    <inertial>
+      <origin xyz="0.00012063381 0.14620674 0.029740711000000003"/>
+      <mass value="2.1199315"/>
+      <!-- org 3.5, foot un modeled-->
+      <inertia ixx="0.018426402999999997" ixy="2.8226635999999997e-05" ixz="-2.9204348e-05" iyy="0.0047685615" iyz="-0.002640771" izz="0.017793722"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/knee.stl" scale="0.001 0.001 -0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/knee.stl" scale="0.001 0.001 -0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="ankle1_1">
+    <inertial>
+      <origin xyz="0.00012360094 -0.06274712 -0.0049068453"/>
+      <mass value="1.4023368"/>
+      <!-- org 3.5, foot un modeled-->
+      <inertia ixx="0.0049928829" ixy="5.762355699999999e-06" ixz="1.9606394e-05" iyy="0.0017621464" iyz="0.00045315091" izz="0.0048706328"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/ankle-pitch.stl" scale="0.001 0.001 -0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/ankle-pitch.stl" scale="0.001 0.001 -0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="ankle2_1">
+    <inertial>
+      <origin xyz="0.00018512273 0.0017369405000000001 -0.091340683"/>
+      <mass value="2.440929"/>
+      <!-- org 3.5, foot un modeled-->
+      <inertia ixx="0.017711499" ixy="1.0713141e-05" ixz="-2.2094697e-05" iyy="0.014263485" iyz="0.00033426702" izz="0.0080747557"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/mesh_ankle-yaw_wheel23.stl" scale="0.001 -0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <!-- <collision>
+                <origin xyz="0.0 0.0 0.0" rpy="0.0 0.0 0.0" />
+                <geometry>
+                    <mesh
+                        filename="${MESH_PATH}/v2/collision/mesh_ankle-yaw_wheel23.stl"
+                        scale="0.001 ${mirror*0.001} 0.001" />
+                </geometry>
+            </collision> -->
+  </link>
+  <joint name="j_wheel_1" type="continuous">
+    <parent link="ankle2_1"/>
+    <child link="wheel_1"/>
+    <origin rpy="1.57079632679 0.0 0.0" xyz="0.0 -0.0 -0.14435"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="35" velocity="20.0"/>
+    <!-- 23 Nm, 1
+            rad/s-->
+  </joint>
+  <link name="wheel_1">
+    <inertial>
+      <origin xyz="0.0 0.0 0.00015403767"/>
+      <mass value="1.8569388"/>
+      <inertia ixx="0.008637359099999999" ixy="0.0" ixz="-0.0" iyy="0.0086373644" iyz="-0.0" izz="0.016376551"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/mesh_wheel23.stl" scale="0.001 0.001 -0.001"/>
+      </geometry>
+      <material name="black"/>
+    </visual>
+    <collision>
+      <origin rpy="0 0 0" xyz="0 0 0"/>
+      <geometry>
+        <cylinder length="0.03" radius="0.124"/>
+      </geometry>
+      <!-- <geometry>
+                    <mesh filename="${MESH_PATH}/v2/collision/wheel23_donut.stl"
+                        scale="0.001 0.001 ${mirror*0.001}" />
+                </geometry> -->
+      <material name="black"/>
+    </collision>
+  </link>
+  <!-- *********** CONTACT FRAME ******************-->
+  <link name="contact_1"/>
+  <joint name="contact_joint_1" type="fixed">
+    <parent link="ankle2_1"/>
+    <child link="contact_1"/>
+    <origin rpy="0 0 0" xyz="0 0 -0.26835"/>
+  </joint>
+  <!-- ******************************** -->
+  <joint name="hip_yaw_1" type="revolute">
+    <parent link="pelvis"/>
+    <child link="hip1_1"/>
+    <origin rpy="0.0 0.0 0.0" xyz="0.25 0.125 -0.056600000000000004"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="304" lower="-2.03" upper="2.51" velocity="8.8"/>
+    <!-- 200 Nm, 1 rad/s-->
+  </joint>
+  <joint name="hip_pitch_1" type="revolute">
+    <parent link="hip1_1"/>
+    <child link="hip2_1"/>
+    <origin rpy="1.57079632679 0.0 0.0" xyz="0.0 0.1146 -0.0625"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="304" lower="-2.0226" upper="2.02748" velocity="8.8"/>
+    <!-- 200 Nm, 1 rad/s-->
+  </joint>
+  <joint name="knee_pitch_1" type="revolute">
+    <parent link="hip2_1"/>
+    <child link="knee_1"/>
+    <origin rpy="3.14159265359 0.0 0.0" xyz="0.0 -0.3 0.1182"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="304" lower="-2.4056" upper="2.3994" velocity="8.8"/>
+    <!-- 200 Nm, 1 rad/s-->
+  </joint>
+  <joint name="ankle_pitch_1" type="revolute">
+    <parent link="knee_1"/>
+    <child link="ankle1_1"/>
+    <origin rpy="3.14159265359 0.0 0.0" xyz="0.0 0.2 0.10015"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="147" lower="-2.4056" upper="2.3184" velocity="8.1"/>
+    <!-- 98 Nm, 1 rad/s-->
+  </joint>
+  <joint name="ankle_yaw_1" type="revolute">
+    <parent link="ankle1_1"/>
+    <child link="ankle2_1"/>
+    <origin rpy="-1.57079632679 0.0 0.0" xyz="0.0 -0.139 -0.00085"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="35" lower="-2.5626" upper="2.5384" velocity="20.0"/>
+    <!-- 23 Nm, 1 rad/s-->
+  </joint>
+  <link name="hip1_2">
+    <inertial>
+      <origin xyz="0.0003595855 -0.033338818000000006 -0.059190309"/>
+      <mass value="2.3349153"/>
+      <!-- org 2 -->
+      <inertia ixx="0.0068405374" ixy="-1.8572633e-05" ixz="-4.8325972e-05" iyy="0.0026261815" iyz="-0.0001564273" izz="0.0066933354"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/hip-yaw.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/hip-yaw.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="hip2_2">
+    <inertial>
+      <origin xyz="-4.6285952e-06 -0.2214019 -0.03316685"/>
+      <mass value="3.4572118"/>
+      <!-- org 4.5 -->
+      <inertia ixx="0.050469853" ixy="-7.784938699999999e-06" ixz="9.220580999999999e-06" iyy="0.008181596999999999" iyz="-0.0059796848" izz="0.048213808999999996"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/hip-pitch.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/hip-pitch.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="knee_2">
+    <inertial>
+      <origin xyz="0.00012063381 0.14620674 -0.029740711000000003"/>
+      <mass value="2.1199315"/>
+      <!-- org 3.5, foot un modeled-->
+      <inertia ixx="0.018426402999999997" ixy="2.8226635999999997e-05" ixz="2.9204348e-05" iyy="0.0047685615" iyz="0.002640771" izz="0.017793722"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/knee.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/knee.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="ankle1_2">
+    <inertial>
+      <origin xyz="0.00012360094 -0.06274712 0.0049068453"/>
+      <mass value="1.4023368"/>
+      <!-- org 3.5, foot un modeled-->
+      <inertia ixx="0.0049928829" ixy="5.762355699999999e-06" ixz="-1.9606394e-05" iyy="0.0017621464" iyz="-0.00045315091" izz="0.0048706328"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/ankle-pitch.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/ankle-pitch.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="ankle2_2">
+    <inertial>
+      <origin xyz="0.00018512273 -0.0017369405000000001 -0.091340683"/>
+      <mass value="2.440929"/>
+      <!-- org 3.5, foot un modeled-->
+      <inertia ixx="0.017711499" ixy="-1.0713141e-05" ixz="-2.2094697e-05" iyy="0.014263485" iyz="-0.00033426702" izz="0.0080747557"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/mesh_ankle-yaw_wheel23.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <!-- <collision>
+                <origin xyz="0.0 0.0 0.0" rpy="0.0 0.0 0.0" />
+                <geometry>
+                    <mesh
+                        filename="${MESH_PATH}/v2/collision/mesh_ankle-yaw_wheel23.stl"
+                        scale="0.001 ${mirror*0.001} 0.001" />
+                </geometry>
+            </collision> -->
+  </link>
+  <joint name="j_wheel_2" type="continuous">
+    <parent link="ankle2_2"/>
+    <child link="wheel_2"/>
+    <origin rpy="1.57079632679 0.0 0.0" xyz="0.0 0.0 -0.14435"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="35" velocity="20.0"/>
+    <!-- 23 Nm, 1
+            rad/s-->
+  </joint>
+  <link name="wheel_2">
+    <inertial>
+      <origin xyz="0.0 0.0 -0.00015403767"/>
+      <mass value="1.8569388"/>
+      <inertia ixx="0.008637359099999999" ixy="0.0" ixz="0.0" iyy="0.0086373644" iyz="0.0" izz="0.016376551"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/mesh_wheel23.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+      <material name="black"/>
+    </visual>
+    <collision>
+      <origin rpy="0 0 0" xyz="0 0 0"/>
+      <geometry>
+        <cylinder length="0.03" radius="0.124"/>
+      </geometry>
+      <!-- <geometry>
+                    <mesh filename="${MESH_PATH}/v2/collision/wheel23_donut.stl"
+                        scale="0.001 0.001 ${mirror*0.001}" />
+                </geometry> -->
+      <material name="black"/>
+    </collision>
+  </link>
+  <!-- *********** CONTACT FRAME ******************-->
+  <link name="contact_2"/>
+  <joint name="contact_joint_2" type="fixed">
+    <parent link="ankle2_2"/>
+    <child link="contact_2"/>
+    <origin rpy="0 0 0" xyz="0 0 -0.26835"/>
+  </joint>
+  <!-- ******************************** -->
+  <joint name="hip_yaw_2" type="revolute">
+    <parent link="pelvis"/>
+    <child link="hip1_2"/>
+    <origin rpy="0.0 0.0 0.0" xyz="0.25 -0.125 -0.056600000000000004"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="304" lower="-2.51" upper="2.03" velocity="8.8"/>
+    <!-- 200 Nm, 1 rad/s-->
+  </joint>
+  <joint name="hip_pitch_2" type="revolute">
+    <parent link="hip1_2"/>
+    <child link="hip2_2"/>
+    <origin rpy="1.57079632679 0.0 0.0" xyz="0.0 -0.1146 -0.0625"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="304" lower="-2.0156" upper="2.0354" velocity="8.8"/>
+    <!-- 200 Nm, 1 rad/s-->
+  </joint>
+  <joint name="knee_pitch_2" type="revolute">
+    <parent link="hip2_2"/>
+    <child link="knee_2"/>
+    <origin rpy="3.14159265359 0.0 0.0" xyz="0.0 -0.3 -0.1182"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="304" lower="-2.4006" upper="2.4034" velocity="8.8"/>
+    <!-- 200 Nm, 1 rad/s-->
+  </joint>
+  <joint name="ankle_pitch_2" type="revolute">
+    <parent link="knee_2"/>
+    <child link="ankle1_2"/>
+    <origin rpy="3.14159265359 0.0 0.0" xyz="0.0 0.2 -0.10015"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="147" lower="-2.3266" upper="2.3914" velocity="8.1"/>
+    <!-- 98 Nm, 1 rad/s-->
+  </joint>
+  <joint name="ankle_yaw_2" type="revolute">
+    <parent link="ankle1_2"/>
+    <child link="ankle2_2"/>
+    <origin rpy="-1.57079632679 0.0 0.0" xyz="0.0 -0.139 0.00085"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="35" lower="-2.5546" upper="2.5484" velocity="20.0"/>
+    <!-- 23 Nm, 1 rad/s-->
+  </joint>
+  <link name="hip1_3">
+    <inertial>
+      <origin xyz="0.0003595855 0.033338818000000006 -0.059190309"/>
+      <mass value="2.3349153"/>
+      <!-- org 2 -->
+      <inertia ixx="0.0068405374" ixy="1.8572633e-05" ixz="-4.8325972e-05" iyy="0.0026261815" iyz="0.0001564273" izz="0.0066933354"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/hip-yaw.stl" scale="0.001 -0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/hip-yaw.stl" scale="0.001 -0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="hip2_3">
+    <inertial>
+      <origin xyz="-4.6285952e-06 -0.2214019 0.03316685"/>
+      <mass value="3.4572118"/>
+      <!-- org 4.5 -->
+      <inertia ixx="0.050469853" ixy="-7.784938699999999e-06" ixz="-9.220580999999999e-06" iyy="0.008181596999999999" iyz="0.0059796848" izz="0.048213808999999996"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/hip-pitch.stl" scale="0.001 0.001 -0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/hip-pitch.stl" scale="0.001 0.001 -0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="knee_3">
+    <inertial>
+      <origin xyz="0.00012063381 0.14620674 0.029740711000000003"/>
+      <mass value="2.1199315"/>
+      <!-- org 3.5, foot un modeled-->
+      <inertia ixx="0.018426402999999997" ixy="2.8226635999999997e-05" ixz="-2.9204348e-05" iyy="0.0047685615" iyz="-0.002640771" izz="0.017793722"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/knee.stl" scale="0.001 0.001 -0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/knee.stl" scale="0.001 0.001 -0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="ankle1_3">
+    <inertial>
+      <origin xyz="0.00012360094 -0.06274712 -0.0049068453"/>
+      <mass value="1.4023368"/>
+      <!-- org 3.5, foot un modeled-->
+      <inertia ixx="0.0049928829" ixy="5.762355699999999e-06" ixz="1.9606394e-05" iyy="0.0017621464" iyz="0.00045315091" izz="0.0048706328"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/ankle-pitch.stl" scale="0.001 0.001 -0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/ankle-pitch.stl" scale="0.001 0.001 -0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="ankle2_3">
+    <inertial>
+      <origin xyz="0.00018512273 0.0017369405000000001 -0.091340683"/>
+      <mass value="2.440929"/>
+      <!-- org 3.5, foot un modeled-->
+      <inertia ixx="0.017711499" ixy="1.0713141e-05" ixz="-2.2094697e-05" iyy="0.014263485" iyz="0.00033426702" izz="0.0080747557"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/mesh_ankle-yaw_wheel23.stl" scale="0.001 -0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <!-- <collision>
+                <origin xyz="0.0 0.0 0.0" rpy="0.0 0.0 0.0" />
+                <geometry>
+                    <mesh
+                        filename="${MESH_PATH}/v2/collision/mesh_ankle-yaw_wheel23.stl"
+                        scale="0.001 ${mirror*0.001} 0.001" />
+                </geometry>
+            </collision> -->
+  </link>
+  <joint name="j_wheel_3" type="continuous">
+    <parent link="ankle2_3"/>
+    <child link="wheel_3"/>
+    <origin rpy="1.57079632679 0.0 0.0" xyz="0.0 -0.0 -0.14435"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="35" velocity="20.0"/>
+    <!-- 23 Nm, 1
+            rad/s-->
+  </joint>
+  <link name="wheel_3">
+    <inertial>
+      <origin xyz="0.0 0.0 0.00015403767"/>
+      <mass value="1.8569388"/>
+      <inertia ixx="0.008637359099999999" ixy="0.0" ixz="-0.0" iyy="0.0086373644" iyz="-0.0" izz="0.016376551"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/mesh_wheel23.stl" scale="0.001 0.001 -0.001"/>
+      </geometry>
+      <material name="black"/>
+    </visual>
+    <collision>
+      <origin rpy="0 0 0" xyz="0 0 0"/>
+      <geometry>
+        <cylinder length="0.03" radius="0.124"/>
+      </geometry>
+      <!-- <geometry>
+                    <mesh filename="${MESH_PATH}/v2/collision/wheel23_donut.stl"
+                        scale="0.001 0.001 ${mirror*0.001}" />
+                </geometry> -->
+      <material name="black"/>
+    </collision>
+  </link>
+  <!-- *********** CONTACT FRAME ******************-->
+  <link name="contact_3"/>
+  <joint name="contact_joint_3" type="fixed">
+    <parent link="ankle2_3"/>
+    <child link="contact_3"/>
+    <origin rpy="0 0 0" xyz="0 0 -0.26835"/>
+  </joint>
+  <!-- ******************************** -->
+  <joint name="hip_yaw_3" type="revolute">
+    <parent link="pelvis"/>
+    <child link="hip1_3"/>
+    <origin rpy="0.0 0.0 0.0" xyz="-0.25 0.125 -0.056600000000000004"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="304" lower="-2.51" upper="2.03" velocity="8.8"/>
+    <!-- 200 Nm, 1 rad/s-->
+  </joint>
+  <joint name="hip_pitch_3" type="revolute">
+    <parent link="hip1_3"/>
+    <child link="hip2_3"/>
+    <origin rpy="1.57079632679 0.0 0.0" xyz="0.0 0.1146 -0.0625"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="304" lower="-2.0126" upper="2.0384" velocity="8.8"/>
+    <!-- 200 Nm, 1 rad/s-->
+  </joint>
+  <joint name="knee_pitch_3" type="revolute">
+    <parent link="hip2_3"/>
+    <child link="knee_3"/>
+    <origin rpy="3.14159265359 0.0 0.0" xyz="0.0 -0.3 0.1182"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="304" lower="-2.4056" upper="2.4064" velocity="8.8"/>
+    <!-- 200 Nm, 1 rad/s-->
+  </joint>
+  <joint name="ankle_pitch_3" type="revolute">
+    <parent link="knee_3"/>
+    <child link="ankle1_3"/>
+    <origin rpy="3.14159265359 0.0 0.0" xyz="0.0 0.2 0.10015"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="147" lower="-2.4696" upper="2.4464" velocity="8.1"/>
+    <!-- 98 Nm, 1 rad/s-->
+  </joint>
+  <joint name="ankle_yaw_3" type="revolute">
+    <parent link="ankle1_3"/>
+    <child link="ankle2_3"/>
+    <origin rpy="-1.57079632679 0.0 0.0" xyz="0.0 -0.139 -0.00085"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="35" lower="-2.5606" upper="2.5454" velocity="20.0"/>
+    <!-- 23 Nm, 1 rad/s-->
+  </joint>
+  <link name="hip1_4">
+    <inertial>
+      <origin xyz="0.0003595855 -0.033338818000000006 -0.059190309"/>
+      <mass value="2.3349153"/>
+      <!-- org 2 -->
+      <inertia ixx="0.0068405374" ixy="-1.8572633e-05" ixz="-4.8325972e-05" iyy="0.0026261815" iyz="-0.0001564273" izz="0.0066933354"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/hip-yaw.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/hip-yaw.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="hip2_4">
+    <inertial>
+      <origin xyz="-4.6285952e-06 -0.2214019 -0.03316685"/>
+      <mass value="3.4572118"/>
+      <!-- org 4.5 -->
+      <inertia ixx="0.050469853" ixy="-7.784938699999999e-06" ixz="9.220580999999999e-06" iyy="0.008181596999999999" iyz="-0.0059796848" izz="0.048213808999999996"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/hip-pitch.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/hip-pitch.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="knee_4">
+    <inertial>
+      <origin xyz="0.00012063381 0.14620674 -0.029740711000000003"/>
+      <mass value="2.1199315"/>
+      <!-- org 3.5, foot un modeled-->
+      <inertia ixx="0.018426402999999997" ixy="2.8226635999999997e-05" ixz="2.9204348e-05" iyy="0.0047685615" iyz="0.002640771" izz="0.017793722"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/knee.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/knee.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="ankle1_4">
+    <inertial>
+      <origin xyz="0.00012360094 -0.06274712 0.0049068453"/>
+      <mass value="1.4023368"/>
+      <!-- org 3.5, foot un modeled-->
+      <inertia ixx="0.0049928829" ixy="5.762355699999999e-06" ixz="-1.9606394e-05" iyy="0.0017621464" iyz="-0.00045315091" izz="0.0048706328"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/ankle-pitch.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </visual>
+    <collision>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/collision/ankle-pitch.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+    </collision>
+  </link>
+  <link name="ankle2_4">
+    <inertial>
+      <origin xyz="0.00018512273 -0.0017369405000000001 -0.091340683"/>
+      <mass value="2.440929"/>
+      <!-- org 3.5, foot un modeled-->
+      <inertia ixx="0.017711499" ixy="-1.0713141e-05" ixz="-2.2094697e-05" iyy="0.014263485" iyz="-0.00033426702" izz="0.0080747557"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/mesh_ankle-yaw_wheel23.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+      <material name="dark_grey"/>
+    </visual>
+    <!-- <collision>
+                <origin xyz="0.0 0.0 0.0" rpy="0.0 0.0 0.0" />
+                <geometry>
+                    <mesh
+                        filename="${MESH_PATH}/v2/collision/mesh_ankle-yaw_wheel23.stl"
+                        scale="0.001 ${mirror*0.001} 0.001" />
+                </geometry>
+            </collision> -->
+  </link>
+  <joint name="j_wheel_4" type="continuous">
+    <parent link="ankle2_4"/>
+    <child link="wheel_4"/>
+    <origin rpy="1.57079632679 0.0 0.0" xyz="0.0 0.0 -0.14435"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="35" velocity="20.0"/>
+    <!-- 23 Nm, 1
+            rad/s-->
+  </joint>
+  <link name="wheel_4">
+    <inertial>
+      <origin xyz="0.0 0.0 -0.00015403767"/>
+      <mass value="1.8569388"/>
+      <inertia ixx="0.008637359099999999" ixy="0.0" ixz="0.0" iyy="0.0086373644" iyz="0.0" izz="0.016376551"/>
+    </inertial>
+    <visual>
+      <origin rpy="0.0 0.0 0.0" xyz="0.0 0.0 0.0"/>
+      <geometry>
+        <mesh filename="/root/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/meshes/v2/mesh_wheel23.stl" scale="0.001 0.001 0.001"/>
+      </geometry>
+      <material name="black"/>
+    </visual>
+    <collision>
+      <origin rpy="0 0 0" xyz="0 0 0"/>
+      <geometry>
+        <cylinder length="0.03" radius="0.124"/>
+      </geometry>
+      <!-- <geometry>
+                    <mesh filename="${MESH_PATH}/v2/collision/wheel23_donut.stl"
+                        scale="0.001 0.001 ${mirror*0.001}" />
+                </geometry> -->
+      <material name="black"/>
+    </collision>
+  </link>
+  <!-- *********** CONTACT FRAME ******************-->
+  <link name="contact_4"/>
+  <joint name="contact_joint_4" type="fixed">
+    <parent link="ankle2_4"/>
+    <child link="contact_4"/>
+    <origin rpy="0 0 0" xyz="0 0 -0.26835"/>
+  </joint>
+  <!-- ******************************** -->
+  <joint name="hip_yaw_4" type="revolute">
+    <parent link="pelvis"/>
+    <child link="hip1_4"/>
+    <origin rpy="0.0 0.0 0.0" xyz="-0.25 -0.125 -0.056600000000000004"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="304" lower="-2.03" upper="2.51" velocity="8.8"/>
+    <!-- 200 Nm, 1 rad/s-->
+  </joint>
+  <joint name="hip_pitch_4" type="revolute">
+    <parent link="hip1_4"/>
+    <child link="hip2_4"/>
+    <origin rpy="1.57079632679 0.0 0.0" xyz="0.0 -0.1146 -0.0625"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="304" lower="-2.0326" upper="2.0524" velocity="8.8"/>
+    <!-- 200 Nm, 1 rad/s-->
+  </joint>
+  <joint name="knee_pitch_4" type="revolute">
+    <parent link="hip2_4"/>
+    <child link="knee_4"/>
+    <origin rpy="3.14159265359 0.0 0.0" xyz="0.0 -0.3 -0.1182"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="304" lower="-2.4176" upper="2.3884" velocity="8.8"/>
+    <!-- 200 Nm, 1 rad/s-->
+  </joint>
+  <joint name="ankle_pitch_4" type="revolute">
+    <parent link="knee_4"/>
+    <child link="ankle1_4"/>
+    <origin rpy="3.14159265359 0.0 0.0" xyz="0.0 0.2 -0.10015"/>
+    <axis xyz="0 0 1"/>
+    <limit effort="147" lower="-2.3426" upper="2.3934" velocity="8.1"/>
+    <!-- 98 Nm, 1 rad/s-->
+  </joint>
+  <joint name="ankle_yaw_4" type="revolute">
+    <parent link="ankle1_4"/>
+    <child link="ankle2_4"/>
+    <origin rpy="-1.57079632679 0.0 0.0" xyz="0.0 -0.139 0.00085"/>
+    <axis xyz="0 0 -1"/>
+    <limit effort="35" lower="-2.6046" upper="2.5514" velocity="20.0"/>
+    <!-- 23 Nm, 1 rad/s-->
+  </joint>
+  <!-- ************ Control frames ************ -->
+</robot>

centauro_rhc.py

@@ -0,0 +1,145 @@
+from aug_mpc.controllers.rhc.horizon_based.hybrid_quad_rhc import HybridQuadRhc
+
+import numpy as np
+
+from typing import Dict 
+
+from centaurohybridmpc.utils.sysutils import PathsGetter
+
+class CentauroRhc(HybridQuadRhc):
+
+    def __init__(self, 
+            srdf_path: str,
+            urdf_path: str,
+            robot_name: str, # used for shared memory namespaces
+            codegen_dir: str,
+            n_nodes: float = 31,
+            dt: float = 0.03,
+            injection_node: int = 10,
+            max_solver_iter = 1, # defaults to rt-iteration
+            open_loop: bool = True,
+            close_loop_all: bool = False,
+            dtype = np.float32, 
+            verbose = False, 
+            debug = False,
+            refs_in_hor_frame = True,
+            timeout_ms: int = 60000,
+            custom_opts: Dict = {}
+            ):
+
+        paths = PathsGetter()
+        self._files_suffix=""
+        if open_loop:
+            self._files_suffix="_open"
+        
+        self._add_upper_body=False
+        if ("add_upper_body" in custom_opts) and \
+            (custom_opts["add_upper_body"]):
+            self._add_upper_body=True
+            self._files_suffix+="_ub"
+
+        config_path=paths.RHCCONFIGPATH_NO_WHEELS+self._files_suffix+".yaml"
+        
+        super().__init__(srdf_path=srdf_path,
+            urdf_path=urdf_path,
+            config_path=config_path,
+            robot_name=robot_name, # used for shared memory namespaces
+            codegen_dir=codegen_dir, 
+            n_nodes=n_nodes,
+            dt=dt,
+            injection_node=injection_node,
+            max_solver_iter=max_solver_iter, # defaults to rt-iteration
+            open_loop=open_loop,
+            close_loop_all=close_loop_all,
+            dtype=dtype,
+            verbose=verbose, 
+            debug=debug,
+            refs_in_hor_frame=refs_in_hor_frame,
+            timeout_ms=timeout_ms,
+            custom_opts=custom_opts)
+        
+        self._fail_idx_scale=1e-9
+        self._fail_idx_thresh_open_loop=1e0
+        self._fail_idx_thresh_close_loop=10
+        
+        if open_loop:
+            self._fail_idx_thresh=self._fail_idx_thresh_open_loop
+        else:
+            self._fail_idx_thresh=self._fail_idx_thresh_close_loop
+
+        # adding some additional config files for jnt imp control
+        self._rhc_fpaths.append(paths.JNT_IMP_CONFIG_XBOT+".yaml")
+        self._rhc_fpaths.append(paths.JNT_IMP_CONFIG+".yaml")
+        
+    def _set_rhc_pred_idx(self):
+        self._pred_node_idx=round((self._n_nodes-1)*2/3)
+
+    def _set_rhc_cmds_idx(self):
+        self._rhc_cmds_node_idx=2
+
+    def _config_override(self):
+        paths = PathsGetter()
+        if ("control_wheels" in self._custom_opts):
+            if self._custom_opts["control_wheels"]:
+                self.config_path = paths.RHCCONFIGPATH_WHEELS+self._files_suffix+".yaml"
+                if ("fix_yaw" in self._custom_opts) and \
+                    (self._custom_opts["fix_yaw"]):
+                    self.config_path = paths.RHCCONFIGPATH_WHEELS_NO_YAW+self._files_suffix+".yaml"
+                if ("replace_continuous_joints" in self._custom_opts) and \
+                    (not self._custom_opts["replace_continuous_joints"]):
+                    # use continuous joints -> different config
+                    self.config_path = paths.RHCCONFIGPATH_WHEELS_CONTINUOUS+self._files_suffix+".yaml"
+                    if ("fix_yaw" in self._custom_opts) and \
+                        (self._custom_opts["fix_yaw"]):
+                        self.config_path = paths.RHCCONFIGPATH_WHEELS_CONTINUOUS_NO_YAW+self._files_suffix+".yaml"
+
+        else:
+            self._custom_opts["control_wheels"]=False
+
+        if not self._custom_opts["control_wheels"]:
+            self._fixed_jnt_patterns=self._fixed_jnt_patterns+\
+                ["j_wheel", 
+                "ankle_yaw"]
+            self._custom_opts["replace_continuous_joints"]=True
+        
+    def _init_problem(self):
+        
+        if not self._custom_opts["control_wheels"]:
+            self._yaw_vertical_weight=120.0
+        else:
+            self._yaw_vertical_weight=50.0
+
+        fixed_jnts_patterns=[
+            "d435_head",
+            "velodyne_joint",
+            "dagana"]
+        
+        if not self._add_upper_body:
+            fixed_jnts_patterns.append("j_arm")
+            fixed_jnts_patterns.append("torso")
+
+        if ("fix_yaw" in self._custom_opts) and \
+            (self._custom_opts["fix_yaw"]):
+            fixed_jnts_patterns.append("ankle_yaw")
+
+        flight_duration_sec=0.5 # [s]
+        flight_duration=int(flight_duration_sec/self._dt)
+        post_flight_duration_sec=0.2 # [s]
+        post_flight_duration=int(post_flight_duration_sec/self._dt)
+
+        step_height=0.1
+        if ("step_height" in self._custom_opts):
+            step_height=self._custom_opts["step_height"]
+
+        super()._init_problem(fixed_jnt_patterns=fixed_jnts_patterns,
+            wheels_patterns=["wheel_"],
+            foot_linkname="wheel_1",
+            flight_duration=flight_duration,
+            post_flight_stance=post_flight_duration,
+            step_height=step_height,
+            keep_yaw_vert=True,
+            yaw_vertical_weight=self._yaw_vertical_weight,
+            vertical_landing=True,
+            vertical_land_weight=10.0,
+            phase_force_reg=5e-2,
+            vel_bounds_weight=1.0)

centauro_rhc_wheels_continuous_no_yaw_ub.yaml

@@ -0,0 +1,344 @@
+solver:
+  type: ilqr
+  ipopt.linear_solver: ma57
+  ipopt.tol: 0.1
+#  ilqr.merit_der_threshold: 1e-3
+#  ilqr.defect_norm_threshold: 1e-3
+  ipopt.constr_viol_tol: 0.01
+  ilqr.constraint_violation_threshold: 1e-2
+#  ipopt.hessian_approximation: exact
+  ipopt.print_level: 5
+  ipopt.suppress_all_output: 'yes'
+  ipopt.sb: 'yes'
+  ilqr.suppress_all_output: 'yes'
+  ilqr.codegen_enabled: true
+  ilqr.codegen_workdir: /tmp/tyhio
+  ilqr.enable_gn: true
+  ilqr.hxx_reg_base: 0.0
+  ilqr.n_threads: 0
+  print_time: 0
+
+constraints:
+  - contact_1
+  - contact_2
+  - contact_3
+  - contact_4
+
+costs:
+  - z_contact_1
+  - z_contact_2
+  - z_contact_3
+  - z_contact_4
+  # - vz_contact_1
+  # - vz_contact_2
+  # - vz_contact_3
+  # - vz_contact_4
+  # - xy_contact_1
+  # - xy_contact_2
+  # - xy_contact_3
+  # - xy_contact_4
+  # - vxy_contact_1
+  # - vxy_contact_2
+  # - vxy_contact_3
+  # - vxy_contact_4
+  - base_lin_velxy
+  - base_lin_velz
+  - base_omega
+  - base_capture
+  - joint_posture_capture
+  - v_regularization
+  - a_regularization
+  # - force_regularization
+
+.define:
+  - &w_base_omega 15.
+  - &w_base_vxy 20.
+  - &w_base_vz 20.
+  - &w_base_z 15.
+  - &w_contact_z 225.0
+  - &w_contact_vz 250.0
+  - &w_contact_xy 60.0
+  - &w_contact_vxy 50.0
+  - &w_base_capture 200.
+  - &wheel_radius 0.124
+
+base_lin_velxy:
+  type: Cartesian
+  distal_link: base_link
+  indices: [0, 1]
+  nodes: ${range(1, N-5)}
+  cartesian_type: velocity
+  weight: *w_base_vxy
+
+base_lin_velz:
+  type: Cartesian
+  distal_link: base_link
+  indices: [2]
+  nodes: ${range(1, N-5)}
+  cartesian_type: velocity
+  weight: *w_base_vz
+
+base_omega:
+  type: Cartesian
+  distal_link: base_link
+  indices: [3, 4, 5]
+  nodes: ${range(1, N-5)}
+  cartesian_type: velocity
+  weight: *w_base_omega
+
+base_capture:
+  type: Cartesian
+  distal_link: base_link
+  indices: [0, 1, 2, 3, 4, 5]
+  nodes: ${range(N-5, N+1)}
+  cartesian_type: velocity
+  weight: *w_base_capture
+
+# ===============================
+
+rolling_contact_1:
+  type: Rolling
+  frame: wheel_1
+  radius: *wheel_radius
+
+rolling_contact_2:
+  type: Rolling
+  frame: wheel_2
+  radius: *wheel_radius
+
+rolling_contact_3:
+  type: Rolling
+  frame: wheel_3
+  radius: *wheel_radius
+
+rolling_contact_4:
+  type: Rolling
+  frame: wheel_4
+  radius: *wheel_radius
+
+# ==================================
+
+interaction_contact_1:
+  type: VertexForce
+  frame: contact_1
+  fn_min: 100.0
+  enable_fc: true
+  friction_coeff: 0.5
+  vertex_frames:
+    - wheel_1
+
+interaction_contact_2:
+  type: VertexForce
+  frame: contact_2
+  fn_min: 100.0
+  enable_fc: true
+  friction_coeff: 0.5
+  vertex_frames:
+    - wheel_2
+
+interaction_contact_3:
+  type: VertexForce
+  frame: contact_3
+  fn_min: 100.0
+  enable_fc: true
+  friction_coeff: 0.5
+  vertex_frames:
+    - wheel_3
+
+interaction_contact_4:
+  type: VertexForce
+  frame: contact_4
+  fn_min: 100.0
+  enable_fc: true
+  friction_coeff: 0.5
+  vertex_frames:
+    - wheel_4
+
+contact_1:
+  type: Contact
+  subtask: [interaction_contact_1, rolling_contact_1]
+
+contact_2:
+  type: Contact
+  subtask: [interaction_contact_2, rolling_contact_2]
+
+contact_3:
+  type: Contact
+  subtask: [interaction_contact_3, rolling_contact_3]
+
+contact_4:
+  type: Contact
+  subtask: [interaction_contact_4, rolling_contact_4]
+
+joint_posture_capture:
+  type: Postural
+  weight: [45.0, 10.0, 10.0, 10.0,
+          45.0, 10.0, 10.0, 10.0,
+          45.0, 10.0, 10.0, 10.0,
+          45.0, 10.0, 10.0, 10.0,
+          55.0,
+          25.0, 25.0, 25.0, 25.0, 25.0, 25.0,
+          25.0, 25.0, 25.0, 25.0, 25.0, 25.0]
+  indices: [0,  1,  2,  3,
+            6, 7,  8,  9,  
+            12, 13, 14, 15, 
+            18, 19, 20, 21,
+            24,
+            25, 26, 27, 28, 29, 30,
+            31, 32, 33, 34, 35, 36]
+  nodes: ${range(N-5, N+1)}
+
+v_regularization:
+  type: Regularization
+  variable_name: v
+  nodes: ${range(0, N+1)}
+  indices: [0, 1, 2, 3, 4, 5,
+    6, 7, 8, 9, 10,
+    11, 12, 13, 14, 15,
+    16, 17, 18, 19, 20,
+    21, 21, 23, 24, 25,
+    26, 
+    27, 28, 29, 30, 31, 32,
+    33, 34, 35, 36, 37, 38]
+  weight: [2e-1, 2e-1, 2e-1, 2e-1, 2e-1, 2e-1,
+        1e0, 8e-1, 8e-1, 8e-1, 1e0,
+        1e0, 8e-1, 8e-1, 8e-1, 1e0,
+        1e0, 8e-1, 8e-1, 8e-1, 1e0,
+        1e0, 8e-1, 8e-1, 8e-1, 1e0,
+        2e0,
+        2e0, 2e0, 2e0, 2e0, 2e0, 2e0,
+        2e0, 2e0, 2e0, 2e0, 2e0, 2e0]
+
+a_regularization:
+  type: Regularization
+  variable_name: a
+  nodes: ${range(0, N+1)}
+  indices: [0, 1, 2, 3, 4, 5,
+    6, 7, 8, 9, 10,
+    11, 12, 13, 14, 15,
+    16, 17, 18, 19, 20,
+    21, 21, 23, 24, 25,
+    26, 
+    27, 28, 29, 30, 31, 32,
+    33, 34, 35, 36, 37, 38]
+  weight: [4e-1, 4e-1, 4e-1, 4e-1, 4e-1, 4e-1,
+        4e-1, 4e-1, 4e-1, 4e-1, 8e-1,
+        4e-1, 4e-1, 4e-1, 4e-1, 8e-1,
+        4e-1, 4e-1, 4e-1, 4e-1, 8e-1,
+        4e-1, 4e-1, 4e-1, 4e-1, 8e-1,
+        4e-1,
+        4e-1, 4e-1, 4e-1, 4e-1, 4e-1, 4e-1,
+        4e-1, 4e-1, 4e-1, 4e-1, 4e-1, 4e-1]
+
+# flight phase traj tracking
+z_contact_1:
+  type: Cartesian
+  distal_link: contact_1
+  indices: [2]
+  cartesian_type: position
+  weight: *w_contact_z
+
+z_contact_2:
+  type: Cartesian
+  distal_link: contact_2
+  indices: [2]
+  cartesian_type: position
+  weight: *w_contact_z
+
+z_contact_3:
+  type: Cartesian
+  distal_link: contact_3
+  indices: [2]
+  cartesian_type: position
+  weight: *w_contact_z
+
+z_contact_4:
+  type: Cartesian
+  distal_link: contact_4
+  indices: [2]
+  cartesian_type: position
+  weight: *w_contact_z
+
+xy_contact_1:
+  type: Cartesian
+  distal_link: contact_1
+  indices: [0, 1]
+  cartesian_type: position
+  weight: *w_contact_xy
+
+xy_contact_2:
+  type: Cartesian
+  distal_link: contact_2
+  indices: [0, 1]
+  cartesian_type: position
+  weight: *w_contact_xy
+
+xy_contact_3:
+  type: Cartesian
+  distal_link: contact_3
+  indices: [0, 1]
+  cartesian_type: position
+  weight: *w_contact_xy
+
+xy_contact_4:
+  type: Cartesian
+  distal_link: contact_4
+  indices: [0, 1]
+  cartesian_type: position
+  weight: *w_contact_xy
+
+vz_contact_1:
+  type: Cartesian
+  distal_link: contact_1
+  indices: [2]
+  cartesian_type: velocity
+  weight: *w_contact_vz
+
+vz_contact_2:
+  type: Cartesian
+  distal_link: contact_2
+  indices: [2]
+  cartesian_type: velocity
+  weight: *w_contact_vz
+
+vz_contact_3:
+  type: Cartesian
+  distal_link: contact_3
+  indices: [2]
+  cartesian_type: velocity
+  weight: *w_contact_vz
+
+vz_contact_4:
+  type: Cartesian
+  distal_link: contact_4
+  indices: [2]
+  cartesian_type: velocity
+  weight: *w_contact_vz
+
+vxy_contact_1:
+  type: Cartesian
+  distal_link: contact_1
+  indices: [0, 1]
+  cartesian_type: velocity
+  weight: *w_contact_vxy
+
+vxy_contact_2:
+  type: Cartesian
+  distal_link: contact_2
+  indices: [0, 1]
+  cartesian_type: velocity
+  weight: *w_contact_vxy
+
+vxy_contact_3:
+  type: Cartesian
+  distal_link: contact_3
+  indices: [0, 1]
+  cartesian_type: velocity
+  weight: *w_contact_vxy
+
+vxy_contact_4:
+  type: Cartesian
+  distal_link: contact_4
+  indices: [0, 1]
+  cartesian_type: velocity
+  weight: *w_contact_vxy

d2026_03_07_h19_m22_s30-CentauroCloopPartialNoYawUb_FakePosTrackingEnv_model

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:1b1bb0640bc5b6daf899d3b2fa5be08e9a243dd72b8eeeae32e7ad121ed3deb4
+size 132

fake_pos_tracking_env.py

@@ -0,0 +1,202 @@
+import os
+
+from typing import Dict
+
+import torch
+
+from EigenIPC.PyEigenIPC import VLevel
+
+from mpc_hive.utilities.math_utils_torch import world2base_frame
+
+from aug_mpc_envs.training_envs.twist_tracking_env import TwistTrackingEnv
+
+class FakePosTrackingEnv(TwistTrackingEnv):
+    """Converts random planar position goals into twist references so the agent learns to drive the robot toward targets while managing contact scheduling."""
+
+    def __init__(self,
+            namespace: str,
+            actions_dim: int = 10,
+            verbose: bool = False,
+            vlevel: VLevel = VLevel.V1,
+            use_gpu: bool = True,
+            dtype: torch.dtype = torch.float32,
+            debug: bool = True,
+            override_agent_refs: bool = False,
+            timeout_ms: int = 60000,
+            env_opts: Dict = {}):
+
+        self._add_env_opt(env_opts, "max_distance", default=5.0) # [m] 
+        self._add_env_opt(env_opts, "min_distance", default=0.0) # [m]
+        self._add_env_opt(env_opts, "max_vref", default=1.0) # [m/s]
+        self._add_env_opt(env_opts, "max_dp", default=5.0) # [m] after this, v ref saturates
+        self._add_env_opt(env_opts, "max_dt", default=env_opts["max_dp"]/ env_opts["max_vref"])
+
+        TwistTrackingEnv.__init__(self, 
+            namespace=namespace,
+            actions_dim=actions_dim, # twist + contact flags
+            verbose=verbose,
+            vlevel=vlevel,
+            use_gpu=use_gpu,
+            dtype=dtype,
+            debug=debug,
+            override_agent_refs=override_agent_refs,
+            timeout_ms=timeout_ms,
+            env_opts=env_opts)
+
+    def get_file_paths(self):
+        paths=TwistTrackingEnv.get_file_paths(self)
+        paths.append(os.path.abspath(__file__))        
+        return paths
+    
+    def _custom_post_init(self):
+        TwistTrackingEnv._custom_post_init(self)
+        
+        # position targets to be reached (wrt robot's pos at ep start)
+        self._p_trgt_w=self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu)[:, 0:2].detach().clone()
+        self._p_delta_w=self._p_trgt_w.detach().clone()
+        self._dp_norm=torch.zeros((self._n_envs, 1),dtype=self._dtype,device=self._device)
+        self._dp_versor=self._p_trgt_w.detach().clone()
+
+        self._trgt_d=torch.zeros((self._n_envs, 1),dtype=self._dtype,device=self._device)
+        self._trgt_theta=torch.zeros((self._n_envs, 1),dtype=self._dtype,device=self._device)
+
+    def _update_loc_twist_refs(self):
+        # this is called at each env substep
+        
+        self._compute_twist_ref_w()
+    
+        if not self._override_agent_refs:
+            agent_p_ref_current=self._agent_refs.rob_refs.root_state.get(data_type="p",
+            gpu=self._use_gpu)
+            agent_p_ref_current[:, 0:2]=self._p_trgt_w
+
+        # then convert it to base ref local for the agent
+        robot_q = self._robot_state.root_state.get(data_type="q",gpu=self._use_gpu)
+        # rotate agent ref from world to robot base
+        world2base_frame(t_w=self._agent_twist_ref_current_w, q_b=robot_q, 
+            t_out=self._agent_twist_ref_current_base_loc)
+        # write it to agent refs tensors
+        self._agent_refs.rob_refs.root_state.set(data_type="twist", data=self._agent_twist_ref_current_base_loc,
+                                            gpu=self._use_gpu)
+
+    def _compute_twist_ref_w(self, env_indxs: torch.Tensor = None):
+        
+        # angular refs are not altered
+        if env_indxs is None:
+            # we update the position error using the current base position
+            self._p_delta_w[:, :]=self._p_trgt_w-\
+                self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu)[:, 0:2]
+                
+            self._dp_norm[:, :]=self._p_delta_w.norm(dim=1,keepdim=True)+1e-6
+            self._dp_versor[:, :]=self._p_delta_w/self._dp_norm
+
+            # apply for vref saturation
+            to_be_saturated=self._dp_norm[:, :]>self._env_opts["max_dp"]
+            self._dp_norm[to_be_saturated.flatten(), :]=self._env_opts["max_dp"]
+
+            # we compute the twist refs for the agent depending of the position error
+            self._agent_twist_ref_current_w[:, 0:2]=self._dp_norm*self._dp_versor/self._env_opts["max_dt"]
+            self._agent_twist_ref_current_w[:, 2:3]=0 # no vertical vel
+
+            # apply pof0 using last value of bernoully coeffs
+            self._agent_twist_ref_current_w[:, 0:3] = self._agent_twist_ref_current_w[:, 0:3]*self._bernoulli_coeffs_linvel # linvel
+            self._agent_twist_ref_current_w[:, 3:6] = self._agent_twist_ref_current_w[:, 3:6]*self._bernoulli_coeffs_omega # omega
+        else:
+            self._p_delta_w[env_indxs, :]=self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu)[env_indxs, 0:2] -\
+                self._p_trgt_w[env_indxs, :]
+            
+            # apply for vref saturation
+            to_be_saturated=torch.logical_and((self._dp_norm[:, :]>self._env_opts["max_dp"]).flatten(),env_indxs)
+            self._dp_norm[to_be_saturated.flatten(), :]=self._env_opts["max_dp"]
+
+            self._dp_norm[env_indxs, :]=self._p_delta_w[env_indxs, :].norm(dim=1,keepdim=True)+1e-6
+            self._dp_versor[env_indxs, :]=self._p_delta_w[env_indxs, :]/self._dp_norm[env_indxs, :]
+
+            self._agent_twist_ref_current_w[env_indxs, 0:2]=self._dp_norm[env_indxs, :]*self._dp_versor[env_indxs, :]/self._env_opts["max_dt"]        
+            self._agent_twist_ref_current_w[env_indxs, 2:3]=0 # no vertical vel
+
+            # apply pof0 using last value of bernoully coeffs
+            self._agent_twist_ref_current_w[env_indxs, 0:3] = self._agent_twist_ref_current_w[env_indxs, 0:3]*self._bernoulli_coeffs_linvel[env_indxs, :]
+            self._agent_twist_ref_current_w[env_indxs, 3:6] = self._agent_twist_ref_current_w[env_indxs, 3:6]*self._bernoulli_coeffs_omega[env_indxs, :] # omega
+
+    def _override_refs(self,
+            env_indxs: torch.Tensor = None):
+        
+        # runs at every post_step
+        self._agent_refs.rob_refs.root_state.synch_all(read=True,retry=True) # first read from mem
+        if self._use_gpu:
+            # copies latest refs to GPU 
+            self._agent_refs.rob_refs.root_state.synch_mirror(from_gpu=False,non_blocking=False) 
+
+        agent_p_ref_current=self._agent_refs.rob_refs.root_state.get(data_type="p",
+                gpu=self._use_gpu)
+        
+        agent_yaw_omega_ref_current=self._agent_refs.rob_refs.root_state.get(data_type="omega",
+                gpu=self._use_gpu)
+        
+        # self._p_trgt_w[:, :]=self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu)[:, 0:2] + \
+        #     agent_p_ref_current[:, 0:2]
+        self._p_trgt_w[:, :]=agent_p_ref_current[:, 0:2] # set p target target from shared mem
+
+        self._agent_twist_ref_current_w[:, 5:6]=agent_yaw_omega_ref_current[:, 2:3] # set yaw ang. vel target from shared mem
+    
+    def _debug_agent_refs(self):
+        if self._use_gpu:
+            self._agent_refs.rob_refs.root_state.synch_mirror(from_gpu=True,non_blocking=False)
+        self._agent_refs.rob_refs.root_state.synch_all(read=False, retry = True)
+
+    def _randomize_task_refs(self,
+        env_indxs: torch.Tensor = None):
+
+        # we randomize the target position/omega in world frame
+        if env_indxs is None:
+            self._trgt_d.uniform_(self._env_opts["min_distance"], self._env_opts["max_distance"])
+            self._trgt_theta.uniform_(0.0, 2*torch.pi)
+
+            self._p_trgt_w[:, :]=self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu)[:, 0:2] +\
+                torch.cat((self._trgt_d*torch.cos(self._trgt_theta)
+                           ,self._trgt_d*torch.sin(self._trgt_theta)), dim=1)
+            
+            # randomize just omega
+            random_uniform=torch.full_like(self._agent_twist_ref_current_w[:, 3:6], fill_value=0.0)
+            torch.nn.init.uniform_(random_uniform, a=-1, b=1)
+            self._agent_twist_ref_current_w[:, 3:6] = random_uniform*self._twist_ref_scale[:, 3:6] + self._twist_ref_offset[:, 3:6]
+
+            # sample for all envs pof0
+            if self._env_opts["use_pof0"]: # sample from bernoulli distribution and update coefficients
+                torch.bernoulli(input=self._pof1_b_linvel,out=self._bernoulli_coeffs_linvel) # by default bernoulli_coeffs are 1 if not self._env_opts["use_pof0"]
+                torch.bernoulli(input=self._pof1_b_omega,out=self._bernoulli_coeffs_omega)
+                           
+        else:
+
+            if env_indxs.any():
+                integer_idxs=torch.nonzero(env_indxs).flatten()
+                
+                trgt_d_selected=self._trgt_d[integer_idxs, :]
+                trgt_d_selected.uniform_(self._env_opts["min_distance"], self._env_opts["max_distance"])
+                self._trgt_d[integer_idxs, :]=trgt_d_selected
+
+                trgt_theta_selected=self._trgt_theta[integer_idxs, :]
+                trgt_theta_selected.uniform_(0.0, 2*torch.pi)
+                self._trgt_theta[integer_idxs, :]=trgt_theta_selected
+
+                self._p_trgt_w[integer_idxs, 0:1]=self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu)[integer_idxs, 0:1] +\
+                    self._trgt_d[integer_idxs, :]*torch.cos(self._trgt_theta[integer_idxs, :])
+                self._p_trgt_w[integer_idxs, 1:2]=self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu)[integer_idxs, 1:2] +\
+                    self._trgt_d[integer_idxs, :]*torch.sin(self._trgt_theta[integer_idxs, :])
+
+                # randomize just omega
+                random_uniform=torch.full_like(self._agent_twist_ref_current_w[env_indxs, 3:6], fill_value=0.0)
+                torch.nn.init.uniform_(random_uniform, a=-1, b=1)
+                self._agent_twist_ref_current_w[env_indxs, 3:6] = random_uniform*self._twist_ref_scale[:, 3:6] + self._twist_ref_offset[:, 3:6]
+
+                # sample for all envs pof0, then reset to 1 for envs which are not to be randomized
+                if self._env_opts["use_pof0"]: # sample from bernoulli distribution and update coefficients
+                    torch.bernoulli(input=self._pof1_b_linvel,out=self._bernoulli_coeffs_linvel) # by default bernoulli_coeffs are 1 if not self._env_opts["use_pof0"]
+                    torch.bernoulli(input=self._pof1_b_omega,out=self._bernoulli_coeffs_omega)
+                    self._bernoulli_coeffs_linvel[~env_indxs, :]=1
+                    self._bernoulli_coeffs_omega[~env_indxs, :]=1
+
+        self._compute_twist_ref_w(env_indxs=env_indxs) # update linear vel twist refs based on pos error
+
+        

hybrid_quad_rhc.py

@@ -0,0 +1,1324 @@
+from mpc_hive.controllers.rhc import RHController
+
+from aug_mpc.controllers.rhc.horizon_based.horizon_imports import *
+
+from aug_mpc.controllers.rhc.horizon_based.hybrid_quad_rhc_refs import HybridQuadRhcRefs
+from aug_mpc.controllers.rhc.horizon_based.gait_manager import GaitManager
+
+from EigenIPC.PyEigenIPC import VLevel
+from EigenIPC.PyEigenIPC import Journal, LogType
+
+import numpy as np
+
+import os
+
+import time
+
+from typing import Dict, List
+
+class HybridQuadRhc(RHController):
+
+    def __init__(self, 
+            srdf_path: str,
+            urdf_path: str,
+            config_path: str,
+            robot_name: str, # used for shared memory namespaces
+            codegen_dir: str, 
+            n_nodes:float = 25,
+            injection_node:int = 10,
+            dt: float = 0.02,
+            max_solver_iter = 1, # defaults to rt-iteration
+            open_loop: bool = True,
+            close_loop_all: bool = False,
+            dtype = np.float32,
+            verbose = False, 
+            debug = False,
+            refs_in_hor_frame = True,
+            timeout_ms: int = 60000,
+            custom_opts: Dict = {}):
+
+        self._refs_in_hor_frame = refs_in_hor_frame
+
+        self._injection_node = injection_node
+
+        self._open_loop = open_loop
+        self._close_loop_all = close_loop_all
+
+        self._codegen_dir = codegen_dir
+        if not os.path.exists(self._codegen_dir):
+            os.makedirs(self._codegen_dir)
+        # else:
+        #     # Directory already exists, delete it and recreate
+        #     shutil.rmtree(self._codegen_dir)
+        #     os.makedirs(self._codegen_dir)
+
+        self.step_counter = 0
+        self.sol_counter = 0
+    
+        self.max_solver_iter = max_solver_iter
+        
+        self._timer_start = time.perf_counter()
+        self._prb_update_time = time.perf_counter()
+        self._phase_shift_time = time.perf_counter()
+        self._task_ref_update_time = time.perf_counter()
+        self._rti_time = time.perf_counter()
+
+        self.robot_name = robot_name
+        
+        self.config_path = config_path
+
+        self.urdf_path = urdf_path
+        # read urdf and srdf files
+        with open(self.urdf_path, 'r') as file:
+            self.urdf = file.read()
+        self._base_init = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0])
+
+        self._c_timelines = dict()
+        self._f_reg_timelines = dict()
+        
+        self._custom_opts={"replace_continuous_joints": False,
+            "use_force_feedback": False,
+            "lin_a_feedback": False,
+            "is_open_loop": self._open_loop, # fully open (just for db)
+            "fully_closed": False, # closed loop with full feedback (just for db)
+            "closed_partial": False, # closed loop with partial feedback
+            "adaptive_is": True, # closed loop with adaptation
+            "estimate_v_root": False, # when adaptive_is or closed_partial, estimate vbase
+            "alpha_from_outside": False, # alpha set ext. from shared memory
+            "alpha_half": 1.0, 
+            "only_vel_wheels": True, # whether wheels (if present) are just vel controlled
+            "use_jnt_v_feedback": False
+            }
+        
+        self._custom_opts.update(custom_opts)
+        
+        self._alpha_half=self._custom_opts["alpha_half"]
+
+        if self._open_loop:
+            self._custom_opts["fully_closed"]=False
+            self._custom_opts["adaptive_is"]=False
+            self._custom_opts["closed_partial"]=False
+        else:
+            self._custom_opts["is_open_loop"]=False
+            if self._custom_opts["fully_closed"]:
+                self._custom_opts["adaptive_is"]=False
+                self._custom_opts["closed_partial"]=False
+                self._custom_opts["lin_a_feedback"]=False
+            if self._custom_opts["closed_partial"]:
+                self._custom_opts["adaptive_is"]=False
+                self._custom_opts["fully_closed"]=False
+                self._custom_opts["lin_a_feedback"]=False
+            if self._custom_opts["adaptive_is"]:
+                self._custom_opts["closed_partial"]=False
+                self._custom_opts["fully_closed"]=False
+
+        super().__init__(srdf_path=srdf_path,
+                        n_nodes=n_nodes,
+                        dt=dt,
+                        namespace=self.robot_name,
+                        dtype=dtype,
+                        verbose=verbose, 
+                        debug=debug,
+                        timeout_ms=timeout_ms)
+
+        self._rhc_fpaths.append(self.config_path)
+
+    def _config_override(self):
+        pass
+
+    def _post_problem_init(self):
+
+        self.rhc_costs={}
+        self.rhc_constr={}
+
+        self._fail_idx_scale=0.0
+        self._expl_idx_window_size=int(1*self._n_nodes) 
+        self._explosion_idx_buffer=np.zeros((1,self._expl_idx_window_size))
+        self._expl_idx_counter=0
+        self._expl_idx_buffer_counter=0
+
+        self._pred_node_idx=self._n_nodes-1
+
+        self._nq=self.nq()
+        self._nq_jnts=self._nq-7# assuming floating base
+        self._nv=self.nv()
+        self._nv_jnts=self._nv-6
+
+        self._alphas_q_root=np.zeros((7, 1), dtype=self._dtype)
+        self._alphas_q_jnts=np.zeros((self._nq_jnts, 1), dtype=self._dtype)
+        self._alphas_twist_root=np.zeros((6, 1), dtype=self._dtype)
+        self._alphas_v_jnts=np.zeros((self._nv_jnts, 1), dtype=self._dtype)
+        self._alphas_a=np.zeros((self._nv, 1), dtype=self._dtype)
+        self._alphas_q_root[:, :]=1.0 # default to all open
+        self._alphas_q_jnts[:, :]=1.0 
+        self._alphas_twist_root[:, :]=1.0 
+        self._alphas_v_jnts[:, :]=1.0 
+        self._alphas_a[:, :]=1.0
+
+    def _init_problem(self,
+            fixed_jnt_patterns: List[str] = None,
+            wheels_patterns: List[str] = None,
+            foot_linkname: str = None,
+            flight_duration: int = 10,
+            post_flight_stance: int = 3,
+            step_height: float = 0.12,
+            keep_yaw_vert: bool = False,
+            yaw_vertical_weight: float = 2.0,
+            vertical_landing: bool = False,
+            vertical_land_weight: float = 1.0,
+            phase_force_reg: float = 1e-2,
+            vel_bounds_weight: float = 1.0):
+        
+        self._fixed_jnt_patterns=fixed_jnt_patterns
+
+        self._config_override()
+        
+        Journal.log(self.__class__.__name__,
+            "_init_problem",
+            f" Will use horizon config file at {self.config_path}",
+            LogType.INFO,
+            throw_when_excep=True)
+        
+        self._vel_bounds_weight=vel_bounds_weight
+        self._phase_force_reg=phase_force_reg
+        self._yaw_vertical_weight=yaw_vertical_weight
+        self._vertical_land_weight=vertical_land_weight
+        # overrides parent
+        self._prb = Problem(self._n_intervals, 
+                        receding=True, 
+                        casadi_type=cs.SX)
+        self._prb.setDt(self._dt)
+
+        if "replace_continuous_joints" in self._custom_opts:
+            # continous joints are parametrized in So2
+            if self._custom_opts["replace_continuous_joints"]:
+                self.urdf = self.urdf.replace('continuous', 'revolute')
+        else:
+            self.urdf = self.urdf.replace('continuous', 'revolute')
+        
+        self._kin_dyn = casadi_kin_dyn.CasadiKinDyn(self.urdf) # used for getting joint names 
+        self._assign_controller_side_jnt_names(jnt_names=self._get_robot_jnt_names())
+        
+        self._init_robot_homer()
+
+        # handle fixed joints
+        fixed_joint_map={}
+        if self._fixed_jnt_patterns is not None:
+            for jnt_name in self._get_robot_jnt_names():
+                for fixed_jnt_pattern in self._fixed_jnt_patterns:
+                    if fixed_jnt_pattern in jnt_name: 
+                        fixed_joint_map.update({f"{jnt_name}":
+                            self._homer.get_homing_val(jnt_name=jnt_name)})
+                        break # do not search for other pattern matches
+        
+        if not len(fixed_joint_map)==0: # we need to recreate kin dyn and homers
+            Journal.log(self.__class__.__name__,
+                "_init_problem",
+                f"Will fix following joints: \n{str(fixed_joint_map)}",
+                LogType.INFO,
+                throw_when_excep=True)
+            # with the fixed joint map
+            self._kin_dyn = casadi_kin_dyn.CasadiKinDyn(self.urdf,fixed_joints=fixed_joint_map)
+            # assign again controlled joints names
+            self._assign_controller_side_jnt_names(jnt_names=self._get_robot_jnt_names())
+            # updated robot homer for controlled joints
+            self._init_robot_homer()
+
+        # handle continuous joints (need to change homing and retrieve
+        # cont jnts indexes) and homing
+        self._continuous_joints=self._get_continuous_jnt_names()
+        # reduced
+        self._continuous_joints_idxs=[]
+        self._continuous_joints_idxs_cos=[]
+        self._continuous_joints_idxs_sin=[]
+        self._continuous_joints_idxs_red=[]
+        self._rev_joints_idxs=[]
+        self._rev_joints_idxs_red=[]
+        # qfull
+        self._continuous_joints_idxs_qfull=[]
+        self._continuous_joints_idxs_cos_qfull=[]
+        self._continuous_joints_idxs_sin_qfull=[]
+        self._continuous_joints_idxs_red_qfull=[]
+        self._rev_joints_idxs_qfull=[]
+        self._rev_joints_idxs_red_qfull=[]
+        jnt_homing=[""]*(len(self._homer.get_homing())+len(self._continuous_joints))
+        jnt_names=self._get_robot_jnt_names()
+        for i in range(len(jnt_names)):
+            jnt=jnt_names[i]
+            index=self._get_jnt_id(jnt)# accounting for floating joint
+            homing_idx=index-7 # homing is only for actuated joints
+            homing_value=self._homer.get_homing_val(jnt)
+            if jnt in self._continuous_joints:
+                jnt_homing[homing_idx]=np.cos(homing_value).item()
+                jnt_homing[homing_idx+1]=np.sin(homing_value).item()
+                # just actuated joints
+                self._continuous_joints_idxs.append(homing_idx) # cos
+                self._continuous_joints_idxs.append(homing_idx+1) # sin
+                self._continuous_joints_idxs_cos.append(homing_idx)
+                self._continuous_joints_idxs_sin.append(homing_idx+1)
+                self._continuous_joints_idxs_red.append(i)
+                # q full
+                self._continuous_joints_idxs_qfull.append(index) # cos
+                self._continuous_joints_idxs_qfull.append(index+1) # sin
+                self._continuous_joints_idxs_cos_qfull.append(index)
+                self._continuous_joints_idxs_sin_qfull.append(index+1)
+                self._continuous_joints_idxs_red_qfull.append(i+7)
+            else:
+                jnt_homing[homing_idx]=homing_value
+                # just actuated joints
+                self._rev_joints_idxs.append(homing_idx) 
+                self._rev_joints_idxs_red.append(i) 
+                # q full
+                self._rev_joints_idxs_qfull.append(index) 
+                self._rev_joints_idxs_red_qfull.append(i+7) 
+
+        self._jnts_q_reduced=None
+        if not len(self._continuous_joints)==0: 
+            cont_joints=", ".join(self._continuous_joints)
+
+            Journal.log(self.__class__.__name__,
+                "_init_problem",
+                f"The following continuous joints were found: \n{cont_joints}",
+                LogType.INFO,
+                throw_when_excep=True)
+            # preallocating data 
+            self._jnts_q_reduced=np.zeros((1,self.nv()-6),dtype=self._dtype)
+            self._jnts_q_expanded=np.zeros((self.nq()-7,1),dtype=self._dtype)
+            self._full_q_reduced=np.zeros((7+len(jnt_names), self._n_nodes),dtype=self._dtype)
+            self._jnts_q_delta_expanded=np.zeros((self.nq()-7,1),dtype=self._dtype)
+        else:
+            self._custom_opts["replace_continuous_joints"]=True
+            Journal.log(self.__class__.__name__,
+                "_init_problem",
+                f"No continuous joints were found.",
+                LogType.INFO,
+                throw_when_excep=True)
+
+        # retrieve wheels indexes (not considering continuous joints, 
+        # ok just for v, eff, etc..)
+        self._wheel_patterns=wheels_patterns
+        self._wheels_idxs_v=self._get_wheels_jnt_v_idxs(wheel_patterns=self._wheel_patterns)
+        self._f0 = [0, 0, self._kin_dyn.mass()/4*9.81]
+        
+        # we can create an init for the base
+        init = self._base_init.tolist() + jnt_homing
+
+        if foot_linkname is not None:
+            FK = self._kin_dyn.fk(foot_linkname) # just to get robot reference height
+            ground_level = FK(q=init)['ee_pos']
+            self._base_init[2] = -ground_level[2]  # override init
+        
+        self._model = FullModelInverseDynamics(problem=self._prb,
+            kd=self._kin_dyn,
+            q_init=self._homer.get_homing_map(),
+            base_init=self._base_init)
+
+        self._ti = TaskInterface(prb=self._prb, 
+                            model=self._model, 
+                            max_solver_iter=self.max_solver_iter,
+                            debug = self._debug,
+                            verbose = self._verbose, 
+                            codegen_workdir = self._codegen_dir)
+        self._ti.setTaskFromYaml(self.config_path)
+        
+        # setting initial base pos ref if exists
+        base_pos = self._ti.getTask('base_height')
+        if base_pos is not None:
+            base_pos.setRef(np.atleast_2d(self._base_init).T)
+
+        self._pm = pymanager.PhaseManager(self._n_nodes, debug=False) # intervals or nodes?????
+
+        self._gm = GaitManager(self._ti, 
+            self._pm, 
+            self._injection_node,
+            keep_yaw_vert=keep_yaw_vert,
+            yaw_vertical_weight=self._yaw_vertical_weight,
+            vertical_landing=vertical_landing,
+            landing_vert_weight=self._vertical_land_weight,
+            phase_force_reg=self._phase_force_reg,
+            custom_opts=self._custom_opts,
+            flight_duration=flight_duration,
+            post_flight_stance=post_flight_stance,
+            step_height=step_height,
+            dh=0.0)
+                
+        self._ti.model.q.setBounds(self._ti.model.q0, self._ti.model.q0, nodes=0)
+        self._ti.model.v.setBounds(self._ti.model.v0, self._ti.model.v0, nodes=0)
+        self._ti.model.q.setInitialGuess(self._ti.model.q0)
+        self._ti.model.v.setInitialGuess(self._ti.model.v0)
+        for _, cforces in self._ti.model.cmap.items():
+            n_contact_f=len(cforces)
+            for c in cforces:
+                c.setInitialGuess(np.array(self._f0)/n_contact_f)        
+
+        vel_lims = self._model.kd.velocityLimits()
+        import horizon.utils as utils
+        self._prb.createResidual('vel_lb_barrier', self._vel_bounds_weight*utils.utils.barrier(vel_lims[7:] - self._model.v[7:]))
+        self._prb.createResidual('vel_ub_barrier', self._vel_bounds_weight*utils.utils.barrier1(-1 * vel_lims[7:] - self._model.v[7:]))
+
+        self._meas_lin_a_par=None
+        # if self._custom_opts["lin_a_feedback"]:
+        #     # acceleration feedback on first node
+        #     self._meas_lin_a_par=self._prb.createParameter(name="lin_a_feedback",
+        #         dim=3, nodes=0)   
+        #     base_lin_a_prb=self._prb.getInput().getVars()[0:3] 
+        #     self._prb.createConstraint('lin_acceleration_feedback', base_lin_a_prb - self._meas_lin_a_par, 
+        #             nodes=[0])
+
+        # if not self._open_loop:
+        #     # we create a residual cost to be used as an attractor to the measured state on the first node
+        #     # hard constraints injecting meas. states are pure EVIL!
+        #     prb_state=self._prb.getState()
+        #     full_state=prb_state.getVars()
+        #     state_dim=prb_state.getBounds()[0].shape[0]
+        #     meas_state=self._prb.createParameter(name="measured_state",
+        #         dim=state_dim, nodes=0)     
+        #     self._prb.createResidual('meas_state_attractor', meas_state_attractor_weight * (full_state - meas_state), 
+        #                 nodes=[0])
+
+        self._ti.finalize()
+        self._ti.bootstrap()
+
+        self._ti.init_inv_dyn_for_res() # we initialize some objects for sol. postprocessing purposes
+        self._ti.load_initial_guess()
+
+        self.n_dofs = self._get_ndofs() # after loading the URDF and creating the controller we
+        # know n_dofs -> we assign it (by default = None)
+
+        self.n_contacts = len(self._model.cmap.keys())
+        
+        # remove variables bounds (before they were necessary to have a good
+        # quality bootstrap solution)
+        self._q_inf=np.zeros((self.nq(), 1))
+        self._q_inf[:, :]=np.inf
+        self._v_inf=np.zeros((self.nv(), 1))
+        self._v_inf[:, :]=np.inf
+        self._ti.model.q.setBounds(-self._q_inf, self._q_inf, nodes=0)
+        self._ti.model.v.setBounds(-self._v_inf, self._v_inf, nodes=0)
+
+        # self.horizon_anal = analyzer.ProblemAnalyzer(self._prb)
+
+    def get_file_paths(self):
+        # can be overriden by child
+        paths = super().get_file_paths()
+        return paths
+    
+    def _get_quat_remap(self):
+        # overrides parent
+        return [1, 2, 3, 0] # mapping from robot quat. to Horizon's quaternion convention
+    
+    def _zmp(self, model):
+
+        num = cs.SX([0, 0])
+        den = cs.SX([0])
+        pos_contact = dict()
+        force_val = dict()
+
+        q = cs.SX.sym('q', model.nq)
+        v = cs.SX.sym('v', model.nv)
+        a = cs.SX.sym('a', model.nv)
+
+        com = model.kd.centerOfMass()(q=q, v=v, a=a)['com']
+
+        n = cs.SX([0, 0, 1])
+        for c in model.fmap.keys():
+            pos_contact[c] = model.kd.fk(c)(q=q)['ee_pos']
+            force_val[c] = cs.SX.sym('force_val', 3)
+            num += (pos_contact[c][0:2] - com[0:2]) * cs.dot(force_val[c], n)
+            den += cs.dot(force_val[c], n)
+
+        zmp = com[0:2] + (num / den)
+        input_list = []
+        input_list.append(q)
+        input_list.append(v)
+        input_list.append(a)
+
+        for elem in force_val.values():
+            input_list.append(elem)
+
+        f = cs.Function('zmp', input_list, [zmp])
+        return f
+    
+    def _add_zmp(self):
+
+        input_zmp = []
+
+        input_zmp.append(self._model.q)
+        input_zmp.append(self._model.v)
+        input_zmp.append(self._model.a)
+
+        for f_var in self._model.fmap.values():
+            input_zmp.append(f_var)
+
+        c_mean = cs.SX([0, 0, 0])
+        for c_name, f_var in self._model.fmap.items():
+            fk_c_pos = self._kin_dyn.fk(c_name)(q=self._model.q)['ee_pos']
+            c_mean += fk_c_pos
+
+        c_mean /= len(self._model.cmap.keys())
+
+        zmp_nominal_weight = 10.
+        zmp_fun = self._zmp(self._model)(*input_zmp)
+
+        if 'wheel_joint_1' in self._model.kd.joint_names():
+            zmp_residual = self._prb.createIntermediateResidual('zmp',  zmp_nominal_weight * (zmp_fun[0:2] - c_mean[0:2]))
+
+    def _quaternion_multiply(self, 
+                    q1, q2):
+        x1, y1, z1, w1 = q1
+        x2, y2, z2, w2 = q2
+        w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
+        x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
+        y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
+        z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
+        return np.array([x, y, z, w])
+    
+    def _get_continuous_jnt_names(self):
+        import xml.etree.ElementTree as ET
+        root = ET.fromstring(self.urdf)
+        continuous_joints = []
+        for joint in root.findall('joint'):
+            joint_type = joint.get('type')
+            if joint_type == 'continuous':
+                joint_name = joint.get('name')
+                continuous_joints.append(joint_name)
+        return continuous_joints
+    
+    def _get_wheels_jnt_v_idxs(self, wheel_patterns: List[str]):
+        jnt_names=self._get_robot_jnt_names()
+        wheels_idxs=[]
+        for i in range(len(jnt_names)):
+            jnt_name=jnt_names[i]
+            for wheel_pattern in wheel_patterns:
+                if wheel_pattern in jnt_name:
+                    wheels_idxs.append(i)
+                    break
+        return wheels_idxs
+    
+    def _get_jnt_id(self, jnt_name):
+        return self._kin_dyn.joint_iq(jnt_name)
+    
+    def _init_rhc_task_cmds(self):
+        
+        rhc_refs = HybridQuadRhcRefs(gait_manager=self._gm,
+            robot_index_shm=self.controller_index,
+            robot_index_view=0, # when using optimize_mem the view if always of shape 1x...
+            namespace=self.namespace,
+            safe=False, 
+            verbose=self._verbose,
+            vlevel=VLevel.V2,
+            use_force_feedback=self._custom_opts["use_force_feedback"],
+            optimize_mem=True)
+        
+        rhc_refs.run()
+
+        rhc_refs.rob_refs.set_jnts_remapping(jnts_remapping=self._to_controller)
+        rhc_refs.rob_refs.set_q_remapping(q_remapping=self._get_quat_remap())
+        
+        rhc_refs.set_default_refs(p_ref=np.atleast_2d(self._base_init)[:, 0:3],
+            q_ref=np.atleast_2d(self._base_init)[:, 3:7])
+        
+        return rhc_refs
+    
+    def get_vertex_fnames_from_ti(self):
+        tasks=self._ti.task_list
+        contact_f_names=[]
+        for task in tasks:
+            if isinstance(task, ContactTask):
+                interaction_task=task.dynamics_tasks[0]
+                contact_f_names.append(interaction_task.vertex_frames[0])
+        return contact_f_names
+        
+    def _get_contact_names(self):
+        # should get contact names from vertex frames
+        # list(self._ti.model.cmap.keys())
+        return self.get_vertex_fnames_from_ti()
+    
+    def _get_robot_jnt_names(self):
+
+        joints_names = self._kin_dyn.joint_names()
+        to_be_removed = ["universe", 
+                        "reference", 
+                        "world", 
+                        "floating", 
+                        "floating_base"]
+        for name in to_be_removed:
+            if name in joints_names:
+                joints_names.remove(name)
+
+        return joints_names
+    
+    def _get_ndofs(self):
+        return len(self._model.joint_names)
+    
+    def nq(self):
+        return self._kin_dyn.nq()
+    
+    def nv(self):
+        return self._kin_dyn.nv()
+    
+    def _get_robot_mass(self):
+
+        return self._kin_dyn.mass()
+
+    def _get_root_full_q_from_sol(self, node_idx=1):
+
+        root_q_full=self._ti.solution['q'][0:7, node_idx].reshape(1, 7).astype(self._dtype)
+
+        np.nan_to_num(root_q_full, nan=1e3, posinf=1e3, neginf=-1e3, copy=False)
+        np.clip(a=root_q_full, a_min=-1e3, a_max=1e3, out=root_q_full)
+
+        return root_q_full
+    
+    def _get_full_q_from_sol(self, node_idx=1):
+
+        return self._ti.solution['q'][:, node_idx].reshape(1, -1).astype(self._dtype)
+    
+    def _get_root_twist_from_sol(self, node_idx=1):
+        # provided in world frame
+        twist_base_local=self._get_v_from_sol()[0:6, node_idx].reshape(1, 6)
+        # if world_aligned:
+        #     q_root_rhc = self._get_root_full_q_from_sol(node_idx=node_idx)[:, 0:4]
+        #     r_base_rhc=Rotation.from_quat(q_root_rhc.flatten()).as_matrix()
+        #     twist_base_local[:, 0:3] = r_base_rhc @ twist_base_local[0, 0:3]
+        #     twist_base_local[:, 3:6] = r_base_rhc @ twist_base_local[0, 3:6]
+        return twist_base_local
+
+    def _get_root_a_from_sol(self, node_idx=0):
+        # provided in world frame
+        a_base_local=self._get_a_from_sol()[0:6, node_idx].reshape(1, 6)
+        # if world_aligned:
+        #     q_root_rhc = self._get_root_full_q_from_sol(node_idx=node_idx)[:, 0:4]
+        #     r_base_rhc=Rotation.from_quat(q_root_rhc.flatten()).as_matrix()
+        #     a_base_local[:, 0:3] = r_base_rhc @ a_base_local[0, 0:3]
+        #     a_base_local[:, 3:6] = r_base_rhc @ v[0, 3:6]
+        return a_base_local
+    
+    def _get_jnt_q_from_sol(self, node_idx=0, 
+            reduce: bool = True,
+            clamp: bool = True):
+        
+        full_jnts_q=self._ti.solution['q'][7:, node_idx:node_idx+1].reshape(1,-1).astype(self._dtype)
+        
+        np.nan_to_num(full_jnts_q, nan=1e3, posinf=1e3, neginf=-1e3, copy=False) # in place
+        np.clip(a=full_jnts_q, a_max=1e3, a_min=-1e3, out=full_jnts_q) # in place
+
+        if self._custom_opts["replace_continuous_joints"] or (not reduce):
+            if clamp:
+                return np.fmod(full_jnts_q, 2*np.pi)
+            else:
+                return full_jnts_q
+        else:
+            cos_sin=full_jnts_q[:,self._continuous_joints_idxs].reshape(-1,2)
+            # copy rev joint vals
+            self._jnts_q_reduced[:, self._rev_joints_idxs_red]=np.fmod(full_jnts_q[:, self._rev_joints_idxs], 2*np.pi).reshape(1, -1)
+            # and continuous
+            self._jnts_q_reduced[:, self._continuous_joints_idxs_red]=np.arctan2(cos_sin[:, 1], cos_sin[:, 0]).reshape(1,-1)
+            return self._jnts_q_reduced
+        
+    def _get_jnt_v_from_sol(self, node_idx=1):
+        
+        jnt_v_sol=self._get_v_from_sol()[6:, node_idx].reshape(1,  
+                    self._nv_jnts)
+        np.nan_to_num(jnt_v_sol, nan=1e5, posinf=1e5, neginf=-1e5, copy=False) # in place
+        # np.clip(a=jnt_v_sol, a_max=1e5, a_min=-1e5, out=jnt_v_sol) # in place
+
+        return jnt_v_sol
+
+    def _get_jnt_a_from_sol(self, node_idx=1):
+
+        return self._get_a_from_sol()[6:, node_idx].reshape(1,
+                    self._nv_jnts)
+
+    def _get_jnt_eff_from_sol(self, node_idx=0):
+        
+        efforts_on_node = self._ti.eval_efforts_on_node(node_idx=node_idx)
+        
+        # if self._custom_opts["only_vel_wheels"]:
+
+        jnt_efforts=efforts_on_node[6:, 0]
+        
+        if self._custom_opts["only_vel_wheels"] and self._wheels_idxs_v:
+            jnt_efforts[self._wheels_idxs_v]=0.0
+
+        return jnt_efforts.reshape(1,  
+                self._nv_jnts).astype(self._dtype)
+    
+    def _get_rhc_cost(self):
+
+        return self._ti.solution["opt_cost"]
+    
+    def _get_rhc_constr_viol(self):
+        
+        return self._ti.solution["residual_norm"]
+    
+    def _get_rhc_nodes_cost(self):
+
+        cost = self._ti.solver_rti.getCostValOnNodes()
+        return cost.reshape((1, -1))
+    
+    def _get_rhc_nodes_constr_viol(self):
+        
+        constr_viol = self._ti.solver_rti.getConstrValOnNodes()
+        return constr_viol.reshape((1, -1))
+    
+    def _get_rhc_niter_to_sol(self):
+
+        return self._ti.solution["n_iter2sol"]
+    
+    def _set_ig_bootstrap(self,
+            q_state: np.ndarray = None,
+            v_state: np.ndarray = None):
+
+        xig = self._ti.solution['x_opt'].copy()
+        uig = self._ti.solution['u_opt'].copy()
+
+        # Normalize and keep quaternion in the same hemisphere as the previous
+        # solution to avoid artificial 180-deg jumps in the bootstrap warm start.
+        q_state_boot = q_state.copy()
+        q_prev = xig[3:7, 0]
+        q_now = q_state_boot[3:7, 0]
+
+        q_now_norm = np.linalg.norm(q_now)
+        if q_now_norm > 1e-9:
+            q_state_boot[3:7, :] /= q_now_norm
+        else:
+            q_state_boot[3:7, :] = np.array([[0.0], [0.0], [0.0], [1.0]], dtype=self._dtype)
+
+        q_prev_norm = np.linalg.norm(q_prev)
+        if q_prev_norm > 1e-9:
+            q_prev = q_prev / q_prev_norm
+
+        q_now = q_state_boot[3:7, 0]
+        if np.dot(q_prev, q_now) < 0.0:
+            q_state_boot[3:7, :] *= -1.0
+        
+        xig[0:self._nq, :] = q_state_boot
+        xig[self._nq:self._nq + self._nv, :] = 0.0 # 0 velocity on first nodes 
+        uig[0:self._nv, :]=0.0 # 0 acceleration
+
+        # assigning ig
+        self._prb.getState().setInitialGuess(xig)
+        self._prb.getInput().setInitialGuess(uig)
+        # self._prb.getVariables("a").setInitialGuess(np.zeros((self._nv, 1), dtype=self._dtype))
+        for _, cforces in self._ti.model.cmap.items():
+            n_contact_f = len(cforces)
+            if n_contact_f == 0:
+                continue
+            f_guess = np.array(self._f0, dtype=self._dtype) / n_contact_f
+            for c in cforces:
+                c.setInitialGuess(f_guess)
+
+        # print("initial guesses")
+        # print(self._nq)
+        # print(self._nv)
+        # print("q")
+        # qig=self._ti.model.q.getInitialGuess()
+        # print(qig.shape)
+        # print(qig)
+        # print("v")
+        # print(self._ti.model.v.getInitialGuess())
+        # print("a")
+        # print(self._ti.model.a.getInitialGuess())
+        # for _, cforces in self._ti.model.cmap.items():
+        #     for c in cforces:
+        #         print("force")
+        #         print(c.getInitialGuess())
+        
+        return xig, uig
+    
+    def _set_ig(self):
+
+        shift_num = -1 # shift data by one node
+
+        x_opt = self._ti.solution['x_opt']
+        u_opt = self._ti.solution['u_opt']
+
+        # building ig for state
+        xig = np.roll(x_opt, 
+                shift_num, axis=1) # rolling state sol.
+        for i in range(abs(shift_num)):
+            # state on last node is copied to the elements
+            # which are "lost" during the shift operation
+            xig[:, -1 - i] = x_opt[:, -1]
+        # building ig for inputs
+        uig = np.roll(u_opt, 
+                shift_num, axis=1) # rolling state sol.
+        for i in range(abs(shift_num)):
+            # state on last node is copied to the elements
+            # which are "lost" during the shift operation
+            uig[:, -1 - i] = u_opt[:, -1]
+
+        # assigning ig
+        self._prb.getState().setInitialGuess(xig)
+        self._prb.getInput().setInitialGuess(uig)
+        
+        return xig, uig
+    
+    def _update_open_loop(self,
+            bootstrap: bool = False):
+
+        q_state, v_state, a_state=self._set_is_open()
+
+        if not bootstrap:
+            self._set_ig()
+        else:
+            self._set_ig_bootstrap(q_state=q_state, v_state=v_state)
+
+        # robot_state=xig[:, 0]
+        # # open loop update:
+        # self._prb.setInitialState(x0=robot_state) # (xig has been shifted, so node 0
+        # # is node 1 in the last opt solution)
+
+        return q_state, v_state, a_state
+    
+    def _update_closed_loop(self,
+            bootstrap: bool = False):
+        
+        # set initial state
+        q_state=None
+        v_state=None
+        a_state=None
+        if self._custom_opts["adaptive_is"]:
+            # adaptive closed loop
+            q_state, v_state, a_state=self._set_is_adaptive()
+        elif self._custom_opts["fully_closed"]:
+            q_state, v_state, a_state=self._set_is_full()
+        elif self._custom_opts["closed_partial"]:
+            q_state, v_state, a_state=self._set_is_partial()
+        else:
+            Journal.log(self.__class__.__name__,
+                    "_update_closed_loop",
+                    "Neither adaptive_is, fully_closed, or closed_partial.",
+                    LogType.EXCEP,
+                    throw_when_excep = False)
+            q_state, v_state, a_state=self._set_is()
+
+        # set initial guess for controller
+        if not bootstrap:
+            self._set_ig()
+        else:
+            self._set_ig_bootstrap(q_state=q_state, v_state=v_state)
+
+        return q_state, v_state, a_state
+    
+    def _set_is_open(self):
+        
+        # overriding states with rhc data
+        q_full_root=self._get_root_full_q_from_sol(node_idx=1).reshape(-1, 1)
+        q_jnts=self._get_jnt_q_from_sol(node_idx=1, reduce=False).reshape(-1, 1)
+        
+        twist_root=self._get_root_twist_from_sol(node_idx=1).reshape(-1, 1)
+        v_jnts=self._get_jnt_v_from_sol(node_idx=1).reshape(-1, 1)
+
+        # rhc variables to be set
+        q=self._prb.getVariables("q") # .setBounds()
+        root_q_full_rhc=q[0:7] # root full q
+        jnts_q_rhc=q[7:] # jnts q
+        vel=self._prb.getVariables("v")
+        root_twist_rhc=vel[0:6] # lin v.
+        jnts_v_rhc=vel[6:] # jnts v
+
+        self.rhc_refs.set_alpha(alpha=1.0) # fully open
+
+        # close state on known quantities
+        root_q_full_rhc.setBounds(lb=q_full_root,
+            ub=q_full_root, nodes=0)
+        jnts_q_rhc.setBounds(lb=q_jnts, 
+            ub=q_jnts, nodes=0)
+        root_twist_rhc.setBounds(lb=twist_root, 
+            ub=twist_root, nodes=0)
+        jnts_v_rhc.setBounds(lb=v_jnts, 
+            ub=v_jnts, nodes=0)
+        
+        # return state used for feedback
+        q_state=np.concatenate((q_full_root, q_jnts),
+                axis=0)
+        v_state=np.concatenate((twist_root, v_jnts),
+                axis=0)
+        
+        return (q_state, v_state, None)
+    
+    def _set_is_full(self):
+        
+        # measurements
+        q_full_root = self.robot_state.root_state.get(data_type="q_full", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        v_root = self.robot_state.root_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        omega = self.robot_state.root_state.get(data_type="omega", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        a_root = self.robot_state.root_state.get(data_type="a_full", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        
+        q_jnts = self.robot_state.jnts_state.get(data_type="q", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        v_jnts = self.robot_state.jnts_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        a_jnts = self.robot_state.jnts_state.get(data_type="a", robot_idxs=self.controller_index_np).reshape(-1, 1)
+
+        if (not len(self._continuous_joints)==0): # we need do expand some meas. rev jnts to So2
+            self._jnts_q_expanded[self._rev_joints_idxs, :]=q_jnts[self._rev_joints_idxs_red ,:]
+            self._jnts_q_expanded[self._continuous_joints_idxs_cos, :]=np.cos(q_jnts[self._continuous_joints_idxs_red, :]) # cos
+            self._jnts_q_expanded[self._continuous_joints_idxs_sin, :]=np.sin(q_jnts[self._continuous_joints_idxs_red, :]) # sin
+            q_jnts=self._jnts_q_expanded.reshape(-1,1)
+            
+        # rhc variables to be set
+        q=self._prb.getVariables("q") # .setBounds()
+        root_full_q_rhc=q[0:7] # root p
+        jnts_q_rhc=q[7:] # jnts q
+        vel=self._prb.getVariables("v")
+        root_v_rhc=vel[0:3] # lin v.
+        root_omega_rhc=vel[3:6] # omega
+        jnts_v_rhc=vel[6:] # jnts v
+        acc=self._prb.getVariables("a")
+        lin_a_prb=acc[0:3] # lin acc
+        
+        self.rhc_refs.set_alpha(alpha=0.0) # non-adaptive
+
+        root_full_q_rhc.setBounds(lb=q_full_root,
+            ub=q_full_root, nodes=0)
+        jnts_q_rhc.setBounds(lb=q_jnts, 
+            ub=q_jnts, nodes=0)
+        root_v_rhc.setBounds(lb=v_root, 
+            ub=v_root, nodes=0) # leaving lin v of the base free (estimated from constraints)
+        root_omega_rhc.setBounds(lb=omega, 
+            ub=omega, nodes=0)
+        jnts_v_rhc.setBounds(lb=v_jnts, 
+            ub=v_jnts, nodes=0)
+        if self._custom_opts["lin_a_feedback"]:
+            # write base lin 13793197 from meas
+            lin_a_prb.setBounds(lb=a_root[0:3, :], 
+                ub=a_root[0:3, :], 
+                nodes=0)
+
+        # return state used for feedback
+        q_state=np.concatenate((q_full_root, q_jnts),
+                axis=0)
+        v_state=np.concatenate((v_root, omega, v_jnts),
+                axis=0)
+        a_state=np.concatenate((a_root, a_jnts),
+                axis=0)
+        
+        return (q_state, v_state, a_state)
+    
+    def _set_is_partial(self):
+        
+        # measurements
+        p_root = self.robot_state.root_state.get(data_type="p", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        q_root = self.robot_state.root_state.get(data_type="q", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        v_root = self.robot_state.root_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        omega = self.robot_state.root_state.get(data_type="omega", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        a_root = self.robot_state.root_state.get(data_type="a_full", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        
+        q_jnts = self.robot_state.jnts_state.get(data_type="q", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        v_jnts = self.robot_state.jnts_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        a_jnts = self.robot_state.jnts_state.get(data_type="a", robot_idxs=self.controller_index_np).reshape(-1, 1)
+
+        if (not len(self._continuous_joints)==0): # we need do expand some meas. rev jnts to So2
+            self._jnts_q_expanded[self._rev_joints_idxs, :]=q_jnts[self._rev_joints_idxs_red ,:]
+            self._jnts_q_expanded[self._continuous_joints_idxs_cos, :]=np.cos(q_jnts[self._continuous_joints_idxs_red, :]) # cos
+            self._jnts_q_expanded[self._continuous_joints_idxs_sin, :]=np.sin(q_jnts[self._continuous_joints_idxs_red, :]) # sin
+            q_jnts=self._jnts_q_expanded.reshape(-1,1)
+
+        # overriding states with rhc data (-> all overridden state are open looop)
+        root_q_full_from_rhc=self._get_root_full_q_from_sol(node_idx=1).reshape(-1, 1)
+        root_p_from_rhc=root_q_full_from_rhc[0:3, :]
+        p_root[:, :]=root_p_from_rhc # position is always open loop
+        
+        if not self._custom_opts["estimate_v_root"]:
+            v_root[:, :]=self._get_root_twist_from_sol(node_idx=1)[:, 0:3].reshape(-1, 1)
+            # override v jnts with the ones from controller
+            if not self._custom_opts["use_jnt_v_feedback"]:
+                v_jnts[:, :]=self._get_jnt_v_from_sol(node_idx=1).reshape(-1, 1)
+            # v_jnts[:, :]=self._get_jnt_v_from_sol(node_idx=1).reshape(-1, 1)
+        # root_twist_from_rhc=self._get_root_twist_from_sol(node_idx=1)
+        # root_v_from_rhc=root_twist_from_rhc[:, 0:3].reshape(-1, 1)
+        # root_omega_from_rhc=root_twist_from_rhc[:, 3:6].reshape(-1, 1)
+        # jnt_q_from_rhc=self._get_jnt_q_from_sol(node_idx=1,reduce=False,clamp=False).reshape(-1, 1)
+        # jnt_v_from_rhc=self._get_jnt_v_from_sol(node_idx=1).reshape(-1, 1)
+            
+        # rhc variables to be set
+        q=self._prb.getVariables("q") # .setBounds()
+        root_p_rhc=q[0:3] # root p
+        root_q_rhc=q[3:7] # root orientation
+        jnts_q_rhc=q[7:] # jnts q
+        vel=self._prb.getVariables("v")
+        root_v_rhc=vel[0:3] # lin v.
+        root_omega_rhc=vel[3:6] # omega
+        jnts_v_rhc=vel[6:] # jnts v
+        acc=self._prb.getVariables("a")
+        lin_a_prb=acc[0:3] # lin acc
+        
+        self.rhc_refs.set_alpha(alpha=0.0) # non-adaptive
+
+        root_p_rhc.setBounds(lb=p_root,
+            ub=p_root, nodes=0)
+        root_q_rhc.setBounds(lb=q_root, 
+            ub=q_root, nodes=0)
+        jnts_q_rhc.setBounds(lb=q_jnts, 
+            ub=q_jnts, nodes=0)
+        
+        if self._custom_opts["estimate_v_root"]:
+            root_v_rhc.setBounds(lb=-self._v_inf[0:3], 
+                ub=self._v_inf[0:3], nodes=0) # leaving lin v of the base free (estimated from constraints)
+        else: # get it from controller 
+            root_v_rhc.setBounds(lb=v_root, 
+                ub=v_root, nodes=0)
+        root_omega_rhc.setBounds(lb=omega, 
+            ub=omega, nodes=0)
+        jnts_v_rhc.setBounds(lb=v_jnts, 
+            ub=v_jnts, nodes=0)
+        if self._custom_opts["lin_a_feedback"]:
+            # write base lin 13793197 from meas
+            lin_a_prb.setBounds(lb=a_root[0:3, :], 
+                ub=a_root[0:3, :], 
+                nodes=0)
+
+        # return state used for feedback
+        q_state=np.concatenate((p_root, q_root, q_jnts),
+                axis=0)
+        v_state=np.concatenate((v_root, omega, v_jnts),
+                axis=0)
+        a_state=np.concatenate((a_root, a_jnts),
+                axis=0)
+        
+        return (q_state, v_state, a_state)
+    
+    def _set_is_adaptive(self):
+        
+        # measurements
+        p_root = self.robot_state.root_state.get(data_type="p", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        q_root = self.robot_state.root_state.get(data_type="q", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        v_root = self.robot_state.root_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        omega = self.robot_state.root_state.get(data_type="omega", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        a_root = self.robot_state.root_state.get(data_type="a_full", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        
+        q_jnts = self.robot_state.jnts_state.get(data_type="q", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        v_jnts = self.robot_state.jnts_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        a_jnts = self.robot_state.jnts_state.get(data_type="a", robot_idxs=self.controller_index_np).reshape(-1, 1)
+            
+        # rhc variables to be set
+        q=self._prb.getVariables("q") # .setBounds()
+        root_p_rhc=q[0:3] # root p
+        root_q_rhc=q[3:7] # root orientation
+        jnts_q_rhc=q[7:] # jnts q
+        vel=self._prb.getVariables("v")
+        root_v_rhc=vel[0:3] # lin v.
+        root_omega_rhc=vel[3:6] # omega
+        jnts_v_rhc=vel[6:] # jnts v
+        acc=self._prb.getVariables("a")
+        lin_a_prb=acc[0:3] # lin acc
+        
+        # getting prediction defects
+        root_q_delta=self.rhc_pred_delta.root_state.get(data_type="q", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        jnt_q_delta=self.rhc_pred_delta.jnts_state.get(data_type="q", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        jnt_v_delta=self.rhc_pred_delta.jnts_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        v_root_delta = self.rhc_pred_delta.root_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        omega_root_delta = self.rhc_pred_delta.root_state.get(data_type="omega", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        a_root_delta = self.rhc_pred_delta.root_state.get(data_type="a_full", robot_idxs=self.controller_index_np).reshape(-1, 1)
+
+        # close state on known quantities, estimate some (e.g. lin vel) and
+        # open loop if thing start to explode
+        alpha_now=1.0
+        delta=0.0 
+        if self._custom_opts["alpha_from_outside"]:
+            alpha_now=self.rhc_refs.get_alpha()
+        else: # "autotuned" alpha
+            if self._custom_opts["estimate_v_root"]: # we copmute delta based on jnt v (since we use meas.)
+                delta=np.max(np.abs(jnt_v_delta))
+            else:
+                delta=np.max(np.abs(omega_root_delta))
+            # fail_idx=self._get_failure_index()
+            # fail_idx=self._get_explosion_idx()/self._fail_idx_thresh
+            alpha_now=(np.tanh(2*self._alpha_half*(delta-self._alpha_half))+1)/2.0
+
+        bound_relaxation=-np.log(1-alpha_now+1e-16) # [~0, ~inf] if alpha_now [0, 1]
+        self.rhc_refs.set_alpha(alpha=alpha_now) # also writes on shared mem for db
+        self.rhc_refs.set_bound_relax(bound_relax=bound_relaxation) # also writes on shared mem for db
+
+        self._alphas_q_root[:]=alpha_now # for now single alpha for everything
+        self._alphas_q_jnts[:]=alpha_now
+        self._alphas_twist_root[:]=alpha_now
+        self._alphas_v_jnts[:]=alpha_now
+        self._alphas_a[:]=alpha_now
+        if not self._custom_opts["estimate_v_root"]:
+            self._alphas_twist_root[0:3]=1.0 # open
+            self._alphas_v_jnts[:]=1.0 # open
+
+        # position is always open loop
+        root_q_full_from_rhc=self._get_root_full_q_from_sol(node_idx=1).reshape(-1, 1)
+        root_p_from_rhc=root_q_full_from_rhc[0:3, :]
+        p_root[:, :]=root_p_from_rhc 
+    
+        # expaning meas q if continuous joints
+        if (not len(self._continuous_joints)==0): # we need do expand some meas. rev jnts to So2
+            self._jnts_q_expanded[self._rev_joints_idxs, :]=q_jnts[self._rev_joints_idxs_red ,:]
+            self._jnts_q_expanded[self._continuous_joints_idxs_cos, :]=np.cos(q_jnts[self._continuous_joints_idxs_red, :]) # cos
+            self._jnts_q_expanded[self._continuous_joints_idxs_sin, :]=np.sin(q_jnts[self._continuous_joints_idxs_red, :]) # sin
+            
+            # continous joints position is always open loop, but we need a delta vector of matching dimension
+            q_jnts_from_rhc=self._get_jnt_q_from_sol(node_idx=1).reshape(-1, 1)
+
+            self._jnts_q_delta_expanded[self._rev_joints_idxs, :]=jnt_q_delta[self._rev_joints_idxs_red ,:]
+
+            self._jnts_q_delta_expanded[self._continuous_joints_idxs_cos, :]=\
+                np.cos(q_jnts_from_rhc[self._continuous_joints_idxs_red, :]) - \
+                    np.cos(q_jnts[self._continuous_joints_idxs_red, :])
+            self._jnts_q_delta_expanded[self._continuous_joints_idxs_sin, :]=\
+                np.sin(q_jnts_from_rhc[self._continuous_joints_idxs_red, :]) - \
+                    np.sin(q_jnts[self._continuous_joints_idxs_red, :])
+
+            q_jnts=self._jnts_q_expanded.reshape(-1,1) # overriting with expanded jnts
+            jnt_q_delta=self._jnts_q_delta_expanded.reshape(-1, 1) # overriting with expanded jnts
+
+            self._alphas_q_jnts[self._continuous_joints_idxs_cos, :]=1.0 # open loop
+            self._alphas_q_jnts[self._continuous_joints_idxs_sin, :]=1.0 # open loop
+
+            # self._alphas_v_jnts[self._continuous_joints_idxs_red, :]=0.0 # open loop
+
+        root_p_rhc.setBounds(lb=p_root,
+            ub=p_root, nodes=0)
+        root_q_rhc.setBounds(lb=q_root+self._alphas_q_root[3:7]*root_q_delta, 
+            ub=q_root+self._alphas_q_root[3:7]*root_q_delta, nodes=0)
+        jnts_q_rhc.setBounds(lb=q_jnts+self._alphas_q_jnts*jnt_q_delta, 
+            ub=q_jnts+self._alphas_q_jnts*jnt_q_delta, nodes=0)
+        if self._custom_opts["estimate_v_root"]:
+            root_v_rhc.setBounds(lb=-self._v_inf[0:3], 
+                ub=self._v_inf[0:3], nodes=0)
+        else:
+            root_v_rhc.setBounds(lb=v_root+self._alphas_twist_root[0:3, :]*v_root_delta, 
+                ub=v_root+self._alphas_twist_root[0:3, :]*v_root_delta, nodes=0)
+        root_omega_rhc.setBounds(lb=omega+self._alphas_twist_root[3:6, :]*omega_root_delta, 
+            ub=omega+self._alphas_twist_root[3:6, :]*omega_root_delta, nodes=0)
+        jnts_v_rhc.setBounds(lb=v_jnts+self._alphas_v_jnts*jnt_v_delta, 
+            ub=v_jnts+self._alphas_v_jnts*jnt_v_delta, nodes=0)
+        if self._custom_opts["lin_a_feedback"]:
+            lin_a_prb.setBounds(lb=a_root[0:3, :]+self._alphas_a[0:3]*a_root_delta[0:3, :], 
+                ub=a_root[0:3, :]+self._alphas_a[0:3]*a_root_delta[0:3, :], 
+                nodes=0)
+            
+        # return state used for feedback
+        q_state=np.concatenate((p_root, q_root, q_jnts),
+                axis=0)
+        v_state=np.concatenate((v_root, omega, v_jnts),
+                axis=0)
+        a_state=np.concatenate((a_root, a_jnts),
+                axis=0)
+        
+        return (q_state, v_state, a_state)
+
+    def _solve(self):
+        
+        if self._debug:
+            return self._db_solve(bootstrap=False)
+        else:
+            return self._min_solve(bootstrap=False)
+
+    def _bootstrap(self):
+
+        if self._debug:
+            return self._db_solve(bootstrap=True)
+        else:
+            return self._min_solve(bootstrap=True)
+        
+    def _min_solve(self, bootstrap: bool = False):
+        # minimal solve version -> no debug 
+        robot_qstate=None
+        robot_vstate=None
+        robot_astate=None
+        if self._open_loop:
+            robot_qstate, robot_vstate, robot_astate = self._update_open_loop(bootstrap=bootstrap) # updates the TO ig and 
+            # initial conditions using data from the solution itself
+        else: 
+            robot_qstate, robot_vstate, robot_astate = self._update_closed_loop(bootstrap=bootstrap) # updates the TO ig and 
+            # initial conditions using robot measurements
+    
+        self._pm.shift() # shifts phases of one dt
+        if self._refs_in_hor_frame:
+            # q_base=self.robot_state.root_state.get(data_type="q", 
+            #     robot_idxs=self.controller_index).reshape(-1, 1)
+            # q_full=self._get_full_q_from_sol(node_idx=1).reshape(-1, 1)
+            # using internal base pose from rhc. in case of closed loop, it will be the meas state
+            force_norm=None
+            if self._custom_opts["use_force_feedback"]:
+                contact_forces=self.robot_state.contact_wrenches.get(data_type="f", 
+                    robot_idxs=self.controller_index_np,
+                    contact_name=None).reshape(self.n_contacts,3)
+                force_norm=np.linalg.norm(contact_forces, axis=1)
+            self.rhc_refs.step(qstate=robot_qstate, vstate=robot_vstate,
+                force_norm=force_norm)
+        else:
+            self.rhc_refs.step()
+            
+        try:
+            if not bootstrap:
+                converged = self._ti.rti() # RTI step
+            else:
+                converged = self._ti.bootstrap() # full solve (to convergence)
+            self.sol_counter = self.sol_counter + 1
+            return not self._check_rhc_failure()
+        except Exception as e: # fail in case of exceptions
+            return False
+    
+    def _db_solve(self, bootstrap: bool = False):
+
+        self._timer_start = time.perf_counter()
+
+        robot_qstate=None
+        robot_vstate=None
+        robot_astate=None
+        if self._open_loop:
+            robot_qstate, robot_vstate, robot_astate = self._update_open_loop(bootstrap=bootstrap) # updates the TO ig and 
+            # initial conditions using data from the solution itself
+        else: 
+            robot_qstate, robot_vstate, robot_astate = self._update_closed_loop(bootstrap=bootstrap) # updates the TO ig and 
+            # initial conditions using robot measurements
+
+        self._prb_update_time = time.perf_counter() 
+        self._pm.shift() # shifts phases of one dt
+        self._phase_shift_time = time.perf_counter()
+
+        if self._refs_in_hor_frame:
+            # q_base=self.robot_state.root_state.get(data_type="q", 
+            #     robot_idxs=self.controller_index).reshape(-1, 1)
+            # q_full=self._get_full_q_from_sol(node_idx=1).reshape(-1, 1)
+            # using internal base pose from rhc. in case of closed loop, it will be the meas state
+            force_norm=None
+            if self._custom_opts["use_force_feedback"]:
+                contact_forces=self.robot_state.contact_wrenches.get(data_type="f", 
+                    robot_idxs=self.controller_index_np,
+                    contact_name=None).reshape(self.n_contacts,3)
+                force_norm=np.linalg.norm(contact_forces, axis=1)
+            self.rhc_refs.step(qstate=robot_qstate, vstate=robot_vstate,
+                force_norm=force_norm)
+        else:
+            self.rhc_refs.step()
+             
+        self._task_ref_update_time = time.perf_counter() 
+    
+        try:
+            if not bootstrap:
+                converged = self._ti.rti() # RTI step
+            else:
+                converged = self._ti.bootstrap() # full solve bootstrap
+            self._rti_time = time.perf_counter() 
+            self.sol_counter = self.sol_counter + 1
+            self._update_db_data()
+            return not self._check_rhc_failure()
+        except Exception as e: # fail in case of exceptions
+            if self._verbose:
+                solve_mode = "RTI" if not bootstrap else "Bootstrap"
+                exception = f"{solve_mode}() for controller {self.controller_index} failed" + \
+                f" with exception {type(e).__name__}"
+                Journal.log(self.__class__.__name__,
+                    "solve",
+                    exception,
+                    LogType.EXCEP,
+                    throw_when_excep = False)
+            self._update_db_data()
+            return False
+    
+    def _get_fail_idx(self):
+        
+        self._explosion_idx_buffer[:, self._expl_idx_buffer_counter]=self._get_explosion_idx()
+        self._expl_idx_buffer_counter+=1
+        self._expl_idx_counter+=1
+        if self._expl_idx_counter%self._expl_idx_window_size==0:
+            self._expl_idx_buffer_counter=0 # restart from 0
+        
+        running_avrg=np.mean(self._explosion_idx_buffer).item()
+        
+        return running_avrg
+    
+    def _get_explosion_idx(self):
+        explosion_index = self._get_rhc_constr_viol() + self._get_rhc_cost()*self._fail_idx_scale
+        return explosion_index
+
+    def _update_db_data(self):
+
+        self._profiling_data_dict["problem_update_dt"] = self._prb_update_time - self._timer_start
+        self._profiling_data_dict["phases_shift_dt"] = self._phase_shift_time - self._prb_update_time
+        self._profiling_data_dict["task_ref_update"] = self._task_ref_update_time - self._phase_shift_time
+        self._profiling_data_dict["rti_solve_dt"] = self._rti_time - self._task_ref_update_time
+        self.rhc_costs.update(self._ti.solver_rti.getCostsValues())
+        self.rhc_constr.update(self._ti.solver_rti.getConstraintsValues())
+
+    def _reset(self):
+        
+        # reset task interface (ig, solvers, etc..) + 
+        # phase manager and sets bootstap as solution
+        self._gm.reset()
+        self._explosion_idx_buffer[:, :]=self._get_explosion_idx() # reset with data from reset solution
+        self._expl_idx_counter=0.0
+        self._expl_idx_buffer_counter=0
+
+    def _get_cost_info(self):
+        
+        cost_dict = self._ti.solver_rti.getCostsValues()
+        cost_names = list(cost_dict.keys())
+        cost_dims = [1] * len(cost_names) # costs are always scalar
+        return cost_names, cost_dims
+    
+    def _get_constr_info(self):
+        
+        constr_dict = self._ti.solver_rti.getConstraintsValues()
+        constr_names = list(constr_dict.keys())
+        constr_dims = [-1] * len(constr_names)
+        i = 0
+        for constr in constr_dict:
+            constr_val = constr_dict[constr]
+            constr_shape = constr_val.shape
+            constr_dims[i] = constr_shape[0]
+            i+=1
+        return constr_names, constr_dims
+    
+    def _get_q_from_sol(self):
+        full_q=self._ti.solution['q'].astype(self._dtype)
+        if self._custom_opts["replace_continuous_joints"]:
+            return full_q
+        else:
+            cont_jnts=full_q[self._continuous_joints_idxs_qfull, :]
+            cos=cont_jnts[::2, :]
+            sin=cont_jnts[1::2, :]
+            # copy root
+            self._full_q_reduced[0:7, :]=full_q[0:7, :]
+            # copy rev joint vals
+            self._full_q_reduced[self._rev_joints_idxs_red_qfull, :]=full_q[self._rev_joints_idxs_qfull, :]
+            # and continuous
+            angle=np.arctan2(sin, cos)
+            self._full_q_reduced[self._continuous_joints_idxs_red_qfull, :]=angle
+            return self._full_q_reduced
+
+    def _get_v_from_sol(self):
+        return self._ti.solution['v'].astype(self._dtype)
+    
+    def _get_a_from_sol(self):
+        return self._ti.solution['a'].astype(self._dtype)
+    
+    def _get_a_dot_from_sol(self):
+        return None
+    
+    def _get_f_from_sol(self):
+        # to be overridden by child class
+        contact_names =self._get_contacts() # we use controller-side names
+        try: 
+            data=[]
+            for key in contact_names:
+                contact_f=self._ti.solution["f_" + key].astype(self._dtype)
+                np.nan_to_num(contact_f, nan=1e6, posinf=1e6, neginf=-1e6, copy=False)
+                np.clip(a=contact_f, a_max=1e6, a_min=-1e6, out=contact_f)
+                data.append(contact_f)
+            return np.concatenate(data, axis=0)
+        except:
+            return None
+            
+    def _get_f_dot_from_sol(self):
+        # to be overridden by child class
+        return None
+    
+    def _get_eff_from_sol(self):
+        # to be overridden by child class
+        return None
+    
+    def _get_cost_from_sol(self,
+                    cost_name: str):
+        return self.rhc_costs[cost_name]
+    
+    def _get_constr_from_sol(self,
+                    constr_name: str):
+        return self.rhc_constr[constr_name]

ibrido_run__2026_03_07_19_20_05_ID/training_cfg_centauro_no_yaw_ub_cloop.sh

@@ -0,0 +1,85 @@
+#!/bin/bash
+export EVAL=0
+export DET_EVAL=1
+export EVAL_ON_CPU=1
+export OVERRIDE_ENV=1
+export OVERRIDE_AGENT_REFS=1
+export MPATH="/root/training_data/"
+export MNAME=""
+
+export WANDB_KEY="25f235316292344cea6dfa68e7c95409b3374d03"
+export SHM_NS="centauro_big_wheels_no_yaw_ub" # shared mem namespace used for all shared data on CPU 
+export N_ENVS=800 # number of env to run in parallel
+export RNAME="CentauroCloopPartialNoYawUb" # a descriptive base name for this run
+export SEED=98 # random n generator seed to be used for this run
+export REMOTE_STEPPING=1
+export COMPRESSION_RATIO=0.6
+export ACTOR_LWIDTH=128
+export ACTOR_DEPTH=3
+export CRITIC_LWIDTH=256
+export CRITIC_DEPTH=4
+export OBS_NORM=1
+export OBS_RESCALING=0
+export CRITIC_ACTION_RESCALE=1
+export WEIGHT_NORM=1
+export LAYER_NORM=0
+export BATCH_NORM=0
+export IS_CLOSED_LOOP=1
+export DEBUG=1
+export RMDEBUG=1
+export DUMP_ENV_CHECKPOINTS=1
+export DEMO_STOP_THRESH=10.0
+export TOT_STEPS=20000000
+export USE_RND=0
+export DEMO_ENVS_PERC=0.0
+export EXPL_ENVS_PERC=0.0
+export ACTION_REPEAT=3
+export USE_SAC=1
+export USE_DUMMY=0
+export DISCOUNT_FACTOR=0.99
+export USE_PERIOD_RESETS=0
+export COMMENT='centauro big wheels (fixed ankle yaw) with upper body CLOOP, a bit of pushes' # any training comment
+export URDF_PATH="${HOME}/ibrido_ws/src/iit-centauro-ros-pkg/centauro_urdf/urdf/centauro.urdf.xacro" # name of the description package for the robot
+export SRDF_PATH="${HOME}/ibrido_ws/src/iit-centauro-ros-pkg/centauro_srdf/srdf/centauro.srdf.xacro" # base path where the description package for the robot are located
+export JNT_IMP_CF_PATH="${HOME}/ibrido_ws/src/CentauroHybridMPC/centaurohybridmpc/config/jnt_imp_config_no_yaw_ub_open.yaml" # path to yaml file for jnt imp configuration
+if (( $IS_CLOSED_LOOP )); then
+  export JNT_IMP_CF_PATH="${HOME}/ibrido_ws/src/CentauroHybridMPC/centaurohybridmpc/config/jnt_imp_config_no_yaw_ub.yaml"
+fi
+
+export CLUSTER_CL_FNAME="centaurohybridmpc.controllers.horizon_based.centauro_rhc_cluster_client" # base path where the description package for the robot are located
+export CLUSTER_DT=0.04
+export N_NODES=25
+export CLUSTER_DB=1
+export PHYSICS_DT=0.0005
+export USE_GPU_SIM=1
+# export CODEGEN_OVERRIDE_BDIR="none"
+export CODEGEN_OVERRIDE_BDIR="${HOME}/aux_data/CentauroRHCLusterClient_${SHM_NS}/CodeGen/${SHM_NS}Rhc"
+# export TRAIN_ENV_FNAME="twist_tracking_env"
+# export TRAIN_ENV_CNAME="TwistTrackingEnv"
+export TRAIN_ENV_FNAME="fake_pos_tracking_env"
+export TRAIN_ENV_CNAME="FakePosTrackingEnv"
+# export TRAIN_ENV_FNAME="fake_pos_track_env_phase_control"
+# export TRAIN_ENV_CNAME="FakePosTrackEnvPhaseControl"
+# export TRAIN_ENV_FNAME="fake_pos_track_env_phase_control_with_demo"
+# export TRAIN_ENV_CNAME="FakePosTrackEnvPhaseControlWithDemo"
+# export TRAIN_ENV_FNAME="fake_pos_tracking_with_demo"
+# export TRAIN_ENV_CNAME="FakePosTrackingEnvWithDemo"
+# export TRAIN_ENV_FNAME="linvel_env_with_demo"
+# export TRAIN_ENV_CNAME="TwistTrackingEnvWithDemo"
+export PUB_HEIGHTMAP=0
+export BAG_SDT=90.0
+export BRIDGE_DT=0.1
+export DUMP_DT=50.0
+export ENV_IDX_BAG=4
+export ENV_IDX_BAG_DEMO=-1
+export ENV_IDX_BAG_EXPL=-1
+export SRDF_PATH_ROSBAG="${HOME}/aux_data/CentauroRHClusterClient_${SHM_NS}/$SHM_NS.srdf" # base path where the description package for the robot are located
+export CUSTOM_ARGS_NAMES="use_random_pertub use_jnt_v_feedback step_height control_wheels fixed_flights adaptive_is lin_a_feedback closed_partial fix_yaw use_flat_ground estimate_v_root self_collide add_upper_body"
+export CUSTOM_ARGS_DTYPE="bool bool float bool bool bool bool bool bool bool bool bool bool "
+export CUSTOM_ARGS_VALS="true false 0.1 true true true false true true true false false true"
+# export CUSTOM_ARGS_NAMES+=" contact_prims"
+# export CUSTOM_ARGS_DTYPE+=" strlist"
+# export CUSTOM_ARGS_VALS+=" wheel_1,wheel_2,wheel_3,wheel_4"
+export SET_ULIM=1 
+export ULIM_N=131072  # maximum number of open file descriptors for each process (shared memory)
+export TIMEOUT_MS=120000 # timeout after which each script autokills ([ms])

jnt_imp_config.yaml

@@ -0,0 +1,46 @@
+XBotInterface: # just used for retrieving homing in sim with xbot_mujoco
+  urdf_path: $PWD/xmj_env_files/centauro.urdf
+  srdf_path: $PWD/xmj_env_files/centauro.srdf
+
+# defaults
+motor_pd:
+  "j_arm*_1": [500, 10]
+  "j_arm*_2": [500, 10]
+  "j_arm*_3": [500, 10]
+  "j_arm*_4": [500, 10]
+  "j_arm*_5": [100, 5]
+  "j_arm*_6": [100, 5]
+  "j_arm*_7": [100, 5]
+  "hip_yaw_*": [3000, 30]
+  "hip_pitch_*": [3000, 30]
+  "knee_pitch_*": [3000, 30]
+  "ankle_pitch_*": [1000, 10]
+  "ankle_yaw_*": [300, 10]
+  "neck_pitch": [10, 1]
+  "neck_yaw": [10, 1]
+  "torso_yaw": [1000, 30]
+  "j_wheel_*": [0, 30]
+  "velodyne_*": [10, 1]
+  "d435_*": [10, 1]
+  "dagana_*": [50, 1]
+
+startup_motor_pd: # iannis combo
+  "j_arm*_1": [500, 30]
+  "j_arm*_2": [500, 30]
+  "j_arm*_3": [500, 30]
+  "j_arm*_4": [500, 30]
+  "j_arm*_5": [100, 5]
+  "j_arm*_6": [100, 5]
+  "j_arm*_7": [100, 5]
+  "hip_yaw_*": [200, 60]
+  "hip_pitch_*": [200, 60]
+  "knee_pitch_*": [200, 60]
+  "ankle_pitch_*": [200, 60]
+  "ankle_yaw_*": [120, 30]
+  "neck_pitch": [10, 1]
+  "neck_yaw": [10, 1]
+  "torso_yaw": [1000, 30]
+  "j_wheel_*": [0, 30]
+  "velodyne_*": [10, 1]
+  "d435_*": [10, 1]
+  "dagana_*": [50, 1]

jnt_imp_config_no_yaw_ub.yaml

@@ -0,0 +1,46 @@
+XBotInterface: # just used for retrieving homing in sim with xbot_mujoco
+  urdf_path: $PWD/xmj_env_files/centauro.urdf
+  srdf_path: $PWD/xmj_env_files/centauro.srdf
+
+# defaults
+motor_pd:
+  "j_arm*_1": [500, 10]
+  "j_arm*_2": [500, 10]
+  "j_arm*_3": [500, 10]
+  "j_arm*_4": [500, 10]
+  "j_arm*_5": [100, 5]
+  "j_arm*_6": [100, 5]
+  "j_arm*_7": [100, 5]
+  "hip_yaw_*": [3000, 30]
+  "hip_pitch_*": [3000, 30]
+  "knee_pitch_*": [3000, 30]
+  "ankle_pitch_*": [1000, 10]
+  "ankle_yaw_*": [300, 10]
+  "neck_pitch": [10, 1]
+  "neck_yaw": [10, 1]
+  "torso_yaw": [1000, 30]
+  "j_wheel_*": [0, 30]
+  "velodyne_*": [10, 1]
+  "d435_*": [10, 1]
+  "dagana_*": [50, 1]
+
+startup_motor_pd: # iannis combo
+  "j_arm*_1": [100, 25]
+  "j_arm*_2": [100, 25]
+  "j_arm*_3": [100, 25]
+  "j_arm*_4": [100, 25]
+  "j_arm*_5": [20, 8]
+  "j_arm*_6": [20, 8]
+  "j_arm*_7": [20, 8]
+  "hip_yaw_*": [200, 60]
+  "hip_pitch_*": [200, 60]
+  "knee_pitch_*": [200, 60]
+  "ankle_pitch_*": [200, 60]
+  "ankle_yaw_*": [600, 10]
+  "neck_pitch": [10, 1]
+  "neck_yaw": [10, 1]
+  "torso_yaw": [400, 60]
+  "j_wheel_*": [0, 30]
+  "velodyne_*": [10, 1]
+  "d435_*": [10, 1]
+  "dagana_*": [50, 1]

launch_control_cluster.py

@@ -0,0 +1,106 @@
+import os
+import argparse
+import multiprocessing as mp
+import importlib.util
+import inspect
+
+from aug_mpc.utils.custom_arg_parsing import generate_custom_arg_dict
+
+from EigenIPC.PyEigenIPC import Journal, LogType
+
+this_script_name = os.path.splitext(os.path.basename(os.path.abspath(__file__)))[0]
+
+# Function to dynamically import a module from a specific file path
+def import_env_module(env_path):
+    spec = importlib.util.spec_from_file_location("env_module", env_path)
+    env_module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(env_module)
+    return env_module
+
+if __name__ == "__main__":  
+
+    # Parse command line arguments for CPU affinity
+    parser = argparse.ArgumentParser(description="Set CPU affinity for the script.")
+    parser.add_argument('--ns', type=str, help='Namespace to be used for cluster shared memory')
+    parser.add_argument('--urdf_path', type=str, help='Robot description package path for URDF ')
+    parser.add_argument('--srdf_path', type=str, help='Robot description package path for SRDF ')
+    parser.add_argument('--size', type=int, help='cluster size', default=1)
+
+    # Replacing argparse.BooleanOptionalAction with 'store_true' and 'store_false' for compatibility with Python 3.8
+    parser.add_argument('--cloop',action='store_true', help='whether to use RHC controllers in closed loop mode')
+    parser.add_argument('--cluster_dt', type=float, default=0.05, help='dt used by MPC controllers for discretization')
+    parser.add_argument('--n_nodes', type=int, default=31, help='n nodes used by MPC controllers')
+
+    parser.add_argument('--verbose',action='store_true', help='run in verbose mode')
+
+    parser.add_argument('--enable_debug',action='store_true', help='enable debug mode for cluster client and all controllers')
+
+    parser.add_argument('--dmpdir', type=str, help='directory where data is dumped', default="/root/aux_data")
+
+    parser.add_argument('--no_mp_fork',action='store_true', help='whether to multiprocess with forkserver context')
+
+    parser.add_argument('--set_affinity',action='store_true', help='set affinity to a core for each controller')
+
+    parser.add_argument('--comment', type=str, help='Any useful comment associated with this run', default="")
+    parser.add_argument('--timeout_ms', type=int, help='connection timeout after which the script self-terminates', default=60000)
+    parser.add_argument('--codegen_override_dir', type=str, help='Path to base dir where codegen is to be loaded', default="")
+
+    parser.add_argument('--custom_args_names', nargs='+', default=None,
+                            help='list of custom arguments names')
+    parser.add_argument('--custom_args_vals', nargs='+', default=None,
+                            help='list of custom arguments values')
+    parser.add_argument('--custom_args_dtype', nargs='+', default=None,
+                            help='list of custom arguments data types')
+    
+    parser.add_argument('--cluster_client_fname', type=str, 
+        default="aug_mpc.controllers.rhc.hybrid_quad_client",
+        help="cluster client file import pattern (without extension)")
+
+    args = parser.parse_args()
+    
+    # Ensure custom_args_names and custom_args_vals have the same length
+    custom_opts = generate_custom_arg_dict(args=args)
+    custom_opts.update({"cloop": args.cloop, 
+        "cluster_dt": args.cluster_dt, 
+        "n_nodes": args.n_nodes, 
+        "codegen_override_dir": args.codegen_override_dir})
+    if args.no_mp_fork: # this needs to be in the main
+        mp.set_start_method('spawn')
+    else:
+        # mp.set_start_method('forkserver')
+        mp.set_start_method('fork')
+        
+    cluster_module=importlib.import_module(args.cluster_client_fname)
+    # Get all classes defined in the module
+    classes_in_module = [name for name, obj in inspect.getmembers(cluster_module, inspect.isclass) 
+                        if obj.__module__ == cluster_module.__name__]
+    if len(classes_in_module) == 1:
+        cluster_classname=classes_in_module[0]
+        ClusterClient = getattr(cluster_module, cluster_classname)
+        cluster_client = ClusterClient(namespace=args.ns, 
+            cluster_size=args.size,
+            urdf_xacro_path=args.urdf_path,
+            srdf_xacro_path=args.srdf_path,
+            open_loop=not args.cloop,
+            use_mp_fork = not args.no_mp_fork,
+            verbose=args.verbose,
+            debug=args.enable_debug,
+            base_dump_dir=args.dmpdir,
+            timeout_ms=args.timeout_ms,
+            custom_opts=custom_opts,
+            codegen_override=args.codegen_override_dir,
+            set_affinity=args.set_affinity)
+        cluster_client.run()
+        
+    else:
+        class_list_str = ", ".join(classes_in_module)
+        Journal.log("launch_control_cluster.py",
+            "",
+            f"Found more than one class in cluster client file {args.cluster_client_fname}. Found: {class_list_str}",
+            LogType.EXCEP,
+            throw_when_excep = False)
+        exit()
+
+    # control_cluster_client = 
+    # control_cluster_client.run() # spawns the controllers on separate processes (blocking)
+

launch_train_env.py

@@ -0,0 +1,358 @@
+from aug_mpc.utils.determinism import deterministic_run
+
+from mpc_hive.utilities.shared_data.sim_data import SharedEnvInfo
+from mpc_hive.utilities.shared_data.cluster_data import SharedClusterInfo
+
+from EigenIPC.PyEigenIPC import VLevel, Journal, LogType
+from EigenIPC.PyEigenIPC import StringTensorServer
+
+import os, argparse, sys, types, inspect
+
+from perf_sleep.pyperfsleep import PerfSleep
+
+import importlib.util
+import torch
+import signal
+
+algo = None  # global to make it accessible by signal handler
+exit_request=False
+dummy_step_exit_req=False
+
+def handle_sigint(signum, frame):
+    global exit_request, dummy_step_exit_req
+    Journal.log("launch_train_env.py",
+        "",
+        f"Received sigint. Will stop training.",
+        LogType.WARN)
+    exit_request=True
+    dummy_step_exit_req=True # in case dummy_step_loop was used
+    
+# Function to dynamically import a module from a specific file path
+# def import_env_module(env_path):
+#     spec = importlib.util.spec_from_file_location("env_module", env_path)
+#     env_module = importlib.util.module_from_spec(spec)
+#     spec.loader.exec_module(env_module)
+#     return env_module
+
+def import_env_module(env_path, local_env_root: str = None):
+    """
+    env_path: full path to the child env .py file to exec
+    local_env_root: directory where local copies of aug_mpc_envs.training_envs modules live
+    """
+    if local_env_root is not None:
+        local_env_root = os.path.abspath(local_env_root)
+        # override aug_mpc_envs.training_envs package to point to the local_env_root
+        pkg_name = "aug_mpc_envs.training_envs"
+        if pkg_name not in sys.modules:
+            mod = types.ModuleType(pkg_name)
+            mod.__path__ = [local_env_root]  # tell Python to look here first
+            sys.modules[pkg_name] = mod
+        else:
+            existing = getattr(sys.modules[pkg_name], "__path__", None)
+            if existing is None:
+                sys.modules[pkg_name].__path__ = [local_env_root]
+            elif local_env_root not in existing:
+                existing.insert(0, local_env_root)
+
+    # load the module as usual
+    spec = importlib.util.spec_from_file_location("env_module", env_path)
+    env_module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(env_module)
+    return env_module
+
+def log_env_hierarchy(env_class, env_path, env_type="training"):
+    """
+    Logs the env class, its file, and full inheritance hierarchy with file paths.
+    env_class: the child env class
+    env_path: file path where the child class was loaded from
+    env_type: string label, e.g., "training", "evaluation", "resumed_training"
+    """
+    def get_bases_recursive(cls):
+        """Recursively get all base classes with their file paths."""
+        info = []
+        for base in cls.__bases__:
+            try:
+                file = inspect.getfile(base)
+            except TypeError:
+                file = "built-in or unknown"
+            info.append(f"{base.__name__} (from {file})")
+            # Recurse unless it's object
+            if base is not object:
+                info.extend(get_bases_recursive(base))
+        return info
+
+    hierarchy_info = get_bases_recursive(env_class)
+    hierarchy_str = " -> ".join(hierarchy_info) if hierarchy_info else "No parents"
+
+    Journal.log(
+        "launch_train_env.py",
+        "",
+        f"loading {env_type} env {env_class.__name__} (from {env_path}) "
+        f"with hierarchy: {hierarchy_str}",
+        LogType.INFO,
+        throw_when_excep=True
+    )
+
+def dummy_step_loop(env):
+    global dummy_step_exit_req
+    while True:
+        if dummy_step_exit_req: 
+            return True
+        step_ok=env.step(action=env.safe_action) # not a busy loop because of MPC in the step
+        if not step_ok:
+            return False
+
+if __name__ == "__main__":  
+
+    signal.signal(signal.SIGINT, handle_sigint)
+
+    # Parse command line arguments for CPU affinity
+    parser = argparse.ArgumentParser(description="Set CPU affinity for the script.")
+
+    parser.add_argument('--run_name', type=str, default=None, help='Name of training run')
+    parser.add_argument('--ns', type=str, help='Namespace to be used for shared memory')
+    parser.add_argument('--timeout_ms', type=int, help='Connection timeout after which the script self-terminates', default=60000)
+    parser.add_argument('--drop_dir', type=str, help='Directory root where all run data will be dumped')
+    parser.add_argument('--comment', type=str, help='Any useful comment associated with this run', default="")
+    parser.add_argument('--seed', type=int, help='Seed', default=1)
+    parser.add_argument('--use_cpu',action='store_true', help='If set, all the training (data included) will be performed on CPU')
+
+    parser.add_argument('--db',action='store_true', help='Whether to enable local data logging for the algorithm (reward metrics, etc.)')
+    parser.add_argument('--env_db',action='store_true', help='Whether to enable env db data logging on shared mem (e.g. reward metrics are not available for reading anymore)')
+    parser.add_argument('--full_env_db',action='store_true', help='Whether to enable detailed episodic data storage (data over single transitions)')
+    parser.add_argument('--rmdb',action='store_true', help='Whether to enable remote debug (e.g. data logging on remote servers)')
+
+    parser.add_argument('--tot_tsteps', type=int, help='Total number of timesteps to be collected', default=int(30e6))
+    parser.add_argument('--action_repeat', type=int, help='Frame skipping (1-> no skip)', default=1)
+    parser.add_argument('--discount_factor', type=float, help='', default=0.99)
+    parser.add_argument('--obs_norm',action='store_true', help='Whether to enable the use of running normalizer in agent')
+    parser.add_argument('--obs_rescale',action='store_true', help='Whether to rescale observation depending on their expected range')
+    parser.add_argument('--add_weight_norm',action='store_true', help='Whether to add weight normalization to agent interal llayers')
+    parser.add_argument('--add_layer_norm',action='store_true', help='Whether to add layer normalization to agent internal llayers')
+    parser.add_argument('--add_batch_norm',action='store_true', help='Whether to add batch normalization to agent internal llayers')
+
+    parser.add_argument('--act_rescale_critic',action='store_true', help='Whether to rescale actions provided to critic (if SAC) to be in range [-1, 1]')
+    parser.add_argument('--use_period_resets',action='store_true', help='')
+
+    parser.add_argument('--sac',action='store_true', help='Use SAC, otherwise PPO, unless dummy is set')
+    parser.add_argument('--dummy',action='store_true', help='Use dummy agent (useful for testing and debugging environments)')
+
+    parser.add_argument('--dump_checkpoints',action='store_true', help='Whether to dump model checkpoints during training')
+
+    parser.add_argument('--demo_envs_perc', type=float, help='[0, 1]', default=0.0)
+    parser.add_argument('--demo_stop_thresh', type=float, default=None, 
+        help='Performance hreshold above which demonstration envs should be deactivated.')
+    
+    parser.add_argument('--expl_envs_perc', type=float, help='[0, 1]', default=0)
+    
+    parser.add_argument('--use_rnd',action='store_true', help='Whether to use RND for exploration')
+
+    parser.add_argument('--eval',action='store_true', help='Whether to perform an evaluation run')
+    parser.add_argument('--n_eval_timesteps', type=int, help='Total number of timesteps to be evaluated', default=int(1e6))
+    parser.add_argument('--det_eval',action='store_true', help='Whether to perform a deterministic eval (only action mean is used). Only valid if --eval.')
+    parser.add_argument('--allow_expl_during_eval',action='store_true', help='Whether to allow expl envs during evaluation (useful to tune exploration)')
+    
+    parser.add_argument('--resume',action='store_true', help='Resume a previous training using a checkpoint')
+    parser.add_argument('--mpath', type=str, help='Model path to be used for policy evaluation', default=None)
+    parser.add_argument('--mname', type=str, help='Model name', default=None)
+    parser.add_argument('--override_env',action='store_true', help='Whether to override env when running evaluation')
+    
+    parser.add_argument('--anomaly_detect',action='store_true', help='Whether to enable anomaly detection (useful for debug)')
+
+    parser.add_argument('--compression_ratio', type=float,
+        help='If e.g. 0.8, the fist layer will be of dimension [input_features_size x (input_features_size*compression_ratio)]', default=-1.0)
+    parser.add_argument('--actor_lwidth', type=int, help='Actor network layer width', default=128)
+    parser.add_argument('--critic_lwidth', type=int, help='Critic network layer width', default=256)
+    parser.add_argument('--actor_n_hlayers', type=int, help='Actor network size', default=3)
+    parser.add_argument('--critic_n_hlayers', type=int, help='Critic network size', default=4)
+
+    parser.add_argument('--env_fname', type=str, default="twist_tracking_env", help='Training env file name (without extension)')
+    parser.add_argument('--env_classname', type=str, default="TwistTrackingEnv", help='Training env class name')
+    parser.add_argument('--override_agent_actions',action='store_true', help='Whether to override agent actions with custom ones from shared mem (useful for db)')
+    parser.add_argument('--override_agent_refs',action='store_true', help='Whether to override automatically generated agent refs (useful for debug)')
+    
+    parser.add_argument('--step_while_setup',action='store_true', help='Continue stepping env with default actions while setting up agent, etc..')
+    parser.add_argument('--reset_on_init',action='store_true', help='Whether to reset the environment on initialization')
+
+    args = parser.parse_args()
+    args_dict = vars(args)
+
+    if args.eval and args.resume:
+        Journal.log("launch_train_env.py",
+            "",
+            f"Cannot set both --eval and --resume flags. Exiting.",
+            LogType.EXCEP,
+            throw_when_excep = True)
+
+    deterministic_run(seed=args.seed, torch_det_algos=False)
+
+    anomaly_detect=False
+    if args.anomaly_detect:
+        torch.autograd.set_detect_anomaly(True)
+
+    if (not args.mpath is None) and (not args.mname is None):
+        mpath_full = os.path.join(args.mpath, args.mname)
+    else:
+        mpath_full=None
+    
+    env_fname=args.env_fname
+    env_classname = args.env_classname
+    env_path=""
+    env_module=None
+    if (not args.eval and not args.resume) or (args.override_env):
+        # if starting a fresh traning or overriding env, load from a fresh env from aug_mpc
+        env_path = f"aug_mpc_envs.training_envs.{env_fname}"
+        env_module = importlib.import_module(env_path)
+    else:
+        if args.mpath is None:
+            Journal.log("launch_train_env.py",
+                "",
+                f"no mpath provided! Cannot load env. Either provide a mpath or run with --override_env",
+                LogType.EXCEP,
+                throw_when_excep = True)
+
+        env_path = os.path.join(args.mpath, env_fname + ".py")
+        env_module = import_env_module(env_path, local_env_root=args.mpath)
+       
+    EnvClass = getattr(env_module, env_classname)
+    env_type = "training" if not args.eval else "evaluation"
+    if args.resume:
+        env_type = "resumed_training"
+    log_env_hierarchy(EnvClass, env_path, env_type) # db print of env class
+    
+    env = EnvClass(namespace=args.ns,
+            verbose=True,
+            vlevel=VLevel.V2,
+            use_gpu=not args.use_cpu,
+            debug=args.env_db,
+            override_agent_refs=args.override_agent_refs,
+            timeout_ms=args.timeout_ms,
+            env_opts=args_dict)
+    if not env.is_ready(): # something went wrong
+        exit()
+    
+    dummy_step_thread = None
+    if args.step_while_setup:
+        import threading
+        # spawn step thread (we don't true parallelization, thread is fine)
+        # start the dummy stepping in a separate thread so setup can continue concurrently
+        dummy_step_thread = threading.Thread(target=dummy_step_loop, args=(env,), daemon=True)
+        dummy_step_thread.start()
+    
+    # getting some sim info for debugging
+    sim_data = {}
+    sim_info_shared = SharedEnvInfo(namespace=args.ns,
+                is_server=False,
+                safe=False)
+    sim_info_shared.run()
+    sim_info_keys = sim_info_shared.param_keys
+    sim_info_data = sim_info_shared.get().flatten()
+    for i in range(len(sim_info_keys)):
+        sim_data[sim_info_keys[i]] = sim_info_data[i]
+    
+    # getting come cluster info for debugging
+    cluster_data={}
+    cluste_info_shared = SharedClusterInfo(namespace=args.ns,
+                is_server=False,
+                safe=False)
+    cluste_info_shared.run()
+    cluster_info_keys = cluste_info_shared.param_keys
+    cluster_info_data = cluste_info_shared.get().flatten()
+    for i in range(len(cluster_info_keys)):
+        cluster_data[cluster_info_keys[i]] = cluster_info_data[i]
+
+    custom_args={}
+    custom_args["uname_host"]="user_host"
+    try:
+        username = os.getlogin() # add machine info to db data
+        hostname = os.uname().nodename
+        user_host = f"{username}@{hostname}"
+        custom_args["uname_host"]=user_host
+    except:
+        pass
+    
+    algo=None
+    if not args.dummy:
+        if args.sac:
+            from aug_mpc.training_algs.sac.sac import SAC
+
+            algo = SAC(env=env, 
+                debug=args.db, 
+                remote_db=args.rmdb,
+                seed=args.seed)
+        else:
+            from aug_mpc.training_algs.ppo.ppo import PPO
+
+            algo = PPO(env=env, 
+                debug=args.db, 
+                remote_db=args.rmdb,
+                seed=args.seed)
+    else:
+        from aug_mpc.training_algs.dummy.dummy import Dummy
+
+        algo=Dummy(env=env, 
+                debug=args.db, 
+                remote_db=args.rmdb,
+                seed=args.seed)
+
+    custom_args.update(args_dict)
+    custom_args.update(cluster_data)
+    custom_args.update(sim_data)
+
+    run_name=env_classname if args.run_name is None else args.run_name
+    algo.setup(run_name=run_name, 
+        ns=args.ns,
+        verbose=True,
+        drop_dir_name=args.drop_dir,
+        custom_args=custom_args,
+        comment=args.comment,
+        eval=args.eval,
+        resume=args.resume,
+        model_path=mpath_full,
+        n_eval_timesteps=args.n_eval_timesteps,
+        dump_checkpoints=args.dump_checkpoints,
+        norm_obs=args.obs_norm,
+        rescale_obs=args.obs_rescale)
+    
+    full_drop_dir=algo.drop_dir()
+    shared_drop_dir = StringTensorServer(length=1, 
+        basename="SharedTrainingDropDir", 
+        name_space=args.ns,
+        verbose=True, 
+        vlevel=VLevel.V2, 
+        force_reconnection=True)
+    shared_drop_dir.run()
+    
+    while True:
+        if not shared_drop_dir.write_vec([full_drop_dir], 0):
+            ns=1000000000
+            PerfSleep.thread_sleep(ns)
+            continue
+        else:
+            break
+        
+    if args.step_while_setup:
+        # stop dummy step thread and give algo authority on step
+        dummy_step_exit_req=True
+        # wait for thread to join
+        if dummy_step_thread is not None:
+            dummy_step_thread.join()
+        Journal.log("launch_train_env.py",
+            "",
+            f"Dummy env step thread joined. Moving step authority to algo.",
+            LogType.INFO)
+
+    eval=args.eval
+    if args.override_agent_actions:
+        eval=True
+    if not eval:
+        while not exit_request:
+            if not algo.learn():
+                break
+    else: # eval phase
+        with torch.no_grad(): # no need for grad computation
+            while not exit_request:
+                if not algo.eval():
+                    break
+    
+    algo.done() # make sure to terminate training properly

launch_world_interface.py

@@ -0,0 +1,207 @@
+import os
+import argparse
+import importlib.util
+import inspect
+
+from aug_mpc.utils.rt_factor import RtFactor
+from aug_mpc.utils.custom_arg_parsing import generate_custom_arg_dict
+from aug_mpc.utils.determinism import deterministic_run
+
+from mpc_hive.utilities.shared_data.sim_data import SharedEnvInfo
+
+from EigenIPC.PyEigenIPC import VLevel
+from EigenIPC.PyEigenIPC import Journal, LogType
+
+script_name = os.path.splitext(os.path.basename(os.path.abspath(__file__)))[0]
+
+# Function to dynamically import a module from a specific file path
+def import_world_module(env_path):
+    spec = importlib.util.spec_from_file_location("world_module", env_path)
+    world_module = importlib.util.module_from_spec(spec)
+    spec.loader.exec_module(world_module)
+    return world_module
+
+if __name__ == '__main__':
+
+    parser = argparse.ArgumentParser(description="Sim. env launcher")
+    # Add arguments
+    parser.add_argument('--robot_name', type=str, help='Alias to be used for the robot and also shared memory')
+    parser.add_argument('--urdf_path', type=str, help='path to the URDF file description for each robot')
+    parser.add_argument('--srdf_path', type=str, help='path to the SRDF file description for each robot (used for homing)')
+    parser.add_argument('--jnt_imp_config_path', type=str, help='path to a valid YAML file containing information on jnt impedance gains')
+    parser.add_argument('--num_envs', type=int, default=1)
+    parser.add_argument('--n_contacts', type=int, default=4)
+    parser.add_argument('--cluster_dt', type=float, default=0.03, help='dt at which the control cluster runs')
+    parser.add_argument('--dmpdir', type=str, help='directory where data is dumped', default="/root/aux_data")
+    parser.add_argument('--remote_stepping',action='store_true',
+                help='Whether to use remote stepping for cluster triggering (to be set during training)')
+    
+    # Replacing argparse.BooleanOptionalAction with 'store_true' and 'store_false' for compatibility with Python 3.8
+    parser.add_argument('--use_gpu',action='store_true', help='Whether to use gpu simulation')
+
+    parser.add_argument('--enable_debug',action='store_true', help='Whether to enable debug mode (may introduce significant overhead)')
+
+    parser.add_argument('--headless',action='store_true', help='Whether to run simulation in headless mode')
+
+    parser.add_argument('--verbose',action='store_true', help='Enable verbose mode')
+
+    parser.add_argument('--comment', type=str, help='Any useful comment associated with this run', default="")
+    parser.add_argument('--timeout_ms', type=int, help='connection timeout after which the script self-terminates', default=60000)
+    parser.add_argument('--physics_dt', type=float, default=5e-4, help='')
+
+    parser.add_argument('--use_custom_jnt_imp',action='store_true', 
+        help='Whether to override the default PD controller with a custom one')
+
+    parser.add_argument('--diff_vels',action='store_true', 
+        help='Whether to obtain velocities by differentiation or not')
+    
+    parser.add_argument('--init_timesteps', type=int, help='initialization timesteps', default=None)
+    parser.add_argument('--seed', type=int, help='seed', default=0)
+
+    parser.add_argument('--custom_args_names', nargs='+', default=None,
+                            help='list of custom arguments names')
+    parser.add_argument('--custom_args_vals', nargs='+', default=None,
+                            help='list of custom arguments values')
+    parser.add_argument('--custom_args_dtype', nargs='+', default=None,
+                            help='list of custom arguments data types')
+    
+    parser.add_argument('--world_iface_fname', type=str, 
+        default="aug_mpc_envs.world_interfaces.isaac_world_interface",
+        help="world interface file import pattern (without extension)")
+    
+    args = parser.parse_args()
+    
+    deterministic_run(seed=args.seed, torch_det_algos=False)
+
+    default_init_duration=3.0 # [s]
+    default_init_tsteps=int(default_init_duration/args.physics_dt)
+    init_tsteps=args.init_timesteps 
+    if init_tsteps is None:
+        init_tsteps=default_init_tsteps
+    # Ensure custom_args_names, custom_args_vals, and custom_args_dtype have the same length
+    custom_opt = generate_custom_arg_dict(args=args)
+
+    Journal.log("launch_world_interface.py",
+            "",
+            f"Will warmup world interface for {default_init_duration}s ({default_init_tsteps} physics steps)",
+            LogType.STAT)
+
+    robot_names = [args.robot_name]
+    robot_urdf_paths = [args.urdf_path]
+    robot_srdf_paths = [args.srdf_path]
+    control_clust_dts = [float(args.cluster_dt)]
+    use_remote_stepping = [args.remote_stepping]
+    n_contacts = [args.n_contacts]
+    jnt_imp_config_paths = [args.jnt_imp_config_path]
+    num_envs = args.num_envs
+    control_clust_dt = args.cluster_dt # [s]. Dt at which RHC controllers run 
+    headless = args.headless
+
+    # simulation parameters
+    remote_env_params = {}
+    remote_env_params["physics_dt"] = args.physics_dt # physics_dt?
+    remote_env_params["n_envs"] = num_envs 
+    remote_env_params["use_gpu"] =  args.use_gpu 
+    remote_env_params["substepping_dt"] = control_clust_dts[0]
+    remote_env_params["headless"] = headless
+    remote_env_params["debug_enabled"] = args.enable_debug
+    remote_env_params["seed"] = args.seed
+    remote_env_params.update(custom_opt)
+    # sim info to be broadcasted on shared memory
+    # adding some data to dict for debugging
+
+    shared_sim_infos = []
+    for i in range(len(robot_names)):
+        shared_sim_infos.append(SharedEnvInfo(
+            namespace=robot_names[i],
+            is_server=True, 
+            env_params_dict=remote_env_params,
+            verbose=True,
+            vlevel=VLevel.V2,
+            force_reconnection=True))
+        shared_sim_infos[i].run()
+
+    world_module=importlib.import_module(args.world_iface_fname)
+    classes_in_module = [name for name, obj in inspect.getmembers(world_module, inspect.isclass) 
+                        if obj.__module__ == world_module.__name__]
+    if len(classes_in_module) == 1:
+        cluster_classname=classes_in_module[0]
+        WorldInterface = getattr(world_module, cluster_classname)
+    else:
+        class_list_str = ", ".join(classes_in_module)
+        Journal.log("launch_world_interface.py",
+            "",
+            f"Found more than one class in world file {args.world_iface_fname}. Found: {class_list_str}",
+            LogType.EXCEP,
+            throw_when_excep = False)
+        exit()
+
+    world_interface = WorldInterface(robot_names=robot_names,
+        robot_urdf_paths=robot_urdf_paths,
+        robot_srdf_paths=robot_srdf_paths,
+        cluster_dt=control_clust_dts,
+        jnt_imp_config_paths=jnt_imp_config_paths,
+        n_contacts=n_contacts,
+        use_remote_stepping=use_remote_stepping,
+        name=classes_in_module[0],
+        num_envs=num_envs,
+        debug=args.enable_debug,
+        verbose=args.verbose,
+        vlevel=VLevel.V2,
+        n_init_step=init_tsteps,
+        timeout_ms=args.timeout_ms,
+        env_opts=remote_env_params,
+        use_gpu=args.use_gpu,
+        override_low_lev_controller=args.use_custom_jnt_imp) # create environment
+    # reset_ok=world_interface.reset(reset_sim=True)
+    # if not reset_ok:
+    #     world_interface.close()
+    #     exit()
+
+    rt_factor = RtFactor(dt_nom=world_interface.physics_dt(),
+                window_size=100)
+    
+    while True:
+        
+        if rt_factor.reset_due():
+            rt_factor.reset()
+
+        step_ok=world_interface.step() 
+
+        if not step_ok:
+            break
+
+        rt_factor.update()
+
+        for i in range(len(robot_names)):
+            robot_name=robot_names[i]
+            n_steps = world_interface.cluster_sim_step_counters[robot_name]
+            sol_counter = world_interface.cluster_servers[robot_name].solution_counter()
+            trigger_counter = world_interface.cluster_servers[robot_name].trigger_counter()
+            shared_sim_infos[i].write(dyn_info_name=["sim_rt_factor", 
+                                                "total_rt_factor", 
+                                                "env_stepping_dt",
+                                                "world_stepping_dt",
+                                                "time_to_get_states_from_env",
+                                                "cluster_state_update_dt",
+                                                "cluster_sol_time",
+                                                "n_sim_steps",
+                                                "n_cluster_trigger_steps",
+                                                "n_cluster_sol_steps",
+                                                "sim_time",
+                                                "cluster_time"],
+                                val=[rt_factor.get(), 
+                                    rt_factor.get() * num_envs,
+                                    rt_factor.get_avrg_step_time(),
+                                    world_interface.debug_data["time_to_step_world"],
+                                    world_interface.debug_data["time_to_get_states_from_env"],
+                                    world_interface.debug_data["cluster_state_update_dt"][robot_name],
+                                    world_interface.debug_data["cluster_sol_time"][robot_name],
+                                    n_steps,
+                                    trigger_counter,
+                                    sol_counter,
+                                    world_interface.debug_data["sim_time"][robot_name],
+                                    sol_counter*world_interface.cluster_servers[robot_name].cluster_dt()
+                                    ])
+            
+    world_interface.close()

other/augmpc_cluster_client.py

@@ -0,0 +1,123 @@
+from mpc_hive.cluster_client.control_cluster_client import ControlClusterClient
+from aug_mpc.utils.xrdf_gen import generate_srdf, generate_urdf
+from aug_mpc.utils.hybrid_quad_xrdf_gen import get_xrdf_cmds
+from aug_mpc.utils.custom_arg_parsing import extract_custom_xacro_args, merge_xacro_cmds
+
+from EigenIPC.PyEigenIPC import Journal, LogType
+
+from typing import List, Dict
+
+import os 
+
+from abc import abstractmethod
+
+class AugMpcClusterClient(ControlClusterClient):
+    
+    def _import_aux_libs(self):
+        super()._import_aux_libs()
+        
+    def __init__(self, 
+            namespace: str, 
+            urdf_xacro_path: str,
+            srdf_xacro_path: str,
+            cluster_size: int,
+            set_affinity: bool = False,
+            use_mp_fork: bool = False,
+            isolated_cores_only: bool = False,
+            core_ids_override_list: List[int] = None,
+            verbose: bool = False,
+            debug: bool = False,
+            codegen_base_dirname: str = "CodeGen",
+            base_dump_dir: str = "/tmp",
+            codegen_override: str = None,
+            custom_opts: Dict = {}):
+               
+        self._base_dump_dir = base_dump_dir
+    
+        self._temp_path = base_dump_dir + "/" + f"{self.__class__.__name__}" + f"_{namespace}"
+        
+        self._codegen_base_dirname = codegen_base_dirname
+        self._codegen_basedir = self._temp_path + "/" + self._codegen_base_dirname
+
+        self._codegen_override = codegen_override # can be used to manually override
+        # the default codegen dir 
+
+        if not os.path.exists(self._temp_path):
+            os.makedirs(self._temp_path)
+        if not os.path.exists(self._codegen_basedir):
+            os.makedirs(self._codegen_basedir)
+
+        self._urdf_xacro_path = urdf_xacro_path
+        self._srdf_xacro_path = srdf_xacro_path
+        self._urdf_path=""
+        self._srdf_path=""
+        
+        super().__init__(namespace = namespace, 
+                        cluster_size=cluster_size,
+                        isolated_cores_only = isolated_cores_only,
+                        set_affinity = set_affinity,
+                        use_mp_fork = use_mp_fork,
+                        core_ids_override_list = core_ids_override_list,
+                        verbose = verbose,
+                        debug = debug,
+                        custom_opts=custom_opts)
+        self._generate_srdf(namespace=namespace)
+
+        self._generate_urdf(namespace=namespace)
+
+    
+    def codegen_dir(self):
+
+        return self._codegen_basedir
+    
+    def codegen_dir_override(self):
+
+        return self._codegen_override
+    
+    def _generate_srdf(self,namespace:str):
+        
+        custom_xacro_args=extract_custom_xacro_args(self._custom_opts)
+        cmds=merge_xacro_cmds(prev_cmds=self._xrdf_cmds(),
+            new_cmds=custom_xacro_args)
+
+        # call _xrdf_cmds_override in case some cmds need to be overridden
+        override_cmds=self._xrdf_cmds_override()
+        cmds=merge_xacro_cmds(prev_cmds=cmds,
+            new_cmds=override_cmds)
+        
+        self._srdf_path=generate_srdf(robot_name=namespace,
+            xacro_path=self._srdf_xacro_path,
+            dump_path=self._temp_path,
+            xrdf_cmds=cmds)
+    
+    def _generate_urdf(self,namespace:str):
+        
+        custom_xacro_args=extract_custom_xacro_args(self._custom_opts)
+        cmds=merge_xacro_cmds(prev_cmds=self._xrdf_cmds(),
+            new_cmds=custom_xacro_args)
+
+        # call _xrdf_cmds_override in case some cmds need to be overridden
+        override_cmds=self._xrdf_cmds_override()
+        cmds=merge_xacro_cmds(prev_cmds=cmds,
+            new_cmds=override_cmds)
+        
+        self._urdf_path=generate_urdf(robot_name=namespace,
+            xacro_path=self._urdf_xacro_path,
+            dump_path=self._temp_path,
+            xrdf_cmds=cmds)
+            
+    @abstractmethod
+    def _xrdf_cmds(self):
+        
+        # to be implemented by parent class (
+        # for an example have a look at utils/centauro_xrdf_gen.py)
+
+        pass    
+
+    def _xrdf_cmds_override(self):
+        
+        # to be overridden by parent class 
+
+        to_be_overridden = ["dummy_cmd:=true"]
+
+        return to_be_overridden    

other/augmpc_cluster_server.py

@@ -0,0 +1,43 @@
+from mpc_hive.cluster_server.control_cluster_server import ControlClusterServer
+from typing import List
+from EigenIPC.PyEigenIPC import VLevel
+
+class AugMpcClusterServer(ControlClusterServer):
+
+    def __init__(self, 
+            robot_name: str,
+            cluster_size: int, 
+            cluster_dt: float, 
+            control_dt: float, 
+            jnt_names: List[str],
+            n_contacts: int,
+            contact_linknames: List[str] = None,
+            verbose: bool = False, 
+            vlevel: VLevel = VLevel.V1,
+            debug: bool = False,
+            use_gpu: bool = True,
+            force_reconnection: bool = True,
+            timeout_ms: int = 60000,
+            enable_height_sensor: bool = False,
+            height_grid_size: int = None,
+            height_grid_resolution: float = None):
+
+        self.robot_name = robot_name
+                
+        super().__init__( 
+            namespace=self.robot_name,
+            cluster_size=cluster_size, 
+            cluster_dt=cluster_dt, 
+            control_dt=control_dt, 
+            jnt_names=jnt_names,
+            n_contacts = n_contacts,
+            contact_linknames = contact_linknames, 
+            verbose=verbose, 
+            vlevel=vlevel,
+            debug=debug,
+            use_gpu=use_gpu,
+            force_reconnection=force_reconnection,
+            timeout_ms=timeout_ms,
+            enable_height_sensor=enable_height_sensor,
+            height_grid_size=height_grid_size,
+            height_grid_resolution=height_grid_resolution)

other/horizon_based/config/rhc_horizon_config.yaml

@@ -0,0 +1,200 @@
+# A dummy example of a rhc controller configuration built on top of horizon
+# and iLQR
+
+solver:
+  type: ilqr
+  ilqr.tol: 0.01
+  ilqr.constr_viol_tol: 0.01
+  ilqr.suppress_all_output: 'yes'
+  ilqr.codegen_enabled: true
+#  ilqr.codegen_workdir: /tmp/tyhio
+  ilqr.enable_gn: true
+  ilqr.hxx_reg_base: 0.0
+  ilqr.n_threads: 0
+  print_time: 0
+
+constraints:
+  - ball_1_contact
+  - ball_2_contact
+  - ball_3_contact
+  - ball_4_contact
+  - z_ball_1
+  - z_ball_2
+  - z_ball_3
+  - z_ball_4
+
+costs:
+  - joint_regularization
+  - joint_posture
+  - base_position
+  - base_orientation
+
+.define:
+    - &w_base_pos 10
+    - &w_base_ori 1
+    - &w_ball_z 1
+    # - &w_postural 0.0001
+    - &w_jnt_v_reg 0.01
+    - &w_jnt_a_reg 0.001
+    - &w_jnt_f_reg 0.0001
+    - &wheel_radius 0.124
+
+base_position: 
+  type: Cartesian
+  distal_link: base_link
+  indices: [0, 1, 2]
+  nodes: ${N}
+  weight: *w_base_pos
+
+base_orientation:
+  type: Cartesian
+  distal_link: base_link
+  indices: [3, 4, 5]
+  nodes: ${N}
+  weight: *w_base_ori
+
+# ===============================
+
+rolling_contact_1:
+  type: Rolling
+  frame: wheel_1
+  radius: *wheel_radius
+
+rolling_contact_2:
+  type: Rolling
+  frame: wheel_2
+  radius: *wheel_radius
+
+rolling_contact_3:
+  type: Rolling
+  frame: wheel_3
+  radius: *wheel_radius
+
+rolling_contact_4:
+  type: Rolling
+  frame: wheel_4
+  radius: *wheel_radius
+
+# contact_1:
+#   type: Cartesian
+#   distal_link: ball_1
+#   indices: [0, 1, 2]
+#   cartesian_type: velocity
+
+# contact_2:
+#   type: Cartesian
+#   distal_link: ball_2
+#   indices: [0, 1, 2]
+#   cartesian_type: velocity
+
+# contact_3:
+#   type: Cartesian
+#   distal_link: ball_3
+#   indices: [0, 1, 2]
+#   cartesian_type: velocity
+
+# contact_4:
+#   type: Cartesian
+#   distal_link: ball_4
+#   indices: [0, 1, 2]
+#   cartesian_type: velocity
+
+# ==================================
+
+interaction_wheel_1:
+  type: VertexForce
+  frame: ball_1
+  fn_min: 10.0
+  enable_fc: true
+  friction_coeff: 0.5
+  vertex_frames:
+    - wheel_1
+
+interaction_wheel_2:
+  type: VertexForce
+  frame: ball_2
+  fn_min: 10.0
+  enable_fc: true
+  friction_coeff: 0.5
+  vertex_frames:
+    - wheel_2
+
+interaction_wheel_3:
+  type: VertexForce
+  frame: ball_3
+  fn_min: 10.0
+  enable_fc: true
+  friction_coeff: 0.5
+  vertex_frames:
+    - wheel_3
+
+interaction_wheel_4:
+  type: VertexForce
+  frame: ball_4
+  fn_min: 10.0
+  enable_fc: true
+  friction_coeff: 0.5
+  vertex_frames:
+    - wheel_4
+
+ball_1_contact:
+  type: Contact
+  subtask: [interaction_wheel_1, rolling_contact_1]
+
+ball_2_contact:
+  type: Contact
+  subtask: [interaction_wheel_2, rolling_contact_2]
+
+ball_3_contact:
+  type: Contact
+  subtask: [interaction_wheel_3, rolling_contact_3]
+
+ball_4_contact:
+  type: Contact
+  subtask: [interaction_wheel_4, rolling_contact_4]
+
+# joint_posture:
+#   type: Postural
+#   weight: *w_postural
+#   indices: [0, 1, 2,
+#             4, 5, 6,
+#             8, 9, 10,
+#             12, 13, 14]
+#   nodes: ${range(N-8, N)}
+
+# todo: wrong, as the order COUNTS. If I add the contacts after the joint regularization, they wont get considered.
+joint_regularization:
+  type: Regularization
+  nodes: all # maybe not on first nodes??
+  weight:
+    velocity: *w_jnt_v_reg
+    acceleration: *w_jnt_a_reg 
+    force: *w_jnt_f_reg
+
+z_ball_1:
+  type: Cartesian
+  distal_link: ball_1
+  indices: [2]
+  cartesian_type: position
+  weight: *w_ball_z
+
+z_ball_2:
+  type: Cartesian
+  distal_link: ball_2
+  indices: [2]
+  cartesian_type: position
+  weight: *w_ball_z
+
+z_ball_3:
+  type: Cartesian
+  distal_link: ball_3
+  indices: [2]
+  cartesian_type: position
+  weight: *w_ball_z
+
+z_ball_4:
+  type: Cartesian
+  distal_link: ball_4
+  indices: [2]
+  cartesian_type: position
+  weight: *w_ball_z

other/horizon_based/gait_manager.py

@@ -0,0 +1,566 @@
+import numpy as np
+
+from aug_mpc.controllers.rhc.horizon_based.horizon_imports import *
+
+from EigenIPC.PyEigenIPC import VLevel
+from EigenIPC.PyEigenIPC import Journal, LogType
+
+from typing import Dict
+
+class GaitManager:
+
+    def __init__(self, 
+            task_interface: TaskInterface, 
+            phase_manager: pymanager.PhaseManager, 
+            injection_node: int = None,
+            keep_yaw_vert: bool = False,
+            yaw_vertical_weight: float = None,
+            vertical_landing: bool = False,
+            landing_vert_weight: float = None,
+            phase_force_reg: float = None,
+            flight_duration: int = 15,
+            post_flight_stance: int = 3,
+            step_height: float = 0.1,
+            dh: float = 0.0,
+            custom_opts: Dict = {}):
+
+        self._custom_opts=custom_opts
+
+        self._is_open_loop=False
+        if "is_open_loop" in self._custom_opts:
+            self._is_open_loop=self._custom_opts["is_open_loop"]
+
+        self.task_interface = task_interface
+        self._n_nodes_prb=self.task_interface.prb.getNNodes()
+        
+        self._phase_manager = phase_manager
+        self._model=self.task_interface.model
+        self._q0=self._model.q0
+        self._kin_dyn=self.task_interface.model.kd
+        
+        # phase weights and regs
+        self._keep_yaw_vert=keep_yaw_vert
+        self._yaw_vertical_weight=yaw_vertical_weight
+        self._vertical_landing=vertical_landing
+        self._landing_vert_weight=landing_vert_weight
+        self._phase_force_reg=phase_force_reg
+        self._total_weight = np.atleast_2d(np.array([0, 0, self._kin_dyn.mass() * 9.81])).T 
+        
+        self._f_reg_ref={}
+
+        # flight parameters
+        self._post_flight_stance=post_flight_stance
+        self._flight_info_now=None 
+        self._flight_duration_max=self._n_nodes_prb-(injection_node+1)
+        self._flight_duration_min=3
+        self._flight_duration_default=flight_duration 
+        # apex bounds/defaults
+        self._step_height_default=step_height
+        self._step_height_min=0.0
+        self._step_height_max=0.5
+        # end height bounds/defaults
+        self._dh_default=dh
+        self._dh_min=0.0
+        self._dh_max=0.5
+        # landing dx, dy bounds/defaults
+        self._land_dx_default=0.0
+        self._land_dx_min=-0.5
+        self._land_dx_max=0.5
+        self._land_dy_default=0.0       
+        self._land_dy_min=-0.5
+        self._land_dy_max=0.5
+        
+        # timeline data
+        self._contact_timelines = dict()
+        self.timeline_names=[]
+
+        self._flight_phases = {}
+        self._touchdown_phases = {}
+        self._contact_phases = {}
+        self._fk_contacts = {}
+        self._fkd_contacts = {}
+        self._f_reg_ref = {}
+
+        # reference traj
+        self._tg = trajectoryGenerator.TrajectoryGenerator()
+        self._traj_der= [None, 0, 0]
+        self._traj_second_der=[None, 0, 0]
+        self._third_traj_der=[None, None, 0]
+
+        self._ref_trjs = {}
+        self._ref_vtrjs = {}
+
+        if injection_node is None:
+            self._injection_node = round(self.task_interface.prb.getNNodes()/2.0)
+        else:
+            self._injection_node = injection_node
+        
+        self._init_contact_timelines()  
+        
+        self._reset_contact_timelines()
+
+    def _init_contact_timelines(self):
+
+        short_stance_duration=1
+        flight_phase_short_duration=1
+
+        self.n_contacts=len(self._model.cmap.keys()) 
+        self._dt=float(self.task_interface.prb.getDt())
+
+        self._name_to_idx_map={}
+
+        j=0
+        for contact in self._model.cmap.keys():
+            
+            self._fk_contacts[contact]=self._kin_dyn.fk(contact)
+            self._fkd_contacts[contact]=self._kin_dyn.frameVelocity(contact, self._model.kd_frame)
+            self.timeline_names.append(contact)
+            self._contact_timelines[contact]=self._phase_manager.createTimeline(f'{contact}_timeline')
+            # stances
+            self._contact_phases[contact] = self._contact_timelines[contact].createPhase(short_stance_duration, 
+                                    f'stance_{contact}_short')
+                
+            if self.task_interface.getTask(f'{contact}') is not None:
+                self._contact_phases[contact].addItem(self.task_interface.getTask(f'{contact}'))
+            else:
+                Journal.log(self.__class__.__name__,
+                    "_init_contact_timelines",
+                    f"contact task {contact} not found",
+                    LogType.EXCEP,
+                    throw_when_excep=True)
+            i=0
+            self._f_reg_ref[contact]=[]
+            for force in self._model.cmap[contact]:
+                f_ref=self.task_interface.prb.createParameter(name=f"{contact}_force_reg_f{i}_ref",
+                    dim=3)
+                force_reg=self.task_interface.prb.createResidual(f'{contact}_force_reg_f{i}', self._phase_force_reg * (force - f_ref), 
+                    nodes=[])
+                self._f_reg_ref[contact].append(f_ref)
+                self.set_f_reg(contact_name=contact, scale=4)                
+                self._contact_phases[contact].addCost(force_reg, nodes=list(range(0, short_stance_duration)))
+                i+=1
+            
+            # flights
+            self._flight_phases[contact]=self._contact_timelines[contact].createPhase(flight_phase_short_duration, 
+                                    f'flight_{contact}_short')
+
+            # sanity checks (z trajectory)
+            self._zpos_task_found=True
+            self._zvel_task_found=True
+            self._xypos_task_found=True
+            self._xyvel_task_found=True
+            if self.task_interface.getTask(f'z_{contact}') is None:
+                self._zpos_task_found=False
+            if self.task_interface.getTask(f'vz_{contact}') is None:
+                self._zvel_task_found=False
+            if self.task_interface.getTask(f'xy_{contact}') is None:
+                self._xypos_task_found=False
+            if self.task_interface.getTask(f'vxy_{contact}') is None:
+                self._xyvel_task_found=False
+            if (not self._zpos_task_found) and (not self._zvel_task_found):
+                Journal.log(self.__class__.__name__,
+                    "_init_contact_timelines",
+                    f"neither pos or vel task for contacts were found! Aborting.",
+                    LogType.EXCEP,
+                    throw_when_excep=True)
+            if (not self._zpos_task_found) and self._is_open_loop:
+                Journal.log(self.__class__.__name__,
+                    "_init_contact_timelines",
+                    f"Running in open loop, but no contact pos task found. Aborting.",
+                    LogType.EXCEP,
+                    throw_when_excep=True)
+            if self._zpos_task_found and self._xyvel_task_found:
+                Journal.log(self.__class__.__name__,
+                    "_init_contact_timelines",
+                    f"Both pos and vel task for contact {contact} found! This is not allowed, aborting.",
+                    LogType.EXCEP,
+                    throw_when_excep=True)
+            if self._zvel_task_found and self._xypos_task_found:
+                Journal.log(self.__class__.__name__,
+                    "_init_contact_timelines",
+                    f"Both pos and vel task for contact {contact} found! This is not allowed, aborting.",
+                    LogType.EXCEP,
+                    throw_when_excep=True)
+            if (not self._xypos_task_found) and (not self._xyvel_task_found):
+                Journal.log(self.__class__.__name__,
+                    "_init_contact_timelines",
+                    f"neither pos or vel task for contact {contact} xy were found! Will proceed without xy landing constraints.",
+                    LogType.WARN)
+            # if (not self._zvel_task_found) and (not self._is_open_loop):
+            #     Journal.log(self.__class__.__name__,
+            #         "_init_contact_timelines",
+            #         f"Running in closed loop, but contact vel task not found. Aborting",
+            #         LogType.EXCEP,
+            #         throw_when_excep=True)
+            
+            self._ref_trjs[contact]=None
+            self._ref_vtrjs[contact]=None
+            self._touchdown_phases[contact]=None
+
+            if self._zpos_task_found: # we use pos trajectory
+                self._ref_trjs[contact]=np.zeros(shape=[7, self.task_interface.prb.getNNodes()])
+                init_z_foot = self._fk_contacts[contact](q=self._q0)['ee_pos'].elements()[2]
+                if self._is_open_loop:
+                    self._ref_trjs[contact][2, :] = np.atleast_2d(init_z_foot)
+                else:
+                    self._ref_trjs[contact][2, :] = 0.0 # place foot at ground level initially ()
+                
+                # z
+                self._flight_phases[contact].addItemReference(self.task_interface.getTask(f'z_{contact}'), 
+                    self._ref_trjs[contact][2, 0:1], 
+                    nodes=list(range(0, flight_phase_short_duration)))
+
+                if self._xypos_task_found: # xy, we add a landing phase of unit duration to enforce landing pos costs
+                    
+                    self._touchdown_phases[contact]=self._contact_timelines[contact].createPhase(flight_phase_short_duration, 
+                                    f'touchdown_{contact}_short') 
+
+                    self._touchdown_phases[contact].addItemReference(self.task_interface.getTask(f'xy_{contact}'), 
+                        self._ref_trjs[contact][0:2, 0:1], 
+                        nodes=list(range(0, short_stance_duration)))
+                    
+            else: # foot traj in velocity
+                # ref vel traj
+                self._ref_vtrjs[contact]=np.zeros(shape=[7, self.task_interface.prb.getNNodes()]) # allocate traj
+                # of max length eual to number of nodes
+                self._ref_vtrjs[contact][2, :] = np.atleast_2d(0) 
+
+                # z
+                self._flight_phases[contact].addItemReference(self.task_interface.getTask(f'vz_{contact}'), 
+                    self._ref_vtrjs[contact][2, 0:1], 
+                    nodes=list(range(0, flight_phase_short_duration)))
+                
+                if self._xyvel_task_found: # xy (when in vel the xy vel is set on the whole flight phase)
+                    self._flight_phases[contact].addItemReference(self.task_interface.getTask(f'vxy_{contact}'), 
+                        self._ref_vtrjs[contact][0:2, 0:1], 
+                        nodes=list(range(0, flight_phase_short_duration)))
+                    
+                if self._vertical_landing: # add touchdown phase for vertical landing
+                    self._touchdown_phases[contact]=self._contact_timelines[contact].createPhase(flight_phase_short_duration, 
+                                        f'touchdown_{contact}_short')
+
+            if self._vertical_landing and self._touchdown_phases[contact] is not None:
+                v_xy=self._fkd_contacts[contact](q=self._model.q, qdot=self._model.v)['ee_vel_linear'][0:2]
+                vertical_landing=self.task_interface.prb.createResidual(f'{contact}_only_vert_v', 
+                    self._landing_vert_weight * v_xy, 
+                    nodes=[])
+                self._touchdown_phases[contact].addCost(vertical_landing, nodes=list(range(0, short_stance_duration)))
+
+            if self._keep_yaw_vert:
+                # keep ankle vertical over the whole horizon (can be useful with wheeled robots)
+                c_ori = self._model.kd.fk(contact)(q=self._model.q)['ee_rot'][2, :]
+                cost_ori = self.task_interface.prb.createResidual(f'{contact}_ori', self._yaw_vertical_weight * (c_ori.T - np.array([0, 0, 1])))
+                # flight_phase.addCost(cost_ori, nodes=list(range(0, flight_duration+post_landing_stance)))
+            
+            self._name_to_idx_map[contact]=j
+
+            j+=1
+
+        # current pos [c0, c1, ....], current length, nominal length, nom. apex, no. landing height, landing dx, landing dy (local world aligned) 
+        self._flight_info_now=np.zeros(shape=(7*self.n_contacts))
+
+        self.update()
+
+    def _reset_contact_timelines(self):
+
+        for contact in self._model.cmap.keys():
+            
+            idx=self._name_to_idx_map[contact]
+            # we follow same order as in shm for more efficient writing 
+            self._flight_info_now[idx]= -1.0 # pos [nodes]
+            self._flight_info_now[idx+1*self.n_contacts]= -1.0 # duration (remaining) [nodes]
+            self._flight_info_now[idx+2*self.n_contacts]=self._flight_duration_default # [nodes]
+            self._flight_info_now[idx+3*self.n_contacts]=self._step_height_default
+            self._flight_info_now[idx+4*self.n_contacts]=self._dh_default
+            self._flight_info_now[idx+5*self.n_contacts]=self._land_dx_default
+            self._flight_info_now[idx+6*self.n_contacts]=self._land_dy_default
+            # fill timeline with stances
+            contact_timeline=self._contact_timelines[contact]
+            contact_timeline.clear() # remove phases
+            short_stance_phase = contact_timeline.getRegisteredPhase(f'stance_{contact}_short')
+            while contact_timeline.getEmptyNodes() > 0:
+                contact_timeline.addPhase(short_stance_phase)   
+            
+            self.update()
+            
+    def reset(self):
+        # self.phase_manager.clear()
+        self.task_interface.reset()
+        self._reset_contact_timelines()
+    
+    def set_f_reg(self, 
+        contact_name,
+        scale: float = 4.0):
+        f_refs=self._f_reg_ref[contact_name]
+        for force in f_refs:
+            ref=self._total_weight/(scale*len(f_refs))
+            force.assign(ref)
+    
+    def set_flight_duration(self, contact_name, val: float):
+        self._flight_info_now[self._name_to_idx_map[contact_name]+2*self.n_contacts]=val
+    
+    def get_flight_duration(self, contact_name):
+        return self._flight_info_now[self._name_to_idx_map[contact_name]+2*self.n_contacts]
+    
+    def set_step_apexdh(self, contact_name, val: float):
+        self._flight_info_now[self._name_to_idx_map[contact_name]+3*self.n_contacts]=val
+    
+    def get_step_apexdh(self, contact_name):
+        return self._flight_info_now[self._name_to_idx_map[contact_name]+3*self.n_contacts]
+    
+    def set_step_enddh(self, contact_name, val: float):
+        self._flight_info_now[self._name_to_idx_map[contact_name]+4*self.n_contacts]=val
+    
+    def get_step_enddh(self, contact_name):
+        return self._flight_info_now[self._name_to_idx_map[contact_name]+4*self.n_contacts]
+    
+    def get_step_landing_dx(self, contact_name):
+        return self._flight_info_now[self._name_to_idx_map[contact_name]+5*self.n_contacts]
+    
+    def set_step_landing_dx(self, contact_name, val: float):
+        self._flight_info_now[self._name_to_idx_map[contact_name]+5*self.n_contacts]=val
+    
+    def get_step_landing_dy(self, contact_name):
+        return self._flight_info_now[self._name_to_idx_map[contact_name]+6*self.n_contacts]     
+    
+    def set_step_landing_dy(self, contact_name, val: float):
+        self._flight_info_now[self._name_to_idx_map[contact_name]+6*self.n_contacts]=val
+
+    def add_stand(self, contact_name):
+        # always add stand at the end of the horizon
+        timeline = self._contact_timelines[contact_name]
+        if timeline.getEmptyNodes() > 0:
+            timeline.addPhase(timeline.getRegisteredPhase(f'stance_{contact_name}_short'))
+    
+    def add_flight(self, contact_name,
+        robot_q: np.ndarray):
+
+        timeline = self._contact_timelines[contact_name]
+
+        flights_on_horizon=self._contact_timelines[contact_name].getPhaseIdx(self._flight_phases[contact_name]) 
+        
+        last_flight_idx=self._injection_node-1 # default to make things work
+        if not len(flights_on_horizon)==0: # some flight phases are there
+            last_flight_idx=flights_on_horizon[-1]+self._post_flight_stance
+
+        if last_flight_idx<self._injection_node: # allow injecting
+            
+            flight_duration_req=int(self.get_flight_duration(contact_name=contact_name))
+            flight_apex_req=self.get_step_apexdh(contact_name=contact_name)
+            flight_enddh_req=self.get_step_enddh(contact_name=contact_name)
+            flight_land_dx_req=self.get_step_landing_dx(contact_name=contact_name)
+            flight_land_dy_req=self.get_step_landing_dy(contact_name=contact_name)
+            if not flight_duration_req>1:
+                Journal.log(self.__class__.__name__,
+                    "add_flight",
+                    f"Got flight duration {flight_duration_req} < 1!",
+                    LogType.WARN,
+                    throw_when_excep=True)
+
+            # process requests to ensure flight params are valid
+            # duration
+            if flight_duration_req<self._flight_duration_min:
+                flight_duration_req=self._flight_duration_min
+            if flight_duration_req>self._flight_duration_max:
+                flight_duration_req=self._flight_duration_max
+            # apex height
+            if flight_apex_req<self._step_height_min:
+                flight_apex_req=self._step_height_min
+            if flight_apex_req>self._step_height_max:
+                flight_apex_req=self._step_height_max
+            # landing height
+            if flight_enddh_req<self._dh_min:
+                flight_enddh_req=self._dh_min
+            if flight_enddh_req>self._dh_max:
+                flight_enddh_req=self._dh_max
+            # landing dx
+            if flight_land_dx_req<self._land_dx_min:
+                flight_land_dx_req=self._land_dx_min
+            if flight_land_dx_req>self._land_dx_max:
+                flight_land_dx_req=self._land_dx_max
+            # landing dy                
+            if flight_land_dy_req<self._land_dy_min:
+                flight_land_dy_req=self._land_dy_min
+            if flight_land_dy_req>self._land_dy_max:
+                flight_land_dy_req=self._land_dy_max   
+
+            land_dx_w = flight_land_dx_req
+            land_dy_w = flight_land_dy_req
+            if self._xypos_task_found or self._xyvel_task_found:
+                # landing dx/dy are specified in horizontal frame; rotate into world aligned frame
+                q_base = robot_q[3:7]
+                if q_base.ndim == 1:
+                    q_base = q_base.reshape(-1, 1)
+                q_w = q_base[3, 0]
+                q_x = q_base[0, 0]
+                q_y = q_base[1, 0]
+                q_z = q_base[2, 0]
+                r11 = 1 - 2 * (q_y * q_y + q_z * q_z)
+                r21 = 2 * (q_x * q_y + q_z * q_w)
+                norm = np.hypot(r11, r21)
+                if norm > 0.0:
+                    cos_yaw = r11 / norm
+                    sin_yaw = r21 / norm
+                else:
+                    cos_yaw = 1.0
+                    sin_yaw = 0.
+                    
+                land_dx_w = flight_land_dx_req * cos_yaw - flight_land_dy_req * sin_yaw
+                land_dy_w = flight_land_dx_req * sin_yaw + flight_land_dy_req * cos_yaw
+
+            if self._ref_vtrjs[contact_name] is not None and \
+                self._ref_trjs[contact_name] is not None: # only allow one mode (pos/velocity traj)
+                Journal.log(self.__class__.__name__,
+                    "add_flight",
+                    f"Both pos and vel traj for contact {contact_name} are not None! This is not allowed, aborting.",
+                    LogType.EXCEP,
+                    throw_when_excep=True)
+                
+            # inject pos traj if pos mode
+            if self._ref_trjs[contact_name] is not None:
+                # recompute trajectory online (just needed if using pos traj)
+                foot_pos=self._fk_contacts[contact_name](q=robot_q)['ee_pos'].elements()
+                starting_pos=foot_pos[2] # compute foot traj (local world aligned)
+                starting_x_pos=foot_pos[0]
+                starting_y_pos=foot_pos[1]
+                # starting_pos=0.0
+                self._ref_trjs[contact_name][2, 0:flight_duration_req]=np.atleast_2d(self._tg.from_derivatives(
+                    flight_duration_req, 
+                    p_start=starting_pos, 
+                    p_goal=starting_pos+flight_enddh_req, 
+                    clearance=flight_apex_req,
+                    derivatives=self._traj_der,
+                    second_der=self._traj_second_der,
+                    third_der=self._third_traj_der)
+                    )
+                if self._xypos_task_found: # we use _ref_trjs to write xy pos references
+                    self._ref_trjs[contact_name][0, -1]=starting_x_pos+land_dx_w
+                    self._ref_trjs[contact_name][1, -1]=starting_y_pos+land_dy_w
+
+                for i in range(flight_duration_req):
+                    res, phase_token_flight=timeline.addPhase(self._flight_phases[contact_name], 
+                        pos=self._injection_node+i, 
+                        absolute_position=True)
+                    phase_token_flight.setItemReference(f'z_{contact_name}',
+                        self._ref_trjs[contact_name][:, i])
+                    
+                if self._touchdown_phases[contact_name] is not None:
+                    # add touchdown phase after flight
+                    res, phase_token_touchdown=timeline.addPhase(self._touchdown_phases[contact_name], 
+                            pos=self._injection_node+flight_duration_req, 
+                            absolute_position=True)    
+                    if self._xypos_task_found:
+                        phase_token_touchdown.setItemReference(f'xy_{contact_name}',
+                            self._ref_trjs[contact_name][:, -1])                    
+                    
+            # inject vel traj if vel mode
+            if self._ref_vtrjs[contact_name] is not None:
+                self._ref_vtrjs[contact_name][2, 0:flight_duration_req]=np.atleast_2d(self._tg.derivative_of_trajectory(
+                    flight_duration_req, 
+                    p_start=0.0, 
+                    p_goal=flight_enddh_req, 
+                    clearance=flight_apex_req,
+                    derivatives=self._traj_der,
+                    second_der=self._traj_second_der,
+                    third_der=self._third_traj_der))
+                if self._xyvel_task_found: # compute vel reference using problem dt and flight length
+                    flight_duration_sec=float(flight_duration_req)*self._dt
+                    self._ref_vtrjs[contact_name][0, 0:flight_duration_req]=land_dx_w/flight_duration_sec
+                    self._ref_vtrjs[contact_name][1, 0:flight_duration_req]=land_dy_w/flight_duration_sec
+                    
+                for i in range(flight_duration_req):
+                    res, phase_token=timeline.addPhase(self._flight_phases[contact_name], 
+                        pos=self._injection_node+i, 
+                        absolute_position=True)
+                    phase_token.setItemReference(f'vz_{contact_name}',
+                        self._ref_vtrjs[contact_name][2:3, i:i+1])
+                if self._touchdown_phases[contact_name] is not None:
+                    # add touchdown phase for forcing vertical landing
+                    res, phase_token=timeline.addPhase(self._touchdown_phases[contact_name], 
+                            pos=self._injection_node+flight_duration_req, 
+                            absolute_position=True)       
+
+        if timeline.getEmptyNodes() > 0: # fill empty nodes at the end of the horizon, if any, with stance
+            timeline.addPhase(timeline.getRegisteredPhase(f'stance_{contact_name}_short'))
+    
+    def update(self):
+        self._phase_manager.update()
+        
+    def update_flight_info(self, timeline_name):
+
+        # get current position and remaining duration of flight phases over the horizon for a single contact
+
+        # phase indexes over timeline
+        phase_idxs=self._contact_timelines[timeline_name].getPhaseIdx(self._flight_phases[timeline_name])
+        
+        if not len(phase_idxs)==0: # at least one flight phase on horizon -> read info from timeline
+
+            # all active phases on timeline
+            active_phases=self._contact_timelines[timeline_name].getActivePhases()
+
+            phase_idx_start=phase_idxs[0]
+            # active_nodes_start=active_phases[phase_idx_start].getActiveNodes()
+            pos_start=active_phases[phase_idx_start].getPosition()
+            # n_nodes_start=active_phases[phase_idx_start].getNNodes()
+
+            phase_idx_last=phase_idxs[-1] # just get info for last phase on the horizon
+            # active_nodes_last=active_phases[phase_idx_last].getActiveNodes()
+            pos_last=active_phases[phase_idx_last].getPosition()
+            # n_nodes_last=active_phases[phase_idx_last].getNNodes()
+            
+            # write to info
+            self._flight_info_now[self._name_to_idx_map[timeline_name]+0*self.n_contacts]=pos_last
+            self._flight_info_now[self._name_to_idx_map[timeline_name]+1*self.n_contacts]=pos_last - pos_start
+
+            return True
+        
+        return False
+    
+    def get_flight_info(self, timeline_name):
+        return (self._flight_info_now[self._name_to_idx_map[timeline_name]+0*self.n_contacts], 
+            self._flight_info_now[self._name_to_idx_map[timeline_name]+1*self.n_contacts],
+            self._flight_info_now[self._name_to_idx_map[timeline_name]+2*self.n_contacts],
+            self._flight_info_now[self._name_to_idx_map[timeline_name]+3*self.n_contacts],
+            self._flight_info_now[self._name_to_idx_map[timeline_name]+4*self.n_contacts],
+            self._flight_info_now[self._name_to_idx_map[timeline_name]+5*self.n_contacts],
+            self._flight_info_now[self._name_to_idx_map[timeline_name]+6*self.n_contacts])
+    
+    def get_flight_info_all(self):
+        return self._flight_info_now
+    
+    def set_ref_pos(self,
+        timeline_name:str,
+        ref_height: np.array = None,
+        threshold: float = 0.05):
+        
+        if ref_height is not None:
+            self._ref_trjs[timeline_name][2, :]=ref_height
+            if ref_height>threshold:
+                self.add_flight(timeline_name=timeline_name)
+                this_flight_token_idx=self._contact_timelines[timeline_name].getPhaseIdx(self._flight_phases[timeline_name])[-1]
+                active_phases=self._contact_timelines[timeline_name].getActivePhases()
+                active_phases[this_flight_token_idx].setItemReference(f'z_{timeline_name}',
+                    self._ref_trjs[timeline_name])
+            else:
+                self.add_stand(timeline_name=timeline_name)
+
+    def set_force_feedback(self,
+        timeline_name: str,
+        force_norm: float):
+        
+        flight_tokens=self._contact_timelines[timeline_name].getPhaseIdx(self._flight_phases[timeline_name])
+        contact_tokens=self._contact_phases[timeline_name].getPhaseIdx(self._contact_phases[timeline_name])
+        if not len(flight_tokens)==0:
+            first_flight=flight_tokens[0]
+            first_flight
+    
+    def check_horizon_full(self,
+        timeline_name):
+        timeline = self._contact_timelines[timeline_name]
+        if timeline.getEmptyNodes() > 0:
+            error = f"Empty nodes detected over the horizon for timeline {timeline}! Make sure to fill the whole horizon with valid phases!!"
+            Journal.log(self.__class__.__name__,
+                "check_horizon_full",
+                error,
+                LogType.EXCEP,
+                throw_when_excep = True)

other/horizon_based/horizon_imports.py

@@ -0,0 +1,18 @@
+# robot modeling and automatic differentiation
+import casadi_kin_dyn.py3casadi_kin_dyn as casadi_kin_dyn
+import casadi as cs
+
+# horizon stuff
+import horizon.utils.kin_dyn as kd
+from horizon.problem import Problem
+from horizon.rhc.model_description import FullModelInverseDynamics
+from horizon.rhc.taskInterface import TaskInterface
+from horizon.rhc.tasks.interactionTask import VertexContact
+from horizon.rhc.tasks.contactTask import ContactTask
+from horizon.utils import trajectoryGenerator, utils
+
+# phase managing
+import phase_manager.pymanager as pymanager
+import phase_manager.pyphase as pyphase
+import phase_manager.pytimeline as pytimeline
+

other/horizon_based/horizon_imports_glob.py

@@ -0,0 +1,28 @@
+"""
+Dynamically import all necessary Horizon and related dependencies.
+This function is intended to be used within the import_child_lib method
+of a class, enabling the parent process to load all required libraries.
+"""
+def import_horizon_global():
+    # Global imports to make modules accessible in child processes
+    global casadi_kin_dyn, cs, kd, Problem, FullModelInverseDynamics
+    global TaskInterface, VertexContact, ContactTask, trajectoryGenerator, utils
+    global pymanager, pyphase, pytimeline
+
+    # robot modeling and automatic differentiation
+    import casadi_kin_dyn.py3casadi_kin_dyn as casadi_kin_dyn
+    import casadi as cs
+
+    # horizon stuff
+    import horizon.utils.kin_dyn as kd
+    from horizon.problem import Problem
+    from horizon.rhc.model_description import FullModelInverseDynamics
+    from horizon.rhc.taskInterface import TaskInterface
+    from horizon.rhc.tasks.interactionTask import VertexContact
+    from horizon.rhc.tasks.contactTask import ContactTask
+    from horizon.utils import trajectoryGenerator, utils
+
+    # phase managing
+    import phase_manager.pymanager as pymanager
+    import phase_manager.pyphase as pyphase
+    import phase_manager.pytimeline as pytimeline

other/horizon_based/hybrid_quad_rhc.py

@@ -0,0 +1,1324 @@
+from mpc_hive.controllers.rhc import RHController
+
+from aug_mpc.controllers.rhc.horizon_based.horizon_imports import *
+
+from aug_mpc.controllers.rhc.horizon_based.hybrid_quad_rhc_refs import HybridQuadRhcRefs
+from aug_mpc.controllers.rhc.horizon_based.gait_manager import GaitManager
+
+from EigenIPC.PyEigenIPC import VLevel
+from EigenIPC.PyEigenIPC import Journal, LogType
+
+import numpy as np
+
+import os
+
+import time
+
+from typing import Dict, List
+
+class HybridQuadRhc(RHController):
+
+    def __init__(self, 
+            srdf_path: str,
+            urdf_path: str,
+            config_path: str,
+            robot_name: str, # used for shared memory namespaces
+            codegen_dir: str, 
+            n_nodes:float = 25,
+            injection_node:int = 10,
+            dt: float = 0.02,
+            max_solver_iter = 1, # defaults to rt-iteration
+            open_loop: bool = True,
+            close_loop_all: bool = False,
+            dtype = np.float32,
+            verbose = False, 
+            debug = False,
+            refs_in_hor_frame = True,
+            timeout_ms: int = 60000,
+            custom_opts: Dict = {}):
+
+        self._refs_in_hor_frame = refs_in_hor_frame
+
+        self._injection_node = injection_node
+
+        self._open_loop = open_loop
+        self._close_loop_all = close_loop_all
+
+        self._codegen_dir = codegen_dir
+        if not os.path.exists(self._codegen_dir):
+            os.makedirs(self._codegen_dir)
+        # else:
+        #     # Directory already exists, delete it and recreate
+        #     shutil.rmtree(self._codegen_dir)
+        #     os.makedirs(self._codegen_dir)
+
+        self.step_counter = 0
+        self.sol_counter = 0
+    
+        self.max_solver_iter = max_solver_iter
+        
+        self._timer_start = time.perf_counter()
+        self._prb_update_time = time.perf_counter()
+        self._phase_shift_time = time.perf_counter()
+        self._task_ref_update_time = time.perf_counter()
+        self._rti_time = time.perf_counter()
+
+        self.robot_name = robot_name
+        
+        self.config_path = config_path
+
+        self.urdf_path = urdf_path
+        # read urdf and srdf files
+        with open(self.urdf_path, 'r') as file:
+            self.urdf = file.read()
+        self._base_init = np.array([0.0, 0.0, 0.0, 0.0, 0.0, 0.0, 1.0])
+
+        self._c_timelines = dict()
+        self._f_reg_timelines = dict()
+        
+        self._custom_opts={"replace_continuous_joints": False,
+            "use_force_feedback": False,
+            "lin_a_feedback": False,
+            "is_open_loop": self._open_loop, # fully open (just for db)
+            "fully_closed": False, # closed loop with full feedback (just for db)
+            "closed_partial": False, # closed loop with partial feedback
+            "adaptive_is": True, # closed loop with adaptation
+            "estimate_v_root": False, # when adaptive_is or closed_partial, estimate vbase
+            "alpha_from_outside": False, # alpha set ext. from shared memory
+            "alpha_half": 1.0, 
+            "only_vel_wheels": True, # whether wheels (if present) are just vel controlled
+            "use_jnt_v_feedback": False
+            }
+        
+        self._custom_opts.update(custom_opts)
+        
+        self._alpha_half=self._custom_opts["alpha_half"]
+
+        if self._open_loop:
+            self._custom_opts["fully_closed"]=False
+            self._custom_opts["adaptive_is"]=False
+            self._custom_opts["closed_partial"]=False
+        else:
+            self._custom_opts["is_open_loop"]=False
+            if self._custom_opts["fully_closed"]:
+                self._custom_opts["adaptive_is"]=False
+                self._custom_opts["closed_partial"]=False
+                self._custom_opts["lin_a_feedback"]=False
+            if self._custom_opts["closed_partial"]:
+                self._custom_opts["adaptive_is"]=False
+                self._custom_opts["fully_closed"]=False
+                self._custom_opts["lin_a_feedback"]=False
+            if self._custom_opts["adaptive_is"]:
+                self._custom_opts["closed_partial"]=False
+                self._custom_opts["fully_closed"]=False
+
+        super().__init__(srdf_path=srdf_path,
+                        n_nodes=n_nodes,
+                        dt=dt,
+                        namespace=self.robot_name,
+                        dtype=dtype,
+                        verbose=verbose, 
+                        debug=debug,
+                        timeout_ms=timeout_ms)
+
+        self._rhc_fpaths.append(self.config_path)
+
+    def _config_override(self):
+        pass
+
+    def _post_problem_init(self):
+
+        self.rhc_costs={}
+        self.rhc_constr={}
+
+        self._fail_idx_scale=0.0
+        self._expl_idx_window_size=int(1*self._n_nodes) 
+        self._explosion_idx_buffer=np.zeros((1,self._expl_idx_window_size))
+        self._expl_idx_counter=0
+        self._expl_idx_buffer_counter=0
+
+        self._pred_node_idx=self._n_nodes-1
+
+        self._nq=self.nq()
+        self._nq_jnts=self._nq-7# assuming floating base
+        self._nv=self.nv()
+        self._nv_jnts=self._nv-6
+
+        self._alphas_q_root=np.zeros((7, 1), dtype=self._dtype)
+        self._alphas_q_jnts=np.zeros((self._nq_jnts, 1), dtype=self._dtype)
+        self._alphas_twist_root=np.zeros((6, 1), dtype=self._dtype)
+        self._alphas_v_jnts=np.zeros((self._nv_jnts, 1), dtype=self._dtype)
+        self._alphas_a=np.zeros((self._nv, 1), dtype=self._dtype)
+        self._alphas_q_root[:, :]=1.0 # default to all open
+        self._alphas_q_jnts[:, :]=1.0 
+        self._alphas_twist_root[:, :]=1.0 
+        self._alphas_v_jnts[:, :]=1.0 
+        self._alphas_a[:, :]=1.0
+
+    def _init_problem(self,
+            fixed_jnt_patterns: List[str] = None,
+            wheels_patterns: List[str] = None,
+            foot_linkname: str = None,
+            flight_duration: int = 10,
+            post_flight_stance: int = 3,
+            step_height: float = 0.12,
+            keep_yaw_vert: bool = False,
+            yaw_vertical_weight: float = 2.0,
+            vertical_landing: bool = False,
+            vertical_land_weight: float = 1.0,
+            phase_force_reg: float = 1e-2,
+            vel_bounds_weight: float = 1.0):
+        
+        self._fixed_jnt_patterns=fixed_jnt_patterns
+
+        self._config_override()
+        
+        Journal.log(self.__class__.__name__,
+            "_init_problem",
+            f" Will use horizon config file at {self.config_path}",
+            LogType.INFO,
+            throw_when_excep=True)
+        
+        self._vel_bounds_weight=vel_bounds_weight
+        self._phase_force_reg=phase_force_reg
+        self._yaw_vertical_weight=yaw_vertical_weight
+        self._vertical_land_weight=vertical_land_weight
+        # overrides parent
+        self._prb = Problem(self._n_intervals, 
+                        receding=True, 
+                        casadi_type=cs.SX)
+        self._prb.setDt(self._dt)
+
+        if "replace_continuous_joints" in self._custom_opts:
+            # continous joints are parametrized in So2
+            if self._custom_opts["replace_continuous_joints"]:
+                self.urdf = self.urdf.replace('continuous', 'revolute')
+        else:
+            self.urdf = self.urdf.replace('continuous', 'revolute')
+        
+        self._kin_dyn = casadi_kin_dyn.CasadiKinDyn(self.urdf) # used for getting joint names 
+        self._assign_controller_side_jnt_names(jnt_names=self._get_robot_jnt_names())
+        
+        self._init_robot_homer()
+
+        # handle fixed joints
+        fixed_joint_map={}
+        if self._fixed_jnt_patterns is not None:
+            for jnt_name in self._get_robot_jnt_names():
+                for fixed_jnt_pattern in self._fixed_jnt_patterns:
+                    if fixed_jnt_pattern in jnt_name: 
+                        fixed_joint_map.update({f"{jnt_name}":
+                            self._homer.get_homing_val(jnt_name=jnt_name)})
+                        break # do not search for other pattern matches
+        
+        if not len(fixed_joint_map)==0: # we need to recreate kin dyn and homers
+            Journal.log(self.__class__.__name__,
+                "_init_problem",
+                f"Will fix following joints: \n{str(fixed_joint_map)}",
+                LogType.INFO,
+                throw_when_excep=True)
+            # with the fixed joint map
+            self._kin_dyn = casadi_kin_dyn.CasadiKinDyn(self.urdf,fixed_joints=fixed_joint_map)
+            # assign again controlled joints names
+            self._assign_controller_side_jnt_names(jnt_names=self._get_robot_jnt_names())
+            # updated robot homer for controlled joints
+            self._init_robot_homer()
+
+        # handle continuous joints (need to change homing and retrieve
+        # cont jnts indexes) and homing
+        self._continuous_joints=self._get_continuous_jnt_names()
+        # reduced
+        self._continuous_joints_idxs=[]
+        self._continuous_joints_idxs_cos=[]
+        self._continuous_joints_idxs_sin=[]
+        self._continuous_joints_idxs_red=[]
+        self._rev_joints_idxs=[]
+        self._rev_joints_idxs_red=[]
+        # qfull
+        self._continuous_joints_idxs_qfull=[]
+        self._continuous_joints_idxs_cos_qfull=[]
+        self._continuous_joints_idxs_sin_qfull=[]
+        self._continuous_joints_idxs_red_qfull=[]
+        self._rev_joints_idxs_qfull=[]
+        self._rev_joints_idxs_red_qfull=[]
+        jnt_homing=[""]*(len(self._homer.get_homing())+len(self._continuous_joints))
+        jnt_names=self._get_robot_jnt_names()
+        for i in range(len(jnt_names)):
+            jnt=jnt_names[i]
+            index=self._get_jnt_id(jnt)# accounting for floating joint
+            homing_idx=index-7 # homing is only for actuated joints
+            homing_value=self._homer.get_homing_val(jnt)
+            if jnt in self._continuous_joints:
+                jnt_homing[homing_idx]=np.cos(homing_value).item()
+                jnt_homing[homing_idx+1]=np.sin(homing_value).item()
+                # just actuated joints
+                self._continuous_joints_idxs.append(homing_idx) # cos
+                self._continuous_joints_idxs.append(homing_idx+1) # sin
+                self._continuous_joints_idxs_cos.append(homing_idx)
+                self._continuous_joints_idxs_sin.append(homing_idx+1)
+                self._continuous_joints_idxs_red.append(i)
+                # q full
+                self._continuous_joints_idxs_qfull.append(index) # cos
+                self._continuous_joints_idxs_qfull.append(index+1) # sin
+                self._continuous_joints_idxs_cos_qfull.append(index)
+                self._continuous_joints_idxs_sin_qfull.append(index+1)
+                self._continuous_joints_idxs_red_qfull.append(i+7)
+            else:
+                jnt_homing[homing_idx]=homing_value
+                # just actuated joints
+                self._rev_joints_idxs.append(homing_idx) 
+                self._rev_joints_idxs_red.append(i) 
+                # q full
+                self._rev_joints_idxs_qfull.append(index) 
+                self._rev_joints_idxs_red_qfull.append(i+7) 
+
+        self._jnts_q_reduced=None
+        if not len(self._continuous_joints)==0: 
+            cont_joints=", ".join(self._continuous_joints)
+
+            Journal.log(self.__class__.__name__,
+                "_init_problem",
+                f"The following continuous joints were found: \n{cont_joints}",
+                LogType.INFO,
+                throw_when_excep=True)
+            # preallocating data 
+            self._jnts_q_reduced=np.zeros((1,self.nv()-6),dtype=self._dtype)
+            self._jnts_q_expanded=np.zeros((self.nq()-7,1),dtype=self._dtype)
+            self._full_q_reduced=np.zeros((7+len(jnt_names), self._n_nodes),dtype=self._dtype)
+            self._jnts_q_delta_expanded=np.zeros((self.nq()-7,1),dtype=self._dtype)
+        else:
+            self._custom_opts["replace_continuous_joints"]=True
+            Journal.log(self.__class__.__name__,
+                "_init_problem",
+                f"No continuous joints were found.",
+                LogType.INFO,
+                throw_when_excep=True)
+
+        # retrieve wheels indexes (not considering continuous joints, 
+        # ok just for v, eff, etc..)
+        self._wheel_patterns=wheels_patterns
+        self._wheels_idxs_v=self._get_wheels_jnt_v_idxs(wheel_patterns=self._wheel_patterns)
+        self._f0 = [0, 0, self._kin_dyn.mass()/4*9.81]
+        
+        # we can create an init for the base
+        init = self._base_init.tolist() + jnt_homing
+
+        if foot_linkname is not None:
+            FK = self._kin_dyn.fk(foot_linkname) # just to get robot reference height
+            ground_level = FK(q=init)['ee_pos']
+            self._base_init[2] = -ground_level[2]  # override init
+        
+        self._model = FullModelInverseDynamics(problem=self._prb,
+            kd=self._kin_dyn,
+            q_init=self._homer.get_homing_map(),
+            base_init=self._base_init)
+
+        self._ti = TaskInterface(prb=self._prb, 
+                            model=self._model, 
+                            max_solver_iter=self.max_solver_iter,
+                            debug = self._debug,
+                            verbose = self._verbose, 
+                            codegen_workdir = self._codegen_dir)
+        self._ti.setTaskFromYaml(self.config_path)
+        
+        # setting initial base pos ref if exists
+        base_pos = self._ti.getTask('base_height')
+        if base_pos is not None:
+            base_pos.setRef(np.atleast_2d(self._base_init).T)
+
+        self._pm = pymanager.PhaseManager(self._n_nodes, debug=False) # intervals or nodes?????
+
+        self._gm = GaitManager(self._ti, 
+            self._pm, 
+            self._injection_node,
+            keep_yaw_vert=keep_yaw_vert,
+            yaw_vertical_weight=self._yaw_vertical_weight,
+            vertical_landing=vertical_landing,
+            landing_vert_weight=self._vertical_land_weight,
+            phase_force_reg=self._phase_force_reg,
+            custom_opts=self._custom_opts,
+            flight_duration=flight_duration,
+            post_flight_stance=post_flight_stance,
+            step_height=step_height,
+            dh=0.0)
+                
+        self._ti.model.q.setBounds(self._ti.model.q0, self._ti.model.q0, nodes=0)
+        self._ti.model.v.setBounds(self._ti.model.v0, self._ti.model.v0, nodes=0)
+        self._ti.model.q.setInitialGuess(self._ti.model.q0)
+        self._ti.model.v.setInitialGuess(self._ti.model.v0)
+        for _, cforces in self._ti.model.cmap.items():
+            n_contact_f=len(cforces)
+            for c in cforces:
+                c.setInitialGuess(np.array(self._f0)/n_contact_f)        
+
+        vel_lims = self._model.kd.velocityLimits()
+        import horizon.utils as utils
+        self._prb.createResidual('vel_lb_barrier', self._vel_bounds_weight*utils.utils.barrier(vel_lims[7:] - self._model.v[7:]))
+        self._prb.createResidual('vel_ub_barrier', self._vel_bounds_weight*utils.utils.barrier1(-1 * vel_lims[7:] - self._model.v[7:]))
+
+        self._meas_lin_a_par=None
+        # if self._custom_opts["lin_a_feedback"]:
+        #     # acceleration feedback on first node
+        #     self._meas_lin_a_par=self._prb.createParameter(name="lin_a_feedback",
+        #         dim=3, nodes=0)   
+        #     base_lin_a_prb=self._prb.getInput().getVars()[0:3] 
+        #     self._prb.createConstraint('lin_acceleration_feedback', base_lin_a_prb - self._meas_lin_a_par, 
+        #             nodes=[0])
+
+        # if not self._open_loop:
+        #     # we create a residual cost to be used as an attractor to the measured state on the first node
+        #     # hard constraints injecting meas. states are pure EVIL!
+        #     prb_state=self._prb.getState()
+        #     full_state=prb_state.getVars()
+        #     state_dim=prb_state.getBounds()[0].shape[0]
+        #     meas_state=self._prb.createParameter(name="measured_state",
+        #         dim=state_dim, nodes=0)     
+        #     self._prb.createResidual('meas_state_attractor', meas_state_attractor_weight * (full_state - meas_state), 
+        #                 nodes=[0])
+
+        self._ti.finalize()
+        self._ti.bootstrap()
+
+        self._ti.init_inv_dyn_for_res() # we initialize some objects for sol. postprocessing purposes
+        self._ti.load_initial_guess()
+
+        self.n_dofs = self._get_ndofs() # after loading the URDF and creating the controller we
+        # know n_dofs -> we assign it (by default = None)
+
+        self.n_contacts = len(self._model.cmap.keys())
+        
+        # remove variables bounds (before they were necessary to have a good
+        # quality bootstrap solution)
+        self._q_inf=np.zeros((self.nq(), 1))
+        self._q_inf[:, :]=np.inf
+        self._v_inf=np.zeros((self.nv(), 1))
+        self._v_inf[:, :]=np.inf
+        self._ti.model.q.setBounds(-self._q_inf, self._q_inf, nodes=0)
+        self._ti.model.v.setBounds(-self._v_inf, self._v_inf, nodes=0)
+
+        # self.horizon_anal = analyzer.ProblemAnalyzer(self._prb)
+
+    def get_file_paths(self):
+        # can be overriden by child
+        paths = super().get_file_paths()
+        return paths
+    
+    def _get_quat_remap(self):
+        # overrides parent
+        return [1, 2, 3, 0] # mapping from robot quat. to Horizon's quaternion convention
+    
+    def _zmp(self, model):
+
+        num = cs.SX([0, 0])
+        den = cs.SX([0])
+        pos_contact = dict()
+        force_val = dict()
+
+        q = cs.SX.sym('q', model.nq)
+        v = cs.SX.sym('v', model.nv)
+        a = cs.SX.sym('a', model.nv)
+
+        com = model.kd.centerOfMass()(q=q, v=v, a=a)['com']
+
+        n = cs.SX([0, 0, 1])
+        for c in model.fmap.keys():
+            pos_contact[c] = model.kd.fk(c)(q=q)['ee_pos']
+            force_val[c] = cs.SX.sym('force_val', 3)
+            num += (pos_contact[c][0:2] - com[0:2]) * cs.dot(force_val[c], n)
+            den += cs.dot(force_val[c], n)
+
+        zmp = com[0:2] + (num / den)
+        input_list = []
+        input_list.append(q)
+        input_list.append(v)
+        input_list.append(a)
+
+        for elem in force_val.values():
+            input_list.append(elem)
+
+        f = cs.Function('zmp', input_list, [zmp])
+        return f
+    
+    def _add_zmp(self):
+
+        input_zmp = []
+
+        input_zmp.append(self._model.q)
+        input_zmp.append(self._model.v)
+        input_zmp.append(self._model.a)
+
+        for f_var in self._model.fmap.values():
+            input_zmp.append(f_var)
+
+        c_mean = cs.SX([0, 0, 0])
+        for c_name, f_var in self._model.fmap.items():
+            fk_c_pos = self._kin_dyn.fk(c_name)(q=self._model.q)['ee_pos']
+            c_mean += fk_c_pos
+
+        c_mean /= len(self._model.cmap.keys())
+
+        zmp_nominal_weight = 10.
+        zmp_fun = self._zmp(self._model)(*input_zmp)
+
+        if 'wheel_joint_1' in self._model.kd.joint_names():
+            zmp_residual = self._prb.createIntermediateResidual('zmp',  zmp_nominal_weight * (zmp_fun[0:2] - c_mean[0:2]))
+
+    def _quaternion_multiply(self, 
+                    q1, q2):
+        x1, y1, z1, w1 = q1
+        x2, y2, z2, w2 = q2
+        w = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
+        x = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
+        y = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
+        z = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
+        return np.array([x, y, z, w])
+    
+    def _get_continuous_jnt_names(self):
+        import xml.etree.ElementTree as ET
+        root = ET.fromstring(self.urdf)
+        continuous_joints = []
+        for joint in root.findall('joint'):
+            joint_type = joint.get('type')
+            if joint_type == 'continuous':
+                joint_name = joint.get('name')
+                continuous_joints.append(joint_name)
+        return continuous_joints
+    
+    def _get_wheels_jnt_v_idxs(self, wheel_patterns: List[str]):
+        jnt_names=self._get_robot_jnt_names()
+        wheels_idxs=[]
+        for i in range(len(jnt_names)):
+            jnt_name=jnt_names[i]
+            for wheel_pattern in wheel_patterns:
+                if wheel_pattern in jnt_name:
+                    wheels_idxs.append(i)
+                    break
+        return wheels_idxs
+    
+    def _get_jnt_id(self, jnt_name):
+        return self._kin_dyn.joint_iq(jnt_name)
+    
+    def _init_rhc_task_cmds(self):
+        
+        rhc_refs = HybridQuadRhcRefs(gait_manager=self._gm,
+            robot_index_shm=self.controller_index,
+            robot_index_view=0, # when using optimize_mem the view if always of shape 1x...
+            namespace=self.namespace,
+            safe=False, 
+            verbose=self._verbose,
+            vlevel=VLevel.V2,
+            use_force_feedback=self._custom_opts["use_force_feedback"],
+            optimize_mem=True)
+        
+        rhc_refs.run()
+
+        rhc_refs.rob_refs.set_jnts_remapping(jnts_remapping=self._to_controller)
+        rhc_refs.rob_refs.set_q_remapping(q_remapping=self._get_quat_remap())
+        
+        rhc_refs.set_default_refs(p_ref=np.atleast_2d(self._base_init)[:, 0:3],
+            q_ref=np.atleast_2d(self._base_init)[:, 3:7])
+        
+        return rhc_refs
+    
+    def get_vertex_fnames_from_ti(self):
+        tasks=self._ti.task_list
+        contact_f_names=[]
+        for task in tasks:
+            if isinstance(task, ContactTask):
+                interaction_task=task.dynamics_tasks[0]
+                contact_f_names.append(interaction_task.vertex_frames[0])
+        return contact_f_names
+        
+    def _get_contact_names(self):
+        # should get contact names from vertex frames
+        # list(self._ti.model.cmap.keys())
+        return self.get_vertex_fnames_from_ti()
+    
+    def _get_robot_jnt_names(self):
+
+        joints_names = self._kin_dyn.joint_names()
+        to_be_removed = ["universe", 
+                        "reference", 
+                        "world", 
+                        "floating", 
+                        "floating_base"]
+        for name in to_be_removed:
+            if name in joints_names:
+                joints_names.remove(name)
+
+        return joints_names
+    
+    def _get_ndofs(self):
+        return len(self._model.joint_names)
+    
+    def nq(self):
+        return self._kin_dyn.nq()
+    
+    def nv(self):
+        return self._kin_dyn.nv()
+    
+    def _get_robot_mass(self):
+
+        return self._kin_dyn.mass()
+
+    def _get_root_full_q_from_sol(self, node_idx=1):
+
+        root_q_full=self._ti.solution['q'][0:7, node_idx].reshape(1, 7).astype(self._dtype)
+
+        np.nan_to_num(root_q_full, nan=1e3, posinf=1e3, neginf=-1e3, copy=False)
+        np.clip(a=root_q_full, a_min=-1e3, a_max=1e3, out=root_q_full)
+
+        return root_q_full
+    
+    def _get_full_q_from_sol(self, node_idx=1):
+
+        return self._ti.solution['q'][:, node_idx].reshape(1, -1).astype(self._dtype)
+    
+    def _get_root_twist_from_sol(self, node_idx=1):
+        # provided in world frame
+        twist_base_local=self._get_v_from_sol()[0:6, node_idx].reshape(1, 6)
+        # if world_aligned:
+        #     q_root_rhc = self._get_root_full_q_from_sol(node_idx=node_idx)[:, 0:4]
+        #     r_base_rhc=Rotation.from_quat(q_root_rhc.flatten()).as_matrix()
+        #     twist_base_local[:, 0:3] = r_base_rhc @ twist_base_local[0, 0:3]
+        #     twist_base_local[:, 3:6] = r_base_rhc @ twist_base_local[0, 3:6]
+        return twist_base_local
+
+    def _get_root_a_from_sol(self, node_idx=0):
+        # provided in world frame
+        a_base_local=self._get_a_from_sol()[0:6, node_idx].reshape(1, 6)
+        # if world_aligned:
+        #     q_root_rhc = self._get_root_full_q_from_sol(node_idx=node_idx)[:, 0:4]
+        #     r_base_rhc=Rotation.from_quat(q_root_rhc.flatten()).as_matrix()
+        #     a_base_local[:, 0:3] = r_base_rhc @ a_base_local[0, 0:3]
+        #     a_base_local[:, 3:6] = r_base_rhc @ v[0, 3:6]
+        return a_base_local
+    
+    def _get_jnt_q_from_sol(self, node_idx=0, 
+            reduce: bool = True,
+            clamp: bool = True):
+        
+        full_jnts_q=self._ti.solution['q'][7:, node_idx:node_idx+1].reshape(1,-1).astype(self._dtype)
+        
+        np.nan_to_num(full_jnts_q, nan=1e3, posinf=1e3, neginf=-1e3, copy=False) # in place
+        np.clip(a=full_jnts_q, a_max=1e3, a_min=-1e3, out=full_jnts_q) # in place
+
+        if self._custom_opts["replace_continuous_joints"] or (not reduce):
+            if clamp:
+                return np.fmod(full_jnts_q, 2*np.pi)
+            else:
+                return full_jnts_q
+        else:
+            cos_sin=full_jnts_q[:,self._continuous_joints_idxs].reshape(-1,2)
+            # copy rev joint vals
+            self._jnts_q_reduced[:, self._rev_joints_idxs_red]=np.fmod(full_jnts_q[:, self._rev_joints_idxs], 2*np.pi).reshape(1, -1)
+            # and continuous
+            self._jnts_q_reduced[:, self._continuous_joints_idxs_red]=np.arctan2(cos_sin[:, 1], cos_sin[:, 0]).reshape(1,-1)
+            return self._jnts_q_reduced
+        
+    def _get_jnt_v_from_sol(self, node_idx=1):
+        
+        jnt_v_sol=self._get_v_from_sol()[6:, node_idx].reshape(1,  
+                    self._nv_jnts)
+        np.nan_to_num(jnt_v_sol, nan=1e5, posinf=1e5, neginf=-1e5, copy=False) # in place
+        # np.clip(a=jnt_v_sol, a_max=1e5, a_min=-1e5, out=jnt_v_sol) # in place
+
+        return jnt_v_sol
+
+    def _get_jnt_a_from_sol(self, node_idx=1):
+
+        return self._get_a_from_sol()[6:, node_idx].reshape(1,
+                    self._nv_jnts)
+
+    def _get_jnt_eff_from_sol(self, node_idx=0):
+        
+        efforts_on_node = self._ti.eval_efforts_on_node(node_idx=node_idx)
+        
+        # if self._custom_opts["only_vel_wheels"]:
+
+        jnt_efforts=efforts_on_node[6:, 0]
+        
+        if self._custom_opts["only_vel_wheels"] and self._wheels_idxs_v:
+            jnt_efforts[self._wheels_idxs_v]=0.0
+
+        return jnt_efforts.reshape(1,  
+                self._nv_jnts).astype(self._dtype)
+    
+    def _get_rhc_cost(self):
+
+        return self._ti.solution["opt_cost"]
+    
+    def _get_rhc_constr_viol(self):
+        
+        return self._ti.solution["residual_norm"]
+    
+    def _get_rhc_nodes_cost(self):
+
+        cost = self._ti.solver_rti.getCostValOnNodes()
+        return cost.reshape((1, -1))
+    
+    def _get_rhc_nodes_constr_viol(self):
+        
+        constr_viol = self._ti.solver_rti.getConstrValOnNodes()
+        return constr_viol.reshape((1, -1))
+    
+    def _get_rhc_niter_to_sol(self):
+
+        return self._ti.solution["n_iter2sol"]
+    
+    def _set_ig_bootstrap(self,
+            q_state: np.ndarray = None,
+            v_state: np.ndarray = None):
+
+        xig = self._ti.solution['x_opt'].copy()
+        uig = self._ti.solution['u_opt'].copy()
+
+        # Normalize and keep quaternion in the same hemisphere as the previous
+        # solution to avoid artificial 180-deg jumps in the bootstrap warm start.
+        q_state_boot = q_state.copy()
+        q_prev = xig[3:7, 0]
+        q_now = q_state_boot[3:7, 0]
+
+        q_now_norm = np.linalg.norm(q_now)
+        if q_now_norm > 1e-9:
+            q_state_boot[3:7, :] /= q_now_norm
+        else:
+            q_state_boot[3:7, :] = np.array([[0.0], [0.0], [0.0], [1.0]], dtype=self._dtype)
+
+        q_prev_norm = np.linalg.norm(q_prev)
+        if q_prev_norm > 1e-9:
+            q_prev = q_prev / q_prev_norm
+
+        q_now = q_state_boot[3:7, 0]
+        if np.dot(q_prev, q_now) < 0.0:
+            q_state_boot[3:7, :] *= -1.0
+        
+        xig[0:self._nq, :] = q_state_boot
+        xig[self._nq:self._nq + self._nv, :] = 0.0 # 0 velocity on first nodes 
+        uig[0:self._nv, :]=0.0 # 0 acceleration
+
+        # assigning ig
+        self._prb.getState().setInitialGuess(xig)
+        self._prb.getInput().setInitialGuess(uig)
+        # self._prb.getVariables("a").setInitialGuess(np.zeros((self._nv, 1), dtype=self._dtype))
+        for _, cforces in self._ti.model.cmap.items():
+            n_contact_f = len(cforces)
+            if n_contact_f == 0:
+                continue
+            f_guess = np.array(self._f0, dtype=self._dtype) / n_contact_f
+            for c in cforces:
+                c.setInitialGuess(f_guess)
+
+        # print("initial guesses")
+        # print(self._nq)
+        # print(self._nv)
+        # print("q")
+        # qig=self._ti.model.q.getInitialGuess()
+        # print(qig.shape)
+        # print(qig)
+        # print("v")
+        # print(self._ti.model.v.getInitialGuess())
+        # print("a")
+        # print(self._ti.model.a.getInitialGuess())
+        # for _, cforces in self._ti.model.cmap.items():
+        #     for c in cforces:
+        #         print("force")
+        #         print(c.getInitialGuess())
+        
+        return xig, uig
+    
+    def _set_ig(self):
+
+        shift_num = -1 # shift data by one node
+
+        x_opt = self._ti.solution['x_opt']
+        u_opt = self._ti.solution['u_opt']
+
+        # building ig for state
+        xig = np.roll(x_opt, 
+                shift_num, axis=1) # rolling state sol.
+        for i in range(abs(shift_num)):
+            # state on last node is copied to the elements
+            # which are "lost" during the shift operation
+            xig[:, -1 - i] = x_opt[:, -1]
+        # building ig for inputs
+        uig = np.roll(u_opt, 
+                shift_num, axis=1) # rolling state sol.
+        for i in range(abs(shift_num)):
+            # state on last node is copied to the elements
+            # which are "lost" during the shift operation
+            uig[:, -1 - i] = u_opt[:, -1]
+
+        # assigning ig
+        self._prb.getState().setInitialGuess(xig)
+        self._prb.getInput().setInitialGuess(uig)
+        
+        return xig, uig
+    
+    def _update_open_loop(self,
+            bootstrap: bool = False):
+
+        q_state, v_state, a_state=self._set_is_open()
+
+        if not bootstrap:
+            self._set_ig()
+        else:
+            self._set_ig_bootstrap(q_state=q_state, v_state=v_state)
+
+        # robot_state=xig[:, 0]
+        # # open loop update:
+        # self._prb.setInitialState(x0=robot_state) # (xig has been shifted, so node 0
+        # # is node 1 in the last opt solution)
+
+        return q_state, v_state, a_state
+    
+    def _update_closed_loop(self,
+            bootstrap: bool = False):
+        
+        # set initial state
+        q_state=None
+        v_state=None
+        a_state=None
+        if self._custom_opts["adaptive_is"]:
+            # adaptive closed loop
+            q_state, v_state, a_state=self._set_is_adaptive()
+        elif self._custom_opts["fully_closed"]:
+            q_state, v_state, a_state=self._set_is_full()
+        elif self._custom_opts["closed_partial"]:
+            q_state, v_state, a_state=self._set_is_partial()
+        else:
+            Journal.log(self.__class__.__name__,
+                    "_update_closed_loop",
+                    "Neither adaptive_is, fully_closed, or closed_partial.",
+                    LogType.EXCEP,
+                    throw_when_excep = False)
+            q_state, v_state, a_state=self._set_is()
+
+        # set initial guess for controller
+        if not bootstrap:
+            self._set_ig()
+        else:
+            self._set_ig_bootstrap(q_state=q_state, v_state=v_state)
+
+        return q_state, v_state, a_state
+    
+    def _set_is_open(self):
+        
+        # overriding states with rhc data
+        q_full_root=self._get_root_full_q_from_sol(node_idx=1).reshape(-1, 1)
+        q_jnts=self._get_jnt_q_from_sol(node_idx=1, reduce=False).reshape(-1, 1)
+        
+        twist_root=self._get_root_twist_from_sol(node_idx=1).reshape(-1, 1)
+        v_jnts=self._get_jnt_v_from_sol(node_idx=1).reshape(-1, 1)
+
+        # rhc variables to be set
+        q=self._prb.getVariables("q") # .setBounds()
+        root_q_full_rhc=q[0:7] # root full q
+        jnts_q_rhc=q[7:] # jnts q
+        vel=self._prb.getVariables("v")
+        root_twist_rhc=vel[0:6] # lin v.
+        jnts_v_rhc=vel[6:] # jnts v
+
+        self.rhc_refs.set_alpha(alpha=1.0) # fully open
+
+        # close state on known quantities
+        root_q_full_rhc.setBounds(lb=q_full_root,
+            ub=q_full_root, nodes=0)
+        jnts_q_rhc.setBounds(lb=q_jnts, 
+            ub=q_jnts, nodes=0)
+        root_twist_rhc.setBounds(lb=twist_root, 
+            ub=twist_root, nodes=0)
+        jnts_v_rhc.setBounds(lb=v_jnts, 
+            ub=v_jnts, nodes=0)
+        
+        # return state used for feedback
+        q_state=np.concatenate((q_full_root, q_jnts),
+                axis=0)
+        v_state=np.concatenate((twist_root, v_jnts),
+                axis=0)
+        
+        return (q_state, v_state, None)
+    
+    def _set_is_full(self):
+        
+        # measurements
+        q_full_root = self.robot_state.root_state.get(data_type="q_full", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        v_root = self.robot_state.root_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        omega = self.robot_state.root_state.get(data_type="omega", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        a_root = self.robot_state.root_state.get(data_type="a_full", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        
+        q_jnts = self.robot_state.jnts_state.get(data_type="q", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        v_jnts = self.robot_state.jnts_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        a_jnts = self.robot_state.jnts_state.get(data_type="a", robot_idxs=self.controller_index_np).reshape(-1, 1)
+
+        if (not len(self._continuous_joints)==0): # we need do expand some meas. rev jnts to So2
+            self._jnts_q_expanded[self._rev_joints_idxs, :]=q_jnts[self._rev_joints_idxs_red ,:]
+            self._jnts_q_expanded[self._continuous_joints_idxs_cos, :]=np.cos(q_jnts[self._continuous_joints_idxs_red, :]) # cos
+            self._jnts_q_expanded[self._continuous_joints_idxs_sin, :]=np.sin(q_jnts[self._continuous_joints_idxs_red, :]) # sin
+            q_jnts=self._jnts_q_expanded.reshape(-1,1)
+            
+        # rhc variables to be set
+        q=self._prb.getVariables("q") # .setBounds()
+        root_full_q_rhc=q[0:7] # root p
+        jnts_q_rhc=q[7:] # jnts q
+        vel=self._prb.getVariables("v")
+        root_v_rhc=vel[0:3] # lin v.
+        root_omega_rhc=vel[3:6] # omega
+        jnts_v_rhc=vel[6:] # jnts v
+        acc=self._prb.getVariables("a")
+        lin_a_prb=acc[0:3] # lin acc
+        
+        self.rhc_refs.set_alpha(alpha=0.0) # non-adaptive
+
+        root_full_q_rhc.setBounds(lb=q_full_root,
+            ub=q_full_root, nodes=0)
+        jnts_q_rhc.setBounds(lb=q_jnts, 
+            ub=q_jnts, nodes=0)
+        root_v_rhc.setBounds(lb=v_root, 
+            ub=v_root, nodes=0) # leaving lin v of the base free (estimated from constraints)
+        root_omega_rhc.setBounds(lb=omega, 
+            ub=omega, nodes=0)
+        jnts_v_rhc.setBounds(lb=v_jnts, 
+            ub=v_jnts, nodes=0)
+        if self._custom_opts["lin_a_feedback"]:
+            # write base lin 13793197 from meas
+            lin_a_prb.setBounds(lb=a_root[0:3, :], 
+                ub=a_root[0:3, :], 
+                nodes=0)
+
+        # return state used for feedback
+        q_state=np.concatenate((q_full_root, q_jnts),
+                axis=0)
+        v_state=np.concatenate((v_root, omega, v_jnts),
+                axis=0)
+        a_state=np.concatenate((a_root, a_jnts),
+                axis=0)
+        
+        return (q_state, v_state, a_state)
+    
+    def _set_is_partial(self):
+        
+        # measurements
+        p_root = self.robot_state.root_state.get(data_type="p", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        q_root = self.robot_state.root_state.get(data_type="q", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        v_root = self.robot_state.root_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        omega = self.robot_state.root_state.get(data_type="omega", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        a_root = self.robot_state.root_state.get(data_type="a_full", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        
+        q_jnts = self.robot_state.jnts_state.get(data_type="q", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        v_jnts = self.robot_state.jnts_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        a_jnts = self.robot_state.jnts_state.get(data_type="a", robot_idxs=self.controller_index_np).reshape(-1, 1)
+
+        if (not len(self._continuous_joints)==0): # we need do expand some meas. rev jnts to So2
+            self._jnts_q_expanded[self._rev_joints_idxs, :]=q_jnts[self._rev_joints_idxs_red ,:]
+            self._jnts_q_expanded[self._continuous_joints_idxs_cos, :]=np.cos(q_jnts[self._continuous_joints_idxs_red, :]) # cos
+            self._jnts_q_expanded[self._continuous_joints_idxs_sin, :]=np.sin(q_jnts[self._continuous_joints_idxs_red, :]) # sin
+            q_jnts=self._jnts_q_expanded.reshape(-1,1)
+
+        # overriding states with rhc data (-> all overridden state are open looop)
+        root_q_full_from_rhc=self._get_root_full_q_from_sol(node_idx=1).reshape(-1, 1)
+        root_p_from_rhc=root_q_full_from_rhc[0:3, :]
+        p_root[:, :]=root_p_from_rhc # position is always open loop
+        
+        if not self._custom_opts["estimate_v_root"]:
+            v_root[:, :]=self._get_root_twist_from_sol(node_idx=1)[:, 0:3].reshape(-1, 1)
+            # override v jnts with the ones from controller
+            if not self._custom_opts["use_jnt_v_feedback"]:
+                v_jnts[:, :]=self._get_jnt_v_from_sol(node_idx=1).reshape(-1, 1)
+            # v_jnts[:, :]=self._get_jnt_v_from_sol(node_idx=1).reshape(-1, 1)
+        # root_twist_from_rhc=self._get_root_twist_from_sol(node_idx=1)
+        # root_v_from_rhc=root_twist_from_rhc[:, 0:3].reshape(-1, 1)
+        # root_omega_from_rhc=root_twist_from_rhc[:, 3:6].reshape(-1, 1)
+        # jnt_q_from_rhc=self._get_jnt_q_from_sol(node_idx=1,reduce=False,clamp=False).reshape(-1, 1)
+        # jnt_v_from_rhc=self._get_jnt_v_from_sol(node_idx=1).reshape(-1, 1)
+            
+        # rhc variables to be set
+        q=self._prb.getVariables("q") # .setBounds()
+        root_p_rhc=q[0:3] # root p
+        root_q_rhc=q[3:7] # root orientation
+        jnts_q_rhc=q[7:] # jnts q
+        vel=self._prb.getVariables("v")
+        root_v_rhc=vel[0:3] # lin v.
+        root_omega_rhc=vel[3:6] # omega
+        jnts_v_rhc=vel[6:] # jnts v
+        acc=self._prb.getVariables("a")
+        lin_a_prb=acc[0:3] # lin acc
+        
+        self.rhc_refs.set_alpha(alpha=0.0) # non-adaptive
+
+        root_p_rhc.setBounds(lb=p_root,
+            ub=p_root, nodes=0)
+        root_q_rhc.setBounds(lb=q_root, 
+            ub=q_root, nodes=0)
+        jnts_q_rhc.setBounds(lb=q_jnts, 
+            ub=q_jnts, nodes=0)
+        
+        if self._custom_opts["estimate_v_root"]:
+            root_v_rhc.setBounds(lb=-self._v_inf[0:3], 
+                ub=self._v_inf[0:3], nodes=0) # leaving lin v of the base free (estimated from constraints)
+        else: # get it from controller 
+            root_v_rhc.setBounds(lb=v_root, 
+                ub=v_root, nodes=0)
+        root_omega_rhc.setBounds(lb=omega, 
+            ub=omega, nodes=0)
+        jnts_v_rhc.setBounds(lb=v_jnts, 
+            ub=v_jnts, nodes=0)
+        if self._custom_opts["lin_a_feedback"]:
+            # write base lin 13793197 from meas
+            lin_a_prb.setBounds(lb=a_root[0:3, :], 
+                ub=a_root[0:3, :], 
+                nodes=0)
+
+        # return state used for feedback
+        q_state=np.concatenate((p_root, q_root, q_jnts),
+                axis=0)
+        v_state=np.concatenate((v_root, omega, v_jnts),
+                axis=0)
+        a_state=np.concatenate((a_root, a_jnts),
+                axis=0)
+        
+        return (q_state, v_state, a_state)
+    
+    def _set_is_adaptive(self):
+        
+        # measurements
+        p_root = self.robot_state.root_state.get(data_type="p", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        q_root = self.robot_state.root_state.get(data_type="q", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        v_root = self.robot_state.root_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        omega = self.robot_state.root_state.get(data_type="omega", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        a_root = self.robot_state.root_state.get(data_type="a_full", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        
+        q_jnts = self.robot_state.jnts_state.get(data_type="q", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        v_jnts = self.robot_state.jnts_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        a_jnts = self.robot_state.jnts_state.get(data_type="a", robot_idxs=self.controller_index_np).reshape(-1, 1)
+            
+        # rhc variables to be set
+        q=self._prb.getVariables("q") # .setBounds()
+        root_p_rhc=q[0:3] # root p
+        root_q_rhc=q[3:7] # root orientation
+        jnts_q_rhc=q[7:] # jnts q
+        vel=self._prb.getVariables("v")
+        root_v_rhc=vel[0:3] # lin v.
+        root_omega_rhc=vel[3:6] # omega
+        jnts_v_rhc=vel[6:] # jnts v
+        acc=self._prb.getVariables("a")
+        lin_a_prb=acc[0:3] # lin acc
+        
+        # getting prediction defects
+        root_q_delta=self.rhc_pred_delta.root_state.get(data_type="q", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        jnt_q_delta=self.rhc_pred_delta.jnts_state.get(data_type="q", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        jnt_v_delta=self.rhc_pred_delta.jnts_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        v_root_delta = self.rhc_pred_delta.root_state.get(data_type="v", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        omega_root_delta = self.rhc_pred_delta.root_state.get(data_type="omega", robot_idxs=self.controller_index_np).reshape(-1, 1)
+        a_root_delta = self.rhc_pred_delta.root_state.get(data_type="a_full", robot_idxs=self.controller_index_np).reshape(-1, 1)
+
+        # close state on known quantities, estimate some (e.g. lin vel) and
+        # open loop if thing start to explode
+        alpha_now=1.0
+        delta=0.0 
+        if self._custom_opts["alpha_from_outside"]:
+            alpha_now=self.rhc_refs.get_alpha()
+        else: # "autotuned" alpha
+            if self._custom_opts["estimate_v_root"]: # we copmute delta based on jnt v (since we use meas.)
+                delta=np.max(np.abs(jnt_v_delta))
+            else:
+                delta=np.max(np.abs(omega_root_delta))
+            # fail_idx=self._get_failure_index()
+            # fail_idx=self._get_explosion_idx()/self._fail_idx_thresh
+            alpha_now=(np.tanh(2*self._alpha_half*(delta-self._alpha_half))+1)/2.0
+
+        bound_relaxation=-np.log(1-alpha_now+1e-16) # [~0, ~inf] if alpha_now [0, 1]
+        self.rhc_refs.set_alpha(alpha=alpha_now) # also writes on shared mem for db
+        self.rhc_refs.set_bound_relax(bound_relax=bound_relaxation) # also writes on shared mem for db
+
+        self._alphas_q_root[:]=alpha_now # for now single alpha for everything
+        self._alphas_q_jnts[:]=alpha_now
+        self._alphas_twist_root[:]=alpha_now
+        self._alphas_v_jnts[:]=alpha_now
+        self._alphas_a[:]=alpha_now
+        if not self._custom_opts["estimate_v_root"]:
+            self._alphas_twist_root[0:3]=1.0 # open
+            self._alphas_v_jnts[:]=1.0 # open
+
+        # position is always open loop
+        root_q_full_from_rhc=self._get_root_full_q_from_sol(node_idx=1).reshape(-1, 1)
+        root_p_from_rhc=root_q_full_from_rhc[0:3, :]
+        p_root[:, :]=root_p_from_rhc 
+    
+        # expaning meas q if continuous joints
+        if (not len(self._continuous_joints)==0): # we need do expand some meas. rev jnts to So2
+            self._jnts_q_expanded[self._rev_joints_idxs, :]=q_jnts[self._rev_joints_idxs_red ,:]
+            self._jnts_q_expanded[self._continuous_joints_idxs_cos, :]=np.cos(q_jnts[self._continuous_joints_idxs_red, :]) # cos
+            self._jnts_q_expanded[self._continuous_joints_idxs_sin, :]=np.sin(q_jnts[self._continuous_joints_idxs_red, :]) # sin
+            
+            # continous joints position is always open loop, but we need a delta vector of matching dimension
+            q_jnts_from_rhc=self._get_jnt_q_from_sol(node_idx=1).reshape(-1, 1)
+
+            self._jnts_q_delta_expanded[self._rev_joints_idxs, :]=jnt_q_delta[self._rev_joints_idxs_red ,:]
+
+            self._jnts_q_delta_expanded[self._continuous_joints_idxs_cos, :]=\
+                np.cos(q_jnts_from_rhc[self._continuous_joints_idxs_red, :]) - \
+                    np.cos(q_jnts[self._continuous_joints_idxs_red, :])
+            self._jnts_q_delta_expanded[self._continuous_joints_idxs_sin, :]=\
+                np.sin(q_jnts_from_rhc[self._continuous_joints_idxs_red, :]) - \
+                    np.sin(q_jnts[self._continuous_joints_idxs_red, :])
+
+            q_jnts=self._jnts_q_expanded.reshape(-1,1) # overriting with expanded jnts
+            jnt_q_delta=self._jnts_q_delta_expanded.reshape(-1, 1) # overriting with expanded jnts
+
+            self._alphas_q_jnts[self._continuous_joints_idxs_cos, :]=1.0 # open loop
+            self._alphas_q_jnts[self._continuous_joints_idxs_sin, :]=1.0 # open loop
+
+            # self._alphas_v_jnts[self._continuous_joints_idxs_red, :]=0.0 # open loop
+
+        root_p_rhc.setBounds(lb=p_root,
+            ub=p_root, nodes=0)
+        root_q_rhc.setBounds(lb=q_root+self._alphas_q_root[3:7]*root_q_delta, 
+            ub=q_root+self._alphas_q_root[3:7]*root_q_delta, nodes=0)
+        jnts_q_rhc.setBounds(lb=q_jnts+self._alphas_q_jnts*jnt_q_delta, 
+            ub=q_jnts+self._alphas_q_jnts*jnt_q_delta, nodes=0)
+        if self._custom_opts["estimate_v_root"]:
+            root_v_rhc.setBounds(lb=-self._v_inf[0:3], 
+                ub=self._v_inf[0:3], nodes=0)
+        else:
+            root_v_rhc.setBounds(lb=v_root+self._alphas_twist_root[0:3, :]*v_root_delta, 
+                ub=v_root+self._alphas_twist_root[0:3, :]*v_root_delta, nodes=0)
+        root_omega_rhc.setBounds(lb=omega+self._alphas_twist_root[3:6, :]*omega_root_delta, 
+            ub=omega+self._alphas_twist_root[3:6, :]*omega_root_delta, nodes=0)
+        jnts_v_rhc.setBounds(lb=v_jnts+self._alphas_v_jnts*jnt_v_delta, 
+            ub=v_jnts+self._alphas_v_jnts*jnt_v_delta, nodes=0)
+        if self._custom_opts["lin_a_feedback"]:
+            lin_a_prb.setBounds(lb=a_root[0:3, :]+self._alphas_a[0:3]*a_root_delta[0:3, :], 
+                ub=a_root[0:3, :]+self._alphas_a[0:3]*a_root_delta[0:3, :], 
+                nodes=0)
+            
+        # return state used for feedback
+        q_state=np.concatenate((p_root, q_root, q_jnts),
+                axis=0)
+        v_state=np.concatenate((v_root, omega, v_jnts),
+                axis=0)
+        a_state=np.concatenate((a_root, a_jnts),
+                axis=0)
+        
+        return (q_state, v_state, a_state)
+
+    def _solve(self):
+        
+        if self._debug:
+            return self._db_solve(bootstrap=False)
+        else:
+            return self._min_solve(bootstrap=False)
+
+    def _bootstrap(self):
+
+        if self._debug:
+            return self._db_solve(bootstrap=True)
+        else:
+            return self._min_solve(bootstrap=True)
+        
+    def _min_solve(self, bootstrap: bool = False):
+        # minimal solve version -> no debug 
+        robot_qstate=None
+        robot_vstate=None
+        robot_astate=None
+        if self._open_loop:
+            robot_qstate, robot_vstate, robot_astate = self._update_open_loop(bootstrap=bootstrap) # updates the TO ig and 
+            # initial conditions using data from the solution itself
+        else: 
+            robot_qstate, robot_vstate, robot_astate = self._update_closed_loop(bootstrap=bootstrap) # updates the TO ig and 
+            # initial conditions using robot measurements
+    
+        self._pm.shift() # shifts phases of one dt
+        if self._refs_in_hor_frame:
+            # q_base=self.robot_state.root_state.get(data_type="q", 
+            #     robot_idxs=self.controller_index).reshape(-1, 1)
+            # q_full=self._get_full_q_from_sol(node_idx=1).reshape(-1, 1)
+            # using internal base pose from rhc. in case of closed loop, it will be the meas state
+            force_norm=None
+            if self._custom_opts["use_force_feedback"]:
+                contact_forces=self.robot_state.contact_wrenches.get(data_type="f", 
+                    robot_idxs=self.controller_index_np,
+                    contact_name=None).reshape(self.n_contacts,3)
+                force_norm=np.linalg.norm(contact_forces, axis=1)
+            self.rhc_refs.step(qstate=robot_qstate, vstate=robot_vstate,
+                force_norm=force_norm)
+        else:
+            self.rhc_refs.step()
+            
+        try:
+            if not bootstrap:
+                converged = self._ti.rti() # RTI step
+            else:
+                converged = self._ti.bootstrap() # full solve (to convergence)
+            self.sol_counter = self.sol_counter + 1
+            return not self._check_rhc_failure()
+        except Exception as e: # fail in case of exceptions
+            return False
+    
+    def _db_solve(self, bootstrap: bool = False):
+
+        self._timer_start = time.perf_counter()
+
+        robot_qstate=None
+        robot_vstate=None
+        robot_astate=None
+        if self._open_loop:
+            robot_qstate, robot_vstate, robot_astate = self._update_open_loop(bootstrap=bootstrap) # updates the TO ig and 
+            # initial conditions using data from the solution itself
+        else: 
+            robot_qstate, robot_vstate, robot_astate = self._update_closed_loop(bootstrap=bootstrap) # updates the TO ig and 
+            # initial conditions using robot measurements
+
+        self._prb_update_time = time.perf_counter() 
+        self._pm.shift() # shifts phases of one dt
+        self._phase_shift_time = time.perf_counter()
+
+        if self._refs_in_hor_frame:
+            # q_base=self.robot_state.root_state.get(data_type="q", 
+            #     robot_idxs=self.controller_index).reshape(-1, 1)
+            # q_full=self._get_full_q_from_sol(node_idx=1).reshape(-1, 1)
+            # using internal base pose from rhc. in case of closed loop, it will be the meas state
+            force_norm=None
+            if self._custom_opts["use_force_feedback"]:
+                contact_forces=self.robot_state.contact_wrenches.get(data_type="f", 
+                    robot_idxs=self.controller_index_np,
+                    contact_name=None).reshape(self.n_contacts,3)
+                force_norm=np.linalg.norm(contact_forces, axis=1)
+            self.rhc_refs.step(qstate=robot_qstate, vstate=robot_vstate,
+                force_norm=force_norm)
+        else:
+            self.rhc_refs.step()
+             
+        self._task_ref_update_time = time.perf_counter() 
+    
+        try:
+            if not bootstrap:
+                converged = self._ti.rti() # RTI step
+            else:
+                converged = self._ti.bootstrap() # full solve bootstrap
+            self._rti_time = time.perf_counter() 
+            self.sol_counter = self.sol_counter + 1
+            self._update_db_data()
+            return not self._check_rhc_failure()
+        except Exception as e: # fail in case of exceptions
+            if self._verbose:
+                solve_mode = "RTI" if not bootstrap else "Bootstrap"
+                exception = f"{solve_mode}() for controller {self.controller_index} failed" + \
+                f" with exception {type(e).__name__}"
+                Journal.log(self.__class__.__name__,
+                    "solve",
+                    exception,
+                    LogType.EXCEP,
+                    throw_when_excep = False)
+            self._update_db_data()
+            return False
+    
+    def _get_fail_idx(self):
+        
+        self._explosion_idx_buffer[:, self._expl_idx_buffer_counter]=self._get_explosion_idx()
+        self._expl_idx_buffer_counter+=1
+        self._expl_idx_counter+=1
+        if self._expl_idx_counter%self._expl_idx_window_size==0:
+            self._expl_idx_buffer_counter=0 # restart from 0
+        
+        running_avrg=np.mean(self._explosion_idx_buffer).item()
+        
+        return running_avrg
+    
+    def _get_explosion_idx(self):
+        explosion_index = self._get_rhc_constr_viol() + self._get_rhc_cost()*self._fail_idx_scale
+        return explosion_index
+
+    def _update_db_data(self):
+
+        self._profiling_data_dict["problem_update_dt"] = self._prb_update_time - self._timer_start
+        self._profiling_data_dict["phases_shift_dt"] = self._phase_shift_time - self._prb_update_time
+        self._profiling_data_dict["task_ref_update"] = self._task_ref_update_time - self._phase_shift_time
+        self._profiling_data_dict["rti_solve_dt"] = self._rti_time - self._task_ref_update_time
+        self.rhc_costs.update(self._ti.solver_rti.getCostsValues())
+        self.rhc_constr.update(self._ti.solver_rti.getConstraintsValues())
+
+    def _reset(self):
+        
+        # reset task interface (ig, solvers, etc..) + 
+        # phase manager and sets bootstap as solution
+        self._gm.reset()
+        self._explosion_idx_buffer[:, :]=self._get_explosion_idx() # reset with data from reset solution
+        self._expl_idx_counter=0.0
+        self._expl_idx_buffer_counter=0
+
+    def _get_cost_info(self):
+        
+        cost_dict = self._ti.solver_rti.getCostsValues()
+        cost_names = list(cost_dict.keys())
+        cost_dims = [1] * len(cost_names) # costs are always scalar
+        return cost_names, cost_dims
+    
+    def _get_constr_info(self):
+        
+        constr_dict = self._ti.solver_rti.getConstraintsValues()
+        constr_names = list(constr_dict.keys())
+        constr_dims = [-1] * len(constr_names)
+        i = 0
+        for constr in constr_dict:
+            constr_val = constr_dict[constr]
+            constr_shape = constr_val.shape
+            constr_dims[i] = constr_shape[0]
+            i+=1
+        return constr_names, constr_dims
+    
+    def _get_q_from_sol(self):
+        full_q=self._ti.solution['q'].astype(self._dtype)
+        if self._custom_opts["replace_continuous_joints"]:
+            return full_q
+        else:
+            cont_jnts=full_q[self._continuous_joints_idxs_qfull, :]
+            cos=cont_jnts[::2, :]
+            sin=cont_jnts[1::2, :]
+            # copy root
+            self._full_q_reduced[0:7, :]=full_q[0:7, :]
+            # copy rev joint vals
+            self._full_q_reduced[self._rev_joints_idxs_red_qfull, :]=full_q[self._rev_joints_idxs_qfull, :]
+            # and continuous
+            angle=np.arctan2(sin, cos)
+            self._full_q_reduced[self._continuous_joints_idxs_red_qfull, :]=angle
+            return self._full_q_reduced
+
+    def _get_v_from_sol(self):
+        return self._ti.solution['v'].astype(self._dtype)
+    
+    def _get_a_from_sol(self):
+        return self._ti.solution['a'].astype(self._dtype)
+    
+    def _get_a_dot_from_sol(self):
+        return None
+    
+    def _get_f_from_sol(self):
+        # to be overridden by child class
+        contact_names =self._get_contacts() # we use controller-side names
+        try: 
+            data=[]
+            for key in contact_names:
+                contact_f=self._ti.solution["f_" + key].astype(self._dtype)
+                np.nan_to_num(contact_f, nan=1e6, posinf=1e6, neginf=-1e6, copy=False)
+                np.clip(a=contact_f, a_max=1e6, a_min=-1e6, out=contact_f)
+                data.append(contact_f)
+            return np.concatenate(data, axis=0)
+        except:
+            return None
+            
+    def _get_f_dot_from_sol(self):
+        # to be overridden by child class
+        return None
+    
+    def _get_eff_from_sol(self):
+        # to be overridden by child class
+        return None
+    
+    def _get_cost_from_sol(self,
+                    cost_name: str):
+        return self.rhc_costs[cost_name]
+    
+    def _get_constr_from_sol(self,
+                    constr_name: str):
+        return self.rhc_constr[constr_name]

other/horizon_based/hybrid_quad_rhc_refs.py

@@ -0,0 +1,381 @@
+from aug_mpc.controllers.rhc.horizon_based.gait_manager import GaitManager
+from aug_mpc.controllers.rhc.horizon_based.utils.math_utils import hor2w_frame
+
+from mpc_hive.utilities.shared_data.rhc_data import RhcRefs
+
+from EigenIPC.PyEigenIPC import VLevel
+from EigenIPC.PyEigenIPC import LogType
+from EigenIPC.PyEigenIPC import Journal
+
+from typing import Union
+
+import numpy as np
+
+class HybridQuadRhcRefs(RhcRefs):
+
+    def __init__(self, 
+            gait_manager: GaitManager, 
+            robot_index_shm: int,
+            robot_index_view: int,
+            namespace: str, # namespace used for shared mem
+            verbose: bool = True,
+            vlevel: bool = VLevel.V2,
+            safe: bool = True,
+            use_force_feedback: bool = False,
+            optimize_mem: bool = False):
+        
+        self.robot_index = robot_index_shm
+        self.robot_index_view = robot_index_view
+        self.robot_index_np_view = np.array(self.robot_index_view)
+
+        self._step_idx = 0
+        self._print_frequency = 100
+
+        self._verbose = verbose
+
+        self._use_force_feedback=use_force_feedback
+
+        if optimize_mem:
+            super().__init__( 
+                    is_server=False,
+                    with_gpu_mirror=False,
+                    namespace=namespace,
+                    safe=safe,
+                    verbose=verbose,
+                    vlevel=vlevel,
+                    optimize_mem=optimize_mem,
+                    n_robots=1, # we just need the row corresponding to this controller
+                    n_jnts=None, # got from server
+                    n_contacts=None # got from server
+                    )
+        else:
+            super().__init__( 
+                is_server=False,
+                with_gpu_mirror=False,
+                namespace=namespace,
+                safe=safe,
+                verbose=verbose,
+                vlevel=vlevel)
+            
+        if not isinstance(gait_manager, GaitManager):
+            exception = f"Provided gait_manager argument should be of GaitManager type!"
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+               
+        self.gait_manager = gait_manager
+
+        self.timeline_names = self.gait_manager.timeline_names
+
+        # task interfaces from horizon for setting commands to rhc
+        self._get_tasks()
+
+        self._p_ref_default=np.zeros((1, 3))
+        self._q_ref_default=np.zeros((1, 4))
+        self._q_ref_default[0, 0]=1
+
+    def _get_tasks(self):
+        # can be overridden by child
+        # cartesian tasks are in LOCAL_WORLD_ALIGNED (frame centered at distal link, oriented as WORLD)
+        self.base_lin_velxy = self.gait_manager.task_interface.getTask('base_lin_velxy')
+        self.base_lin_velz = self.gait_manager.task_interface.getTask('base_lin_velz')
+        self.base_omega = self.gait_manager.task_interface.getTask('base_omega')
+        self.base_height = self.gait_manager.task_interface.getTask('base_height')
+
+    def run(self):
+
+        super().run()
+        if not (self.robot_index < self.rob_refs.n_robots()):
+            exception = f"Provided \(0-based\) robot index {self.robot_index} exceeds number of " + \
+                " available robots {self.rob_refs.n_robots()}."
+            Journal.log(self.__class__.__name__,
+                "run",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+        contact_names = list(self.gait_manager.task_interface.model.cmap.keys())
+        if not (self.n_contacts() == len(contact_names)):
+            exception = f"N of contacts within problem {len(contact_names)} does not match n of contacts {self.n_contacts()}"
+            Journal.log(self.__class__.__name__,
+                "run",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+        
+        # set some defaults from gait manager
+        for i in range(self.n_contacts()):
+            self.flight_settings_req.set(data=self.gait_manager.get_flight_duration(contact_name=contact_names[i]),
+                data_type="len_remain",
+                robot_idxs=self.robot_index_np_view,
+                contact_idx=i)
+            self.flight_settings_req.set(data=self.gait_manager.get_step_apexdh(contact_name=contact_names[i]),
+                data_type="apex_dpos",
+                robot_idxs=self.robot_index_np_view,
+                contact_idx=i)
+            self.flight_settings_req.set(data=self.gait_manager.get_step_enddh(contact_name=contact_names[i]),
+                data_type="end_dpos",
+                robot_idxs=self.robot_index_np_view,
+                contact_idx=i)
+        
+        self.flight_settings_req.synch_retry(row_index=self.robot_index,
+            col_index=0,
+            row_index_view=self.robot_index_view,
+            n_rows=1,
+            n_cols=self.flight_settings_req.n_cols,
+            read=False)
+
+    def step(self, qstate: np.ndarray = None,
+        vstate: np.ndarray = None,
+        force_norm: np.ndarray = None):
+
+        if self.is_running():
+            
+            # updates robot refs from shared mem
+            self.rob_refs.synch_from_shared_mem(robot_idx=self.robot_index, robot_idx_view=self.robot_index_view)
+            self.phase_id.synch_all(read=True, retry=True,
+                        row_index=self.robot_index,
+                        row_index_view=self.robot_index_view)
+            self.contact_flags.synch_all(read=True, retry=True,
+                        row_index=self.robot_index,
+                        row_index_view=self.robot_index_view)
+            self.flight_settings_req.synch_all(read=True, retry=True,
+                        row_index=self.robot_index,
+                        row_index_view=self.robot_index_view)
+            self._set_contact_phases(q_full=qstate)
+
+            # updated internal references with latest available ones
+            q_base=qstate[3:7,0:1] # quaternion
+            self._apply_refs_to_tasks(q_base=q_base)
+            
+            # if self._use_force_feedback:
+            #     self._set_force_feedback(force_norm=force_norm)
+
+            self._step_idx +=1
+        
+        else:
+            exception = f"{self.__class__.__name__} is not running"
+            Journal.log(self.__class__.__name__,
+                "step",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+    
+    def _set_contact_phases(self,
+        q_full: np.ndarray):
+
+        # phase_id = self.phase_id.read_retry(row_index=self.robot_index,
+        #                         col_index=0)[0]
+        
+        contact_flags_refs = self.contact_flags.get_numpy_mirror()[self.robot_index_np_view, :]
+        target_n_limbs_in_contact=np.sum(contact_flags_refs).item()
+        if target_n_limbs_in_contact==0:
+            target_n_limbs_in_contact=4
+
+        is_contact = contact_flags_refs.flatten().tolist() 
+        n_contacts=len(is_contact)
+
+        for i in range(n_contacts): # loop through contact timelines
+            timeline_name = self.timeline_names[i]
+            
+            self.gait_manager.set_f_reg(contact_name=timeline_name,
+                scale=target_n_limbs_in_contact)
+
+            if is_contact[i]==False: # release contact
+
+                # flight parameters requests are set only when inserting a flight phase
+                len_req_now=int(self.flight_settings_req.get(data_type="len_remain",
+                    robot_idxs=self.robot_index_np_view,
+                    contact_idx=i).item())
+                apex_req_now=self.flight_settings_req.get(data_type="apex_dpos",
+                    robot_idxs=self.robot_index_np_view,
+                    contact_idx=i).item()
+                end_req_now=self.flight_settings_req.get(data_type="end_dpos",
+                    robot_idxs=self.robot_index_np_view,
+                    contact_idx=i).item()
+                landing_dx_req_now=self.flight_settings_req.get(data_type="land_dx",
+                    robot_idxs=self.robot_index_np_view,
+                    contact_idx=i).item()
+                landing_dy_req_now=self.flight_settings_req.get(data_type="land_dy",
+                    robot_idxs=self.robot_index_np_view,
+                    contact_idx=i).item()
+                
+                # set flight phase properties depending on last value on shared memory
+                self.gait_manager.set_flight_duration(contact_name=timeline_name,
+                    val=len_req_now)
+                self.gait_manager.set_step_apexdh(contact_name=timeline_name,
+                    val=apex_req_now)
+                self.gait_manager.set_step_enddh(contact_name=timeline_name,
+                    val=end_req_now)
+                self.gait_manager.set_step_landing_dx(contact_name=timeline_name,
+                    val=landing_dx_req_now)
+                self.gait_manager.set_step_landing_dy(contact_name=timeline_name,
+                    val=landing_dy_req_now)   
+                # insert flight phase over the horizon
+                self.gait_manager.add_flight(contact_name=timeline_name,
+                    robot_q=q_full)
+                                
+            else: # contact phase
+                self.gait_manager.add_stand(contact_name=timeline_name)
+
+            at_least_one_flight=self.gait_manager.update_flight_info(timeline_name)
+            # flight_info=self.gait_manager.get_flight_info(timeline_name)
+            
+            self.gait_manager.check_horizon_full(timeline_name=timeline_name)
+        
+        # write flight info to shared mem for all contacts in one shot (we follow same order as in flight_info shm)
+        all_flight_info=self.gait_manager.get_flight_info_all()
+        flight_info_shared=self.flight_info.get_numpy_mirror()
+        flight_info_shared[self.robot_index_np_view, :]=all_flight_info
+        self.flight_info.synch_retry(row_index=self.robot_index, 
+                                col_index=0, 
+                                row_index_view=self.robot_index_np_view,
+                                n_rows=1, n_cols=self.flight_info.n_cols,
+                                read=False)
+                                     
+        self.gait_manager.update()
+      
+    def _apply_refs_to_tasks(self, q_base = None):
+        # overrides parent
+        if q_base is not None: # rhc refs are assumed to be specified in the so called "horizontal" 
+            # frame, i.e. a vertical frame, with the x axis aligned with the projection of the base x axis
+            # onto the plane
+            root_pose = self.rob_refs.root_state.get(data_type = "q_full", 
+                                robot_idxs=self.robot_index_np_view).reshape(-1, 1) # this should also be
+            # rotated into the horizontal frame (however, for now only the z componet is used, so it's ok)
+            
+            root_twist_ref = self.rob_refs.root_state.get(data_type="twist", 
+                                robot_idxs=self.robot_index_np_view).reshape(-1, 1)
+
+            root_twist_ref_h = root_twist_ref.copy() 
+
+            hor2w_frame(root_twist_ref, q_base, root_twist_ref_h) # horizon works in local world aligned frame
+            
+            if self.base_lin_velxy is not None:
+                self.base_lin_velxy.setRef(root_twist_ref_h[0:2, :])
+            if self.base_omega is not None:
+                self.base_omega.setRef(root_twist_ref_h[3:, :])
+            if self.base_lin_velz is not None:
+                self.base_lin_velz.setRef(root_twist_ref_h[2:3, :])
+            if self.base_height is not None:
+                self.base_height.setRef(root_pose) 
+        else:
+            root_pose = self.rob_refs.root_state.get(data_type = "q_full", 
+                                robot_idxs=self.robot_index_np_view).reshape(-1, 1)
+            root_twist_ref = self.rob_refs.root_state.get(data_type="twist", 
+                                robot_idxs=self.robot_index_np_view).reshape(-1, 1)
+
+            if self.base_lin_velxy is not None:
+                self.base_lin_velxy.setRef(root_twist_ref[0:2, :])
+            if self.base_omega is not None:
+                self.base_omega.setRef(root_twist_ref[3:, :])
+            if self.base_lin_velz is not None:
+                self.base_lin_velz.setRef(root_twist_ref[2:3, :])
+            if self.base_height is not None:
+                self.base_height.setRef(root_pose)
+    
+    # def _set_force_feedback(self,
+    #         force_norm: np.ndarray = None):
+    
+    #     is_contact=force_norm>1.0
+
+    #     for i in range(len(is_contact)):
+    #         timeline_name = self._timeline_names[i]
+    #         self.gait_manager.set_force_feedback(timeline_name=timeline_name,
+    #             force_norm=force_norm[i])
+
+    #         if not is_contact[i]:
+
+
+    def set_default_refs(self,
+        p_ref: np.ndarray,
+        q_ref: np.ndarray):
+
+        self._p_ref_default[:, :]=p_ref
+        self._q_ref_default[:, :]=q_ref
+    
+    def set_alpha(self, alpha:float):
+        # set provided value
+        alpha_shared=self.alpha.get_numpy_mirror()
+        alpha_shared[self.robot_index_np_view, :] = alpha
+        self.alpha.synch_retry(row_index=self.robot_index, col_index=0, 
+                row_index_view=self.robot_index_view,
+                n_rows=1, n_cols=self.alpha.n_cols,
+                read=False)
+            
+    def get_alpha(self):
+        self.alpha.synch_retry(row_index=self.robot_index, col_index=0, 
+                    row_index_view=self.robot_index_view,
+                    n_rows=1, n_cols=self.alpha.n_cols,
+                    read=True)
+        alpha=self.alpha.get_numpy_mirror()[self.robot_index_np_view, :].item()
+        return alpha
+
+    def set_bound_relax(self, bound_relax:float):
+        # set provided value
+        bound_rel_shared=self.bound_rel.get_numpy_mirror()
+        bound_rel_shared[self.robot_index_np_view, :] = bound_relax
+        self.bound_rel.synch_retry(row_index=self.robot_index, col_index=0, n_rows=1, 
+            row_index_view=self.robot_index_view,
+            n_cols=self.alpha.n_cols,
+            read=False)
+
+    def reset(self):
+
+        if self.is_running():
+
+            # resets shared mem
+            contact_flags_current = self.contact_flags.get_numpy_mirror()
+            phase_id_current = self.phase_id.get_numpy_mirror()
+            contact_flags_current[self.robot_index_np_view, :] = np.full((1, self.n_contacts()), dtype=np.bool_, fill_value=True)
+            phase_id_current[self.robot_index_np_view, :] = -1 # defaults to custom phase id
+
+            contact_pos_current=self.rob_refs.contact_pos.get_numpy_mirror()
+            contact_pos_current[self.robot_index_np_view, :] = 0.0
+
+            flight_info_current=self.flight_info.get_numpy_mirror()
+            flight_info_current[self.robot_index_np_view, :] = 0.0
+
+            alpha=self.alpha.get_numpy_mirror()
+            alpha[self.robot_index_np_view, :] = 0.0
+
+            self.rob_refs.root_state.set(data_type="p", data=self._p_ref_default, robot_idxs=self.robot_index_np_view)
+            self.rob_refs.root_state.set(data_type="q", data=self._q_ref_default, robot_idxs=self.robot_index_np_view)
+            self.rob_refs.root_state.set(data_type="twist", data=np.zeros((1, 6)), robot_idxs=self.robot_index_np_view)
+                                           
+            self.contact_flags.synch_retry(row_index=self.robot_index, col_index=0, 
+                                    row_index_view=self.robot_index_view,
+                                    n_rows=1, n_cols=self.contact_flags.n_cols,
+                                    read=False)
+            self.phase_id.synch_retry(row_index=self.robot_index, col_index=0, 
+                                    row_index_view=self.robot_index_view,
+                                    n_rows=1, n_cols=self.phase_id.n_cols,
+                                    read=False)
+            self.rob_refs.root_state.synch_retry(row_index=self.robot_index, col_index=0, 
+                                    row_index_view=self.robot_index_view,
+                                    n_rows=1, n_cols=self.rob_refs.root_state.n_cols,
+                                    read=False)
+
+            self.rob_refs.contact_pos.synch_retry(row_index=self.robot_index, col_index=0, 
+                                    row_index_view=self.robot_index_view,
+                                    n_rows=1, n_cols=self.rob_refs.contact_pos.n_cols,
+                                    read=False)
+            
+            self.flight_info.synch_retry(row_index=self.robot_index, 
+                                col_index=0, 
+                                row_index_view=self.robot_index_view,
+                                n_rows=1, n_cols=self.flight_info.n_cols,
+                                read=False)
+            
+            # should also empty the timeline for stepping phases
+            self._step_idx = 0
+
+        else:
+            exception = f"Cannot call reset() since run() was not called!"
+            Journal.log(self.__class__.__name__,
+                "reset",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+

other/horizon_based/utils/math_utils.py

@@ -0,0 +1,165 @@
+
+import numpy as np
+
+def w2hor_frame(t_w: np.ndarray,
+        q_b: np.ndarray,
+        t_out: np.ndarray):
+    """
+    Transforms a twist vector expressed in WORLD frame to
+    an "horizontal" frame (z aligned as world, x aligned as the projection
+    of the x-axis of the base frame described by q_b). This is useful for specifying locomotion
+    references in a "game"-like fashion.
+    t_out will hold the result
+    """
+    # q_b = q_b / q_b.norm(dim=1, keepdim=True)
+    q_w, q_i, q_j, q_k = q_b[3, :], q_b[0, :], q_b[1, :], q_b[2, :]
+    
+    R_11 = 1 - 2 * (q_j ** 2 + q_k ** 2)
+    R_21 = 2 * (q_i * q_j + q_k * q_w)
+    
+    norm = np.sqrt(R_11 ** 2 + R_21 ** 2)
+    x_proj_x = R_11 / norm
+    x_proj_y = R_21 / norm
+    
+    y_proj_x = -x_proj_y
+    y_proj_y = x_proj_x
+        
+    t_out[0, :] = t_w[0, :] * x_proj_x + t_w[1, :] * x_proj_y
+    t_out[1, :] = t_w[0, :] * y_proj_x + t_w[1, :] * y_proj_y
+    t_out[2, :] = t_w[2, :]  # z-component remains the same
+    t_out[3, :] = t_w[3, :] * x_proj_x + t_w[4, :] * x_proj_y
+    t_out[4, :] = t_w[3, :] * y_proj_x + t_w[4, :] * y_proj_y
+    t_out[5, :] = t_w[5, :]  # z-component remains the same
+
+def hor2w_frame(t_h: np.ndarray,
+        q_b: np.ndarray,
+        t_out: np.ndarray):
+    """
+    Transforms a velocity vector expressed in "horizontal" frame to WORLD
+    v_out will hold the result
+    """
+
+    # Extract quaternion components
+    q_w, q_i, q_j, q_k = q_b[3, :], q_b[0, :], q_b[1, :], q_b[2, :]
+    
+    # Compute rotation matrix elements
+    R_11 = 1 - 2 * (q_j ** 2 + q_k ** 2)
+    R_21 = 2 * (q_i * q_j + q_k * q_w)
+    
+    # Normalize to get projection components
+    norm = np.sqrt(R_11 ** 2 + R_21 ** 2)
+    x_proj_x = R_11 / norm
+    x_proj_y = R_21 / norm
+    
+    # Orthogonal vector components
+    y_proj_x = -x_proj_y
+    y_proj_y = x_proj_x
+    
+    # Transform velocity vector components from horizontal to world frame
+    t_out[0, :] = t_h[0, :] * x_proj_x + t_h[1, :] * y_proj_x
+    t_out[1, :] = t_h[0, :] * x_proj_y + t_h[1, :] * y_proj_y
+    t_out[2, :] = t_h[2, :]  # z-component remains the same
+    t_out[3, :] = t_h[3, :] * x_proj_x + t_h[4, :] * y_proj_x
+    t_out[4, :] = t_h[3, :] * x_proj_y + t_h[4, :] * y_proj_y
+    t_out[5, :] = t_h[5, :]  # z-component remains the same
+
+def base2world_frame(t_b: np.ndarray, 
+                q_b: np.ndarray,
+                t_out: np.ndarray):
+    """
+    Transforms a velocity vector expressed in the base frame to
+    the WORLD frame using the given quaternion that describes the orientation
+    of the base with respect to the world frame. The result is written in v_out.
+    """
+    # q_b = q_b / q_b.norm(dim=1, keepdim=True)
+    q_w, q_i, q_j, q_k = q_b[3, :], q_b[0, :], q_b[1, :], q_b[2, :]
+    
+    R_00 = 1 - 2 * (q_j ** 2 + q_k ** 2)
+    R_01 = 2 * (q_i * q_j - q_k * q_w)
+    R_02 = 2 * (q_i * q_k + q_j * q_w)
+    
+    R_10 = 2 * (q_i * q_j + q_k * q_w)
+    R_11 = 1 - 2 * (q_i ** 2 + q_k ** 2)
+    R_12 = 2 * (q_j * q_k - q_i * q_w)
+    
+    R_20 = 2 * (q_i * q_k - q_j * q_w)
+    R_21 = 2 * (q_j * q_k + q_i * q_w)
+    R_22 = 1 - 2 * (q_i ** 2 + q_j ** 2)
+    
+    # Transform the velocity to the world frame
+    t_out[0, :] = t_b[0, :] * R_00 + t_b[1, :] * R_01 + t_b[2, :] * R_02
+    t_out[1, :] = t_b[0, :] * R_10 + t_b[1, :] * R_11 + t_b[2, :] * R_12
+    t_out[2, :] = t_b[0, :] * R_20 + t_b[1, :] * R_21 + t_b[2, :] * R_22
+    t_out[3, :] = t_b[3, :] * R_00 + t_b[4, :] * R_01 + t_b[5, :] * R_02
+    t_out[4, :] = t_b[3, :] * R_10 + t_b[4, :] * R_11 + t_b[5, :] * R_12
+    t_out[5, :] = t_b[3, :] * R_20 + t_b[4, :] * R_21 + t_b[5, :] * R_22
+
+def world2base_frame(t_w: np.ndarray, 
+                q_b: np.ndarray, 
+                t_out: np.ndarray):
+    """
+    Transforms a velocity vector expressed in the WORLD frame to
+    the base frame using the given quaternion that describes the orientation
+    of the base with respect to the world frame. The result is written in v_out.
+    """
+    # q_b = q_b / q_b.norm(dim=1, keepdim=True)
+    q_w, q_i, q_j, q_k = q_b[3, :], q_b[0, :], q_b[1, :], q_b[2, :]
+    
+    R_00 = 1 - 2 * (q_j ** 2 + q_k ** 2)
+    R_01 = 2 * (q_i * q_j - q_k * q_w)
+    R_02 = 2 * (q_i * q_k + q_j * q_w)
+    
+    R_10 = 2 * (q_i * q_j + q_k * q_w)
+    R_11 = 1 - 2 * (q_i ** 2 + q_k ** 2)
+    R_12 = 2 * (q_j * q_k - q_i * q_w)
+    
+    R_20 = 2 * (q_i * q_k - q_j * q_w)
+    R_21 = 2 * (q_j * q_k + q_i * q_w)
+    R_22 = 1 - 2 * (q_i ** 2 + q_j ** 2)
+        
+    # Transform the velocity to the base frame using the transpose of the rotation matrix
+    t_out[0, :] = t_w[0, :] * R_00 + t_w[1, :] * R_10 + t_w[2, :] * R_20
+    t_out[1, :] = t_w[0, :] * R_01 + t_w[1, :] * R_11 + t_w[2, :] * R_21
+    t_out[2, :] = t_w[0, :] * R_02 + t_w[1, :] * R_12 + t_w[2, :] * R_22
+    t_out[3, :] = t_w[3, :] * R_00 + t_w[4, :] * R_10 + t_w[5, :] * R_20
+    t_out[4, :] = t_w[3, :] * R_01 + t_w[4, :] * R_11 + t_w[5, :] * R_21
+    t_out[5, :] = t_w[3, :] * R_02 + t_w[4, :] * R_12 + t_w[5, :] * R_22
+
+if __name__ == "__main__":  
+
+    n_envs = 5000
+    t_b = np.random.rand(6, n_envs)
+
+    q_b = np.random.rand(4, n_envs)
+    q_b_norm = q_b / np.linalg.norm(q_b, axis=0, keepdims=True)
+
+    t_w = np.zeros_like(t_b)  # To hold horizontal frame velocities
+    t_b_recovered = np.zeros_like(t_b)  # To hold recovered world frame velocities
+    base2world_frame(t_b, q_b_norm, t_w)
+    world2base_frame(t_w, q_b_norm, t_b_recovered)
+    assert np.allclose(t_b, t_b_recovered, atol=1e-6), "Test failed: t_w_recovered does not match t_b"
+    print("Forward test passed: t_b_recovered matches t_b")
+    
+    t_b2 = np.zeros_like(t_b)  # To hold horizontal frame velocities
+    t_w_recovered = np.zeros_like(t_b)
+    world2base_frame(t_b, q_b_norm, t_b2)
+    base2world_frame(t_b2, q_b_norm, t_w_recovered)
+    assert np.allclose(t_b, t_w_recovered, atol=1e-6), "Test failed: t_w_recovered does not match t_b"
+    print("Backward test passed: t_w_recovered matches t_w")
+    
+    # test transf. world-horizontal frame
+    v_h = np.zeros_like(t_b)  # To hold horizontal frame velocities
+    v_recovered = np.zeros_like(t_b)
+    w2hor_frame(t_b, q_b_norm, v_h)
+    hor2w_frame(v_h, q_b_norm, v_recovered)
+    assert np.allclose(t_b, v_recovered, atol=1e-6), "Test failed: v_recovered does not match t_b"
+    print("horizontal forward frame test passed:  matches ")
+
+    t_w = np.zeros_like(t_b)  # To hold horizontal frame velocities
+    v_h_recovered = np.zeros_like(t_b)
+    hor2w_frame(t_b, q_b_norm, t_w)
+    w2hor_frame(t_w, q_b_norm, v_h_recovered)
+    assert np.allclose(t_b, v_h_recovered, atol=1e-6), "Test failed: v_h_recovered does not match t_b"
+    print("horizontal backward frame test passed:  matches ")
+
+    

other/hybrid_quad_client.py

@@ -0,0 +1,95 @@
+from aug_mpc.controllers.rhc.augmpc_cluster_client import AugMpcClusterClient
+
+from aug_mpc.controllers.rhc.horizon_based.hybrid_quad_rhc import HybridQuadRhc
+from aug_mpc.utils.hybrid_quad_xrdf_gen import get_xrdf_cmds_horizon
+from aug_mpc.utils.sys_utils import PathsGetter
+
+from typing import List, Dict
+
+class HybridQuadrupedClusterClient(AugMpcClusterClient):
+    
+    def _import_aux_libs(self):
+        super()._import_aux_libs()
+        # Import Horizon and related dependencies as global libs
+        from aug_mpc.controllers.rhc.horizon_based.horizon_imports_glob import import_horizon_global
+        import_horizon_global()
+        
+    def __init__(self, 
+            namespace: str, 
+            urdf_xacro_path: str,
+            srdf_xacro_path: str,
+            cluster_size: int,
+            set_affinity: bool = False,
+            use_mp_fork: bool = False,
+            isolated_cores_only: bool = False,
+            core_ids_override_list: List[int] = None,
+            verbose: bool = False,
+            debug: bool = False,
+            open_loop: bool = True,
+            base_dump_dir: str = "/tmp",
+            timeout_ms: int = 60000,
+            codegen_override: str = None,
+            custom_opts: Dict = {}):
+        
+        self._open_loop = open_loop
+
+        self._paths = PathsGetter()
+
+        self._codegen_dir_name = namespace
+
+        self._timeout_ms = timeout_ms
+
+        super().__init__(namespace = namespace, 
+                        urdf_xacro_path = urdf_xacro_path,
+                        srdf_xacro_path = srdf_xacro_path,
+                        cluster_size=cluster_size,
+                        set_affinity = set_affinity,
+                        use_mp_fork = use_mp_fork,
+                        isolated_cores_only = isolated_cores_only,
+                        core_ids_override_list = core_ids_override_list,
+                        verbose = verbose,
+                        debug = debug,
+                        base_dump_dir=base_dump_dir,
+                        codegen_override=codegen_override,
+                        custom_opts=custom_opts)
+
+        self._n_nodes = 31 if not ("n_nodes" in self._custom_opts) else self._custom_opts["n_nodes"]
+        self._dt = 0.05 if not ("cluster_dt" in self._custom_opts) else self._custom_opts["cluster_dt"]
+        
+    def _xrdf_cmds(self):
+        parts = self._urdf_path.split('/')
+        urdf_descr_root_path = '/'.join(parts[:-2])
+        cmds = get_xrdf_cmds_horizon(urdf_descr_root_path=urdf_descr_root_path)
+        return cmds
+
+    def _process_codegen_dir(self,idx:int):
+
+        codegen_dir = self.codegen_dir() + f"/{self._codegen_dir_name}Rhc{idx}"
+        codegen_dir_ovveride = self.codegen_dir_override()
+        if not (codegen_dir_ovveride=="" or \
+                codegen_dir_ovveride=="none" or \
+                codegen_dir_ovveride=="None" or \
+                (codegen_dir_ovveride is None)): # if overrde was provided
+            codegen_dir = f"{codegen_dir_ovveride}{idx}"# override
+        
+        return codegen_dir
+        
+    def _generate_controller(self,
+                        idx: int):
+        
+        codegen_dir=self._process_codegen_dir(idx=idx)
+
+        controller = HybridQuadRhc(
+                urdf_path=self._urdf_path, 
+                srdf_path=self._srdf_path,
+                config_path = self._paths.CONFIGPATH,
+                robot_name=self._namespace,
+                codegen_dir=codegen_dir,
+                n_nodes=self._n_nodes, 
+                dt=self._dt,
+                max_solver_iter = 1, # rti
+                open_loop = self._open_loop,
+                verbose = self._verbose, 
+                debug = self._debug)
+        
+        return controller 

rhc.py

@@ -0,0 +1,1262 @@
+# Copyright (C) 2023  Andrea Patrizi (AndrePatri, andreapatrizi1b6e6@gmail.com)
+# 
+# This file is part of MPCHive and distributed under the General Public License version 2 license.
+# 
+# MPCHive is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 2 of the License, or
+# (at your option) any later version.
+# 
+# MPCHive is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+# 
+# You should have received a copy of the GNU General Public License
+# along with MPCHive.  If not, see <http://www.gnu.org/licenses/>.
+# 
+from abc import ABC, abstractmethod
+# from perf_sleep.pyperfsleep import PerfSleep
+# from mpc_hive.utilities.cpu_utils.core_utils import get_memory_usage
+
+import time 
+
+import numpy as np
+
+from mpc_hive.utilities.shared_data.rhc_data import RobotState
+from mpc_hive.utilities.shared_data.rhc_data import RhcCmds, RhcPred, RhcPredDelta
+from mpc_hive.utilities.shared_data.rhc_data import RhcStatus
+from mpc_hive.utilities.shared_data.rhc_data import RhcInternal
+from mpc_hive.utilities.shared_data.cluster_profiling import RhcProfiling
+from mpc_hive.utilities.remote_triggering import RemoteTriggererClnt
+
+from mpc_hive.utilities.homing import RobotHomer
+
+from mpc_hive.utilities.math_utils import world2base_frame
+
+from EigenIPC.PyEigenIPC import VLevel
+from EigenIPC.PyEigenIPC import Journal, LogType
+from EigenIPC.PyEigenIPCExt.wrappers.shared_data_view import SharedTWrapper
+from EigenIPC.PyEigenIPC import dtype
+
+from typing import List
+# from typing import TypeVar, Union
+
+import signal
+import os
+import inspect
+
+class RHController(ABC):
+
+    def __init__(self, 
+            srdf_path: str,
+            n_nodes: int,
+            dt: float,
+            namespace: str, # shared mem namespace
+            dtype = np.float32, 
+            verbose = False, 
+            debug = False,
+            timeout_ms: int = 60000,
+            allow_less_jnts: bool = True):
+    
+        signal.signal(signal.SIGINT, self._handle_sigint)
+
+        self._allow_less_jnts = allow_less_jnts # whether to allow less joints in rhc controller than the ones on the robot
+        # (e.g. some joints might not be desirable for control purposes)
+
+        self.namespace = namespace
+        self._dtype = dtype
+        self._verbose = verbose
+        self._debug = debug
+
+        # if not self._debug:
+        np.seterr(over='ignore') # ignoring overflows
+
+        self._n_nodes = n_nodes
+        self._dt = dt
+        self._n_intervals = self._n_nodes - 1 
+        self._t_horizon = self._n_intervals * dt
+        self._set_rhc_pred_idx() # prection is by default written on last node
+        self._set_rhc_cmds_idx() # default to idx 2 (i.e. cmds to get to third node)
+        self.controller_index = None # will be assigned upon registration to a cluster
+        self.controller_index_np = None 
+        self._robot_mass=1.0
+
+        self.srdf_path = srdf_path # using for parsing robot homing
+
+        self._registered = False
+        self._closed = False 
+        
+        self._allow_triggering_when_failed = True
+        
+        self._profiling_data_dict = {}
+        self._profiling_data_dict["full_solve_dt"] = np.nan
+        self._profiling_data_dict["rti_solve_dt"] = np.nan
+        self._profiling_data_dict["problem_update_dt"] = np.nan
+        self._profiling_data_dict["phases_shift_dt"] = np.nan
+        self._profiling_data_dict["task_ref_update"] = np.nan
+        
+        self.n_dofs = None
+        self.n_contacts = None
+        
+        # shared mem
+        self.robot_state = None 
+        self.rhc_status = None
+        self.rhc_internal = None
+        self.cluster_stats = None
+        self.robot_cmds = None
+        self.robot_pred = None
+        self.rhc_pred_delta = None
+        self.rhc_refs = None
+        self._remote_triggerer = None
+        self._remote_triggerer_timeout = timeout_ms # [ms]
+        
+        # remote termination
+        self._remote_term = None
+        self._term_req_received = False
+
+        # jnt names
+        self._env_side_jnt_names = []
+        self._controller_side_jnt_names = []
+        self._got_jnt_names_from_controllers = False
+
+        # data maps
+        self._to_controller = []
+        self._quat_remap = [0, 1, 2, 3] # defaults to no remap (to be overridden)
+
+        self._got_contact_names = False
+
+        self._received_trigger = False # used for proper termination
+
+        self._n_resets = 0
+        self._n_fails = 0
+        self._fail_idx_thresh = 5e3
+        self._contact_f_thresh = 1e5
+
+        self._failed = False
+
+        self._start_time = time.perf_counter() # used for profiling when in debug mode
+
+        self._homer = None # robot homing manager
+        self._homer_env = None # used for setting homing to jnts not contained in rhc prb
+
+        self._class_name_base = f"{self.__class__.__name__}"
+        self._class_name = self._class_name_base # will append controller index upon registration
+
+        self._contact_force_base_loc_aux=np.zeros((1,3),dtype=self._dtype)
+        self._norm_grav_vector_w=np.zeros((1,3),dtype=self._dtype)
+        self._norm_grav_vector_w[:, 2]=-1.0
+        self._norm_grav_vector_base_loc=np.zeros((1,3),dtype=self._dtype)
+
+        self._init() # initialize controller
+
+        if not hasattr(self, '_rhc_fpaths'):
+            self._rhc_fpaths = []
+        self._rhc_fpaths.append(os.path.abspath(__file__))
+        
+    def __init_subclass__(cls, **kwargs):
+        super().__init_subclass__(**kwargs)
+        # Get the file path of the class being initialized and append it to the list
+        if not hasattr(cls, '_rhc_fpaths'):
+            cls._rhc_fpaths = []
+        child_class_file_path = os.path.abspath(inspect.getfile(cls))
+        cls._rhc_fpaths.append(child_class_file_path)
+
+    def this_paths(self):
+        return self._rhc_fpaths
+    
+    def __del__(self):
+        self._close()
+
+    def _handle_sigint(self, signum, frame):
+        if self._verbose:
+            Journal.log(self._class_name,
+                    "_handle_sigint",
+                    "SIGINT received",
+                    LogType.WARN)
+        self._term_req_received = True
+    
+    def _set_rhc_pred_idx(self):
+        # default to last node
+        self._pred_node_idx=self._n_nodes-1
+    
+    def _set_rhc_cmds_idx(self):
+        # use index 2 by default to compensate for
+        # the inevitable action delay 
+        # (get_state, trigger sol -> +dt -> apply sol )
+        # if we apply action from second node (idenx 1) 
+        # that action will already be one dt old. We assume we were already
+        # applying the right action to get to the state of node idx 1 and get the 
+        # cmds for reaching the state at idx 1
+        self._rhc_cmds_node_idx=2
+
+    def _close(self):
+        if not self._closed:
+            self._unregister_from_cluster()
+            if self.robot_cmds is not None:
+                self.robot_cmds.close()
+            if self.robot_pred is not None:
+                self.robot_pred.close()
+            if self.rhc_pred_delta is not None:
+                self.rhc_pred_delta.close()
+            if self.robot_state is not None:
+                self.robot_state.close()
+            if self.rhc_status is not None:
+                self.rhc_status.close()
+            if self.rhc_internal is not None:
+                self.rhc_internal.close()
+            if self.cluster_stats is not None:
+                self.cluster_stats.close()
+            if self._remote_triggerer is not None:
+                self._remote_triggerer.close()
+            if self._remote_term is not None:
+                self._remote_term.close()
+            self._closed = True
+
+    def close(self):
+        self._close()
+
+    def get_file_paths(self):
+        # can be overriden by child
+        base_paths = []
+        base_paths.append(self._this_path)
+        return base_paths
+    
+    def init_rhc_task_cmds(self):
+        
+        self.rhc_refs = self._init_rhc_task_cmds()
+        self.rhc_refs.reset()
+        
+    def _init_states(self):
+        
+        quat_remap = self._get_quat_remap()
+        self.robot_state = RobotState(namespace=self.namespace,
+                                is_server=False,
+                                q_remapping=quat_remap, # remapping from environment to controller
+                                with_gpu_mirror=False,
+                                with_torch_view=False, 
+                                safe=False,
+                                verbose=self._verbose,
+                                vlevel=VLevel.V2,
+                                optimize_mem=True,
+                                n_robots=1, # we just need the row corresponding to this controller
+                                n_jnts=None, # got from server
+                                n_contacts=None # got from server
+                                ) 
+        self.robot_state.run()
+        self.robot_cmds = RhcCmds(namespace=self.namespace,
+                                is_server=False,
+                                q_remapping=quat_remap, # remapping from environment to controller
+                                with_gpu_mirror=False,
+                                with_torch_view=False, 
+                                safe=False,
+                                verbose=self._verbose,
+                                vlevel=VLevel.V2,
+                                optimize_mem=True,
+                                n_robots=1, # we just need the row corresponding to this controller
+                                n_jnts=None, # got from server
+                                n_contacts=None # got from server
+                                ) 
+        self.robot_cmds.run()
+        self.robot_pred = RhcPred(namespace=self.namespace,
+                                is_server=False,
+                                q_remapping=quat_remap, # remapping from environment to controller
+                                with_gpu_mirror=False,
+                                with_torch_view=False, 
+                                safe=False,
+                                verbose=self._verbose,
+                                vlevel=VLevel.V2,
+                                optimize_mem=True,
+                                n_robots=1, # we just need the row corresponding to this controller
+                                n_jnts=None, # got from server
+                                n_contacts=None # got from server
+                                )
+        self.robot_pred.run()
+        self.rhc_pred_delta = RhcPredDelta(namespace=self.namespace,
+                                is_server=False,
+                                q_remapping=quat_remap, # remapping from environment to controller
+                                with_gpu_mirror=False,
+                                with_torch_view=False, 
+                                safe=False,
+                                verbose=self._verbose,
+                                vlevel=VLevel.V2,
+                                optimize_mem=True,
+                                n_robots=1, # we just need the row corresponding to this controller
+                                n_jnts=None, # got from server
+                                n_contacts=None # got from server
+                                )
+        self.rhc_pred_delta.run()
+
+    def _rhc(self, rti: bool = True):
+        if self._debug:
+            self._rhc_db(rti=rti)
+        else:
+            self._rhc_min(rti=rti)
+    
+    def _rhc_db(self, rti: bool = True):
+        # rhc with debug data
+        self._start_time = time.perf_counter()
+
+        self.robot_state.synch_from_shared_mem(robot_idx=self.controller_index, robot_idx_view=self.controller_index_np) # updates robot state with
+        # latest data on shared mem
+
+        self._compute_pred_delta()
+
+        if not self.failed():
+            # we can solve only if not in failure state
+            if rti:
+                self._failed = not self._solve() # solve actual TO with RTI
+            else:
+                self._failed = not self._bootstrap() # full bootstrap solve
+            if (self._failed): 
+                # perform failure procedure
+                self._on_failure()                       
+        else:
+            if not self._allow_triggering_when_failed:
+                Journal.log(self._class_name,
+                    "solve",
+                    f"Received solution req, but controller is in failure state. " + \
+                        " You should have reset() the controller!",
+                    LogType.EXCEP,
+                    throw_when_excep = True)
+            else: 
+                if self._verbose:
+                    Journal.log(self._class_name,
+                        "solve",
+                        f"Received solution req, but controller is in failure state. No solution will be performed. " + \
+                            " Use the reset() method to continue solving!",
+                        LogType.WARN)
+            
+        self._write_cmds_from_sol() # we update update the views of the cmds
+        # from the latest solution
+    
+        # in debug, rhc internal state is streamed over 
+        # shared mem.
+        self._update_rhc_internal()
+        self._profiling_data_dict["full_solve_dt"] = time.perf_counter() - self._start_time
+        self._update_profiling_data() # updates all profiling data
+        if self._verbose:
+            Journal.log(self._class_name,
+                "solve",
+                f"RHC full solve loop execution time  -> " + str(self._profiling_data_dict["full_solve_dt"]),
+                LogType.INFO,
+                throw_when_excep = True) 
+
+    def _rhc_min(self, rti: bool = True):
+
+        self.robot_state.synch_from_shared_mem(robot_idx=self.controller_index, robot_idx_view=self.controller_index_np) # updates robot state with
+        # latest data on shared mem
+
+        self._compute_pred_delta()
+
+        if not self.failed():
+            # we can solve only if not in failure state
+            if rti:
+                self._failed = not self._solve() # solve actual TO with RTI
+            else:
+                self._failed = not self._bootstrap() # full bootstrap solve
+            if (self._failed):  
+                # perform failure procedure
+                self._on_failure()                       
+        else:
+            if not self._allow_triggering_when_failed:
+                Journal.log(self._class_name,
+                    "solve",
+                    f"Received solution req, but controller is in failure state. " + \
+                        " You should have reset() the controller!",
+                    LogType.EXCEP,
+                    throw_when_excep = True)
+            else: 
+                if self._verbose:
+                    Journal.log(self._class_name,
+                        "solve",
+                        f"Received solution req, but controller is in failure state. No solution will be performed. " + \
+                            " Use the reset() method to continue solving!",
+                        LogType.WARN)
+                    
+        self._write_cmds_from_sol() # we update the views of the cmds
+        # from the latest solution even if failed
+        
+    def solve_once(self):
+        # run a single iteration of the solve loop (used for pooling)
+        if self._term_req_received:
+            return False
+
+        if not self._remote_triggerer.wait(self._remote_triggerer_timeout):
+            Journal.log(self._class_name,
+                f"solve",
+                "Didn't receive any remote trigger req within timeout!",
+                LogType.EXCEP,
+                throw_when_excep = False)
+            return False
+
+        self._received_trigger = True
+
+        if self.rhc_status.resets.read_retry(row_index=self.controller_index,
+                                col_index=0,
+                                row_index_view=0)[0]:
+            self.reset() # rhc is reset
+
+        if self.rhc_status.trigger.read_retry(row_index=self.controller_index,
+                    col_index=0,
+                    row_index_view=0)[0]:
+            rti_solve = self.rhc_status.rti_solve.read_retry(row_index=self.controller_index,
+                        col_index=0,
+                        row_index_view=0)[0]
+            self._rhc(rti=rti_solve) # run solution with requested solve mode
+            self.rhc_status.trigger.write_retry(False, 
+                row_index=self.controller_index,
+                col_index=0,
+                row_index_view=0) # allow next solution trigger 
+        
+        self._remote_triggerer.ack() # send ack signal to server
+        self._received_trigger = False
+        
+        self._term_req_received = self._term_req_received or self._remote_term.read_retry(row_index=0,
+                                                        col_index=0,
+                                                        row_index_view=0)[0]
+        
+        if self._term_req_received:
+            self.close()
+            return False
+
+        return True
+
+    def solve(self):
+        
+        # run the solution loop and wait for trigger signals
+        # using cond. variables (efficient)
+        while True:
+            if not self.solve_once():
+                break
+        
+        self.close() # is not stricly necessary
+
+    def reset(self):
+        
+        if not self._closed:
+            self.reset_rhc_data()
+            self._failed = False # allow triggering
+            self._n_resets += 1
+            self.rhc_status.fails.write_retry(False, 
+                                    row_index=self.controller_index,
+                                    col_index=0,
+                                    row_index_view=0)
+            self.rhc_status.resets.write_retry(False, 
+                                    row_index=self.controller_index,
+                                    col_index=0,
+                                    row_index_view=0)
+
+    def _create_jnt_maps(self):
+        
+        # retrieve env-side joint names from shared mem
+        self._env_side_jnt_names = self.robot_state.jnt_names()
+        self._check_jnt_names_compatibility() # will raise exception if not self._allow_less_jnts
+        if not self._got_jnt_names_from_controllers:
+            exception = f"Cannot run the solve(). assign_env_side_jnt_names() was not called!"
+            Journal.log(self._class_name,
+                    "_create_jnt_maps",
+                    exception,
+                    LogType.EXCEP,
+                    throw_when_excep = True)
+
+        # robot homer guarantees that _controller_side_jnt_names is at least contained in self._env_side_jnt_names (or equal)
+        self._to_controller = [self._env_side_jnt_names.index(element) for element in self._controller_side_jnt_names]
+        # set joint remappings for shared data from env data to controller
+        # all shared data is by convention specified according to the env (jnts are ordered that way)
+        # the remapping is used so that when data is returned, its a remapped view from env to controller,
+        # so that data can be assigned direclty from the rhc controller without any issues
+        self.robot_state.set_jnts_remapping(jnts_remapping=self._to_controller)
+        self.robot_cmds.set_jnts_remapping(jnts_remapping=self._to_controller)
+        self.robot_pred.set_jnts_remapping(jnts_remapping=self._to_controller)
+        self.rhc_pred_delta.set_jnts_remapping(jnts_remapping=self._to_controller)
+
+        return True
+
+    def reset_rhc_data(self):
+        
+        self._reset() # custom reset (e.g. it should set the current solution to some
+        # default solution, like a bootstrap)
+
+        self.rhc_refs.reset() # reset rhc refs to default (has to be called after _reset)
+
+        self._write_cmds_from_sol() # use latest solution (e.g. from bootstrap if called before running
+        # the first solve) as default state
+    
+    def failed(self):
+        return self._failed
+
+    def robot_mass(self):
+        return self._robot_mass
+    
+    def _assign_cntrl_index(self, reg_state: np.ndarray):
+        state = reg_state.flatten() # ensure 1D tensor
+        free_spots = np.nonzero(~state.flatten())[0]
+        return free_spots[0].item()  # just return the first free spot
+    
+    def _register_to_cluster(self):
+        
+        self.rhc_status = RhcStatus(is_server=False,
+            namespace=self.namespace, 
+            verbose=self._verbose, 
+            vlevel=VLevel.V2,
+            with_torch_view=False, 
+            with_gpu_mirror=False,
+            optimize_mem=True,
+            cluster_size=1, # we just need the row corresponding to this controller
+            n_contacts=None, # we get this from server
+            n_nodes=None # we get this from server
+            )
+        self.rhc_status.run() # rhc status (reg. flags, failure, tot cost, tot cnstrl viol, etc...)
+
+        # acquire semaphores since we have to perform non-atomic operations
+        # on the whole memory views
+        self.rhc_status.registration.data_sem_acquire()
+        self.rhc_status.controllers_counter.data_sem_acquire()
+        self.rhc_status.controllers_counter.synch_all(retry = True,
+                                                read = True)
+
+        available_spots = self.rhc_status.cluster_size
+        # from here on all pre registration ops can be done
+
+        # incrementing cluster controllers counter
+        controllers_counter = self.rhc_status.controllers_counter.get_numpy_mirror()
+        if controllers_counter[0, 0] + 1 > available_spots: # no space left -> return 
+            self.rhc_status.controllers_counter.data_sem_release()
+            self.rhc_status.registration.data_sem_release()
+            exception = "Cannot register to cluster. No space left " + \
+                f"({controllers_counter[0, 0]} controllers already registered)"
+            Journal.log(self._class_name,
+                    "_register_to_cluster",
+                    exception,
+                    LogType.EXCEP,
+                    throw_when_excep = True) 
+        
+        # now we can register 
+
+        # increment controllers counter
+        controllers_counter += 1 
+        self.controller_index = controllers_counter.item() -1 
+        
+        # actually register to cluster
+        self.rhc_status.controllers_counter.synch_all(retry = True,
+            read = False) # writes to shared mem (just one for all controllers (i.e n_rows = 1))
+        
+        # read current registration state
+        self.rhc_status.registration.synch_all(retry = True,
+                                        read = True,
+                                        row_index=self.controller_index,
+                                        row_index_view=0)
+        registrations = self.rhc_status.registration.get_numpy_mirror()
+        # self.controller_index = self._assign_cntrl_index(registrations)
+        
+
+        self._class_name_base = self._class_name_base+str(self.controller_index)
+        # self.controller_index_np = np.array(self.controller_index)
+        self.controller_index_np = np.array(0) # given that we use optimize_mem, the shared mem copy has shape 1 x n_cols (we can write and read at [0, :])
+        
+        registrations[self.controller_index_np, 0] = True
+        self.rhc_status.registration.synch_all(retry = True,
+                                        read = False,
+                                        row_index=self.controller_index,
+                                        col_index=0,
+                                        row_index_view=0) 
+
+        # now all heavy stuff that would otherwise make the registration slow
+        self._remote_term = SharedTWrapper(namespace=self.namespace,
+            basename="RemoteTermination",
+            is_server=False,
+            verbose = self._verbose, 
+            vlevel = VLevel.V2,
+            with_gpu_mirror=False,
+            with_torch_view=False,
+            dtype=dtype.Bool)
+        self._remote_term.run()
+
+        # other initializations
+        
+        self.cluster_stats = RhcProfiling(is_server=False, 
+                                    name=self.namespace,
+                                    verbose=self._verbose,
+                                    vlevel=VLevel.V2,
+                                    safe=True,
+                                    optimize_mem=True,
+                                    cluster_size=1 # we just need the row corresponding to this controller
+                                    ) # profiling data
+        self.cluster_stats.run()
+        self.cluster_stats.synch_info()
+    
+        self._create_jnt_maps()
+        self.init_rhc_task_cmds() # initializes rhc interface to external commands (defined by child class)
+        self._consinstency_checks() # sanity checks
+
+        if self._homer is None:
+            self._init_robot_homer() # call this in case it wasn't called by child
+    
+        self._robot_mass = self._get_robot_mass() # uses child class implemented method
+        self._contact_f_scale = self._get_robot_mass() * 9.81
+
+        # writing some static info about this controller
+        # self.rhc_status.rhc_static_info.synch_all(retry = True,
+        #     read = True,
+        #     row_index=self.controller_index,
+        #     col_index=0) # first read current static info from other controllers (not necessary if optimize_mem==True)
+        self.rhc_status.rhc_static_info.set(data=np.array(self._dt),
+            data_type="dts",
+            rhc_idxs=self.controller_index_np,
+            gpu=False)
+        self.rhc_status.rhc_static_info.set(data=np.array(self._t_horizon),
+            data_type="horizons",
+            rhc_idxs=self.controller_index_np,
+            gpu=False)
+        self.rhc_status.rhc_static_info.set(data=np.array(self._n_nodes),
+            data_type="nnodes",
+            rhc_idxs=self.controller_index_np,
+            gpu=False)
+        # writing some static rhc info which is only available after problem init
+        self.rhc_status.rhc_static_info.set(data=np.array(len(self._get_contacts())),
+            data_type="ncontacts",
+            rhc_idxs=self.controller_index_np,
+            gpu=False)
+        self.rhc_status.rhc_static_info.set(data=np.array(self.robot_mass()),
+            data_type="robot_mass",
+            rhc_idxs=self.controller_index_np,
+            gpu=False)
+        self.rhc_status.rhc_static_info.set(data=np.array(self._pred_node_idx),
+            data_type="pred_node_idx",
+            rhc_idxs=self.controller_index_np,
+            gpu=False)
+        
+        self.rhc_status.rhc_static_info.synch_retry(row_index=self.controller_index, 
+            col_index=0,
+            row_index_view=0,
+            n_rows=1, n_cols=self.rhc_status.rhc_static_info.n_cols,
+            read=False)
+        
+        # we set homings also for joints which are not in the rhc homing map
+        # since this is usually required on server side
+    
+        homing_full = self._homer_env.get_homing().reshape(1, 
+                        self.robot_cmds.n_jnts())
+    
+        null_action = np.zeros((1, self.robot_cmds.n_jnts()), 
+                        dtype=self._dtype)
+
+        self.robot_cmds.jnts_state.set(data=homing_full, data_type="q", 
+                            robot_idxs=self.controller_index_np,
+                            no_remap=True)
+        self.robot_cmds.jnts_state.set(data=null_action, data_type="v", 
+                            robot_idxs=self.controller_index_np,
+                            no_remap=True)
+        self.robot_cmds.jnts_state.set(data=null_action, data_type="eff", 
+                            robot_idxs=self.controller_index_np,
+                            no_remap=True)
+        
+        # write all joints (including homing for env-only ones)
+        self.robot_cmds.jnts_state.synch_retry(row_index=self.controller_index, col_index=0,
+                                row_index_view=0,
+                                n_rows=1, n_cols=self.robot_cmds.jnts_state.n_cols,
+                                read=False) # only write data corresponding to this controller
+        
+        self.reset() # reset controller
+        self._n_resets=0
+        
+        # for last we create the trigger client
+        self._remote_triggerer = RemoteTriggererClnt(namespace=self.namespace,
+                                        verbose=self._verbose,
+                                        vlevel=VLevel.V2) # remote triggering
+        self._remote_triggerer.run()
+
+        if self._debug:
+            # internal solution is published on shared mem
+            # we assume the user has made available the cost
+            # and constraint data at this point (e.g. through
+            # the solution of a bootstrap)
+            cost_data = self._get_cost_info()
+            constr_data = self._get_constr_info()
+            config = RhcInternal.Config(is_server=True, 
+                        enable_q= True, 
+                        enable_v=True, 
+                        enable_a=True, 
+                        enable_a_dot=False, 
+                        enable_f=True,
+                        enable_f_dot=False, 
+                        enable_eff=False, 
+                        cost_names=cost_data[0], 
+                        cost_dims=cost_data[1],
+                        constr_names=constr_data[0],
+                        constr_dims=constr_data[1],
+                        )
+            self.rhc_internal = RhcInternal(config=config, 
+                                    namespace=self.namespace,
+                                    rhc_index = self.controller_index,
+                                    n_contacts=self.n_contacts,
+                                    n_jnts=self.n_dofs,
+                                    jnt_names=self._controller_side_jnt_names,
+                                    n_nodes=self._n_nodes,
+                                    verbose = self._verbose,
+                                    vlevel=VLevel.V2,
+                                    force_reconnection=True,
+                                    safe=True)
+            self.rhc_internal.run()
+
+        self._class_name = self._class_name + f"-{self.controller_index}"
+
+        Journal.log(self._class_name,
+                    "_register_to_cluster",
+                    f"controller registered",
+                    LogType.STAT,
+                    throw_when_excep = True)
+        
+        # we can now release everything so that other controllers can register
+        self.rhc_status.controllers_counter.data_sem_release()
+        self.rhc_status.registration.data_sem_release()
+
+        self._registered = True
+
+    def _unregister_from_cluster(self):
+        
+        if self._received_trigger:
+            # received interrupt during solution --> 
+            # send ack signal to server anyway
+            self._remote_triggerer.ack() 
+        if self._registered:
+            # acquire semaphores since we have to perform operations
+            # on the whole memory views
+            self.rhc_status.registration.data_sem_acquire()
+            self.rhc_status.controllers_counter.data_sem_acquire()
+            self.rhc_status.registration.write_retry(False, 
+                                    row_index=self.controller_index,
+                                    col_index=0,
+                                    row_index_view=0)
+            self._deactivate()
+            # decrementing controllers counter
+            self.rhc_status.controllers_counter.synch_all(retry = True,
+                                                    read = True)
+            controllers_counter = self.rhc_status.controllers_counter.get_numpy_mirror()
+            controllers_counter -= 1 
+            self.rhc_status.controllers_counter.synch_all(retry = True,
+                                                    read = False)
+            Journal.log(self._class_name,
+                    "_unregister_from_cluster",
+                    "Done",
+                    LogType.STAT,
+                    throw_when_excep = True)
+            # we can now release everything
+            self.rhc_status.registration.data_sem_release()
+            self.rhc_status.controllers_counter.data_sem_release()
+            self._registered = False
+
+    def _get_quat_remap(self):
+        # to be overridden by child class if necessary
+        return [0, 1, 2, 3]
+    
+    def _consinstency_checks(self):
+        
+        # check controller dt
+        server_side_cluster_dt = self.cluster_stats.get_info(info_name="cluster_dt")
+        if not (abs(server_side_cluster_dt - self._dt) < 1e-4):
+            exception = f"Trying to initialize a controller with control dt {self._dt}, which" + \
+                f"does not match the cluster control dt {server_side_cluster_dt}"
+            Journal.log(self._class_name,
+                        "_consinstency_checks",
+                        exception,
+                        LogType.EXCEP,
+                        throw_when_excep = True)
+        # check contact names
+        
+        server_side_contact_names = set(self.robot_state.contact_names())
+        control_side_contact_names = set(self._get_contacts())
+
+        if (not server_side_contact_names == control_side_contact_names) and self._verbose:
+            warn = f"Controller-side contact names do not match server-side names!" + \
+                f"\nServer: {self.robot_state.contact_names()}\n Controller: {self._get_contacts()}"
+            Journal.log(self._class_name,
+                        "_consinstency_checks",
+                        warn,
+                        LogType.WARN,
+                        throw_when_excep = True)
+        if not len(self.robot_state.contact_names()) == len(self._get_contacts()):
+            # at least, we need the n of contacts to match!
+            exception = f"Controller-side n contacts {self._get_contacts()} do not match " + \
+                f"server-side n contacts {len(self.robot_state.contact_names())}!"
+            Journal.log(self._class_name,
+                        "_consinstency_checks",
+                        exception,
+                        LogType.EXCEP,
+                        throw_when_excep = True)
+            
+    def _init(self):
+
+        stat = f"Trying to initialize RHC controller " + \
+            f"with dt: {self._dt} s, t_horizon: {self._t_horizon} s, n_intervals: {self._n_intervals}"
+        Journal.log(self._class_name,
+                    "_init",
+                    stat,
+                    LogType.STAT,
+                    throw_when_excep = True)
+        
+        self._init_states() # initializes shared mem. states 
+
+        self._init_problem() # we call the child's initialization method for the actual problem
+        self._post_problem_init()
+
+        self._register_to_cluster() # registers the controller to the cluster
+
+        Journal.log(self._class_name,
+                    "_init",
+                    f"RHC controller initialized with cluster index {self.controller_index} on process {os.getpid()}",
+                    LogType.STAT,
+                    throw_when_excep = True)
+
+    def _deactivate(self):
+        # signal controller deactivation over shared mem
+        self.rhc_status.activation_state.write_retry(False, 
+                                row_index=self.controller_index,
+                                col_index=0,
+                                row_index_view=0)
+        # also set cmds to homing for safety
+        # self.reset_rhc_data()
+
+    def _on_failure(self):
+        
+        self.rhc_status.fails.write_retry(True, 
+                                    row_index=self.controller_index,
+                                    col_index=0,
+                                    row_index_view=0)
+        self._deactivate()
+        self._n_fails += 1
+        self.rhc_status.controllers_fail_counter.write_retry(self._n_fails,
+                                    row_index=self.controller_index,
+                                    col_index=0,
+                                    row_index_view=0)
+
+    def _init_robot_homer(self):
+        self._homer = RobotHomer(srdf_path=self.srdf_path, 
+                            jnt_names=self._controller_side_jnt_names,
+                            verbose=self._verbose)
+        
+        self._homer_env = RobotHomer(srdf_path=self.srdf_path, 
+                            jnt_names=self.robot_state.jnt_names(),
+                            verbose=self._verbose)
+        
+    def _update_profiling_data(self):
+
+        # updated debug data on shared memory
+        # with the latest info available
+        self.cluster_stats.solve_loop_dt.write_retry(self._profiling_data_dict["full_solve_dt"], 
+                                                            row_index=self.controller_index,
+                                                            col_index=0,
+                                                            row_index_view=0)
+        self.cluster_stats.rti_sol_time.write_retry(self._profiling_data_dict["rti_solve_dt"], 
+                                                            row_index=self.controller_index,
+                                                            col_index=0,
+                                                            row_index_view=0)
+        self.cluster_stats.prb_update_dt.write_retry(self._profiling_data_dict["problem_update_dt"], 
+                                                            row_index=self.controller_index,
+                                                            col_index=0,
+                                                            row_index_view=0)
+        self.cluster_stats.phase_shift_dt.write_retry(self._profiling_data_dict["phases_shift_dt"], 
+                                                            row_index=self.controller_index,
+                                                            col_index=0,
+                                                            row_index_view=0)
+        self.cluster_stats.task_ref_update_dt.write_retry(self._profiling_data_dict["task_ref_update"], 
+                                                            row_index=self.controller_index,
+                                                            col_index=0,
+                                                            row_index_view=0)
+    
+    def _write_cmds_from_sol(self):
+
+        # gets data from the solution and updates the view on the shared data
+
+        # cmds for robot
+        # jnts
+        self.robot_cmds.jnts_state.set(data=self._get_jnt_q_from_sol(node_idx=self._rhc_cmds_node_idx), data_type="q", robot_idxs=self.controller_index_np)
+        self.robot_cmds.jnts_state.set(data=self._get_jnt_v_from_sol(node_idx=self._rhc_cmds_node_idx), data_type="v", robot_idxs=self.controller_index_np)
+        self.robot_cmds.jnts_state.set(data=self._get_jnt_a_from_sol(node_idx=self._rhc_cmds_node_idx-1), data_type="a", robot_idxs=self.controller_index_np)
+        self.robot_cmds.jnts_state.set(data=self._get_jnt_eff_from_sol(node_idx=self._rhc_cmds_node_idx-1), data_type="eff", robot_idxs=self.controller_index_np)
+        # root
+        self.robot_cmds.root_state.set(data=self._get_root_full_q_from_sol(node_idx=self._rhc_cmds_node_idx), data_type="q_full", robot_idxs=self.controller_index_np)
+        self.robot_cmds.root_state.set(data=self._get_root_twist_from_sol(node_idx=self._rhc_cmds_node_idx), data_type="twist", robot_idxs=self.controller_index_np)
+        self.robot_cmds.root_state.set(data=self._get_root_a_from_sol(node_idx=self._rhc_cmds_node_idx-1), data_type="a_full", robot_idxs=self.controller_index_np)
+        self.robot_cmds.root_state.set(data=self._get_norm_grav_vector_from_sol(node_idx=self._rhc_cmds_node_idx-1), data_type="gn", robot_idxs=self.controller_index_np)
+        f_contact = self._get_f_from_sol()
+        if f_contact is not None:
+            contact_names = self.robot_state.contact_names()
+            node_idx_f_estimate=self._rhc_cmds_node_idx-1 # we always write the force to reach the desired state (prev node) 
+            rhc_q_estimate=self._get_root_full_q_from_sol(node_idx=node_idx_f_estimate)[:, 3:7]
+            for i in range(len(contact_names)):
+                contact = contact_names[i]
+                contact_idx = i*3
+                contact_force_rhc_world=f_contact[contact_idx:(contact_idx+3), node_idx_f_estimate:node_idx_f_estimate+1].T
+                world2base_frame(v_w=contact_force_rhc_world, 
+                    q_b=rhc_q_estimate, 
+                    v_out=self._contact_force_base_loc_aux,
+                    is_q_wijk=False # horizon q is ijkw
+                    )
+                self.robot_cmds.contact_wrenches.set(data=self._contact_force_base_loc_aux, 
+                    data_type="f", 
+                    robot_idxs=self.controller_index_np,
+                    contact_name=contact)
+        
+        # prediction data from MPC horizon
+        self.robot_pred.jnts_state.set(data=self._get_jnt_q_from_sol(node_idx=self._pred_node_idx), data_type="q", robot_idxs=self.controller_index_np)
+        self.robot_pred.jnts_state.set(data=self._get_jnt_v_from_sol(node_idx=self._pred_node_idx), data_type="v", robot_idxs=self.controller_index_np)
+        self.robot_pred.jnts_state.set(data=self._get_jnt_a_from_sol(node_idx=self._pred_node_idx-1), data_type="a", robot_idxs=self.controller_index_np)
+        self.robot_pred.jnts_state.set(data=self._get_jnt_eff_from_sol(node_idx=self._pred_node_idx-1), data_type="eff", robot_idxs=self.controller_index_np)
+        self.robot_pred.root_state.set(data=self._get_root_full_q_from_sol(node_idx=self._pred_node_idx), data_type="q_full", robot_idxs=self.controller_index_np)
+        self.robot_pred.root_state.set(data=self._get_root_twist_from_sol(node_idx=self._pred_node_idx), data_type="twist", robot_idxs=self.controller_index_np)
+        self.robot_pred.root_state.set(data=self._get_root_a_from_sol(node_idx=self._pred_node_idx-1), data_type="a_full", robot_idxs=self.controller_index_np)
+        self.robot_pred.root_state.set(data=self._get_norm_grav_vector_from_sol(node_idx=self._pred_node_idx-1), data_type="gn", robot_idxs=self.controller_index_np)
+
+        # write robot commands
+        self.robot_cmds.jnts_state.synch_retry(row_index=self.controller_index, col_index=0, 
+                                row_index_view=0,
+                                n_rows=1, n_cols=self.robot_cmds.jnts_state.n_cols,
+                                read=False) # jnt state
+        self.robot_cmds.root_state.synch_retry(row_index=self.controller_index, col_index=0, 
+                                row_index_view=0,
+                                n_rows=1, n_cols=self.robot_cmds.root_state.n_cols,
+                                read=False) # root state, in case it was written
+        self.robot_cmds.contact_wrenches.synch_retry(row_index=self.controller_index, col_index=0, 
+                                row_index_view=0,
+                                n_rows=1, n_cols=self.robot_cmds.contact_wrenches.n_cols,
+                                read=False) # contact state
+        
+        # write robot pred
+        self.robot_pred.jnts_state.synch_retry(row_index=self.controller_index, col_index=0, 
+                                row_index_view=0,
+                                n_rows=1, n_cols=self.robot_cmds.jnts_state.n_cols,
+                                read=False)
+        self.robot_pred.root_state.synch_retry(row_index=self.controller_index, col_index=0, 
+                                row_index_view=0,
+                                n_rows=1, n_cols=self.robot_cmds.root_state.n_cols,
+                                read=False)
+        
+        # we also fill other data (cost, constr. violation, etc..)
+        self.rhc_status.rhc_cost.write_retry(self._get_rhc_cost(), 
+                                    row_index=self.controller_index,
+                                    col_index=0,
+                                    row_index_view=0)
+        self.rhc_status.rhc_constr_viol.write_retry(self._get_rhc_constr_viol(), 
+                                    row_index=self.controller_index,
+                                    col_index=0,
+                                    row_index_view=0)
+        self.rhc_status.rhc_n_iter.write_retry(self._get_rhc_niter_to_sol(), 
+                                    row_index=self.controller_index,
+                                    col_index=0,
+                                    row_index_view=0)
+        self.rhc_status.rhc_nodes_cost.write_retry(data=self._get_rhc_nodes_cost(), 
+                                    row_index=self.controller_index, 
+                                    col_index=0,
+                                    row_index_view=0)
+        self.rhc_status.rhc_nodes_constr_viol.write_retry(data=self._get_rhc_nodes_constr_viol(), 
+                                    row_index=self.controller_index, 
+                                    col_index=0,
+                                    row_index_view=0)
+        self.rhc_status.rhc_fail_idx.write_retry(self._get_failure_index(), 
+                                    row_index=self.controller_index,
+                                    col_index=0,
+                                    row_index_view=0) # write idx  on shared mem
+
+    def _compute_pred_delta(self):
+        
+        # measurements
+        q_full_root_meas = self.robot_state.root_state.get(data_type="q_full", robot_idxs=self.controller_index_np)
+        twist_root_meas = self.robot_state.root_state.get(data_type="twist", robot_idxs=self.controller_index_np)
+        a_root_meas = self.robot_state.root_state.get(data_type="a_full", robot_idxs=self.controller_index_np)
+        g_vec_root_meas = self.robot_state.root_state.get(data_type="gn", robot_idxs=self.controller_index_np)
+
+        q_jnts_meas = self.robot_state.jnts_state.get(data_type="q", robot_idxs=self.controller_index_np)
+        v_jnts_meas = self.robot_state.jnts_state.get(data_type="v", robot_idxs=self.controller_index_np)
+        a_jnts_meas = self.robot_state.jnts_state.get(data_type="a", robot_idxs=self.controller_index_np)
+        eff_jnts_meas = self.robot_state.jnts_state.get(data_type="eff", robot_idxs=self.controller_index_np)
+
+        # prediction from rhc 
+        delta_root_q_full=self._get_root_full_q_from_sol(node_idx=1)-q_full_root_meas
+        delta_root_twist=self._get_root_twist_from_sol(node_idx=1)-twist_root_meas
+        delta_root_a=self._get_root_a_from_sol(node_idx=0)-a_root_meas
+        delta_g_vec=self._get_norm_grav_vector_from_sol(node_idx=0)-g_vec_root_meas
+
+        delta_jnts_q=self._get_jnt_q_from_sol(node_idx=1)-q_jnts_meas
+        delta_jnts_v=self._get_jnt_v_from_sol(node_idx=1)-v_jnts_meas
+        delta_jnts_a=self._get_jnt_a_from_sol(node_idx=0)-a_jnts_meas
+        delta_jnts_eff=self._get_jnt_eff_from_sol(node_idx=0)-eff_jnts_meas
+
+        # writing pred. errors
+        self.rhc_pred_delta.root_state.set(data=delta_root_q_full, data_type="q_full", robot_idxs=self.controller_index_np)
+        self.rhc_pred_delta.root_state.set(data=delta_root_twist,data_type="twist", robot_idxs=self.controller_index_np)
+        self.rhc_pred_delta.root_state.set(data=delta_root_a,data_type="a_full", robot_idxs=self.controller_index_np)
+        self.rhc_pred_delta.root_state.set(data=delta_g_vec, data_type="gn", robot_idxs=self.controller_index_np)
+
+        self.rhc_pred_delta.jnts_state.set(data=delta_jnts_q,data_type="q", robot_idxs=self.controller_index_np)
+        self.rhc_pred_delta.jnts_state.set(data=delta_jnts_v,data_type="v", robot_idxs=self.controller_index_np)
+        self.rhc_pred_delta.jnts_state.set(data=delta_jnts_a,data_type="a", robot_idxs=self.controller_index_np)
+        self.rhc_pred_delta.jnts_state.set(data=delta_jnts_eff, data_type="eff", robot_idxs=self.controller_index_np)
+
+        # write on shared memory
+        self.rhc_pred_delta.jnts_state.synch_retry(row_index=self.controller_index, 
+                                                   col_index=0, 
+                                                   n_rows=1, 
+                                                   row_index_view=0,
+                                                   n_cols=self.robot_cmds.jnts_state.n_cols,
+                                read=False) # jnt state
+        self.rhc_pred_delta.root_state.synch_retry(row_index=self.controller_index, 
+                                                    col_index=0,
+                                                    n_rows=1, 
+                                                    row_index_view=0,
+                                                    n_cols=self.robot_cmds.root_state.n_cols,
+                                read=False) # root state
+    
+    def _assign_controller_side_jnt_names(self, 
+                        jnt_names: List[str]):
+
+        self._controller_side_jnt_names = jnt_names
+        self._got_jnt_names_from_controllers = True
+
+    def _check_jnt_names_compatibility(self):
+
+        set_rhc = set(self._controller_side_jnt_names)
+        set_env  = set(self._env_side_jnt_names)
+        
+        if not set_rhc == set_env:
+            rhc_is_missing=set_env-set_rhc
+            env_is_missing=set_rhc-set_env
+
+            msg_type=LogType.WARN
+            message=""
+            if not len(rhc_is_missing)==0: # allowed
+                message = "\nSome env-side joint names are missing on rhc client-side!\n" + \
+                "RHC-SERVER-SIDE-> \n" + \
+                " ".join(self._env_side_jnt_names) + "\n" +\
+                "RHC-CLIENT-SIDE -> \n" + \
+                " ".join(self._controller_side_jnt_names) + "\n" \
+                "\MISSING -> \n" + \
+                " ".join(list(rhc_is_missing)) + "\n"
+                if not self._allow_less_jnts: # raise exception
+                    msg_type=LogType.EXCEP
+
+            if not len(env_is_missing)==0: # not allowed
+                message = "\nSome rhc-side joint names are missing on rhc server-side!\n" + \
+                "RHC-SERVER-SIDE-> \n" + \
+                " ".join(self._env_side_jnt_names) + \
+                "RHC-CLIENT-SIDE -> \n" + \
+                " ".join(self._controller_side_jnt_names) + "\n" \
+                "\nmissing -> \n" + \
+                " ".join(list(env_is_missing))
+                msg_type=LogType.EXCEP
+            
+            if msg_type==LogType.WARN and not self._verbose:
+                return
+            
+            Journal.log(self._class_name,
+                    "_check_jnt_names_compatibility",
+                    message,
+                    msg_type,
+                    throw_when_excep = True)
+    
+    def _get_cost_info(self):
+        # to be overridden by child class
+        return None, None
+    
+    def _get_constr_info(self):
+        # to be overridden by child class
+        return None, None
+    
+    def _get_fail_idx(self):
+        # to be overriden by parent
+        return 0.0
+    
+    def _get_failure_index(self):
+        fail_idx=self._get_fail_idx()/self._fail_idx_thresh
+        if fail_idx>1.0:
+            fail_idx=1.0
+        return fail_idx
+    
+    def _check_rhc_failure(self):
+        # we use fail idx viol to detect failures
+        idx = self._get_failure_index()
+        return idx>=1.0
+    
+    def _update_rhc_internal(self):
+        # data which is not enabled in the config is not actually 
+        # written so overhead is minimal for non-enabled data
+        self.rhc_internal.write_q(data= self._get_q_from_sol(),
+                            retry=True)
+        self.rhc_internal.write_v(data= self._get_v_from_sol(),
+                            retry=True)
+        self.rhc_internal.write_a(data= self._get_a_from_sol(),
+                            retry=True)
+        self.rhc_internal.write_a_dot(data= self._get_a_dot_from_sol(),
+                            retry=True)
+        self.rhc_internal.write_f(data= self._get_f_from_sol(),
+                            retry=True)
+        self.rhc_internal.write_f_dot(data= self._get_f_dot_from_sol(),
+                            retry=True)
+        self.rhc_internal.write_eff(data= self._get_eff_from_sol(),
+                            retry=True)
+        for cost_idx in range(self.rhc_internal.config.n_costs):
+            # iterate over all costs and update all values
+            cost_name = self.rhc_internal.config.cost_names[cost_idx]
+            self.rhc_internal.write_cost(data= self._get_cost_from_sol(cost_name = cost_name),
+                                cost_name = cost_name,
+                                retry=True)
+        for constr_idx in range(self.rhc_internal.config.n_constr):
+            # iterate over all constraints and update all values
+            constr_name = self.rhc_internal.config.constr_names[constr_idx]
+            self.rhc_internal.write_constr(data= self._get_constr_from_sol(constr_name=constr_name),
+                                constr_name = constr_name,
+                                retry=True)
+    
+    def _get_contacts(self): 
+        contact_names = self._get_contact_names()
+        self._got_contact_names = True
+        return contact_names
+    
+    def _get_q_from_sol(self):
+        # to be overridden by child class
+        return None
+
+    def _get_v_from_sol(self):
+        # to be overridden by child class
+        return None
+    
+    def _get_a_from_sol(self):
+        # to be overridden by child class
+        return None
+    
+    def _get_a_dot_from_sol(self):
+        # to be overridden by child class
+        return None
+    
+    def _get_f_from_sol(self):
+        # to be overridden by child class
+        return None
+    
+    def _get_f_dot_from_sol(self):
+        # to be overridden by child class
+        return None
+    
+    def _get_eff_from_sol(self):
+        # to be overridden by child class
+        return None
+    
+    def _get_cost_from_sol(self,
+                    cost_name: str):
+        # to be overridden by child class
+        return None
+    
+    def _get_constr_from_sol(self,
+                    constr_name: str):
+        # to be overridden by child class
+        return None
+    
+    @abstractmethod
+    def _reset(self):   
+        pass
+
+    @abstractmethod
+    def _init_rhc_task_cmds(self):
+        pass
+
+    @abstractmethod
+    def _get_robot_jnt_names(self):
+        pass
+    
+    @abstractmethod
+    def _get_contact_names(self):
+        pass
+
+    @abstractmethod
+    def _get_jnt_q_from_sol(self, node_idx=1) -> np.ndarray:
+        pass
+
+    @abstractmethod
+    def _get_jnt_v_from_sol(self, node_idx=1) -> np.ndarray:
+        pass
+    
+    @abstractmethod
+    def _get_jnt_a_from_sol(self, node_idx=0) -> np.ndarray:
+        pass
+
+    @abstractmethod
+    def _get_jnt_eff_from_sol(self, node_idx=0) -> np.ndarray:
+        pass
+    
+    @abstractmethod
+    def _get_root_full_q_from_sol(self, node_idx=1) -> np.ndarray:
+        pass
+    
+    @abstractmethod
+    def _get_full_q_from_sol(self, node_idx=1) -> np.ndarray:
+        pass
+
+    @abstractmethod
+    def _get_root_twist_from_sol(self, node_idx=1) -> np.ndarray:
+        pass
+    
+    @abstractmethod
+    def _get_root_a_from_sol(self, node_idx=0) -> np.ndarray:
+        pass
+
+    def _get_norm_grav_vector_from_sol(self, node_idx=1) -> np.ndarray:
+        rhc_q=self._get_root_full_q_from_sol(node_idx=node_idx)[:, 3:7]
+        world2base_frame(v_w=self._norm_grav_vector_w,q_b=rhc_q,v_out=self._norm_grav_vector_base_loc,
+            is_q_wijk=False)
+        return self._norm_grav_vector_base_loc
+    
+    def _get_rhc_cost(self):
+        # to be overridden
+        return np.nan
+    
+    def _get_rhc_constr_viol(self):
+        # to be overridden
+        return np.nan
+
+    def _get_rhc_nodes_cost(self):
+        # to be overridden
+        return np.zeros((1,self.rhc_status.n_nodes), dtype=self._dtype)
+    
+    def _get_rhc_nodes_constr_viol(self):
+        # to be overridden
+        return np.zeros((1,self.rhc_status.n_nodes), dtype=self._dtype)
+    
+    def _get_rhc_niter_to_sol(self) -> np.ndarray:
+        # to be overridden
+        return np.nan
+    
+    @abstractmethod
+    def _update_open_loop(self):
+        # updates rhc controller 
+        # using the internal state 
+        pass
+    
+    @abstractmethod
+    def _update_closed_loop(self):
+        # uses meas. from robot
+        pass
+
+    @abstractmethod
+    def _solve(self) -> bool:
+        pass
+
+    @abstractmethod
+    def _bootstrap(self) -> bool:
+        pass
+            
+    @abstractmethod
+    def _get_ndofs(self):
+        pass
+    
+    abstractmethod
+    def _get_robot_mass(self):
+        pass
+
+    @abstractmethod
+    def _init_problem(self):
+        # initialized horizon's TO problem
+        pass
+
+    @abstractmethod
+    def _post_problem_init(self):
+        pass

sac.py

@@ -0,0 +1,680 @@
+import torch 
+import torch.nn as nn
+from torch.distributions.normal import Normal
+import math
+
+from aug_mpc.utils.nn.normalization_utils import RunningNormalizer 
+from aug_mpc.utils.nn.layer_utils import llayer_init 
+
+from typing import List
+
+from EigenIPC.PyEigenIPC import LogType
+from EigenIPC.PyEigenIPC import Journal
+from EigenIPC.PyEigenIPC import VLevel
+
+class SACAgent(nn.Module):
+    def __init__(self,
+            obs_dim: int, 
+            actions_dim: int,
+            obs_ub: List[float] = None,
+            obs_lb: List[float] = None,
+            actions_ub: List[float] = None,
+            actions_lb: List[float] = None,
+            rescale_obs: bool = False,
+            norm_obs: bool = True,
+            use_action_rescale_for_critic: bool = True,
+            device:str="cuda",
+            dtype=torch.float32,
+            is_eval:bool=False,
+            load_qf:bool=False,
+            epsilon:float=1e-8,
+            debug:bool=False,
+            compression_ratio:float=-1.0, # > 0; if [0, 1] compression, >1 "expansion"
+            layer_width_actor:int=256,
+            n_hidden_layers_actor:int=2,
+            layer_width_critic:int=512,
+            n_hidden_layers_critic:int=4,
+            torch_compile: bool = False,
+            add_weight_norm: bool = False,
+            add_layer_norm: bool = False,
+            add_batch_norm: bool = False):
+
+        super().__init__()
+
+        self._use_torch_compile=torch_compile
+
+        self._layer_width_actor=layer_width_actor
+        self._layer_width_critic=layer_width_critic
+        self._n_hidden_layers_actor=n_hidden_layers_actor
+        self._n_hidden_layers_critic=n_hidden_layers_critic
+
+        self._obs_dim=obs_dim
+        self._actions_dim=actions_dim
+        self._actions_ub=actions_ub
+        self._actions_lb=actions_lb
+
+        self._add_weight_norm=add_weight_norm
+        self._add_layer_norm=add_layer_norm
+        self._add_batch_norm=add_batch_norm
+
+        self._is_eval=is_eval
+        self._load_qf=load_qf
+
+        self._epsilon=epsilon
+
+        if compression_ratio > 0.0:
+            self._layer_width_actor=int(compression_ratio*obs_dim)
+            self._layer_width_critic=int(compression_ratio*(obs_dim+actions_dim))
+        
+        self._normalize_obs = norm_obs
+        self._rescale_obs=rescale_obs
+        if self._rescale_obs and self._normalize_obs:
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                f"Both running normalization and obs rescaling is enabled. Was this intentional?",
+                LogType.WARN,
+                throw_when_excep = True)
+        
+        self._use_action_rescale_for_critic=use_action_rescale_for_critic
+
+        self._rescaling_epsi=1e-9
+
+        self._debug = debug
+
+        self._torch_device = device
+        self._torch_dtype = dtype
+
+        # obs scale and bias
+        if obs_ub is None:
+            obs_ub = [1] * obs_dim
+        if obs_lb is None:
+            obs_lb = [-1] * obs_dim
+        if (len(obs_ub) != obs_dim):
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                f"Observations ub list length should be equal to {obs_dim}, but got {len(obs_ub)}",
+                LogType.EXCEP,
+                throw_when_excep = True)
+        if (len(obs_lb) != obs_dim):
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                f"Observations lb list length should be equal to {obs_dim}, but got {len(obs_lb)}",
+                LogType.EXCEP,
+                throw_when_excep = True)
+
+        self._obs_ub = torch.tensor(obs_ub, dtype=self._torch_dtype, 
+                                device=self._torch_device)
+        self._obs_lb = torch.tensor(obs_lb, dtype=self._torch_dtype,
+                                device=self._torch_device)
+        obs_scale = torch.full((obs_dim, ),
+                            fill_value=0.0,
+                            dtype=self._torch_dtype,
+                            device=self._torch_device)
+        obs_scale[:] = (self._obs_ub-self._obs_lb)/2.0
+        self.register_buffer(
+            "obs_scale", obs_scale
+        )
+        obs_bias = torch.full((obs_dim, ),
+                            fill_value=0.0,
+                            dtype=self._torch_dtype,
+                            device=self._torch_device)
+        obs_bias[:] = (self._obs_ub+self._obs_lb)/2.0
+        self.register_buffer(
+            "obs_bias", obs_bias)
+        
+        self._build_nets()
+        
+        self._init_obs_norm()
+
+        msg=f"Created SAC agent with actor [{self._layer_width_actor}, {self._n_hidden_layers_actor}]\
+        and critic [{self._layer_width_critic}, {self._n_hidden_layers_critic}] sizes.\
+        \n Runningobs normalizer: {type(self.obs_running_norm)} \
+        \n Batch normalization: {self._add_batch_norm} \
+        \n Layer normalization: {self._add_layer_norm} \
+        \n Weight normalization: {self._add_weight_norm} \
+        \n Critic input actions are descaled: {self._use_action_rescale_for_critic}"
+        Journal.log(self.__class__.__name__,
+            "__init__",
+            msg,
+            LogType.INFO)
+    
+    def _init_obs_norm(self):
+        
+        self.obs_running_norm=None
+        if self._normalize_obs:
+            self.obs_running_norm = RunningNormalizer((self._obs_dim,), 
+                                        epsilon=self._epsilon, 
+                                        device=self._torch_device, dtype=self._torch_dtype, 
+                                        freeze_stats=True, # always start with freezed stats
+                                        debug=self._debug)
+            self.obs_running_norm.type(self._torch_dtype) # ensuring correct dtype for whole module
+
+    def _build_nets(self):
+
+        if self._add_weight_norm:
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                f"Will use weight normalization reparametrization\n",
+                LogType.INFO)
+
+        self.actor=None
+        self.qf1=None
+        self.qf2=None
+        self.qf1_target=None
+        self.qf2_target=None
+        
+        self.actor = Actor(obs_dim=self._obs_dim,
+                    actions_dim=self._actions_dim,
+                    actions_ub=self._actions_ub,
+                    actions_lb=self._actions_lb,
+                    device=self._torch_device,
+                    dtype=self._torch_dtype,
+                    layer_width=self._layer_width_actor,
+                    n_hidden_layers=self._n_hidden_layers_actor,
+                    add_weight_norm=self._add_weight_norm,
+                    add_layer_norm=self._add_layer_norm,
+                    add_batch_norm=self._add_batch_norm,
+                    )
+
+        if (not self._is_eval) or self._load_qf: # just needed for training or during eval
+            # for debug, if enabled
+            self.qf1 = CriticQ(obs_dim=self._obs_dim,
+                    actions_dim=self._actions_dim,
+                    device=self._torch_device,
+                    dtype=self._torch_dtype,
+                    layer_width=self._layer_width_critic,
+                    n_hidden_layers=self._n_hidden_layers_critic,
+                    add_weight_norm=self._add_weight_norm,
+                    add_layer_norm=self._add_layer_norm,
+                    add_batch_norm=self._add_batch_norm)
+            self.qf1_target = CriticQ(obs_dim=self._obs_dim,
+                    actions_dim=self._actions_dim,
+                    device=self._torch_device,
+                    dtype=self._torch_dtype,
+                    layer_width=self._layer_width_critic,
+                    n_hidden_layers=self._n_hidden_layers_critic,
+                    add_weight_norm=self._add_weight_norm,
+                    add_layer_norm=self._add_layer_norm,
+                    add_batch_norm=self._add_batch_norm)
+            
+            self.qf2 = CriticQ(obs_dim=self._obs_dim,
+                    actions_dim=self._actions_dim,
+                    device=self._torch_device,
+                    dtype=self._torch_dtype,
+                    layer_width=self._layer_width_critic,
+                    n_hidden_layers=self._n_hidden_layers_critic,
+                    add_weight_norm=self._add_weight_norm,
+                    add_layer_norm=self._add_layer_norm,
+                    add_batch_norm=self._add_batch_norm)
+            self.qf2_target = CriticQ(obs_dim=self._obs_dim,
+                    actions_dim=self._actions_dim,
+                    device=self._torch_device,
+                    dtype=self._torch_dtype,
+                    layer_width=self._layer_width_critic,
+                    n_hidden_layers=self._n_hidden_layers_critic,
+                    add_weight_norm=self._add_weight_norm,
+                    add_layer_norm=self._add_layer_norm,
+                    add_batch_norm=self._add_batch_norm)
+        
+            self.qf1_target.load_state_dict(self.qf1.state_dict())
+            self.qf2_target.load_state_dict(self.qf2.state_dict())
+
+        if self._use_torch_compile:
+            self.obs_running_norm=torch.compile(self.obs_running_norm)
+            self.actor = torch.compile(self.actor)
+            if (not self._is_eval) or self._load_qf:
+                self.qf1 = torch.compile(self.qf1)
+                self.qf2 = torch.compile(self.qf2)
+                self.qf1_target = torch.compile(self.qf1_target)
+                self.qf2_target = torch.compile(self.qf2_target)
+            
+    def reset(self, reset_stats: bool = False):
+        # we should just reinitialize the parameters, but for easiness
+        # we recreate the networks
+
+        # force deallocation of objects
+        import gc
+        del self.actor
+        del self.qf1
+        del self.qf2
+        del self.qf1_target
+        del self.qf2_target
+        gc.collect()
+
+        self._build_nets()
+
+        if reset_stats: # we also reinitialize obs norm
+            self._init_obs_norm()
+
+        # self.obs_running_norm.reset()
+
+    def layer_width_actor(self):
+        return self._layer_width_actor
+
+    def n_hidden_layers_actor(self):
+        return self._n_hidden_layers_actor
+
+    def layer_width_critic(self):
+        return self._layer_width_critic
+
+    def n_hidden_layers_critic(self):
+        return self._n_hidden_layers_critic
+
+    def get_impl_path(self):
+        import os 
+        return os.path.abspath(__file__)
+    
+    def update_obs_bnorm(self, x):
+        self.obs_running_norm.unfreeze()
+        self.obs_running_norm.manual_stat_update(x)
+        self.obs_running_norm.freeze()
+
+    def _obs_scaling_layer(self, x):
+        x=(x-self.obs_bias)
+        x=x/(self.obs_scale+self._rescaling_epsi)
+        return x
+    
+    def _preprocess_obs(self, x):
+        if self._rescale_obs:
+            x = self._obs_scaling_layer(x)
+        if self.obs_running_norm is not None:
+            x = self.obs_running_norm(x)
+        return x
+
+    def _preprocess_actions(self, a):
+        if self._use_action_rescale_for_critic:
+            a=self.actor.remove_scaling(a=a) # rescale to be in range [-1, 1]
+        return a
+
+    def get_action(self, x):
+        x = self._preprocess_obs(x)
+        return self.actor.get_action(x)
+    
+    def get_qf1_val(self, x, a):
+        x = self._preprocess_obs(x)
+        a = self._preprocess_actions(a)
+        return self.qf1(x, a)
+
+    def get_qf2_val(self, x, a):
+        x = self._preprocess_obs(x)
+        a = self._preprocess_actions(a)
+        return self.qf2(x, a)
+    
+    def get_qf1t_val(self, x, a):
+        x = self._preprocess_obs(x)
+        a = self._preprocess_actions(a)
+        return self.qf1_target(x, a)
+    
+    def get_qf2t_val(self, x, a):
+        x = self._preprocess_obs(x)
+        a = self._preprocess_actions(a)
+        return self.qf2_target(x, a)
+
+    def load_state_dict(self, param_dict):
+
+        missing, unexpected = super().load_state_dict(param_dict,
+            strict=False)
+        if not len(missing)==0:
+            Journal.log(self.__class__.__name__,
+                "load_state_dict",
+                f"These parameters are missing from the provided state dictionary: {str(missing)}\n",
+                LogType.EXCEP,
+                throw_when_excep = True)
+        if not len(unexpected)==0:
+            Journal.log(self.__class__.__name__,
+                "load_state_dict",
+                f"These parameters present in the provided state dictionary are not needed: {str(unexpected)}\n",
+                LogType.WARN)
+        
+        # sanity check on running normalizer
+        import re
+        running_norm_pattern = r"running_norm"
+        error=f"Found some keys in model state dictionary associated with a running normalizer. Are you running the agent with norm_obs=True?\n"
+        if any(re.match(running_norm_pattern, key) for key in unexpected):
+            Journal.log(self.__class__.__name__,
+                "load_state_dict",
+                error,
+                LogType.EXCEP,
+                throw_when_excep=True)
+
+class CriticQ(nn.Module):
+    def __init__(self,
+        obs_dim: int, 
+        actions_dim: int,
+        device: str = "cuda",
+        dtype = torch.float32,
+        layer_width: int = 512,
+        n_hidden_layers: int = 4,
+        add_weight_norm: bool = False,
+        add_layer_norm: bool = False,
+        add_batch_norm: bool = False):
+
+        super().__init__()
+
+        self._lrelu_slope=0.01
+
+        self._torch_device = device
+        self._torch_dtype = dtype
+
+        self._obs_dim = obs_dim
+        self._actions_dim = actions_dim
+        self._q_net_dim = self._obs_dim + self._actions_dim
+
+        self._first_hidden_layer_width=self._q_net_dim # fist layer fully connected and of same dim
+        # as input
+
+        init_type="kaiming_uniform" # maintains the variance of activations throughout the network
+        nonlinearity="leaky_relu" # suited for kaiming
+
+        # Input layer        
+        layers=llayer_init(
+            layer=nn.Linear(self._q_net_dim, self._first_hidden_layer_width),
+            init_type=init_type,
+            nonlinearity=nonlinearity,
+            a_leaky_relu=self._lrelu_slope,
+            device=self._torch_device,
+            dtype=self._torch_dtype,
+            add_weight_norm=add_weight_norm,
+            add_layer_norm=add_layer_norm,
+            add_batch_norm=add_batch_norm,
+            uniform_biases=False, # constant bias init
+            bias_const=0.0
+        )
+        layers.extend([nn.LeakyReLU(negative_slope=self._lrelu_slope)])
+        
+        # Hidden layers
+        layers.extend(
+            llayer_init(
+                layer=nn.Linear(self._first_hidden_layer_width, layer_width),
+                init_type=init_type,
+                nonlinearity=nonlinearity,
+                a_leaky_relu=self._lrelu_slope,
+                device=self._torch_device,
+                dtype=self._torch_dtype,
+                add_weight_norm=add_weight_norm,
+                add_layer_norm=add_layer_norm,
+                add_batch_norm=add_batch_norm,
+                uniform_biases=False, # constant bias init
+                bias_const=0.0
+            )
+        )
+        layers.extend([nn.LeakyReLU(negative_slope=self._lrelu_slope)])
+
+        for _ in range(n_hidden_layers - 1):
+            layers.extend(
+                llayer_init(
+                    layer=nn.Linear(layer_width, layer_width),
+                    init_type=init_type,
+                    nonlinearity=nonlinearity,
+                    a_leaky_relu=self._lrelu_slope,
+                    device=self._torch_device,
+                    dtype=self._torch_dtype,
+                    add_weight_norm=add_weight_norm,
+                    add_layer_norm=add_layer_norm,
+                    add_batch_norm=add_batch_norm,
+                    uniform_biases=False, # constant bias init
+                    bias_const=0.0
+                )
+            )
+            layers.extend([nn.LeakyReLU(negative_slope=self._lrelu_slope)])
+
+        # Output layer
+        layers.extend(
+            llayer_init(
+                layer=nn.Linear(layer_width, 1),
+                init_type="uniform",
+                uniform_biases=False, # contact biases
+                bias_const=-0.1, # negative to prevent overestimation
+                scale_weight=1e-2,
+                device=self._torch_device,
+                dtype=self._torch_dtype,
+                add_weight_norm=False,
+                add_layer_norm=False,
+                add_batch_norm=False
+            )
+        )
+
+        # Creating the full sequential network
+        self._q_net = nn.Sequential(*layers)
+        self._q_net.to(self._torch_device).type(self._torch_dtype)
+
+        print("Critic architecture")
+        print(self._q_net)
+
+    def get_n_params(self):
+        return sum(p.numel() for p in self.parameters())
+
+    def forward(self, x, a):
+        x = torch.cat([x, a], dim=1)
+        return self._q_net(x)
+        
+class Actor(nn.Module):
+    def __init__(self,
+        obs_dim: int, 
+        actions_dim: int,
+        actions_ub: List[float] = None,
+        actions_lb: List[float] = None,
+        device: str = "cuda",
+        dtype = torch.float32,
+        layer_width: int = 256,
+        n_hidden_layers: int = 2,
+        add_weight_norm: bool = False,
+        add_layer_norm: bool = False,
+        add_batch_norm: bool = False):
+    
+        super().__init__()
+
+        self._lrelu_slope=0.01
+        
+        self._torch_device = device
+        self._torch_dtype = dtype
+
+        self._obs_dim = obs_dim
+        self._actions_dim = actions_dim
+        
+        self._first_hidden_layer_width=self._obs_dim # fist layer fully connected and of same dim
+    
+        # Action scale and bias
+        if actions_ub is None:
+            actions_ub = [1] * actions_dim
+        if actions_lb is None:
+            actions_lb = [-1] * actions_dim
+        if (len(actions_ub) != actions_dim):
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                f"Actions ub list length should be equal to {actions_dim}, but got {len(actions_ub)}",
+                LogType.EXCEP,
+                throw_when_excep = True)
+        if (len(actions_lb) != actions_dim):
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                f"Actions lb list length should be equal to {actions_dim}, but got {len(actions_lb)}",
+                LogType.EXCEP,
+                throw_when_excep = True)
+
+        self._actions_ub = torch.tensor(actions_ub, dtype=self._torch_dtype, 
+                                device=self._torch_device)
+        self._actions_lb = torch.tensor(actions_lb, dtype=self._torch_dtype,
+                                device=self._torch_device)
+        action_scale = torch.full((actions_dim, ),
+                            fill_value=0.0,
+                            dtype=self._torch_dtype,
+                            device=self._torch_device)
+        action_scale[:] = (self._actions_ub-self._actions_lb)/2.0
+        self.register_buffer(
+            "action_scale", action_scale
+        )
+        actions_bias = torch.full((actions_dim, ),
+                            fill_value=0.0,
+                            dtype=self._torch_dtype,
+                            device=self._torch_device)
+        actions_bias[:] = (self._actions_ub+self._actions_lb)/2.0
+        self.register_buffer(
+            "action_bias", actions_bias)
+        
+        # Network configuration
+        self.LOG_STD_MAX = 2
+        self.LOG_STD_MIN = -5
+
+        # Input layer followed by hidden layers
+        layers=llayer_init(nn.Linear(self._obs_dim, self._first_hidden_layer_width), 
+                    init_type="kaiming_uniform",
+                    nonlinearity="leaky_relu",
+                    a_leaky_relu=self._lrelu_slope,
+                    device=self._torch_device, 
+                    dtype=self._torch_dtype,
+                    add_weight_norm=add_weight_norm,
+                    add_layer_norm=add_layer_norm,
+                    add_batch_norm=add_batch_norm,
+                    uniform_biases=False, # constant bias init
+                    bias_const=0.0
+                    )
+        layers.extend([nn.LeakyReLU(negative_slope=self._lrelu_slope)])
+    
+        # Hidden layers
+        layers.extend(
+            llayer_init(nn.Linear(self._first_hidden_layer_width, layer_width), 
+                init_type="kaiming_uniform",
+                nonlinearity="leaky_relu",
+                a_leaky_relu=self._lrelu_slope,
+                device=self._torch_device,
+                dtype=self._torch_dtype,
+                add_weight_norm=add_weight_norm,
+                add_layer_norm=add_layer_norm,
+                add_batch_norm=add_batch_norm,
+                uniform_biases=False, # constant bias init
+                bias_const=0.0)
+        )
+        layers.extend([nn.LeakyReLU(negative_slope=self._lrelu_slope)])
+        
+        for _ in range(n_hidden_layers - 1):
+            layers.extend(
+                llayer_init(nn.Linear(layer_width, layer_width), 
+                    init_type="kaiming_uniform",
+                    nonlinearity="leaky_relu",
+                    a_leaky_relu=self._lrelu_slope,
+                    device=self._torch_device,
+                    dtype=self._torch_dtype,
+                    add_weight_norm=add_weight_norm,
+                    add_layer_norm=add_layer_norm,
+                    add_batch_norm=add_batch_norm,
+                    uniform_biases=False, # constant bias init
+                    bias_const=0.0)            
+            )
+            layers.extend([nn.LeakyReLU(negative_slope=self._lrelu_slope)])
+        
+        # Sequential layers for the feature extractor
+        self._fc12 = nn.Sequential(*layers)
+
+        # Mean and log_std layers
+        out_fc_mean=llayer_init(nn.Linear(layer_width, self._actions_dim), 
+                        init_type="uniform",
+                        uniform_biases=False, # constant bias init
+                        bias_const=0.0,
+                        scale_weight=1e-3, # scaling (output layer)
+                        scale_bias=1.0,
+                        device=self._torch_device, 
+                        dtype=self._torch_dtype,
+                        add_weight_norm=False,
+                        add_layer_norm=False,
+                        add_batch_norm=False
+                        )
+        self.fc_mean = nn.Sequential(*out_fc_mean)
+        out_fc_logstd= llayer_init(nn.Linear(layer_width, self._actions_dim), 
+                        init_type="uniform",
+                        uniform_biases=False,
+                        bias_const=math.log(0.5),
+                        scale_weight=1e-3, # scaling (output layer)
+                        scale_bias=1.0,
+                        device=self._torch_device, 
+                        dtype=self._torch_dtype,
+                        add_weight_norm=False,
+                        add_layer_norm=False,
+                        add_batch_norm=False,
+                        )
+        self.fc_logstd = nn.Sequential(*out_fc_logstd)
+
+        # Move all components to the specified device and dtype
+        self._fc12.to(device=self._torch_device, dtype=self._torch_dtype)
+        self.fc_mean.to(device=self._torch_device, dtype=self._torch_dtype)
+        self.fc_logstd.to(device=self._torch_device, dtype=self._torch_dtype)
+
+        print("Actor architecture")
+        print(self._fc12)
+        print(self.fc_mean)
+        print(self.fc_logstd)
+
+    def get_n_params(self):
+        return sum(p.numel() for p in self.parameters())
+    
+    def forward(self, x):
+        x = self._fc12(x)
+        mean = self.fc_mean(x)
+        log_std = self.fc_logstd(x)
+        log_std = torch.tanh(log_std)
+        log_std = self.LOG_STD_MIN + 0.5 * (self.LOG_STD_MAX - self.LOG_STD_MIN) * (log_std + 1)  # From SpinUp / Denis Yarats
+        return mean, log_std
+
+    def get_action(self, x):
+        mean, log_std = self(x)
+        std = log_std.exp()
+        normal = torch.distributions.Normal(mean, std)
+        x_t = normal.rsample()  # Reparameterization trick (for SAC we neex action
+        # to be differentible since we use Q nets. Using sample() would break the
+        # comp. graph and not allow gradients to flow)
+        y_t = torch.tanh(x_t)
+        action = y_t * self.action_scale + self.action_bias
+        log_prob_vec = normal.log_prob(x_t) # per-dimension log prob before tanh
+        log_prob_vec = log_prob_vec - torch.log(self.action_scale * (1 - y_t.pow(2)) + 1e-6) # tanh Jacobian + scaling
+        log_prob_sum = log_prob_vec.sum(1, keepdim=True)
+        mean = torch.tanh(mean) * self.action_scale + self.action_bias
+        return action, (log_prob_sum, log_prob_vec), mean
+    
+    def remove_scaling(self, a):
+        return (a - self.action_bias)/self.action_scale
+
+if __name__ == "__main__":  
+    device = "cpu"  # or "cpu"
+    import time
+    obs_dim = 273
+    agent = SACAgent(
+        obs_dim=obs_dim,
+        actions_dim=10,
+        actions_lb=None,
+        actions_ub=None,
+        obs_lb=None,
+        obs_ub=None,
+        rescale_obs=False,
+        norm_obs=True,
+        use_action_rescale_for_critic=True,
+        is_eval=True,
+        compression_ratio=0.6,
+        layer_width_actor=128,
+        layer_width_critic=256,
+        n_hidden_layers_actor=3,
+        n_hidden_layers_critic=3,
+        device=device,
+        dtype=torch.float32,
+        add_weight_norm=True,
+        add_layer_norm=False,
+        add_batch_norm=False
+    )
+
+    n_samples = 10000
+    random_obs = torch.rand((1, obs_dim), dtype=torch.float32, device=device)
+
+    if device == "cuda":
+        torch.cuda.synchronize()
+    start = time.time()
+
+    for i in range(n_samples):
+        actions, _, mean = agent.get_action(x=random_obs)
+        actions = actions.detach()
+        actions[:, :] = mean.detach()
+
+    if device == "cuda":
+        torch.cuda.synchronize()
+    end = time.time()
+
+    avrg_eval_time = (end - start) / n_samples
+    print(f"Average policy evaluation time on {device}: {avrg_eval_time:.6f} s")

training_env_base.py

@@ -0,0 +1,2002 @@
+import torch
+import math
+from aug_mpc.utils.math_utils import quaternion_to_angular_velocity, quaternion_difference
+
+from mpc_hive.utilities.shared_data.rhc_data import RobotState
+from mpc_hive.utilities.shared_data.rhc_data import RhcCmds, RhcPred
+from mpc_hive.utilities.shared_data.rhc_data import RhcRefs
+from mpc_hive.utilities.shared_data.rhc_data import RhcStatus
+from mpc_hive.utilities.shared_data.sim_data import SharedEnvInfo
+
+from aug_mpc.utils.shared_data.remote_stepping import RemoteStepperSrvr
+from aug_mpc.utils.shared_data.remote_stepping import RemoteResetSrvr
+from aug_mpc.utils.shared_data.remote_stepping import RemoteResetRequest
+
+from aug_mpc.utils.shared_data.agent_refs import AgentRefs
+from aug_mpc.utils.shared_data.training_env import SharedTrainingEnvInfo
+
+from aug_mpc.utils.shared_data.training_env import Observations, NextObservations
+from aug_mpc.utils.shared_data.training_env import TotRewards
+from aug_mpc.utils.shared_data.training_env import SubRewards
+from aug_mpc.utils.shared_data.training_env import Actions
+from aug_mpc.utils.shared_data.training_env import Terminations, SubTerminations
+from aug_mpc.utils.shared_data.training_env import Truncations, SubTruncations
+from aug_mpc.utils.shared_data.training_env import EpisodesCounter,TaskRandCounter,SafetyRandResetsCounter,RandomTruncCounter,SubStepAbsCounter
+
+from aug_mpc.utils.episodic_rewards import EpisodicRewards
+from aug_mpc.utils.episodic_data import EpisodicData
+from aug_mpc.utils.episodic_data import MemBuffer
+from aug_mpc.utils.signal_smoother import ExponentialSignalSmoother
+from aug_mpc.utils.math_utils import check_capsize
+
+from mpc_hive.utilities.math_utils_torch import world2base_frame
+
+from EigenIPC.PyEigenIPC import VLevel
+from EigenIPC.PyEigenIPC import LogType
+from EigenIPC.PyEigenIPC import Journal
+from EigenIPC.PyEigenIPC import StringTensorClient
+
+from perf_sleep.pyperfsleep import PerfSleep
+
+from abc import abstractmethod, ABC
+
+import os
+from typing import List, Dict
+
+class AugMPCTrainingEnvBase(ABC):
+
+    """Base class for a remote training environment tailored to Learning-based Receding Horizon Control"""
+
+    def __init__(self,
+            namespace: str,
+            obs_dim: int,
+            actions_dim: int,
+            env_name: str = "",
+            verbose: bool = False,
+            vlevel: VLevel = VLevel.V1,
+            debug: bool = True,
+            use_gpu: bool = True,
+            dtype: torch.dtype = torch.float32,
+            override_agent_refs: bool = False,
+            timeout_ms: int = 60000,
+            env_opts: Dict = {}):
+
+        self._this_path = os.path.abspath(__file__)
+
+        self.custom_db_data = None
+
+        self._random_reset_active=False
+
+        self._action_smoother_continuous=None
+        self._action_smoother_discrete=None
+
+        self._closed = False
+        self._ready=False
+
+        self._namespace = namespace
+        self._with_gpu_mirror = True
+        self._safe_shared_mem = False
+        
+        self._obs_dim = obs_dim
+        self._actions_dim = actions_dim
+
+        self._use_gpu = use_gpu
+        if self._use_gpu:
+            self._device="cuda"
+        else:
+            self._device="cpu"
+
+        self._dtype = dtype
+
+        self._verbose = verbose
+        self._vlevel = vlevel
+
+        self._is_debug = debug
+
+        self._env_name = env_name
+
+        self._override_agent_refs = override_agent_refs
+
+        self._substep_dt=1.0 # dt [s] between each substep
+
+        self._env_opts={}
+        self._env_opts.update(env_opts)   
+        self._process_env_opts()
+
+        self._robot_state = None
+        self._rhc_cmds = None
+        self._rhc_pred = None
+        self._rhc_refs = None
+        self._rhc_status = None
+
+        self._remote_stepper = None
+        self._remote_resetter = None
+        self._remote_reset_req = None
+
+        self._agent_refs = None
+
+        self._n_envs = 0
+        
+        self._ep_timeout_counter = None
+        self._task_rand_counter = None
+        self._rand_safety_reset_counter = None
+        self._rand_trunc_counter = None
+
+        self._actions_map={} # to be used to hold info like action idxs
+        self._obs_map={}
+
+        self._obs = None
+        self._obs_ub = None
+        self._obs_lb = None
+        self._next_obs = None
+        self._actions = None
+        self._actual_actions = None
+        self._actions_ub = None
+        self._actions_lb = None
+        self._tot_rewards = None
+        self._sub_rewards = None
+        self._sub_terminations = None
+        self._sub_truncations = None
+        self._terminations = None
+        self._truncations = None
+        self._act_mem_buffer = None
+
+        self._episodic_rewards_metrics = None
+        
+        self._timeout = timeout_ms
+    
+        self._height_grid_size = None
+        self._height_flat_dim = 0
+
+        self._attach_to_shared_mem()
+
+        self._init_obs(obs_dim)
+        self._init_actions(actions_dim)
+        self._init_rewards()
+        self._init_terminations()
+        self._init_truncations()
+        self._init_custom_db_data()
+        
+        self._demo_setup() # setup for demo envs
+
+        # to ensure maps are properly initialized
+        _ = self._get_action_names()
+        _ = self._get_obs_names()
+        _ = self._get_sub_trunc_names()
+        _ = self._get_sub_term_names()
+
+        self._set_substep_rew()
+        self._set_substep_obs()
+
+        self._custom_post_init()
+
+        # update actions scale and offset in case it was modified in _custom_post_init
+        self._actions_scale = (self._actions_ub - self._actions_lb)/2.0
+        self._actions_offset = (self._actions_ub + self._actions_lb)/2.0
+
+        if self._env_opts["use_action_smoothing"]:
+            self._init_action_smoothing()
+
+        self._ready=self._init_step(reset_on_init=self._env_opts["reset_on_init"])
+
+    def _add_env_opt(self,
+        opts: Dict,
+        name: str,
+        default):
+
+        if not name in opts:
+            opts[name]=default
+
+    def _process_env_opts(self, ):
+
+        self._check_for_env_opts("episode_timeout_lb", int)
+        self._check_for_env_opts("episode_timeout_ub", int)
+        self._check_for_env_opts("n_steps_task_rand_lb", int)
+        self._check_for_env_opts("n_steps_task_rand_ub", int)
+        self._check_for_env_opts("use_random_trunc", bool)
+        self._check_for_env_opts("random_trunc_freq", int)
+        self._check_for_env_opts("random_trunc_freq_delta", int)
+        self._check_for_env_opts("use_random_safety_reset", bool)
+        self._check_for_env_opts("random_reset_freq", int)
+
+        self._check_for_env_opts("action_repeat", int)
+
+        self._check_for_env_opts("n_preinit_steps", int)
+
+        self._check_for_env_opts("demo_envs_perc", float)
+
+        self._check_for_env_opts("vec_ep_freq_metrics_db", int)
+
+        self._check_for_env_opts("srew_drescaling", bool)
+
+        self._check_for_env_opts("use_action_history", bool)
+        self._check_for_env_opts("actions_history_size", int)
+        
+        self._check_for_env_opts("use_action_smoothing", bool)
+        self._check_for_env_opts("smoothing_horizon_c", float)
+        self._check_for_env_opts("smoothing_horizon_d", float)
+
+        self._check_for_env_opts("add_heightmap_obs", bool)
+
+        self._check_for_env_opts("reset_on_init", bool)
+        
+        # parse action repeat opt + get some sim information
+        if self._env_opts["action_repeat"] <=0: 
+            self._env_opts["action_repeat"] = 1
+        self._action_repeat=self._env_opts["action_repeat"]
+        # parse remote sim info
+        sim_info = {}
+        sim_info_shared = SharedEnvInfo(namespace=self._namespace,
+                    is_server=False,
+                    safe=False,
+                    verbose=self._verbose,
+                    vlevel=self._vlevel)
+        sim_info_shared.run()
+        sim_info_keys = sim_info_shared.param_keys
+        sim_info_data = sim_info_shared.get().flatten()
+        for i in range(len(sim_info_keys)):
+            sim_info[sim_info_keys[i]] = sim_info_data[i]
+        if "substepping_dt" in sim_info_keys:
+            self._substep_dt=sim_info["substepping_dt"]
+        self._env_opts.update(sim_info)
+
+        self._env_opts["substep_dt"]=self._substep_dt
+
+        self._env_opts["override_agent_refs"]=self._override_agent_refs
+
+        self._env_opts["episode_timeout_lb"] = round(self._env_opts["episode_timeout_lb"]/self._action_repeat) 
+        self._env_opts["episode_timeout_ub"] = round(self._env_opts["episode_timeout_ub"]/self._action_repeat)
+
+        self._env_opts["n_steps_task_rand_lb"] = round(self._env_opts["n_steps_task_rand_lb"]/self._action_repeat)
+        self._env_opts["n_steps_task_rand_ub"] = round(self._env_opts["n_steps_task_rand_ub"]/self._action_repeat)
+        
+        if self._env_opts["random_reset_freq"] <=0:
+            self._env_opts["use_random_safety_reset"]=False
+            self._env_opts["random_reset_freq"]=-1
+        self._random_reset_active=self._env_opts["use_random_safety_reset"]
+
+        self._env_opts["random_trunc_freq"] = round(self._env_opts["random_trunc_freq"]/self._action_repeat) 
+        self._env_opts["random_trunc_freq_delta"] = round(self._env_opts["random_trunc_freq_delta"]/self._action_repeat) 
+
+        if self._env_opts["random_trunc_freq"] <=0:
+            self._env_opts["use_random_trunc"]=False
+            self._env_opts["random_trunc_freq"]=-1
+
+        self._full_db=False
+        if "full_env_db" in self._env_opts:
+            self._full_db=self._env_opts["full_env_db"]
+
+    def _check_for_env_opts(self, 
+            name: str,
+            expected_type):
+        if not (name in self._env_opts):
+            Journal.log(self.__class__.__name__,
+                "_check_for_env_opts",
+                f"Required option {name} missing for env opts!",
+                LogType.EXCEP,
+                throw_when_excep=True)
+        if not isinstance(self._env_opts[name], expected_type):
+            Journal.log(self.__class__.__name__,
+                "_check_for_env_opts",
+                f"Option {name} in env opts is not of type {expected_type} (got {type(self._env_opts[name])})!",
+                LogType.EXCEP,
+                throw_when_excep=True)
+            
+    def __del__(self):
+
+        self.close()
+
+    def _demo_setup(self):
+
+        self._demo_envs_idxs=None
+        self._demo_envs_idxs_bool=None
+        self._n_demo_envs=round(self._env_opts["demo_envs_perc"]*self._n_envs)
+        self._add_demos=False
+        if not self._n_demo_envs >0:
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                "will not use demo environments",
+                LogType.INFO,
+                throw_when_excep=False)
+        else:
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                f"Will run with {self._n_demo_envs} demonstration envs.",
+                LogType.INFO)
+            self._demo_envs_idxs = torch.randperm(self._n_envs, device=self._device)[:self._n_demo_envs]
+            self._demo_envs_idxs_bool=torch.full((self._n_envs, ), dtype=torch.bool, device=self._device,
+                                        fill_value=False)
+            self._demo_envs_idxs_bool[self._demo_envs_idxs]=True
+
+            self._init_demo_envs() # custom logic
+
+            demo_idxs_str=", ".join(map(str, self._demo_envs_idxs.tolist()))
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                f"Demo env. indexes are [{demo_idxs_str}]",
+                LogType.INFO)
+    
+    def env_opts(self):
+        return self._env_opts
+    
+    def demo_env_idxs(self, get_bool: bool=False):
+        if get_bool:
+            return self._demo_envs_idxs_bool
+        else:
+            return self._demo_envs_idxs
+        
+    def _init_demo_envs(self):
+        pass
+    
+    def n_demo_envs(self):
+        return self._n_demo_envs
+
+    def demo_active(self):
+        return self._add_demos
+    
+    def switch_demo(self, active: bool = False):
+        if self._demo_envs_idxs is not None:
+            self._add_demos=active
+        else:
+            Journal.log(self.__class__.__name__,
+                "switch_demo",
+                f"Cannot switch demostrations on. No demo envs available!",
+                LogType.EXCEP,
+                throw_when_excep=True)
+
+    def _get_this_file_path(self):
+        return self._this_path
+    
+    def episode_timeout_bounds(self):
+        return self._env_opts["episode_timeout_lb"], self._env_opts["episode_timeout_ub"]
+    
+    def task_rand_timeout_bounds(self):
+        return self._env_opts["n_steps_task_rand_lb"], self._env_opts["n_steps_task_rand_ub"]
+    
+    def n_action_reps(self):
+        return self._action_repeat
+    
+    def get_file_paths(self):
+        from aug_mpc.utils.sys_utils import PathsGetter
+        path_getter = PathsGetter()
+        base_paths = []
+        base_paths.append(self._get_this_file_path())
+        base_paths.append(path_getter.REMOTENVPATH)
+        for script_path in path_getter.SCRIPTSPATHS:
+            base_paths.append(script_path)
+
+        # rhc files
+        from EigenIPC.PyEigenIPC import StringTensorClient
+        from perf_sleep.pyperfsleep import PerfSleep
+        shared_rhc_shared_files = StringTensorClient(
+            basename="SharedRhcFilesDropDir", 
+            name_space=self._namespace,
+            verbose=self._verbose, 
+            vlevel=VLevel.V2)
+        shared_rhc_shared_files.run()
+        shared_rhc_files_vals=[""]*shared_rhc_shared_files.length()
+        while not shared_rhc_shared_files.read_vec(shared_rhc_files_vals, 0):
+            nsecs =  1000000000 # 1 sec
+            PerfSleep.thread_sleep(nsecs) # we just keep it alive
+        rhc_list=[]
+        for rhc_files in shared_rhc_files_vals:
+            file_list = rhc_files.split(", ")
+            rhc_list.extend(file_list)
+        rhc_list = list(set(rhc_list)) # removing duplicates
+        base_paths.extend(rhc_list)
+        
+        # world interface files
+        get_world_interface_paths = self.get_world_interface_paths()
+        base_paths.extend(get_world_interface_paths)
+        return base_paths
+
+    def get_world_interface_paths(self):
+        paths = []
+        shared_world_iface_files = StringTensorClient(
+            basename="SharedWorldInterfaceFilesDropDir", 
+            name_space=self._namespace,
+            verbose=self._verbose, 
+            vlevel=VLevel.V2)
+        shared_world_iface_files.run()
+        world_iface_vals=[""]*shared_world_iface_files.length()
+        while not shared_world_iface_files.read_vec(world_iface_vals, 0):
+            nsecs =  1000000000 # 1 sec
+            PerfSleep.thread_sleep(nsecs) # keep alive while waiting
+        shared_world_iface_files.close()
+        for files in world_iface_vals:
+            if files == "":
+                continue
+            file_list = files.split(", ")
+            for f in file_list:
+                if f not in paths:
+                    paths.append(f)
+        return paths
+
+    def get_aux_dir(self):
+        empty_list = []
+        return empty_list
+
+    def _init_step(self, reset_on_init: bool = True):
+        
+        self._check_controllers_registered(retry=True)
+        self._activate_rhc_controllers()
+
+        # just an auxiliary tensor
+        initial_reset_aux = self._terminations.get_torch_mirror(gpu=self._use_gpu).clone()
+        initial_reset_aux[:, :] = reset_on_init # we reset all sim envs first
+        init_step_ok=True
+        init_step_ok=self._remote_sim_step() and init_step_ok
+        if not init_step_ok:
+            return False
+        init_step_ok=self._remote_reset(reset_mask=initial_reset_aux) and init_step_ok
+        if not init_step_ok:
+            return False
+            
+        for i in range(self._env_opts["n_preinit_steps"]): # perform some
+            # dummy remote env stepping to make sure to have meaningful 
+            # initializations (doesn't increment step counter)
+            init_step_ok=self._remote_sim_step() and init_step_ok # 1 remote sim. step
+            if not init_step_ok:
+                return False
+            init_step_ok=self._send_remote_reset_req() and init_step_ok # fake reset request 
+            if not init_step_ok:
+                return False
+            
+        self.reset()
+
+        return init_step_ok
+
+    def _debug(self):
+
+        if self._use_gpu:
+            # using non_blocking which is not safe when GPU->CPU
+            self._obs.synch_mirror(from_gpu=True,non_blocking=True) # copy data from gpu to cpu view
+            self._next_obs.synch_mirror(from_gpu=True,non_blocking=True)
+            self._actions.synch_mirror(from_gpu=True,non_blocking=True)
+            self._truncations.synch_mirror(from_gpu=True,non_blocking=True) 
+            self._sub_truncations.synch_mirror(from_gpu=True,non_blocking=True)
+            self._terminations.synch_mirror(from_gpu=True,non_blocking=True)
+            self._sub_terminations.synch_mirror(from_gpu=True,non_blocking=True)
+            self._tot_rewards.synch_mirror(from_gpu=True,non_blocking=True)
+            self._sub_rewards.synch_mirror(from_gpu=True,non_blocking=True)
+            # if we want reliable db data then we should synchronize data streams
+            torch.cuda.synchronize()
+
+        # copy CPU view on shared memory
+        self._obs.synch_all(read=False, retry=True) 
+        self._next_obs.synch_all(read=False, retry=True)
+        self._actions.synch_all(read=False, retry=True) 
+        self._tot_rewards.synch_all(read=False, retry=True)
+        self._sub_rewards.synch_all(read=False, retry=True)
+        self._truncations.synch_all(read=False, retry = True) 
+        self._sub_truncations.synch_all(read=False, retry = True)
+        self._terminations.synch_all(read=False, retry = True) 
+        self._sub_terminations.synch_all(read=False, retry = True)
+        
+        self._debug_agent_refs()
+        
+    def _debug_agent_refs(self):
+        if self._use_gpu:
+            if not self._override_agent_refs:
+                self._agent_refs.rob_refs.root_state.synch_mirror(from_gpu=True,non_blocking=False)
+        if not self._override_agent_refs:
+            self._agent_refs.rob_refs.root_state.synch_all(read=False, retry = True)
+
+    def _remote_sim_step(self):
+
+        self._remote_stepper.trigger() # triggers simulation + RHC
+        if not self._remote_stepper.wait_ack_from(1, self._timeout):
+            Journal.log(self.__class__.__name__,
+            "_remote_sim_step",
+            "Remote sim. env step ack not received within timeout",
+            LogType.EXCEP,
+            throw_when_excep = False)
+            return False
+        return True
+
+    def _remote_reset(self,
+                reset_mask: torch.Tensor = None):
+
+        reset_reqs = self._remote_reset_req.get_torch_mirror()
+        if reset_mask is None: # just send the signal to allow stepping, but do not reset any of
+            # the remote envs
+            reset_reqs[:, :] = False
+        else:
+            reset_reqs[:, :] = reset_mask # remotely reset envs corresponding to
+            # the mask (True--> to be reset)
+        self._remote_reset_req.synch_all(read=False, retry=True) # write on shared buffer
+        remote_reset_ok = self._send_remote_reset_req() # process remote request
+
+        if reset_mask is not None:
+            self._synch_state(gpu=self._use_gpu) # if some env was reset, we use _obs
+            # to hold the states, including resets, while _next_obs will always hold the 
+            # state right after stepping the sim env
+            # (could be a bit more efficient, since in theory we only need to read the envs
+            # corresponding to the reset_mask)
+            
+            
+        return remote_reset_ok
+    
+    def _send_remote_reset_req(self):
+
+        self._remote_resetter.trigger()
+        if not self._remote_resetter.wait_ack_from(1, self._timeout): # remote reset completed
+            Journal.log(self.__class__.__name__,
+                "_post_step",
+                "Remote reset did not complete within the prescribed timeout!",
+                LogType.EXCEP,
+                throw_when_excep = False)
+            return False
+        return True
+    
+    def step(self, 
+            action):
+
+        actions_norm = action.detach() # IMPORTANT: assumes actions are already normalized in [-1, 1]
+
+        actions = self._actions.get_torch_mirror(gpu=self._use_gpu) # will hold agent actions (real range)
+
+        # scale normalized actions to physical space before interfacing with controllers
+        actions[:, :] = actions_norm*self._actions_scale + self._actions_offset
+
+        self._override_actions_with_demo() # if necessary override some actions with expert demonstrations
+        # (getting actions with get_actions will return the modified actions tensor)
+
+        actions.clamp_(self._actions_lb, self._actions_ub) # just to be safe
+
+        if self._act_mem_buffer is not None: # store norm actions in memory buffer
+            self._act_mem_buffer.update(new_data=actions_norm)
+
+        if self._env_opts["use_action_smoothing"]:
+            self._apply_actions_smoothing() # smooth actions if enabled (the tensor returned by 
+            # get_actions does not contain smoothing and can be safely employed for experience collection)
+
+        self._apply_actions_to_rhc() # apply last agent actions to rhc controller
+
+        stepping_ok = True
+        tot_rewards = self._tot_rewards.get_torch_mirror(gpu=self._use_gpu)
+        sub_rewards = self._sub_rewards.get_torch_mirror(gpu=self._use_gpu)
+        next_obs = self._next_obs.get_torch_mirror(gpu=self._use_gpu)
+        tot_rewards.zero_()
+        sub_rewards.zero_()
+        self._substep_rewards.zero_()
+        next_obs.zero_() # necessary for substep obs
+
+        for i in range(0, self._action_repeat):
+
+            self._pre_substep() # custom logic @ substep freq
+
+            stepping_ok = stepping_ok and self._check_controllers_registered(retry=False) # does not make sense to run training
+            # if we lost some controllers
+            stepping_ok = stepping_ok and self._remote_sim_step() # blocking, 
+
+            # no sim substepping is allowed to fail
+            self._synch_state(gpu=self._use_gpu) # read state from shared mem (done in substeps also, 
+            # since substeps rewards will need updated substep obs)
+            
+            self._custom_post_substp_pre_rew() # custom substepping logic
+            self._compute_substep_rewards()
+            self._assemble_substep_rewards() # includes rewards clipping
+            self._custom_post_substp_post_rew() # custom substepping logic
+
+            # fill substep obs
+            self._fill_substep_obs(self._substep_obs)
+            self._assemble_substep_obs()
+            if not i==(self._action_repeat-1):
+                # sends reset signal to complete remote step sequence,
+                # but does not reset any remote env
+                stepping_ok = stepping_ok and self._remote_reset(reset_mask=None) 
+            else: # last substep
+                
+                self._fill_step_obs(next_obs) # update next obs
+                self._clamp_obs(next_obs) # good practice
+                obs = self._obs.get_torch_mirror(gpu=self._use_gpu)
+                obs[:, :] = next_obs # start from next observation, unless reset (handled in post_step())
+
+                self._compute_step_rewards() # implemented by child
+
+                tot_rewards = self._tot_rewards.get_torch_mirror(gpu=self._use_gpu)
+                sub_rewards = self._sub_rewards.get_torch_mirror(gpu=self._use_gpu)
+                self._clamp_rewards(sub_rewards) # clamp all sub rewards
+                
+                tot_rewards[:, :] = torch.sum(sub_rewards, dim=1, keepdim=True)
+
+                scale=1 # scale tot rew by the number of action repeats
+                if self._env_opts["srew_drescaling"]: # scale rewards depending on the n of subrewards
+                    scale*=sub_rewards.shape[1] # n. dims rescaling
+                tot_rewards.mul_(1/scale)
+
+            self._substep_abs_counter.increment() # @ substep freq
+
+            if not stepping_ok:
+                return False
+            
+        stepping_ok =  stepping_ok and self._post_step() # post sub-stepping operations
+        # (if action_repeat > 1, then just the db data at the last substep is logged)
+        # also, if a reset of an env occurs, obs will hold the reset state
+
+        return stepping_ok 
+    
+    def _post_step(self):
+        
+        # first increment counters
+        self._ep_timeout_counter.increment() # episode timeout
+        self._task_rand_counter.increment() # task randomization
+        if self._rand_trunc_counter is not None: # random truncations (for removing temp. correlations)
+            self._rand_trunc_counter.increment()
+
+        # check truncation and termination conditions 
+        self._check_truncations() # defined in child env
+        self._check_terminations()
+        terminated = self._terminations.get_torch_mirror(gpu=self._use_gpu)
+        truncated = self._truncations.get_torch_mirror(gpu=self._use_gpu)
+        ignore_ep_end=None
+        if self._rand_trunc_counter is not None:
+            ignore_ep_end=self._rand_trunc_counter.time_limits_reached()
+            if self._use_gpu:
+                ignore_ep_end=ignore_ep_end.cuda()
+        
+            truncated = torch.logical_or(truncated, 
+                ignore_ep_end) # add truncation (sub truncations defined in child env
+            # remain untouched)
+        
+        episode_finished = torch.logical_or(terminated,
+                            truncated)
+        episode_finished_cpu = episode_finished.cpu()
+
+        if self._rand_safety_reset_counter is not None and self._random_reset_active:
+            self._rand_safety_reset_counter.increment(to_be_incremented=episode_finished_cpu.flatten())
+            # truncated[:,:] = torch.logical_or(truncated,
+            #     self._rand_safety_reset_counter.time_limits_reached().cuda())
+
+        if self._act_mem_buffer is not None:
+            self._act_mem_buffer.reset(to_be_reset=episode_finished.flatten(),
+                            init_data=self._normalize_actions(self.default_action))
+
+        if self._action_smoother_continuous is not None:
+            self._action_smoother_continuous.reset(to_be_reset=episode_finished.flatten(),
+                reset_val=self.default_action[:, self._is_continuous_actions])
+        if self._action_smoother_discrete is not None:
+            self._action_smoother_discrete.reset(to_be_reset=episode_finished.flatten(),
+                reset_val=self.default_action[:, ~self._is_continuous_actions])
+
+        # debug step if required (IMPORTANT: must be before remote reset so that we always db
+        # actual data from the step and not after reset)
+        if self._is_debug:
+            self._debug() # copies db data on shared memory
+            ignore_ep_end_cpu=ignore_ep_end if not self._use_gpu else ignore_ep_end.cpu()
+            self._update_custom_db_data(episode_finished=episode_finished_cpu, 
+                    ignore_ep_end=ignore_ep_end_cpu # ignore data if random trunc
+                    )           
+            self._episodic_rewards_metrics.update(rewards = self._sub_rewards.get_torch_mirror(gpu=False),
+                    ep_finished=episode_finished_cpu,
+                    ignore_ep_end=ignore_ep_end_cpu # ignore data if random trunc
+                    )
+
+        # remotely reset envs
+        to_be_reset=self._to_be_reset()
+        to_be_reset_custom=self._custom_reset()
+        if to_be_reset_custom is not None:
+            to_be_reset[:, :] = torch.logical_or(to_be_reset,to_be_reset_custom)
+        rm_reset_ok = self._remote_reset(reset_mask=to_be_reset)
+        
+        self._custom_post_step(episode_finished=episode_finished) # any additional logic from child env  
+        # here, before actual reset taskes place  (at this point the state is the reset one)
+
+        # updating also prev pos and orientation in case some env was reset
+        self._prev_root_p_substep[:, :]=self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu)
+        self._prev_root_q_substep[:, :]=self._robot_state.root_state.get(data_type="q",gpu=self._use_gpu)
+
+        obs = self._obs.get_torch_mirror(gpu=self._use_gpu)
+        self._fill_step_obs(obs)
+        self._clamp_obs(obs)
+
+        # updating prev step quantities
+        self._prev_root_p_step[:, :]=self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu)
+        self._prev_root_q_step[:, :]=self._robot_state.root_state.get(data_type="q",gpu=self._use_gpu)
+
+        # synchronize and reset counters for finished episodes
+        self._ep_timeout_counter.reset(to_be_reset=episode_finished)
+        self._task_rand_counter.reset(to_be_reset=episode_finished)
+        self._substep_abs_counter.reset(to_be_reset=torch.logical_or(terminated,to_be_reset),
+            randomize_offsets=True # otherwise timers across envs would be strongly correlated
+            ) # reset only if resetting environment or if terminal
+
+        if self._rand_trunc_counter is not None:
+            # only reset when safety truncation was is triggered   
+            self._rand_trunc_counter.reset(to_be_reset=self._rand_trunc_counter.time_limits_reached(),
+                randomize_limits=True, # we need to randomize otherwise the other counters will synchronize
+                # with the episode counters
+                randomize_offsets=False # always restart at 0
+                )
+        # safety reset counter is only when it reches its reset interval (just to keep
+        # the counter bounded)
+        if self._rand_safety_reset_counter is not None and self._random_reset_active:
+            self._rand_safety_reset_counter.reset(to_be_reset=self._rand_safety_reset_counter.time_limits_reached())
+
+        return rm_reset_ok
+    
+    def _to_be_reset(self):
+        # always reset if a termination occurred or if there's a random safety reset
+        # request
+        terminated = self._terminations.get_torch_mirror(gpu=self._use_gpu)
+        to_be_reset=terminated.clone()
+        if (self._rand_safety_reset_counter is not None) and self._random_reset_active:
+            to_be_reset=torch.logical_or(to_be_reset,
+                self._rand_safety_reset_counter.time_limits_reached())
+
+        return to_be_reset
+
+    def _custom_reset(self):
+        # can be overridden by child
+        return None
+    
+    def _apply_actions_smoothing(self):
+
+        actions = self._actions.get_torch_mirror(gpu=self._use_gpu)
+        actual_actions=self.get_actual_actions() # will write smoothed actions here
+        if self._action_smoother_continuous is not None:
+            self._action_smoother_continuous.update(new_signal=
+                    actions[:, self._is_continuous_actions])
+            actual_actions[:, self._is_continuous_actions]=self._action_smoother_continuous.get()
+        if self._action_smoother_discrete is not None:
+            self._action_smoother_discrete.update(new_signal=
+                    actions[:, ~self._is_continuous_actions])
+            actual_actions[:, ~self._is_continuous_actions]=self._action_smoother_discrete.get()
+
+    def _update_custom_db_data(self,
+                    episode_finished,
+                    ignore_ep_end):
+
+        # update defaults
+        self.custom_db_data["RhcRefsFlag"].update(new_data=self._rhc_refs.contact_flags.get_torch_mirror(gpu=False), 
+                                    ep_finished=episode_finished,
+                                    ignore_ep_end=ignore_ep_end) # before potentially resetting the flags, get data
+        self.custom_db_data["Actions"].update(new_data=self._actions.get_torch_mirror(gpu=False), 
+                                    ep_finished=episode_finished,
+                                    ignore_ep_end=ignore_ep_end)
+        self.custom_db_data["Obs"].update(new_data=self._obs.get_torch_mirror(gpu=False), 
+                                    ep_finished=episode_finished,
+                                    ignore_ep_end=ignore_ep_end)
+        
+        self.custom_db_data["SubTerminations"].update(new_data=self._sub_terminations.get_torch_mirror(gpu=False), 
+                                    ep_finished=episode_finished,
+                                    ignore_ep_end=ignore_ep_end)
+        self.custom_db_data["SubTruncations"].update(new_data=self._sub_truncations.get_torch_mirror(gpu=False), 
+                                    ep_finished=episode_finished,
+                                    ignore_ep_end=ignore_ep_end)
+        
+        self.custom_db_data["Terminations"].update(new_data=self._terminations.get_torch_mirror(gpu=False), 
+                                    ep_finished=episode_finished,
+                                    ignore_ep_end=ignore_ep_end)
+        self.custom_db_data["Truncations"].update(new_data=self._truncations.get_torch_mirror(gpu=False), 
+                                    ep_finished=episode_finished,
+                                    ignore_ep_end=ignore_ep_end)
+
+        
+        self._get_custom_db_data(episode_finished=episode_finished, ignore_ep_end=ignore_ep_end)
+
+    def reset_custom_db_data(self, keep_track: bool = False):
+        # to be called periodically to reset custom db data stat. collection 
+        for custom_db_data in self.custom_db_data.values():
+            custom_db_data.reset(keep_track=keep_track)
+
+    def _assemble_substep_rewards(self):
+        # by default assemble  substep rewards by averaging
+        sub_rewards = self._sub_rewards.get_torch_mirror(gpu=self._use_gpu)
+        
+        # average over substeps depending on scale
+        # sub_rewards[:, self._is_substep_rew] = sub_rewards[:, self._is_substep_rew] + \
+        #     self._substep_rewards[:, self._is_substep_rew]/self._action_repeat
+        sub_rewards[:, self._is_substep_rew] += self._substep_rewards[:, self._is_substep_rew]/self._action_repeat
+    
+    def _assemble_substep_obs(self):
+        next_obs = self._next_obs.get_torch_mirror(gpu=self._use_gpu)        
+        next_obs[:, self._is_substep_obs] += self._substep_obs[:, self._is_substep_obs]/self._action_repeat
+
+    def randomize_task_refs(self,
+                env_indxs: torch.Tensor = None):
+                    
+        if self._override_agent_refs:
+            self._override_refs(env_indxs=env_indxs)
+        else:
+            self._randomize_task_refs(env_indxs=env_indxs)
+            
+    def reset(self):
+        
+        self.randomize_task_refs(env_indxs=None) # randomize all refs across envs
+
+        self._obs.reset()
+        self._actions.reset()
+        self._next_obs.reset()
+        self._sub_rewards.reset()
+        self._tot_rewards.reset()
+        self._terminations.reset()
+        self._sub_terminations.reset()
+        self._truncations.reset()
+        self._sub_truncations.reset()
+
+        self._ep_timeout_counter.reset(randomize_offsets=True)
+        self._task_rand_counter.reset()
+        self._task_rand_counter.sync_counters(other_counter=self._ep_timeout_counter)
+        if self._rand_safety_reset_counter is not None:
+            self._rand_safety_reset_counter.reset()
+        self._substep_abs_counter.reset()
+
+        if self._act_mem_buffer is not None:
+            self._act_mem_buffer.reset_all(init_data=self._normalize_actions(self.default_action))
+
+        if self._action_smoother_continuous is not None:
+            self._action_smoother_continuous.reset(reset_val=self.default_action[:, self._is_continuous_actions])
+        if self._action_smoother_discrete is not None:
+            self._action_smoother_discrete.reset(reset_val=self.default_action[:, ~self._is_continuous_actions])
+
+        self._synch_state(gpu=self._use_gpu) # read obs from shared mem
+
+        # just calling custom post step to ensure tak refs are updated 
+        terminated = self._terminations.get_torch_mirror(gpu=self._use_gpu)
+        truncated = self._truncations.get_torch_mirror(gpu=self._use_gpu)
+        episode_finished = torch.logical_or(terminated,
+                            truncated)
+        self._custom_post_step(episode_finished=episode_finished)
+
+        obs = self._obs.get_torch_mirror(gpu=self._use_gpu)
+        next_obs = self._next_obs.get_torch_mirror(gpu=self._use_gpu)
+        self._fill_step_obs(obs) # initialize observations 
+        self._clamp_obs(obs) # to avoid bad things
+        self._fill_step_obs(next_obs) # and next obs
+        self._clamp_obs(next_obs)
+
+        self.reset_custom_db_data(keep_track=False)
+        self._episodic_rewards_metrics.reset(keep_track=False)
+
+        self._prev_root_p_step[:, :]=self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu)
+        self._prev_root_q_step[:, :]=self._robot_state.root_state.get(data_type="q",gpu=self._use_gpu)
+        self._prev_root_p_substep[:, :]=self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu)
+        self._prev_root_q_substep[:, :]=self._robot_state.root_state.get(data_type="q",gpu=self._use_gpu)
+
+    def is_ready(self):
+        return self._ready
+    
+    def close(self):
+        
+        if not self._closed:
+
+            # close all shared mem. clients
+            self._robot_state.close()
+            self._rhc_cmds.close()
+            self._rhc_pred.close()
+            self._rhc_refs.close()
+            self._rhc_status.close()
+            
+            self._remote_stepper.close()
+            
+            self._ep_timeout_counter.close()
+            self._task_rand_counter.close()
+            if self._rand_safety_reset_counter is not None:
+                self._rand_safety_reset_counter.close()
+
+            # closing env.-specific shared data
+            self._obs.close()
+            self._next_obs.close()
+            self._actions.close()
+            if self._actual_actions is not None:
+                self._actual_actions.close()
+            self._sub_rewards.close()
+            self._tot_rewards.close()
+
+            self._terminations.close()
+            self._sub_terminations.close()
+            self._truncations.close()
+            self._sub_truncations.close()
+
+            self._closed = True
+
+    def get_obs(self, clone:bool=False):
+        if clone:
+            return self._obs.get_torch_mirror(gpu=self._use_gpu).detach().clone()
+        else:
+            return self._obs.get_torch_mirror(gpu=self._use_gpu).detach()
+
+    def get_next_obs(self, clone:bool=False):
+        if clone:
+            return self._next_obs.get_torch_mirror(gpu=self._use_gpu).detach().clone()
+        else:
+            return self._next_obs.get_torch_mirror(gpu=self._use_gpu).detach()
+        
+    def get_actions(self, clone:bool=False, normalized: bool = False):
+        actions = self._actions.get_torch_mirror(gpu=self._use_gpu).detach()
+        if normalized:
+            normalized_actions = self._normalize_actions(actions)
+            return normalized_actions.clone() if clone else normalized_actions
+        return actions.clone() if clone else actions
+    
+    def get_actual_actions(self, clone:bool=False, normalized: bool = False):
+        if self._env_opts["use_action_smoothing"]:
+            actions = self._actual_actions.get_torch_mirror(gpu=self._use_gpu).detach()
+        else: # actual action coincides with the one from the agent + possible modif.
+            actions = self.get_actions(clone=False, normalized=False)
+        if normalized:
+            normalized_actions = self._normalize_actions(actions)
+            return normalized_actions.clone() if clone else normalized_actions
+        return actions.clone() if clone else actions
+
+    def _normalize_actions(self, actions: torch.Tensor):
+        scale = torch.where(self._actions_scale == 0.0,
+            torch.ones_like(self._actions_scale),
+            self._actions_scale)
+        normalized = (actions - self._actions_offset)/scale
+        zero_scale_mask = torch.eq(self._actions_scale, 0.0).squeeze(0)
+        if torch.any(zero_scale_mask):
+            normalized[:, zero_scale_mask] = 0.0
+        return normalized
+        
+    def get_rewards(self, clone:bool=False):
+        if clone:
+            return self._tot_rewards.get_torch_mirror(gpu=self._use_gpu).detach().clone()
+        else:
+            return self._tot_rewards.get_torch_mirror(gpu=self._use_gpu).detach()
+        
+    def get_terminations(self, clone:bool=False):
+        if clone:
+            return self._terminations.get_torch_mirror(gpu=self._use_gpu).detach().clone()
+        else:
+            return self._terminations.get_torch_mirror(gpu=self._use_gpu).detach()
+    
+    def get_truncations(self, clone:bool=False):
+        if clone:
+            return self._truncations.get_torch_mirror(gpu=self._use_gpu).detach().clone()
+        else:
+            return self._truncations.get_torch_mirror(gpu=self._use_gpu).detach()
+        
+    def obs_dim(self):
+
+        return self._obs_dim
+    
+    def actions_dim(self):
+
+        return self._actions_dim
+    
+    def ep_rewards_metrics(self):
+
+        return self._episodic_rewards_metrics
+    
+    def using_gpu(self):
+
+        return self._use_gpu
+
+    def name(self):
+
+        return self._env_name
+
+    def n_envs(self):
+
+        return self._n_envs
+
+    def dtype(self):
+                                    
+        return self._dtype 
+    
+    def obs_names(self):
+        return self._get_obs_names()
+    
+    def action_names(self):
+        return self._get_action_names()
+
+    def sub_rew_names(self):
+        return self._get_rewards_names()
+    
+    def sub_term_names(self):
+        return self._get_sub_term_names()
+    
+    def sub_trunc_names(self):
+        return self._get_sub_trunc_names()
+    
+    def _get_obs_names(self):
+        # to be overridden by child class
+        return None
+    
+    def get_robot_jnt_names(self):
+        return self._robot_state.jnt_names()
+
+    def _get_action_names(self):
+        # to be overridden by child class
+        return None
+    
+    def _get_rewards_names(self):
+        # to be overridden by child class
+        return None
+    
+    def _get_sub_term_names(self):
+        # to be overridden by child class
+        sub_term_names = []
+        sub_term_names.append("rhc_failure")
+        sub_term_names.append("robot_capsize")
+        sub_term_names.append("rhc_capsize")
+
+        return sub_term_names
+    
+    def _get_sub_trunc_names(self):
+        # to be overridden by child class
+        sub_trunc_names = []
+        sub_trunc_names.append("ep_timeout")
+
+        return sub_trunc_names
+    
+    def _get_custom_db_data(self, episode_finished):
+        # to be overridden by child class
+        pass
+    
+    def set_observed_joints(self):
+        # ny default observe all joints available 
+        return self._robot_state.jnt_names()
+    
+    def _set_jnts_blacklist_pattern(self):
+        self._jnt_q_blacklist_patterns=[]
+
+    def get_observed_joints(self):
+        return self._observed_jnt_names
+    
+    def _init_obs(self, obs_dim: int):
+        
+        device = "cuda" if self._use_gpu else "cpu"
+        
+        obs_threshold_default = 1e3
+        self._obs_threshold_lb = -obs_threshold_default # used for clipping observations
+        self._obs_threshold_ub = obs_threshold_default
+        
+        self._obs_ub = torch.full((1, obs_dim), dtype=self._dtype, device=device,
+                                        fill_value=1.0) 
+        self._obs_lb = torch.full((1, obs_dim), dtype=self._dtype, device=device,
+                                        fill_value=-1.0)
+        self._obs_scale = (self._obs_ub - self._obs_lb)/2.0
+        self._obs_offset = (self._obs_ub + self._obs_lb)/2.0
+
+        if not self._obs_dim==len(self._get_obs_names()):
+            error=f"obs dim {self._obs_dim} does not match obs names length {len(self._get_obs_names())}!!"
+            Journal.log(self.__class__.__name__,
+                "_init_obs",
+                error,
+                LogType.EXCEP,
+                throw_when_excep = True)
+
+        self._obs = Observations(namespace=self._namespace,
+                            n_envs=self._n_envs,
+                            obs_dim=self._obs_dim,
+                            obs_names=self._get_obs_names(),
+                            env_names=None,
+                            is_server=True,
+                            verbose=self._verbose,
+                            vlevel=self._vlevel,
+                            safe=True,
+                            force_reconnection=True,
+                            with_gpu_mirror=self._use_gpu,
+                            fill_value=0.0)
+        
+        self._next_obs = NextObservations(namespace=self._namespace,
+                            n_envs=self._n_envs,
+                            obs_dim=self._obs_dim,
+                            obs_names=self._get_obs_names(),
+                            env_names=None,
+                            is_server=True,
+                            verbose=self._verbose,
+                            vlevel=self._vlevel,
+                            safe=True,
+                            force_reconnection=True,
+                            with_gpu_mirror=self._use_gpu,
+                            fill_value=0.0)
+
+        self._obs.run()
+        self._next_obs.run()
+
+        self._is_substep_obs = torch.zeros((self.obs_dim(),), dtype=torch.bool, device=device)
+        self._is_substep_obs.fill_(False) # default to all step obs
+
+        # not super memory efficient
+        self._substep_obs=torch.full_like(self._obs.get_torch_mirror(gpu=self._use_gpu), fill_value=0.0)
+        
+    def _init_actions(self, actions_dim: int):
+        
+        device = "cuda" if self._use_gpu else "cpu"
+        # action scalings to be applied to agent's output
+        self._actions_ub = torch.full((1, actions_dim), dtype=self._dtype, device=device,
+                                        fill_value=1.0) 
+        self._actions_lb = torch.full((1, actions_dim), dtype=self._dtype, device=device,
+                                        fill_value=-1.0)
+        self._actions_scale = (self._actions_ub - self._actions_lb)/2.0
+        self._actions_offset = (self._actions_ub + self._actions_lb)/2.0
+
+        if not self._actions_dim==len(self._get_action_names()):
+            error=f"action dim {self._actions_dim} does not match action names length {len(self._get_action_names())}!!"
+            Journal.log(self.__class__.__name__,
+                "_init_actions",
+                error,
+                LogType.EXCEP,
+                throw_when_excep = True)
+        self._actions = Actions(namespace=self._namespace,
+            n_envs=self._n_envs,
+            action_dim=self._actions_dim,
+            action_names=self._get_action_names(),
+            env_names=None,
+            is_server=True,
+            verbose=self._verbose,
+            vlevel=self._vlevel,
+            safe=True,
+            force_reconnection=True,
+            with_gpu_mirror=self._use_gpu,
+            fill_value=0.0)
+
+        self._actions.run()
+
+        self.default_action = torch.full_like(input=self.get_actions(),fill_value=0.0)
+        self.safe_action = torch.full_like(input=self.get_actions(),fill_value=0.0)
+
+        if self._env_opts["use_action_history"]:
+            self._act_mem_buffer=MemBuffer(name="ActionMemBuf",
+                data_tensor=self._actions.get_torch_mirror(),
+                data_names=self._get_action_names(),
+                debug=self._debug,
+                horizon=self._env_opts["actions_history_size"],
+                dtype=self._dtype,
+                use_gpu=self._use_gpu)
+        
+        # default to all continuous actions (changes the way noise is added)
+        self._is_continuous_actions=torch.full((actions_dim, ), 
+            dtype=torch.bool, device=device,
+            fill_value=True) 
+    
+    def _init_action_smoothing(self):
+            
+        continuous_actions=self.get_actions()[:, self._is_continuous_actions]
+        discrete_actions=self.get_actions()[:, ~self._is_continuous_actions]
+        self._action_smoother_continuous=ExponentialSignalSmoother(signal=continuous_actions,
+            update_dt=self._substep_dt*self._action_repeat, # rate at which actions are decided by agent
+            smoothing_horizon=self._env_opts["smoothing_horizon_c"],
+            target_smoothing=0.5, 
+            debug=self._debug,
+            dtype=self._dtype,
+            use_gpu=self._use_gpu,
+            name="ActionSmootherContinuous")
+        self._action_smoother_discrete=ExponentialSignalSmoother(signal=discrete_actions,
+            update_dt=self._substep_dt*self._action_repeat, # rate at which actions are decided by agent
+            smoothing_horizon=self._env_opts["smoothing_horizon_d"],
+            target_smoothing=0.5,
+            debug=self._debug,
+            dtype=self._dtype,
+            use_gpu=self._use_gpu,
+            name="ActionSmootherDiscrete")
+        
+        # we also need somewhere to keep the actual actions after smoothing
+        self._actual_actions = Actions(namespace=self._namespace+"_actual",
+            n_envs=self._n_envs,
+            action_dim=self._actions_dim,
+            action_names=self._get_action_names(),
+            env_names=None,
+            is_server=True,
+            verbose=self._verbose,
+            vlevel=self._vlevel,
+            safe=True,
+            force_reconnection=True,
+            with_gpu_mirror=self._use_gpu,
+            fill_value=0.0)
+        self._actual_actions.run()
+            
+    def _init_rewards(self):
+        
+        reward_thresh_default = 1.0
+        n_sub_rewards = len(self._get_rewards_names())
+        device = "cuda" if self._use_gpu else "cpu"
+        self._reward_thresh_lb = torch.full((1, n_sub_rewards), dtype=self._dtype, fill_value=-reward_thresh_default, device=device) # used for clipping rewards
+        self._reward_thresh_ub = torch.full((1, n_sub_rewards), dtype=self._dtype, fill_value=reward_thresh_default, device=device) 
+
+        self._sub_rewards = SubRewards(namespace=self._namespace,
+                            n_envs=self._n_envs,
+                            n_rewards=n_sub_rewards,
+                            reward_names=self._get_rewards_names(),
+                            env_names=None,
+                            is_server=True,
+                            verbose=self._verbose,
+                            vlevel=self._vlevel,
+                            safe=True,
+                            force_reconnection=True,
+                            with_gpu_mirror=self._use_gpu,
+                            fill_value=0.0)
+        
+        self._tot_rewards = TotRewards(namespace=self._namespace,
+                            n_envs=self._n_envs,
+                            reward_names=["total_reward"],
+                            env_names=None,
+                            is_server=True,
+                            verbose=self._verbose,
+                            vlevel=self._vlevel,
+                            safe=True,
+                            force_reconnection=True,
+                            with_gpu_mirror=self._use_gpu,
+                            fill_value=0.0)
+        
+        self._sub_rewards.run()
+        self._tot_rewards.run()
+
+        self._substep_rewards = self._sub_rewards.get_torch_mirror(gpu=self._use_gpu).detach().clone() 
+        # used to hold substep rewards (not super mem. efficient)
+        self._is_substep_rew = torch.zeros((self._substep_rewards.shape[1],),dtype=torch.bool,device=device)
+        self._is_substep_rew.fill_(True) # default to all substep rewards
+   
+        self._episodic_rewards_metrics = EpisodicRewards(reward_tensor=self._sub_rewards.get_torch_mirror(),
+                                        reward_names=self._get_rewards_names(),
+                                        ep_vec_freq=self._env_opts["vec_ep_freq_metrics_db"],
+                                        store_transitions=self._full_db,
+                                        max_ep_duration=self._max_ep_length())
+        self._episodic_rewards_metrics.set_constant_data_scaling(scaling=self._get_reward_scaling())
+        
+    def _get_reward_scaling(self):
+        # to be overridden by child (default to no scaling)
+        return 1
+    
+    def _max_ep_length(self):
+        #.should be overriden by child 
+        return self._env_opts["episode_timeout_ub"]
+    
+    def _init_custom_db_data(self):
+
+        self.custom_db_data = {}
+        # by default always log this contact data
+        rhc_latest_contact_ref = self._rhc_refs.contact_flags.get_torch_mirror()
+        contact_names = self._rhc_refs.rob_refs.contact_names()
+        stepping_data = EpisodicData("RhcRefsFlag", rhc_latest_contact_ref, contact_names,
+            ep_vec_freq=self._env_opts["vec_ep_freq_metrics_db"],
+            store_transitions=self._full_db,
+            max_ep_duration=self._max_ep_length())
+        self._add_custom_db_data(db_data=stepping_data)
+        
+        # log also action data
+        actions = self._actions.get_torch_mirror()
+        action_names = self._get_action_names()
+        action_data = EpisodicData("Actions", actions, action_names,
+            ep_vec_freq=self._env_opts["vec_ep_freq_metrics_db"],
+            store_transitions=self._full_db,
+            max_ep_duration=self._max_ep_length())
+        self._add_custom_db_data(db_data=action_data)
+        
+        # and observations
+        observations = self._obs.get_torch_mirror()
+        observations_names = self._get_obs_names()
+        obs_data = EpisodicData("Obs", observations, observations_names,
+            ep_vec_freq=self._env_opts["vec_ep_freq_metrics_db"],
+            store_transitions=self._full_db,
+            max_ep_duration=self._max_ep_length())
+        self._add_custom_db_data(db_data=obs_data)
+
+        # log sub-term and sub-truncations data
+        t_scaling=1 # 1 so that we log an interpretable data in terms of why the episode finished
+        data_scaling = torch.full((self._n_envs, 1),
+                    fill_value=t_scaling,
+                    dtype=torch.int32,device="cpu")
+        sub_term = self._sub_terminations.get_torch_mirror()
+        term = self._terminations.get_torch_mirror()
+        sub_termination_names = self.sub_term_names()
+    
+        sub_term_data = EpisodicData("SubTerminations", sub_term, sub_termination_names,
+            ep_vec_freq=self._env_opts["vec_ep_freq_metrics_db"],
+            store_transitions=self._full_db,
+            max_ep_duration=self._max_ep_length())
+        sub_term_data.set_constant_data_scaling(enable=True,scaling=data_scaling)
+        self._add_custom_db_data(db_data=sub_term_data)
+        term_data = EpisodicData("Terminations", term, ["terminations"],
+            ep_vec_freq=self._env_opts["vec_ep_freq_metrics_db"],
+            store_transitions=self._full_db,
+            max_ep_duration=self._max_ep_length())
+        term_data.set_constant_data_scaling(enable=True,scaling=data_scaling)
+        self._add_custom_db_data(db_data=term_data)
+        
+        sub_trunc = self._sub_truncations.get_torch_mirror()
+        trunc = self._truncations.get_torch_mirror()
+        sub_truncations_names = self.sub_trunc_names()
+        sub_trunc_data = EpisodicData("SubTruncations", sub_trunc, sub_truncations_names,
+            ep_vec_freq=self._env_opts["vec_ep_freq_metrics_db"],
+            store_transitions=self._full_db,
+            max_ep_duration=self._max_ep_length())
+        sub_trunc_data.set_constant_data_scaling(enable=True,scaling=data_scaling)
+        self._add_custom_db_data(db_data=sub_trunc_data)
+        trunc_data = EpisodicData("Truncations", trunc, ["truncations"],
+            ep_vec_freq=self._env_opts["vec_ep_freq_metrics_db"],
+            store_transitions=self._full_db,
+            max_ep_duration=self._max_ep_length())
+        trunc_data.set_constant_data_scaling(enable=True,scaling=data_scaling)
+        self._add_custom_db_data(db_data=trunc_data)
+
+    def _add_custom_db_data(self, db_data: EpisodicData):
+        self.custom_db_data[db_data.name()] = db_data
+
+    def _init_terminations(self):
+
+        # Boolean array indicating whether each environment episode has terminated after 
+        # the current step. An episode termination could occur based on predefined conditions
+        # in the environment, such as reaching a goal or exceeding a time limit.
+
+        self._terminations = Terminations(namespace=self._namespace,
+                            n_envs=self._n_envs,
+                            is_server=True,
+                            verbose=self._verbose,
+                            vlevel=self._vlevel,
+                            safe=True,
+                            force_reconnection=True,
+                            with_gpu_mirror=self._use_gpu,
+                            fill_value=False) 
+        self._terminations.run()
+
+        sub_t_names = self.sub_term_names()
+        self._sub_terminations = SubTerminations(namespace=self._namespace,
+                n_envs=self._n_envs,
+                n_term=len(sub_t_names),
+                term_names=sub_t_names,
+                is_server=True,
+                verbose=self._verbose,
+                vlevel=self._vlevel,
+                safe=True,
+                force_reconnection=True,
+                with_gpu_mirror=self._use_gpu,
+                fill_value=False)
+        self._sub_terminations.run()
+
+        device = "cuda" if self._use_gpu else "cpu"
+        self._is_capsized=torch.zeros((self._n_envs,1), 
+            dtype=torch.bool, device=device)
+        self._is_rhc_capsized=torch.zeros((self._n_envs,1), 
+            dtype=torch.bool, device=device)
+        self._max_pitch_angle=60.0*math.pi/180.0
+    
+    def _init_truncations(self):
+        
+        self._truncations = Truncations(namespace=self._namespace,
+                            n_envs=self._n_envs,
+                            is_server=True,
+                            verbose=self._verbose,
+                            vlevel=self._vlevel,
+                            safe=True,
+                            force_reconnection=True,
+                            with_gpu_mirror=self._use_gpu,
+                            fill_value=False) 
+        
+        self._truncations.run()
+
+        sub_trc_names = self.sub_trunc_names()
+        self._sub_truncations = SubTruncations(namespace=self._namespace,
+                n_envs=self._n_envs,
+                n_trunc=len(sub_trc_names),
+                truc_names=sub_trc_names,
+                is_server=True,
+                verbose=self._verbose,
+                vlevel=self._vlevel,
+                safe=True,
+                force_reconnection=True,
+                with_gpu_mirror=self._use_gpu,
+                fill_value=False)
+        self._sub_truncations.run()
+    
+    def _update_jnt_blacklist(self):
+        device = "cuda" if self._use_gpu else "cpu"
+        all_available_jnts=self.get_observed_joints()        
+        blacklist=[]
+        for i in range(len(all_available_jnts)):
+            for pattern in self._jnt_q_blacklist_patterns:
+                if pattern in all_available_jnts[i]:
+                    # stop at first pattern match
+                    blacklist.append(i)
+                    break
+        if not len(blacklist)==0:
+            self._jnt_q_blacklist_idxs=torch.tensor(blacklist, dtype=torch.int, device=device)
+            
+    def _attach_to_shared_mem(self):
+
+        # runs shared mem clients for getting observation and setting RHC commands
+
+        # remote stepping data
+        self._remote_stepper = RemoteStepperSrvr(namespace=self._namespace,
+                            verbose=self._verbose,
+                            vlevel=self._vlevel,
+                            force_reconnection=True)
+        self._remote_stepper.run()
+        self._remote_resetter = RemoteResetSrvr(namespace=self._namespace,
+                            verbose=self._verbose,
+                            vlevel=self._vlevel,
+                            force_reconnection=True)
+        self._remote_resetter.run()
+        self._remote_reset_req = RemoteResetRequest(namespace=self._namespace,
+                                            is_server=False, 
+                                            verbose=self._verbose,
+                                            vlevel=self._vlevel,
+                                            safe=True)
+        self._remote_reset_req.run()
+
+        self._jnts_remapping=None
+        self._jnt_q_blacklist_idxs=None
+
+        self._robot_state = RobotState(namespace=self._namespace,
+                                is_server=False, 
+                                safe=self._safe_shared_mem,
+                                verbose=self._verbose,
+                                vlevel=self._vlevel,
+                                with_gpu_mirror=self._use_gpu,
+                                with_torch_view=True,
+                                enable_height_sensor=self._env_opts["add_heightmap_obs"])
+        
+        self._rhc_cmds = RhcCmds(namespace=self._namespace,
+                                is_server=False, 
+                                safe=self._safe_shared_mem,
+                                verbose=self._verbose,
+                                vlevel=self._vlevel,
+                                with_gpu_mirror=self._use_gpu,
+                                with_torch_view=True)
+        
+        self._rhc_pred = RhcPred(namespace=self._namespace,
+                                is_server=False, 
+                                safe=self._safe_shared_mem,
+                                verbose=self._verbose,
+                                vlevel=self._vlevel,
+                                with_gpu_mirror=self._use_gpu,
+                                with_torch_view=True)
+
+        self._rhc_refs = RhcRefs(namespace=self._namespace,
+                            is_server=False,
+                            safe=self._safe_shared_mem,
+                            verbose=self._verbose,
+                            vlevel=self._vlevel,
+                            with_gpu_mirror=self._use_gpu,
+                            with_torch_view=True)
+
+        self._rhc_status = RhcStatus(namespace=self._namespace,
+                                is_server=False,
+                                verbose=self._verbose,
+                                vlevel=self._vlevel,
+                                with_gpu_mirror=self._use_gpu,
+                                with_torch_view=True)
+        
+        self._robot_state.run()
+        self._n_envs = self._robot_state.n_robots()
+        self._n_jnts = self._robot_state.n_jnts()
+        self._n_contacts = self._robot_state.n_contacts() # we assume same n contacts for all rhcs for now
+
+        self._rhc_cmds.run()
+        self._rhc_pred.run()
+        self._rhc_refs.run()
+        self._rhc_status.run()
+        # we read rhc info now and just this time, since it's assumed to be static 
+        self._check_controllers_registered(retry=True) # blocking
+        # (we need controllers to be connected to read meaningful data)
+
+        self._rhc_status.rhc_static_info.synch_all(read=True,retry=True)
+        if self._use_gpu:
+            self._rhc_status.rhc_static_info.synch_mirror(from_gpu=False,non_blocking=False)
+        rhc_horizons=self._rhc_status.rhc_static_info.get("horizons",gpu=self._use_gpu)
+        rhc_nnodes=self._rhc_status.rhc_static_info.get("nnodes",gpu=self._use_gpu)
+        rhc_dts=self._rhc_status.rhc_static_info.get("dts",gpu=self._use_gpu)
+
+        # height sensor metadata (client side)
+        if self._env_opts["add_heightmap_obs"]:
+            self._height_grid_size = self._robot_state.height_sensor.grid_size
+            self._height_flat_dim = self._robot_state.height_sensor.n_cols
+        rhc_ncontacts=self._rhc_status.rhc_static_info.get("ncontacts",gpu=self._use_gpu)
+        robot_mass=self._rhc_status.rhc_static_info.get("robot_mass",gpu=self._use_gpu)
+        pred_node_idxs_rhc=self._rhc_status.rhc_static_info.get("pred_node_idx",gpu=self._use_gpu)
+
+        self._n_nodes_rhc=torch.round(rhc_nnodes) # we assume nodes are static during an env lifetime
+        self._rhc_horizons=rhc_horizons
+        self._rhc_dts=rhc_dts
+        self._n_contacts_rhc=rhc_ncontacts
+        self._rhc_robot_masses=robot_mass
+        if (self._rhc_robot_masses == 0).any():
+            zero_indices = torch.nonzero(self._rhc_robot_masses == 0, as_tuple=True)
+            print(zero_indices)  # This will print the indices of zero elements
+            Journal.log(self.__class__.__name__,
+                "_attach_to_shared_mem",
+                "Found at least one robot with 0 mass from RHC static info!!",
+                LogType.EXCEP,
+                throw_when_excep=True)
+
+        self._rhc_robot_weight=robot_mass*9.81
+        self._pred_node_idxs_rhc=pred_node_idxs_rhc
+        self._pred_horizon_rhc=self._pred_node_idxs_rhc*self._rhc_dts
+
+        # run server for agent commands
+        self._agent_refs = AgentRefs(namespace=self._namespace,
+                                is_server=True,
+                                n_robots=self._n_envs,
+                                n_jnts=self._robot_state.n_jnts(),
+                                n_contacts=self._robot_state.n_contacts(),
+                                contact_names=self._robot_state.contact_names(),
+                                q_remapping=None,
+                                with_gpu_mirror=self._use_gpu,
+                                force_reconnection=True,
+                                safe=False,
+                                verbose=self._verbose,
+                                vlevel=self._vlevel,
+                                fill_value=0)
+        self._agent_refs.run()
+        q_init_agent_refs=torch.full_like(self._robot_state.root_state.get(data_type="q", gpu=self._use_gpu),fill_value=0.0)
+        q_init_agent_refs[:, 0]=1.0
+        self._agent_refs.rob_refs.root_state.set(data_type="q", data=q_init_agent_refs,
+                gpu=self._use_gpu)
+        if self._use_gpu:
+            self._agent_refs.rob_refs.root_state.synch_mirror(from_gpu=True,non_blocking=True) 
+        self._agent_refs.rob_refs.root_state.synch_all(read=False, retry=True)
+        # episode steps counters (for detecting episode truncations for 
+        # time limits) 
+        self._ep_timeout_counter = EpisodesCounter(namespace=self._namespace,
+                            n_envs=self._n_envs,
+                            n_steps_lb=self._env_opts["episode_timeout_lb"],
+                            n_steps_ub=self._env_opts["episode_timeout_ub"],
+                            randomize_offsets_at_startup=True, # this has to be randomized
+                            is_server=True,
+                            verbose=self._verbose,
+                            vlevel=self._vlevel,
+                            safe=True,
+                            force_reconnection=True,
+                            with_gpu_mirror=self._use_gpu,
+                            debug=self._debug) # handles step counter through episodes and through envs
+        self._ep_timeout_counter.run()
+        self._task_rand_counter = TaskRandCounter(namespace=self._namespace,
+                            n_envs=self._n_envs,
+                            n_steps_lb=self._env_opts["n_steps_task_rand_lb"],
+                            n_steps_ub=self._env_opts["n_steps_task_rand_ub"],
+                            randomize_offsets_at_startup=False, # not necessary since it will be synched with the timeout counter
+                            is_server=True,
+                            verbose=self._verbose,
+                            vlevel=self._vlevel,
+                            safe=True,
+                            force_reconnection=True,
+                            with_gpu_mirror=self._use_gpu,
+                            debug=self._debug) # handles step counter through episodes and through envs
+        self._task_rand_counter.run()
+        self._task_rand_counter.sync_counters(other_counter=self._ep_timeout_counter)
+        if self._env_opts["use_random_trunc"]:
+            self._rand_trunc_counter=RandomTruncCounter(namespace=self._namespace,
+                            n_envs=self._n_envs,
+                            n_steps_lb=self._env_opts["random_trunc_freq"]-self._env_opts["random_trunc_freq_delta"],
+                            n_steps_ub=self._env_opts["random_trunc_freq"],
+                            randomize_offsets_at_startup=True,
+                            is_server=True,
+                            verbose=self._verbose,
+                            vlevel=self._vlevel,
+                            safe=True,
+                            force_reconnection=True,
+                            with_gpu_mirror=self._use_gpu,
+                            debug=False)
+            self._rand_trunc_counter.run()
+            # self._rand_trunc_counter.sync_counters(other_counter=self._ep_timeout_counter)
+        if self._env_opts["use_random_safety_reset"]:
+            self._rand_safety_reset_counter=SafetyRandResetsCounter(namespace=self._namespace,
+                            n_envs=self._n_envs,
+                            n_steps_lb=self._env_opts["random_reset_freq"],
+                            n_steps_ub=self._env_opts["random_reset_freq"],
+                            randomize_offsets_at_startup=True,
+                            is_server=True,
+                            verbose=self._verbose,
+                            vlevel=self._vlevel,
+                            safe=True,
+                            force_reconnection=True,
+                            with_gpu_mirror=self._use_gpu,
+                            debug=False)
+            self._rand_safety_reset_counter.run()
+            # self._rand_safety_reset_counter.sync_counters(other_counter=self._ep_timeout_counter)
+
+        # timer to track abs time in each env (reset logic to be implemented in child)
+        self._substep_abs_counter = SubStepAbsCounter(namespace=self._namespace,
+                            n_envs=self._n_envs,
+                            n_steps_lb=1e9,
+                            n_steps_ub=1e9,
+                            randomize_offsets_at_startup=True, # randomizing startup offsets
+                            is_server=True,
+                            verbose=self._verbose,
+                            vlevel=self._vlevel,
+                            safe=True,
+                            force_reconnection=True,
+                            with_gpu_mirror=self._use_gpu,
+                            debug=self._debug)
+        self._substep_abs_counter.run()
+
+        # debug data servers
+        traing_env_param_dict = {}
+        traing_env_param_dict["use_gpu"] = self._use_gpu
+        traing_env_param_dict["debug"] = self._is_debug
+        traing_env_param_dict["n_preinit_steps"] = self._env_opts["n_preinit_steps"]
+        traing_env_param_dict["n_preinit_steps"] = self._n_envs
+        
+        self._training_sim_info = SharedTrainingEnvInfo(namespace=self._namespace,
+                is_server=True, 
+                training_env_params_dict=traing_env_param_dict,
+                safe=False,
+                force_reconnection=True,
+                verbose=self._verbose,
+                vlevel=self._vlevel)
+        self._training_sim_info.run()
+
+        self._observed_jnt_names=self.set_observed_joints()
+        self._set_jnts_blacklist_pattern()
+        self._update_jnt_blacklist()
+
+        self._prev_root_p_substep=self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu).clone()
+        self._prev_root_q_substep=self._robot_state.root_state.get(data_type="q",gpu=self._use_gpu).clone()
+        self._prev_root_p_step=self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu).clone()
+        self._prev_root_q_step=self._robot_state.root_state.get(data_type="q",gpu=self._use_gpu).clone()
+        
+    def _activate_rhc_controllers(self):
+        self._rhc_status.activation_state.get_torch_mirror()[:, :] = True
+        self._rhc_status.activation_state.synch_all(read=False, retry=True) # activates all controllers
+    
+    def _synch_state(self,
+            gpu: bool = True):
+
+        # read from shared memory on CPU
+        # robot state
+        self._robot_state.root_state.synch_all(read = True, retry = True)
+        self._robot_state.jnts_state.synch_all(read = True, retry = True)
+        # rhc cmds
+        self._rhc_cmds.root_state.synch_all(read = True, retry = True)
+        self._rhc_cmds.jnts_state.synch_all(read = True, retry = True)
+        self._rhc_cmds.contact_wrenches.synch_all(read = True, retry = True)
+        # rhc pred
+        self._rhc_pred.root_state.synch_all(read = True, retry = True)
+        # self._rhc_pred.jnts_state.synch_all(read = True, retry = True)
+        # self._rhc_pred.contact_wrenches.synch_all(read = True, retry = True)
+        # refs for root link and contacts
+        self._rhc_refs.rob_refs.root_state.synch_all(read = True, retry = True)
+        self._rhc_refs.contact_flags.synch_all(read = True, retry = True)
+        self._rhc_refs.flight_info.synch_all(read = True, retry = True)
+        self._rhc_refs.flight_settings_req.synch_all(read = True, retry = True)
+        self._rhc_refs.rob_refs.contact_pos.synch_all(read = True, retry = True)
+        # rhc cost
+        self._rhc_status.rhc_cost.synch_all(read = True, retry = True)
+        # rhc constr. violations
+        self._rhc_status.rhc_constr_viol.synch_all(read = True, retry = True)
+        # failure states
+        self._rhc_status.fails.synch_all(read = True, retry = True)
+        # tot cost and cnstr viol on nodes + step variable
+        self._rhc_status.rhc_nodes_cost.synch_all(read = True, retry = True)
+        self._rhc_status.rhc_nodes_constr_viol.synch_all(read = True, retry = True)
+        self._rhc_status.rhc_fcn.synch_all(read = True, retry = True)
+        self._rhc_status.rhc_fail_idx.synch_all(read = True, retry = True)
+        if self._env_opts["add_heightmap_obs"]:
+            self._robot_state.height_sensor.synch_all(read=True, retry=True)
+        if gpu:
+            # copies data to "mirror" on GPU --> we can do it non-blocking since
+            # in this direction it should be safe
+            self._robot_state.root_state.synch_mirror(from_gpu=False,non_blocking=True) # copies shared data on GPU
+            self._robot_state.jnts_state.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_cmds.root_state.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_cmds.jnts_state.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_cmds.contact_wrenches.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_pred.root_state.synch_mirror(from_gpu=False,non_blocking=True)
+            # self._rhc_pred.jnts_state.synch_mirror(from_gpu=False,non_blocking=True)
+            # self._rhc_pred.contact_wrenches.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_refs.rob_refs.root_state.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_refs.contact_flags.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_refs.rob_refs.contact_pos.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_refs.flight_info.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_refs.flight_settings_req.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_status.rhc_cost.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_status.rhc_constr_viol.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_status.fails.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_status.rhc_nodes_cost.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_status.rhc_nodes_constr_viol.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_status.rhc_fcn.synch_mirror(from_gpu=False,non_blocking=True)
+            self._rhc_status.rhc_fail_idx.synch_mirror(from_gpu=False,non_blocking=True)
+            if self._env_opts["add_heightmap_obs"]:
+                self._robot_state.height_sensor.synch_mirror(from_gpu=False, non_blocking=True)
+            torch.cuda.synchronize() # ensuring that all the streams on the GPU are completed \
+            # before the CPU continues execution
+    
+    def _override_refs(self,
+            env_indxs: torch.Tensor = None):
+
+        # just used for setting agent refs externally (i.e. from shared mem on CPU)
+        self._agent_refs.rob_refs.root_state.synch_all(read=True,retry=True) # first read from mem
+        if self._use_gpu:
+            # copies latest refs to GPU 
+            self._agent_refs.rob_refs.root_state.synch_mirror(from_gpu=False,non_blocking=False) 
+
+    def _clamp_obs(self, 
+            obs: torch.Tensor):
+        if self._is_debug:
+            self._check_finite(obs, "observations", False)
+        torch.nan_to_num(input=obs, out=obs, nan=self._obs_threshold_ub, 
+            posinf=self._obs_threshold_ub, 
+            neginf=self._obs_threshold_lb) # prevent nans
+
+        obs.clamp_(self._obs_threshold_lb, self._obs_threshold_ub)
+    
+    def _clamp_rewards(self, 
+            rewards: torch.Tensor):
+        if self._is_debug:
+            self._check_finite(rewards, "rewards", False)
+        torch.nan_to_num(input=rewards, out=rewards, nan=0.0, 
+            posinf=None, 
+            neginf=None) # prevent nans
+        rewards.clamp_(self._reward_thresh_lb, self._reward_thresh_ub)
+
+    def get_actions_lb(self):
+        return self._actions_lb
+
+    def get_actions_ub(self):
+        return self._actions_ub
+    
+    def get_actions_scale(self):
+        return self._actions_scale
+    
+    def get_actions_offset(self):
+        return self._actions_offset
+    
+    def get_obs_lb(self):
+        return self._obs_lb
+
+    def get_obs_ub(self):
+        return self._obs_ub
+    
+    def get_obs_scale(self):
+        self._obs_scale = (self._obs_ub - self._obs_lb)/2.0
+        return self._obs_scale
+    
+    def get_obs_offset(self):
+        self._obs_offset = (self._obs_ub + self._obs_lb)/2.0
+        return self._obs_offset
+    
+    def switch_random_reset(self, on: bool = True):
+        self._random_reset_active=on
+
+    def set_jnts_remapping(self, 
+        remapping: List = None):
+
+        self._jnts_remapping=remapping
+        if self._jnts_remapping is not None:
+            self._robot_state.set_jnts_remapping(jnts_remapping=self._jnts_remapping)
+            self._rhc_cmds.set_jnts_remapping(jnts_remapping=self._jnts_remapping)
+            self._rhc_pred.set_jnts_remapping(jnts_remapping=self._jnts_remapping)
+            # we need to also update the list of observed joints to match
+            available_joints=self._robot_state.jnt_names()
+            # the remapping ordering 
+            self._observed_jnt_names=[]
+            for i in range(len(available_joints)):
+                self._observed_jnt_names.append(available_joints[self._jnts_remapping[i]])
+
+            self._update_jnt_blacklist()
+
+            updated_obs_names=self._get_obs_names() # get updated obs names (should use get_observed_joints
+            # internally, so that jnt names are updated)
+
+            # also update jnt obs names on shared memory
+            names_old=self._obs.get_obs_names()
+            names_old_next=self._next_obs.get_obs_names()
+            names_old[:]=updated_obs_names
+            names_old_next[:]=updated_obs_names
+            self._obs.update_names()
+            self._next_obs.update_names()
+
+            # also update 
+            if "Obs" in self.custom_db_data:
+                db_obs_names=self.custom_db_data["Obs"].data_names()
+                db_obs_names[:]=updated_obs_names
+
+    def _check_finite(self, 
+                tensor: torch.Tensor,
+                name: str, 
+                throw: bool = False):
+        if not torch.isfinite(tensor).all().item():
+            exception = f"Found nonfinite elements in {name} tensor!!"            
+            non_finite_idxs=torch.nonzero(~torch.isfinite(tensor))
+            n_nonf_elems=non_finite_idxs.shape[0]
+
+            if name=="observations":
+                for i in range(n_nonf_elems):
+                    db_msg=f"{self.obs_names()[non_finite_idxs[i,1]]} (env. {non_finite_idxs[i,0]}):" + \
+                        f" {tensor[non_finite_idxs[i,0],non_finite_idxs[i,1]].item()}"
+                    print(db_msg)
+            if name=="rewards":
+                for i in range(n_nonf_elems):
+                    db_msg=f"{self.sub_rew_names()[non_finite_idxs[i,1]]} (env. {non_finite_idxs[i,0]}):" + \
+                        f" {tensor[non_finite_idxs[i,0],non_finite_idxs[i,1]].item()}"
+                    print(db_msg)
+            print(tensor)
+            Journal.log(self.__class__.__name__,
+                "_check_finite",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = throw)
+            
+    def _check_controllers_registered(self, 
+                retry: bool = False):
+
+        if retry:
+            self._rhc_status.controllers_counter.synch_all(read=True, retry=True)
+            n_connected_controllers = self._rhc_status.controllers_counter.get_torch_mirror()[0, 0].item()
+            while not (n_connected_controllers == self._n_envs):
+                warn = f"Expected {self._n_envs} controllers to be connected during training, " + \
+                    f"but got {n_connected_controllers}. Will wait for all to be connected..."
+                Journal.log(self.__class__.__name__,
+                    "_check_controllers_registered",
+                    warn,
+                    LogType.WARN,
+                    throw_when_excep = False)
+                nsecs = int(2 * 1000000000)
+                PerfSleep.thread_sleep(nsecs) 
+                self._rhc_status.controllers_counter.synch_all(read=True, retry=True)
+                n_connected_controllers = self._rhc_status.controllers_counter.get_torch_mirror()[0, 0].item()
+            info = f"All {n_connected_controllers} controllers connected!"
+            Journal.log(self.__class__.__name__,
+                "_check_controllers_registered",
+                info,
+                LogType.INFO,
+                throw_when_excep = False)
+            return True
+        else:
+            self._rhc_status.controllers_counter.synch_all(read=True, retry=True)
+            n_connected_controllers = self._rhc_status.controllers_counter.get_torch_mirror()[0, 0].item()
+            if not (n_connected_controllers == self._n_envs):
+                exception = f"Expected {self._n_envs} controllers to be connected during training, " + \
+                    f"but got {n_connected_controllers}. Aborting..."
+                Journal.log(self.__class__.__name__,
+                    "_check_controllers_registered",
+                    exception,
+                    LogType.EXCEP,
+                    throw_when_excep = False)
+                return False
+            return True
+    
+    def _check_truncations(self):
+        
+        self._check_sub_truncations()
+        sub_truncations = self._sub_truncations.get_torch_mirror(gpu=self._use_gpu)
+        truncations = self._truncations.get_torch_mirror(gpu=self._use_gpu)
+        truncations[:, :] = torch.any(sub_truncations,dim=1,keepdim=True)
+
+    def _check_terminations(self):
+        
+        self._check_sub_terminations()
+        sub_terminations = self._sub_terminations.get_torch_mirror(gpu=self._use_gpu)
+        terminations = self._terminations.get_torch_mirror(gpu=self._use_gpu)
+        terminations[:, :] = torch.any(sub_terminations,dim=1,keepdim=True)
+
+    def _check_sub_truncations(self):
+        # default behaviour-> to be overriden by child
+        sub_truncations = self._sub_truncations.get_torch_mirror(gpu=self._use_gpu)
+        sub_truncations[:, 0:1]=self._ep_timeout_counter.time_limits_reached()
+
+    def _check_sub_terminations(self):
+        # default behaviour-> to be overriden by child
+        sub_terminations = self._sub_terminations.get_torch_mirror(gpu=self._use_gpu)
+        robot_q_meas = self._robot_state.root_state.get(data_type="q",gpu=self._use_gpu)
+        robot_q_pred = self._rhc_cmds.root_state.get(data_type="q",gpu=self._use_gpu)
+
+        # terminate when either the real robot or the prediction from the MPC are capsized
+        check_capsize(quat=robot_q_meas,max_angle=self._max_pitch_angle,
+            output_t=self._is_capsized)
+        check_capsize(quat=robot_q_pred,max_angle=self._max_pitch_angle,
+            output_t=self._is_rhc_capsized)
+        
+        sub_terminations[:, 0:1] = self._rhc_status.fails.get_torch_mirror(gpu=self._use_gpu)
+        sub_terminations[:, 1:2] = self._is_capsized
+        sub_terminations[:, 2:3] = self._is_rhc_capsized
+
+    def is_action_continuous(self):
+        return self._is_continuous_actions
+    
+    def is_action_discrete(self):
+        return ~self._is_continuous_actions
+
+    @abstractmethod
+    def _pre_substep(self):
+        pass
+    
+    @abstractmethod
+    def _custom_post_step(self,episode_finished):
+        pass
+    
+    @abstractmethod
+    def _custom_post_substp_post_rew(self):
+        pass
+    
+    @abstractmethod
+    def _custom_post_substp_pre_rew(self):
+        pass
+
+    @abstractmethod
+    def _apply_actions_to_rhc(self):
+        pass
+    
+    def _override_actions_with_demo(self):
+        pass
+
+    @abstractmethod
+    def _compute_substep_rewards(self):
+        pass
+    
+    @abstractmethod
+    def _set_substep_rew(self):
+        pass
+
+    @abstractmethod
+    def _set_substep_obs(self):
+        pass
+
+    @abstractmethod
+    def _compute_step_rewards(self):
+        pass
+    
+    @abstractmethod
+    def _fill_substep_obs(self,
+            obs: torch.Tensor):
+        pass
+    
+    @abstractmethod
+    def _fill_step_obs(self,
+            obs: torch.Tensor):
+        pass
+
+    @abstractmethod
+    def _randomize_task_refs(self,
+                env_indxs: torch.Tensor = None):
+        pass
+
+    def _custom_post_init(self):
+        pass
+    
+    def _get_avrg_substep_root_twist(self, 
+            out: torch.Tensor,
+            base_loc: bool = True):
+        # to be called at each substep
+        robot_p_meas = self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu)
+        robot_q_meas = self._robot_state.root_state.get(data_type="q",gpu=self._use_gpu)
+
+        root_v_avrg_w=(robot_p_meas-self._prev_root_p_substep)/self._substep_dt
+        root_omega_avrg_w=quaternion_to_angular_velocity(q_diff=quaternion_difference(self._prev_root_q_substep,robot_q_meas),\
+            dt=self._substep_dt)
+        twist_w=torch.cat((root_v_avrg_w, 
+            root_omega_avrg_w), 
+            dim=1)
+        if not base_loc:
+            self._prev_root_p_substep[:, :]=robot_p_meas
+            self._prev_root_q_substep[:, :]=robot_q_meas
+            out[:, :]=twist_w
+        # rotate using the current (end-of-substep) orientation for consistency with other signals
+        world2base_frame(t_w=twist_w, q_b=robot_q_meas, t_out=out)
+        self._prev_root_p_substep[:, :]=robot_p_meas
+        self._prev_root_q_substep[:, :]=robot_q_meas
+
+    def _get_avrg_step_root_twist(self, 
+            out: torch.Tensor,
+            base_loc: bool = True):
+        # to be called after substeps of actions repeats
+        robot_p_meas = self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu)
+        robot_q_meas = self._robot_state.root_state.get(data_type="q",gpu=self._use_gpu)
+
+        dt=self._substep_dt*self._action_repeat # accounting for frame skipping
+        root_v_avrg_w=(robot_p_meas-self._prev_root_p_step)/(dt)
+        root_omega_avrg_w=quaternion_to_angular_velocity(q_diff=quaternion_difference(self._prev_root_q_step,robot_q_meas),\
+            dt=dt)
+        twist_w=torch.cat((root_v_avrg_w, 
+            root_omega_avrg_w), 
+            dim=1)
+        if not base_loc:
+            out[:, :]=twist_w
+        # rotate using the current (end-of-step) orientation for consistency with other signals
+        world2base_frame(t_w=twist_w, q_b=robot_q_meas, t_out=out)
+
+    def _get_avrg_rhc_root_twist(self,
+            out: torch.Tensor,
+            base_loc: bool = True):
+        
+        rhc_root_p =self._rhc_cmds.root_state.get(data_type="p",gpu=self._use_gpu)
+        rhc_root_q =self._rhc_cmds.root_state.get(data_type="q",gpu=self._use_gpu)
+        rhc_root_p_pred =self._rhc_pred.root_state.get(data_type="p",gpu=self._use_gpu)
+        rhc_root_q_pred =self._rhc_pred.root_state.get(data_type="q",gpu=self._use_gpu)
+
+        rhc_root_v_avrg_rhc_w=(rhc_root_p_pred-rhc_root_p)/self._pred_horizon_rhc
+        rhc_root_omega_avrg_rhc_w=quaternion_to_angular_velocity(q_diff=quaternion_difference(rhc_root_q,rhc_root_q_pred),\
+            dt=self._pred_horizon_rhc)
+    
+        rhc_pred_avrg_twist_rhc_w = torch.cat((rhc_root_v_avrg_rhc_w, 
+            rhc_root_omega_avrg_rhc_w), 
+            dim=1)
+        if not base_loc:
+            out[:, :]=rhc_pred_avrg_twist_rhc_w
+        # to rhc base frame (using first node as reference)
+        world2base_frame(t_w=rhc_pred_avrg_twist_rhc_w, q_b=rhc_root_q, t_out=out)

twist_tracking_env.py

@@ -0,0 +1,1396 @@
+from typing import Dict
+
+import os
+
+import torch
+
+from EigenIPC.PyEigenIPC import VLevel, LogType, Journal
+
+from mpc_hive.utilities.shared_data.rhc_data import RobotState, RhcStatus, RhcRefs
+from mpc_hive.utilities.math_utils_torch import world2base_frame, base2world_frame, w2hor_frame
+
+from aug_mpc.utils.sys_utils import PathsGetter
+from aug_mpc.utils.timers import PeriodicTimer
+from aug_mpc.utils.episodic_data import EpisodicData
+from aug_mpc.utils.signal_smoother import ExponentialSignalSmoother
+from aug_mpc.utils.math_utils import check_capsize
+from aug_mpc.training_envs.training_env_base import AugMPCTrainingEnvBase
+
+class TwistTrackingEnv(AugMPCTrainingEnvBase):
+    """Base AugMPC training env that tracks commanded twists by pushing velocity and contact targets into the RHC controller while handling locomotion rewards/resets."""
+
+    def __init__(self,
+            namespace: str,
+            actions_dim: int = 10,
+            verbose: bool = False,
+            vlevel: VLevel = VLevel.V1,
+            use_gpu: bool = True,
+            dtype: torch.dtype = torch.float32,
+            debug: bool = True,
+            override_agent_refs: bool = False,
+            timeout_ms: int = 60000,
+            env_opts: Dict = {}):
+        
+        env_name = "LinVelTrack"
+        device = "cuda" if use_gpu else "cpu"
+
+        self._add_env_opt(env_opts, "srew_drescaling", 
+            False)
+        
+        self._add_env_opt(env_opts, "step_thresh", 0.) # when step action < thresh, a step is requested
+
+        # counters settings
+        self._add_env_opt(env_opts, "single_task_ref_per_episode", 
+            True # if True, the task ref is constant over the episode (ie
+            # episodes are truncated when task is changed) 
+            )
+        self._add_env_opt(env_opts, "add_angvel_ref_rand", default=True) # randomize also agular vel ref (just z component)
+
+        self._add_env_opt(env_opts, "episode_timeout_lb", 
+            1024)
+        self._add_env_opt(env_opts, "episode_timeout_ub", 
+            1024)
+        self._add_env_opt(env_opts, "n_steps_task_rand_lb", 
+            512)
+        self._add_env_opt(env_opts, "n_steps_task_rand_ub", 
+            512)
+        self._add_env_opt(env_opts, "use_random_safety_reset", 
+            True)
+        self._add_env_opt(env_opts, "random_reset_freq", 
+            10) # a random reset once every n-episodes (per env)
+        self._add_env_opt(env_opts, "use_random_trunc", 
+            True)
+        self._add_env_opt(env_opts, "random_trunc_freq", 
+            env_opts["episode_timeout_ub"]*5) # to remove temporal correlations between envs
+        self._add_env_opt(env_opts, "random_trunc_freq_delta", 
+            env_opts["episode_timeout_ub"]*2)  # to randomize trunc frequency between envs
+    
+        if not env_opts["single_task_ref_per_episode"]:
+            env_opts["random_reset_freq"]=int(env_opts["random_reset_freq"]/\
+                (env_opts["episode_timeout_lb"]/float(env_opts["n_steps_task_rand_lb"])))
+        
+        self._add_env_opt(env_opts, "action_repeat", 1) # frame skipping (different agent action every action_repeat
+        # env substeps)
+        
+        self._add_env_opt(env_opts, "n_preinit_steps", 1) # n steps of the controllers to properly initialize everything
+        
+        self._add_env_opt(env_opts, "vec_ep_freq_metrics_db", 1) # n eps over which debug metrics are reported
+        self._add_env_opt(env_opts, "demo_envs_perc", 0.0)
+        self._add_env_opt(env_opts, "max_cmd_v", 1.5) # maximum cmd v for lin v actions (single component)
+        self._add_env_opt(env_opts, "max_cmd_omega", 1.0) # maximum cmd v for omega v actions (single component)
+
+        # action smoothing
+        self._add_env_opt(env_opts, "use_action_smoothing", False)
+        self._add_env_opt(env_opts, "smoothing_horizon_c", 0.01)
+        self._add_env_opt(env_opts, "smoothing_horizon_d", 0.03)
+
+        # whether to smooth vel error signal
+        self._add_env_opt(env_opts, "use_track_reward_smoother", False)
+        self._add_env_opt(env_opts, "smoothing_horizon_vel_err", 0.08)
+        self._add_env_opt(env_opts, "track_rew_smoother", None)
+
+        # rewards
+        self._reward_map={}
+        self._reward_lb_map={}
+
+        self._add_env_opt(env_opts, "reward_lb_default", -0.5)
+        self._add_env_opt(env_opts, "reward_ub_default", 1e6)
+
+        self._add_env_opt(env_opts, "task_error_reward_lb", -0.5)
+        self._add_env_opt(env_opts, "CoT_reward_lb", -0.5)
+        self._add_env_opt(env_opts, "power_reward_lb", -0.5)
+        self._add_env_opt(env_opts, "action_rate_reward_lb", -0.5)
+        self._add_env_opt(env_opts, "jnt_vel_reward_lb", -0.5)
+        self._add_env_opt(env_opts, "rhc_avrg_vel_reward_lb", -0.5)
+
+        self._add_env_opt(env_opts, "add_power_reward", False)
+        self._add_env_opt(env_opts, "add_CoT_reward", True)
+        self._add_env_opt(env_opts, "use_CoT_wrt_ref", False)
+        self._add_env_opt(env_opts, "add_action_rate_reward", True)
+        self._add_env_opt(env_opts, "add_jnt_v_reward", False)
+
+        self._add_env_opt(env_opts, "use_rhc_avrg_vel_tracking", False)
+
+        # task tracking
+        self._add_env_opt(env_opts, "use_relative_error", default=False) # use relative vel error (wrt current task norm)
+        self._add_env_opt(env_opts, "directional_tracking", default=True) # whether to compute tracking error based on reference direction
+        # if env_opts["add_angvel_ref_rand"]:
+        #     env_opts["directional_tracking"]=False
+
+        self._add_env_opt(env_opts, "use_L1_norm", default=True) # whether to use L1 norm for the error (otherwise L2)
+        self._add_env_opt(env_opts, "use_exp_track_rew", default=True) # whether to use a reward of the form A*e^(B*x), 
+        # otherwise A*(1-B*x)
+
+        self._add_env_opt(env_opts, "use_fail_idx_weight", default=False)
+        self._add_env_opt(env_opts, "task_track_offset_exp", default=1.0)
+        self._add_env_opt(env_opts, "task_track_scale_exp", default=5.0)
+        self._add_env_opt(env_opts, "task_track_offset", default=1.0)
+        self._add_env_opt(env_opts, "task_track_scale", default=1.5)
+        self._add_env_opt(env_opts, "task_track_front_weight", default=1.0)
+        self._add_env_opt(env_opts, "task_track_lat_weight", default=0.05)
+        self._add_env_opt(env_opts, "task_track_vert_weight", default=0.05)
+        self._add_env_opt(env_opts, "task_track_omega_z_weight", default=0.4)
+        self._add_env_opt(env_opts, "task_track_omega_x_weight", default=0.05)
+        self._add_env_opt(env_opts, "task_track_omega_y_weight", default=0.05)
+        # if env_opts["add_angvel_ref_rand"]:
+        #     env_opts["task_track_omega_x_weight"]=0.0
+        #     env_opts["task_track_omega_y_weight"]=0.0
+        #     env_opts["task_track_omega_z_weight"]=1.0
+
+        # task pred tracking
+        self._add_env_opt(env_opts, "task_pred_track_offset", default=1.0)
+        self._add_env_opt(env_opts, "task_pred_track_scale", default=3.0)
+
+        # energy penalties
+        self._add_env_opt(env_opts, "CoT_offset", default=0.3)
+        self._add_env_opt(env_opts, "CoT_scale", default=0.5)
+        self._add_env_opt(env_opts, "power_offset", default=0.1)
+        self._add_env_opt(env_opts, "power_scale", default=8e-4)
+
+        # action rate penalty
+        self._add_env_opt(env_opts, "action_rate_offset", default=0.1)
+        self._add_env_opt(env_opts, "action_rate_scale", default=2.0)
+        self._add_env_opt(env_opts, "action_rate_rew_d_weight", default=0.1)
+        self._add_env_opt(env_opts, "action_rate_rew_c_weight", default=1.0)
+
+        # jnt vel penalty
+        self._add_env_opt(env_opts, "jnt_vel_offset", default=0.1)
+        self._add_env_opt(env_opts, "jnt_vel_scale", default=2.0)
+
+        # terminations
+        self._add_env_opt(env_opts, "add_term_mpc_capsize", default=False) # add termination based on mpc capsizing prediction
+
+        # observations
+        self._add_env_opt(env_opts, "rhc_fail_idx_scale", default=1.0)
+        self._add_env_opt(env_opts, "use_action_history", default=True) # whether to add information on past actions to obs
+        self._add_env_opt(env_opts, "add_prev_actions_stats_to_obs", default=False) # add actions std, mean + last action over a horizon to obs (if self._use_action_history True)
+        self._add_env_opt(env_opts, "actions_history_size", default=3)
+        
+        self._add_env_opt(env_opts, "add_mpc_contact_f_to_obs", default=True) # add estimate vertical contact f to obs
+        self._add_env_opt(env_opts, "add_fail_idx_to_obs", default=True) # we need to obserse mpc failure idx to correlate it with terminations
+        
+        self._add_env_opt(env_opts, "use_linvel_from_rhc", default=True) # no lin vel meas available, we use est. from mpc
+        self._add_env_opt(env_opts, "add_flight_info", default=True) # add feedback info on pos, remamining duration, length, 
+        # apex and landing height of flight phases
+        self._add_env_opt(env_opts, "add_flight_settings", default=False) # add feedback info on current flight requests for mpc
+
+        self._add_env_opt(env_opts, "use_prob_based_stepping", default=False) # interpret actions as stepping prob (never worked)
+
+        self._add_env_opt(env_opts, "add_rhc_cmds_to_obs", default=True) # add the rhc cmds which are being applied now to the robot
+
+        if not "add_periodic_clock_to_obs" in env_opts:
+            # add a sin/cos clock to obs (useful if task is explicitly
+            # time-dependent)
+            self._add_env_opt(env_opts, "add_periodic_clock_to_obs", default=False) 
+
+        self._add_env_opt(env_opts, "add_heightmap_obs", default=False)         
+
+        # temporarily creating robot state client to get some data
+        robot_state_tmp = RobotState(namespace=namespace,
+                                is_server=False, 
+                                safe=False,
+                                verbose=verbose,
+                                vlevel=vlevel,
+                                with_gpu_mirror=False,
+                                with_torch_view=False,
+                                enable_height_sensor=env_opts["add_heightmap_obs"])
+        robot_state_tmp.run()
+        rhc_status_tmp = RhcStatus(is_server=False,
+                        namespace=namespace, 
+                        verbose=verbose, 
+                        vlevel=vlevel,
+                        with_torch_view=False, 
+                        with_gpu_mirror=False)
+        rhc_status_tmp.run()
+        rhc_refs_tmp = RhcRefs(namespace=namespace,
+                            is_server=False,
+                            safe=False,
+                            verbose=verbose,
+                            vlevel=vlevel,
+                            with_gpu_mirror=False,
+                            with_torch_view=False)
+        rhc_refs_tmp.run()
+        n_jnts = robot_state_tmp.n_jnts()
+        self._contact_names = robot_state_tmp.contact_names()
+        self._n_contacts = len(self._contact_names)
+        self._flight_info_size=rhc_refs_tmp.flight_info.n_cols
+        self._flight_setting_size=rhc_refs_tmp.flight_settings_req.n_cols
+        # height sensor metadata (if present)
+        self._height_grid_size = None
+        self._height_flat_dim = 0
+        if env_opts["add_heightmap_obs"]:
+            self._height_grid_size = robot_state_tmp.height_sensor.grid_size
+            self._height_flat_dim = robot_state_tmp.height_sensor.n_cols
+
+        robot_state_tmp.close()
+        rhc_status_tmp.close()
+        rhc_refs_tmp.close()
+
+        # defining obs dimension
+        obs_dim=3 # normalized gravity vector in base frame
+        obs_dim+=6 # meas twist in base frame
+        obs_dim+=2*n_jnts # joint pos + vel
+        if env_opts["add_mpc_contact_f_to_obs"]:
+            obs_dim+=3*self._n_contacts
+        obs_dim+=6 # twist reference in base frame frame
+        if env_opts["add_fail_idx_to_obs"]:
+            obs_dim+=1 # rhc controller failure index
+        if env_opts["add_term_mpc_capsize"]: 
+            obs_dim+=3 # gravity vec from mpc
+        if env_opts["use_rhc_avrg_vel_tracking"]:
+            obs_dim+=6 # mpc avrg twist
+        if env_opts["add_flight_info"]: # contact pos, remaining duration, length, apex, landing height, landing dx, dy
+            obs_dim+=self._flight_info_size
+        if env_opts["add_flight_settings"]:
+            obs_dim+=self._flight_setting_size
+        if env_opts["add_rhc_cmds_to_obs"]:
+            obs_dim+=3*n_jnts 
+        if env_opts["use_action_history"]:
+            if env_opts["add_prev_actions_stats_to_obs"]:
+                obs_dim+=3*actions_dim # previous agent actions statistics (mean, std + last action)
+            else: # full action history
+                obs_dim+=env_opts["actions_history_size"]*actions_dim
+        if env_opts["use_action_smoothing"]:
+            obs_dim+=actions_dim # it's better to also add the smoothed actions as obs
+        if env_opts["add_periodic_clock_to_obs"]:
+            obs_dim+=2
+        if env_opts["add_heightmap_obs"]:
+            obs_dim+=self._height_flat_dim
+        # Agent task reference
+        self._add_env_opt(env_opts, "use_pof0", default=True) # with some prob, references will be null
+        self._add_env_opt(env_opts, "pof0_linvel", default=0.3) # [0, 1] prob of both linvel and omega refs being null(from bernoulli distr)
+        self._add_env_opt(env_opts, "pof0_omega", default=0.3) # [0, 1] prob of both linvel and omega refs being null(from bernoulli distr)
+        self._add_env_opt(env_opts, "max_linvel_ref", default=0.3) # m/s
+        self._add_env_opt(env_opts, "max_angvel_ref", default=0.0) # rad/s
+        if env_opts["add_angvel_ref_rand"]:   
+            env_opts["max_angvel_ref"]=0.4
+
+        # ready to init base class
+        self._this_child_path = os.path.abspath(__file__)
+        AugMPCTrainingEnvBase.__init__(self,
+                    namespace=namespace,
+                    obs_dim=obs_dim,
+                    actions_dim=actions_dim,
+                    env_name=env_name,
+                    verbose=verbose,
+                    vlevel=vlevel,
+                    use_gpu=use_gpu,
+                    dtype=dtype,
+                    debug=debug,
+                    override_agent_refs=override_agent_refs,
+                    timeout_ms=timeout_ms,
+                    env_opts=env_opts)
+
+    def _custom_post_init(self):
+
+        device = "cuda" if self._use_gpu else "cpu"
+
+        self._update_jnt_blacklist() # update blacklist for joints
+
+        # constant base-frame unit vectors (reuse to avoid per-call allocations)
+        self._base_x_dir = torch.zeros((self._n_envs, 3), dtype=self._dtype, device=device)
+        self._base_x_dir[:, 0] = 1.0
+        self._base_y_dir = torch.zeros((self._n_envs, 3), dtype=self._dtype, device=device)
+        self._base_y_dir[:, 1] = 1.0
+
+        self._twist_ref_lb = torch.full((1, 6), dtype=self._dtype, device=device,
+                            fill_value=-1.5) 
+        self._twist_ref_ub = torch.full((1, 6), dtype=self._dtype, device=device,
+                            fill_value=1.5)
+        
+        # task reference parameters (world frame)
+        # lin vel
+        self._twist_ref_lb[0, 0] = -self._env_opts["max_linvel_ref"]
+        self._twist_ref_lb[0, 1] = -self._env_opts["max_linvel_ref"]
+        self._twist_ref_lb[0, 2] = 0.0
+        self._twist_ref_ub[0, 0] = self._env_opts["max_linvel_ref"]
+        self._twist_ref_ub[0, 1] = self._env_opts["max_linvel_ref"]
+        self._twist_ref_ub[0, 2] = 0.0
+        # angular vel
+        self._twist_ref_lb[0, 3] = 0.0
+        self._twist_ref_lb[0, 4] = 0.0
+        self._twist_ref_lb[0, 5] = -self._env_opts["max_angvel_ref"]
+        self._twist_ref_ub[0, 3] = 0.0
+        self._twist_ref_ub[0, 4] = 0.0
+        self._twist_ref_ub[0, 5] = self._env_opts["max_angvel_ref"]
+
+        self._twist_ref_offset = (self._twist_ref_ub + self._twist_ref_lb)/2.0
+        self._twist_ref_scale = (self._twist_ref_ub - self._twist_ref_lb)/2.0
+
+        # adding some custom db info 
+        agent_twist_ref = self._agent_refs.rob_refs.root_state.get(data_type="twist",gpu=False)
+        agent_twist_ref_data = EpisodicData("AgentTwistRefs", agent_twist_ref, 
+            ["v_x", "v_y", "v_z", "omega_x", "omega_y", "omega_z"],
+            ep_vec_freq=self._env_opts["vec_ep_freq_metrics_db"],
+            store_transitions=self._full_db,
+            max_ep_duration=self._max_ep_length())
+        rhc_fail_idx = EpisodicData("RhcFailIdx", self._rhc_fail_idx(gpu=False), ["rhc_fail_idx"],
+            ep_vec_freq=self._env_opts["vec_ep_freq_metrics_db"],
+            store_transitions=self._full_db,
+            max_ep_duration=self._max_ep_length())
+        
+        f_names=[]
+        for contact in self._contact_names:
+            f_names.append(f"fc_{contact}_x_base_loc")
+            f_names.append(f"fc_{contact}_y_base_loc")
+            f_names.append(f"fc_{contact}_z_base_loc")
+        rhc_contact_f = EpisodicData("RhcContactForces", 
+            self._rhc_cmds.contact_wrenches.get(data_type="f",gpu=False), 
+            f_names,
+            ep_vec_freq=self._env_opts["vec_ep_freq_metrics_db"],
+            store_transitions=self._full_db,
+            max_ep_duration=self._max_ep_length())
+
+        self._pow_db_data=torch.full(size=(self._n_envs,2),
+                dtype=self._dtype, device="cpu",
+                fill_value=-1.0)
+        power_db = EpisodicData("Power", 
+            self._pow_db_data, 
+            ["CoT", "W"],
+            ep_vec_freq=self._env_opts["vec_ep_freq_metrics_db"],
+            store_transitions=self._full_db,
+            max_ep_duration=self._max_ep_length())
+        
+        self._track_error_db=torch.full_like(agent_twist_ref, fill_value=0.0)
+        task_err_db = EpisodicData("TrackingError", 
+            agent_twist_ref, 
+            ["e_vx", "e_vy", "e_vz", "e_omegax", "e_omegay", "e_omegaz"],
+            ep_vec_freq=self._env_opts["vec_ep_freq_metrics_db"],
+            store_transitions=self._full_db,
+            max_ep_duration=self._max_ep_length())
+
+        self._add_custom_db_data(db_data=agent_twist_ref_data)
+        self._add_custom_db_data(db_data=rhc_fail_idx)
+        self._add_custom_db_data(db_data=rhc_contact_f)
+        self._add_custom_db_data(db_data=power_db)
+        self._add_custom_db_data(db_data=task_err_db)
+
+        # rewards
+        self._task_err_weights = torch.full((1, 6), dtype=self._dtype, device=device,
+                            fill_value=0.0) 
+        if self._env_opts["directional_tracking"]:
+            self._task_err_weights[0, 0] = self._env_opts["task_track_front_weight"] # frontal
+            self._task_err_weights[0, 1] = self._env_opts["task_track_lat_weight"] # lateral
+            self._task_err_weights[0, 2] = self._env_opts["task_track_vert_weight"] # vertical
+            self._task_err_weights[0, 3] = self._env_opts["task_track_omega_x_weight"]
+            self._task_err_weights[0, 4] = self._env_opts["task_track_omega_y_weight"]
+            self._task_err_weights[0, 5] = self._env_opts["task_track_omega_z_weight"]
+        else:
+            self._task_err_weights[0, 0] = self._env_opts["task_track_front_weight"]
+            self._task_err_weights[0, 1] = self._env_opts["task_track_front_weight"]
+            self._task_err_weights[0, 2] = 0.1*self._env_opts["task_track_front_weight"]
+            self._task_err_weights[0, 3] = self._env_opts["task_track_omega_x_weight"]
+            self._task_err_weights[0, 4] = self._env_opts["task_track_omega_y_weight"]
+            self._task_err_weights[0, 5] = self._env_opts["task_track_omega_z_weight"]
+            
+        self._task_pred_err_weights = torch.full((1, 6), dtype=self._dtype, device=device,
+                            fill_value=0.0) 
+        if self._env_opts["directional_tracking"]:
+            self._task_pred_err_weights[0, 0] = self._env_opts["task_track_front_weight"]
+            self._task_pred_err_weights[0, 1] = self._env_opts["task_track_lat_weight"]
+            self._task_pred_err_weights[0, 2] = self._env_opts["task_track_vert_weight"]
+            self._task_pred_err_weights[0, 3] = self._env_opts["task_track_omega_x_weight"]
+            self._task_pred_err_weights[0, 4] = self._env_opts["task_track_omega_y_weight"]
+            self._task_pred_err_weights[0, 5] = self._env_opts["task_track_omega_z_weight"]
+        else:
+            self._task_pred_err_weights[0, 0] = self._env_opts["task_track_front_weight"]
+            self._task_pred_err_weights[0, 1] = self._env_opts["task_track_front_weight"]
+            self._task_pred_err_weights[0, 2] = 0.1*self._env_opts["task_track_front_weight"]
+            self._task_pred_err_weights[0, 3] = self._env_opts["task_track_omega_x_weight"]
+            self._task_pred_err_weights[0, 4] = self._env_opts["task_track_omega_y_weight"]
+            self._task_pred_err_weights[0, 5] = self._env_opts["task_track_omega_z_weight"]
+
+        self._power_penalty_weights = torch.full((1, self._n_jnts), dtype=self._dtype, device=device,
+                            fill_value=1.0)
+        self._power_penalty_weights_sum = torch.sum(self._power_penalty_weights).item()
+        subr_names=self._get_rewards_names() # initializes
+        
+        # reward clipping
+        self._reward_thresh_lb[:, :] = self._env_opts["reward_lb_default"]
+        self._reward_thresh_ub[:, :]= self._env_opts["reward_ub_default"]
+
+        for reward_name, env_opt_key in self._reward_lb_map.items():
+            if reward_name in self._reward_map:
+                self._reward_thresh_lb[:, self._reward_map[reward_name]] = self._env_opts[env_opt_key]
+
+        # obs bounds
+        self._obs_threshold_lb = -1e3 # used for clipping observations
+        self._obs_threshold_ub = 1e3
+
+        # actions
+        if not self._env_opts["use_prob_based_stepping"]:
+            self._is_continuous_actions[6:10]=False
+
+        v_cmd_max = self._env_opts["max_cmd_v"]
+        omega_cmd_max = self._env_opts["max_cmd_omega"]
+        self._actions_lb[:, 0:3] = -v_cmd_max 
+        self._actions_ub[:, 0:3] = v_cmd_max  
+        self._actions_lb[:, 3:6] = -omega_cmd_max # twist cmds
+        self._actions_ub[:, 3:6] = omega_cmd_max  
+        if "contact_flag_start" in self._actions_map:
+            idx=self._actions_map["contact_flag_start"]
+            if self._env_opts["use_prob_based_stepping"]:
+                self._actions_lb[:, idx:idx+self._n_contacts] = 0.0 # contact flags
+                self._actions_ub[:, idx:idx+self._n_contacts] = 1.0 
+            else:
+                self._actions_lb[:, idx:idx+self._n_contacts] = -1.0 
+                self._actions_ub[:, idx:idx+self._n_contacts] = 1.0 
+        
+        self.default_action[:, :] = (self._actions_ub+self._actions_lb)/2.0
+        # self.default_action[:, ~self._is_continuous_actions] = 1.0
+        self.safe_action[:, :] = self.default_action
+        if "contact_flag_start" in self._actions_map: # safe actions for contacts is 1 (keep contact)
+            idx=self._actions_map["contact_flag_start"]
+            self.safe_action[:, idx:idx+self._n_contacts] = 1.0
+
+        # assign obs bounds (useful if not using automatic obs normalization)
+        obs_names=self._get_obs_names()
+        obs_patterns=["gn",
+            "linvel",
+            "omega",
+            "q_jnt",
+            "v_jnt",
+            "fc",
+            "rhc_fail",
+            "rhc_cmd_q",
+            "rhc_cmd_v",
+            "rhc_cmd_eff",
+            "flight_pos"
+            ]
+        obs_ubs=[1.0,
+            5*v_cmd_max,
+            5*omega_cmd_max,
+            2*torch.pi,
+            30.0,
+            2.0,
+            1.0,
+            2*torch.pi,
+            30.0,
+            200.0,
+            self._n_nodes_rhc.mean().item()]
+        obs_lbs=[-1.0,
+            -5*v_cmd_max,
+            -5*omega_cmd_max,
+            -2*torch.pi,
+            -30.0,
+            -2.0,
+            0.0,
+            -2*torch.pi,
+            -30.0,
+            -200.0,
+            0.0]
+        obs_bounds = {name: (lb, ub) for name, lb, ub in zip(obs_patterns, obs_lbs, obs_ubs)}
+        
+        for i in range(len(obs_names)):
+            obs_name=obs_names[i]
+            for pattern in obs_patterns:
+                if pattern in obs_name:
+                    lb=obs_bounds[pattern][0]
+                    ub=obs_bounds[pattern][1]
+                    self._obs_lb[:, i]=lb
+                    self._obs_ub[:, i]=ub
+                    break
+        
+        # handle action memory buffer in obs
+        if self._env_opts["use_action_history"]: # just history stats
+            if self._env_opts["add_prev_actions_stats_to_obs"]:
+                i=0
+                prev_actions_idx = next((i for i, s in enumerate(obs_names) if "_prev_act" in s), None)
+                prev_actions_mean_idx=next((i for i, s in enumerate(obs_names) if "_avrg_act" in s), None)
+                prev_actions_std_idx=next((i for i, s in enumerate(obs_names) if "_std_act" in s), None)
+                
+                # assume actions are always normalized in [-1, 1] by agent
+                if prev_actions_idx is not None:
+                    self._obs_lb[:, prev_actions_idx:prev_actions_idx+self.actions_dim()]=-1.0
+                    self._obs_ub[:, prev_actions_idx:prev_actions_idx+self.actions_dim()]=1.0
+                if prev_actions_mean_idx is not None:
+                    self._obs_lb[:, prev_actions_mean_idx:prev_actions_mean_idx+self.actions_dim()]=-1.0
+                    self._obs_ub[:, prev_actions_mean_idx:prev_actions_mean_idx+self.actions_dim()]=1.0
+                if prev_actions_std_idx is not None:
+                    self._obs_lb[:, prev_actions_std_idx:prev_actions_std_idx+self.actions_dim()]=0
+                    self._obs_ub[:, prev_actions_std_idx:prev_actions_std_idx+self.actions_dim()]=1.0
+                
+            else: # full history
+                i=0
+                first_action_mem_buffer_idx = next((i for i, s in enumerate(obs_names) if "_m1_act" in s), None)
+                if first_action_mem_buffer_idx is not None:
+                    action_idx_start_idx_counter=first_action_mem_buffer_idx
+                    for j in range(self._env_opts["actions_history_size"]):
+                        self._obs_lb[:, action_idx_start_idx_counter:action_idx_start_idx_counter+self.actions_dim()]=-1.0
+                        self._obs_ub[:, action_idx_start_idx_counter:action_idx_start_idx_counter+self.actions_dim()]=1.0
+                        action_idx_start_idx_counter+=self.actions_dim()
+
+        # some aux data to avoid allocations at training runtime
+        self._rhc_twist_cmd_rhc_world=self._robot_state.root_state.get(data_type="twist",gpu=self._use_gpu).detach().clone()
+        self._rhc_twist_cmd_rhc_h=self._rhc_twist_cmd_rhc_world.detach().clone()
+        self._agent_twist_ref_current_w=self._rhc_twist_cmd_rhc_world.detach().clone()
+        self._agent_twist_ref_current_base_loc=self._rhc_twist_cmd_rhc_world.detach().clone()
+        self._substep_avrg_root_twist_base_loc=self._rhc_twist_cmd_rhc_world.detach().clone()
+        self._step_avrg_root_twist_base_loc=self._rhc_twist_cmd_rhc_world.detach().clone()
+        self._root_twist_avrg_rhc_base_loc=self._rhc_twist_cmd_rhc_world.detach().clone()
+        self._root_twist_avrg_rhc_base_loc_next=self._rhc_twist_cmd_rhc_world.detach().clone()
+        
+        self._random_thresh_contacts=torch.rand((self._n_envs,self._n_contacts), device=device)
+        # aux data
+        self._task_err_scaling = torch.zeros((self._n_envs, 1),dtype=self._dtype,device=device)
+
+        self._pof1_b_linvel= torch.full(size=(self._n_envs,1),dtype=self._dtype,device=device,fill_value=1-self._env_opts["pof0_linvel"])
+        self._pof1_b_omega = torch.full(size=(self._n_envs,1),dtype=self._dtype,device=device,fill_value=1-self._env_opts["pof0_omega"])
+        self._bernoulli_coeffs_linvel = self._pof1_b_linvel.clone()
+        self._bernoulli_coeffs_linvel[:, :] = 1.0
+        self._bernoulli_coeffs_omega = self._pof1_b_omega.clone()
+        self._bernoulli_coeffs_omega[:, :] = 1.0
+
+        # smoothing
+        self._track_rew_smoother=None
+        if self._env_opts["use_track_reward_smoother"]:
+            sub_reward_proxy=self._sub_rewards.get_torch_mirror(gpu=self._use_gpu)[:, 0:1]
+            smoothing_dt=self._substep_dt
+            if not self._is_substep_rew[self._reward_map["task_error"]]: # assuming first reward is tracking
+                smoothing_dt=self._substep_dt*self._action_repeat
+            self._track_rew_smoother=ExponentialSignalSmoother(
+                name=self.__class__.__name__+"VelErrorSmoother",
+                signal=sub_reward_proxy, # same dimension of vel error
+                update_dt=smoothing_dt,
+                smoothing_horizon=self._env_opts["smoothing_horizon_vel_err"],
+                target_smoothing=0.5,
+                debug=self._is_debug,
+                dtype=self._dtype,
+                use_gpu=self._use_gpu)
+
+        # if we need the action rate, we also need the action history
+        if self._env_opts["add_action_rate_reward"]:
+            if not self._env_opts["use_action_history"]:
+                Journal.log(self.__class__.__name__,
+                    "_custom_post_init",
+                    "add_action_rate_reward is True, but ",
+                    LogType.EXCEP,
+                    throw_when_excep=True)
+            
+            history_size=self._env_opts["actions_history_size"]
+            if history_size < 2:
+                Journal.log(self.__class__.__name__,
+                    "_custom_post_init",
+                    f"add_action_rate_reward  requires actions history ({history_size}) to be >=2!",
+                    LogType.EXCEP,
+                    throw_when_excep=True)
+        
+        # add periodic timer if required
+        self._periodic_clock=None
+        if self._env_opts["add_periodic_clock_to_obs"]:
+            self._add_env_opt(self._env_opts, "clock_period", 
+                default=int(1.5*self._action_repeat*self.task_rand_timeout_bounds()[1])) # correcting with n substeps
+            # (we are using the _substep_abs_counter counter)
+            self._periodic_clock=PeriodicTimer(counter=self._substep_abs_counter,
+                                    period=self._env_opts["clock_period"], 
+                                    dtype=self._dtype,
+                                    device=self._device)
+
+    def get_file_paths(self):
+        paths=AugMPCTrainingEnvBase.get_file_paths(self)
+        paths.append(self._this_child_path)        
+        return paths
+
+    def get_aux_dir(self):
+        aux_dirs = []
+        path_getter = PathsGetter()
+        aux_dirs.append(path_getter.RHCDIR)
+        return aux_dirs
+
+    def _get_reward_scaling(self):
+        if self._env_opts["single_task_ref_per_episode"]:
+            return self._env_opts["n_steps_task_rand_ub"]
+        else:
+            return self._env_opts["episode_timeout_ub"]
+    
+    def _max_ep_length(self):
+        if self._env_opts["single_task_ref_per_episode"]:
+            return self._env_opts["n_steps_task_rand_ub"]
+        else:
+            return self._env_opts["episode_timeout_ub"]
+    
+    def _check_sub_truncations(self):
+        # overrides parent
+        sub_truncations = self._sub_truncations.get_torch_mirror(gpu=self._use_gpu)
+        sub_truncations[:, 0:1] = self._ep_timeout_counter.time_limits_reached()
+        if self._env_opts["single_task_ref_per_episode"]:
+            sub_truncations[:, 1:2] = self._task_rand_counter.time_limits_reached()
+    
+    def _check_sub_terminations(self):
+        # default behaviour-> to be overriden by child
+        sub_terminations = self._sub_terminations.get_torch_mirror(gpu=self._use_gpu)
+        
+        # terminate if mpc just failed
+        sub_terminations[:, 0:1] = self._rhc_status.fails.get_torch_mirror(gpu=self._use_gpu)
+
+        # check if robot is capsizing
+        robot_q_meas = self._robot_state.root_state.get(data_type="q",gpu=self._use_gpu)
+        check_capsize(quat=robot_q_meas,max_angle=self._max_pitch_angle,
+            output_t=self._is_capsized)
+        sub_terminations[:, 1:2] = self._is_capsized
+        
+        if self._env_opts["add_term_mpc_capsize"]:
+            # check if robot is about to capsize accordin to MPC
+            robot_q_pred = self._rhc_cmds.root_state.get(data_type="q",gpu=self._use_gpu)
+            check_capsize(quat=robot_q_pred,max_angle=self._max_pitch_angle,
+                output_t=self._is_rhc_capsized)
+            sub_terminations[:, 2:3] = self._is_rhc_capsized
+
+    def _custom_reset(self):
+        return None
+    
+    def reset(self):
+        AugMPCTrainingEnvBase.reset(self)
+
+    def _pre_substep(self): 
+        pass
+
+    def _custom_post_step(self,episode_finished):
+        # executed after checking truncations and terminations and remote env reset
+        if self._use_gpu:
+            time_to_rand_or_ep_finished = torch.logical_or(self._task_rand_counter.time_limits_reached().cuda(),episode_finished)
+            self.randomize_task_refs(env_indxs=time_to_rand_or_ep_finished.flatten())
+        else:
+            time_to_rand_or_ep_finished = torch.logical_or(self._task_rand_counter.time_limits_reached(),episode_finished)
+            self.randomize_task_refs(env_indxs=time_to_rand_or_ep_finished.flatten())
+        # task refs are randomized in world frame -> we rotate them in base local
+        # (not super efficient, we should do it just for the finished envs)
+        self._update_loc_twist_refs()
+
+        if self._track_rew_smoother is not None: # reset smoother
+            self._track_rew_smoother.reset_all(to_be_reset=episode_finished.flatten(), 
+                    value=0.0)
+
+    def _custom_post_substp_pre_rew(self):
+        self._update_loc_twist_refs()
+        
+    def _custom_post_substp_post_rew(self):
+        pass
+    
+    def _update_loc_twist_refs(self):
+        # get fresh robot orientation
+        if not self._override_agent_refs:
+            robot_q = self._robot_state.root_state.get(data_type="q",gpu=self._use_gpu)
+            # rotate agent ref from world to robot base
+            world2base_frame(t_w=self._agent_twist_ref_current_w, q_b=robot_q, 
+                t_out=self._agent_twist_ref_current_base_loc)
+            # write it to agent refs tensors
+            self._agent_refs.rob_refs.root_state.set(data_type="twist", data=self._agent_twist_ref_current_base_loc,
+                                                gpu=self._use_gpu)
+        
+    def _apply_actions_to_rhc(self):
+        
+        self._set_rhc_refs()
+
+        self._write_rhc_refs()
+
+    def _set_rhc_refs(self):
+
+        action_to_be_applied = self.get_actual_actions() # see _get_action_names() to get 
+        # the meaning of each component of this tensor
+
+        rhc_latest_twist_cmd = self._rhc_refs.rob_refs.root_state.get(data_type="twist", gpu=self._use_gpu)
+        rhc_latest_contact_ref = self._rhc_refs.contact_flags.get_torch_mirror(gpu=self._use_gpu)
+        rhc_latest_pos_ref = self._rhc_refs.rob_refs.contact_pos.get(data_type="p_z", gpu=self._use_gpu)
+        rhc_q=self._rhc_cmds.root_state.get(data_type="q",gpu=self._use_gpu) # this is always 
+        # avaialble
+
+        # reference twist for MPC is assumed to always be specified in MPC's 
+        # horizontal frame, while agent actions are interpreted as in MPC's
+        # base frame -> we need to rotate the actions into the horizontal frame
+        base2world_frame(t_b=action_to_be_applied[:, 0:6],q_b=rhc_q,t_out=self._rhc_twist_cmd_rhc_world)
+        w2hor_frame(t_w=self._rhc_twist_cmd_rhc_world,q_b=rhc_q,t_out=self._rhc_twist_cmd_rhc_h)
+
+        rhc_latest_twist_cmd[:, 0:6] = self._rhc_twist_cmd_rhc_h
+        
+        # self._rhc_refs.rob_refs.root_state.set(data_type="p", data=rhc_latest_p_ref,
+        #                                     gpu=self._use_gpu)
+        self._rhc_refs.rob_refs.root_state.set(data_type="twist", data=rhc_latest_twist_cmd,
+            gpu=self._use_gpu) 
+        
+        # contact flags
+        idx=self._actions_map["contact_flag_start"]
+        if self._env_opts["use_prob_based_stepping"]:
+            # encode actions as probs
+            self._random_thresh_contacts.uniform_() # random values in-place between 0 and 1
+            rhc_latest_contact_ref[:, :] = action_to_be_applied[:, idx:idx+self._n_contacts] >= self._random_thresh_contacts  # keep contact with 
+            # probability action_to_be_applied[:, 6:10]
+        else: # just use a threshold
+            rhc_latest_contact_ref[:, :] = action_to_be_applied[:, idx:idx+self._n_contacts] > self._env_opts["step_thresh"]
+        # actually apply actions to controller
+        
+    def _write_rhc_refs(self):
+
+        if self._use_gpu:
+            # GPU->CPU --> we cannot use asynchronous data transfer since it's unsafe
+            self._rhc_refs.rob_refs.root_state.synch_mirror(from_gpu=True,non_blocking=False) # write from gpu to cpu mirror
+            self._rhc_refs.contact_flags.synch_mirror(from_gpu=True,non_blocking=False)
+            self._rhc_refs.rob_refs.contact_pos.synch_mirror(from_gpu=True,non_blocking=False)
+
+        self._rhc_refs.rob_refs.root_state.synch_all(read=False, retry=True) # write mirror to shared mem
+        self._rhc_refs.contact_flags.synch_all(read=False, retry=True)
+        self._rhc_refs.rob_refs.contact_pos.synch_all(read=False, retry=True)
+    
+    def _override_refs(self,
+            env_indxs: torch.Tensor = None):
+        
+        # runs at every post_step
+        self._agent_refs.rob_refs.root_state.synch_all(read=True,retry=True) # first read from mem
+        if self._use_gpu:
+            # copies latest refs to GPU 
+            self._agent_refs.rob_refs.root_state.synch_mirror(from_gpu=False,non_blocking=False) 
+
+        agent_linvel_ref_current=self._agent_refs.rob_refs.root_state.get(data_type="v",
+                gpu=self._use_gpu)
+        
+        agent_yaw_omega_ref_current=self._agent_refs.rob_refs.root_state.get(data_type="omega",
+                gpu=self._use_gpu)
+        
+        # self._p_trgt_w[:, :]=self._robot_state.root_state.get(data_type="p",gpu=self._use_gpu)[:, 0:2] + \
+        #     agent_p_ref_current[:, 0:2]
+        self._agent_twist_ref_current_w[:, 0:3]=agent_linvel_ref_current # set linvel target
+        
+        self._agent_twist_ref_current_w[:, 5:6]=agent_yaw_omega_ref_current[:, 2:3] # set yaw ang. vel target from shared mem
+
+    def _fill_substep_obs(self,
+            obs: torch.Tensor):
+
+        # measured stuff
+        robot_twist_meas_base_loc = self._robot_state.root_state.get(data_type="twist",gpu=self._use_gpu)
+        robot_jnt_v_meas = self._robot_state.jnts_state.get(data_type="v",gpu=self._use_gpu)
+        
+        if self._env_opts["use_linvel_from_rhc"]:
+            # twist estimate from mpc
+            robot_twist_rhc_base_loc_next = self._rhc_cmds.root_state.get(data_type="twist",gpu=self._use_gpu)
+            obs[:, self._obs_map["linvel_meas"]:(self._obs_map["linvel_meas"]+3)] = robot_twist_rhc_base_loc_next[:, 0:3]
+        else:
+            obs[:, self._obs_map["linvel_meas"]:(self._obs_map["linvel_meas"]+3)] = robot_twist_meas_base_loc[:, 0:3]
+        obs[:, self._obs_map["omega_meas"]:(self._obs_map["omega_meas"]+3)] = robot_twist_meas_base_loc[:, 3:6]
+
+        obs[:, self._obs_map["v_jnt"]:(self._obs_map["v_jnt"]+self._n_jnts)] = robot_jnt_v_meas 
+
+    def _fill_step_obs(self,
+            obs: torch.Tensor):
+
+        # measured stuff
+        robot_gravity_norm_base_loc = self._robot_state.root_state.get(data_type="gn",gpu=self._use_gpu)
+        robot_twist_meas_base_loc = self._robot_state.root_state.get(data_type="twist",gpu=self._use_gpu)
+        robot_jnt_q_meas = self._robot_state.jnts_state.get(data_type="q",gpu=self._use_gpu)
+        if self._jnt_q_blacklist_idxs is not None: # we don't want to read joint pos from blacklist
+            robot_jnt_q_meas[:, self._jnt_q_blacklist_idxs]=0.0
+        robot_jnt_v_meas = self._robot_state.jnts_state.get(data_type="v",gpu=self._use_gpu)
+        
+        # twist estimate from mpc
+        robot_twist_rhc_base_loc_next = self._rhc_cmds.root_state.get(data_type="twist",gpu=self._use_gpu)
+        # cmds for jnt imp to be applied next
+        robot_jnt_q_rhc_applied_next=self._rhc_cmds.jnts_state.get(data_type="q",gpu=self._use_gpu)
+        robot_jnt_v_rhc_applied_next=self._rhc_cmds.jnts_state.get(data_type="v",gpu=self._use_gpu)
+        robot_jnt_eff_rhc_applied_next=self._rhc_cmds.jnts_state.get(data_type="eff",gpu=self._use_gpu)
+
+        flight_info_now = self._rhc_refs.flight_info.get(data_type="all",gpu=self._use_gpu)
+        flight_settings_now = self._rhc_refs.flight_settings_req.get(data_type="all",gpu=self._use_gpu)
+        
+        # refs
+        agent_twist_ref = self._agent_refs.rob_refs.root_state.get(data_type="twist",gpu=self._use_gpu)
+
+        obs[:, self._obs_map["gn_base"]:(self._obs_map["gn_base"]+3)] = robot_gravity_norm_base_loc # norm. gravity vector in base frame
+        
+        obs[:, self._obs_map["q_jnt"]:(self._obs_map["q_jnt"]+self._n_jnts)] = robot_jnt_q_meas # meas jnt pos
+        obs[:, self._obs_map["twist_ref"]:(self._obs_map["twist_ref"]+6)] = agent_twist_ref # high lev agent refs to be tracked
+
+        if self._env_opts["add_mpc_contact_f_to_obs"]:
+            n_forces=3*len(self._contact_names)
+            obs[:, self._obs_map["contact_f_mpc"]:(self._obs_map["contact_f_mpc"]+n_forces)] = self._rhc_cmds.contact_wrenches.get(data_type="f",gpu=self._use_gpu)
+        if self._env_opts["add_fail_idx_to_obs"]:
+            obs[:, self._obs_map["rhc_fail_idx"]:(self._obs_map["rhc_fail_idx"]+1)] = self._rhc_fail_idx(gpu=self._use_gpu)
+        if self._env_opts["add_term_mpc_capsize"]:
+            obs[:, self._obs_map["gn_base_mpc"]:(self._obs_map["gn_base_mpc"]+3)] = self._rhc_cmds.root_state.get(data_type="gn",gpu=self._use_gpu)
+        if self._env_opts["use_rhc_avrg_vel_tracking"]:
+            self._get_avrg_rhc_root_twist(out=self._root_twist_avrg_rhc_base_loc, base_loc=True)
+            obs[:, self._obs_map["avrg_twist_mpc"]:(self._obs_map["avrg_twist_mpc"]+6)] = self._root_twist_avrg_rhc_base_loc
+        if self._env_opts["add_flight_info"]:
+            obs[:, self._obs_map["flight_info"]:(self._obs_map["flight_info"]+self._flight_info_size)] = flight_info_now
+        if self._env_opts["add_flight_settings"]:
+            obs[:, self._obs_map["flight_settings_req"]:(self._obs_map["flight_settings_req"]+self._flight_setting_size)] = \
+                flight_settings_now
+
+        if self._env_opts["add_rhc_cmds_to_obs"]:
+            obs[:, self._obs_map["rhc_cmds_q"]:(self._obs_map["rhc_cmds_q"]+self._n_jnts)] = robot_jnt_q_rhc_applied_next
+            obs[:, self._obs_map["rhc_cmds_v"]:(self._obs_map["rhc_cmds_v"]+self._n_jnts)] = robot_jnt_v_rhc_applied_next
+            obs[:, self._obs_map["rhc_cmds_eff"]:(self._obs_map["rhc_cmds_eff"]+self._n_jnts)] = robot_jnt_eff_rhc_applied_next
+        if self._env_opts["use_action_history"]:
+            if self._env_opts["add_prev_actions_stats_to_obs"]: # just add last, std and mean to obs
+                obs[:, self._obs_map["action_history_prev"]:(self._obs_map["action_history_prev"]+self.actions_dim())]=self._act_mem_buffer.get(idx=0)
+                obs[:, self._obs_map["action_history_avrg"]:(self._obs_map["action_history_avrg"]+self.actions_dim())]=self._act_mem_buffer.mean(clone=False)
+                obs[:, self._obs_map["action_history_std"]:(self._obs_map["action_history_std"]+self.actions_dim())]=self._act_mem_buffer.std(clone=False)
+            else: # add whole memory buffer to obs
+                next_idx=self._obs_map["action_history"]
+                for i in range(self._env_opts["actions_history_size"]):
+                    obs[:, next_idx:(next_idx+self.actions_dim())]=self._act_mem_buffer.get(idx=i) # get all (n_envs x (obs_dim x horizon))
+                    next_idx+=self.actions_dim()
+
+        if self._env_opts["use_action_smoothing"]: # adding smoothed actions
+            obs[:, self._obs_map["action_smoothing"]:(self._obs_map["action_smoothing"]+self.actions_dim())]=self.get_actual_actions(normalized=True)
+            next_idx+=self.actions_dim()
+        
+        if self._env_opts["add_periodic_clock_to_obs"]:
+            obs[:, next_idx:(next_idx+2)]=self._periodic_clock.get()
+            next_idx+=2
+        if self._env_opts["add_heightmap_obs"]:
+            hm = self._robot_state.height_sensor.get(gpu=self._use_gpu)
+            obs[:, self._obs_map["heightmap"]:(self._obs_map["heightmap"]+self._height_flat_dim)] = hm
+
+    def _get_custom_db_data(self, 
+            episode_finished,
+            ignore_ep_end):
+        episode_finished = episode_finished.cpu()
+        self.custom_db_data["AgentTwistRefs"].update(
+                new_data=self._agent_refs.rob_refs.root_state.get(data_type="twist", gpu=False), 
+                ep_finished=episode_finished,
+                ignore_ep_end=ignore_ep_end)
+        self.custom_db_data["RhcFailIdx"].update(new_data=self._rhc_fail_idx(gpu=False), 
+                ep_finished=episode_finished,
+                ignore_ep_end=ignore_ep_end)
+        self.custom_db_data["RhcContactForces"].update(
+                new_data=self._rhc_cmds.contact_wrenches.get(data_type="f",gpu=False), 
+                ep_finished=episode_finished,
+                ignore_ep_end=ignore_ep_end)
+        self.custom_db_data["Power"].update(
+                new_data=self._pow_db_data, 
+                ep_finished=episode_finished,
+                ignore_ep_end=ignore_ep_end)
+        self.custom_db_data["TrackingError"].update(
+                new_data=self._track_error_db, 
+                ep_finished=episode_finished,
+                ignore_ep_end=ignore_ep_end)
+
+    # reward functions
+    def _action_rate(self):
+        continuous_actions=self._is_continuous_actions
+        discrete_actions=~self._is_continuous_actions
+        n_c_actions=continuous_actions.sum().item()
+        n_d_actions=discrete_actions.sum().item()
+        actions_prev=self._act_mem_buffer.get(idx=1) 
+        actions_now=self._act_mem_buffer.get(idx=0)
+        actions_rate=(actions_now-actions_prev) # actions already normalized
+        actions_rate_c=actions_rate[:, continuous_actions]
+        actions_rate_d=actions_rate[:, discrete_actions]
+
+        actions_rate_sqrd=None # assuming n_c_actions > 0 always
+        actions_rate_sqrd=self._env_opts["action_rate_rew_c_weight"]*torch.sum(actions_rate_c*actions_rate_c, dim=1, keepdim=True)/n_c_actions
+        if discrete_actions.any():
+            actions_rate_sqrd+=self._env_opts["action_rate_rew_d_weight"]*torch.sum(actions_rate_d*actions_rate_d, dim=1, keepdim=True)/n_d_actions
+        return actions_rate_sqrd
+
+    def _mech_pow(self, jnts_vel, jnts_effort, autoscaled: bool = False, drained: bool = True):
+        mech_pow_jnts=(jnts_effort*jnts_vel)*self._power_penalty_weights
+        if drained:
+            mech_pow_jnts.clamp_(0.0,torch.inf) # do not account for regenerative power
+        mech_pow_tot = torch.sum(mech_pow_jnts, dim=1, keepdim=True)
+        self._pow_db_data[:, 1:2]=mech_pow_tot.cpu()
+        if autoscaled:
+            mech_pow_tot=mech_pow_tot/self._power_penalty_weights_sum
+        return mech_pow_tot
+
+    def _cost_of_transport(self, jnts_vel, jnts_effort, v_norm, mass_weight: bool = False):
+        drained_mech_pow=self._mech_pow(jnts_vel=jnts_vel,
+            jnts_effort=jnts_effort, 
+            drained=True)
+        CoT=drained_mech_pow/(v_norm+1e-2)
+        if mass_weight:
+            robot_weight=self._rhc_robot_weight
+            CoT=CoT/robot_weight
+        
+        # add to db metrics
+        self._pow_db_data[:, 0:1]=CoT.cpu()
+        self._pow_db_data[:, 1:2]=drained_mech_pow.cpu()
+
+        return CoT
+
+    def _jnt_vel_penalty(self, jnts_vel):
+        weighted_jnt_vel = torch.sum(jnts_vel*jnts_vel, dim=1, keepdim=True)/self._n_jnts
+        return weighted_jnt_vel
+    
+    def _rhc_fail_idx(self, gpu: bool):
+        rhc_fail_idx = self._rhc_status.rhc_fail_idx.get_torch_mirror(gpu=gpu)
+        return self._env_opts["rhc_fail_idx_scale"]*rhc_fail_idx
+    
+    # basic L1 and L2 error functions
+    def _track_err_wmse(self, task_ref, task_meas, scaling, weights):
+        # weighted mean-squared error computation 
+        task_error = (task_meas-task_ref)
+        # add to db metrics
+        self._track_error_db[:, :]=torch.abs(task_error)
+        scaled_error=task_error/scaling
+        
+        task_wmse = torch.sum(scaled_error*scaled_error*weights, dim=1, keepdim=True)/torch.sum(weights).item()
+        return task_wmse # weighted mean square error (along task dimension)
+    
+    def _track_err_dir_wmse(self, task_ref, task_meas, scaling, weights):
+        # weighted DIRECTIONAL mean-squared error computation 
+        task_error = (task_meas-task_ref)
+        # add to db metrics
+        self._track_error_db[:, :]=torch.abs(task_error)
+        task_error=task_error/scaling
+
+        # projection along commanded direction and gravity, matching paper formulation
+        v_ref=task_ref[:, 0:3]
+        delta_v=task_error[:, 0:3]
+
+        v_ref_norm=torch.norm(v_ref, dim=1, keepdim=True)
+        cmd_dir=v_ref/(v_ref_norm+1e-8)
+        # fallback to measured direction if command is (near) zero to avoid degenerate projection
+        meas_dir=task_meas[:, 0:3]
+        meas_dir=meas_dir/(torch.norm(meas_dir, dim=1, keepdim=True)+1e-8)
+        cmd_dir=torch.where((v_ref_norm>1e-6), cmd_dir, meas_dir)
+
+        gravity_dir = self._robot_state.root_state.get(data_type="gn",gpu=self._use_gpu) # normalized gravity in base frame
+        gravity_dir = gravity_dir/(torch.norm(gravity_dir, dim=1, keepdim=True)+1e-8)
+
+        forward_error=torch.sum(delta_v*cmd_dir, dim=1, keepdim=True)
+        vertical_error=torch.sum(delta_v*gravity_dir, dim=1, keepdim=True)
+        lateral_vec=delta_v - vertical_error*gravity_dir - forward_error*cmd_dir
+        lateral_error=torch.norm(lateral_vec, dim=1, keepdim=True)
+
+        # angular directional components: use gravity as vertical, project base x onto the world xy plane for roll, and close the triad with pitch
+        base_x = self._base_x_dir
+        base_y = self._base_y_dir
+
+        roll_dir = base_x - torch.sum(base_x*gravity_dir, dim=1, keepdim=True)*gravity_dir
+        roll_norm = torch.norm(roll_dir, dim=1, keepdim=True)
+        roll_dir_alt = base_y - torch.sum(base_y*gravity_dir, dim=1, keepdim=True)*gravity_dir # fallback if base x is almost aligned with gravity
+        roll_norm_alt = torch.norm(roll_dir_alt, dim=1, keepdim=True)
+        use_alt_roll = roll_norm < 1e-6
+        roll_dir = torch.where(use_alt_roll, roll_dir_alt, roll_dir)
+        roll_norm = torch.where(use_alt_roll, roll_norm_alt, roll_norm)
+        roll_dir = roll_dir/(roll_norm+1e-8)
+
+        pitch_dir = torch.cross(gravity_dir, roll_dir, dim=1)
+        pitch_dir = pitch_dir/(torch.norm(pitch_dir, dim=1, keepdim=True)+1e-8)
+
+        delta_omega = task_error[:, 3:6]
+        omega_roll_error = torch.sum(delta_omega*roll_dir, dim=1, keepdim=True)
+        omega_pitch_error = torch.sum(delta_omega*pitch_dir, dim=1, keepdim=True)
+        omega_vertical_error = torch.sum(delta_omega*gravity_dir, dim=1, keepdim=True)
+
+        full_error=torch.cat((forward_error, lateral_error, vertical_error, omega_roll_error, omega_pitch_error, omega_vertical_error), dim=1)
+        task_wmse_dir = torch.sum(full_error*full_error*weights, dim=1, keepdim=True)/torch.sum(weights).item()
+        return task_wmse_dir # weighted mean square error (along task dimension)
+    
+    # L2 errors
+    def _tracking_err_rel_wmse(self, task_ref, task_meas, weights, directional: bool = False):
+        ref_norm = task_ref.norm(dim=1,keepdim=True) # norm of the full twist reference
+        self._task_err_scaling[:, :] = ref_norm+1e-2
+        if directional:
+            task_rel_err_wmse=self._track_err_dir_wmse(task_ref=task_ref, task_meas=task_meas, 
+                scaling=self._task_err_scaling, weights=weights)
+        else:
+            task_rel_err_wmse=self._track_err_wmse(task_ref=task_ref, task_meas=task_meas, 
+                scaling=self._task_err_scaling, weights=weights)
+        return task_rel_err_wmse
+    
+    def _tracking_err_wmse(self, task_ref, task_meas, weights, directional: bool = False):
+        self._task_err_scaling[:, :] = 1
+        if directional:
+            task_err_wmse = self._track_err_dir_wmse(task_ref=task_ref, 
+                task_meas=task_meas, scaling=self._task_err_scaling, weights=weights)
+        else:
+            task_err_wmse = self._track_err_wmse(task_ref=task_ref, 
+                task_meas=task_meas, scaling=self._task_err_scaling, weights=weights)
+        return task_err_wmse
+    
+    # L1 errors
+    def _tracking_err_rel_lin(self, task_ref, task_meas, weights, directional):
+        task_rel_err_wmse = self._tracking_err_rel_wmse(task_ref=task_ref, 
+            task_meas=task_meas, weights=weights, directional=directional)
+        return task_rel_err_wmse.sqrt()
+    
+    def _tracking_err_lin(self, task_ref, task_meas, weights, directional: bool = False):
+        self._task_err_scaling[:, :] = 1
+        task_err_wmse=self._tracking_err_wmse(task_ref=task_ref,
+            task_meas=task_meas, weights=weights, directional=directional)
+        return task_err_wmse.sqrt()
+    
+    # reward computation over steps/substeps
+    def _compute_step_rewards(self):
+        
+        sub_rewards = self._sub_rewards.get_torch_mirror(gpu=self._use_gpu)
+
+        # tracking reward
+        if self._env_opts["use_L1_norm"]: # linear errors
+            task_error_fun = self._tracking_err_lin
+            if self._env_opts["use_relative_error"]:
+                task_error_fun = self._tracking_err_rel_lin
+        else: # quadratic error
+            task_error_fun = self._tracking_err_wmse
+            if self._env_opts["use_relative_error"]:
+                task_error_fun = self._tracking_err_rel_wmse
+                
+        agent_task_ref_base_loc = self._agent_refs.rob_refs.root_state.get(data_type="twist",gpu=self._use_gpu) # high level agent refs (hybrid twist)
+        self._get_avrg_step_root_twist(out=self._step_avrg_root_twist_base_loc, base_loc=True)
+        task_error = task_error_fun(task_meas=self._step_avrg_root_twist_base_loc, 
+            task_ref=agent_task_ref_base_loc,
+            weights=self._task_err_weights,
+            directional=self._env_opts["directional_tracking"])
+
+        idx=self._reward_map["task_error"]
+        if self._env_opts["use_exp_track_rew"]:
+            sub_rewards[:, idx:(idx+1)] =  \
+                self._env_opts["task_track_offset_exp"]*torch.exp(-self._env_opts["task_track_scale_exp"]*task_error)
+        else: # simple linear reward
+            sub_rewards[:, idx:(idx+1)] = \
+                self._env_opts["task_track_offset"]*(1.0-self._env_opts["task_track_scale"]*task_error)
+
+        if self._env_opts["use_fail_idx_weight"]: # add weight based on fail idx
+            fail_idx=self._rhc_fail_idx(gpu=self._use_gpu)
+            sub_rewards[:, idx:(idx+1)]=(1-fail_idx)*sub_rewards[:, idx:(idx+1)]
+        if self._track_rew_smoother is not None: # smooth reward if required
+            self._track_rew_smoother.update(new_signal=sub_rewards[:, 0:1])
+            sub_rewards[:, idx:(idx+1)]=self._track_rew_smoother.get()
+
+        # action rate
+        if self._env_opts["add_action_rate_reward"]:
+            action_rate=self._action_rate()
+            idx=self._reward_map["action_rate"]
+            sub_rewards[:, idx:(idx+1)] = self._env_opts["action_rate_offset"]*(1.0-self._env_opts["action_rate_scale"]*action_rate)
+
+        # mpc vel tracking
+        if self._env_opts["use_rhc_avrg_vel_tracking"]:
+            self._get_avrg_rhc_root_twist(out=self._root_twist_avrg_rhc_base_loc_next,base_loc=True) # get estimated avrg vel 
+            # from MPC after stepping
+            task_pred_error=task_error_fun(task_meas=self._root_twist_avrg_rhc_base_loc_next, 
+                task_ref=agent_task_ref_base_loc,
+                weights=self._task_pred_err_weights,
+                directional=self._env_opts["directional_tracking"])
+            idx=self._reward_map["rhc_avrg_vel_error"]
+            sub_rewards[:, idx:(idx+1)] = self._env_opts["task_pred_track_offset"]*torch.exp(-self._env_opts["task_pred_track_scale"]*task_pred_error)
+        
+    def _compute_substep_rewards(self):
+        
+        sub_rewards = self._sub_rewards.get_torch_mirror(gpu=self._use_gpu)
+
+        if self._env_opts["add_CoT_reward"] or self._env_opts["add_power_reward"]:
+            jnts_vel = self._robot_state.jnts_state.get(data_type="v",gpu=self._use_gpu)
+            jnts_effort = self._robot_state.jnts_state.get(data_type="eff",gpu=self._use_gpu)
+
+            if self._env_opts["add_CoT_reward"]:
+                if self._env_opts["use_CoT_wrt_ref"]: # uses v ref norm for computing cot
+                    agent_task_ref_base_loc = self._agent_refs.rob_refs.root_state.get(data_type="twist",gpu=self._use_gpu)
+                    v_norm=torch.norm(agent_task_ref_base_loc, dim=1, keepdim=True)
+                else: # uses measured velocity
+                    robot_twist_meas_base_loc = self._robot_state.root_state.get(data_type="twist",gpu=self._use_gpu)
+                    v_norm=torch.norm(robot_twist_meas_base_loc[:,0:3], dim=1, keepdim=True)
+                CoT=self._cost_of_transport(jnts_vel=jnts_vel,jnts_effort=jnts_effort,v_norm=v_norm, 
+                    mass_weight=True
+                    )
+                idx=self._reward_map["CoT"]
+                sub_rewards[:, idx:(idx+1)] = self._env_opts["CoT_offset"]*(1-self._env_opts["CoT_scale"]*CoT)
+            if self._env_opts["add_power_reward"]:
+                weighted_mech_power=self._mech_pow(jnts_vel=jnts_vel,jnts_effort=jnts_effort, drained=True)
+                idx=self._reward_map["mech_pow"]
+                sub_rewards[:, idx:(idx+1)] = self._env_opts["power_offset"]*(1-self._env_opts["power_scale"]*weighted_mech_power)
+        
+        if self._env_opts["add_jnt_v_reward"]:
+            jnts_vel = self._robot_state.jnts_state.get(data_type="v",gpu=self._use_gpu)
+            jnt_v=self._jnt_vel_penalty(jnts_vel=jnts_vel)
+            idx=self._reward_map["jnt_v"]
+            sub_rewards[:, idx:(idx+1)] = self._env_opts["jnt_vel_offset"]*(1-self._env_opts["jnt_vel_scale"]*jnt_v)
+
+    def _randomize_task_refs(self,
+        env_indxs: torch.Tensor = None):
+
+        # we randomize the reference in world frame, since it's much more intuitive 
+        # (it will be rotated in base frame when provided to the agent and used for rew 
+        # computation)
+        
+        if self._env_opts["use_pof0"]: # sample from bernoulli distribution
+            torch.bernoulli(input=self._pof1_b_linvel,out=self._bernoulli_coeffs_linvel) # by default bernoulli_coeffs are 1 if not self._env_opts["use_pof0"]
+            torch.bernoulli(input=self._pof1_b_omega,out=self._bernoulli_coeffs_omega)
+        if env_indxs is None:
+            random_uniform=torch.full_like(self._agent_twist_ref_current_w, fill_value=0.0)
+            torch.nn.init.uniform_(random_uniform, a=-1, b=1)
+            self._agent_twist_ref_current_w[:, :] = random_uniform*self._twist_ref_scale + self._twist_ref_offset
+            self._agent_twist_ref_current_w[:, 0:3] = self._agent_twist_ref_current_w[:, 0:3]*self._bernoulli_coeffs_linvel # linvel
+            self._agent_twist_ref_current_w[:, 3:6] = self._agent_twist_ref_current_w[:, 3:6]*self._bernoulli_coeffs_omega # omega
+        else:
+            random_uniform=torch.full_like(self._agent_twist_ref_current_w[env_indxs, :], fill_value=0.0)
+            torch.nn.init.uniform_(random_uniform, a=-1, b=1)
+            self._agent_twist_ref_current_w[env_indxs, :] = random_uniform * self._twist_ref_scale + self._twist_ref_offset
+            self._agent_twist_ref_current_w[env_indxs, 0:3] = self._agent_twist_ref_current_w[env_indxs, 0:3]*self._bernoulli_coeffs_linvel[env_indxs, :]
+            self._agent_twist_ref_current_w[env_indxs, 3:6] = self._agent_twist_ref_current_w[env_indxs, 3:6]*self._bernoulli_coeffs_omega[env_indxs, :] # omega
+    
+    def _get_obs_names(self):
+
+        obs_names = [""] * self.obs_dim()
+
+        # proprioceptive stream of obs
+        next_idx=0
+
+        self._obs_map["gn_base"]=next_idx
+        obs_names[0] = "gn_x_base_loc"
+        obs_names[1] = "gn_y_base_loc"
+        obs_names[2] = "gn_z_base_loc"
+        next_idx+=3
+
+        self._obs_map["linvel_meas"]=next_idx
+        obs_names[next_idx] = "linvel_x_base_loc"
+        obs_names[next_idx+1] = "linvel_y_base_loc"
+        obs_names[next_idx+2] = "linvel_z_base_loc"
+        next_idx+=3
+
+        self._obs_map["omega_meas"]=next_idx
+        obs_names[next_idx] = "omega_x_base_loc"
+        obs_names[next_idx+1] = "omega_y_base_loc"
+        obs_names[next_idx+2] = "omega_z_base_loc"
+        next_idx+=3
+
+        jnt_names=self.get_observed_joints()
+        self._obs_map["q_jnt"]=next_idx
+        for i in range(self._n_jnts): # jnt obs (pos):
+            obs_names[next_idx+i] = f"q_jnt_{jnt_names[i]}"
+        next_idx+=self._n_jnts
+        
+        self._obs_map["v_jnt"]=next_idx
+        for i in range(self._n_jnts): # jnt obs (v):
+            obs_names[next_idx+i] = f"v_jnt_{jnt_names[i]}"
+        next_idx+=self._n_jnts
+        
+        # references
+        self._obs_map["twist_ref"]=next_idx
+        obs_names[next_idx] = "linvel_x_ref_base_loc"
+        obs_names[next_idx+1] = "linvel_y_ref_base_loc"
+        obs_names[next_idx+2] = "linvel_z_ref_base_loc"
+        obs_names[next_idx+3] = "omega_x_ref_base_loc"
+        obs_names[next_idx+4] = "omega_y_ref_base_loc"
+        obs_names[next_idx+5] = "omega_z_ref_base_loc"
+        next_idx+=6
+        
+        # contact forces
+        if self._env_opts["add_mpc_contact_f_to_obs"]:
+            i = 0
+            self._obs_map["contact_f_mpc"]=next_idx
+            for contact in self._contact_names:
+                obs_names[next_idx+i] = f"fc_{contact}_x_base_loc"
+                obs_names[next_idx+i+1] = f"fc_{contact}_y_base_loc"
+                obs_names[next_idx+i+2] = f"fc_{contact}_z_base_loc"
+                i+=3        
+            next_idx+=3*len(self._contact_names)
+            
+        # data directly from MPC
+        if self._env_opts["add_fail_idx_to_obs"]:
+            self._obs_map["rhc_fail_idx"]=next_idx
+            obs_names[next_idx] = "rhc_fail_idx"
+            next_idx+=1
+        if self._env_opts["add_term_mpc_capsize"]:
+            self._obs_map["gn_base_mpc"]=next_idx
+            obs_names[next_idx] = "gn_x_rhc_base_loc"
+            obs_names[next_idx+1] = "gn_y_rhc_base_loc"
+            obs_names[next_idx+2] = "gn_z_rhc_base_loc"
+            next_idx+=3
+        if self._env_opts["use_rhc_avrg_vel_tracking"]:
+            self._obs_map["avrg_twist_mpc"]=next_idx
+            obs_names[next_idx] = "linvel_x_avrg_rhc"
+            obs_names[next_idx+1] = "linvel_y_avrg_rhc"
+            obs_names[next_idx+2] = "linvel_z_avrg_rhc"
+            obs_names[next_idx+3] = "omega_x_avrg_rhc"
+            obs_names[next_idx+4] = "omega_y_avrg_rhc"
+            obs_names[next_idx+5] = "omega_z_avrg_rhc"
+            next_idx+=6
+        
+        if self._env_opts["add_flight_info"]:
+            self._obs_map["flight_info"]=next_idx
+            for i in range(len(self._contact_names)):
+                obs_names[next_idx+i] = "flight_pos_"+ self._contact_names[i]
+            next_idx+=len(self._contact_names)
+            for i in range(len(self._contact_names)):
+                obs_names[next_idx+i] = "flight_len_remaining_"+ self._contact_names[i]
+            next_idx+=len(self._contact_names)
+            for i in range(len(self._contact_names)):
+                obs_names[next_idx+i] = "flight_len_nominal_"+ self._contact_names[i]
+            next_idx+=len(self._contact_names)
+            for i in range(len(self._contact_names)):
+                obs_names[next_idx+i] = "flight_apex_nominal_"+ self._contact_names[i]
+            next_idx+=len(self._contact_names)
+            for i in range(len(self._contact_names)):
+                obs_names[next_idx+i] = "flight_end_nominal_"+ self._contact_names[i]
+            next_idx+=len(self._contact_names)
+        
+        if self._env_opts["add_flight_settings"]:
+            self._obs_map["flight_settings_req"]=next_idx
+            for i in range(len(self._contact_names)):
+                obs_names[next_idx+i] = "flight_len_req_"+ self._contact_names[i]
+            next_idx+=len(self._contact_names)
+            for i in range(len(self._contact_names)):
+                obs_names[next_idx+i] = "flight_apex_req_"+ self._contact_names[i]
+            next_idx+=len(self._contact_names)
+            for i in range(len(self._contact_names)):
+                obs_names[next_idx+i] = "flight_end_req_"+ self._contact_names[i]
+            next_idx+=len(self._contact_names)
+
+        if self._env_opts["add_rhc_cmds_to_obs"]:
+            self._obs_map["rhc_cmds_q"]=next_idx
+            for i in range(self._n_jnts): # jnt obs (pos):
+                obs_names[next_idx+i] = f"rhc_cmd_q_{jnt_names[i]}"
+            next_idx+=self._n_jnts
+            self._obs_map["rhc_cmds_v"]=next_idx
+            for i in range(self._n_jnts): # jnt obs (pos):
+                obs_names[next_idx+i] = f"rhc_cmd_v_{jnt_names[i]}"
+            next_idx+=self._n_jnts
+            self._obs_map["rhc_cmds_eff"]=next_idx
+            for i in range(self._n_jnts): # jnt obs (pos):
+                obs_names[next_idx+i] = f"rhc_cmd_eff_{jnt_names[i]}"
+            next_idx+=self._n_jnts
+
+        # previous actions info
+        if self._env_opts["use_action_history"]:
+            self._obs_map["action_history"]=next_idx
+            action_names = self._get_action_names()
+            if self._env_opts["add_prev_actions_stats_to_obs"]:
+                self._obs_map["action_history_prev"]=next_idx
+                for act_idx in range(self.actions_dim()):
+                    obs_names[next_idx+act_idx] = action_names[act_idx]+f"_prev_act"
+                next_idx+=self.actions_dim()
+                self._obs_map["action_history_avrg"]=next_idx
+                for act_idx in range(self.actions_dim()):
+                    obs_names[next_idx+act_idx] = action_names[act_idx]+f"_avrg_act"
+                next_idx+=self.actions_dim()
+                self._obs_map["action_history_std"]=next_idx
+                for act_idx in range(self.actions_dim()):
+                    obs_names[next_idx+act_idx] = action_names[act_idx]+f"_std_act"
+                next_idx+=self.actions_dim()
+            else:
+                for i in range(self._env_opts["actions_history_size"]):
+                    for act_idx in range(self.actions_dim()):
+                        obs_names[next_idx+act_idx] = action_names[act_idx]+f"_m{i+1}_act"
+                    next_idx+=self.actions_dim()
+
+        if self._env_opts["use_action_smoothing"]:
+            self._obs_map["action_smoothing"]=next_idx
+            for smoothed_action in range(self.actions_dim()):
+                obs_names[next_idx+smoothed_action] = action_names[smoothed_action]+f"_smoothed"
+            next_idx+=self.actions_dim()
+
+        if self._env_opts["add_periodic_clock_to_obs"]:
+            self._obs_map["clock"]=next_idx
+            obs_names[next_idx] = "clock_cos"
+            obs_names[next_idx+1] = "clock_sin"
+            next_idx+=2
+        if self._env_opts["add_heightmap_obs"] and self._height_grid_size is not None:
+            self._obs_map["heightmap"]=next_idx
+            gs = self._height_grid_size
+            for r in range(gs):
+                for c in range(gs):
+                    obs_names[next_idx] = f"height_r{r}_c{c}"
+                    next_idx += 1
+
+        return obs_names
+
+    def _set_substep_obs(self):
+        # which obs are to be averaged over substeps?
+
+        self._is_substep_obs[self._obs_map["linvel_meas"]:self._obs_map["linvel_meas"]+3]=True
+        self._is_substep_obs[self._obs_map["omega_meas"]:self._obs_map["omega_meas"]+3]=True
+        self._is_substep_obs[self._obs_map["v_jnt"]:self._obs_map["v_jnt"]+self._n_jnts]=True # also good for noise
+
+        # self._is_substep_obs[self._obs_map["contact_f_mpc"]:self._obs_map["contact_f_mpc"]+3*len(self._contact_names)]=True
+
+    def _get_action_names(self):
+        
+        action_names = [""] * self.actions_dim()
+        action_names[0] = "vx_cmd" # twist commands from agent to RHC controller
+        action_names[1] = "vy_cmd"
+        action_names[2] = "vz_cmd"
+        action_names[3] = "roll_omega_cmd"
+        action_names[4] = "pitch_omega_cmd"
+        action_names[5] = "yaw_omega_cmd"
+
+        next_idx=6
+        
+        self._actions_map["contact_flag_start"]=next_idx
+        for i in range(len(self._contact_names)):
+            contact=self._contact_names[i]
+            action_names[next_idx] = f"contact_flag_{contact}"
+            next_idx+=1
+
+        return action_names
+    
+    def _set_substep_rew(self):
+
+        # which rewards are to be computed at substeps frequency?
+        self._is_substep_rew[self._reward_map["task_error"]]=False
+        if self._env_opts["add_CoT_reward"]:
+            self._is_substep_rew[self._reward_map["CoT"]]=True
+        if self._env_opts["add_power_reward"]:
+            self._is_substep_rew[self._reward_map["mech_pow"]]=True
+        if self._env_opts["add_action_rate_reward"]:
+            self._is_substep_rew[self._reward_map["action_rate"]]=False
+        if self._env_opts["add_jnt_v_reward"]:
+            self._is_substep_rew[self._reward_map["jnt_v"]]=True
+
+        if self._env_opts["use_rhc_avrg_vel_tracking"]:
+            self._is_substep_rew[self._reward_map["rhc_avrg_vel_error"]]=False
+
+    def _get_rewards_names(self):
+        
+        counter=0
+        reward_names = []
+
+        # adding rewards
+        reward_names.append("task_error")
+        self._reward_map["task_error"]=counter
+        self._reward_lb_map["task_error"]="task_error_reward_lb"
+        counter+=1
+        if self._env_opts["add_power_reward"] and self._env_opts["add_CoT_reward"]:
+            Journal.log(self.__class__.__name__,
+                    "__init__",
+                    "Only one between CoT and power reward can be used!",
+                    LogType.EXCEP,
+                    throw_when_excep=True)
+        if self._env_opts["add_CoT_reward"]:
+            reward_names.append("CoT")
+            self._reward_map["CoT"]=counter
+            self._reward_lb_map["CoT"]="CoT_reward_lb"
+            counter+=1
+        if self._env_opts["add_power_reward"]:
+            reward_names.append("mech_pow")
+            self._reward_map["mech_pow"]=counter
+            self._reward_lb_map["mech_pow"]="power_reward_lb"
+            counter+=1
+        if self._env_opts["add_action_rate_reward"]:
+            reward_names.append("action_rate")   
+            self._reward_map["action_rate"]=counter
+            self._reward_lb_map["action_rate"]="action_rate_reward_lb"
+            counter+=1   
+        if self._env_opts["add_jnt_v_reward"]:
+            reward_names.append("jnt_v")   
+            self._reward_map["jnt_v"]=counter
+            self._reward_lb_map["jnt_v"]="jnt_vel_reward_lb"
+            counter+=1   
+        if self._env_opts["use_rhc_avrg_vel_tracking"]:
+            reward_names.append("rhc_avrg_vel_error")   
+            self._reward_map["rhc_avrg_vel_error"]=counter
+            self._reward_lb_map["rhc_avrg_vel_error"]="rhc_avrg_vel_reward_lb"
+            counter+=1   
+
+        return reward_names
+
+    def _get_sub_trunc_names(self):
+        sub_trunc_names = []
+        sub_trunc_names.append("ep_timeout")
+        if self._env_opts["single_task_ref_per_episode"]:
+            sub_trunc_names.append("task_ref_rand")
+        return sub_trunc_names
+
+    def _get_sub_term_names(self):
+        # to be overridden by child class
+        sub_term_names = []
+        sub_term_names.append("rhc_failure")
+        sub_term_names.append("robot_capsize")
+        if self._env_opts["add_term_mpc_capsize"]:
+            sub_term_names.append("rhc_capsize")
+
+        return sub_term_names
+
+    def _set_jnts_blacklist_pattern(self):
+        # used to exclude pos measurement from wheels
+        self._jnt_q_blacklist_patterns=["wheel"]

world_interface_base.py

@@ -0,0 +1,1719 @@
+from aug_mpc.controllers.rhc.augmpc_cluster_server import AugMpcClusterServer
+from aug_mpc.utils.shared_data.remote_stepping import RemoteStepperClnt
+from aug_mpc.utils.shared_data.remote_stepping import RemoteResetClnt
+from aug_mpc.utils.shared_data.remote_stepping import RemoteResetRequest
+from aug_mpc.utils.jnt_imp_control_base import JntImpCntrlBase
+from aug_mpc.utils.hybrid_quad_xrdf_gen import get_xrdf_cmds
+from aug_mpc.utils.xrdf_gen import generate_srdf, generate_urdf
+from aug_mpc.utils.math_utils import quaternion_difference
+from aug_mpc.utils.custom_arg_parsing import extract_custom_xacro_args, merge_xacro_cmds
+
+from aug_mpc.utils.filtering import FirstOrderFilter
+
+from mpc_hive.utilities.homing import RobotHomer
+from mpc_hive.utilities.shared_data.jnt_imp_control import JntImpCntrlData
+
+from EigenIPC.PyEigenIPC import VLevel, Journal, LogType, dtype
+from EigenIPC.PyEigenIPC import StringTensorServer
+from EigenIPC.PyEigenIPCExt.wrappers.shared_data_view import SharedTWrapper
+
+from typing import List, Dict, TypeVar
+
+import os
+import inspect
+import signal
+import time
+
+import numpy as np
+import torch
+
+from abc import ABC, abstractmethod
+
+JntImpCntrlChild = TypeVar('JntImpCntrlChild', bound='JntImpCntrlBase')
+
+class AugMPCWorldInterfaceBase(ABC):
+
+    def __init__(self,
+                robot_names: List[str],
+                robot_urdf_paths: List[str],
+                robot_srdf_paths: List[str],
+                jnt_imp_config_paths: List[str],
+                n_contacts: List[int],
+                cluster_dt: List[float],
+                use_remote_stepping: List[bool],
+                name: str = "AugMPCWorldInterfaceBase",
+                num_envs: int = 1,
+                debug = False,
+                verbose: bool = False,
+                vlevel: VLevel = VLevel.V1,
+                n_init_step: int = 0,
+                timeout_ms: int = 60000,
+                env_opts: Dict = None,
+                use_gpu: bool = True,
+                dtype: torch.dtype = torch.float32,
+                dump_basepath: str = "/tmp",
+                override_low_lev_controller: bool = False):
+
+        # checks on input args
+        # type checks
+        if not isinstance(robot_names, List):
+            exception = "robot_names must be a list!"
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+        if not isinstance(robot_urdf_paths, List):
+            exception = "robot_urdf_paths must be a list!"
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+        if not isinstance(robot_srdf_paths, List):
+            exception = "robot_srdf_paths must be a list!"
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+        if not isinstance(cluster_dt, List):
+            exception = "cluster_dt must be a list!"
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+        if not isinstance(use_remote_stepping, List):
+            exception = "use_remote_stepping must be a list!"
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+        if not isinstance(n_contacts, List):
+            exception = "n_contacts must be a list (of integers)!"
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+        if not isinstance(jnt_imp_config_paths, List):
+            exception = "jnt_imp_config_paths must be a list paths!"
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+            
+        # dim checks
+        if not len(robot_urdf_paths) == len(robot_names):
+            exception = f"robot_urdf_paths has len {len(robot_urdf_paths)}" + \
+             f" while robot_names {len(robot_names)}"
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+        if not len(robot_srdf_paths) == len(robot_names):
+            exception = f"robot_srdf_paths has len {len(robot_srdf_paths)}" + \
+             f" while robot_names {len(robot_names)}"
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+        if not len(cluster_dt) == len(robot_names):
+            exception = f"cluster_dt has len {len(cluster_dt)}" + \
+             f" while robot_names {len(robot_names)}"
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+        if not len(use_remote_stepping) == len(robot_names):
+            exception = f"use_remote_stepping has len {len(use_remote_stepping)}" + \
+             f" while robot_names {len(robot_names)}"
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+        if not len(robot_srdf_paths) == len(robot_names):
+            exception = f"robot_srdf_paths has len {len(robot_srdf_paths)}" + \
+             f" while robot_names {len(robot_names)}"
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+        if not len(jnt_imp_config_paths) == len(robot_names):
+            exception = f"jnt_imp_config_paths has len {len(jnt_imp_config_paths)}" + \
+             f" while robot_names {len(robot_names)}"
+            Journal.log(self.__class__.__name__,
+                "__init__",
+                exception,
+                LogType.EXCEP,
+                throw_when_excep = True)
+        
+        self._remote_exit_flag=None
+
+        self._name=name
+        self._num_envs=num_envs
+        self._debug=debug
+        self._verbose=verbose
+        self._vlevel=vlevel
+        self._force_reconnection=True
+        self._timeout_ms=timeout_ms
+        self._use_gpu=use_gpu
+        self._device = "cuda" if self._use_gpu else "cpu"
+        self._dtype=dtype
+        self._robot_names=robot_names
+        self._env_opts={}
+        self._env_opts["deact_when_failure"]=True
+        self._env_opts["filter_jnt_vel"]=False
+        self._env_opts["filter_cutoff_freq"]=10.0 # [Hz]
+        self._env_opts["filter_sampling_rate"]=100 # rate at which state is filtered [Hz]
+        self._env_opts["add_remote_exit_flag"]=False # add shared data server to trigger a remote exit
+        self._env_opts["wheel_joint_patterns"]=["wheel"]
+        self._env_opts["filter_wheel_pos_ref"]=True
+        self._env_opts["zero_wheel_eff_ref"]=True
+
+        self._env_opts["enable_height_sensor"]=False
+        self._env_opts["height_sensor_resolution"]=0.16
+        self._env_opts["height_sensor_pixels"]=10
+        self._env_opts["height_sensor_lateral_offset"]=0.0
+        self._env_opts["height_sensor_forward_offset"]=0.0
+
+        self._env_opts["run_cluster_bootstrap"] = False
+        
+        self._filter_step_ssteps_freq=None
+
+        self._env_opts.update(env_opts)
+
+        self.step_counter = 0 # global step counter
+        self._n_init_steps = n_init_step # n steps to be performed before applying solutions from control clusters
+        self._srdf_dump_paths = robot_srdf_paths
+        self._homers = {} 
+        self._homing = None
+        self._jnt_imp_cntrl_shared_data = {}
+        self._jnt_imp_controllers = {}
+        self._jnt_imp_config_paths = {}
+
+        # control cluster data
+        self.cluster_sim_step_counters = {}
+        self.cluster_servers = {}
+        self._trigger_sol = {}
+        self._wait_sol = {}
+        self._cluster_dt = {}
+        self._robot_urdf_paths={}
+        self._robot_srdf_paths={}
+        self._contact_names={}
+        self._num_contacts={}
+
+        for i in range(len(self._robot_names)):
+            robot_name = self._robot_names[i]
+            self._cluster_dt[robot_name]=cluster_dt[i]
+            self._robot_urdf_paths[robot_name]=robot_urdf_paths[i]
+            self._robot_srdf_paths[robot_name]=robot_srdf_paths[i]
+            self._contact_names[robot_name]=None
+            self._num_contacts[robot_name]=n_contacts[i]
+            self._jnt_imp_config_paths[robot_name]=jnt_imp_config_paths[i]
+        # db data
+        self.debug_data = {}
+        self.debug_data["time_to_step_world"] = np.nan
+        self.debug_data["time_to_get_states_from_env"] = np.nan
+        self.debug_data["cluster_sol_time"] = {}
+        self.debug_data["cluster_state_update_dt"] = {}
+        self.debug_data["sim_time"] = {}
+        self.debug_data["cluster_time"] = {}
+        
+        self._env_timer = time.perf_counter()
+
+        # remote sim stepping options
+        self._timeout = timeout_ms # timeout for remote stepping
+        self._use_remote_stepping = use_remote_stepping
+        # should use remote stepping
+        self._remote_steppers = {}
+        self._remote_resetters = {}
+        self._remote_reset_requests = {}
+        self._is_first_trigger = {}
+
+        self._closed = False
+             
+        self._this_child_path=os.path.abspath(inspect.getfile(self.__class__))
+        self._descr_dump_path=dump_basepath+"/"+f"{self.__class__.__name__}"
+        self._urdf_dump_paths = {}
+        self._srdf_dump_paths = {}
+        self.xrdf_cmd_vals = [] # by default empty, needs to be overriden by
+        # child class
+        self._world_iface_files_server=None
+
+        self._override_low_lev_controller=override_low_lev_controller
+
+        self._root_p = {}
+        self._root_q = {}
+        self._jnts_q = {} 
+        self._root_p_prev = {} # used for num differentiation
+        self._root_q_prev = {} # used for num differentiation
+        self._jnts_q_prev = {} # used for num differentiation
+        self._root_v_prev = {} # used for num differentiation
+        self._root_omega_prev = {} # used for num differentiation
+        self._root_p_default = {} 
+        self._root_q_default = {}
+        self._jnts_q_default = {}
+        
+        self._gravity_normalized = {}
+        self._gravity_normalized_base_loc = {}
+
+        self._root_v = {}
+        self._root_v_base_loc = {}
+        self._root_v_default = {}
+        self._root_omega = {}
+        self._root_omega_base_loc = {}
+        self._root_omega_default = {}
+        self._root_a = {}
+        self._root_a_base_loc = {}
+        self._root_alpha = {}
+        self._root_alpha_base_loc = {}
+
+        self._jnts_v = {}
+        self._jnt_vel_filter = {}
+        self._jnts_v_default = {}
+        self._jnts_eff = {}
+        self._jnts_eff_default = {}
+
+        self._root_pos_offsets = {} 
+        self._root_q_offsets = {} 
+        self._root_q_offsets_yaw = {}
+        self._root_q_yaw_rel_ws = {}
+
+        self._parse_env_opts()
+
+        self._enable_height_shared = self._env_opts["enable_height_sensor"]
+        self._height_sensor_resolution = self._env_opts["height_sensor_resolution"]
+        self._height_sensor_pixels = self._env_opts["height_sensor_pixels"]
+
+        self._pre_setup() # child's method
+
+        self._init_world() # after this point all info from sim or robot is 
+        # available
+
+        self._publish_world_interface_files()
+        
+        setup_ok=self._setup()
+        if not setup_ok:
+            self.close()
+        
+        self._exit_request=False
+        signal.signal(signal.SIGINT, self.signal_handler)   
+
+    def signal_handler(self, sig, frame):
+        Journal.log(self.__class__.__name__,
+            "signal_handler",
+            "received SIGINT -> cleaning up",
+            LogType.WARN)
+        self._exit_request=True
+    
+    def __del__(self):
+        self.close()
+    
+    def is_closed(self):
+        return self._closed
+    
+    def close(self) -> None:
+        if not self._closed:
+            for i in range(len(self._robot_names)):
+                if self._robot_names[i] in self.cluster_servers:
+                    self.cluster_servers[self._robot_names[i]].close()
+                if self._use_remote_stepping[i]: # remote signaling
+                    if self._robot_names[i] in self._remote_reset_requests:
+                        self._remote_reset_requests[self._robot_names[i]].close()
+                        self._remote_resetters[self._robot_names[i]].close()
+                        self._remote_steppers[self._robot_names[i]].close()
+                if self._robot_names[i] in self._jnt_imp_cntrl_shared_data:
+                    jnt_imp_shared_data=self._jnt_imp_cntrl_shared_data[self._robot_names[i]]
+                    if jnt_imp_shared_data is not None:
+                        jnt_imp_shared_data.close()
+            if self._remote_exit_flag is not None:
+                self._remote_exit_flag.close()
+            if self._world_iface_files_server is not None:
+                self._world_iface_files_server.close()
+            self._close()
+            self._closed=True
+
+    def _collect_world_interface_files(self):
+        files = [self._this_child_path]
+        # prefer generated URDF/SRDF if available, fallback to provided xacros
+        if len(self._urdf_dump_paths) > 0:
+            files.extend(list(self._urdf_dump_paths.values()))
+        else:
+            files.extend(list(self._robot_urdf_paths.values()))
+        if len(self._srdf_dump_paths) > 0:
+            files.extend(list(self._srdf_dump_paths.values()))
+        else:
+            files.extend(list(self._robot_srdf_paths.values()))
+        files.extend(list(self._jnt_imp_config_paths.values()))
+        # remove duplicates while preserving order
+        unique_files=[]
+        for f in files:
+            if f not in unique_files:
+                unique_files.append(f)
+        return unique_files
+
+    def _publish_world_interface_files(self):
+
+        if not any(self._use_remote_stepping):
+            return
+        self._world_iface_files_server=StringTensorServer(length=1,
+            basename="SharedWorldInterfaceFilesDropDir",
+            name_space=self._robot_names[0],
+            verbose=self._verbose,
+            vlevel=self._vlevel,
+            force_reconnection=True)
+        self._world_iface_files_server.run()
+        combined_paths=", ".join(self._collect_world_interface_files())
+        while not self._world_iface_files_server.write_vec([combined_paths], 0):
+            Journal.log(self.__class__.__name__,
+            "_publish_world_interface_files",
+            f"Failed to pub world interface files. Retrying...",
+            LogType.WARN)
+            time.sleep(0.1)
+        Journal.log(self.__class__.__name__,
+            "_publish_world_interface_files",
+            f"World interface files advertised: {combined_paths}",
+            LogType.STAT)
+
+    def _setup(self) -> bool:
+    
+        for i in range(len(self._robot_names)):
+            robot_name = self._robot_names[i]
+
+            # normalized gravity vector
+            self._gravity_normalized[robot_name]=torch.full_like(self._root_v[robot_name], fill_value=0.0)
+            self._gravity_normalized[robot_name][:, 2]=-1.0
+            self._gravity_normalized_base_loc[robot_name]=self._gravity_normalized[robot_name].detach().clone()
+
+            # Pre-allocate yaw-related buffers once and reuse them in root_q_yaw_rel().
+            q_ref = self._root_q[robot_name]
+            self._root_q_offsets_yaw[robot_name] = torch.zeros(
+                (self._num_envs,), dtype=q_ref.dtype, device=q_ref.device)
+            self._root_q_yaw_rel_ws[robot_name] = {
+                "yaw_abs": torch.zeros((self._num_envs,), dtype=q_ref.dtype, device=q_ref.device),
+                "yaw_rel": torch.zeros((self._num_envs,), dtype=q_ref.dtype, device=q_ref.device),
+                "yaw_sin": torch.zeros((self._num_envs,), dtype=q_ref.dtype, device=q_ref.device),
+                "yaw_cos": torch.zeros((self._num_envs,), dtype=q_ref.dtype, device=q_ref.device),
+                "q_abs_unit": torch.zeros_like(q_ref),
+                "q_yaw_abs": torch.zeros_like(q_ref),
+                "q_yaw_rel": torch.zeros_like(q_ref),
+                "q_yaw_abs_conj": torch.zeros_like(q_ref),
+                "q_pr": torch.zeros_like(q_ref),
+                "q_rel": torch.zeros_like(q_ref),
+            }
+
+            self.cluster_sim_step_counters[robot_name]=0
+            self._is_first_trigger[robot_name] = True
+            if not isinstance(self._cluster_dt[robot_name], (float)):
+                exception = f"cluster_dt[{i}] should be a float!"
+                Journal.log(self.__class__.__name__,
+                    "_setup",
+                    exception,
+                    LogType.EXCEP,
+                    throw_when_excep = False)
+                return False
+            self._cluster_dt[robot_name] = self._cluster_dt[robot_name]
+            self._trigger_sol[robot_name] = True # allow first trigger
+            self._wait_sol[robot_name] = False
+
+            # initialize a lrhc cluster server for communicating with rhc controllers
+            self.cluster_servers[robot_name] = AugMpcClusterServer(cluster_size=self._num_envs, 
+                        cluster_dt=self._cluster_dt[robot_name], 
+                        control_dt=self.physics_dt(), 
+                        jnt_names=self._robot_jnt_names(robot_name=robot_name), 
+                        n_contacts=self._n_contacts(robot_name=robot_name),
+                        contact_linknames=self._contact_names[robot_name], 
+                        verbose=self._verbose, 
+                        vlevel=self._vlevel,
+                        debug=self._debug, 
+                        robot_name=robot_name,
+                        use_gpu=self._use_gpu,
+                        force_reconnection=self._force_reconnection,
+                        timeout_ms=self._timeout,
+                        enable_height_sensor=self._enable_height_shared,
+                        height_grid_size=self._height_sensor_pixels,
+                        height_grid_resolution=self._height_sensor_resolution)
+            self.cluster_servers[robot_name].run()
+            self.debug_data["cluster_sol_time"][robot_name] = np.nan
+            self.debug_data["cluster_state_update_dt"][robot_name] = np.nan
+            self.debug_data["sim_time"][robot_name] = np.nan
+            # remote sim stepping
+            if self._use_remote_stepping[i]:
+                self._remote_steppers[robot_name] = RemoteStepperClnt(namespace=robot_name,
+                                                            verbose=self._debug,
+                                                            vlevel=self._vlevel)
+                self._remote_resetters[robot_name] = RemoteResetClnt(namespace=robot_name,
+                                                            verbose=self._debug,
+                                                            vlevel=self._vlevel)
+                self._remote_reset_requests[robot_name] = RemoteResetRequest(namespace=robot_name,
+                                                                    n_env=self._num_envs,
+                                                                    is_server=True,
+                                                                    verbose=self._debug,
+                                                                    vlevel=self._vlevel, 
+                                                                    force_reconnection=self._force_reconnection, 
+                                                                    safe=False)
+                self._remote_steppers[robot_name].run()
+                self._remote_resetters[robot_name].run()
+                self._remote_reset_requests[robot_name].run()
+            else:
+                self._remote_steppers[robot_name] = None
+                self._remote_reset_requests[robot_name] = None
+                self._remote_resetters[robot_name] = None
+
+            self._homers[robot_name] = RobotHomer(srdf_path=self._srdf_dump_paths[robot_name], 
+                            jnt_names=self._robot_jnt_names(robot_name=robot_name),
+                            filter=True,
+                            verbose=self._verbose)
+            robot_homing=torch.from_numpy(self._homers[robot_name].get_homing().reshape(1,-1))
+            if "cuda" in self._device:
+                robot_homing=robot_homing.cuda()
+            self._homing=robot_homing.repeat(self._num_envs, 1)
+
+            self._jnts_q_default[robot_name] = self._homing
+            self._set_jnts_to_homing(robot_name=robot_name)
+            self._set_root_to_defconfig(robot_name=robot_name)
+            self._reset_sim()
+
+            self._init_safe_cluster_actions(robot_name=robot_name)
+
+            Journal.log(self.__class__.__name__,
+                "_setup",
+                f"Will use joint impedance config at {self._jnt_imp_config_paths[robot_name]} for {robot_name}",
+                LogType.STAT)
+            
+            self._jnt_imp_controllers[robot_name] = self._generate_jnt_imp_control(robot_name=robot_name)
+            self._jnt_imp_controllers[robot_name].set_velocity_controlled_joints(
+                name_patterns=self._env_opts["wheel_joint_patterns"],
+                filter_pos_ref=self._env_opts["filter_wheel_pos_ref"],
+                zero_eff_ref=self._env_opts["zero_wheel_eff_ref"])
+            self._jnt_imp_cntrl_shared_data[robot_name] = JntImpCntrlData(is_server=True, 
+                                            n_envs=self._num_envs, 
+                                            n_jnts=len(self._robot_jnt_names(robot_name=robot_name)),
+                                            jnt_names=self._robot_jnt_names(robot_name=robot_name),
+                                            namespace=robot_name, 
+                                            verbose=self._verbose, 
+                                            force_reconnection=self._force_reconnection,
+                                            vlevel=self._vlevel,
+                                            use_gpu=self._use_gpu,
+                                            safe=False)
+            self._jnt_imp_cntrl_shared_data[robot_name].run()
+
+            self._jnt_vel_filter[robot_name]=None
+            if self._env_opts["filter_jnt_vel"]:
+                self._jnt_vel_filter[robot_name]=FirstOrderFilter(dt=1.0/self._env_opts["filter_sampling_rate"],
+                    filter_BW=self._env_opts["filter_cutoff_freq"],
+                    rows=self._num_envs,
+                    cols=len(self._robot_jnt_names(robot_name=robot_name)),
+                    device=self._device,
+                    dtype=self._dtype)
+                
+                physics_rate=1.0/self.physics_dt()
+                self._filter_step_ssteps_freq=int(physics_rate/self._env_opts["filter_sampling_rate"])
+                if self._filter_step_ssteps_freq <=0:
+                    Journal.log(self.__class__.__name__,
+                        "_setup",
+                        f"The filter_sampling_rate should be smaller that the physics rate ({physics_rate} Hz)",
+                        LogType.EXCEP,
+                        throw_when_excep=True)
+
+            for n in range(self._n_init_steps): # run some initialization steps
+                if hasattr(self, "_alter_twist_warmup"):
+                    self._alter_twist_warmup(robot_name=robot_name, env_indxs=None)
+                self._step_world()
+               
+            self._read_jnts_state_from_robot(robot_name=robot_name,
+                env_indxs=None)
+            self._read_root_state_from_robot(robot_name=robot_name,
+                    env_indxs=None)
+            # allow child to perform additional warmup validations (e.g., terrain/tilt)
+            # retry_done = False
+            if hasattr(self, "_post_warmup_validation"):
+                failing = self._post_warmup_validation(robot_name=robot_name)
+                if failing is not None and failing.numel() > 0:
+                    # retry: reset only failing envs, rerun warmup, revalidate once
+                    failing = failing.to(self._device)
+                    Journal.log(self.__class__.__name__,
+                        "_setup",
+                        f"Warmup validation failed for {robot_name}, envs indexes {failing.tolist()}",
+                        LogType.EXCEP,
+                        throw_when_excep=True)
+                else:
+                    Journal.log(self.__class__.__name__,
+                        "_setup",
+                        f"Warmup validation passed for {robot_name}",
+                        LogType.INFO)
+            
+            # write some inits for all robots
+            # self._update_root_offsets(robot_name)
+            self._synch_default_root_states(robot_name=robot_name)
+            epsi=0.03 # adding a bit of height to avoid initial penetration
+            self._root_p_default[robot_name][:, 2]=self._root_p_default[robot_name][:, 2]+epsi
+            
+            reset_ok=self._reset(env_indxs=None,
+                robot_name=robot_name,
+                reset_cluster=True,
+                reset_cluster_counter=False,
+                randomize=True,
+                acquire_offsets=True) # resets everything, updates the cluster with fresh reset states 
+            # and acquire offsets
+            if not reset_ok:
+                return False
+            
+            # cluster setup here
+            control_cluster=self.cluster_servers[robot_name]
+
+            control_cluster.pre_trigger()
+            to_be_activated=control_cluster.get_inactive_controllers()
+            if to_be_activated is not None:
+                control_cluster.activate_controllers(
+                    idxs=to_be_activated)   
+
+            if self._env_opts["run_cluster_bootstrap"]:
+                cluster_setup_ok=self._setup_mpc_cluster(robot_name)
+                if not cluster_setup_ok:
+                    return False
+                self._set_cluster_actions(robot_name=robot_name) # write last cmds
+                self._apply_cmds_to_jnt_imp_control(robot_name=robot_name) # apply to robot
+            
+            if self._use_remote_stepping[i]:
+                step_wait_ok = self._wait_for_remote_step_req(robot_name=robot_name)
+                if not step_wait_ok:
+                   return False
+            
+            self._set_startup_jnt_imp_gains(robot_name=robot_name) # set gains to
+            # startup config (usually lower)
+
+            control_cluster.pre_trigger()
+            control_cluster.trigger_solution(bootstrap=False) # trigger first solution (in real-time iteration) before first call to step to ensure that first solution is ready when step is called the first time
+            
+        if self._env_opts["add_remote_exit_flag"]:
+            self._remote_exit_flag=SharedTWrapper(namespace = self._robot_names[0],# use first robot as name
+                basename = "IbridoRemoteEnvExitFlag",
+                is_server = True, 
+                n_rows = 1, 
+                n_cols = 1, 
+                verbose = True, 
+                vlevel = self._vlevel,
+                safe = False,
+                dtype=dtype.Bool,
+                force_reconnection=True,
+                fill_value = False)
+            self._remote_exit_flag.run()
+
+        self._setup_done=True
+
+        return self._setup_done
+
+    def _setup_mpc_cluster(self, robot_name: str):
+
+        control_cluster = self.cluster_servers[robot_name]
+        
+        # self._set_state_to_cluster(robot_name=robot_name)
+        rhc_state = control_cluster.get_state()
+        root_twist=rhc_state.root_state.get(data_type="twist", robot_idxs = None, gpu=self._use_gpu)
+        jnt_v=rhc_state.jnts_state.get(data_type="v", robot_idxs = None, gpu=self._use_gpu) 
+        root_twist[:, :]=0 # override meas state to make sure MPC bootstrap uses zero velocity
+        jnt_v[:, :]=0 
+
+        control_cluster.write_robot_state()
+        
+        # trigger bootstrap solution (solvers will run up to convergence) 
+        control_cluster.trigger_solution(bootstrap=True) # this will trigger the bootstrap solver with the initial state,
+        # which will run until convergence before returning
+        wait_ok=control_cluster.wait_for_solution() # blocking
+        if not wait_ok:
+            return False
+        failed = control_cluster.get_failed_controllers(gpu=self._use_gpu)
+        if failed is not None:
+            failed_idxs = torch.nonzero(failed).squeeze(-1)
+            if failed_idxs.numel() > 0:
+                Journal.log(self.__class__.__name__,
+                    "_setup",
+                    f"Bootstrap solution failed for {robot_name} | n_failed: {failed_idxs.numel()}, idxs: {failed_idxs.cpu().tolist()}",
+                    LogType.EXCEP,
+                    throw_when_excep=False)
+                return False
+
+        return True
+
+    def step(self) -> bool:
+
+        success=False
+
+        if self._remote_exit_flag is not None:
+            # check for exit request
+            self._remote_exit_flag.synch_all(read=True, retry = False)
+            self._exit_request=self._exit_request or \
+                bool(self._remote_exit_flag.get_numpy_mirror()[0, 0].item())
+
+        if self._exit_request:
+            self.close()
+            
+        if self.is_running() and (not self.is_closed()):
+            if self._debug:
+                pre_step_ok=self._pre_step_db()
+                if not pre_step_ok:
+                    return False
+                self._env_timer=time.perf_counter()
+                self._step_world()
+                self.debug_data["time_to_step_world"] = \
+                    time.perf_counter() - self._env_timer
+                self._post_world_step_db()
+                success=True
+            else:
+                pre_step_ok=self._pre_step()
+                if not pre_step_ok:
+                    return False
+                self._step_world()
+                self._post_world_step()
+                success=True
+        
+        return success
+    
+    def render(self, mode:str="human") -> None:
+        self._render_sim(mode)
+
+    def reset(self,
+        env_indxs: torch.Tensor = None,
+        reset_cluster: bool = False,
+        reset_cluster_counter = False,
+        randomize: bool = False,
+        reset_sim: bool = False) -> None:
+
+        for i in range(len(self._robot_names)):
+            robot_name=self._robot_names[i]
+            reset_ok=self._reset(robot_name=robot_name,
+                env_indxs=env_indxs,
+                randomize=randomize,
+                reset_cluster=reset_cluster,
+                reset_cluster_counter=reset_cluster_counter)
+            if not reset_ok:
+                return False    
+            self._set_startup_jnt_imp_gains(robot_name=robot_name,
+                env_indxs=env_indxs)
+            
+        if reset_sim:
+            self._reset_sim()
+        
+        return True
+
+    def _reset_cluster(self,
+            robot_name: str,
+            env_indxs: torch.Tensor = None,
+            reset_cluster_counter: bool = False):
+        
+        control_cluster = self.cluster_servers[robot_name]
+
+        reset_ok=control_cluster.reset_controllers(idxs=env_indxs)
+        if not reset_ok:
+            return False
+        
+        self._set_state_to_cluster(robot_name=robot_name,
+            env_indxs=env_indxs)
+        control_cluster.write_robot_state() # writes to shared memory
+
+        if reset_cluster_counter:
+            self.cluster_sim_step_counters[robot_name] = 0 
+        
+        return True
+
+    def _step_jnt_vel_filter(self,
+            robot_name: str, 
+            env_indxs: torch.Tensor = None):
+        
+        self._jnt_vel_filter[robot_name].update(refk=self.jnts_v(robot_name=robot_name, env_idxs=env_indxs), 
+            idxs=env_indxs)
+
+    def _set_state_to_cluster(self, 
+        robot_name: str, 
+        env_indxs: torch.Tensor = None,
+        base_loc: bool = True):
+        
+        if self._debug:
+            if not isinstance(env_indxs, (torch.Tensor, type(None))):
+                msg = "Provided env_indxs should be a torch tensor of indexes!"
+                raise Exception(f"[{self.__class__.__name__}]" + f"[{self.journal.exception}]: " + msg)
+            
+        control_cluster = self.cluster_servers[robot_name]
+        # floating base
+        rhc_state = control_cluster.get_state()
+        # configuration
+        rhc_state.root_state.set(data=self.root_p_rel(robot_name=robot_name, env_idxs=env_indxs), 
+                data_type="p", robot_idxs = env_indxs, gpu=self._use_gpu)
+        rhc_state.root_state.set(data=self.root_q(robot_name=robot_name, env_idxs=env_indxs), 
+                data_type="q", robot_idxs = env_indxs, gpu=self._use_gpu)
+        # rhc_state.root_state.set(data=self.root_q_yaw_rel(robot_name=robot_name, env_idxs=env_indxs), 
+        #         data_type="q", robot_idxs = env_indxs, gpu=self._use_gpu)
+        # twist
+        rhc_state.root_state.set(data=self.root_v(robot_name=robot_name, env_idxs=env_indxs,base_loc=base_loc), 
+                data_type="v", robot_idxs = env_indxs, gpu=self._use_gpu)
+        rhc_state.root_state.set(data=self.root_omega(robot_name=robot_name, env_idxs=env_indxs,base_loc=base_loc), 
+                data_type="omega", robot_idxs = env_indxs, gpu=self._use_gpu)
+        
+        # angular accc.
+        rhc_state.root_state.set(data=self.root_a(robot_name=robot_name, env_idxs=env_indxs,base_loc=base_loc), 
+                data_type="a", robot_idxs = env_indxs, gpu=self._use_gpu)
+        rhc_state.root_state.set(data=self.root_alpha(robot_name=robot_name, env_idxs=env_indxs,base_loc=base_loc), 
+                data_type="alpha", robot_idxs = env_indxs, gpu=self._use_gpu)
+        
+        # gravity vec
+        rhc_state.root_state.set(data=self.gravity(robot_name=robot_name, env_idxs=env_indxs,base_loc=base_loc), 
+                data_type="gn", robot_idxs = env_indxs, gpu=self._use_gpu)
+        
+        # joints
+        rhc_state.jnts_state.set(data=self.jnts_q(robot_name=robot_name, env_idxs=env_indxs), 
+            data_type="q", robot_idxs = env_indxs, gpu=self._use_gpu)
+        
+        v_jnts=self.jnts_v(robot_name=robot_name, env_idxs=env_indxs)
+        if self._jnt_vel_filter[robot_name] is not None: # apply filtering
+            v_jnts=self._jnt_vel_filter[robot_name].get(idxs=env_indxs)
+        rhc_state.jnts_state.set(data=v_jnts, 
+            data_type="v", robot_idxs = env_indxs, gpu=self._use_gpu) 
+        rhc_state.jnts_state.set(data=self.jnts_eff(robot_name=robot_name, env_idxs=env_indxs), 
+            data_type="eff", robot_idxs = env_indxs, gpu=self._use_gpu) 
+
+        # height map
+        if self._enable_height_shared:
+            hdata = self._height_imgs[robot_name]
+            if env_indxs is not None:
+                hdata = hdata[env_indxs]
+            flat = hdata.reshape(hdata.shape[0], -1)
+            rhc_state.height_sensor.set(data=flat, data_type=None, robot_idxs=env_indxs, gpu=self._use_gpu)
+
+        # Updating contact state for selected contact links
+        self._update_contact_state(robot_name=robot_name, env_indxs=env_indxs)
+    
+    def _update_contact_state(self, 
+            robot_name: str, 
+            env_indxs: torch.Tensor = None):
+        for i in range(0, self.cluster_servers[robot_name].n_contact_sensors()):
+            contact_link = self.cluster_servers[robot_name].contact_linknames()[i]
+            f_contact = self._get_contact_f(robot_name=robot_name,
+                contact_link=contact_link,
+                env_indxs=env_indxs)
+            if f_contact is not None:
+                self.cluster_servers[robot_name].get_state().contact_wrenches.set(data=f_contact, data_type="f",
+                                contact_name=contact_link, 
+                                robot_idxs = env_indxs, 
+                                gpu=self._use_gpu)
+                    
+    def _init_safe_cluster_actions(self,
+                            robot_name: str):
+
+        # this does not actually write on shared memory, 
+        # but it's enough to get safe actions for the simulator before the 
+        # cluster starts to receive data from the controllers
+        control_cluster = self.cluster_servers[robot_name]
+        rhc_cmds = control_cluster.get_actions()
+        n_jnts = rhc_cmds.n_jnts()
+        
+        null_action = torch.zeros((self._num_envs, n_jnts), 
+                        dtype=self._dtype,
+                        device=self._device)
+        rhc_cmds.jnts_state.set(data=self._homing, data_type="q", gpu=self._use_gpu)
+        rhc_cmds.jnts_state.set(data=null_action, data_type="v", gpu=self._use_gpu)
+        rhc_cmds.jnts_state.set(data=null_action, data_type="eff", gpu=self._use_gpu)
+
+    def _pre_step_db(self) -> None:
+        
+        # cluster step logic here
+        for i in range(len(self._robot_names)):
+            robot_name = self._robot_names[i]
+
+            if self._override_low_lev_controller:
+                # if overriding low-lev jnt imp. this has to run at the highest
+                # freq possible
+                start=time.perf_counter()
+                self._read_jnts_state_from_robot(robot_name=robot_name)
+                self.debug_data["time_to_get_states_from_env"]= time.perf_counter()-start
+                
+                self._write_state_to_jnt_imp(robot_name=robot_name)
+                self._apply_cmds_to_jnt_imp_control(robot_name=robot_name)
+
+            if self._jnt_vel_filter[robot_name] is not None and \
+                (self.cluster_sim_step_counters[robot_name]+1) % self._filter_step_ssteps_freq == 0:
+                # filter joint vel at a fixed frequency wrt sim steps
+                if not self._override_low_lev_controller:
+                    # we need a fresh sensor reading
+                    self._read_jnts_state_from_robot(robot_name=robot_name)
+                self._step_jnt_vel_filter(robot_name=robot_name, env_indxs=None)
+
+            control_cluster = self.cluster_servers[robot_name]
+            if control_cluster.is_cluster_instant(self.cluster_sim_step_counters[robot_name]):
+                wait_ok=control_cluster.wait_for_solution() # this is blocking
+                if not wait_ok:
+                    return False
+                failed = control_cluster.get_failed_controllers(gpu=self._use_gpu)
+                self._set_cluster_actions(robot_name=robot_name) # write last cmds to low level control
+                if not self._override_low_lev_controller:
+                    self._apply_cmds_to_jnt_imp_control(robot_name=robot_name) # apply to robot
+                    # we can update the jnt state just at the rate at which the cluster needs it
+                    start=time.perf_counter()
+                    self._read_jnts_state_from_robot(robot_name=robot_name, env_indxs=None)
+                else:
+                    # read state necessary for cluster
+                    start=time.perf_counter()
+                self._read_root_state_from_robot(robot_name=robot_name, 
+                    env_indxs=None)
+                self.debug_data["time_to_get_states_from_env"]= time.perf_counter()-start
+                start=time.perf_counter()
+                self._set_state_to_cluster(robot_name=robot_name, 
+                    env_indxs=None)
+                control_cluster.write_robot_state()
+                self.debug_data["cluster_state_update_dt"][robot_name] = time.perf_counter()-start
+
+                self._update_jnt_imp_cntrl_shared_data() # only if debug_mode_jnt_imp is enabled
+
+                if self._use_remote_stepping[i]:
+                    self._remote_steppers[robot_name].ack() # signal cluster stepping is finished
+                    if failed is not None and self._env_opts["deact_when_failure"]: # deactivate robot completely 
+                        self._deactivate(env_indxs=failed,
+                            robot_name=robot_name)
+                    wait_reset_ok=self._process_remote_reset_req(robot_name=robot_name) # wait for remote reset request (blocking)
+                    wait_step_ok=self._wait_for_remote_step_req(robot_name=robot_name)
+                    if not wait_reset_ok or not wait_step_ok:   
+                        return False
+                else:
+                    if failed is not None:
+                        reset_ok=self._reset(env_indxs=failed,
+                            robot_name=robot_name,
+                            reset_cluster=True,
+                            reset_cluster_counter=False,
+                            randomize=True)
+                        if not reset_ok:
+                            return False
+                        self._set_startup_jnt_imp_gains(robot_name=robot_name,
+                            env_indxs=failed)
+
+                    control_cluster.activate_controllers(idxs=control_cluster.get_inactive_controllers())
+
+                control_cluster.pre_trigger() # performs pre-trigger steps, like retrieving
+                # values of some rhc flags on shared memory
+                control_cluster.trigger_solution() # trigger only active controllers
+
+        return True
+
+    def _pre_step(self) -> None:
+        
+        # cluster step logic here
+        for i in range(len(self._robot_names)):
+            robot_name = self._robot_names[i]
+            if self._override_low_lev_controller:
+                # if overriding low-lev jnt imp. this has to run at the highest
+                # freq possible
+                self._read_jnts_state_from_robot(robot_name=robot_name)
+                self._write_state_to_jnt_imp(robot_name=robot_name)
+                self._apply_cmds_to_jnt_imp_control(robot_name=robot_name)
+
+            if self._jnt_vel_filter[robot_name] is not None and \
+                (self.cluster_sim_step_counters[robot_name]+1) % self._filter_step_ssteps_freq == 0:
+                # filter joint vel at a fixed frequency wrt sim steps
+                if not self._override_low_lev_controller:
+                    # we need a fresh sensor reading
+                    self._read_jnts_state_from_robot(robot_name=robot_name)
+                self._step_jnt_vel_filter(robot_name=robot_name, env_indxs=None)
+
+            control_cluster = self.cluster_servers[robot_name]
+            if control_cluster.is_cluster_instant(self.cluster_sim_step_counters[robot_name]):
+                wait_ok=control_cluster.wait_for_solution() # this is blocking
+                if not wait_ok:
+                    return False
+                failed = control_cluster.get_failed_controllers(gpu=self._use_gpu)
+                self._set_cluster_actions(robot_name=robot_name) # set last cmds to low level control
+                if not self._override_low_lev_controller:
+                    self._apply_cmds_to_jnt_imp_control(robot_name=robot_name) # apply to robot
+                    # we can update the jnt state just at the rate at which the cluster needs it
+                    self._read_jnts_state_from_robot(robot_name=robot_name, env_indxs=None)
+                # read state necessary for cluster
+                self._read_root_state_from_robot(robot_name=robot_name, 
+                    env_indxs=None)
+                # write last robot state to the cluster of controllers
+                self._set_state_to_cluster(robot_name=robot_name, 
+                    env_indxs=None)
+                control_cluster.write_robot_state() # write on shared mem
+
+                if self._use_remote_stepping[i]:
+                    self._remote_steppers[robot_name].ack() # signal cluster stepping is finished
+                    if failed is not None and self._env_opts["deact_when_failure"]:
+                        self._deactivate(env_indxs=failed,
+                            robot_name=robot_name)
+                    wait_reset_ok=self._process_remote_reset_req(robot_name=robot_name) # wait for remote reset request (blocking)
+                    wait_step_ok=self._wait_for_remote_step_req(robot_name=robot_name)
+                    if not wait_reset_ok or not wait_step_ok:   
+                        return False
+                else:
+                    if failed is not None:
+                        reset_ok=self._reset(env_indxs=failed,
+                            robot_name=robot_name,
+                            reset_cluster=True,
+                            reset_cluster_counter=False,
+                            randomize=True)
+                        if not reset_ok:
+                            return False
+                        self._set_startup_jnt_imp_gains(robot_name=robot_name,
+                            env_indxs=failed)
+                    control_cluster.activate_controllers(idxs=control_cluster.get_inactive_controllers())
+
+                control_cluster.pre_trigger() # performs pre-trigger steps, like retrieving
+                # values of some rhc flags on shared memory
+                control_cluster.trigger_solution() # trigger only active controllers
+        
+        return True
+    
+    def _post_world_step_db(self) -> bool:
+
+        for i in range(len(self._robot_names)):
+            robot_name = self._robot_names[i]
+            control_cluster = self.cluster_servers[robot_name]
+            self.cluster_sim_step_counters[robot_name]+=1 # this has to be update with sim freq
+            if self._debug:
+                self.debug_data["sim_time"][robot_name]=self.world_time(robot_name=robot_name)
+                self.debug_data["cluster_sol_time"][robot_name] = \
+                    control_cluster.solution_time()
+                
+        self.step_counter +=1
+    
+    def _post_world_step(self) -> bool:
+
+        for i in range(len(self._robot_names)):
+            robot_name = self._robot_names[i]
+            self.cluster_sim_step_counters[robot_name]+=1
+        self.step_counter +=1
+                    
+    def _reset(self,
+            robot_name: str,
+            env_indxs: torch.Tensor = None,
+            randomize: bool = False,
+            reset_cluster: bool = False,
+            reset_cluster_counter = False,
+            acquire_offsets: bool = False):
+        
+        # resets the state of target robot and env to the defaults
+        self._reset_state(env_indxs=env_indxs, 
+            robot_name=robot_name,
+            randomize=randomize)
+        
+        # and jnt imp. controllers
+        self._reset_jnt_imp_control(robot_name=robot_name,
+                env_indxs=env_indxs)
+        
+        # read reset state
+        self._read_root_state_from_robot(robot_name=robot_name,
+                env_indxs=env_indxs)
+        self._read_jnts_state_from_robot(robot_name=robot_name,
+            env_indxs=env_indxs)
+
+        if self._jnt_vel_filter[robot_name] is not None:
+            self._jnt_vel_filter[robot_name].reset(idxs=env_indxs)
+
+        if acquire_offsets:
+            self._update_root_offsets(robot_name=robot_name,
+                    env_indxs=env_indxs)
+            
+        if reset_cluster: # reset controllers remotely
+            reset_ok=self._reset_cluster(env_indxs=env_indxs,
+                robot_name=robot_name,
+                reset_cluster_counter=reset_cluster_counter)
+            if not reset_ok:
+                return False
+            
+        return True
+        
+    def _randomize_yaw(self,
+            robot_name: str,
+            env_indxs: torch.Tensor = None):
+
+        root_q_default = self._root_q_default[robot_name]
+        if env_indxs is None:
+            env_indxs = torch.arange(root_q_default.shape[0])
+
+        num_indices = env_indxs.shape[0]
+        yaw_angles = torch.rand((num_indices,), 
+                        device=root_q_default.device) * 2 * torch.pi  # uniformly distributed random angles
+        
+        # Compute cos and sin once
+        cos_half = torch.cos(yaw_angles / 2)
+        root_q_default[env_indxs, :] = torch.stack((cos_half, 
+                                torch.zeros_like(cos_half),
+                                torch.zeros_like(cos_half), 
+                                torch.sin(yaw_angles / 2)), dim=1).reshape(num_indices, 4)
+
+    def _deactivate(self,
+        robot_name: str,
+        env_indxs: torch.Tensor = None):
+        
+        # deactivate jnt imp controllers for given robots and envs (makes the robot fall)
+        self._jnt_imp_controllers[robot_name].deactivate(robot_indxs=env_indxs)
+    
+    def _n_contacts(self, robot_name: str) -> List[int]:
+        return self._num_contacts[robot_name]
+    
+    def root_p(self,
+            robot_name: str,
+            env_idxs: torch.Tensor = None):
+
+        if env_idxs is None:
+            return self._root_p[robot_name]
+        else:
+            return self._root_p[robot_name][env_idxs, :]
+
+    def root_p_rel(self,
+            robot_name: str,
+            env_idxs: torch.Tensor = None):
+
+        if env_idxs is None:
+            rel_pos = torch.sub(self.root_p(robot_name=robot_name),
+                self._root_pos_offsets[robot_name])
+        else:
+            rel_pos = torch.sub(self.root_p(robot_name=robot_name,
+                                                env_idxs=env_idxs), 
+                    self._root_pos_offsets[robot_name][env_idxs, :])
+        return rel_pos
+    
+    def root_q(self,
+            robot_name: str,
+            env_idxs: torch.Tensor = None):
+
+        if env_idxs is None:
+            return self._root_q[robot_name]
+        else:
+            return self._root_q[robot_name][env_idxs, :]
+
+    def root_q_rel(self,
+            robot_name: str,
+            env_idxs: torch.Tensor = None):
+
+        if env_idxs is None:
+            return quaternion_difference(self._root_q_offsets[robot_name], 
+                                self.root_q(robot_name=robot_name))
+        rel_q = quaternion_difference(self._root_q_offsets[robot_name][env_idxs, :], 
+                            self.root_q(robot_name=robot_name,
+                                            env_idxs=env_idxs))
+        return rel_q
+
+    def _quat_to_yaw_wxyz(self, q: torch.Tensor, out: torch.Tensor = None):
+        # Quaternion convention is w, x, y, z.
+        w, x, y, z = q[:, 0], q[:, 1], q[:, 2], q[:, 3]
+        num = 2.0 * (w * z + x * y)
+        den = 1.0 - 2.0 * (y * y + z * z)
+        if out is None:
+            return torch.atan2(num, den)
+        return torch.atan2(num, den, out=out)
+
+    def _yaw_to_quat_wxyz(self, yaw: torch.Tensor, like_q: torch.Tensor,
+            out: torch.Tensor = None):
+        q = out
+        if q is None:
+            q = torch.zeros((yaw.shape[0], 4), dtype=like_q.dtype, device=like_q.device)
+        else:
+            q.zero_()
+        q[:, 0] = torch.cos(yaw / 2.0)
+        q[:, 3] = torch.sin(yaw / 2.0)
+        return q
+
+    def _quat_conjugate_wxyz(self, q: torch.Tensor, out: torch.Tensor = None):
+        qi = out
+        if qi is None:
+            qi = torch.empty_like(q)
+        qi[:, :] = q
+        qi[:, 1:] = -qi[:, 1:]
+        return qi
+
+    def _quat_multiply_wxyz(self, q1: torch.Tensor, q2: torch.Tensor,
+            out: torch.Tensor = None):
+        q_out = out
+        if q_out is None:
+            q_out = torch.empty_like(q1)
+        w1, x1, y1, z1 = q1[:, 0], q1[:, 1], q1[:, 2], q1[:, 3]
+        w2, x2, y2, z2 = q2[:, 0], q2[:, 1], q2[:, 2], q2[:, 3]
+        q_out[:, 0] = w1 * w2 - x1 * x2 - y1 * y2 - z1 * z2
+        q_out[:, 1] = w1 * x2 + x1 * w2 + y1 * z2 - z1 * y2
+        q_out[:, 2] = w1 * y2 - x1 * z2 + y1 * w2 + z1 * x2
+        q_out[:, 3] = w1 * z2 + x1 * y2 - y1 * x2 + z1 * w2
+        return q_out
+
+    def _normalize_quat_wxyz(self, q: torch.Tensor, out: torch.Tensor = None):
+        q_norm = out
+        if q_norm is None:
+            q_norm = torch.empty_like(q)
+        q_norm[:, :] = q
+        q_norm /= torch.clamp(torch.norm(q_norm, dim=1, keepdim=True), min=1e-9)
+        return q_norm
+
+    def root_q_yaw_rel(self,
+            robot_name: str,
+            env_idxs: torch.Tensor = None):
+        
+        # Return quaternion with startup yaw removed while preserving current pitch/roll.
+        if env_idxs is None:
+            ws = self._root_q_yaw_rel_ws[robot_name]
+            q_abs = self._root_q[robot_name]
+            yaw_start = self._root_q_offsets_yaw[robot_name]
+
+            self._normalize_quat_wxyz(q=q_abs, out=ws["q_abs_unit"])
+            self._quat_to_yaw_wxyz(q=ws["q_abs_unit"], out=ws["yaw_abs"])
+
+            torch.sub(ws["yaw_abs"], yaw_start, out=ws["yaw_rel"])
+            torch.sin(ws["yaw_rel"], out=ws["yaw_sin"])
+            torch.cos(ws["yaw_rel"], out=ws["yaw_cos"])
+            torch.atan2(ws["yaw_sin"], ws["yaw_cos"], out=ws["yaw_rel"])
+
+            # Build pure-yaw quaternions for:
+            # 1) the current absolute heading and 2) the startup-relative heading.
+            self._yaw_to_quat_wxyz(yaw=ws["yaw_abs"], like_q=ws["q_abs_unit"], out=ws["q_yaw_abs"])
+            self._yaw_to_quat_wxyz(yaw=ws["yaw_rel"], like_q=ws["q_abs_unit"], out=ws["q_yaw_rel"])
+
+            # Isolate pitch/roll by removing the absolute yaw from the current orientation.
+            # For unit quaternions q_pr = q_yaw_abs^{-1} * q_abs.
+            self._quat_conjugate_wxyz(q=ws["q_yaw_abs"], out=ws["q_yaw_abs_conj"])
+            self._quat_multiply_wxyz(q1=ws["q_yaw_abs_conj"], q2=ws["q_abs_unit"], out=ws["q_pr"])
+            # Recompose orientation with relative yaw + current pitch/roll.
+            self._quat_multiply_wxyz(q1=ws["q_yaw_rel"], q2=ws["q_pr"], out=ws["q_rel"])
+
+            return self._normalize_quat_wxyz(q=ws["q_rel"], out=ws["q_rel"])
+
+        q_abs = self.root_q(robot_name=robot_name, env_idxs=env_idxs)
+        q_abs = self._normalize_quat_wxyz(q=q_abs, out=q_abs)
+
+        yaw_abs = self._quat_to_yaw_wxyz(q_abs)
+        yaw_start = self._root_q_offsets_yaw[robot_name][env_idxs]
+        yaw_rel = yaw_abs - yaw_start
+        yaw_rel = torch.atan2(torch.sin(yaw_rel), torch.cos(yaw_rel))
+
+        q_yaw_abs = self._yaw_to_quat_wxyz(yaw_abs, like_q=q_abs)
+        q_yaw_rel = self._yaw_to_quat_wxyz(yaw_rel, like_q=q_abs)
+        q_pr = self._quat_multiply_wxyz(self._quat_conjugate_wxyz(q_yaw_abs), q_abs)
+        q_rel = self._quat_multiply_wxyz(q_yaw_rel, q_pr)
+
+        return self._normalize_quat_wxyz(q_rel)
+    
+    def root_v(self,
+            robot_name: str,
+            env_idxs: torch.Tensor = None,
+            base_loc: bool = True):
+
+        root_v=self._root_v[robot_name]
+        if base_loc:
+            root_v=self._root_v_base_loc[robot_name]
+        if env_idxs is None:
+            return root_v
+        else:
+            return root_v[env_idxs, :]
+    
+    def root_omega(self,
+            robot_name: str,
+            env_idxs: torch.Tensor = None,
+            base_loc: bool = True):
+
+        root_omega=self._root_omega[robot_name]
+        if base_loc:
+            root_omega=self._root_omega_base_loc[robot_name]
+        if env_idxs is None:
+            return root_omega
+        else:
+            return root_omega[env_idxs, :]
+    
+    def root_a(self,
+            robot_name: str,
+            env_idxs: torch.Tensor = None,
+            base_loc: bool = True):
+
+        root_a=self._root_a[robot_name]
+        if base_loc:
+            root_a=self._root_a_base_loc[robot_name]
+        if env_idxs is None:
+            return root_a
+        else:
+            return root_a[env_idxs, :]
+    
+    def root_alpha(self,
+            robot_name: str,
+            env_idxs: torch.Tensor = None,
+            base_loc: bool = True):
+
+        root_alpha=self._root_alpha[robot_name]
+        if base_loc:
+            root_alpha=self._root_alpha_base_loc[robot_name]
+        if env_idxs is None:
+            return root_alpha
+        else:
+            return root_alpha[env_idxs, :]
+        
+    def gravity(self,
+            robot_name: str,
+            env_idxs: torch.Tensor = None,
+            base_loc: bool = True):
+
+        gravity_loc=self._gravity_normalized[robot_name]
+        if base_loc:
+            gravity_loc=self._gravity_normalized_base_loc[robot_name]
+        if env_idxs is None:
+            return gravity_loc
+        else:
+            return gravity_loc[env_idxs, :]
+    
+    def jnts_q(self,
+            robot_name: str,
+            env_idxs: torch.Tensor = None):
+        
+        if env_idxs is None:
+            return self._jnts_q[robot_name]
+        else:
+            return self._jnts_q[robot_name][env_idxs, :]
+
+    def jnts_v(self,
+            robot_name: str,
+            env_idxs: torch.Tensor = None):
+
+        if env_idxs is None:
+            return self._jnts_v[robot_name]
+        else:
+            return self._jnts_v[robot_name][env_idxs, :]
+
+    def jnts_eff(self,
+            robot_name: str,
+            env_idxs: torch.Tensor = None): # (measured) efforts
+
+        if env_idxs is None:
+            return self._jnts_eff[robot_name]
+        else:
+            return self._jnts_eff[robot_name][env_idxs, :]
+
+    def _wait_for_remote_step_req(self,
+            robot_name: str):
+        if not self._remote_steppers[robot_name].wait(self._timeout):
+            Journal.log(self.__class__.__name__,
+                "_wait_for_remote_step_req",
+                "Didn't receive any remote step req within timeout!",
+                LogType.EXCEP,
+                throw_when_excep = False)
+            return False
+        return True
+    
+    def _process_remote_reset_req(self,
+            robot_name: str):
+        
+        if not self._remote_resetters[robot_name].wait(self._timeout):
+            Journal.log(self.__class__.__name__,
+                "_process_remote_reset_req",
+                "Didn't receive any remote reset req within timeout!",
+                LogType.EXCEP,
+                throw_when_excep = False)
+            return False
+            
+        reset_requests = self._remote_reset_requests[robot_name]
+        reset_requests.synch_all(read=True, retry=True) # read reset requests from shared mem
+        to_be_reset = reset_requests.to_be_reset(gpu=self._use_gpu)
+        if to_be_reset is not None:
+            reset_ok=self._reset(env_indxs=to_be_reset,
+                robot_name=robot_name,
+                reset_cluster=True,
+                reset_cluster_counter=False,
+                randomize=True)
+            if not reset_ok:
+                return False
+            self._set_startup_jnt_imp_gains(robot_name=robot_name,
+                env_indxs=to_be_reset) # set gains to startup config (usually lower gains)
+        
+        control_cluster = self.cluster_servers[robot_name]
+        control_cluster.activate_controllers(idxs=to_be_reset) # activate controllers
+        # (necessary if failed)
+
+        self._remote_resetters[robot_name].ack() # signal reset performed
+
+        return True
+    
+    def _update_jnt_imp_cntrl_shared_data(self):
+        if self._debug:
+            for i in range(0, len(self._robot_names)):
+                robot_name = self._robot_names[i]
+                # updating all the jnt impedance data - > this may introduce some overhead
+                imp_data = self._jnt_imp_cntrl_shared_data[robot_name].imp_data_view
+                # set data
+                imp_data.set(data_type="pos_err",
+                        data=self._jnt_imp_controllers[robot_name].pos_err(),
+                        gpu=self._use_gpu)
+                imp_data.set(data_type="vel_err",
+                        data=self._jnt_imp_controllers[robot_name].vel_err(),
+                        gpu=self._use_gpu)
+                imp_data.set(data_type="pos_gains",
+                        data=self._jnt_imp_controllers[robot_name].pos_gains(),
+                        gpu=self._use_gpu)
+                imp_data.set(data_type="vel_gains",
+                        data=self._jnt_imp_controllers[robot_name].vel_gains(),
+                        gpu=self._use_gpu)
+                imp_data.set(data_type="eff_ff",
+                        data=self._jnt_imp_controllers[robot_name].eff_ref(),
+                        gpu=self._use_gpu)
+                imp_data.set(data_type="pos",
+                        data=self._jnt_imp_controllers[robot_name].pos(),
+                        gpu=self._use_gpu)
+                imp_data.set(data_type="pos_ref",
+                        data=self._jnt_imp_controllers[robot_name].pos_ref(),
+                        gpu=self._use_gpu)
+                imp_data.set(data_type="vel",
+                        data=self._jnt_imp_controllers[robot_name].vel(),
+                        gpu=self._use_gpu)
+                imp_data.set(data_type="vel_ref",
+                        data=self._jnt_imp_controllers[robot_name].vel_ref(),
+                        gpu=self._use_gpu)
+                imp_data.set(data_type="eff",
+                        data=self._jnt_imp_controllers[robot_name].eff(),
+                        gpu=self._use_gpu)
+                imp_data.set(data_type="imp_eff",
+                        data=self._jnt_imp_controllers[robot_name].imp_eff(),
+                        gpu=self._use_gpu)
+                # copy from GPU to CPU if using gpu
+                if self._use_gpu:
+                    imp_data.synch_mirror(from_gpu=True,non_blocking=True)
+                    # even if it's from GPU->CPu we can use non-blocking since it's just for db 
+                    # purposes
+                # write copies to shared memory
+                imp_data.synch_all(read=False, retry=False)
+
+    def _set_startup_jnt_imp_gains(self,
+            robot_name:str, 
+            env_indxs: torch.Tensor = None):
+        
+        startup_p_gains=self._jnt_imp_controllers[robot_name].startup_p_gains()
+        startup_d_gains=self._jnt_imp_controllers[robot_name].startup_d_gains()
+
+        if env_indxs is not None:
+            self._jnt_imp_controllers[robot_name].set_gains(robot_indxs=env_indxs,
+                pos_gains=startup_p_gains[env_indxs, :], 
+                vel_gains=startup_d_gains[env_indxs, :])
+        else:
+            self._jnt_imp_controllers[robot_name].set_gains(robot_indxs=env_indxs,
+                pos_gains=startup_p_gains[:, :], 
+                vel_gains=startup_d_gains[:, :])
+            
+    def _write_state_to_jnt_imp(self,
+        robot_name: str):
+     
+        # always update ,imp. controller internal state (jnt imp control is supposed to be
+        # always running)
+        self._jnt_imp_controllers[robot_name].update_state(pos=self.jnts_q(robot_name=robot_name), 
+            vel = self.jnts_v(robot_name=robot_name),
+            eff = self.jnts_eff(robot_name=robot_name))
+
+    def _set_cluster_actions(self, 
+        robot_name):
+        control_cluster = self.cluster_servers[robot_name]
+        actions=control_cluster.get_actions()
+        active_controllers=control_cluster.get_active_controllers(gpu=self._use_gpu)
+        
+        if active_controllers is not None:
+            self._jnt_imp_controllers[robot_name].set_refs(
+                pos_ref=actions.jnts_state.get(data_type="q", gpu=self._use_gpu)[active_controllers, :], 
+                vel_ref=actions.jnts_state.get(data_type="v", gpu=self._use_gpu)[active_controllers, :], 
+                eff_ref=actions.jnts_state.get(data_type="eff", gpu=self._use_gpu)[active_controllers, :],
+                robot_indxs=active_controllers)            
+    
+    def _jnt_imp_reset_overrride(self, robot_name:str):
+        # to be overriden
+        pass
+
+    def _apply_cmds_to_jnt_imp_control(self, robot_name:str):
+
+        self._jnt_imp_controllers[robot_name].apply_cmds()
+
+    def _update_root_offsets(self, 
+                    robot_name: str,
+                    env_indxs: torch.Tensor = None):
+        
+        if self._debug:
+            for_robots = ""
+            if env_indxs is not None:
+                if not isinstance(env_indxs, torch.Tensor):                
+                    msg = "Provided env_indxs should be a torch tensor of indexes!"
+                    Journal.log(self.__class__.__name__,
+                        "update_root_offsets",
+                        msg,
+                        LogType.EXCEP,
+                        throw_when_excep = True)    
+                if self._use_gpu:
+                    if not env_indxs.device.type == "cuda":
+                            error = "Provided env_indxs should be on GPU!"
+                            Journal.log(self.__class__.__name__,
+                            "_step_jnt_imp_control",
+                            error,
+                            LogType.EXCEP,
+                            True)
+                else:
+                    if not env_indxs.device.type == "cpu":
+                        error = "Provided env_indxs should be on CPU!"
+                        Journal.log(self.__class__.__name__,
+                            "_step_jnt_imp_control",
+                            error,
+                            LogType.EXCEP,
+                            True)
+                for_robots = f"for robot {robot_name}, indexes: " + str(env_indxs.tolist())
+            if self._verbose:
+                Journal.log(self.__class__.__name__,
+                    "update_root_offsets",
+                    f"updating root offsets " + for_robots,
+                    LogType.STAT,
+                    throw_when_excep = True)
+
+        # only planar position used
+        if env_indxs is None:
+            self._root_pos_offsets[robot_name][:, 0:2]  = self._root_p[robot_name][:, 0:2]
+            self._normalize_quat_wxyz(q=self._root_q[robot_name], out=self._root_q_offsets[robot_name])
+            self._quat_to_yaw_wxyz(q=self._root_q_offsets[robot_name],
+                out=self._root_q_offsets_yaw[robot_name])
+            
+        else:
+            self._root_pos_offsets[robot_name][env_indxs, 0:2]  = self._root_p[robot_name][env_indxs, 0:2]
+            q_root_norm=self._normalize_quat_wxyz(self._root_q[robot_name][env_indxs, :])
+            self._root_q_offsets[robot_name][env_indxs, :]  = q_root_norm
+            self._root_q_offsets_yaw[robot_name][env_indxs] = self._quat_to_yaw_wxyz(q=q_root_norm)
+
+    def _reset_jnt_imp_control(self, 
+        robot_name: str,
+        env_indxs: torch.Tensor = None):
+        
+        if self._debug:
+            for_robots = ""
+            if env_indxs is not None:
+                if not isinstance(env_indxs, torch.Tensor):
+                    Journal.log(self.__class__.__name__,
+                        "reset_jnt_imp_control",
+                        "Provided env_indxs should be a torch tensor of indexes!",
+                        LogType.EXCEP,
+                        throw_when_excep = True)
+                if self._use_gpu:
+                    if not env_indxs.device.type == "cuda":
+                            error = "Provided env_indxs should be on GPU!"
+                            Journal.log(self.__class__.__name__,
+                            "_step_jnt_imp_control",
+                            error,
+                            LogType.EXCEP,
+                            True)
+                else:
+                    if not env_indxs.device.type == "cpu":
+                        error = "Provided env_indxs should be on CPU!"
+                        Journal.log(self.__class__.__name__,
+                            "_step_jnt_imp_control",
+                            error,
+                            LogType.EXCEP,
+                            True)
+                for_robots = f"for robot {robot_name}, indexes: " + str(env_indxs)
+                                
+            if self._verbose:
+                Journal.log(self.__class__.__name__,
+                    "reset_jnt_imp_control",
+                    f"resetting joint impedances " + for_robots,
+                    LogType.STAT,
+                    throw_when_excep = True)
+
+        # resets all internal data, refs to defaults
+        self._jnt_imp_controllers[robot_name].reset(robot_indxs=env_indxs)
+        
+        #restore jnt imp refs to homing            
+        if env_indxs is None:                               
+            self._jnt_imp_controllers[robot_name].set_refs(pos_ref=self._homing[:, :],
+                robot_indxs = None)
+        else:
+            self._jnt_imp_controllers[robot_name].set_refs(pos_ref=self._homing[env_indxs, :],
+                robot_indxs = env_indxs)
+
+        # self._write_state_to_jnt_imp(robot_name=robot_name)
+        # actually applies reset commands to the articulation
+        self._write_state_to_jnt_imp(robot_name=robot_name)
+        self._jnt_imp_reset_overrride(robot_name=robot_name)
+        self._apply_cmds_to_jnt_imp_control(robot_name=robot_name)
+
+    def _synch_default_root_states(self,
+            robot_name: str,
+            env_indxs: torch.Tensor = None):
+
+        if self._debug:
+            for_robots = ""
+            if env_indxs is not None:
+                if not isinstance(env_indxs, torch.Tensor):
+                    msg = "Provided env_indxs should be a torch tensor of indexes!"
+                    Journal.log(self.__class__.__name__,
+                        "synch_default_root_states",
+                        msg,
+                        LogType.EXCEP,
+                        throw_when_excep = True)  
+                if self._use_gpu:
+                    if not env_indxs.device.type == "cuda":
+                            error = "Provided env_indxs should be on GPU!"
+                            Journal.log(self.__class__.__name__,
+                            "_step_jnt_imp_control",
+                            error,
+                            LogType.EXCEP,
+                            True)
+                else:
+                    if not env_indxs.device.type == "cpu":
+                        error = "Provided env_indxs should be on CPU!"
+                        Journal.log(self.__class__.__name__,
+                            "_step_jnt_imp_control",
+                            error,
+                            LogType.EXCEP,
+                            True)  
+                for_robots = f"for robot {robot_name}, indexes: " + str(env_indxs.tolist())
+            if self._verbose:
+                Journal.log(self.__class__.__name__,
+                            "synch_default_root_states",
+                            f"updating default root states " + for_robots,
+                            LogType.STAT,
+                            throw_when_excep = True)
+
+        if env_indxs is None:
+            self._root_p_default[robot_name][:, :] = self._root_p[robot_name]
+            self._root_q_default[robot_name][:, :] = self._root_q[robot_name]
+        else:
+            self._root_p_default[robot_name][env_indxs, :] = self._root_p[robot_name][env_indxs, :]
+            self._root_q_default[robot_name][env_indxs, :] = self._root_q[robot_name][env_indxs, :]
+
+    def _generate_rob_descriptions(self, 
+                    robot_name: str, 
+                    urdf_path: str,
+                    srdf_path: str):
+        
+        custom_xacro_args=extract_custom_xacro_args(self._env_opts)
+        Journal.log(self.__class__.__name__,
+                    "_generate_rob_descriptions",
+                    "generating URDF for robot "+ f"{robot_name}, from URDF {urdf_path}...",
+                    LogType.STAT,
+                    throw_when_excep = True)
+        xrdf_cmds=self._xrdf_cmds(robot_name=robot_name)
+        xrdf_cmds=merge_xacro_cmds(prev_cmds=xrdf_cmds,
+            new_cmds=custom_xacro_args)
+        self._urdf_dump_paths[robot_name]=generate_urdf(robot_name=robot_name, 
+            xacro_path=urdf_path,
+            dump_path=self._descr_dump_path,
+            xrdf_cmds=xrdf_cmds)
+        Journal.log(self.__class__.__name__,
+                    "_generate_rob_descriptions",
+                    "generating SRDF for robot "+ f"{robot_name}, from SRDF {srdf_path}...",
+                    LogType.STAT,
+                    throw_when_excep = True)
+        # we also generate SRDF files, which are useful for control
+        self._srdf_dump_paths[robot_name]=generate_srdf(robot_name=robot_name, 
+            xacro_path=srdf_path,
+            dump_path=self._descr_dump_path,
+            xrdf_cmds=xrdf_cmds)
+    
+    def _xrdf_cmds(self, robot_name:str):
+        urdfpath=self._robot_urdf_paths[robot_name]
+        # we assume directory tree of the robot package is like
+        # robot-ros-pkg/robot_urdf/urdf/robot.urdf.xacro
+        parts = urdfpath.split('/')
+        urdf_descr_root_path = '/'.join(parts[:-2])
+        cmds = get_xrdf_cmds(urdf_descr_root_path=urdf_descr_root_path) 
+        return cmds
+    
+    @abstractmethod
+    def current_tstep(self) -> int:
+        pass
+    
+    @abstractmethod
+    def world_time(self, robot_name: str) -> float:
+        return self.cluster_sim_step_counters[robot_name]*self.physics_dt()
+    
+    @abstractmethod
+    def is_running(self) -> bool:
+        pass
+    
+    @abstractmethod
+    def _get_contact_f(self, 
+        robot_name: str, 
+        contact_link: str,
+        env_indxs: torch.Tensor) -> torch.Tensor:
+        return None
+    
+    @abstractmethod
+    def physics_dt(self) -> float:
+        pass
+    
+    @abstractmethod
+    def rendering_dt(self) -> float:
+        pass
+    
+    @abstractmethod
+    def set_physics_dt(self, physics_dt:float):
+        pass
+
+    @abstractmethod
+    def set_rendering_dt(self, rendering_dt:float):
+        pass
+
+    @abstractmethod
+    def _robot_jnt_names(self, robot_name: str) -> List[str]:
+        pass
+    
+    @abstractmethod
+    def _read_root_state_from_robot(self,
+        robot_name: str,
+        env_indxs: torch.Tensor = None):
+        # IMPORTANT: Child interfaces should provide root quaternions in w, x, y, z convention.
+        pass
+    
+    @abstractmethod
+    def _read_jnts_state_from_robot(self,
+        robot_name: str,
+        env_indxs: torch.Tensor = None):
+        pass
+
+    @abstractmethod
+    def _init_robots_state(self):
+        pass
+
+    @abstractmethod
+    def _reset_state(self,
+            robot_name: str,
+            env_indxs: torch.Tensor = None,
+            randomize: bool = False):
+        pass
+    
+    @abstractmethod
+    def _init_world(self):
+        pass
+
+    @abstractmethod
+    def _reset_sim(self) -> None:
+        pass
+    
+    @abstractmethod
+    def _set_jnts_to_homing(self, robot_name: str):
+        pass
+
+    @abstractmethod
+    def _set_root_to_defconfig(self, robot_name: str):
+        pass
+
+    @abstractmethod
+    def _parse_env_opts(self):
+        pass
+    
+    @abstractmethod
+    def _pre_setup(self):
+        pass
+
+    @abstractmethod
+    def _generate_jnt_imp_control(self) -> JntImpCntrlChild:
+        pass
+
+    @abstractmethod
+    def _render_sim(self, mode:str="human") -> None:
+        pass
+
+    @abstractmethod
+    def _close(self) -> None:
+        pass
+
+    @abstractmethod
+    def _step_world(self) -> None:
+        pass

xbot2_basic.yaml

@@ -0,0 +1,86 @@
+XBotInterface:
+  urdf_path: $PWD/centauro.urdf
+  srdf_path: $PWD/centauro_old.srdf
+  # urdf_path: $PWD/centauro_dagana_right.urdf
+  # srdf_path: $PWD/centauro_dagana_right.srdf
+
+ModelInterface:
+  model_type: RBDL
+  is_model_floating_base: true
+
+motor_pd:
+  "j_arm*_1": [500, 10]
+  "j_arm*_2": [500, 10]
+  "j_arm*_3": [500, 10]
+  "j_arm*_4": [500, 10]
+  "j_arm*_5": [100, 5]
+  "j_arm*_6": [100, 5]
+  "j_arm*_7": [100, 5]
+  "hip_yaw_*": [3000, 30]
+  "hip_pitch_*": [3000, 30]
+  "knee_pitch_*": [3000, 30]
+  "ankle_pitch_*": [1000, 10]
+  "ankle_yaw_*": [300, 10]
+  "neck_pitch": [10, 1]
+  "neck_yaw": [10, 1]
+  "torso_yaw": [1000, 30]
+  "j_wheel_*": [0, 30]
+  "velodyne_*": [10, 1]
+  "d435_*": [10, 1]
+  "dagana_*": [50, 1]
+
+startup_motor_pd:
+  "j_arm*_1": [500, 30]
+  "j_arm*_2": [500, 30]
+  "j_arm*_3": [500, 30]
+  "j_arm*_4": [500, 30]
+  "j_arm*_5": [100, 5]
+  "j_arm*_6": [100, 5]
+  "j_arm*_7": [100, 5]
+  "hip_yaw_*": [800, 80]
+  "hip_pitch_*": [800, 80]
+  "knee_pitch_*": [800, 80]
+  "ankle_pitch_*": [800, 80]
+  "ankle_yaw_*": [480, 50]
+  "neck_pitch": [10, 1]
+  "neck_yaw": [10, 1]
+  "torso_yaw": [800, 80]
+  "j_wheel_*": [0, 30]
+  "velodyne_*": [10, 1]
+  "d435_*": [10, 1]
+  "dagana_*": [50, 1]
+
+motor_vel:
+  j_wheel_*: [1]
+  neck_velodyne: [1]
+  
+# hal
+xbotcore_device_configs:
+    sim: $PWD/hal/centauro_gz.yaml
+    dummy: $PWD/hal/centauro_dummy.yaml
+
+# threads
+xbotcore_threads:
+    rt_main:  {sched: fifo , prio: 60, period: 0.001}
+    nrt_main: {sched: other, prio: 0 , period: 0.005}
+
+# plugins
+xbotcore_plugins:
+
+    homing:
+        thread: rt_main
+        type: homing
+
+    ros_io: {thread: nrt_main, type: ros_io}
+
+    ros_control: {thread: nrt_main, type: ros_control, params: {autostart: {type: bool, value: true}}}
+
+# global parameters
+xbotcore_param:
+  /xbot/hal/joint_safety/delta_check_dt: 0.02
+  /xbot/hal/joint_safety/filter_autostart: {value: true, type: bool}
+  # /xbot/hal/joint_safety/filter_cutoff_hz: {value: 1.0, type: double}
+  /xbot/hal/joint_safety/filter_safe_cutoff_hz: {value: 5.0, type: double}
+  /xbot/hal/joint_safety/filter_medium_cutoff_hz: {value: 15.0, type: double}
+  /xbot/hal/joint_safety/filter_fast_cutoff_hz: {value: 25.0, type: double}
+  /xbot/hal/enable_safety: {value: true, type: bool}