Init

legged_gym/utils/__init__.py

@@ -0,0 +1,5 @@
+from .helpers import class_to_dict, get_load_path, get_args, set_seed, update_class_from_dict
+from .task_registry import task_registry
+from .logger import Logger
+from .math import *
+from .terrain import Terrain

legged_gym/utils/exporter.py

@@ -0,0 +1,268 @@
+# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
+# All rights reserved.
+#
+# SPDX-License-Identifier: BSD-3-Clause
+
+import copy
+import os
+import torch
+from torch import nn
+from typing import Optional
+
+
+def export_policy_as_jit(policy: object, path: str, normalizer: Optional[object] = None, filename="policy.pt"):
+    """Export policy into a Torch JIT file.
+
+    Args:
+        policy: The policy torch module.
+        normalizer: The empirical normalizer module. If None, Identity is used.
+        path: The path to the saving directory.
+        filename: The name of exported JIT file. Defaults to "policy.pt".
+    """
+    policy_exporter = _TorchPolicyExporter(policy, normalizer)
+    policy_exporter.export(path, filename)
+
+
+def export_policy_as_onnx(
+    policy: object, path: str, normalizer: Optional[object] = None, filename="policy.onnx", verbose=False
+):
+    """Export policy into a Torch ONNX file.
+
+    Args:
+        policy: The policy torch module.
+        normalizer: The empirical normalizer module. If None, Identity is used.
+        path: The path to the saving directory.
+        filename: The name of exported ONNX file. Defaults to "policy.onnx".
+        verbose: Whether to print the model summary. Defaults to False.
+    """
+    if not os.path.exists(path):
+        os.makedirs(path, exist_ok=True)
+    policy_exporter = _OnnxPolicyExporter(policy, normalizer, verbose)
+    policy_exporter.export(path, filename)
+
+
+
+def export_policy_as_pkl(
+    policy: nn.Module, path: str, filename="policy.pkl"
+):
+    """Export policy into a Torch pkl file.
+
+    Args:
+        policy: The policy torch module.
+        normalizer: The empirical normalizer module. If None, Identity is used.
+        path: The path to the saving directory.
+        filename: The name of exported pkl file. Defaults to "policy.pkl".
+    """
+    if not os.path.exists(path):
+        os.makedirs(path, exist_ok=True)
+    model_dict = policy.state_dict()
+    torch.save(model_dict, os.path.join(path, filename))
+
+
+"""
+Helper Classes - Private.
+"""
+
+
+class _TorchPolicyExporter(torch.nn.Module):
+    """Exporter of actor-critic into JIT file."""
+
+    def __init__(self, policy, normalizer=None):
+        super().__init__()
+        self.is_recurrent = policy.is_recurrent
+        # copy policy parameters
+        if hasattr(policy, "student_encoder"):
+            self.student_encoder = copy.deepcopy(policy.student_encoder).cpu()
+            self.history = torch.zeros([1, policy.history.shape[1], policy.history.shape[2]], device='cpu')
+            self.forward = self.forward_cts
+        if hasattr(policy, "student_moe_encoder"):
+            self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder).cpu()
+            self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
+            self.history_length = policy.history.shape[1]
+            self.history = torch.zeros([1, policy.history.shape[1], policy.history.shape[2]], device='cpu')
+            self.forward = self.forward_moe_cts
+        if hasattr(policy, "actor"):
+            self.actor = copy.deepcopy(policy.actor)
+            if self.is_recurrent:
+                self.rnn = copy.deepcopy(policy.memory_a.rnn)
+        elif hasattr(policy, "student"):
+            self.actor = copy.deepcopy(policy.student)
+            if self.is_recurrent:
+                self.rnn = copy.deepcopy(policy.memory_s.rnn)
+        else:
+            raise ValueError("Policy does not have an actor/student module.")
+        # set up recurrent network
+        if self.is_recurrent:
+            self.rnn.cpu()
+            self.register_buffer("hidden_state", torch.zeros(self.rnn.num_layers, 1, self.rnn.hidden_size))
+            self.register_buffer("cell_state", torch.zeros(self.rnn.num_layers, 1, self.rnn.hidden_size))
+            self.forward = self.forward_lstm
+            self.reset = self.reset_memory
+        # copy normalizer if exists
+        if normalizer:
+            self.normalizer = copy.deepcopy(normalizer)
+        else:
+            self.normalizer = torch.nn.Identity()
+
+    def forward_lstm(self, x):
+        x = self.normalizer(x)
+        x, (h, c) = self.rnn(x.unsqueeze(0), (self.hidden_state, self.cell_state))
+        self.hidden_state[:] = h
+        self.cell_state[:] = c
+        x = x.squeeze(0)
+        return self.actor(x)
+
+    def forward(self, x):
+        return self.actor(self.normalizer(x))
+    
+    def forward_cts(self, x):  # x is single observations
+        x = self.normalizer(x)
+        self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1)
+        latent = self.student_encoder(self.history.flatten(1))
+        x = torch.cat([latent, x], dim=1)
+        return self.actor(x)
+    
+    def forward_moe_cts(self, x):  # x is single observations
+        x = self.normalizer(x)
+        self.history = torch.cat([self.history[:, 1:], x.unsqueeze(1)], dim=1)
+        history_no_goal = self.history.reshape(1, self.history_length, -1)[:, :, self.obs_no_goal_mask].reshape(1, -1)
+        latent, weights = self.student_moe_encoder(self.history.flatten(1), history_no_goal)
+        x = torch.cat([latent, x], dim=1)
+        return self.actor(x), (weights, latent)
+
+    @torch.jit.export
+    def reset(self):
+        if hasattr(self, 'history'):
+            self.history = torch.zeros_like(self.history)
+
+    def reset_memory(self):
+        self.hidden_state[:] = 0.0
+        self.cell_state[:] = 0.0
+
+    def export(self, path, filename):
+        os.makedirs(path, exist_ok=True)
+        path = os.path.join(path, filename)
+        self.to("cpu")
+        traced_script_module = torch.jit.script(self)
+        traced_script_module.save(path)
+
+
+class _OnnxPolicyExporter(torch.nn.Module):
+    """Exporter of actor-critic into ONNX file."""
+
+    def __init__(self, policy, normalizer=None, verbose=False):
+        super().__init__()
+        self.verbose = verbose
+        self.input_dim = None
+        self.num_actions = 12
+        self.normalizer = torch.nn.Identity()
+        
+        # copy policy parameters
+        if hasattr(policy, 'student_encoder'):
+            self.student_encoder = copy.deepcopy(policy.student_encoder)
+            self.forward = self.forward_cts
+            self.input_dim = self.student_encoder[0].in_features
+            
+        elif hasattr(policy, "student_moe_encoder"):
+            self.student_moe_encoder = copy.deepcopy(policy.student_moe_encoder)
+            self.obs_no_goal_mask = copy.deepcopy(policy.obs_no_goal_mask).cpu()
+            self.history_length = policy.history.shape[1]
+            self.forward = self.forward_moe_cts
+            self.input_dim = self.history_length * policy.history.shape[2]
+            
+        if hasattr(policy, "actor"):
+            self.actor = copy.deepcopy(policy.actor)
+            if hasattr(self, 'is_recurrent') and self.is_recurrent:
+                self.rnn = copy.deepcopy(policy.memory_a.rnn)
+            if self.input_dim is None:
+                 self.input_dim = self.actor[0].in_features
+        else:
+            raise ValueError("Policy does not have an actor/student module.")
+
+    def forward_cts(self, x):  # x is stack observations by terms
+        x = self.normalizer(x)
+        term_dims = [3, 3, 3, self.num_actions, self.num_actions, self.num_actions]
+        obs_dim = sum(term_dims)
+        if x.shape[1] % obs_dim != 0:
+            raise ValueError(f"x.shape[1] ({x.shape[1]}) 不是 obs_dim ({obs_dim}) 的整数倍")
+            
+        frames = x.shape[1] // obs_dim
+        split_sizes = [dim * frames for dim in term_dims]
+        # [B, dim0*frames], [B, dim1*frames], ...
+        term_chunks = torch.split(x, split_sizes, dim=1)
+
+        # [ [B, frames, dim0], [B, frames, dim1], ... ]
+        frame_terms_reshaped = [
+            chunk.view(-1, frames, dim) 
+            for chunk, dim in zip(term_chunks, term_dims)
+        ]
+
+        history_by_frame = []
+        for i in range(frames):
+            # [ [B, dim0], [B, dim1], ... ]
+            terms_for_this_frame = [ftr[:, i, :] for ftr in frame_terms_reshaped]
+            history_by_frame.append(torch.cat(terms_for_this_frame, dim=1))
+        # [B, (Frame0_AllTerms), (Frame1_AllTerms), ...]
+        history = torch.cat(history_by_frame, dim=1)
+
+        last_obs = history[:, -obs_dim:]
+        latent = self.student_encoder(history)
+        x = torch.cat([latent, last_obs], dim=1)
+        
+        return self.actor(x)
+
+    def forward_moe_cts(self, x):
+        x = self.normalizer(x)
+        term_dims = [3, 3, 3, self.num_actions, self.num_actions, self.num_actions]
+        obs_dim = sum(term_dims)
+        if x.shape[1] % obs_dim != 0:
+            raise ValueError(f"x.shape[1] ({x.shape[1]}) 不是 obs_dim ({obs_dim}) 的整数倍")
+            
+        frames = x.shape[1] // obs_dim
+        split_sizes = [dim * frames for dim in term_dims]
+        # [B, dim0*frames], [B, dim1*frames], ...
+        term_chunks = torch.split(x, split_sizes, dim=1)
+
+        # [ [B, frames, dim0], [B, frames, dim1], ... ]
+        frame_terms_reshaped = [
+            chunk.view(-1, frames, dim) 
+            for chunk, dim in zip(term_chunks, term_dims)
+        ]
+
+        history_by_frame = []
+        for i in range(frames):
+            # [ [B, dim0], [B, dim1], ... ]
+            terms_for_this_frame = [ftr[:, i, :] for ftr in frame_terms_reshaped]
+            history_by_frame.append(torch.cat(terms_for_this_frame, dim=1))
+        # [B, (Frame0_AllTerms), (Frame1_AllTerms), ...]
+        history = torch.cat(history_by_frame, dim=1)
+
+        last_obs = history[:, -obs_dim:]
+        history_3d = history.view(-1, self.history_length, obs_dim)
+        history_no_goal = history_3d[:, :, self.obs_no_goal_mask].reshape(x.shape[0], -1)
+
+        latent, weights = self.student_moe_encoder(history, history_no_goal)
+        x = torch.cat([latent, last_obs], dim=1)
+
+        return self.actor(x), weights, latent
+
+    def export(self, path, filename):
+        self.to("cpu")
+        obs = torch.zeros(1, self.input_dim)
+        
+        output_names = ["actions"]
+        if self.forward == self.forward_moe_cts:
+            output_names.append("weights")
+            output_names.append("latent")
+
+        torch.onnx.export(
+            self,
+            obs,
+            os.path.join(path, filename),
+            export_params=True,
+            opset_version=11,
+            verbose=self.verbose,
+            input_names=["obs"],
+            output_names=output_names,
+            dynamic_axes={},
+        )

