Legged gym 使用笔记

入门案例
1.倒立摆

首先在 envs 目录下创建自己的环境 cartpole

这篇主要是看自定义 Cartpole2Task 环境时做的源码笔记。整体流程大致是这样:

  1. envs 目录下放自己的任务代码和 config,比如 cartpole2.pycartpole2_config.py

  2. envs/__init__.py 中注册任务名、环境类、环境配置和 PPO 配置。训练脚本通过任务名找到对应环境。

  3. 环境类继承 BaseTask,核心是实现仿真创建、step、reset、reward、observation 这些函数。

  4. 初始化时先解析配置,再调用父类创建 sim/env,随后初始化 PyTorch buffer,最后准备奖励函数列表。

  5. 训练时算法会不断调用 env.step(actions)。一次 step 里先裁剪 action,再计算 torque,接着推进物理仿真,刷新状态,然后检查终止条件,计算奖励,重置需要结束的环境,最后重新计算 observation 并返回给算法。

下面按函数继续记录。

导入

LEGGED_GYM_ROOT_DIR 是项目根目录路径,通常用于拼 URDF、资源文件路径。envs 文件夹同级下的 __init__.py 会注册任务。Cartpole2Config 是环境配置文件,Cartpole2ConfigPPO 是训练配置文件。

常用导入大致如下:

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
from legged_gym import LEGGED_GYM_ROOT_DIR, envs
from time import time
from warnings import WarningMessage
import numpy as np
import os

from isaacgym.torch_utils import *
from isaacgym import gymtorch, gymapi, gymutil

import torch

# from legged_gym import LEGGED_GYM_ROOT_DIR
from legged_gym.envs.base.base_task import BaseTask
from legged_gym.utils.helpers import class_to_dict #类转化为字典
from .cartpole2_config import Cartpole2Config

其中 class_to_dict 是将配置类的属性转成字典,后面处理 reward scales 时会用到。

此文件中只定义了一个类,Cartpole2Task 继承于 BaseTask

这个倒立摆例子比真正的四足环境简单,大概包含 19 个方法。原始 legged_gym 里的腿足机器人任务会复杂得多,除了基础 step/reset/reward,还会有地形、命令采样、课程学习、接触检测、域随机化等逻辑。

这篇主要看三块:基础生命周期、环境初始化、奖励函数。

注册入口

自定义环境写完以后,需要在 envs/__init__.py 里注册。训练脚本不是直接 import 某个环境文件,而是通过 task name 到 task_registry 里找环境类和配置类。

大概形式如下:

1
2
3
4
5
from legged_gym.envs.cartpole2.cartpole2 import Cartpole2Task
from legged_gym.envs.cartpole2.cartpole2_config import Cartpole2Config, Cartpole2ConfigPPO
from legged_gym.utils.task_registry import task_registry

task_registry.register("cartpole2", Cartpole2Task, Cartpole2Config, Cartpole2ConfigPPO)

这里的 "cartpole2" 就是训练时传给脚本的任务名。后面执行训练时,runner 会根据这个名字创建环境、读取 PPO 配置,再进入 env.step(actions) 的循环。

这一层如果没注册,或者名字拼错,环境类本身写对了也没用,训练入口找不到它。

初始化

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
def __init__(self, cfg: Cartpole2Config, sim_params, physics_engine, sim_device, headless):
"""Parses the provided config file,
calls create_sim() (which creates, simulation, terrain and environments),
initializes PyTorch buffers used during training

Args:
cfg (Dict): Environment config file
sim_params (gymapi.SimParams): simulation parameters
physics_engine (gymapi.SimType): gymapi.SIM_PHYSX (must be PhysX)
device_type (string): 'cuda' or 'cpu'
device_id (int): 0, 1, ...
headless (bool): Run without rendering if True
"""
self.cfg = cfg # 主要配置文件 这应该是个类
self.sim_params = sim_params # 仿真参数配置
self.init_done = False # 是否初始化完成
self._parse_cfg(self.cfg) # 解析配置文件
super().__init__(self.cfg, sim_params, physics_engine, sim_device, headless)
# 初始化父类,和C++不一样,父类的初始化不一定在最开始

if not self.headless: # 设置渲染窗口的相机初始化位置
self.set_camera(self.cfg.viewer.pos, self.cfg.viewer.lookat)
self._init_buffers() # 初始化 PyTorch buffer
self._prepare_reward_function() # 准备奖励函数
self.init_done = True # 初始化完成

cfg: Cartpole2Config是cfg类型的预期类是Cartpole2Config类
解析完参数之后调用父类的初始化方法进行初始化
self.set_camera(self.cfg.viewer.pos, self.cfg.viewer.lookat)是生成相机窗口界面

这里的执行顺序是:self._parse_cfg(self.cfg)BaseTask.__init__() 之前执行,用来先算出 dt、episode 长度、reward scales 等后面会用到的配置。BaseTask.__init__() 里会调用子类实现的 create_sim(),所以子类要提前把创建仿真需要的配置准备好。

初始化过程大致按下面的顺序执行:

1
2
3
4
5
6
7
8
Cartpole2Task.__init__
1. _parse_cfg()
2. BaseTask.__init__()
2.1 create_sim()
2.2 _create_ground_plane()
2.3 _create_envs()
3. _init_buffers()
4. _prepare_reward_function()

仿真步骤

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
def step(self, actions):
"""Apply actions, simulate, call self.post_physics_step()
这个函数是仿真的核心函数,每一步仿真都会调用这个函数.
(1) 该函数首先接受由策略提供的action,并计算出对应的torque;
(2) 然后将torque应用到仿真中,进行一次仿真;
(3) 仿真完成后,调用post_physics_step()函数,检查终止条件,计算奖励,更新观测等。
(4) 最后返回观测,奖励,终止标志等信息。
强化学习仿真环境的编写主要过程包括:
应用action,仿真,计算奖励,检查终止条件,更新观测,返回信息。
Args:
actions (torch.Tensor): Tensor of shape (num_envs, num_actions_per_env)
"""
clip_actions = self.cfg.normalization.clip_actions # action的范围
self.actions = torch.clip(actions, -clip_actions, clip_actions).to(self.device)
# 限制action的范围,防止过大或过小
# step physics and render each frame
self.render() # 渲染
for _ in range(self.cfg.control.decimation):
# 计算torque,将其应用在关节上,并进行一次仿真,decimation是控制频率的参数
self.torques = self._compute_torques(self.actions).view(self.torques.shape) # 计算torque
self.gym.set_dof_actuation_force_tensor(
self.sim, gymtorch.unwrap_tensor(self.torques)
) # 将torque应用到仿真中
self.gym.simulate(self.sim) # 进行一次仿真
if self.device == "cpu":
self.gym.fetch_results(self.sim, True)
self.gym.refresh_dof_state_tensor(self.sim) # 刷新关节状态

self.post_physics_step() # 检查终止条件,计算奖励,更新观测等

# return clipped obs, clipped states (None), rewards, dones and infos
clip_obs = self.cfg.normalization.clip_observations
self.obs_buf = torch.clip(self.obs_buf, -clip_obs, clip_obs) # 对观测进行裁剪,防止过大或过小
if self.privileged_obs_buf is not None: # 本实验中没有使用到特权观测信息
self.privileged_obs_buf = torch.clip(self.privileged_obs_buf, -clip_obs, clip_obs)
return (
self.obs_buf,
self.privileged_obs_buf,
self.rew_buf,
self.reset_buf,
self.extras,
)

先对 action 范围进行裁剪。
decimation用来桥接物理仿真频率和控制频率
.view(self.torques.shape)改变张量形状
gymtorch.unwrap_tensor(tensor) 返回 Isaac Gym 底层接口能识别的 tensor 描述,可以理解为把 PyTorch Tensor 交给 Isaac Gym C++/CUDA 侧使用。这里不是复制一份数据,而是让 Isaac Gym 能直接访问对应 GPU/CPU 内存。
gymtorch.unwrap_tensor(torques) # 转换为 PhysX引擎可识别的格式
在 CPU 模式下,强制等待物理仿真完成并获取最新状态。
simulate执行并行GPU仿真,比如同时开1024个环境,然后self.gym.refresh_dof_state_tensor(self.sim)这里会自动等待所有环境这一步完成后,一起更新状态
吞吐量指的是​​每秒处理的环境步数​ ​吞吐量 = (环境数量 × 每秒每环境步数)
接下来完成观测值裁减和返回训练所需数据
​特权观测Privileged Observations)​​:仅在特定任务中使用的额外信息(如环境物理参数)。
self.rew_buf 奖励值
self.reset_buf, 终止标志
self.extras, 额外信息

后处理步骤
这里好像只刷新了关节状态,没有
self.gym.refresh_actor_root_state_tensor(self.sim)
self.gym.refresh_net_contact_force_tensor(self.sim)
self.gym.refresh_rigid_body_state_tensor(self.sim)

decimation 的含义比较关键。策略网络每输出一次 action,物理仿真会推进 decimation 个小步。假设 sim_params.dt = 0.005decimation = 4,那么控制周期就是 0.02s,也就是 50Hz。self.dt = self.cfg.control.decimation * self.sim_params.dt 后面算 episode 长度和 reward scale 都会用到。

这里的刷新函数要看观测需要什么。这个 Cartpole 例子只用关节位置和速度,所以刷新 dof_state_tensor 就够了。如果观测里需要 base pose、刚体位置、接触力,就要刷新对应 tensor,例如 root state、rigid body state、net contact force。四足环境里通常会刷新更多状态。

仿真后处理

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def post_physics_step(self):
"""check terminations, compute observations and rewards
calls self._post_physics_step_callback() for common computations
calls self._draw_debug_vis() if needed
"""
# in some cases a simulation step might be required to refresh some obs (for example body positions)
self.gym.refresh_dof_state_tensor(self.sim) # 刷新关节状态

self.episode_length_buf += 1 # 记录每个episode的长度
self.common_step_counter += 1 # 记录总的步数

