MaJ

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/manager_based_rl_env.py

@@ -77,6 +77,8 @@ def __init__(self, cfg: ManagerBasedRLEnvCfg, render_mode: str | None = None, **
         self.render_mode = render_mode
 
         # initialize data and constants
+        self.nZ = torch.tensor([0.0,0.0,-1.0], device=self.device)
+        self.Z = torch.tensor([0.0,0.0,1.0], device=self.device)
         # -- counter for curriculum
         self.common_step_counter = 0
         # -- init buffers

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/observations.py

@@ -182,6 +182,13 @@ def body_incoming_wrench(env: ManagerBasedEnv, asset_cfg: SceneEntityCfg) -> tor
     return link_incoming_forces.view(env.num_envs, -1)
 
 
+def binary_contact(env: BaseEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
+	contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
+
+	iscontact = contact_sensor.data.current_contact_time[:,sensor_cfg.body_ids] > 0.0
+	return iscontact.float()
+
+
 """
 Actions.
 """

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/rewards.py

@@ -13,6 +13,7 @@
 
 import torch
 from typing import TYPE_CHECKING
+from math import pi
 
 from omni.isaac.lab.assets import Articulation, RigidObject
 from omni.isaac.lab.managers import SceneEntityCfg
@@ -296,3 +297,197 @@ def track_ang_vel_z_exp(
     # compute the error
     ang_vel_error = torch.square(env.command_manager.get_command(command_name)[:, 2] - asset.data.root_ang_vel_b[:, 2])
     return torch.exp(-ang_vel_error / std**2)
+
+
+
+# linear task rewards (rewards err closer to 0): ---------------------------------------------------------------------------------------------------------
+# [0,1] range : 0 is maxerr, 1 is err=0
+# form :
+#	err = clip(|err|)
+# 	rew = 1 - err / maxerr
+
+DEBUG_REW = False
+if DEBUG_REW : import inspect
+
+def lin(err: float, maxerr: float):
+	if DEBUG_REW : print(inspect.currentframe().f_back.f_code.co_name, torch.amax(err, dim=0))
+
+	# err should be >0
+	err = torch.clip(err, max=maxerr)
+	return 1 - err / maxerr
+def linmax(err: float, maxerr: float):
+	if DEBUG_REW : print(inspect.currentframe().f_back.f_code.co_name, torch.amax(err, dim=0))
+
+	# err should be >0
+	err = torch.clip(err, max=maxerr)
+	return err / maxerr
+
+def r_joint_acc_lin(env: RLTaskEnv, maxerr: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+	asset: Articulation = env.scene[asset_cfg.name]
+
+	data = asset.data.joint_acc[:, asset_cfg.joint_ids]
+	err = torch.sum(torch.abs(data), dim=-1)
+	return lin(err, maxerr)
+def r_joint_vel_lin(env: RLTaskEnv, maxerr: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+	asset: Articulation = env.scene[asset_cfg.name]
+
+	data = asset.data.joint_vel[:, asset_cfg.joint_ids]
+	err = torch.sum(torch.abs(data), dim=-1)
+	return lin(err, maxerr)
+def r_joint_torque_lin(env: RLTaskEnv, maxerr: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+	asset: Articulation = env.scene[asset_cfg.name]
+
+	data = asset.data.applied_torque[:, asset_cfg.joint_ids]
+	err = torch.sum(torch.abs(data), dim=-1)
+	return lin(err, maxerr)
+def r_joint_power_lin(env: RLTaskEnv, maxerr: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+	asset: Articulation = env.scene[asset_cfg.name]
+
+	torque = asset.data.applied_torque[:, asset_cfg.joint_ids]
+	speed = asset.data.joint_vel[:, asset_cfg.joint_ids]
+	err = torch.sum(torch.abs(torque * speed), dim=-1)
+	return lin(err, maxerr)
+
+def r_frame_val_lin(env: RLTaskEnv, maxerr: float, val_name: str, frame_cfg) -> torch.Tensor:
+	frame: FrameTransformer = env.scene[frame_cfg.name]
+
+	data = getattr(frame.data, val_name)
+	err = torch.sum(torch.linalg.norm(data, dim=-1), dim=-1)
+	return lin(err, maxerr)
+def r_frame_vel_lin(env: RLTaskEnv, maxerr: float, frame_cfg) -> torch.Tensor:
+	return r_frame_val_lin(env, maxerr, "target_vel_w", frame_cfg)
+def r_frame_acc_lin(env: RLTaskEnv, maxerr: float, frame_cfg) -> torch.Tensor:
+	return r_frame_val_lin(env, maxerr, "target_acc_w", frame_cfg)
+
+def r_action_rate_lin(env: RLTaskEnv, maxerr: float) -> torch.Tensor:
+	diff = env.action_manager.action - env.action_manager.prev_action
+	err = torch.linalg.norm(diff, dim=-1)
+	return lin(err, maxerr)	
+def r_action_lin(env: RLTaskEnv, maxerr: float) -> torch.Tensor:
+	data = env.action_manager.action
+	err = torch.sum(torch.abs(data))
+	return lin(err, maxerr)	
+
+def r_velz_lin(env: RLTaskEnv, maxerr: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+	asset: RigidObject = env.scene[asset_cfg.name]
+	err = torch.abs(asset.data.root_lin_vel_b[:, 2])
+
+	return lin(err, maxerr)
+def r_acc_lin(env: RLTaskEnv, maxerr: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+	asset: Articulation = env.scene[asset_cfg.name]
+
+	data = asset.data.body_lin_acc_w[:, asset_cfg.body_ids, :]
+	err = torch.sum(torch.norm(data, dim=-1), dim=-1)
+	return lin(err, maxerr)
+
+def r_flat_orientation_lin(env: RLTaskEnv, maxerr: float=2.0, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+	asset: RigidObject = env.scene[asset_cfg.name]
+
+	dir_g = torch.nn.functional.normalize(asset.data.projected_gravity_b, dim=-1)
+	err = torch.linalg.norm(env.nZ - dir_g, dim=-1) # 2 at max
+	return lin(err, maxerr)
+def r_joint_pose_lin(env: RLTaskEnv, maxerr: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+	asset: Articulation = env.scene[asset_cfg.name]
+
+	diff = asset.data.joint_pos[:, asset_cfg.joint_ids] - asset.data.default_joint_pos[:, asset_cfg.joint_ids]
+	err = torch.linalg.norm(diff, dim=-1)
+	return lin(err, maxerr)
+def r_contact_dist_lin(env: RLTaskEnv, maxerr: float, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
+	contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
+
+	err = torch.sum(contact_sensor.data.current_contact_distance[:,sensor_cfg.body_ids], dim=-1)
+	return lin(err, maxerr)
+def r_feetair_velcom_lin(env: RLTaskEnv, target_t: float, maxerr: float, command_name: str, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
+	contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
+	command: torch.Tensor = env.command_manager.get_command(command_name)
+
+	times = contact_sensor.data.current_air_time[:, sensor_cfg.body_ids]
+	err = torch.abs(times - target_t).sum(dim=-1)
+	#err *= torch.norm(command[:, :2], dim=1) > 0.1
+	return lin(err, maxerr)
+def r_nbcontact_lin(env: RLTaskEnv, sensor_cfg: SceneEntityCfg) -> torch.Tensor:
+	contact_sensor: ContactSensor = env.scene.sensors[sensor_cfg.name]
+
+	iscontact = contact_sensor.data.current_contact_time[:,sensor_cfg.body_ids] > 1e-3
+	err = iscontact.float().sum(dim=-1)
+	maxerr = iscontact.size(dim=-1)
+	return lin(err, maxerr)
+
+def r_com_linvel_lin(env: RLTaskEnv, maxerr: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+	asset: RigidObject = env.scene[asset_cfg.name]
+	command: torch.Tensor = env.command_manager.get_command(command_name)
+
+	diff = command[:,:2] - asset.data.root_lin_vel_b[:, :2]
+	err = torch.linalg.norm(diff, dim=-1)
+	return lin(err, maxerr)
+def r_com_angvel_lin(env: RLTaskEnv, maxerr: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+	asset: RigidObject = env.scene[asset_cfg.name]
+	command: torch.Tensor = env.command_manager.get_command(command_name)
+
+	diff = command[:, 2] - asset.data.root_ang_vel_b[:, 2]
+	err = torch.abs(diff)
+	return lin(err, maxerr)
+def r_com_heading_lin(env: RLTaskEnv, command_name: str, maxerr: float = pi, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+	asset: RigidObject = env.scene[asset_cfg.name]
+	command: torch.Tensor = env.command_manager.get_command(command_name)
+
+	err = torch.abs(command[:, 2]) # pi at max
+	return lin(err, maxerr)
+
+def r_com_headingpos_lin(env: RLTaskEnv, command_name: str, maxerr: float = pi, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+	asset: RigidObject = env.scene[asset_cfg.name]
+	command: torch.Tensor = env.command_manager.get_term(command_name).heading_command_b
+
+	err = torch.abs(command[:]) # pi at max
+	return lin(err, maxerr)
+def r_com_eepos_lin(env: RLTaskEnv, maxerr: float, command_name: str, ee_frame_cfg: SceneEntityCfg = SceneEntityCfg("ee_frame")) -> torch.Tensor:
+	ee_frame: FrameTransformer = env.scene[ee_frame_cfg.name]
+	command: torch.Tensor = env.command_manager.get_term(command_name).pos_command_w
+
+	diff = ee_frame.data.target_pos_w[:,0,:] - command[:,:3]
+	err = torch.linalg.norm(diff, dim=-1)
+	return lin(err, maxerr)
+
+def r_com_pos_lin(env: RLTaskEnv, maxerr: float, command_name: str, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+	asset: RigidObject = env.scene[asset_cfg.name]
+	command: torch.Tensor = env.command_manager.get_term(command_name).pos_command_w
+
+	diff = asset.data.root_pos_w[:,:3] - command[:,:3]
+	err = torch.linalg.norm(diff, dim=-1)
+	return lin(err, maxerr) 
+
+def r_com_releepos_lin(env: RLTaskEnv, maxerr: float, command_name: str, 
+	robot_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
+	object_cfg: SceneEntityCfg = SceneEntityCfg("robot", body_names=["gripperStator"])
+) -> torch.Tensor:
+	"""Reward the agent for tracking the goal pose using tanh-kernel."""
+	# extract the used quantities (to enable type-hinting)
+	robot: RigidObject = env.scene[robot_cfg.name]
+	object: RigidObject = env.scene[object_cfg.name]
+	command = env.command_manager.get_command(command_name)
+	# compute the desired position in the world frame
+	des_pos_b = command[:, :3]
+	des_pos_w, _ = combine_frame_transforms(robot.data.root_state_w[:, :3], robot.data.root_state_w[:, 3:7], des_pos_b)
+	# distance of the end-effector to the object: (num_envs,)
+	distance = torch.norm(des_pos_w - object.data.root_pos_w[:, :3], dim=1)
+
+	return lin(distance, maxerr) 
+
+
+# linear max task rewards (rewards err closer to errmax): ---------------------------------------------------------------------------------------------------------
+# [0,1] range : 0 is err=0, 1 is err=errmax
+# form :
+#	err = clip(|err|)
+# 	rew = err / maxerr
+
+
+
+def r_velx_linmax(env: RLTaskEnv, maxerr: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+	asset: RigidObject = env.scene[asset_cfg.name]
+
+	err = torch.abs(asset.data.root_lin_vel_w[:, 0])
+	return linmax(err, maxerr)
+def r_height_linmax(env: RLTaskEnv, maxerr: float, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")) -> torch.Tensor:
+	asset: RigidObject = env.scene[asset_cfg.name]
+	err = torch.clip(asset.data.root_pos_w[:, 2], min=0)
+	return linmax(err, maxerr)

source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/terminations.py

@@ -72,6 +72,19 @@ def root_height_below_minimum(
     return asset.data.root_pos_w[:, 2] < minimum_height
 
 
+def root_out_of_curriculum(
+	env: RLTaskEnv, asset_cfg: SceneEntityCfg = SceneEntityCfg("robot")
+) -> torch.Tensor:
+	"""Terminate when the asset's root position is outside the terrain.
+		Only usable with generator terrain.
+	"""
+
+	# extract the used quantities (to enable type-hinting)
+	asset: RigidObject = env.scene[asset_cfg.name]
+	terrain: TerrainImporter = env.scene.terrain
+	dist = torch.abs(asset.data.root_pos_w[:, :2] - env.scene.env_origins[:,:2])
+	return torch.logical_or(dist[:,0] > terrain.cfg.terrain_generator.size[0]/2, dist[:,1] > terrain.cfg.terrain_generator.size[1]/2)
+
 """
 Joint terminations.
 """

source/extensions/omni.isaac.lab/omni/isaac/lab/terrains/config/rough.py

@@ -6,8 +6,8 @@
 """Configuration for custom terrains."""
 
 import omni.isaac.lab.terrains as terrain_gen
-
-from ..terrain_generator_cfg import TerrainGeneratorCfg
+import math
+from ..terrain_generator_cfg import TerrainGeneratorCfg, FlatPatchSamplingCfg
 
 ROUGH_TERRAINS_CFG = TerrainGeneratorCfg(
     size=(8.0, 8.0),
@@ -50,3 +50,90 @@
     },
 )
 """Rough terrains configuration."""
+
+
+PLATFORM_sz = 2.0
+BORDER_WIDTH = 0.0
+MAX_SLOPE = 33.0*math.pi/180.0
+TERRAIN_SZ = 10.0
+TARGET_SAMPLING_CFG = FlatPatchSamplingCfg(
+	num_patches=8,
+	patch_radius=[0.01],
+	max_height_diff=math.inf,
+	x_range = (-TERRAIN_SZ/2, TERRAIN_SZ/2),
+	y_range = (-TERRAIN_SZ/2, TERRAIN_SZ/2),
+	z_range = (-100, 100), # should be enough
+)
+CUSTOM_TERRAIN_CFG = TerrainGeneratorCfg(
+	size=(TERRAIN_SZ,TERRAIN_SZ),
+	border_width=5.0, # prevent robots from falling over the edge
+	num_rows=10,
+	num_cols=8,
+	use_cache=False,
+	curriculum=True,
+	sub_terrains={
+		"flat": terrain_gen.MeshPlaneTerrainCfg(),
+		"boxes": terrain_gen.MeshRandomGridTerrainCfg(
+			platform_width=PLATFORM_sz,
+			grid_width=0.45,
+			grid_height_range=(0.0, 0.2),
+		),
+		"pyramid_slope": terrain_gen.MeshPyramidSlopeTerrainCfg(
+			platform_width=PLATFORM_sz,
+			border_width=BORDER_WIDTH,
+			slope_angle_range=(0.0, MAX_SLOPE),
+			holes=False,
+		),
+		"pyramid_stairs": terrain_gen.MeshPyramidStairsTerrainCfg(
+			platform_width=PLATFORM_sz,
+			border_width=BORDER_WIDTH,
+			step_height_range=(0.0, 0.3),
+			step_width=0.3,
+			holes=False,
+		),
+		"pyramid_randstairs": terrain_gen.MeshPyramidStairsRandTerrainCfg(
+			platform_width=PLATFORM_sz,
+			border_width=BORDER_WIDTH,
+			step_height_range=(0.0, 0.3),
+			step_width_range=(0.25, 0.7),
+			step_height_maxincr = 0.1,
+			holes=False,
+		),
+		"inv_pyramid_slope": terrain_gen.MeshInvertedPyramidSlopeTerrainCfg(
+			platform_width=PLATFORM_sz,
+			border_width=BORDER_WIDTH,
+			slope_angle_range=(0.0, MAX_SLOPE),
+			holes=False,
+		),
+		"inv_pyramid_stairs": terrain_gen.MeshInvertedPyramidStairsTerrainCfg(
+			platform_width=PLATFORM_sz,
+			border_width=BORDER_WIDTH,
+			step_height_range=(0.0, 0.3),
+			step_width=0.3,
+			holes=False,
+		),
+		"inv_pyramid_randstairs": terrain_gen.MeshInvertedPyramidStairsRandTerrainCfg(
+			platform_width=PLATFORM_sz,
+			border_width=BORDER_WIDTH,
+			step_height_range=(0.0, 0.3),
+			step_width_range=(0.25, 0.7),
+			step_height_maxincr = 0.1,
+			holes=False,
+		),
+	},
+)
+for key in CUSTOM_TERRAIN_CFG.sub_terrains : CUSTOM_TERRAIN_CFG.sub_terrains[key].flat_patch_sampling={'target': TARGET_SAMPLING_CFG}
+
+NOISE_TERRAIN_CFG = TerrainGeneratorCfg(
+	size=(4.0,4.0),
+	difficulty_range=(1.0,1.0),
+	num_rows=1,
+	num_cols=1,
+	use_cache=False,
+	curriculum=False,
+	sub_terrains={
+		"noise": terrain_gen.HfRandomUniformTerrainCfg(
+			noise_range=(0.02, 0.10), noise_step=0.02
+		)
+	}
+)

source/extensions/omni.isaac.lab/omni/isaac/lab/terrains/trimesh/__init__.py

@@ -12,18 +12,4 @@
 efficient than the height-field representation, but it is not as flexible.
 """
 
-from .mesh_terrains_cfg import (
-    MeshBoxTerrainCfg,
-    MeshFloatingRingTerrainCfg,
-    MeshGapTerrainCfg,
-    MeshInvertedPyramidStairsTerrainCfg,
-    MeshPitTerrainCfg,
-    MeshPlaneTerrainCfg,
-    MeshPyramidStairsTerrainCfg,
-    MeshRailsTerrainCfg,
-    MeshRandomGridTerrainCfg,
-    MeshRepeatedBoxesTerrainCfg,
-    MeshRepeatedCylindersTerrainCfg,
-    MeshRepeatedPyramidsTerrainCfg,
-    MeshStarTerrainCfg,
-)
+from .mesh_terrains_cfg import *