环境设计

环境设计#

凭借我们对项目及其结构的理解，我们准备开始修改代码以适应我们的 Jetbot 训练需求。我们的模板设置为 直接 工作流程，这意味着环境类将集中管理所有这些细节。我们需要编写代码来…

  1. 定义机器人

  2. 定义训练模拟并管理克隆

  3. 将智能体的动作应用于机器人

  4. 计算并返回奖励和观测

  5. 管理重置和终止状态

作为第一步，我们的目标将是让环境训练流程加载并运行。在本教程的这一部分，我们将使用一个虚假的奖励信号。你可以在 这里 找到这些修改的代码!

定义机器人#

随着我们项目的壮大，我们可能会有许多想要训练的机器人。我们预谋添加一个名为 robots 的新 module 到我们的教程 extension 中，我们将在其中保留机器人的定义作为单独的 python 脚本。导航到 isaac_lab_tutorial/source/isaac_lab_tutorial/isaac_lab_tutorial ，创建一个名为 robots 的新文件夹。在此文件夹中创建两个文件: __init__.py 和 jetbot.py 。 __init__.py 文件将该目录标记为 python 模块，我们将能够以常规方式导入 jetbot.py 的内容。

jetbot.py 的内容非常简单

import isaaclab.sim as sim_utils
from isaaclab.assets import ArticulationCfg
from isaaclab.actuators import ImplicitActuatorCfg
from isaaclab.utils.assets import ISAAC_NUCLEUS_DIR

JETBOT_CONFIG = ArticulationCfg(
    spawn=sim_utils.UsdFileCfg(usd_path=f"{ISAAC_NUCLEUS_DIR}/Robots/Jetbot/jetbot.usd"),
    actuators={"wheel_acts": ImplicitActuatorCfg(joint_names_expr=[".*"], damping=None, stiffness=None)},
)

此文件的唯一目的是定义一个唯一的范围，用来保存我们的配置。有关机器人配置的详细信息可以在 这个教程 中探索，但对于本教程而言，最值得关注的是 ArticulationCfg 中 spawn 参数的 usd_path 。 Jetbot 资源可通过托管的nucleus服务器公开获得，并且该路径由 ISAAC_NUCLEUS_DIR 定义，但任何到 USD 文件的路径都是有效的，包括本地路径!

环境配置#

导航到环境配置， isaac_lab_tutorial/source/isaac_lab_tutorial/isaac_lab_tutorial/tasks/direct/isaac_lab_tutorial/isaac_lab_tutorial_env_cfg.py ，并用以下内容替换其内容

from isaac_lab_tutorial.robots.jetbot import JETBOT_CONFIG

from isaaclab.assets import ArticulationCfg
from isaaclab.envs import DirectRLEnvCfg
from isaaclab.scene import InteractiveSceneCfg
from isaaclab.sim import SimulationCfg
from isaaclab.utils import configclass

@configclass
class IsaacLabTutorialEnvCfg(DirectRLEnvCfg):
    # env
    decimation = 2
    episode_length_s = 5.0
    # - spaces definition
    action_space = 2
    observation_space = 3
    state_space = 0
    # simulation
    sim: SimulationCfg = SimulationCfg(dt=1 / 120, render_interval=decimation)
    # robot(s)
    robot_cfg: ArticulationCfg = JETBOT_CONFIG.replace(prim_path="/World/envs/env_.*/Robot")
    # scene
    scene: InteractiveSceneCfg = InteractiveSceneCfg(num_envs=100, env_spacing=4.0, replicate_physics=True)
    dof_names = ["left_wheel_joint", "right_wheel_joint"]

这里，我们实际上有与之前相同的环境配置，但是使用 Jetbot 替换了 cartpole。参数 decimation 、 episode_length_s 、 action_space 、 observation_space 和 state_space 都是基类 DirectRLEnvCfg 的成员，并且必须为每个 DirectRLEnv 进行定义。空间参数被解释为给定整数维度的向量，但也可以定义为 gymnasium spaces!

请注意动作和观测空间的差异。作为环境的设计者，我们可以选择这些。对于 Jetbot，我们希望直接控制机器人的关节，其中只有两个被驱动（因此动作空间为两个）。观测空间选择为 3，因为我们现在只是要将 Jetbot 的线速度提供给智能体。随着环境的发展，稍后我们会更改这些。我们的策略不需要维护内部状态，所以我们的状态空间为零。

克隆的攻击#

配置定义完成后，是时候填写环境的细节，从初始化和设置开始。导航到环境定义， isaac_lab_tutorial/source/isaac_lab_tutorial/isaac_lab_tutorial/tasks/direct/isaac_lab_tutorial/isaac_lab_tutorial_env.py ，并用以下内容替换 __init__ 和 _setup_scene 方法的内容。

class IsaacLabTutorialEnv(DirectRLEnv):
    cfg: IsaacLabTutorialEnvCfg

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

        self.dof_idx, _ = self.robot.find_joints(self.cfg.dof_names)

    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)
        # 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)

请注意， _setup_scene 方法没有变化，而 _init__ 方法只是从机器人获取关节索引（记住，setup 是在 super 中调用的）。

