创建直接工作流RL环境

创建直接工作流RL环境#

除了 envs.ManagerBasedRLEnv 类之外，还可以使用配置类来为更模块化的环境提供支持， DirectRLEnv 类允许在环境脚本化中进行更直接的控制。

直接工作流任务实现完全奖励和观察功能的直接任务脚本。这允许在方法的实现中更多地控制，比如使用Pytorch jit功能，并提供一个更少抽象的框架，更容易找到各种代码片段。

在本教程中，我们将使用直接工作流实现来配置cartpole环境，以创建一个平衡杆直立的任务。我们将学习如何使用实现场景创建，动作，重置，奖励和观察的功能来指定任务。

代码#

对于本教程，我们使用 omni.isaac.lab_tasks.direct.cartpole 模块中定义的cartpole环境。

cartpole_env.py的代码
  1# Copyright (c) 2022-2024, The Isaac Lab Project Developers.
  2# All rights reserved.
  3#
  4# SPDX-License-Identifier: BSD-3-Clause
  5
  6from __future__ import annotations
  7
  8import math
  9import torch
 10from collections.abc import Sequence
 11
 12from omni.isaac.lab_assets.cartpole import CARTPOLE_CFG
 13
 14import omni.isaac.lab.sim as sim_utils
 15from omni.isaac.lab.assets import Articulation, ArticulationCfg
 16from omni.isaac.lab.envs import DirectRLEnv, DirectRLEnvCfg
 17from omni.isaac.lab.scene import InteractiveSceneCfg
 18from omni.isaac.lab.sim import SimulationCfg
 19from omni.isaac.lab.sim.spawners.from_files import GroundPlaneCfg, spawn_ground_plane
 20from omni.isaac.lab.utils import configclass
 21from omni.isaac.lab.utils.math import sample_uniform
 22
 23
 24@configclass
 25class CartpoleEnvCfg(DirectRLEnvCfg):
 26    # env
 27    decimation = 2
 28    episode_length_s = 5.0
 29    action_scale = 100.0  # [N]
 30    num_actions = 1
 31    num_observations = 4
 32    num_states = 0
 33
 34    # simulation
 35    sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation)
 36
 37    # robot
 38    robot_cfg: ArticulationCfg = CARTPOLE_CFG.replace(prim_path="/World/envs/env_.*/Robot")
 39    cart_dof_name = "slider_to_cart"
 40    pole_dof_name = "cart_to_pole"
 41
 42    # scene
 43    scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=4096, env_spacing=4.0, replicate_physics=True)
 44
 45    # reset
 46    max_cart_pos = 3.0  # the cart is reset if it exceeds that position [m]
 47    initial_pole_angle_range = [-0.25, 0.25]  # the range in which the pole angle is sampled from on reset [rad]
 48
 49    # reward scales
 50    rew_scale_alive = 1.0
 51    rew_scale_terminated = -2.0
 52    rew_scale_pole_pos = -1.0
 53    rew_scale_cart_vel = -0.01
 54    rew_scale_pole_vel = -0.005
 55
 56
 57class CartpoleEnv(DirectRLEnv):
 58    cfg: CartpoleEnvCfg
 59
 60    def __init__(self, cfg: CartpoleEnvCfg, render_mode: str | None = None, **kwargs):
 61        super().__init__(cfg, render_mode, **kwargs)
 62
 63        self._cart_dof_idx, _ = self.cartpole.find_joints(self.cfg.cart_dof_name)
 64        self._pole_dof_idx, _ = self.cartpole.find_joints(self.cfg.pole_dof_name)
 65        self.action_scale = self.cfg.action_scale
 66
 67        self.joint_pos = self.cartpole.data.joint_pos
 68        self.joint_vel = self.cartpole.data.joint_vel
 69
 70    def _setup_scene(self):
 71        self.cartpole = Articulation(self.cfg.robot_cfg)
 72        # add ground plane
 73        spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg())
 74        # clone, filter, and replicate
 75        self.scene.clone_environments(copy_from_source=False)
 76        self.scene.filter_collisions(global_prim_paths=[])
 77        # add articultion to scene
 78        self.scene.articulations["cartpole"] = self.cartpole
 79        # add lights
 80        light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75))
 81        light_cfg.func("/World/Light", light_cfg)
 82
 83    def _pre_physics_step(self, actions: torch.Tensor) -> None:
 84        self.actions = self.action_scale * actions.clone()
 85
 86    def _apply_action(self) -> None:
 87        self.cartpole.set_joint_effort_target(self.actions, joint_ids=self._cart_dof_idx)
 88
 89    def _get_observations(self) -> dict:
 90        obs = torch.cat(
 91            (
 92                self.joint_pos[:, self._pole_dof_idx[0]].unsqueeze(dim=1),
 93                self.joint_vel[:, self._pole_dof_idx[0]].unsqueeze(dim=1),
 94                self.joint_pos[:, self._cart_dof_idx[0]].unsqueeze(dim=1),
 95                self.joint_vel[:, self._cart_dof_idx[0]].unsqueeze(dim=1),
 96            ),
 97            dim=-1,
 98        )
 99        observations = {"policy": obs}
