why you use different operator for linear velocity and angular velocity?

[hdadong]

https://github.com/CLeARoboticsLab/ssrl/blob/main/ssrl/brax/envs/go1_go_fast.py#L688

Why you use different operator for linear velocity and angular velocity?
In my view, it should use the inv_rotate operator to transform the velocity of world frame to body frame and use the rotate operator to transform the velocity of body frame to world frame.
By the way,
Does pipeline_state.qd.at[self._xml_base_vel_idxs] represent the linear velocity in the world frame?

Does pipeline_state.qd.at[self._xml_rpy_rate_idxs] represent the angular velocity in the body frame?

[jake-levy]

pipeline_state.qd represents both linear and angular velocity in the world frame. I can't find a source right now, but there is mathematical reason for why you do the forward rotation for one, and the inverse for the other. I think I remember it having to do with how brax.math.rotate and inv_rotate are configured. This was a big pain point for me during debugging; configuring it this way fixed a lot of problems for me.

[hdadong]

After investigation, I confirm the following:

1. Linear Velocity Check:

   • The expression (pipeline_state.x.pos[0] - last_pos) / dt is equal to pipeline_state.xd.vel[0].
   • Since pipeline_state.x.pos is in the global coordinate system, pipeline_state.xd.vel represents the linear velocity in the world frame.

2. Base Velocity Confirmation:

   • pipeline_state.xd.vel[0] is equal to pipeline_state.qd[self._xml_base_vel_idxs].
   • Because xd.vel is the linear velocity in the world frame, pipeline_state.qd[self._xml_base_vel_idxs] also represents the linear velocity in the world frame.

3. Angular Velocity Conversion:

   • The result of inv_rotate(pipeline_state.xd.ang[0], quat) is equal to pipeline_state.qd[self._xml_rpy_rate_idxs].
   • Since xd.ang represents the angular velocity in the world frame, the value in pipeline_state.qd[self._xml_rpy_rate_idxs] is expressed in the local frame.

[hdadong]

Another way to verify this is to assume that pipeline_state.qd[self._xml_rpy_rate_idxs] represents the angular velocity in the local frame. You can then convert it to Euler angle rates and compare it with the actual Euler angle rate obtained from finite differences. The result shows that they are equal, which indicates that pipeline_state.qd[self._xml_rpy_rate_idxs] indeed corresponds to the angular velocity in the local frame.

The implementation of function euler_angle_rates follows the method explained in this video: https://www.youtube.com/watch?v=Z8nwjouP58o

@jit
def euler_angle_rates(euler_angles, omega_body):
    roll, pitch, yaw = euler_angles

    T = jp.array([
        [1, jp.sin(roll) * jp.tan(pitch), jp.cos(roll) * jp.tan(pitch)],
        [0, jp.cos(roll), -jp.sin(roll)],
        [0, jp.sin(roll) / jp.cos(pitch), jp.cos(roll) / jp.cos(pitch)]
    ])

    euler_rates = T @ omega_body
    return euler_rates

The full code used for verification is shown below, which loads a trained model, runs the policy, and compares the predicted Euler angle rates from the angular velocity with the finite-difference rates from the actual Euler angles over time. The results are plotted for comparison.

@jit
def euler_angle_rates(euler_angles, omega_body):
    roll, pitch, yaw = euler_angles

    T = jp.array([
        [1, jp.sin(roll) * jp.tan(pitch), jp.cos(roll) * jp.tan(pitch)],
        [0, jp.cos(roll), -jp.sin(roll)],
        [0, jp.sin(roll) / jp.cos(pitch), jp.cos(roll) / jp.cos(pitch)]
    ])

    euler_rates = T @ omega_body
    return euler_rates

def int_multiply(x, y):
    return int(x * y)
import math

OmegaConf.register_new_resolver("int_multiply", int_multiply)