legged_gym/utils/helpers.py

@@ -0,0 +1,193 @@
+import os
+import copy
+import torch
+import numpy as np
+import random
+from pathlib import Path
+from isaacgym import gymapi
+from isaacgym import gymutil
+
+from legged_gym import LEGGED_GYM_ROOT_DIR, LEGGED_GYM_ENVS_DIR
+
+def class_to_dict(obj) -> dict:
+    if not  hasattr(obj,"__dict__"):
+        return obj
+    result = {}
+    for key in dir(obj):
+        if key.startswith("_"):
+            continue
+        element = []
+        val = getattr(obj, key)
+        if isinstance(val, list):
+            for item in val:
+                element.append(class_to_dict(item))
+        else:
+            element = class_to_dict(val)
+        result[key] = element
+    return result
+
+def update_class_from_dict(obj, dict):
+    for key, val in dict.items():
+        attr = getattr(obj, key, None)
+        if isinstance(attr, type):
+            update_class_from_dict(attr, val)
+        else:
+            setattr(obj, key, val)
+    return
+
+def set_seed(seed):
+    if seed == -1:
+        seed = np.random.randint(0, 10000)
+    print("Setting seed: {}".format(seed))
+    
+    random.seed(seed)
+    np.random.seed(seed)
+    torch.manual_seed(seed)
+    os.environ['PYTHONHASHSEED'] = str(seed)
+    torch.cuda.manual_seed(seed)
+    torch.cuda.manual_seed_all(seed)
+
+def parse_sim_params(args, cfg):
+    # code from Isaac Gym Preview 2
+    # initialize sim params
+    sim_params = gymapi.SimParams()
+
+    # set some values from args
+    if args.physics_engine == gymapi.SIM_FLEX:
+        if args.device != "cpu":
+            print("WARNING: Using Flex with GPU instead of PHYSX!")
+    elif args.physics_engine == gymapi.SIM_PHYSX:
+        sim_params.physx.use_gpu = args.use_gpu
+        sim_params.physx.num_subscenes = args.subscenes
+    sim_params.use_gpu_pipeline = args.use_gpu_pipeline
+
+    # if sim options are provided in cfg, parse them and update/override above:
+    if "sim" in cfg:
+        gymutil.parse_sim_config(cfg["sim"], sim_params)
+
+    # Override num_threads if passed on the command line
+    if args.physics_engine == gymapi.SIM_PHYSX and args.num_threads > 0:
+        sim_params.physx.num_threads = args.num_threads
+
+    return sim_params
+
+def get_load_path(root, load_run=-1, checkpoint=-1):
+    try:
+        runs = os.listdir(root)
+        runs = [run for run in runs if len(list((Path(root) / run).glob('model_*.pt'))) > 0]
+        #TODO sort by date to handle change of month
+        runs.sort()
+        if 'exported' in runs: runs.remove('exported')
+        last_run = os.path.join(root, runs[-1])
+    except:
+        raise ValueError("No runs in this directory: " + root)
+    if load_run==-1:
+        load_run = last_run
+    else:
+        load_run = os.path.join(root, load_run)
+
+    if checkpoint==-1:
+        models = [file for file in os.listdir(load_run) if 'model' in file]
+        models.sort(key=lambda m: '{0:0>15}'.format(m))
+        model = models[-1]
+    else:
+        model = "model_{}.pt".format(checkpoint) 
+
+    load_path = os.path.join(load_run, model)
+    return load_path
+
+def update_cfg_from_args(env_cfg, cfg_train, args):
+    # seed
+    if env_cfg is not None:
+        # num envs
+        if args.num_envs is not None:
+            env_cfg.env.num_envs = args.num_envs
+    if cfg_train is not None:
+        if args.seed is not None:
+            cfg_train.seed = args.seed
+        # alg runner parameters
+        if args.max_iterations is not None:
+            cfg_train.runner.max_iterations = args.max_iterations
+        if args.resume:
+            cfg_train.runner.resume = args.resume
+        if args.experiment_name is not None:
+            cfg_train.runner.experiment_name = args.experiment_name
+        if args.run_name is not None:
+            cfg_train.runner.run_name = args.run_name
+        if args.load_run is not None:
+            cfg_train.runner.load_run = args.load_run
+        if args.checkpoint is not None:
+            cfg_train.runner.checkpoint = args.checkpoint
+
+    return env_cfg, cfg_train
+
+def get_args():
+    custom_parameters = [
+        {"name": "--task", "type": str, "default": "go2", "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."},
+    ]
+    # 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)
+
+
+# class PolicyExporterLSTM(torch.nn.Module):
+#     def __init__(self, actor_critic):
+#         super().__init__()
+#         self.actor = copy.deepcopy(actor_critic.actor)
+#         self.is_recurrent = actor_critic.is_recurrent
+#         self.memory = copy.deepcopy(actor_critic.memory_a.rnn)
+#         self.memory.cpu()
+#         self.register_buffer(f'hidden_state', torch.zeros(self.memory.num_layers, 1, self.memory.hidden_size))
+#         self.register_buffer(f'cell_state', torch.zeros(self.memory.num_layers, 1, self.memory.hidden_size))
+
+#     def forward(self, x):
+#         out, (h, c) = self.memory(x.unsqueeze(0), (self.hidden_state, self.cell_state))
+#         self.hidden_state[:] = h
+#         self.cell_state[:] = c
+#         return self.actor(out.squeeze(0))
+
+#     @torch.jit.export
+#     def reset_memory(self):
+#         self.hidden_state[:] = 0.
+#         self.cell_state[:] = 0.
+ 
+#     def export(self, path):
+#         os.makedirs(path, exist_ok=True)
+#         path = os.path.join(path, 'policy_lstm_1.pt')
+#         self.to('cpu')
+#         traced_script_module = torch.jit.script(self)
+#         traced_script_module.save(path)
+
+    