100        return observations
101
102    def _get_rewards(self) -> torch.Tensor:
103        total_reward = compute_rewards(
104            self.cfg.rew_scale_alive,
105            self.cfg.rew_scale_terminated,
106            self.cfg.rew_scale_pole_pos,
107            self.cfg.rew_scale_cart_vel,
108            self.cfg.rew_scale_pole_vel,
109            self.joint_pos[:, self._pole_dof_idx[0]],
110            self.joint_vel[:, self._pole_dof_idx[0]],
111            self.joint_pos[:, self._cart_dof_idx[0]],
112            self.joint_vel[:, self._cart_dof_idx[0]],
113            self.reset_terminated,
114        )
115        return total_reward
116
117    def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]:
118        self.joint_pos = self.cartpole.data.joint_pos
119        self.joint_vel = self.cartpole.data.joint_vel
120
121        time_out = self.episode_length_buf >= self.max_episode_length - 1
122        out_of_bounds = torch.any(torch.abs(self.joint_pos[:, self._cart_dof_idx]) > self.cfg.max_cart_pos, dim=1)
123        out_of_bounds = out_of_bounds | torch.any(torch.abs(self.joint_pos[:, self._pole_dof_idx]) > math.pi / 2, dim=1)
124        return out_of_bounds, time_out
125
126    def _reset_idx(self, env_ids: Sequence[int] | None):
127        if env_ids is None:
128            env_ids = self.cartpole._ALL_INDICES
129        super()._reset_idx(env_ids)
130
131        joint_pos = self.cartpole.data.default_joint_pos[env_ids]
132        joint_pos[:, self._pole_dof_idx] += sample_uniform(
133            self.cfg.initial_pole_angle_range[0] * math.pi,
134            self.cfg.initial_pole_angle_range[1] * math.pi,
135            joint_pos[:, self._pole_dof_idx].shape,
136            joint_pos.device,
137        )
138        joint_vel = self.cartpole.data.default_joint_vel[env_ids]
139
140        default_root_state = self.cartpole.data.default_root_state[env_ids]
141        default_root_state[:, :3] += self.scene.env_origins[env_ids]
142
143        self.joint_pos[env_ids] = joint_pos
144        self.joint_vel[env_ids] = joint_vel
145
146        self.cartpole.write_root_pose_to_sim(default_root_state[:, :7], env_ids)
147        self.cartpole.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids)
148        self.cartpole.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)
149
150
151@torch.jit.script
152def compute_rewards(
153    rew_scale_alive: float,
154    rew_scale_terminated: float,
155    rew_scale_pole_pos: float,
156    rew_scale_cart_vel: float,
157    rew_scale_pole_vel: float,
158    pole_pos: torch.Tensor,
159    pole_vel: torch.Tensor,
160    cart_pos: torch.Tensor,
161    cart_vel: torch.Tensor,
162    reset_terminated: torch.Tensor,
163):
164    rew_alive = rew_scale_alive * (1.0 - reset_terminated.float())
165    rew_termination = rew_scale_terminated * reset_terminated.float()
166    rew_pole_pos = rew_scale_pole_pos * torch.sum(torch.square(pole_pos).unsqueeze(dim=1), dim=-1)
167    rew_cart_vel = rew_scale_cart_vel * torch.sum(torch.abs(cart_vel).unsqueeze(dim=1), dim=-1)
168    rew_pole_vel = rew_scale_pole_vel * torch.sum(torch.abs(pole_vel).unsqueeze(dim=1), dim=-1)
169    total_reward = rew_alive + rew_termination + rew_pole_pos + rew_cart_vel + rew_pole_vel
170    return total_reward