# compute observations, rewards, resets, ...
self.check_termination() # 检查是否需要终止
self.compute_reward() # 计算奖励
env_ids = self.reset_buf.nonzero(as_tuple=False).flatten() # 从reset buffer中找到需要重置的环境的id
self.reset_idx(env_ids) # 重置一些环境
self.compute_observations() # 计算观测

.flatten把向量展平tensor([[1], [3]])展成tensor([1,3])

这个函数名字是 post_physics_step,更准确地说是“物理仿真之后的处理”,不是仿真前处理。它负责把这一轮仿真结果转换成 RL 算法要的东西:是否终止、奖励是多少、哪些环境要 reset、下一步 observation 是什么。

检查终止条件

1
2
3
4
5
6
7
8
9
10
11
12
def check_termination(self):
"""Check if environments need to be reset"""
self.reset_buf = torch.zeros_like(self.episode_length_buf) # 重置buffer
self.time_out_buf = self.episode_length_buf > self.max_episode_length
self.reset_buf |= self.time_out_buf # 如果超过最大episode长度,需要重置

cart_pos = self.dof_pos[:, 0] # 获取小车位置
pole_angle = self.dof_pos[:, 1] # 获取杆1角度
pole_angle2 = self.dof_pos[:, 2] # 获取杆2角度
self.reset_buf |= torch.abs(cart_pos) > 3.0 # 如果小车位置超过3,需要重置
self.reset_buf |= torch.abs(pole_angle) > torch.pi / 2 # 如果杆1角度超过90度,需要重置
self.reset_buf |= torch.abs(pole_angle2) > torch.pi / 2 # 如果杆2角度超过90度,需要重置

episode_length_buf是个[1024]的张量,张量+1,是对每个元素都加1

reset_buf 是按环境编号排列的终止标志,形状通常是 [num_envs]。并行环境里不是所有环境同时结束,所以每一步都要为每个环境单独判断是否 reset。|= 是按位或赋值,只要某个终止条件成立,对应环境的 reset_buf 就会被置为 True。

重置序号

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
def reset_idx(self, env_ids):
"""Reset some environments.
Calls self._reset_dofs(env_ids), self._reset_root_states(env_ids), and self._resample_commands(env_ids)
[Optional] calls self._update_terrain_curriculum(env_ids), self.update_command_curriculum(env_ids) and
Logs episode info
Resets some buffers

Args:
env_ids (list[int]): List of environment ids which must be reset
"""
if len(env_ids) == 0: # 如果没有需要重置的环境,直接返回
return

# reset robot states
self._reset_dofs(env_ids) # 重置关节位置和速度

# reset buffers
self.episode_length_buf[env_ids] = 0 # 重置episode长度
self.reset_buf[env_ids] = 1 # 设置buffer,后续policy需要知道哪些环境被重置了

# fill extras,记录一些奖励相关的信息,供log使用
self.extras["episode"] = {}
for key in self.episode_sums.keys():#不同种类奖励(平衡、速度)的sum
self.extras["episode"]["rew_" + key] = (
torch.mean(self.episode_sums[key][env_ids]) / self.max_episode_length_s
)
self.episode_sums[key][env_ids] = 0.0

# send timeout info to the algorithm
if self.cfg.env.send_timeouts:
self.extras["time_outs"] = self.time_out_buf

计算平均奖励

更新后产生这样的数组
extras[“episode”] = {“rew_balance”: 0.55, “rew_velocity”: 0.225}
表示重置环境的平均奖励
这里有点不太准确,但是毕竟是额外的记录数组,如果不是因为到达最后的episode,则再除以最大episode的话会低估。

这个判断是有意义的。extras["episode"] 主要给 logger 使用,用来观察平均奖励趋势。当前写法用 max_episode_length_s 做归一化,如果环境因为失败提前 reset,确实可能低估这个 episode 的实际平均表现。训练本身使用的是每一步的 rew_bufextras 更多是记录和可视化信息。

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
def compute_reward(self):
"""Compute rewards
Calls each reward function which had a non-zero scale (processed in self._prepare_reward_function())
adds each terms to the episode sums and to the total reward
"""
self.rew_buf[:] = 0.0
# 依次调用 _prepare_reward_function() 收集到的奖励函数。
# cfg.rewards.scales 里有非零 scale 的项,才会进入这个列表。
for i in range(len(self.reward_functions)):
name = self.reward_names[i]
rew = self.reward_functions[i]() * self.reward_scales[name]
self.rew_buf += rew
self.episode_sums[name] += rew

# 同样是计算奖励,这里计算的是智能体被重置时才获取的额外奖励/惩罚
if "termination" in self.reward_scales:
rew = self._reward_termination() * self.reward_scales["termination"]
self.rew_buf += rew
self.episode_sums["termination"] += rew

rew_buf是所有环境的奖励列表[第一个环境,第二个环境,…]这个是即时奖励缓存区,是当前步仿真所得到的奖励,shape 是 [num_envs]
self.episode_sums是回合奖励累计器,记录当前episode的累积奖励。它不是一个总数,而是按奖励项分开的字典,比如 episode_sums["joint_angle"] 也是 [num_envs]