我们的环境需要的下一步是处理动作、观测和奖励的定义。首先，用以下内容替换 _pre_physics_step 和 _apply_action 的内容。

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

def _apply_action(self) -> None:
    self.robot.set_joint_velocity_target(self.actions, joint_ids=self.dof_idx)

这里将将动作应用于环境中的机器人分为两步: _pre_physics_step 和 _apply_action 。物理模拟相对于查询用于动作的策略被降低，这意味着在策略采取的每个动作之间可能发生多次物理步骤。 _pre_physics_step 方法在实际模拟步骤发生之前调用，让我们分离出从正在训练的策略获取数据并应用更新到物理模拟的过程。 _apply_action 方法是将这些动作应用于场景中的机器人的地方，之后模拟会向前推进时间。

接下来是观测和奖励，这将取决于 Jetbot 在机器人本体坐标系中的线速度。用以下内容替换 _get_observations 和 _get_rewards 的内容。

def _get_observations(self) -> dict:
    self.velocity = self.robot.data.root_com_lin_vel_b
    observations = {"policy": self.velocity}
    return observations

def _get_rewards(self) -> torch.Tensor:
    total_reward = torch.linalg.norm(self.velocity, dim=-1, keepdim=True)
    return total_reward

机器人在 Isaac Lab API 中是作为 Articulation 对象存在的。该对象包含一个数据类，即 ArticulationData ，其中包含阶段上 特定 机器人的所有数据。当我们谈论一个像机器人这样的场景实体时，我们可以是广义上的讨论每个场景中存在的机器人实体，也可以是描述舞台上特定单个机器人的。 ArticulationData 包含这些单个克隆的数据。这包括诸如各种运动学向量（如 root_com_lin_vel_b ）和参考向量（如 robot.data.FORWARD_VEC_B ）等。

请注意，在 _apply_action 方法中，我们调用了 self.robot 的方法，这是 Articulation 的一个方法。被应用的动作以 [num_envs, num_actions] 形状的 2D 张量的形式存在。我们同时将动作应用于 所有 舞台上的机器人!在这里，当我们需要获取观测时，我们需要所有舞台上机器人的本体系速度，因此访问 self.robot.data 以获取这些信息。 root_com_lin_vel_b 是 ArticulationData 的一个属性，处理了质心线性速度从世界坐标系到本体坐标系的转换。最后，Isaac Lab 期望观测以字典形式返回，其中 policy 为策略模型定义观测， critic 为评论模型定义观测（在非对称的 actor critic 训练中）。由于我们不执行非对称的 actor critic，我们只需要定义 policy 。

奖励则更为直接。对于每个场景的克隆，我们需要计算一个奖励值，并将其作为形状为 [num_envs, 1] 的张量返回。作为占位符，我们将将奖励设定为 Jetbot 在本体坐标系中的线性速度的大小。有了这个奖励和观测空间，智能体应该会学会在训练开始后不久就随机确定方向，将 Jetbot 往前或往后驱动。

最后，我们可以编写环境的部分来处理终止和重置。 用以下内容替换 _get_dones 和 _reset_idx 的内容。

def _get_dones(self) -> tuple[torch.Tensor, torch.Tensor]:
    time_out = self.episode_length_buf >= self.max_episode_length - 1

    return False, time_out

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]
    default_root_state[:, :3] += self.scene.env_origins[env_ids]

    self.robot.write_root_state_to_sim(default_root_state, env_ids)

与动作类似，终止和重置是分成两部分进行处理的。首先是 _get_dones 方法，其目标仅仅是标记哪些环境需要重置以及原因。传统上，在强化学习中，一个 “episode” 以两种方式结束: 要么是智能体达到一个终止状态，要么是达到最大持续时间。Isaac Lab 对我们很好，因为它在幕后管理所有这些周期长度跟踪。配置参数 episode_length_s 定义了这个最大的 episode 长度（以秒为单位），参数 episode_length_buff 和 max_episode_length 包含单个场景所采取的步数（允许环境异步运行）以及从 episode_length_s 转换过来的最大 episode 长度。计算 time_out 的布尔运算仅仅比较当前缓冲区大小和最大值，如果大于最大值，则返回 true，从而指示哪些场景达到了 episode 长度限制。由于我们当前的环境是一个虚拟的环境，我们没有定义终止状态，所以对于第一个张量返回 False （PyTorch 的强大之处使其自动转换成正确的形状）。

最后， _reset_idx 方法接受一个指示哪些场景需要重置的布尔张量，并对它们进行重置。请注意，这是 DirectRLEnv 的仅存在的另一个直接调用 super 的方法，这样做是为了管理与 episode 长度相关的内部缓冲。对于 env_ids 指示的环境，我们获取根默认状态，并将机器人重置到该状态，同时根据对应场景的原点偏移每个机器人的位置。这是克隆过程的一个结果，该过程从世界坐标系中定义的单个机器人和单个默认状态开始。不要忘记为您自己的自定义环境执行此步骤!

完成这些更改后，当您使用模板 train.py 脚本启动任务时，您应该会看到 Jetbot 慢慢学会向前驾驶。

Jetbot 入侵开始!