legged_gym/utils/isaacgym_utils.py

@@ -0,0 +1,47 @@
+import os
+import numpy as np
+import random
+import torch
+
+@torch.jit.script
+def copysign(a, b):
+    # type: (float, Tensor) -> Tensor
+    a = torch.tensor(a, device=b.device, dtype=torch.float).repeat(b.shape[0])
+    return torch.abs(a) * torch.sign(b)
+def get_euler_xyz(q):
+    qx, qy, qz, qw = 0, 1, 2, 3
+    # roll (x-axis rotation)
+    sinr_cosp = 2.0 * (q[:, qw] * q[:, qx] + q[:, qy] * q[:, qz])
+    cosr_cosp = q[:, qw] * q[:, qw] - q[:, qx] * \
+                q[:, qx] - q[:, qy] * q[:, qy] + q[:, qz] * q[:, qz]
+    roll = torch.atan2(sinr_cosp, cosr_cosp)
+
+    # pitch (y-axis rotation)
+    sinp = 2.0 * (q[:, qw] * q[:, qy] - q[:, qz] * q[:, qx])
+    pitch = torch.where(
+        torch.abs(sinp) >= 1, copysign(np.pi / 2.0, sinp), torch.asin(sinp))
+
+    # yaw (z-axis rotation)
+    siny_cosp = 2.0 * (q[:, qw] * q[:, qz] + q[:, qx] * q[:, qy])
+    cosy_cosp = q[:, qw] * q[:, qw] + q[:, qx] * \
+                q[:, qx] - q[:, qy] * q[:, qy] - q[:, qz] * q[:, qz]
+    yaw = torch.atan2(siny_cosp, cosy_cosp)
+
+    return torch.stack((roll, pitch, yaw), dim=-1)
+
+def sample_disjoint_intervals(env_ids, limit_bound, cfg_min, cfg_max, device):
+    """
+    sample uniform distribution from [cfg_min, -limit_bound] U [limit_bound, cfg_max]
+    """
+    width_neg = torch.nn.functional.relu(-limit_bound - cfg_min)
+    width_pos = torch.nn.functional.relu(cfg_max - limit_bound)
+    
+    total_width = width_neg + width_pos + 1e-6 # 加极小值防除零
+    u = torch.rand(len(env_ids), device=device) * total_width
+    
+    samples = torch.where(
+        u < width_neg, 
+        cfg_min + u, 
+        cfg_max - width_pos + (u - width_neg)
+    )
+    return samples