compute_reward() 先把当前步奖励清零,然后按 _prepare_reward_function() 准备好的函数列表逐项计算奖励,再乘以对应 scale 加到 rew_buf。同时,episode_sums 会累计每个环境当前 episode 内各个奖励项的总和,reset 时再写进 extras

这里有一个容易混的点:奖励函数本身通常返回“未乘权重”的值,正负号和大小主要由 cfg.rewards.scales 控制。比如某个 penalty 函数返回角度平方,函数值是正的,但 scale 如果是负数,最后加到 rew_buf 里就是惩罚项。这样做的好处是 reward 结构写在代码里,权重调参放在 config 里。

计算观测

1
2
3
4
5
6
7
8
9
def compute_observations(self):
"""Computes observations"""
self.obs_buf = torch.cat(
(
(self.dof_pos - self.default_dof_pos) * self.obs_scales.dof_pos, # 关节位置
self.dof_vel * self.obs_scales.dof_vel, # 关节速度
),
dim=-1,
)

状态观测,把传感器信息处理成策略网络可用的输入向量
(self.dof_pos - self.default_dof_pos) * self.obs_scales.dof_pos, # 关节位置 角度减去零位角度乘以缩放值
速度乘以缩放
torch.cat( (pos_obs, vel_obs), dim=-1 )沿最后一个维度(列方向)拼接。

观测的 shape 一般是 [num_envs, num_obs]。这里 dof_posdof_vel 都是 [num_envs, num_dof],沿最后一维拼接后就是 [num_envs, num_dof * 2]。如果以后要加命令、上一帧 action、接触状态等,也是在最后一维继续拼接。

obs_scales 的作用是把不同物理量缩放到比较接近的数值范围。强化学习里 observation 尺度差太大,会让策略网络更难训练。

dim=-1 表示沿最后一维拼接。对二维张量来说,可以理解成在特征维度上拼接;dim=0 则是沿 batch 维拼接,一般不是 observation 这里想要的。

self.create_sim() 需要子类自己实现,BaseTask.__init__() 里会调用它。这里用的是 Python 的方法重写:父类初始化过程中调用 self.create_sim() 时,实际执行的是子类里的 create_sim()。所以子类的 _parse_cfg() 必须放在 super().__init__() 之前,不然创建仿真时可能拿不到配置。

创建仿真环境

1
2
3
4
5
6
7
8
9
10
def create_sim(self):  #这个是初始化的时候调用父类初始化时所调用的方法
"""Creates simulation, terrain and environments"""
self.sim = self.gym.create_sim(
self.sim_device_id,
self.graphics_device_id,
self.physics_engine,
self.sim_params,
) # 创建仿真器handle
self._create_ground_plane() # 创建地面
self._create_envs() # 创建智能体

创建相机位置向量

创建观察目标向量 应用相机设置

创建相机,施加力矩

1
2
3
4
5
def set_camera(self, position, lookat):
"""Set camera position and direction"""
cam_pos = gymapi.Vec3(position[0], position[1], position[2])
cam_target = gymapi.Vec3(lookat[0], lookat[1], lookat[2])
self.gym.viewer_camera_look_at(self.viewer, None, cam_pos, cam_target) # 设置渲染器初始位姿
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
def _compute_torques(self, actions):
"""Compute torques from actions.
Actions can be interpreted as position or velocity targets given to a PD controller, or directly as scaled torques.
[NOTE]: torques must have the same dimension as the number of DOFs, even if some DOFs are not actuated.

Args:
actions (torch.Tensor): Actions

Returns:
[torch.Tensor]: Torques sent to the simulation
"""
# 根据policy输出计算力矩,这里仅简单按比例直出了,训练机器人时可以加一些控制器如pd等
actions_scaled = actions * self.cfg.control.action_scale
torques = torch.zeros_like(self.torques)
torques[:, 0:1] = actions_scaled

return torques

这里是网络输出完action后,输入给电机控制器,然后电机控制器输出实际的力矩,这里充当了电机控制的角色,可以用pd控制
这里先缩放(config 文件里定义)。torques[:, 0:1] = actions_scaled 表示所有环境的第 0 个关节被 action 驱动,不是“前两个关节”。Python 切片 0:1 只包含索引 0,不包含 1。这个例子里小车关节是可驱动的,两个杆关节是无驱动的。

关节状态重置,将所指的id状态重置

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
def _reset_dofs(self, env_ids):
"""Resets DOF position and velocities of selected environments
Positions are randomly selected within 0.5:1.5 x default positions.
Velocities are set to zero.

Args:
env_ids (List[int]): Environment ids
"""
positions = 0.2 * (torch.rand((len(env_ids), self.num_dof), device=self.device) - 0.5)
velocities = 0.5 * (torch.rand((len(env_ids), self.num_dof), device=self.device) - 0.5)
self.dof_pos[env_ids, :] = positions[:]
self.dof_vel[env_ids, :] = velocities[:] # 复位时随机指定位置和速度

