Skip to content

Devel newton #89

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Open
fan-ziqi wants to merge 3 commits into
base: main
Choose a base branch
from
Open

Devel newton #89

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
0325f42
newton
fan-ziqi Aug 27, 2025
345a657
fix: change `nefc_per_env` to `njmax`
fan-ziqi Sep 4, 2025
cfc4d66
update gitignore
fan-ziqi Sep 4, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
9 changes: 9 additions & 0 deletions README.md
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -8,6 +8,15 @@
[![pre-commit](https://img.shields.io/badge/pre--commit-enabled-brightgreen?logo=pre-commit&logoColor=white)](https://pre-commit.com/)
[![License](https://img.shields.io/badge/license-Apache2.0-yellow.svg)](https://opensource.org/license/apache-2-0)

## Newton

```bash
# train
python scripts/reinforcement_learning/rsl_rl/train.py --task=RobotLab-Isaac-Velocity-Flat-Unitree-Go2W-v0 --headless
# play
python scripts/reinforcement_learning/rsl_rl/play.py --task=RobotLab-Isaac-Velocity-Flat-Unitree-Go2W-v0
```

## Overview

**robot_lab** is a RL extension library for robots, based on IsaacLab. It allows you to develop in an isolated environment, outside of the core Isaac Lab repository.
Expand Down
59 changes: 37 additions & 22 deletions scripts/reinforcement_learning/rsl_rl/play.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -40,6 +40,7 @@
help="Use the pre-trained checkpoint from Nucleus.",
)
parser.add_argument("--real-time", action="store_true", default=False, help="Run in real-time, if possible.")
parser.add_argument("--newton_visualizer", action="store_true", default=False, help="Enable Newton rendering.")
parser.add_argument("--keyboard", action="store_true", default=False, help="Whether to use keyboard.")
# append RSL-RL cli arguments
cli_args.add_rsl_rl_args(parser)
Expand All @@ -66,29 +67,43 @@

from rsl_rl.runners import OnPolicyRunner

from isaaclab.devices import Se2Keyboard, Se2KeyboardCfg
from isaaclab.utils.timer import Timer

Timer.enable = False
Timer.enable_display_output = False

# from isaaclab.devices import Se2Keyboard, Se2KeyboardCfg
from isaaclab.envs import (
DirectMARLEnv,
DirectMARLEnvCfg,
# DirectMARLEnv,
# DirectMARLEnvCfg,
DirectRLEnvCfg,
ManagerBasedRLEnvCfg,
multi_agent_to_single_agent,
# multi_agent_to_single_agent,
)
from isaaclab.managers import ObservationTermCfg as ObsTerm
from isaaclab.utils.assets import retrieve_file_path
from isaaclab.utils.dict import print_dict
from isaaclab.utils.pretrained_checkpoint import get_published_pretrained_checkpoint
from isaaclab_rl.rsl_rl import RslRlOnPolicyRunnerCfg, RslRlVecEnvWrapper, export_policy_as_jit, export_policy_as_onnx
from isaaclab_tasks.utils import get_checkpoint_path
from isaaclab_tasks.utils import get_checkpoint_path, parse_env_cfg
from isaaclab_tasks.utils.hydra import hydra_task_config

import robot_lab.tasks # noqa: F401


@hydra_task_config(args_cli.task, args_cli.agent)
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlOnPolicyRunnerCfg):
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg
# | DirectMARLEnvCfg
, agent_cfg: RslRlOnPolicyRunnerCfg):
"""Play with RSL-RL agent."""
task_name = args_cli.task.split(":")[-1]
env_cfg = parse_env_cfg(
args_cli.task,
device=args_cli.device,
num_envs=args_cli.num_envs,
use_fabric=not args_cli.disable_fabric,
newton_visualizer=args_cli.newton_visualizer,
)
# override configurations with non-hydra CLI arguments
agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else 50
Expand All @@ -113,19 +128,19 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
env_cfg.events.push_robot = None
env_cfg.curriculum.command_levels = None

if args_cli.keyboard:
env_cfg.scene.num_envs = 1
env_cfg.terminations.time_out = None
env_cfg.commands.base_velocity.debug_vis = False
config = Se2KeyboardCfg(
v_x_sensitivity=env_cfg.commands.base_velocity.ranges.lin_vel_x[1],
v_y_sensitivity=env_cfg.commands.base_velocity.ranges.lin_vel_y[1],
omega_z_sensitivity=env_cfg.commands.base_velocity.ranges.ang_vel_z[1],
)
controller = Se2Keyboard(config)
env_cfg.observations.policy.velocity_commands = ObsTerm(
func=lambda env: torch.tensor(controller.advance(), dtype=torch.float32).unsqueeze(0).to(env.device),
)
# if args_cli.keyboard:
# env_cfg.scene.num_envs = 1
# env_cfg.terminations.time_out = None
# env_cfg.commands.base_velocity.debug_vis = False
# config = Se2KeyboardCfg(
# v_x_sensitivity=env_cfg.commands.base_velocity.ranges.lin_vel_x[1],
# v_y_sensitivity=env_cfg.commands.base_velocity.ranges.lin_vel_y[1],
# omega_z_sensitivity=env_cfg.commands.base_velocity.ranges.ang_vel_z[1],
# )
# controller = Se2Keyboard(config)
# env_cfg.observations.policy.velocity_commands = ObsTerm(
# func=lambda env: torch.tensor(controller.advance(), dtype=torch.float32).unsqueeze(0).to(env.device),
# )

# specify directory for logging experiments
log_root_path = os.path.join("logs", "rsl_rl", agent_cfg.experiment_name)
Expand All @@ -146,9 +161,9 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
# create isaac environment
env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)