legged_gym/utils/logger.py

@@ -0,0 +1,39 @@
+import numpy as np
+from collections import defaultdict
+from multiprocessing import Process, Value
+
+class Logger:
+    def __init__(self, dt):
+        self.state_log = defaultdict(list)
+        self.rew_log = defaultdict(list)
+        self.dt = dt
+        self.num_episodes = 0
+        self.plot_process = None
+
+    def log_state(self, key, value):
+        self.state_log[key].append(value)
+
+    def log_states(self, dict):
+        for key, value in dict.items():
+            self.log_state(key, value)
+
+    def log_rewards(self, dict, num_episodes):
+        for key, value in dict.items():
+            if 'rew' in key:
+                self.rew_log[key].append(value.item() * num_episodes)
+        self.num_episodes += num_episodes
+
+    def reset(self):
+        self.state_log.clear()
+        self.rew_log.clear()   
+
+    def print_rewards(self):
+        print("Average rewards per second:")
+        for key, values in self.rew_log.items():
+            mean = np.sum(np.array(values)) / self.num_episodes
+            print(f" - {key}: {mean}")
+        print(f"Total number of episodes: {self.num_episodes}")
+    
+    def __del__(self):
+        if self.plot_process is not None:
+            self.plot_process.kill()

legged_gym/utils/math.py