env_ids_int32 = env_ids.to(dtype=torch.int32)
self.gym.set_dof_state_tensor_indexed(
self.sim,
gymtorch.unwrap_tensor(self.dof_state),
gymtorch.unwrap_tensor(env_ids_int32),
len(env_ids_int32),
) # 将新关节位置应用到仿真中


首先缩放到-0.1到0.1范围内的平均数
然后生成随机速度,或者速度可以为0
unwrap_tensor还是将PyTorch张量和底层物理引擎PhysX之间建立数据链接的通道

数组初始化

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
def _init_buffers(self):
"""Initialize torch tensors which will contain simulation states and processed quantities"""
# get gym GPU state tensors

dof_state_tensor = self.gym.acquire_dof_state_tensor(self.sim) # 从仿真器获取关节状态的矩阵

self.gym.refresh_dof_state_tensor(self.sim) # 刷新关节状态

# create some wrapper tensors for different slices
self.dof_state = gymtorch.wrap_tensor(dof_state_tensor) # 将关节状态矩阵包装成 PyTorch tensor
self.dof_pos = self.dof_state.view(self.num_envs, self.num_dof, 2)[..., 0]
self.dof_vel = self.dof_state.view(self.num_envs, self.num_dof, 2)[..., 1] # 将关节状态矩阵分成位置和速度两部分

# the following tensors are NOT used in this experiment, but are kept for compatibility with other tasks
self.commands = torch.zeros(self.num_envs, 3, dtype=torch.float, device=self.device, requires_grad=False)
self.base_lin_vel = torch.zeros(self.num_envs, 3, dtype=torch.float, device=self.device, requires_grad=False)
self.base_ang_vel = torch.zeros(self.num_envs, 3, dtype=torch.float, device=self.device, requires_grad=False)
self.contact_forces = torch.zeros(
self.num_envs, 2, 3, dtype=torch.float, device=self.device, requires_grad=False
)
self.feet_indices = torch.tensor([0, 1], dtype=torch.int32, device=self.device, requires_grad=False)

# initialize some data used later on
self.common_step_counter = 0 # 记录总的步数
self.extras = {} # 用于记录一些额外信息

self.torques = torch.zeros( # 力矩(直接应用于关节)
self.num_envs,
self.num_dof,
dtype=torch.float,
device=self.device,
requires_grad=False,
)

self.actions = torch.zeros( # 动作(网络模型输出)
self.num_envs,
self.num_actions,
dtype=torch.float,
device=self.device,
requires_grad=False,
)

# 关节默认位置为0
self.default_dof_pos = torch.zeros(self.num_dof, dtype=torch.float, device=self.device, requires_grad=False)
self.default_dof_pos = self.default_dof_pos.unsqueeze(0)

dof_state_tensor = self.gym.acquire_dof_state_tensor(self.sim) 从物理引擎获取关节状态的原始 tensor 句柄。
关节顺序由urdf文件中定义的顺序来决定joint的先后顺序

gymtorch.wrap_tensor:将 Isaac Gym tensor 包装成 PyTorch Tensor,后续就可以用 PyTorch 的切片和张量运算处理。

[num_envs, 2, 3] 表示每个环境里有 2 个接触点/脚,每个接触力有 3 个方向分量。
self.default_dof_pos定义的是关节零位角度,所有环境都一样所以没有第一维度
self.base_lin_vel = torch.zeros(self.num_envs, 3, dtype=torch.float, device=self.device, requires_grad=False)
详细解释:dof_state是一个张量结构为:
[env0_joint0_pos, env0_joint0_vel, env0_joint1_pos, env0_joint1_vel, …,
env1_joint0_pos, env1_joint0_vel, …]
[…, 0]和[:, :, 0]等价,位置是在最后一个维度的0维,速度是在最后一个维度的1维
self.torques
获取环境0的所有关节力矩
env0_torques = self.torques[0] # 形状 [num_dof]
获取所有环境的关节0力矩
joint0_torques = self.torques[:, 0] # 形状 [num_envs]
获取环境1的关节2力矩
specific_torque = self.torques[1, 2]

self.envs是一个包含​​所有环境实例句柄​​的列表,可以通过以下代码获取环境0的所有关节的名称
for i in range(self.num_dof):
name = self.gym.get_joint_name(self.envs[0], i)
print(f”关节 {i}: {name}”)

.view()语法,返回具有新形状的张量,​​不复制数据​​(共享内存)
new_tensor = old_tensor.view(shape)
reshaped = self.dof_state.view(self.num_envs, self.num_dof, 2)
为-1的话是自动推断

这里 view(self.num_envs, self.num_dof, 2) 的最后一维 2 分别表示 position 和 velocity。[..., 0] 取所有环境、所有关节的 position;[..., 1] 取 velocity。这个写法比手动按一维数组下标取值更不容易出错。

self.commands是人输入的指令,

这些 buffer 的特点是全部按 num_envs 并行展开。不是一个环境算完再算下一个环境,而是一批环境的状态放在同一个大 tensor 里,策略网络和奖励函数也基本按 batch 一次算完。