解释代码#

与基于管理器的环境类似，为任务定义一个配置类，以保存模拟参数、场景、演员和任务的设置。使用直接工作流实现， envs.DirectRLEnvCfg 类作为配置的基类使用。由于直接工作流实现不使用Action和Observation管理器，任务配置应定义环境的动作和观察数。

@configclass
class CartpoleEnvCfg(DirectRLEnvCfg):
   ...
   num_actions = 1
   num_observations = 4
   num_states = 0

配置类也可以用来定义特定于任务的属性，例如奖励项的缩放和重置条件的阈值。

@configclass
class CartpoleEnvCfg(DirectRLEnvCfg):
   ...
   # reset
   max_cart_pos = 3.0
   initial_pole_angle_range = [-0.25, 0.25]

   # reward scales
   rew_scale_alive = 1.0
   rew_scale_terminated = -2.0
   rew_scale_pole_pos = -1.0
   rew_scale_cart_vel = -0.01
   rew_scale_pole_vel = -0.005

在创建新环境时，代码应该定义一个新的类，该类继承自 DirectRLEnv 。

class CartpoleEnv(DirectRLEnv):
   cfg: CartpoleEnvCfg

   def __init__(self, cfg: CartpoleEnvCfg, render_mode: str | None = None, **kwargs):
     super().__init__(cfg, render_mode, **kwargs)

该类也可以包含所有类函数都可访问的类变量，包括应用动作、计算重置、奖励和观察的功能。

场景创建#

与基于管理器的环境相比，场景创建在直接工作流实现中提供了灵活性，使用户可以实现自己的场景创建功能。这包括将演员添加到场景上，克隆环境，过滤环境之间的碰撞，将演员添加到场景中，以及将其他附加属性添加到场景中，例如地面和灯光。这些操作应该在 _setup_scene(self) 方法中实现。

    def _setup_scene(self):
        self.cartpole = Articulation(self.cfg.robot_cfg)
        # add ground plane
        spawn_ground_plane(prim_path="/World/ground", cfg=GroundPlaneCfg())
        # clone, filter, and replicate
        self.scene.clone_environments(copy_from_source=False)
        self.scene.filter_collisions(global_prim_paths=[])
        # add articultion to scene
        self.scene.articulations["cartpole"] = self.cartpole
        # add lights
        light_cfg = sim_utils.DomeLightCfg(intensity=2000.0, color=(0.75, 0.75, 0.75))
        light_cfg.func("/World/Light", light_cfg)

定义奖励#

奖励函数应该在 _get_rewards(self) API中定义，它将奖励缓冲区作为返回值返回。在这个函数内部，任务可以自由实现奖励函数的逻辑。在这个示例中，我们实现了一个计算奖励函数各个组成部分的Pytorch JIT函数。

def _get_rewards(self) -> torch.Tensor:
     total_reward = compute_rewards(
         self.cfg.rew_scale_alive,
         self.cfg.rew_scale_terminated,
         self.cfg.rew_scale_pole_pos,
         self.cfg.rew_scale_cart_vel,
         self.cfg.rew_scale_pole_vel,
         self.joint_pos[:, self._pole_dof_idx[0]],
         self.joint_vel[:, self._pole_dof_idx[0]],
         self.joint_pos[:, self._cart_dof_idx[0]],
         self.joint_vel[:, self._cart_dof_idx[0]],
         self.reset_terminated,
     )
     return total_reward

@torch.jit.script
def compute_rewards(
    rew_scale_alive: float,
    rew_scale_terminated: float,
    rew_scale_pole_pos: float,
    rew_scale_cart_vel: float,
    rew_scale_pole_vel: float,
    pole_pos: torch.Tensor,
    pole_vel: torch.Tensor,
    cart_pos: torch.Tensor,
    cart_vel: torch.Tensor,
    reset_terminated: torch.Tensor,
):
    rew_alive = rew_scale_alive * (1.0 - reset_terminated.float())
    rew_termination = rew_scale_terminated * reset_terminated.float()
    rew_pole_pos = rew_scale_pole_pos * torch.sum(torch.square(pole_pos), dim=-1)
    rew_cart_vel = rew_scale_cart_vel * torch.sum(torch.abs(cart_vel), dim=-1)
    rew_pole_vel = rew_scale_pole_vel * torch.sum(torch.abs(pole_vel), dim=-1)
    total_reward = rew_alive + rew_termination + rew_pole_pos + rew_cart_vel + rew_pole_vel
    return total_reward