# convert to single-agent instance if required by the RL algorithm
if isinstance(env.unwrapped, DirectMARLEnv):
env = multi_agent_to_single_agent(env)
# # convert to single-agent instance if required by the RL algorithm
# if isinstance(env.unwrapped, DirectMARLEnv):
# env = multi_agent_to_single_agent(env)

# wrap for video recording
if args_cli.video:
Expand Down
23 changes: 16 additions & 7 deletions scripts/reinforcement_learning/rsl_rl/train.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,6 +35,7 @@
parser.add_argument(
"--distributed", action="store_true", default=False, help="Run training with multiple GPUs or nodes."
)
parser.add_argument("--newton_visualizer", action="store_true", default=False, help="Enable Newton rendering.")
# append RSL-RL cli arguments
cli_args.add_rsl_rl_args(parser)
# append AppLauncher cli args
Expand All @@ -61,12 +62,17 @@

from rsl_rl.runners import OnPolicyRunner

from isaaclab.utils.timer import Timer

Timer.enable = False
Timer.enable_display_output = False

from isaaclab.envs import (
DirectMARLEnv,
DirectMARLEnvCfg,
# DirectMARLEnv,
# DirectMARLEnvCfg,
DirectRLEnvCfg,
ManagerBasedRLEnvCfg,
multi_agent_to_single_agent,
# multi_agent_to_single_agent,
)
from isaaclab.utils.dict import print_dict
from isaaclab.utils.io import dump_pickle, dump_yaml
Expand All @@ -83,7 +89,9 @@


@hydra_task_config(args_cli.task, args_cli.agent)
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agent_cfg: RslRlOnPolicyRunnerCfg):
def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg
# | DirectMARLEnvCfg
, agent_cfg: RslRlOnPolicyRunnerCfg):
"""Train with RSL-RL agent."""
# override configurations with non-hydra CLI arguments
agent_cfg = cli_args.update_rsl_rl_cfg(agent_cfg, args_cli)
Expand All @@ -96,6 +104,7 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
# note: certain randomizations occur in the environment initialization so we set the seed here
env_cfg.seed = agent_cfg.seed
env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device
env_cfg.sim.enable_newton_rendering = args_cli.newton_visualizer