unsqueeze(dim) 是在指定位置增加一个长度为 1 的维度。这里 default_dof_pos 原本是 [num_dof]unsqueeze(0) 后变成 [1, num_dof],后面和 [num_envs, num_dof]dof_pos 相减时,就可以通过广播机制自动扩展到所有环境。

广播机制是 PyTorch/Numpy 中自动扩展张量维度的规则。判断时从右向左对齐维度,只要两个维度相等,或者其中一个是 1,就可以逐元素运算。这里并不是手动复制了很多份默认关节角,而是 PyTorch 在运算时按广播规则处理,代码更短,也少占显式存储。

准备奖励函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
def _prepare_reward_function(self):
"""Prepares a list of reward functions, which will be called to compute the total reward.
Looks for self._reward_<REWARD_NAME>, where <REWARD_NAME> are names of all non zero reward scales in the cfg.
"""
# 根据 cfg.rewards.scales 里的名字,收集对应的 _reward_xxx 方法。
# 后面 compute_reward() 只遍历这个列表,不需要手动写一堆 if。
# remove zero scales + multiply non-zero ones by dt
for key in list(self.reward_scales.keys()):
scale = self.reward_scales[key]
if scale == 0:
self.reward_scales.pop(key)
else:
self.reward_scales[key] *= self.dt
# prepare list of functions
self.reward_functions = []
self.reward_names = []
for name, scale in self.reward_scales.items():
if name == "termination":
continue
self.reward_names.append(name)
name = "_reward_" + name
self.reward_functions.append(getattr(self, name))

# reward episode sums
self.episode_sums = {
name: torch.zeros(
self.num_envs,
dtype=torch.float,
device=self.device,
requires_grad=False,
)
for name in self.reward_scales.keys()
}

这里做了三件事。

一是移除 scale 为 0 的奖励项。scale 为 0 就说明当前配置不使用这个奖励,没必要每一步都调用它。

二是把非零 scale 乘以 self.dt。这相当于把每个控制周期内的 reward 做时间尺度补偿,避免控制频率改了以后,单位时间累计奖励跟着大幅变化。比如 decimation 改变后,策略每秒调用 reward 的次数也会变,如果不乘 dt,同一个动作轨迹的累计 reward 可能只是因为控制频率不同而变了。

三是把奖励函数收集到 self.reward_functionsgetattr(self, "_reward_" + name) 是通过字符串找对象方法,比如 config 里有 joint_angle,就会找 self._reward_joint_angle。这要求 config 里的奖励项名字必须和函数名后缀一致,否则运行到这里就会报错。

termination 被跳过,是因为终止奖励要在普通奖励之后单独处理。这样可以让终止惩罚只在 reset 条件触发时生效,不会混进普通的连续奖励项里。

最后的 episode_sums 是字典推导式。每个奖励项都会初始化一个 [num_envs] 的累计数组,方便 reset 时只统计那些刚结束的环境。

创建地平面

1
2
3
4
5
6
7
8
9
def _create_ground_plane(self):
"""Adds a ground plane to the simulation, sets friction and restitution based on the cfg."""
plane_params = gymapi.PlaneParams()
plane_params.normal = gymapi.Vec3(0.0, 0.0, 1.0) # 地面法向量为z轴
plane_params.static_friction = self.cfg.terrain.static_friction
plane_params.dynamic_friction = self.cfg.terrain.dynamic_friction
plane_params.restitution = self.cfg.terrain.restitution
# 设置地面的参数(本实验没有使用到地面的信息,这里只是出于兼容性和完整性考虑保留了)
self.gym.add_ground(self.sim, plane_params)

创建对象gymapi.PlaneParams();设置属性;添加到仿真里
normal:法线方向
static_friction:静摩擦系数
dynamic_friction:动摩擦系数
restitution:弹性恢复系数 恢复系数越高,弹性越高,

创建训练环境

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
def _create_envs(self):
"""Creates environments:
1. loads the robot URDF/MJCF asset,
2. For each environment
2.1 creates the environment,
2.2 calls DOF and Rigid shape properties callbacks,
2.3 create actor with these properties and add them to the env
3. Store indices of different bodies of the robot
"""
asset_path = self.cfg.asset.file.format(LEGGED_GYM_ROOT_DIR=LEGGED_GYM_ROOT_DIR)
asset_root = os.path.dirname(asset_path)
asset_file = os.path.basename(asset_path) # 设置URDF文件名和文件路径

