File size: 4,864 Bytes
0037d53 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122 123 124 125 126 127 128 129 130 | # Copyright (c) 2024-2026, The UW Lab Project Developers. (https://github.com/uw-lab/UWLab/blob/main/CONTRIBUTORS.md).
# All Rights Reserved.
#
# SPDX-License-Identifier: BSD-3-Clause
"""Script to visualize saved reset states from a dataset directory."""
from __future__ import annotations
import argparse
import time
import torch
from typing import cast
from isaaclab.app import AppLauncher
# add argparse arguments
parser = argparse.ArgumentParser(description="Visualize saved reset states from a dataset directory.")
parser.add_argument("--num_envs", type=int, default=1, help="Number of environments to simulate.")
parser.add_argument("--task", type=str, default=None, help="Name of the task.")
parser.add_argument(
"--dataset_dir",
type=str,
default="./Datasets/OmniReset",
help="Base dataset directory (contains Resets/<Pair>/ subdirectories).",
)
parser.add_argument(
"--reset_type",
type=str,
default=None,
help="Single reset type to visualize (e.g. ObjectAnywhereEEAnywhere). If omitted, all four types are loaded.",
)
parser.add_argument("--reset_interval", type=float, default=0.1, help="Time interval between resets in seconds.")
AppLauncher.add_app_launcher_args(parser)
args_cli, remaining_args = parser.parse_known_args()
# launch omniverse app
app_launcher = AppLauncher(headless=args_cli.headless)
simulation_app = app_launcher.app
"""Rest everything else."""
import contextlib
import gymnasium as gym
import inspect
from isaaclab.envs import ManagerBasedRLEnv
from isaaclab.managers import ManagerTermBase
from uwlab_tasks.utils.hydra import hydra_task_compose
torch.backends.cuda.matmul.allow_tf32 = True
torch.backends.cudnn.allow_tf32 = True
torch.backends.cudnn.deterministic = False
torch.backends.cudnn.benchmark = False
@hydra_task_compose(args_cli.task, "env_cfg_entry_point", hydra_args=remaining_args)
def main(env_cfg, agent_cfg) -> None:
# override configurations with non-hydra CLI arguments
env_cfg.scene.num_envs = args_cli.num_envs if args_cli.num_envs is not None else env_cfg.scene.num_envs
env_cfg.sim.device = args_cli.device if args_cli.device is not None else env_cfg.sim.device
# make sure environment is non-deterministic for diverse pose discovery
env_cfg.seed = None
# Override existing MultiResetManager params to use the CLI-specified dataset/types
ALL_RESET_TYPES = [
"ObjectAnywhereEEAnywhere",
"ObjectRestingEEGrasped",
"ObjectAnywhereEEGrasped",
"ObjectPartiallyAssembledEEGrasped",
]
reset_types = [args_cli.reset_type] if args_cli.reset_type else ALL_RESET_TYPES
env_cfg.events.reset_from_reset_states.params["dataset_dir"] = args_cli.dataset_dir
env_cfg.events.reset_from_reset_states.params["reset_types"] = reset_types
env_cfg.events.reset_from_reset_states.params["probs"] = [1.0] * len(reset_types)
# create environment
env = cast(ManagerBasedRLEnv, gym.make(args_cli.task, cfg=env_cfg)).unwrapped
# The EventManager is created before sim.play(), so ManagerTermBase classes
# are deferred to a timeline callback that can silently fail. Force-init any
# class-based event terms that the callback missed.
for mode_cfgs in env.event_manager._mode_term_cfgs.values():
for tc in mode_cfgs:
if inspect.isclass(tc.func) and issubclass(tc.func, ManagerTermBase):
tc.func = tc.func(cfg=tc, env=env)
env.reset()
# Initialize variables
print(f"Starting visualization of saved states from {args_cli.dataset_dir}")
print("Press Ctrl+C to stop")
with contextlib.suppress(KeyboardInterrupt):
while True:
asset = env.unwrapped.scene["robot"]
# specific for robotiq
gripper_joint_positions = asset.data.joint_pos[:, asset.find_joints(["finger_joint"])[0][0]]
gripper_closed_fraction = (
torch.abs(gripper_joint_positions) / env_cfg.actions.gripper.close_command_expr["finger_joint"]
)
gripper_mask = gripper_closed_fraction > 0.1
# Step the simulation
for _ in range(5):
action = torch.zeros(env.action_space.shape, device=env.device, dtype=torch.float32)
action[gripper_mask, -1] = -1.0
action[~gripper_mask, -1] = 1.0
env.step(action)
for _ in range(5):
env.unwrapped.sim.step()
success = env.unwrapped.reward_manager.get_term_cfg("progress_context").func.success
print("Success: ", success)
# Wait for the specified interval
time.sleep(args_cli.reset_interval)
# Reset the environment to load a new state
env.reset()
env.close()
if __name__ == "__main__":
main()
# close sim app
simulation_app.close()
|