Fase 2: Movimiento básico

Objetivos de la fase 2

El objetivo principal de esta fase es que el robot sea capaz de combinar equilibrio y movimiento de forma controlada. De momento no busco mucha velocidad ni precisión en el movimiento, sino un comportamiento estable.

Por lo tanto los objetivos son:

• Que el robot sea capaz de mantenerse en equilibrio cuando el comando de movimiento es cero.

• Que el robot sea capaz de moverse cuando se le solicita, manteniendo el equilibrio durante el desplazamiento.

Estrategia de entrenamiento

Para lograr los objetivos, será necesario agregar nuevas observaciones para incluir información relacionada con el movimiento y los comandos de velocidad, ya que las utilizadas en la Fase 1 no son suficientes para lograr el objetivo del movimiento controlado. En principio, el vector de acción no va a cambiar, ya que solo actuaremos sobre la velocidad de las ruedas igual que en la fase anterior.

Este cambio en el espacio de observaciones hace que la política entrenada en la Fase 1 deje de ser compatible y no pueda reutilizarse directamente (mas info en ¿Entrenar de cero o a partir de una política anterior?). Por este motivo, el primer paso de esta fase es entrenar desde cero una nueva política utilizando el vector de observaciones ampliado, pero manteniendo inicialmente el objetivo de equilibrio.

Una vez obtenga una política base capaz de mantenerse estable con el nuevo vector de observaciones, utilizaré esta política como punto de partida para entrenar el movimiento. De esta forma, el aprendizaje del desplazamiento se apoyará sobre una base de equilibrio ya aprendida, en lugar de intentar aprender ambos comportamientos a la vez.

Entrenamiento del objetivo de estabilidad (Fase 1)

Para este primer entrenamiento los cambios que realizaremos sobre el código van a ser mínimos.

Por un lado, en la funcion _get_observations del simplerobot_env.py agregaremos la componente del eje X de la velocidad lineal del robot y del comando de velocidad al vector de las observaciones.

def _get_observations(self) -> dict:
    self.projected_gravity = self.robot.data.projected_gravity_b # Shape (N,3)
    self.angular_velocity = self.robot.data.root_ang_vel_b # Shape (N,3)
    self.wheel_vel = self.robot.data.joint_vel #Shape (N, num_joints)
    self.lin_vel_b = self.robot.data.root_lin_vel_b[:,0:1]
    self.vx_cmd = self.commands[:,0:1]

    obs = torch.cat(
        [
            self.projected_gravity,
            self.angular_velocity,
            self.wheel_vel,
            self.lin_vel_b,
            self.vx_cmd,
        ],
        dim=-1,
    )

    observations = {"policy": obs}

    return observations

Y, por ende, tendremos que adaptar el tamaño del observation_space del simplerobot_env_cfg.py

observation_space = 3 + 3 + 2 + 1 + 1 # gravity vector (3), angular velocity (3), wheel velocities (2), linear velocity X (1), commanded velocity X (1)

Hechos estos 2 pequeños cambios, lanzo el entrenamiento desde cero para que el robot aprenda a estabilizarse de nuevo.

~/IsaacLab/isaaclab.sh -p ~/SimpleRobot/scripts/rsl_rl/train.py --task=Template-Simplerobot-Direct-v0

Tras aproximadamente 100-150 iteraciones, el robot aprende a estabilizarse sin problema.

Entrenamiento del objetivo de movimiento (Fase 2)

Ahora ya tengo una politica capaz de mantener el robot estable de la que puedo partir para entrenar el movimiento.

Para que la política no se olvide de cómo mantener la estabilidad cuando el comando de velocidad es nulo, tendremos que hacer que no todos los entornos simulados se comporten igual. Algunos robots recibirán comandos de movimiento, mientras que otros no. Esto permite entrenar el desplazamiento sin perder la capacidad de mantenerse estable cuando no el comando de velocidad es cero, evitando que el robot olvide el comportamiento aprendido en el entrenamiento anterior.

A continuación detallaré los cambios realizados sobre el archivo simplerobot_env.py, pero antes dejo el código por aquí:

simplerobot_env.py

# Copyright (c) 2022-2025, The Isaac Lab Project Developers (https://github.com/isaac-sim/IsaacLab/blob/main/CONTRIBUTORS.md).
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