asset_options = gymapi.AssetOptions() # 设置机器人模型的各种属性
# 属性的具体含义可以参考isaacgym的文档
# https://blog.zzshub.cn/legged_gym/api/python/struct_py.html#isaacgym.gymapi.AssetOptions
asset_options.default_dof_drive_mode = self.cfg.asset.default_dof_drive_mode
asset_options.collapse_fixed_joints = self.cfg.asset.collapse_fixed_joints
asset_options.replace_cylinder_with_capsule = self.cfg.asset.replace_cylinder_with_capsule
asset_options.flip_visual_attachments = self.cfg.asset.flip_visual_attachments
asset_options.fix_base_link = self.cfg.asset.fix_base_link
asset_options.density = self.cfg.asset.density
asset_options.angular_damping = self.cfg.asset.angular_damping
asset_options.linear_damping = self.cfg.asset.linear_damping
asset_options.max_angular_velocity = self.cfg.asset.max_angular_velocity
asset_options.max_linear_velocity = self.cfg.asset.max_linear_velocity
asset_options.armature = self.cfg.asset.armature
asset_options.thickness = self.cfg.asset.thickness
asset_options.disable_gravity = self.cfg.asset.disable_gravity

# 加载模型
robot_asset = self.gym.load_asset(self.sim, asset_root, asset_file, asset_options)
self.num_dof = self.gym.get_asset_dof_count(robot_asset) # 获取关节数量
self.num_bodies = self.gym.get_asset_rigid_body_count(robot_asset) # 获取刚体数量
self.dof_names = self.gym.get_asset_dof_names(robot_asset) # 获取关节名称

start_pose = gymapi.Transform()
start_pose.p = gymapi.Vec3(0.0, 0.0, 0.5) # 设置机器人初始位置

env_lower = gymapi.Vec3(0.0, 0.0, 0.0)
env_upper = gymapi.Vec3(self.cfg.env.env_spacing, self.cfg.env.env_spacing, 0.0) # 设置环境边界
self.actor_handles = []
self.envs = []
for i in range(self.num_envs):
# create env instance
# 创建环境实例, Isaac Gym是一个并行化的仿真器,所以这里会创建多个环境实例
# 每个环境有一个env_handler, 每个机器人有一个actor_handle, 每个actor_handle对应一个机器人实例
# 由于这里每个环境只有一个机器人,所以只有一个actor_handle,
# 需要多智能体组合时,可以在一个env_handle中创建多个actor_handle
# 不同 env_handle 之间是隔离的;同一个 env_handle 内多个 actor 可以发生交互
# 前面的asset是机器人模型,首先导入这个模型,并修改模型的相关参数,之后创建机器人实例

# 创建环境
env_handle = self.gym.create_env(self.sim, env_lower, env_upper, int(np.sqrt(self.num_envs)))
# 在环境中创建机器人实例
actor_handle = self.gym.create_actor(
env_handle,
robot_asset,
start_pose,
self.cfg.asset.name,
i,
self.cfg.asset.self_collisions,
0,
)

# 设置关节驱动模式
dof_props = self.gym.get_actor_dof_properties(env_handle, actor_handle) # 获取关节属性
dof_props["driveMode"][0] = gymapi.DOF_MODE_EFFORT # 设置1关节(小车关节)驱动模式为力矩驱动
dof_props["driveMode"][1] = gymapi.DOF_MODE_NONE # 设置2关节(杆1关节)驱动模式为无驱动
dof_props["driveMode"][2] = gymapi.DOF_MODE_NONE # 设置3关节(杆2关节)驱动模式为无驱动
dof_props["stiffness"][:] = 0.0
dof_props["damping"][:] = 0.0 # 设置关节pd为0
self.gym.set_actor_dof_properties(env_handle, actor_handle, dof_props) # 设置关节属性

self.envs.append(env_handle) # 记录环境handle
self.actor_handles.append(actor_handle) # 记录机器人实例handle

self.cfg.asset.file是设置的urdf的位置
string.format(name=value)语法
将字符串中的占位符 {name}替换为指定的 value

设置机器人的各种属性———可以一个一个对应去查
default_dof_drive_mode​​ 设置关节的默认驱动模式
collapse_fixed_joints​​ 是否合并固定关节

这里要区分三个 handle:

  1. sim handle:整个仿真实例。

  2. env handle:一个并行环境实例。num_envs 有多少,就会创建多少个 env handle。

  3. actor handle:某个环境里的一个机器人/物体实例。这个 Cartpole 每个 env 只有一个 actor;如果做多智能体,一个 env 里可以创建多个 actor。

env_spacing 控制不同 env 在物理空间中的间距,避免可视化时重叠。不同 env 之间通常不会互相物理交互;如果希望两个机器人交互,应该把它们放在同一个 env 里作为不同 actor。

仿真设置

1
2
3
4
5
6
7
8
def _parse_cfg(self, cfg):
# 设置仿真时间步长等参数
self.dt = self.cfg.control.decimation * self.sim_params.dt
self.obs_scales = self.cfg.normalization.obs_scales
self.reward_scales = class_to_dict(self.cfg.rewards.scales) #将属性转化为字典
self.max_episode_length_s = self.cfg.env.episode_length_s
self.max_episode_length = np.ceil(self.max_episode_length_s / self.dt)

np.ceil语法:向正无穷方向取整
np.floor:向负无穷方向取整
np.trunc:向零取整
np.round:最近整数