@hydra.main(config_path="configs", config_name="booster_bptt")
def evaluate_action(cfg: DictConfig):

    num_critic_obs = 14
    num_actions = 12
    num_observations = 47
    actor_critic = ActorCritic(num_actions, num_observations, num_critic_obs).to('cpu')
    model_dict = torch.load('xxxl/model_10000.pth', map_location="cpu", weights_only=False)
    print(model_dict)
    actor_critic.load_state_dict(model_dict["model"])

    actor_critic.eval() # switch to evaluation mode (dropout for example)
    
    os.environ["CUDA_DEVICE_ORDER"] = "PCI_BUS_ID"
    os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = ".90"
    if 'CUDA_VISIBLE_DEVICES' not in os.environ:
        os.environ["CUDA_VISIBLE_DEVICES"] = str(cfg.gpus)

    from jax import config
    config.update("jax_enable_x64", True)
    env_kwargs = cfg.env_bptt

    env = BoosterGoFast(backend='generalized' , **env_kwargs)


    @jax.jit
    def jit_step(state, action):
        return env.step(state, action)

    @jax.jit
    def jit_reset(key):
        return env.reset(key)


    def evaluate_and_render(env):
        key = jax.random.PRNGKey(1)
        state = jit_reset(key)

        done = False
        states = []
        i = 0
        start_time = time.time()

        count_lowlevel = 0
        action = jp.zeros(12)
        last_euler_zyx = jp.array([0.0,0.0,0.0])

        euler_rate_list = []          
        euler_diff_list = []           
        euler_list = []

        while not done:
            pipeline_state = state.pipeline_state

            i += 1
            obs = np.zeros([1,num_observations], dtype=np.float32)            

            euler_rate = euler_angle_rates(quat_to_eulerzyx(pipeline_state.x.rot[0]), np.array(pipeline_state.qd[3:6]))
            euler_rate_list.append(np.array(euler_rate))
            euler_diff_list.append(np.array((quat_to_eulerzyx(pipeline_state.x.rot[0]) - last_euler_zyx)/env.dt))

            euler_list.append(np.array(quat_to_eulerzyx(pipeline_state.x.rot[0])))
            last_euler_zyx = quat_to_eulerzyx(pipeline_state.x.rot[0])

            obs[0, :47] =  state.obs[:47]
            dist = actor_critic.act(torch.tensor(obs))
            action = dist.loc.detach()
            action[:] = torch.clip(action, -1.0, 1.0)
            
            action = torch_to_jax(action)[0]
            state = jit_step(state, action)

            count_lowlevel += 10
            states.append(state.pipeline_state) 

            
            if i % 10 == 0:  
                print(f"Step {i}")
                print(f"Time elapsed: {time.time() - start_time:.2f}s")
                start_time = time.time()
            if i > 150:  
                break

        render_html = html.render(env.sys.replace(dt=env.dt), states, height=500) 
        output_path = Path(__file__).parent / 'evaluation_render.html'
        with open(output_path, 'w') as f:
            f.write(render_html)
        print(f"Render saved to {output_path}")
        time_array = np.linspace(0, len(euler_rate_list)*0.02, len(euler_rate_list))

        plt.figure()
        euler_rate_array = np.array(euler_rate_list)
        euler_diff_array = np.array(euler_diff_list)
        euler_array = np.array(euler_list)
        plt.plot(time_array, euler_rate_array[:, 0], label='euler_rate')
        plt.plot(time_array, euler_diff_array[:, 0], label='euler_diff')
        plt.plot(time_array, euler_array[:, 0], label='euler')
        plt.xlabel('Time (s)')
        plt.ylabel('Angular rate (rad/s)')
        plt.title('0')
        plt.legend()
        plt.show()

        plt.figure()
        plt.plot(time_array, euler_rate_array[:, 1], label='euler_rate')
        plt.plot(time_array, euler_diff_array[:, 1], label='euler_diff')
        plt.plot(time_array, euler_array[:, 1], label='euler')

        plt.xlabel('Time (s)')
        plt.ylabel('Euler Rate (rad/s)')
        plt.title(' 1')
        plt.legend()
        plt.show()

        plt.figure()
        plt.plot(time_array, euler_rate_array[:, 2], label='euler_rate')
        plt.plot(time_array, euler_diff_array[:, 2], label='euler_diff')
        plt.plot(time_array, euler_array[:, 2], label='euler')

        plt.xlabel('Time (s)')
        plt.ylabel('Finite Difference Euler Rate (rad/s)')
        plt.title('2')
        plt.legend()
        plt.show()

    evaluate_and_render(env)

if __name__ == '__main__':
    evaluate_action()

The results are as follows:

[hdadong]

I have open a discussion in brax repo:
google/brax#590