from __future__ import annotations

import torch
from collections.abc import Sequence

import isaaclab.sim as sim_utils
from isaaclab.assets import Articulation
from isaaclab.envs import DirectRLEnv
from isaaclab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane

from .simplerobot_env_cfg import SimplerobotEnvCfg
import random


class SimplerobotEnv(DirectRLEnv):
    cfg: SimplerobotEnvCfg

    def __init__(self, cfg: SimplerobotEnvCfg, render_mode: str | None = None, **kwargs):
        super().__init__(cfg, render_mode, **kwargs) # this super call will invoke _setup_scene()

        self.dof_idx, _ = self.robot.find_joints(self.cfg.dof_names) # get the indices of the controlled dofs


    def _setup_scene(self):
        self.robot = Articulation(self.cfg.robot_cfg)
        # add ground plane
        spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg())
        # clone and replicate
        self.scene.clone_environments(copy_from_source=False) # copy_from_source=False will use instanceable references for better performance. 
        # add articulation to scene
        self.scene.articulations["robot"] = self.robot
        # add lights
        light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75))
        light_cfg.func("/World/Light", light_cfg)

        self.commands = torch.zeros((self.cfg.scene.num_envs, 3)).cuda()  # initialize commands buffer

    # Both _pre_physics_step and _apply_action are not called every simulation step, but only at the steps when actions are applied (according to the decimation factor).
    # F.g., if decimation=2, these methods are called every 2 simulation steps: _pre_physics_step -> _apply_action -> physics step -> physics step -> _pre_physics_step -> _apply_action -> physics step -> physics step -> ...
    def _pre_physics_step(self, actions: torch.Tensor) -> None:
        # This method is called before the physics step. We store the actions to be applied later in _apply_action()
        self.actions = actions.clone() # Copy the actions and store them for use in _apply_action(). It acts as a buffer between the policy and the physics step.
        self.actions = self.actions * self.cfg.actions_scale # scale the actions to reasonable values

    def _apply_action(self) -> None:
        # This method is called after the _pre_physics_step() and before the physics step. Here we apply the stored actions to the robot.
        self.robot.set_joint_velocity_target(self.actions, joint_ids=self.dof_idx) # set the wheel velocities according to the actions
    
    def _get_observations(self) -> dict:
        self.projected_gravity = self.robot.data.projected_gravity_b # Shape (N,3)
        self.angular_velocity = self.robot.data.root_ang_vel_b # Shape (N,3)
        self.wheel_vel = self.robot.data.joint_vel #Shape (N, num_joints)
        self.lin_vel_b = self.robot.data.root_lin_vel_b[:,0:1]
        self.vx_cmd = self.commands[:,0:1]

        obs = torch.cat(
            [
                self.projected_gravity,
                self.angular_velocity,
                self.wheel_vel,
                self.lin_vel_b,
                self.vx_cmd,
            ],
            dim=-1,
        )

        observations = {"policy": obs}

        return observations


    def _get_rewards(self) -> torch.Tensor:

        # --- Inclinación ---
        # projected_gravity_b ≈ [0, 0, -1] cuando está vertical
        tilt_error = self.projected_gravity[:, 0]**2 # Solo nos importa el eje X (adelante/atrás)
        upright_reward = torch.exp(-5.0 * tilt_error)

        # --- Velocidad angular (evitar oscilaciones) ---
        ang_vel_penalty = torch.sum(self.angular_velocity ** 2, dim=1)

        # --- Alive ---
        alive_reward = (~self.reset_buf).float()

        # --- Penalización por no moverse cuando hay comando de velocidad no es 0 ---
        cmd = self.vx_cmd.squeeze(-1)
        cmd_mask = torch.abs(cmd) > 1e-3
        vel_error = torch.abs(self.lin_vel_b - self.vx_cmd).squeeze(-1)  # (N,)
        vel_penalty = torch.zeros_like(cmd)
        vel_penalty[cmd_mask] = vel_error[cmd_mask] / torch.abs(cmd[cmd_mask])

        #--- Penalización por no estar quieto cuando el comando es 0 ---
        still_mask = torch.abs(cmd) <= 1e-3
        no_still_penalty = torch.zeros_like(cmd)
        no_still_penalty[still_mask] = torch.abs(self.lin_vel_b.squeeze(-1))[still_mask]

        # --- Penalización por diferencia de velocidad entre ruedas (evitar giros) ---
        diff_penalty = (self.wheel_vel[:, 0] - self.wheel_vel[:, 1])**2

        # --- Reward final ---
        reward = (
            0.2 * upright_reward
            - 0.05 * ang_vel_penalty
            + 0.5 * alive_reward
            - 1.0 * vel_penalty
            - 1.0 * no_still_penalty
            - 0.2 * diff_penalty
        )

        return reward

    def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]:
        time_out = self.episode_length_buf >= self.max_episode_length - 1 # If the episode length buffer exceeds the max length, we time out

        # Si el robot se inclina más de ~50° en el eje X o Y, se considera que ha caído
        fallen = torch.any(torch.abs(self.projected_gravity[:, :2]) > 0.8727, dim=1)

        return fallen, time_out # for now we only terminate episodes on timeout, forgetting about other termination conditions

    def _reset_idx(self, env_ids: Sequence[int] | None):
        if env_ids is None:
            env_ids = self.robot._ALL_INDICES
        super()._reset_idx(env_ids)

        default_root_state = self.robot.data.default_root_state[env_ids] # get the default root state (position and orientation in World frame)
        default_root_state[:, :3] += self.scene.env_origins[env_ids]     # offset the position according to the environment origin
        default_root_state[:, 2] += 0.0  # SUBIR ROBOT

        # pick new commands for reset envs and normalize them just like in the setup
        if random.random() < 0.3:
            self.commands[env_ids] = torch.zeros((len(env_ids), 3)).cuda()
        else:
            self.commands[env_ids, 0] = 0.2
            self.commands[env_ids, 1:] = 0.0

        self.robot.write_root_state_to_sim(default_root_state, env_ids)  # reset the root state of the robot