奖励函数

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
def _reward_joint_angle(self):
# 设置关节角度奖励,这里是一个简单的奖励函数,
# 奖励为关节角度的平方,这里的奖励为正,但是乘上scales后为负,
# 关节角度越大,奖励越小
pole_angle = self.dof_pos[:, 1] #杆0关节角度是车的位置,不care
pole_angle2 = self.dof_pos[:, 2]
reward = pole_angle * pole_angle + pole_angle2 * pole_angle2 * 2
return reward

def _reward_joint_velocity(self):
# 设置关节速度奖励,这里是一个简单的奖励函数,
# 奖励为关节速度的平方,这里的奖励为正,但是乘上scales后为负,
# 关节速度越大,奖励越小,限制关节速度的大小
cart_vel = self.dof_vel[:, 0]
pole_vel = self.dof_vel[:, 1]
pole_vel2 = self.dof_vel[:, 2]

reward = torch.abs(cart_vel) + torch.abs(pole_vel) + torch.abs(pole_vel2)
return reward

def _reward_termination(self):
# 设置终止奖励,当倒立摆倒下的时候需要给一点惩罚
cart_pos = self.dof_pos[:, 0]
pole_angle = self.dof_pos[:, 1]
pole_angle2 = self.dof_pos[:, 2]

reward = torch.zeros(self.num_envs, dtype=torch.float, device=self.device)
reward = torch.where(torch.abs(cart_pos) > 3, torch.ones_like(reward) * -2.0, reward)
reward = torch.where(torch.abs(pole_angle) > torch.pi / 2, torch.ones_like(reward) * -2.0, reward)
reward = torch.where(torch.abs(pole_angle2) > torch.pi / 2, torch.ones_like(reward) * -2.0, reward)

return reward

torch.where(condition, x, y) 是逐元素选择:条件为真取 x,条件为假取 y。这里连续写了三次,意思是只要小车位置越界,或者任意一个杆角度超过阈值,就把对应环境的 termination reward 设成 -2.0

注意这里不是三个失败条件累加惩罚。因为每一次 torch.where 的第三个参数都是当前 reward,如果前面已经变成 -2.0,后面条件再成立也还是写成 -2.0,所以最多就是 -2.0。这和普通 reward 项不同,普通项是逐项加到 rew_buf 里的。

_reward_joint_angle()_reward_joint_velocity() 返回的都是 [num_envs],每个环境一个数。写 reward 时要保持这个习惯,不要返回 Python 标量,也不要把所有环境先求平均,否则并行环境的奖励会混在一起,reset 和 logger 都会出问题。

config 文件

两个类,一个 Cartpole2Config,另一个 Cartpole2ConfigPPO

Cartpole2Config 是环境配置,主要决定仿真和任务本身:

  1. num_envsnum_observationsnum_actions 要和代码里的 observation/action 对上。比如这里只驱动小车一个关节,所以 num_actions 应该是 1;观测里拼了位置和速度,如果是 3 个 dof,就是 6 维。

  2. num_privileged_obs 是特权观测,通常给 teacher 或 asymmetric actor-critic 使用。这个 Cartpole 例子如果没有用到,可以设为 None 或保持项目里的默认写法。

  3. env_spacing 是环境与环境之间的物理间距,主要影响可视化排列和 env 之间的空间隔离。

  4. decimation 是控制频率缩减系数,表示一个策略 action 对应多少个物理仿真小步。它会影响 self.dt,也会影响 reward scale 的时间补偿。

default_dof_drive_mode 是关节驱动模式,常见取值:

  1. 0:无驱动,自由关节。

  2. 1:位置控制。

  3. 2:速度控制。

  4. 3:力矩控制。

这个例子里 _compute_torques() 直接输出力矩,所以驱动模式要和力矩控制匹配。真正做腿足机器人时,如果 asset 里 DOF 属性、控制模式和代码里的 torque 写法不一致,会出现 action 看起来在变,但机器人基本不动,或者仿真表现很奇怪。

ref_env = 0 表示多环境并行仿真中,哪个环境作为相机跟踪的基准。这个参数不影响训练逻辑,主要影响 viewer 看哪个环境。

Cartpole2ConfigPPO 是训练配置,主要决定算法侧的参数,比如 runner、policy、algorithm、训练轮数、保存间隔等。环境代码里一般不会直接用它,但 task registry 会把环境 config 和 PPO config 一起注册,训练脚本根据 task name 取出来创建环境和 runner。

注意事项:

  1. config 里的 reward 名字要和 _reward_xxx 函数后缀一致。scales.joint_angle 对应 _reward_joint_angle(),名字对不上会在 getattr 时出错。

  2. observation 维度改了以后,要同步改 num_observations。这类错误不一定在写代码时立刻暴露,常见表现是 policy 网络输入维度不匹配。

  3. action 维度改了以后,要同步改 num_actions_compute_torques()。如果只改 config,不改 torque 映射,策略输出的 action 可能没有真正作用到想控制的关节上。

  4. reset、reward、observation 都要按 [num_envs, ...] 的 batch 思路写。不要把某一个环境的值拿出来单独算完再塞回去,能用张量批量算就批量算。