@@ -0,0 +1,26 @@
+import torch
+from torch import Tensor
+import numpy as np
+from isaacgym.torch_utils import quat_apply, normalize
+from typing import Tuple
+
+# @ torch.jit.script
+def quat_apply_yaw(quat, vec):
+    quat_yaw = quat.clone().view(-1, 4)
+    quat_yaw[:, :2] = 0.
+    quat_yaw = normalize(quat_yaw)
+    return quat_apply(quat_yaw, vec)
+
+# @ torch.jit.script
+def wrap_to_pi(angles):
+    angles %= 2*np.pi
+    angles -= 2*np.pi * (angles > np.pi)
+    return angles
+
+# @ torch.jit.script
+def torch_rand_sqrt_float(lower, upper, shape, device):
+    # type: (float, float, Tuple[int, int], str) -> Tensor
+    r = 2*torch.rand(*shape, device=device) - 1
+    r = torch.where(r<0., -torch.sqrt(-r), torch.sqrt(r))
+    r =  (r + 1.) / 2.
+    return (upper - lower) * r + lower

legged_gym/utils/task_registry.py

@@ -0,0 +1,129 @@
+import os
+from datetime import datetime
+from typing import Tuple
+import torch
+import numpy as np
+import sys
+
+from rsl_rl.env import VecEnv
+from rsl_rl.runners import OnPolicyRunner, OnPolicyRunnerCTS
+
+from legged_gym import LEGGED_GYM_ROOT_DIR, LEGGED_GYM_ENVS_DIR
+from .helpers import get_args, update_cfg_from_args, class_to_dict, get_load_path, set_seed, parse_sim_params
+from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg, LeggedRobotCfgPPO
+
+class TaskRegistry():
+    def __init__(self):
+        self.task_classes = {}
+        self.env_cfgs = {}
+        self.train_cfgs = {}
+    
+    def register(self, name: str, task_class: VecEnv, env_cfg: LeggedRobotCfg, train_cfg: LeggedRobotCfgPPO):
+        self.task_classes[name] = task_class
+        self.env_cfgs[name] = env_cfg
+        self.train_cfgs[name] = train_cfg
+    
+    def get_task_class(self, name: str) -> VecEnv:
+        return self.task_classes[name]
+    
+    def get_cfgs(self, name) -> Tuple[LeggedRobotCfg, LeggedRobotCfgPPO]:
+        train_cfg = self.train_cfgs[name]
+        env_cfg = self.env_cfgs[name]
+        # copy seed
+        env_cfg.seed = train_cfg.seed
+        return env_cfg, train_cfg
+    
+    def make_env(self, name, args=None, env_cfg=None) -> Tuple[VecEnv, LeggedRobotCfg]:
+        """ Creates an environment either from a registered namme or from the provided config file.
+
+        Args:
+            name (string): Name of a registered env.
+            args (Args, optional): Isaac Gym comand line arguments. If None get_args() will be called. Defaults to None.
+            env_cfg (Dict, optional): Environment config file used to override the registered config. Defaults to None.
+
+        Raises:
+            ValueError: Error if no registered env corresponds to 'name' 
+
+        Returns:
+            isaacgym.VecTaskPython: The created environment
+            Dict: the corresponding config file
+        """
+        # if no args passed get command line arguments
+        if args is None:
+            args = get_args()
+        # check if there is a registered env with that name
+        if name in self.task_classes:
+            task_class = self.get_task_class(name)
+        else:
+            raise ValueError(f"Task with name: {name} was not registered")
+        if env_cfg is None:
+            # load config files
+            env_cfg, _ = self.get_cfgs(name)
+        # override cfg from args (if specified)
+        env_cfg, _ = update_cfg_from_args(env_cfg, None, args)
+        set_seed(env_cfg.seed)
+        # parse sim params (convert to dict first)
+        sim_params = {"sim": class_to_dict(env_cfg.sim)}
+        sim_params = parse_sim_params(args, sim_params)
+        env = task_class(   cfg=env_cfg,
+                            sim_params=sim_params,
+                            physics_engine=args.physics_engine,
+                            sim_device=args.sim_device,
+                            headless=args.headless)
+        return env, env_cfg
+
+    def make_alg_runner(self, env, name=None, args=None, train_cfg=None, log_root="default") -> Tuple[OnPolicyRunner, LeggedRobotCfgPPO]:
+        """ Creates the training algorithm  either from a registered namme or from the provided config file.
+
+        Args:
+            env (isaacgym.VecTaskPython): The environment to train (TODO: remove from within the algorithm)
+            name (string, optional): Name of a registered env. If None, the config file will be used instead. Defaults to None.
+            args (Args, optional): Isaac Gym comand line arguments. If None get_args() will be called. Defaults to None.
+            train_cfg (Dict, optional): Training config file. If None 'name' will be used to get the config file. Defaults to None.
+            log_root (str, optional): Logging directory for Tensorboard. Set to 'None' to avoid logging (at test time for example). 
+                                      Logs will be saved in <log_root>/<date_time>_<run_name>. Defaults to "default"=<path_to_LEGGED_GYM>/logs/<experiment_name>.
+
+        Raises:
+            ValueError: Error if neither 'name' or 'train_cfg' are provided
+            Warning: If both 'name' or 'train_cfg' are provided 'name' is ignored
+
+        Returns:
+            PPO: The created algorithm
+            Dict: the corresponding config file
+        """
+        # if no args passed get command line arguments
+        if args is None:
+            args = get_args()
+        # if config files are passed use them, otherwise load from the name
+        if train_cfg is None:
+            if name is None:
+                raise ValueError("Either 'name' or 'train_cfg' must be not None")
+            # load config files
+            _, train_cfg = self.get_cfgs(name)
+        else:
+            if name is not None:
+                print(f"'train_cfg' provided -> Ignoring 'name={name}'")
+        # override cfg from args (if specified)
+        _, train_cfg = update_cfg_from_args(None, train_cfg, args)
+
+        if log_root=="default":
+            log_root = os.path.join(LEGGED_GYM_ROOT_DIR, 'logs', train_cfg.runner.experiment_name)
+            log_dir = os.path.join(log_root, datetime.now().strftime('%b%d_%H-%M-%S') + '_' + train_cfg.runner.run_name)
+        elif log_root is None:
+            log_dir = None
+        else:
+            log_dir = os.path.join(log_root, datetime.now().strftime('%b%d_%H-%M-%S') + '_' + train_cfg.runner.run_name)
+        
+        train_cfg_dict = class_to_dict(train_cfg)
+        runner = eval(train_cfg.runner_class_name)(env, train_cfg_dict, log_dir, device=args.rl_device)
+        #save resume path before creating a new log_dir
+        resume = train_cfg.runner.resume
+        if resume:
+            # load previously trained model
+            resume_path = get_load_path(log_root, load_run=train_cfg.runner.load_run, checkpoint=train_cfg.runner.checkpoint)
+            print(f"Loading model from: {resume_path}")
+            runner.load(resume_path)
+        return runner, train_cfg
+
+# make global task registry
+task_registry = TaskRegistry()

