Sim2real recommendations for locomotion tasks using a go2

Good afternoon, I hope you're having a great day!

I modified a bit the files for the Unitree Go1 so I could train a policy for the walking locomotion task in a Unitree Go2. When I use the methodology you used to create the videos for testing giving privileged state, normal state and noise, it works quite well, but when I try to test it with mujoco cpu (cpu version) or the gpu version without the noise or privileged state, it works quite poorly (very different to the videos you uploaded to the website :( ). The major issue right now is that it happens one of two things:
1. It tries to make the move according to the command but failing to do it (for example, in case it tries to go forward, it makes some inclination forward and stay there)
2. It starts vibrating and then starts kicking very hard with the joints of the legs.

As context, this what I saw in the simulator after training the policy:

## 1) FR Hip

## 2) FR Thigh

## 3) FR Calf

From what I saw, even though the policy was noisy, the go2 was able to walk properly in simulation, but in reality, as the value of kp and kd are relatively big, it actually try to go to the value and tends to overshoot.

I also noticed that in your paper you mentioned you made a locomotion controller similar to the one used in "Concurrent Training of a Control Policy and a State Estimator for Dynamic and Robust Legged Locomotion (2022) by Gwanghyeon Ji, Juhyeok Mun and Hyeongjun Kim, but as I could understand, they use a neural network to estimate the state, they don't pass it as raw data from sensors. Did you do this as well?? I think maybe this could be the central point that makes my policy work so differently.

Right now I'm thinking maybe I should increase the size of the network or maybe train it for more epochs as it was mentioned in the Go2 support PR in this repository. Please let me know if you have any recommendation or advice for this.

This is the code im using right now for the action server that will give the control commands to the go2:

```
import logging
import os

from matplotlib.pylab import NaN
import rclpy
from rclpy.node import Node, MutuallyExclusiveCallbackGroup
from unitree_go.msg import LowState, LowCmd, JointPosition
import numpy as np
from scipy.spatial.transform import Rotation as R
from rclpy.action import ActionServer
from unitree_go.action import WalkingTask
from rclpy.executors import MultiThreadedExecutor
import jax
from brax.training.agents.ppo import networks as ppo_networks
from brax.io import model
import sys
import time
from nav_msgs.msg import Odometry
from walking_task_go2.utils.motor_crc import get_crc

# Import local mujoco_playground package
sys.path.append("/root/utils/mujoco_playground")
from mujoco_playground.config import locomotion_params
from mujoco_playground import registry

os.environ['JAX_PLATFORMS'] = 'gpu'
jax.config.update('jax_platform_name', 'gpu')
os.environ['XLA_PYTHON_CLIENT_PREALLOCATE'] = 'false'
jax.config.update('jax_disable_jit', False)
# Remove jax debug logging
logging.getLogger('jax').setLevel(logging.WARNING)

class ActionServerWalkingTaskGo2(Node):

def __init__(self):
super().__init__('as_walking_task_go2')
self.callback_group_timer = MutuallyExclusiveCallbackGroup()
self.callback_group_timer_subscribers = MutuallyExclusiveCallbackGroup()
self.callback_group_action_server = MutuallyExclusiveCallbackGroup()
self._action_server = ActionServer(
self,
WalkingTask,
'velocity_command_goal',
self.execute_callback, callback_group=self.callback_group_action_server)
self.low_state_subscriber = self.create_subscription(
LowState,
'lf/lowstate',
self.low_state_callback,
10, callback_group=self.callback_group_timer_subscribers)
self.odometry_sub = self.create_subscription(Odometry,
'odometry/filtered',
self.odometry_callback,
10, callback_group=self.callback_group_timer_subscribers)
self.action_made_publisher = self.create_publisher(JointPosition, 'action_made_joint', 10, callback_group=MutuallyExclusiveCallbackGroup())
self.motor_cmd_publisher = self.create_publisher(LowCmd, 'lowcmd', 10)
self.timer_assign_kp_kd = self.create_timer(1.0, self.assign_parameters_callback, callback_group=self.callback_group_timer)
self.control_dt = 0.02
self.motor_timer = self.create_timer(self.control_dt, self.apply_command_to_go2)
self.gyro = None
self.gravity = None
self.linvel = None
self.have_state = False
self.joint_angles = np.zeros(12).tolist()
self.joint_speeds = np.zeros(12).tolist()

# self.DEFAULT_POSE = [
# 0.0013248026371002197,
# 0.8022780418395996,
# -1.539233922958374,
# 0.0010065734386444092,
# 0.7960106134414673,
# -1.5339735746383667,
# -0.017512917518615723,
# 0.7915239930152893,
# -1.5242899656295776,
# 0.02889728546142578,
# 0.784953773021698,
# -1.5043385028839111,
# ]

self.DEFAULT_POSE = [0.0, 0.9, -1.8, 0.0, 0.9, -1.8, 0.0, 0.9, -1.8, 0.0, 0.9, -1.8]

self.last_action = np.zeros(12).tolist()
self.act_rng, self.rng = jax.random.split(jax.random.PRNGKey(0))
self.env_name = 'Go2JoystickFlatTerrain'
self.env_cfg = registry.get_default_config(self.env_name)
inference = self.import_model(self.env_name, self.env_cfg)
self.inference_fn = jax.jit(inference)

self.low_cmd = LowCmd()
self.low_cmd.head = [0xFE, 0xEF]
self.low_cmd.level_flag = 0xFF
self.action = np.zeros((1, 12))
self.timer_count = 0
self.stand_up = False
self.stand_ramp_duration_s = 2.5
self.stand_ramp_steps = max(1, int(self.stand_ramp_duration_s / self.control_dt))
self.stand_ramp_step = 0
self.stand_start_pos = None
self.stand_kp = 100.0
self.stand_kd = 5.0
self.declare_parameter('policy_kp', 35.0)
self.declare_parameter('policy_kd', 0.6)

for j in range(12):
self.low_cmd.motor_cmd[j].mode = 0x01
self.low_cmd.motor_cmd[j].kp = self.stand_kp
self.low_cmd.motor_cmd[j].kd = self.stand_kd

def low_cmd_callback(self, msg):
"""
Define the new low_cmd message to be published inside the class
Args:
msg (LowCmd): The received LowCmd message containing motor command information.
"""
self.low_cmd = msg

def assign_parameters_callback(self):
"""
Assign the kp and kd parameters from ROS2 parameters to the low_cmd message.
"""
policy_kp = self.get_parameter('policy_kp').get_parameter_value().double_value
policy_kd = self.get_parameter('policy_kd').get_parameter_value().double_value
for i in range(12):
self.low_cmd.motor_cmd[i].kp = policy_kp
self.low_cmd.motor_cmd[i].kd = policy_kd

def import_model(self, env_name='Go2JoystickFlatTerrain', env_cfg=None):
"""
This function imports the PPO trained model for the task 'Go2JoystickFlatTerrain'.

Args:
env_name (str): The name of the environment. Default is 'Go2JoystickFlatTerrain'.
env_cfg: The environment configuration.

Returns:
inference_fnTEST: The inference function for the trained model.
"""


ppo_params = locomotion_params.brax_ppo_config(env_name)
model_params = model.load_params('src/unitree/walking_task_go2/resource/ppo_go2joystick_flatterrain_params_ctrl_004_sim_0004_iter_3_impratio_100_KP_35_KD_0.5_v16')
ppo = ppo_networks.make_ppo_networks(action_size=env_cfg.action_size, observation_size=env_cfg.observation_size, **ppo_params.network_factory)
make_inference = ppo_networks.make_inference_fn(ppo)
inference_fnTEST = make_inference(model_params, deterministic=True)

return inference_fnTEST



def apply_command_to_go2(self):
# Standing up sequence
if not self.stand_up:
if not self.have_state:
return
if self.stand_start_pos is None:
self.stand_start_pos = np.array(self.joint_angles, dtype=float)
self.stand_ramp_step = min(self.stand_ramp_step + 1, self.stand_ramp_steps)
percent = self.stand_ramp_step / self.stand_ramp_steps
target_pose = np.array(self.DEFAULT_POSE, dtype=float)
cmd_pose = (1.0 - percent) * self.stand_start_pos + percent * target_pose
for i in range(12):
self.low_cmd.motor_cmd[i].q = float(cmd_pose[i])
self.low_cmd.motor_cmd[i].kp = self.stand_kp
self.low_cmd.motor_cmd[i].kd = self.stand_kd
self.low_cmd.crc = get_crc(self.low_cmd)
self.motor_cmd_publisher.publish(self.low_cmd)
if self.stand_ramp_step >= self.stand_ramp_steps:
self.stand_up = True
return

# Assigning commands from policy
# timer_assignation = time.time()
for i in range(12):
self.low_cmd.motor_cmd[i].q = float(self.DEFAULT_POSE[i] + self.action[0][i] * 0.5)

# print(f"Time used assigning commands: {time.time() - timer_assignation} seconds")

# Time used calculating CRC
# start_time = time.time()
self.low_cmd.crc = get_crc(self.low_cmd)
# print(f"CRC calculation time: {time.time() - start_time} seconds")
# print(f"Applied command, action: {[self.low_cmd.motor_cmd[i].q for i in range(12)]}")
# print(f"Joint angles: {self.joint_angles}")
joint_position = JointPosition()
joint_position.positions = [self.low_cmd.motor_cmd[i].q for i in range(12)]
self.action_made_publisher.publish(joint_position)

self.motor_cmd_publisher.publish(self.low_cmd)

def execute_callback(self, goal_handle):

try:
# Obtain observations
observation = self.get_observation(commanded_velocity=list(goal_handle.request.commanded_velocity)).tolist()

observation.extend(self.last_action)
# print(f"Commanded velocity: {goal_handle.request.commanded_velocity}")Okay, im having a problem, with kp=35 and kd=5 works well in the policy and it was the value recommended for the walking policy in mujoco playground. The problem is that when I try to put it in stand (make the dog stay in its 4 feet), it doesn't go all the way up, as it didn't have the sufficient force to lock all the joints in the position given, so what I did was increase kp to 60 and decrease kd to 0.5 but that didn't worked as the policy in general didn't work very well (my hypothesis is that it was overreacting) but the standing position was better than right now.
observation.extend(list(goal_handle.request.commanded_velocity))
# Convert observation to jax numpy array
observation = jax.numpy.array(observation).reshape(1, -1)[0]

# Stopping conditions (Will stand)
if goal_handle.request.commanded_velocity[0] == 2.0 \
and goal_handle.request.commanded_velocity[1] == 2.0 \
and goal_handle.request.commanded_velocity[2] == 2.0:
self.action = jax.numpy.zeros((1, 12))
self.DEFAULT_POSE = [0.0, 0.9, -1.8, 0.0, 0.9, -1.8, 0.0, 0.9, -1.8, 0.0, 0.9, -1.8]
# self.DEFAULT_POSE = [
# 0.0013248026371002197,
# 0.8022780418395996,
# -1.539233922958374,
# 0.0010065734386444092,
# 0.7960106134414673,
# -1.5339735746383667,
# -0.017512917518615723,
# 0.7915239930152893,
# -1.5242899656295776,
# 0.02889728546142578,
# 0.784953773021698,
# -1.5043385028839111,
# ]
else:
self.action = self.inference_fn(observation, self.act_rng)
# Saving results
observation = [float(obs) for obs in observation]
# print(f"Observation length: {len(observation)}")
# Compute new seed with jax and save variables
self.act_rng, self.rng = jax.random.split(self.rng)
self.last_action = self.action[0].tolist()
# print(f"Action taken: {self.action[0].tolist()}")
self.timer_count += 1

result = WalkingTask.Result()
result.observations = observation
goal_handle.succeed()

except Exception as e:
raise Exception(f"Error in execute_callback: {e}")
return result

def low_state_callback(self, msg):
"""
This function is called whenever a new LowState message is received.
It updates the gyro, gravity, joint angles, and joint speeds based on the message data.

Args:
msg (LowState): The received LowState message containing IMU and motor state information.
"""
self.gyro = msg.imu_state.gyroscope
rot = R.from_quat([msg.imu_state.quaternion[1],
msg.imu_state.quaternion[2],
msg.imu_state.quaternion[3],
msg.imu_state.quaternion[0]])
g_world = np.array([0.0, 0.0, -1.0])
self.gravity = rot.inv().apply(g_world)
self.have_state = True
for i in range(12):
self.joint_angles[i] = msg.motor_state[i].q
self.joint_speeds[i] = msg.motor_state[i].dq

def odometry_callback(self, msg):
"""Updates the linear velocity whenever a new Odometry message is received.
Args:
msg (Odometry): The received Odometry message containing velocity information.
(among other things)"""
lin = msg.twist.twist.linear
if all(map(np.isfinite, (lin.x, lin.y, lin.z))):
self.linvel = [np.clip(lin.x, -1.0, 1.0), np.clip(lin.y, -1.0, 1.0), np.clip(lin.z, -1.0, 1.0)]

def get_observation(self, commanded_velocity=[0.0, 0.0, 0.0]):
"""
Concat the observation vector from the current state variables.
"""
# print(f"Commanded velocity for observation: {commanded_velocity}")
# print(f"Commanded velocity but with list: {[commanded_velocity[0], commanded_velocity[1], commanded_velocity[2]]}")
observation = []
if self.gyro is not None and self.gravity is not None and self.linvel is not None and self.have_state:
observation.extend([float(self.linvel[0]), float(self.linvel[1]), float(self.linvel[2])])
observation.extend([self.gyro[0], self.gyro[1], self.gyro[2]])
observation.extend(self.gravity.tolist())
observation.extend(self.joint_angles - np.array(self.DEFAULT_POSE))
observation.extend(self.joint_speeds)
# print(f"""Observation constructed:
# Linear Velocity: {[self.linvel[0], self.linvel[1], self.linvel[2]]}
# Gyro: {self.gyro}
# Gravity: {self.gravity}
# Relative Joint Angles: {self.joint_angles - np.array(self.DEFAULT_POSE)}
# Joint Speeds: {self.joint_speeds}
# Last Action: {self.last_action}
# """)

return np.array(observation)

def main(args=None):
rclpy.init(args=args)

as_walking_task_go2 = ActionServerWalkingTaskGo2()
executor = MultiThreadedExecutor()
executor.add_node(as_walking_task_go2)

try:
executor.spin()
except KeyboardInterrupt:
pass
finally:
executor.shutdown()
as_walking_task_go2.destroy_node()
rclpy.shutdown()

if __name__ == '__main__':
main()

```

Thank you in advance and have a great day!