isaac-sim / IsaacLab

Unified framework for robot learning built on NVIDIA Isaac Sim
https://isaac-sim.github.io/IsaacLab
Other
2.32k stars 957 forks source link

[Question] How to access terrain height features without adding height scanner under `ManagerBasedRLEnv` #1366

Open securityCFS opened 3 weeks ago

securityCFS commented 3 weeks ago

Problem

While trying to migrate my RL training codes from legged_gym env to isaaclab env, I found it difficult to access features like height_field in legged_gym to monitor height of certain robot components above ground without sensors. For instance, I tried to use feet height to design a certain reward function that penalizes unbalanced behavior between legs. Since the height information is required by reward function only, adding an additional observation (like height scanner) might be unnecessary. Under flat terrains, I could use scene.data.root_state_w[:,:,2] as replacement, but I couldn't find a way to subtract terrain height under other terrains.

Question

Is there a way to read terrain height, given certain point under world coordinates? Like

asset = env.scene[asset_cfg.name]
coords = asset.data.root_state_w[:, asset_cfg.body_ids, :2]
heights = get_terrain_heights(coords)
KyleM73 commented 3 weeks ago

I do not believe this exists by default, since the coordinates are continuous valued and the terrains are often randomly generated with arbitrary size. For a fixed (cached, not procedurally generated) terrain you could precompute a function like this based on the desired resolution and load the lookup table during training, it sounds like that might be what you're looking for? For a very large terrain with high resolution, using a height scanner is probably your best bet, though. You do not need to add it as an observation, just create the sensors in the scene and pass the corresponding entity config to the reward term.

securityCFS commented 3 weeks ago

I do not believe this exists by default

Actually I have posted the reference to legged gym term (see the link legged_gym which guides to the corresponding term) which is a discrete, sampled matrix of height fields. The approximate terrain height can be then calculated from interpolation.

vassil-atn commented 3 weeks ago

I'm trying to implement this, too. Is there a way to add a single sensor to the simulation (for all of the environments)? I've tried something like:

    env_height_scanner = RayCasterCfg(
        prim_path="/World/ground",
        offset=RayCasterCfg.OffsetCfg(pos=(0.0, 0.0, 20.0)),
        attach_yaw_only=True,
        pattern_cfg=patterns.GridPatternCfg(resolution=0.1, size=[16.,16.]),
        debug_vis=True,
        mesh_prim_paths=["/World/ground"],
    )

But this crashes on resets if you have more than 1 environment. But simulating a height scan like this for every single env is incredibly slow. If I can fix this I can post an example of the function to find the height at specific locations given the height scanner.

vassil-atn commented 3 weeks ago

I've found somewhat of a workaround - instead of creating Raycasters, you can simply use raycast_mesh to essentially get the height at desired locations. Here's a quickly thrown-together simple example:

from omni.isaac.lab.utils.warp import convert_to_warp_mesh, raycast_mesh
import omni.isaac.lab.sim as sim_utils
import warp as wp
from omni.isaac.core.prims import XFormPrimView
from pxr import UsdGeom, UsdPhysics
import omni.log
from omni.isaac.lab.terrains.trimesh.utils import make_plane

def _initiate_meshes(mesh_prim_paths, device): 
    meshes: dict[str, wp.Mesh] = {}
    # check number of mesh prims provided
    if len(mesh_prim_paths) != 1:
        raise NotImplementedError(
            f"RayCaster currently only supports one mesh prim. Received: {len(mesh_prim_paths)}"
        )

    # read prims to ray-cast
    for mesh_prim_path in mesh_prim_paths:
        # check if the prim is a plane - handle PhysX plane as a special case
        # if a plane exists then we need to create an infinite mesh that is a plane
        mesh_prim = sim_utils.get_first_matching_child_prim(
            mesh_prim_path, lambda prim: prim.GetTypeName() == "Plane"
        )
        # if we did not find a plane then we need to read the mesh
        if mesh_prim is None:
            # obtain the mesh prim
            mesh_prim = sim_utils.get_first_matching_child_prim(
                mesh_prim_path, lambda prim: prim.GetTypeName() == "Mesh"
            )
            # check if valid
            if mesh_prim is None or not mesh_prim.IsValid():
                raise RuntimeError(f"Invalid mesh prim path: {mesh_prim_path}")
            # cast into UsdGeomMesh
            mesh_prim = UsdGeom.Mesh(mesh_prim)
            # read the vertices and faces
            points = np.asarray(mesh_prim.GetPointsAttr().Get())
            indices = np.asarray(mesh_prim.GetFaceVertexIndicesAttr().Get())
            wp_mesh = convert_to_warp_mesh(points, indices, device=device)
            # print info
            omni.log.info(
                f"Read mesh prim: {mesh_prim.GetPath()} with {len(points)} vertices and {len(indices)} faces."
            )
        else:
            mesh = make_plane(size=(2e6, 2e6), height=0.0, center_zero=True)
            wp_mesh = convert_to_warp_mesh(mesh.vertices, mesh.faces, device=device)
            # print info
            omni.log.info(f"Created infinite plane mesh prim: {mesh_prim.GetPath()}.")
        # add the warp mesh to the list
        meshes[mesh_prim_path] = wp_mesh

    # throw an error if no meshes are found
    if all([mesh_prim_path not in meshes for mesh_prim_path in mesh_prim_paths]):
        raise RuntimeError(
            f"No meshes found for ray-casting! Please check the mesh prim paths: {mesh_prim_paths}"
        )

    return meshes

def get_height_at_pos(scan_pos_w, env_ids, des_pos_w,device):
    """ Get the height at the desired positions (des_pos_w in (len(env_ids),num_points,2)) in the world frame. The position is computed relative to the scan_pos_w (len(env_ids),3). """

    ray_starts_w = torch.zeros((len(env_ids), des_pos_w.size(1), 3), device=device)
    ray_starts_w += scan_pos_w[:,None,:] # The scans should be relative to the scan_pos_w (for example this can be the env_origins)
    ray_starts_w[:,:,:2] += des_pos_w[:,:,:2] # Add the desired positions in world frame (relative to the env_origin)
    ray_starts_w[:, :, 2] += 20. # offset

    ray_directions_w = torch.zeros_like(ray_starts_w)
    ray_directions_w[:,:,2] = -1.0
    # Get the meshes (can be done once outside this function)
    mesh_prim_paths = ["/World/ground"]
    meshes = _initiate_meshes(mesh_prim_paths, device)

    # ray cast and store the hits
    ray_hits_w = raycast_mesh(
        ray_starts_w,
        ray_directions_w,
        max_dist=1e6,
        mesh=meshes[mesh_prim_paths[0]],
    )[0]

    return ray_hits_w

In this case the des_pos_w are in the world frame of each environment (i.e. (0,0) would be at the origin of the environment), but you can adjust it as you'd like. I'm also not taking rotation into account here as the scan origin is fixed.

This seems to be much faster as you only sample the num_points per env rather than a whole grid. I guess alternatively you could have a RayCaster for every env and only update it at the start of the sim (as long as the terarin is fixed and doesn't change), which might also work.

I think it would be useful to have something like this in a more generalised class, kind of like a RayCaster sensor that only computes the raycasts for the desired locations when called. @Mayankm96, if there is interest in implementing something like this, I can try to put it together.

RandomOakForest commented 2 weeks ago

Thanks for submitting this. We'll follow up for a possible enhancement.