legged_gym/utils/terrain.py

@@ -0,0 +1,188 @@
+from collections import defaultdict
+import numpy as np
+from numpy.random import choice
+from scipy import interpolate
+
+from isaacgym import terrain_utils
+from legged_gym.envs.base.legged_robot_config import LeggedRobotCfg
+
+class Terrain:
+    def __init__(self, cfg: LeggedRobotCfg.terrain, num_robots) -> None:
+
+        self.cfg = cfg
+        self.num_robots = num_robots
+        self.type = cfg.mesh_type
+        if self.type in ["none", 'plane']:
+            return
+        self.env_length = cfg.terrain_length
+        self.env_width = cfg.terrain_width
+        self.proportions = [np.sum(cfg.terrain_proportions[:i+1]) for i in range(len(cfg.terrain_proportions))]
+
+        self.cfg.num_sub_terrains = cfg.num_rows * cfg.num_cols
+        self.env_origins = np.zeros((cfg.num_rows, cfg.num_cols, 3))
+
+        self.width_per_env_pixels = int(self.env_width / cfg.horizontal_scale)
+        self.length_per_env_pixels = int(self.env_length / cfg.horizontal_scale)
+
+        self.spacing = cfg.terrain_spacing
+        self.spacing_pixels = int(self.spacing / cfg.horizontal_scale)
+
+        self.border = int(cfg.border_size/self.cfg.horizontal_scale)
+        self.tot_cols = int(cfg.num_cols * self.width_per_env_pixels + max(0, cfg.num_cols-1) * self.spacing_pixels) + 2 * self.border
+        self.tot_rows = int(cfg.num_rows * self.length_per_env_pixels + max(0, cfg.num_rows-1) * self.spacing_pixels) + 2 * self.border
+        self.name2cols = defaultdict(set)  # terrain type to column index
+        self.cols2id = []  # column index to terrain id
+
+        self.height_field_raw = np.zeros((self.tot_rows , self.tot_cols), dtype=np.int16)
+        if cfg.curriculum:
+            self.curiculum()
+        elif cfg.selected:
+            self.selected_terrain()
+        else:    
+            self.randomized_terrain()   
+        
+        self.heightsamples = self.height_field_raw
+        if self.type=="trimesh":
+            self.vertices, self.triangles = terrain_utils.convert_heightfield_to_trimesh(   self.height_field_raw,
+                                                                                            self.cfg.horizontal_scale,
+                                                                                            self.cfg.vertical_scale,
+                                                                                            self.cfg.slope_treshold)
+    
+    def randomized_terrain(self):
+        for k in range(self.cfg.num_sub_terrains):
+            # Env coordinates in the world
+            (i, j) = np.unravel_index(k, (self.cfg.num_rows, self.cfg.num_cols))
+
+            choice = np.random.uniform(0, 1)
+            difficulty = np.random.choice([0.5, 0.75, 0.9])
+            terrain = self.make_terrain(choice, difficulty)
+            self.add_terrain_to_map(terrain, i, j)
+        
+    def curiculum(self):
+        for j in range(self.cfg.num_cols):
+            for i in range(self.cfg.num_rows):
+                difficulty = i / self.cfg.num_rows
+                choice = j / self.cfg.num_cols + 0.001
+
+                terrain = self.make_terrain(choice, difficulty)
+                self.add_terrain_to_map(terrain, i, j)
+            self.name2cols[terrain.terrain_name].add(j)
+            self.cols2id.append(terrain.terrain_id)
+
+    def selected_terrain(self):
+        terrain_type = self.cfg.terrain_kwargs.pop('type')
+        for k in range(self.cfg.num_sub_terrains):
+            # Env coordinates in the world
+            (i, j) = np.unravel_index(k, (self.cfg.num_rows, self.cfg.num_cols))
+
+            terrain = terrain_utils.SubTerrain("terrain",
+                              width=self.width_per_env_pixels,
+                              length=self.width_per_env_pixels,
+                              vertical_scale=self.vertical_scale,
+                              horizontal_scale=self.horizontal_scale)
+
+            eval(terrain_type)(terrain, **self.cfg.terrain_kwargs.terrain_kwargs)
+            self.add_terrain_to_map(terrain, i, j)
+    
+    def make_terrain(self, choice, difficulty):
+        terrain = terrain_utils.SubTerrain("terrain",
+                                width=self.width_per_env_pixels,
+                                length=self.width_per_env_pixels,
+                                vertical_scale=self.cfg.vertical_scale,
+                                horizontal_scale=self.cfg.horizontal_scale)
+        slope = 0.1 + difficulty * 0.52
+        step_height = 0.05 + 0.23 * difficulty
+        discrete_obstacles_height = 0.05 + difficulty * 0.25
+        stepping_stones_size = 1.5 * (1.05 - difficulty)
+        stone_distance = 0.05 if difficulty==0 else 0.1
+        gap_size = 1. * difficulty
+        # pit_depth = 1. * difficulty  # 删除坑洞
+        amplitude = 0.1 + 0.2 * difficulty
+        
+        if choice < self.proportions[0]:
+            terrain.terrain_name = "wave"
+            terrain.terrain_id = 0
+            terrain_utils.wave_terrain(terrain, num_waves=5, amplitude=amplitude)
+            terrain_utils.random_uniform_terrain(terrain, min_height=-0.05, max_height=0.05, step=0.005, downsampled_scale=0.2)
+        elif choice < self.proportions[1]:  # 平滑坡
+            terrain.terrain_name = "slope"
+            terrain.terrain_id = 1
+            if choice < (self.proportions[0] + self.proportions[1])/ 2:  # 一半正坡, 一半负坡
+                slope *= -1
+            terrain_utils.pyramid_sloped_terrain(terrain, slope=slope, platform_size=3.)
+        elif choice < self.proportions[2]:  # 粗糙坡
+            terrain.terrain_name = "rough_slope"
+            terrain.terrain_id = 2
+            terrain_utils.pyramid_sloped_terrain(terrain, slope=slope, platform_size=3.)
+            terrain_utils.random_uniform_terrain(terrain, min_height=-0.05, max_height=0.05, step=0.005, downsampled_scale=0.2)
+        elif choice < self.proportions[4]:  # 下楼梯
+            terrain.terrain_name = "stairs_down"
+            terrain.terrain_id = 4
+            if choice<self.proportions[3]:  # 上楼梯
+                terrain.terrain_name = "stairs_up"
+                terrain.terrain_id = 3
+                step_height *= -1
+            terrain_utils.pyramid_stairs_terrain(terrain, step_width=0.31, step_height=step_height, platform_size=3.)
+        elif choice < self.proportions[5]:  # 障碍物
+            terrain.terrain_name = "obstacles"
+            terrain.terrain_id = 5
+            num_rectangles = 20
+            rectangle_min_size = 1.
+            rectangle_max_size = 2.
+            terrain_utils.discrete_obstacles_terrain(terrain, discrete_obstacles_height, rectangle_min_size, rectangle_max_size, num_rectangles, platform_size=3.)
+        elif choice < self.proportions[6]:  # 梅花桩
+            terrain.terrain_name = "stepping_stones"
+            terrain.terrain_id = 6
+            terrain_utils.stepping_stones_terrain(terrain, stone_size=stepping_stones_size, stone_distance=stone_distance, max_height=0., platform_size=4.)
+        elif choice < self.proportions[7]:  # 间隙
+            terrain.terrain_name = "gap"
+            terrain.terrain_id = 7
+            gap_terrain(terrain, gap_size=gap_size, platform_size=3.)
+        else:  # 平地
+            terrain.terrain_name = "flat"
+            terrain.terrain_id = 8
+            pit_terrain(terrain, depth=0.0, platform_size=4.)
+        
+        return terrain
+
+    def add_terrain_to_map(self, terrain, row, col):
+        i = row
+        j = col
+        # map coordinate system
+        start_x = self.border + i * (self.length_per_env_pixels + self.spacing_pixels)
+        end_x = start_x + self.length_per_env_pixels
+        start_y = self.border + j * (self.width_per_env_pixels + self.spacing_pixels)
+        end_y = start_y + self.width_per_env_pixels
+        self.height_field_raw[start_x: end_x, start_y:end_y] = terrain.height_field_raw
+
+        env_origin_x = (i + 0.5) * self.env_length + i * self.spacing
+        env_origin_y = (j + 0.5) * self.env_width + j * self.spacing
+        x1 = int((self.env_length/2. - 1) / terrain.horizontal_scale)
+        x2 = int((self.env_length/2. + 1) / terrain.horizontal_scale)
+        y1 = int((self.env_width/2. - 1) / terrain.horizontal_scale)
+        y2 = int((self.env_width/2. + 1) / terrain.horizontal_scale)
+        env_origin_z = np.max(terrain.height_field_raw[x1:x2, y1:y2])*terrain.vertical_scale
+        self.env_origins[i, j] = [env_origin_x, env_origin_y, env_origin_z]
+
+def gap_terrain(terrain, gap_size, platform_size=1.):
+    gap_size = int(gap_size / terrain.horizontal_scale)
+    platform_size = int(platform_size / terrain.horizontal_scale)
+
+    center_x = terrain.length // 2
+    center_y = terrain.width // 2
+    x1 = (terrain.length - platform_size) // 2
+    x2 = x1 + gap_size
+    y1 = (terrain.width - platform_size) // 2
+    y2 = y1 + gap_size
+   
+    terrain.height_field_raw[center_x-x2 : center_x + x2, center_y-y2 : center_y + y2] = -1000
+    terrain.height_field_raw[center_x-x1 : center_x + x1, center_y-y1 : center_y + y1] = 0
+
+def pit_terrain(terrain, depth, platform_size=1.):
+    depth = int(depth / terrain.vertical_scale)
+    platform_size = int(platform_size / terrain.horizontal_scale / 2)
+    x1 = terrain.length // 2 - platform_size
+    x2 = terrain.length // 2 + platform_size
+    y1 = terrain.width // 2 - platform_size
+    y2 = terrain.width // 2 + platform_size
+    terrain.height_field_raw[x1:x2, y1:y2] = -depth