Cambios en el _get_rewards

def _get_rewards(self) -> torch.Tensor:

    # --- Inclinación ---
    # projected_gravity_b ≈ [0, 0, -1] cuando está vertical
    tilt_error = self.projected_gravity[:, 0]**2 # Solo nos importa el eje X (adelante/atrás)
    upright_reward = torch.exp(-5.0 * tilt_error)

    # --- Velocidad angular (evitar oscilaciones) ---
    ang_vel_penalty = torch.sum(self.angular_velocity ** 2, dim=1)

    # --- Alive ---
    alive_reward = (~self.reset_buf).float()

    # --- Penalización por no moverse cuando hay comando de velocidad no es 0 ---
    cmd = self.vx_cmd.squeeze(-1)
    cmd_mask = torch.abs(cmd) > 1e-3
    vel_error = torch.abs(self.lin_vel_b - self.vx_cmd).squeeze(-1)  # (N,)
    vel_penalty = torch.zeros_like(cmd)
    vel_penalty[cmd_mask] = vel_error[cmd_mask] / torch.abs(cmd[cmd_mask])

    #--- Penalización por no estar quieto cuando el comando es 0 ---
    still_mask = torch.abs(cmd) <= 1e-3
    no_still_penalty = torch.zeros_like(cmd)
    no_still_penalty[still_mask] = torch.abs(self.lin_vel_b.squeeze(-1))[still_mask]

    # --- Penalización por diferencia de velocidad entre ruedas (evitar giros) ---
    diff_penalty = (self.wheel_vel[:, 0] - self.wheel_vel[:, 1])**2

    # --- Reward final ---
    reward = (
        0.2 * upright_reward
        - 0.05 * ang_vel_penalty
        + 0.5 * alive_reward
        - 1.0 * vel_penalty
        - 1.0 * no_still_penalty
        - 0.2 * diff_penalty
    )

    return reward

En esta función he creado nuevas penalizaciones:

• vel_penalty: Si el robot debe moverse, penaliza al robot por la diferencia entre su velocidad y la comandada

• no_still_penalty: Penaliza al robot por moverse si no tiene que hacerlo

• diff_penalty: Penaliza al robot por la diferencia entre la velocidad de sus ruedas (evita giros sobre si mismo)

También he ajustado los pesos de las recompensas anteriores.

¿Cómo he llegado a estas recompensas y penalizaciones?

Prueba y error. Los pasos que seguí fueron los siguientes:

1. Empecé agregando la penalización vel_penalty solamente y lancé el entrenamiento

   1. Los robots no aprendían a moverse, se quedaban estáticos

2. Bajé los pesos de las recompensas upright_reward y ang_vel_penalty y lancé otro entrenamiento

   1. Los robots comenzaron a moverse. Todos ellos, incluso los que recibían comandos de velocidad nulos

3. Agregué la penalización no_still_penalty y lancé un nuevo entrenamiento

   1. Los robots aprendieron a moverse y a estarse quietos cuando debían

   2. Alguno de los robots que se movían lo hacían en círculos.

4. Agregué la penalización diff_penalty y lancé un nuevo entrenamiento

   1. Los robots dejaron de moverse en círculos.

Cambios en _reset_idx

def _reset_idx(self, env_ids: Sequence[int] | None):
    if env_ids is None:
        env_ids = self.robot._ALL_INDICES
    super()._reset_idx(env_ids)

    default_root_state = self.robot.data.default_root_state[env_ids] # get the default root state (position and orientation in World frame)
    default_root_state[:, :3] += self.scene.env_origins[env_ids]     # offset the position according to the environment origin
    default_root_state[:, 2] += 0.0  # SUBIR ROBOT

    # pick new commands for reset envs and normalize them just like in the setup
    if random.random() < 0.3:
        self.commands[env_ids] = torch.zeros((len(env_ids), 3)).cuda()
    else:
        self.commands[env_ids, 0] = 0.2
        self.commands[env_ids, 1:] = 0.0

    self.robot.write_root_state_to_sim(default_root_state, env_ids)  # reset the root state of the robot

Ahora al comando de velocidad se le asigna, para el 30% de los entornos una velocidad de 0 en todos los ejes. Para el otro 70% de los entornos al comando de velocidad se le asigna una velocidad en el eje X de 0.2 m/s

¡A entrenar!

Como he detallado en Entrenar a partir de una política anterior (weights-only), utilizaré el script de entrenamiento modificado para cargar los pesos de la política de un checkpoint anterior. Para ello utilizaré el argumento --load_policy <checkpoint_path>

En mi caso, quiero utilizar la política ya entrenada en la fase 1 (200 iteraciones de entrenamiento). Para ello el comando que voy a usar es:

~/IsaacLab/isaaclab.sh -p ~/SimpleRobot/scripts/rsl_rl/train_mod.py --task=Template-Simplerobot-Direct-v0 --load_policy /home/angellm/logs/rsl_rl/simplerobot_direct/2026-01-19_18-39-40/model_200.pt

Nada más empezar el entrenamiento, podemos ver que hay una cantidad de robots que se mantienen de pie sin caerse (los que no han recibido comandos de velocidad) y robots que tratan de moverse y se caen (los que si han recibido comandos de velocidad).

Se puede ver que despues 60 iteraciones los robots ya no parecen caerse durante el entrenamiento, ¡eso es muy buena señal! Como comparación, lanzando el entrenamiento desde cero (sin cargar una política anterior), los robots tardaban unas 50 iteraciones en aprender a estabilizarse y hasta las iteraciones 150-200 no aprendian a moverse sin caerse.

Evaluación

Pese a que alrededor de las 60 iteraciones la política ya parecía haber encontrado una solución estable, continué el entrenamiento hasta las 1000 iteraciones para observar su evolución a largo plazo.

En torno a las 400 iteraciones el robot comenzó a caerse de forma sistemática y la duración de los episodios disminuyó progresivamente. Pero tras otras 100–200 iteraciones, el entrenamiento se recuperó por sí solo y la política volvió a estabilizar el comportamiento del robot.

Este tipo de comportamiento parece ser normal en algoritmos de aprendizaje por refuerzo. Durante el entrenamiento, la política sigue explorando el espacio de soluciones y temporalmente puede alejarse de una solución estable para intentar maximizar la recompensa. Si esa nueva región del espacio de soluciones resulta no ser favorable, el propio proceso de aprendizaje puede reconducir la política hacia soluciones más estables. Este es el resultado:

Cuando al robot se le da un comando de velocidad este lo sigue sin caerse y cuando no se le da, se queda estable estáticamente.

¡Fase 2 completada con éxito! Y además he aprendido a realizar un entrenamiento basado en una política anterior.