【Legged Gym】Legged Gym训练参数详解与自定义任务实现

Legged Gym训练参数详解与自定义任务实现

在进行机器人强化学习训练时，Legged Gym 提供了一套灵活的参数配置系统，以适应不同的训练需求和环境。本文将详细解析 Legged Gym 训练时的关键参数，并特别强调如何通过自定义 task 来实现新任务的训练。同时，也会解释 rl_device 和 sim_device 的赋值方法及其区别。

1. 参数详解

1.1. 任务与实验配置

• --task: 指定训练任务的类型，如 anymal_c_flat。这是训练过程中机器人需要完成的任务。用户可以通过编写自定义的 task 类来实现新的任务。
• --resume: 布尔值，用于决定是否从检查点恢复训练。
• --experiment_name: 实验的名称，用于标识不同的实验配置。
• --run_name: 运行的名称，用于区分同一实验下的不同运行。
• --load_run: 指定恢复训练时加载的运行名称。如果设置为 -1，则加载最近的运行。
• --checkpoint: 指定加载的模型检查点编号。如果设置为 -1，则加载最近的检查点。

1.2. 显示与并行配置

• --headless: 布尔值，强制关闭显示，适用于无头服务器环境。
• --horovod: 布尔值，使用 Horovod 进行多 GPU 训练。

1.3. 设备配置

• --rl_device: 指定强化学习算法使用的设备，如 cuda:0。这通常是一个 GPU，但也可以使用 CPU。
• --num_envs: 要创建的环境数量。
• --seed: 随机种子，用于确保实验的可重复性。
• --max_iterations: 训练的最大迭代次数。
• --sim_device: 指定模拟环境使用的设备，如 cuda:0。这通常是一个 GPU，但也可以使用 CPU。

2. 自定义任务实现

Legged Gym 允许用户通过自定义 task 来实现新的任务。task 类定义了机器人在环境中需要完成的任务目标和评估标准。要创建自定义任务，你需要继承 Legged Gym 的 Task 基类，并实现必要的方法，如 __init__、reset 和 step。这些方法定义了任务的初始化、重置和每个时间步的行为。

3. 参数读取代码（位于legged_gym/legged_gym/utils/helpers.py）

def get_args(custom_args=[]):
    custom_parameters = [
        {"name": "--task", "type": str, "default": "anymal_c_flat", "help": "Resume training or start testing from a checkpoint. Overrides config file if provided."},
        {"name": "--resume", "action": "store_true", "default": False,  "help": "Resume training from a checkpoint"},
        {"name": "--experiment_name", "type": str,  "help": "Name of the experiment to run or load. Overrides config file if provided."},
        {"name": "--run_name", "type": str,  "help": "Name of the run. Overrides config file if provided."},
        {"name": "--load_run", "type": str,  "help": "Name of the run to load when resume=True. If -1: will load the last run. Overrides config file if provided."},
        {"name": "--checkpoint", "type": int,  "help": "Saved model checkpoint number. If -1: will load the last checkpoint. Overrides config file if provided."},
        
        {"name": "--headless", "action": "store_true", "default": False, "help": "Force display off at all times"},
        {"name": "--horovod", "action": "store_true", "default": False, "help": "Use horovod for multi-gpu training"},
        {"name": "--rl_device", "type": str, "default": "cuda:0", "help": 'Device used by the RL algorithm, (cpu, gpu, cuda:0, cuda:1 etc..)'},
        {"name": "--num_envs", "type": int, "help": "Number of environments to create. Overrides config file if provided."},
        {"name": "--seed", "type": int, "help": "Random seed. Overrides config file if provided."},
        {"name": "--max_iterations", "type": int, "help": "Maximum number of training iterations. Overrides config file if provided."},
    ] + custom_args
    # parse arguments
    args = gymutil.parse_arguments(
        description="RL Policy",
        custom_parameters=custom_parameters)

    # name allignment
    args.sim_device_id = args.compute_device_id
    args.sim_device = args.sim_device_type
    if args.sim_device=='cuda':
        args.sim_device += f":{args.sim_device_id}"
    return args

def export_policy_as_jit(actor_critic, path):
    if hasattr(actor_critic, 'memory_a'):
        # assumes LSTM: TODO add GRU
        exporter = PolicyExporterLSTM(actor_critic)
        exporter.export(path)
    else: 
        os.makedirs(path, exist_ok=True)
        path = os.path.join(path, 'policy_1.pt')
        model = copy.deepcopy(actor_critic.actor).to('cpu')
        traced_script_module = torch.jit.script(model)
        traced_script_module.save(path)