定义观察#

观察缓冲区应该在 _get_observations(self) 函数中计算，它为环境构造了观察缓冲区。在这个API的结尾，应该返回一个包含 policy 为键和完整观察缓冲区为值的字典。对于不对称的策略，字典还应该包括 critic 键和状态缓冲区作为值。

    def _get_observations(self) -> dict:
        obs = torch.cat(
            (
                self.joint_pos[:, self._pole_dof_idx[0]].unsqueeze(dim=1),
                self.joint_vel[:, self._pole_dof_idx[0]].unsqueeze(dim=1),
                self.joint_pos[:, self._cart_dof_idx[0]].unsqueeze(dim=1),
                self.joint_vel[:, self._cart_dof_idx[0]].unsqueeze(dim=1),
            ),
            dim=-1,
        )
        observations = {"policy": obs}
        return observations

计算终止和执行重置#

填充 dones 缓冲区应该在 _get_dones(self) 方法中完成。这个方法可以自由实现逻辑，计算哪些环境需要重置，哪些环境已经达到了本集长度限制。这两个结果都应该作为 _get_dones(self) 函数的返回值，以布尔张量的形式返回。

    def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]:
        self.joint_pos = self.cartpole.data.joint_pos
        self.joint_vel = self.cartpole.data.joint_vel

        time_out = self.episode_length_buf >= self.max_episode_length - 1
        out_of_bounds = torch.any(torch.abs(self.joint_pos[:, self._cart_dof_idx]) > self.cfg.max_cart_pos, dim=1)
        out_of_bounds = out_of_bounds | torch.any(torch.abs(self.joint_pos[:, self._pole_dof_idx]) > math.pi / 2, dim=1)
        return out_of_bounds, time_out

一旦计算出需要重置的环境的索引， _reset_idx(self, env_ids) 函数将在这些环境上执行重置操作。在这个函数内部，应该直接将需要重置的环境的新状态设置到模拟中。

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

        joint_pos = self.cartpole.data.default_joint_pos[env_ids]
        joint_pos[:, self._pole_dof_idx] += sample_uniform(
            self.cfg.initial_pole_angle_range[0] * math.pi,
            self.cfg.initial_pole_angle_range[1] * math.pi,
            joint_pos[:, self._pole_dof_idx].shape,
            joint_pos.device,
        )
        joint_vel = self.cartpole.data.default_joint_vel[env_ids]

        default_root_state = self.cartpole.data.default_root_state[env_ids]
        default_root_state[:, :3] += self.scene.env_origins[env_ids]

        self.joint_pos[env_ids] = joint_pos
        self.joint_vel[env_ids] = joint_vel

        self.cartpole.write_root_pose_to_sim(default_root_state[:, :7], env_ids)
        self.cartpole.write_root_velocity_to_sim(default_root_state[:, 7:], env_ids)
        self.cartpole.write_joint_state_to_sim(joint_pos, joint_vel, None, env_ids)

应用动作#

有两个设计用于处理动作的API。 _pre_physics_step(self, actions) 以来于策略的动作作为参数，每个RL步骤只调用一次，在采取任何物理步骤之前。这个函数可以用来处理来自策略的动作缓冲区，并将数据缓存到环境的类变量中。

    def _pre_physics_step(self, actions: torch.Tensor) -> None:
        self.actions = self.action_scale * actions.clone()

_apply_action(self) API被称为 decimation 次数，每个RL步骤之前调用一次，在采取每个物理步骤之前。这为环境提供了更多的灵活性，可以对需要为每个物理步骤应用动作的环境。

    def _apply_action(self) -> None:
        self.cartpole.set_joint_effort_target(self.actions, joint_ids=self._cart_dof_idx)

代码执行#

要运行直接工作流Cartpole环境的训练，我们可以使用以下命令：

./isaaclab.sh -p source/standalone/workflows/rl_games/train.py --task=Isaac-Cartpole-Direct-v0
`train.py` 的结果

所有直接工作流任务的名称都添加了后缀 -Direct 以区分实现风格。

域随机化#

在直接工作流中，域随机化配置使用 configclass 模块来指定包含 EventTermCfg 变量的配置类。

以下是域随机化的配置类示例：

