Skip to content
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

Fix reset functions #1416

Open
wants to merge 13 commits into
base: main
Choose a base branch
from
+70 −10
Open

Fix reset functions #1416

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
13 commits
Select commit Hold shift + click to select a range
ffd9cd1
Add windows configuration to VS code tasks
johnBuffer Sep 9, 2024
9484b0b
Add name to contributors
johnBuffer Sep 9, 2024
e308e2d
Merge branch 'main' into main
johnBuffer Sep 9, 2024
43625bc
Merge branch 'main' into main
johnBuffer Sep 10, 2024
d4856d6
Fix conflicts
johnBuffer Nov 13, 2024
70ef8b3
Add reset function
johnBuffer Nov 13, 2024
b19d96b
Add reset function
johnBuffer Nov 13, 2024
f0d8bfe
Add reset function
johnBuffer Nov 13, 2024
ce7e42b
Fix comments
johnBuffer Nov 13, 2024
3c39a02
Update version and changelog
johnBuffer Nov 13, 2024
562e2b1
Update changelog
johnBuffer Nov 13, 2024
8c1d278
Merge branch 'main' into fix_reset_function
johnBuffer Nov 24, 2024
e890cb1
Merge branch 'main' into fix_reset_function
johnBuffer Nov 27, 2024
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
12 changes: 12 additions & 0 deletions source/extensions/omni.isaac.lab/docs/CHANGELOG.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,18 @@
Changelog
---------

0.27.16 (2024-11-13)
~~~~~~~~~~~~~~~~~~~~

Added
^^^^^

* Added the function :func:`~omni.isaac.lab.envs.mdp.events.reset_joints_to_default` which allows to reset all
joints to their default states.

* Added the possibility for :func:`~omni.isaac.lab.envs.mdp.events.reset_joints_by_scale` and
:func:`~omni.isaac.lab.envs.mdp.events.reset_joints_by_offset` to target specific joints in robots.

0.27.15 (2024-11-09)
~~~~~~~~~~~~~~~~~~~~

Expand Down
52 changes: 42 additions & 10 deletions source/extensions/omni.isaac.lab/omni/isaac/lab/envs/mdp/events.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -797,6 +797,23 @@ def reset_root_state_from_terrain(
asset.write_root_velocity_to_sim(velocities, env_ids=env_ids)


def reset_joints_to_default(
env: ManagerBasedEnv,
env_ids: torch.Tensor,
asset_cfg: SceneEntityCfg = SceneEntityCfg("robot"),
):
"""Reset the robot joints to their default positions and velocities."""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]

# get default joint state
joint_pos = asset.data.default_joint_pos[env_ids].clone()
joint_vel = asset.data.default_joint_vel[env_ids].clone()

# set into the physics simulation
asset.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids)