# multi-gpu training configuration
if args_cli.distributed:
Expand All @@ -122,9 +131,9 @@ def main(env_cfg: ManagerBasedRLEnvCfg | DirectRLEnvCfg | DirectMARLEnvCfg, agen
# create isaac environment
env = gym.make(args_cli.task, cfg=env_cfg, render_mode="rgb_array" if args_cli.video else None)

# convert to single-agent instance if required by the RL algorithm
if isinstance(env.unwrapped, DirectMARLEnv):
env = multi_agent_to_single_agent(env)
# # convert to single-agent instance if required by the RL algorithm
# if isinstance(env.unwrapped, DirectMARLEnv):
# env = multi_agent_to_single_agent(env)

# save resume path before creating a new log_dir
if agent_cfg.resume or agent_cfg.algorithm.class_name == "Distillation":
Expand Down
8 changes: 4 additions & 4 deletions ...t_lab/tasks/manager_based/locomotion/velocity/config/wheeled/unitree_go2w/flat_env_cfg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -13,14 +13,14 @@ def __post_init__(self):
super().__post_init__()

# override rewards
self.rewards.base_height_l2.params["sensor_cfg"] = None
# self.rewards.base_height_l2.params["sensor_cfg"] = None
# change terrain to flat
self.scene.terrain.terrain_type = "plane"
self.scene.terrain.terrain_generator = None
# no height scan
self.scene.height_scanner = None
self.observations.policy.height_scan = None
self.observations.critic.height_scan = None
# self.scene.height_scanner = None
# self.observations.policy.height_scan = None
# self.observations.critic.height_scan = None
# no terrain curriculum
self.curriculum.terrain_levels = None

Expand Down
12 changes: 6 additions & 6 deletions ..._lab/tasks/manager_based/locomotion/velocity/config/wheeled/unitree_go2w/rough_env_cfg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -75,8 +75,8 @@ def __post_init__(self):

# ------------------------------Sence------------------------------
self.scene.robot = UNITREE_GO2W_CFG.replace(prim_path="{ENV_REGEX_NS}/Robot")
self.scene.height_scanner.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name
self.scene.height_scanner_base.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name
# self.scene.height_scanner.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name
# self.scene.height_scanner_base.prim_path = "{ENV_REGEX_NS}/Robot/" + self.base_link_name

# ------------------------------Observations------------------------------
self.observations.policy.joint_pos.func = mdp.joint_pos_rel_without_wheel
Expand All @@ -92,7 +92,7 @@ def __post_init__(self):
self.observations.policy.joint_pos.scale = 1.0
self.observations.policy.joint_vel.scale = 0.05
self.observations.policy.base_lin_vel = None
self.observations.policy.height_scan = None
# self.observations.policy.height_scan = None
self.observations.policy.joint_pos.params["asset_cfg"].joint_names = self.joint_names
self.observations.policy.joint_vel.params["asset_cfg"].joint_names = self.joint_names

Expand Down Expand Up @@ -136,9 +136,9 @@ def __post_init__(self):
self.rewards.lin_vel_z_l2.weight = -2.0
self.rewards.ang_vel_xy_l2.weight = -0.05
self.rewards.flat_orientation_l2.weight = 0
self.rewards.base_height_l2.weight = 0
self.rewards.base_height_l2.params["target_height"] = 0.40
self.rewards.base_height_l2.params["asset_cfg"].body_names = [self.base_link_name]
# self.rewards.base_height_l2.weight = 0
# self.rewards.base_height_l2.params["target_height"] = 0.40
# self.rewards.base_height_l2.params["asset_cfg"].body_names = [self.base_link_name]
self.rewards.body_lin_acc_l2.weight = 0
self.rewards.body_lin_acc_l2.params["asset_cfg"].body_names = [self.base_link_name]

Expand Down
8 changes: 4 additions & 4 deletions source/robot_lab/robot_lab/tasks/manager_based/locomotion/velocity/mdp/events.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -7,7 +7,7 @@
from typing import TYPE_CHECKING, Literal

import isaaclab.utils.math as math_utils
from isaaclab.assets import Articulation, RigidObject
from isaaclab.assets import Articulation
from isaaclab.managers import SceneEntityCfg

if TYPE_CHECKING:
Expand All @@ -33,7 +33,7 @@ def randomize_rigid_body_inertia(
only during the initialization of the environment.
"""
# extract the used quantities (to enable type-hinting)
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
asset: Articulation = env.scene[asset_cfg.name]

# resolve environment ids
if env_ids is None:
Expand Down Expand Up @@ -95,8 +95,8 @@ def randomize_com_positions(
operation (Literal["add", "scale", "abs"]): The operation to apply for randomization.
distribution (Literal["uniform", "log_uniform", "gaussian"]): The distribution to sample random values from.
"""
# Extract the asset (Articulation or RigidObject)
asset: RigidObject | Articulation = env.scene[asset_cfg.name]
# Extract the asset
asset: Articulation = env.scene[asset_cfg.name]

# Resolve environment indices
if env_ids is None:
Expand Down
Loading
Loading