@configclass
class EventCfg:
  robot_physics_material = EventTerm(
      func=mdp.randomize_rigid_body_material,
      mode="reset",
      params={
          "asset_cfg": SceneEntityCfg("robot", body_names=".*"),
          "static_friction_range": (0.7, 1.3),
          "dynamic_friction_range": (1.0, 1.0),
          "restitution_range": (1.0, 1.0),
          "num_buckets": 250,
      },
  )
  robot_joint_stiffness_and_damping = EventTerm(
      func=mdp.randomize_actuator_gains,
      mode="reset",
      params={
          "asset_cfg": SceneEntityCfg("robot", joint_names=".*"),
          "stiffness_distribution_params": (0.75, 1.5),
          "damping_distribution_params": (0.3, 3.0),
          "operation": "scale",
          "distribution": "log_uniform",
      },
  )
  reset_gravity = EventTerm(
      func=mdp.randomize_physics_scene_gravity,
      mode="interval",
      is_global_time=True,
      interval_range_s=(36.0, 36.0),  # time_s = num_steps * (decimation * dt)
      params={
          "gravity_distribution_params": ([0.0, 0.0, 0.0], [0.0, 0.0, 0.4]),
          "operation": "add",
          "distribution": "gaussian",
      },
  )

每个 EventTerm 对象都是 EventTermCfg 类的对象，并接受一个 func 参数，用于指定随机化时调用的函数，一个 mode 参数，可以是 startup、reset 或 interval 。 params 字典应该提供在 func 参数中指定的函数所需的参数。作为 EventTerm 的 func 指定的函数可以在 events 模块中找到。

请注意，作为 "asset_cfg": SceneEntityCfg("robot", body_names=".*") 参数的一部分，提供了actor "robot" 的名称，以及以正则表达式指定的身体或关节名称，这些名称将是应用了随机化的演员和身体/关节。

一旦设置了随机化术语的 configclass ，必须将该类添加到任务的基本配置类中，并分配给变量 events 。

@configclass
class MyTaskConfig:
  events: EventCfg = EventCfg()

动作和观察噪声#

也可以使用 configclass 模块添加动作和观察噪声配置。必须将动作和观察噪声配置添加到主任务配置中，使用 action_noise_model 和 observation_noise_model 变量：

@configclass
class MyTaskConfig:

    # at every time-step add gaussian noise + bias. The bias is a gaussian sampled at reset
    action_noise_model: NoiseModelWithAdditiveBiasCfg = NoiseModelWithAdditiveBiasCfg(
      noise_cfg=GaussianNoiseCfg(mean=0.0, std=0.05, operation="add"),
      bias_noise_cfg=GaussianNoiseCfg(mean=0.0, std=0.015, operation="abs"),
    )

    # at every time-step add gaussian noise + bias. The bias is a gaussian sampled at reset
    observation_noise_model: NoiseModelWithAdditiveBiasCfg = NoiseModelWithAdditiveBiasCfg(
      noise_cfg=GaussianNoiseCfg(mean=0.0, std=0.002, operation="add"),
      bias_noise_cfg=GaussianNoiseCfg(mean=0.0, std=0.0001, operation="abs"),
    )

NoiseModelWithAdditiveBiasCfg 可以用来对每个步骤进行无关噪声的采样，以及在重置时重新采样的相关噪声。

noise_cfg 术语指定了将在每一步中为所有环境进行采样的高斯分布。这个噪声将被添加到相应的动作和观察缓冲区中的每一步。

bias_noise_cfg 术语指定了在重置时为被重置的环境进行采样的相关噪声的高斯分布。对于剩余的环境，这个相同的噪声将在每一步应用，并在下一次重置时重新采样。

如果只希望每步有噪声，可以使用 GaussianNoiseCfg 来指定一个会将采样噪声添加到输入缓冲区中的混合高斯分布。

@configclass
class MyTaskConfig:
  action_noise_model: GaussianNoiseCfg = GaussianNoiseCfg(mean=0.0, std=0.05, operation="add")

在本教程中，我们学习了如何为强化学习创建直接工作流任务环境。我们通过扩展基础环境来包括场景设置、动作、终止、重置、奖励和观察函数来实现这一点。

虽然可以手动创建所需任务的 DirectRLEnv 类的实例，但这种方法不具有可扩展性，因为它需要为每个任务编写专门的脚本。因此，在接下来的教程中，我们将利用 gymnasium.make() 函数来使用gym接口创建环境。

目录