def reset_joints_by_scale(
env: ManagerBasedEnv,
env_ids: torch.Tensor,
Expand All @@ -811,23 +828,31 @@ def reset_joints_by_scale(
"""
# extract the used quantities (to enable type-hinting)
asset: Articulation = env.scene[asset_cfg.name]

# get default joint state
joint_pos = asset.data.default_joint_pos[env_ids].clone()
joint_vel = asset.data.default_joint_vel[env_ids].clone()
joint_pos_all = asset.data.joint_pos[env_ids].clone()
joint_vel_all = asset.data.joint_vel[env_ids].clone()

# isolate relevant joints
joint_pos = joint_pos_all[:, asset_cfg.joint_ids]
joint_vel = joint_vel_all[:, asset_cfg.joint_ids]

# scale these values randomly
joint_pos *= math_utils.sample_uniform(*position_range, joint_pos.shape, joint_pos.device)
joint_vel *= math_utils.sample_uniform(*velocity_range, joint_vel.shape, joint_vel.device)

joint_pos_all[:, asset_cfg.joint_ids] = joint_pos
joint_vel_all[:, asset_cfg.joint_ids] = joint_vel

# clamp joint pos to limits
joint_pos_limits = asset.data.soft_joint_pos_limits[env_ids]
joint_pos = joint_pos.clamp_(joint_pos_limits[..., 0], joint_pos_limits[..., 1])
joint_pos_all = joint_pos_all.clamp_(joint_pos_limits[..., 0], joint_pos_limits[..., 1])
# clamp joint vel to limits
joint_vel_limits = asset.data.soft_joint_vel_limits[env_ids]
joint_vel = joint_vel.clamp_(-joint_vel_limits, joint_vel_limits)
joint_vel_all = joint_vel_all.clamp_(-joint_vel_limits, joint_vel_limits)

# set into the physics simulation
asset.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids)
asset.write_joint_state_to_sim(joint_pos_all, joint_vel_all, env_ids=env_ids)


def reset_joints_by_offset(
Expand All @@ -846,22 +871,29 @@ def reset_joints_by_offset(
asset: Articulation = env.scene[asset_cfg.name]

# get default joint state
joint_pos = asset.data.default_joint_pos[env_ids].clone()
joint_vel = asset.data.default_joint_vel[env_ids].clone()
joint_pos_all = asset.data.joint_pos[env_ids].clone()
joint_vel_all = asset.data.joint_vel[env_ids].clone()

# isolate relevant joints
joint_pos = joint_pos_all[:, asset_cfg.joint_ids]
joint_vel = joint_vel_all[:, asset_cfg.joint_ids]

# bias these values randomly
joint_pos += math_utils.sample_uniform(*position_range, joint_pos.shape, joint_pos.device)
joint_vel += math_utils.sample_uniform(*velocity_range, joint_vel.shape, joint_vel.device)

joint_pos_all[:, asset_cfg.joint_ids] = joint_pos
joint_vel_all[:, asset_cfg.joint_ids] = joint_vel

# clamp joint pos to limits
joint_pos_limits = asset.data.soft_joint_pos_limits[env_ids]
joint_pos = joint_pos.clamp_(joint_pos_limits[..., 0], joint_pos_limits[..., 1])
joint_pos_all = joint_pos_all.clamp_(joint_pos_limits[..., 0], joint_pos_limits[..., 1])
# clamp joint vel to limits
joint_vel_limits = asset.data.soft_joint_vel_limits[env_ids]
joint_vel = joint_vel.clamp_(-joint_vel_limits, joint_vel_limits)
joint_vel_all = joint_vel_all.clamp_(-joint_vel_limits, joint_vel_limits)

# set into the physics simulation
asset.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids)
asset.write_joint_state_to_sim(joint_pos_all, joint_vel_all, env_ids=env_ids)
Copy link
Contributor

@Mayankm96 Mayankm96 Dec 19, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Isn't this a slower operation since you are copying over the data twice (first here by setting into joint_pos_all and then it is again done inside the method). Why not keep it cleaner and just operate over the asset_cfg.joint_indices in the whole thing?

Something like:

    # extract the used quantities (to enable type-hinting)
    asset: Articulation = env.scene[asset_cfg.name]
    command_term: PhaseCommand = env.command_manager.get_term(command_name)

    # broadcast env_ids if needed to allow double indexing
    if env_ids != slice(None) and asset_cfg.joint_ids != slice(None):
        env_ids = env_ids[:, None]

    # read state
    ....

    # clamp joint pos to limits
    joint_pos_limits = asset.data.soft_joint_pos_limits[env_ids, asset_cfg.joint_ids]
    joint_pos = joint_pos.clamp_(joint_pos_limits[..., 0], joint_pos_limits[..., 1])
    # clamp joint vel to limits
    joint_vel_limits = asset.data.soft_joint_vel_limits[env_ids, asset_cfg.joint_ids]
    joint_vel = joint_vel.clamp_(-joint_vel_limits, joint_vel_limits)

    # make sure env ids dimensions are compatible
    if env_ids != slice(None) and asset_cfg.joint_ids != slice(None):
        env_ids = env_ids.squeeze(1)
    # set into the physics simulation
    asset.write_joint_state_to_sim(joint_pos, joint_vel, env_ids=env_ids, joint_ids=asset_cfg.joint_ids)



def reset_nodal_state_uniform(
Expand Down
10 changes: 10 additions & 0 deletions source/extensions/omni.isaac.lab_tasks/docs/CHANGELOG.rst
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -1,6 +1,16 @@
Changelog
---------

0.10.14 (2024-11-13)
~~~~~~~~~~~~~~~~~~~~

Changed
^^^^^^^

* Patched the :class:`~omni.isaac.lab_tasks.manager_based.manager_based.classic.cartpole.cartpole_env_cfg.CartpoleEnvCfg`
to use the `reset_joints_to_default` function.


0.10.13 (2024-10-30)
~~~~~~~~~~~~~~~~~~~~

Expand Down
6 changes: 6 additions & 0 deletions ...i.isaac.lab_tasks/omni/isaac/lab_tasks/manager_based/classic/cartpole/cartpole_env_cfg.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -87,6 +87,12 @@ class EventCfg:
"""Configuration for events."""

# reset
reset_to_default = EventTerm(
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

would this change affect other environments that use reset_joints_by_offset? such as ant, humanoid, cabinet

Copy link
Contributor

@Mayankm96 Mayankm96 Dec 19, 2024 •
edited
Loading

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

This environment file isn't a place to have this feature as a unit test IMO. We're changing the way the environment operates.

func=mdp.reset_joints_to_default,
mode="reset",
params={"asset_cfg": SceneEntityCfg("robot", joint_names=["slider_to_cart"])},
)

reset_cart_position = EventTerm(
func=mdp.reset_joints_by_offset,
mode="reset",
Expand Down