Skip to content

Commit 3d5ea25

Browse files
yijiegkellyguo11
andauthored
Removes wandb logging in AutoMate env (#2912)
# Description wandb logging function is provided in rl_games script. So we remove the wandb logging in task level. Also, we edit the task registration style to help startup perf. ## Type of change - Bug fix (non-breaking change which fixes an issue) - This change requires a documentation update ## Checklist - [ x ] I have run the [`pre-commit` checks](https://pre-commit.com/) with `./isaaclab.sh --format` - [ x ] I have made corresponding changes to the documentation - [ x ] My changes generate no new warnings - [ ] I have added tests that prove my fix is effective or that my feature works - [ ] I have updated the changelog and the corresponding version in the extension's `config/extension.toml` file - [ x ] I have added my name to the `CONTRIBUTORS.md` or my name already exists there --------- Co-authored-by: Kelly Guo <[email protected]>
1 parent 13965cc commit 3d5ea25

File tree

5 files changed

+7
-28
lines changed

5 files changed

+7
-28
lines changed

docs/source/overview/environments.rst

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -225,7 +225,7 @@ We provide environments for both disassembly and assembly.
225225
* |disassembly-link|: The plug starts inserted in the socket. A low-level controller lifts the plug out and moves it to a random position. This process is purely scripted and does not involve any learned policy. Therefore, it does not require policy training or evaluation. The resulting trajectories serve as demonstrations for the reverse process, i.e., learning to assemble. To run disassembly for a specific task: ``python source/isaaclab_tasks/isaaclab_tasks/direct/automate/run_disassembly_w_id.py --assembly_id=ASSEMBLY_ID --disassembly_dir=DISASSEMBLY_DIR``. All generated trajectories are saved to a local directory ``DISASSEMBLY_DIR``.
226226
* |assembly-link|: The goal is to insert the plug into the socket. You can use this environment to train a policy via reinforcement learning or evaluate a pre-trained checkpoint.
227227

228-
* To train an assembly policy, we run the command ``python source/isaaclab_tasks/isaaclab_tasks/direct/automate/run_w_id.py --assembly_id=ASSEMBLY_ID --train``. We can customize the training process using the optional flags: ``--headless`` to run without opening the GUI windows, ``--max_iterations=MAX_ITERATIONS`` to set the number of training iterations, ``--num_envs=NUM_ENVS`` to set the number of parallel environments during training, ``--seed=SEED`` to assign the random seed, ``--wandb`` to enable logging to WandB (requires a WandB account). The policy checkpoints will be saved automatically during training in the directory ``logs/rl_games/Assembly/test``.
228+
* To train an assembly policy, we run the command ``python source/isaaclab_tasks/isaaclab_tasks/direct/automate/run_w_id.py --assembly_id=ASSEMBLY_ID --train``. We can customize the training process using the optional flags: ``--headless`` to run without opening the GUI windows, ``--max_iterations=MAX_ITERATIONS`` to set the number of training iterations, ``--num_envs=NUM_ENVS`` to set the number of parallel environments during training, ``--seed=SEED`` to assign the random seed. The policy checkpoints will be saved automatically during training in the directory ``logs/rl_games/Assembly/test``.
229229
* To evaluate an assembly policy, we run the command ``python source/isaaclab_tasks/isaaclab_tasks/direct/automate/run_w_id.py --assembly_id=ASSEMBLY_ID --checkpoint=CHECKPOINT --log_eval``. The evaluation results are stored in ``evaluation_{ASSEMBLY_ID}.h5``.
230230

231231
.. table::

source/isaaclab_tasks/isaaclab_tasks/direct/automate/__init__.py

Lines changed: 4 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -6,30 +6,28 @@
66
import gymnasium as gym
77

88
from . import agents
9-
from .assembly_env import AssemblyEnv, AssemblyEnvCfg
10-
from .disassembly_env import DisassemblyEnv, DisassemblyEnvCfg
119

1210
##
1311
# Register Gym environments.
1412
##
1513

1614
gym.register(
1715
id="Isaac-AutoMate-Assembly-Direct-v0",
18-
entry_point="isaaclab_tasks.direct.automate:AssemblyEnv",
16+
entry_point=f"{__name__}.assembly_env:AssemblyEnv",
1917
disable_env_checker=True,
2018
kwargs={
21-
"env_cfg_entry_point": AssemblyEnvCfg,
19+
"env_cfg_entry_point": f"{__name__}.assembly_env:AssemblyEnvCfg",
2220
"rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml",
2321
},
2422
)
2523

2624

2725
gym.register(
2826
id="Isaac-AutoMate-Disassembly-Direct-v0",
29-
entry_point="isaaclab_tasks.direct.automate:DisassemblyEnv",
27+
entry_point=f"{__name__}.disassembly_env:DisassemblyEnv",
3028
disable_env_checker=True,
3129
kwargs={
32-
"env_cfg_entry_point": DisassemblyEnvCfg,
30+
"env_cfg_entry_point": f"{__name__}.disassembly_env:DisassemblyEnvCfg",
3331
"rl_games_cfg_entry_point": f"{agents.__name__}:rl_games_ppo_cfg.yaml",
3432
},
3533
)

source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_env.py

Lines changed: 0 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,9 @@
77
import numpy as np
88
import os
99
import torch
10-
from datetime import datetime
1110

1211
import carb
1312
import isaacsim.core.utils.torch as torch_utils
14-
import wandb
1513
import warp as wp
1614

1715
import isaaclab.sim as sim_utils
@@ -71,9 +69,6 @@ def __init__(self, cfg: AssemblyEnvCfg, render_mode: str | None = None, **kwargs
7169
if self.cfg_task.sample_from != "rand":
7270
self._init_eval_loading()
7371

74-
if self.cfg_task.wandb:
75-
wandb.init(project="automate", name=self.cfg_task.assembly_id + "_" + datetime.now().strftime("%m/%d/%Y"))
76-
7772
def _init_eval_loading(self):
7873
eval_held_asset_pose, eval_fixed_asset_pose, eval_success = automate_log.load_log_from_hdf5(
7974
self.cfg_task.eval_filename
@@ -554,9 +549,6 @@ def _get_rewards(self):
554549
rew_buf = self._update_rew_buf(curr_successes)
555550
self.ep_succeeded = torch.logical_or(self.ep_succeeded, curr_successes)
556551

557-
if self.cfg_task.wandb:
558-
wandb.log(self.extras)
559-
560552
# Only log episode success rates at the end of an episode.
561553
if torch.any(self.reset_buf):
562554
self.extras["successes"] = torch.count_nonzero(self.ep_succeeded) / self.num_envs
@@ -579,12 +571,6 @@ def _get_rewards(self):
579571
)
580572

581573
self.extras["curr_max_disp"] = self.curr_max_disp
582-
if self.cfg_task.wandb:
583-
wandb.log({
584-
"success": torch.mean(self.ep_succeeded.float()),
585-
"reward": torch.mean(rew_buf),
586-
"sbc_rwd_scale": sbc_rwd_scale,
587-
})
588574

589575
if self.cfg_task.if_logging_eval:
590576
self.success_log = torch.cat([self.success_log, self.ep_succeeded.reshape((self.num_envs, 1))], dim=0)

source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_tasks_cfg.py

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -138,7 +138,6 @@ class AssemblyTask:
138138
if_logging_eval: bool = False
139139
num_eval_trials: int = 100
140140
eval_filename: str = "evaluation_00015.h5"
141-
wandb: bool = False
142141

143142
# Fine-tuning
144143
sample_from: str = "rand" # gp, gmm, idv, rand

source/isaaclab_tasks/isaaclab_tasks/direct/automate/run_w_id.py

Lines changed: 2 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -9,7 +9,7 @@
99
import sys
1010

1111

12-
def update_task_param(task_cfg, assembly_id, if_sbc, if_log_eval, if_wandb):
12+
def update_task_param(task_cfg, assembly_id, if_sbc, if_log_eval):
1313
# Read the file lines.
1414
with open(task_cfg) as f:
1515
lines = f.readlines()
@@ -21,7 +21,6 @@ def update_task_param(task_cfg, assembly_id, if_sbc, if_log_eval, if_wandb):
2121
if_sbc_pattern = re.compile(r"^(.*if_sbc\s*:\s*bool\s*=\s*).*$")
2222
if_log_eval_pattern = re.compile(r"^(.*if_logging_eval\s*:\s*bool\s*=\s*).*$")
2323
eval_file_pattern = re.compile(r"^(.*eval_filename\s*:\s*str\s*=\s*).*$")
24-
if_wandb_pattern = re.compile(r"^(.*wandb\s*:\s*bool\s*=\s*).*$")
2524

2625
for line in lines:
2726
if "assembly_id =" in line:
@@ -32,8 +31,6 @@ def update_task_param(task_cfg, assembly_id, if_sbc, if_log_eval, if_wandb):
3231
line = if_log_eval_pattern.sub(rf"\1{str(if_log_eval)}", line)
3332
elif "eval_filename: str = " in line:
3433
line = eval_file_pattern.sub(r"\1'{}'".format(f"evaluation_{assembly_id}.h5"), line)
35-
elif "wandb: bool =" in line:
36-
line = if_wandb_pattern.sub(rf"\1{str(if_wandb)}", line)
3734

3835
updated_lines.append(line)
3936

@@ -51,7 +48,6 @@ def main():
5148
default="source/isaaclab_tasks/isaaclab_tasks/direct/automate/assembly_tasks_cfg.py",
5249
)
5350
parser.add_argument("--assembly_id", type=str, help="New assembly ID to set.")
54-
parser.add_argument("--wandb", action="store_true", help="Use wandb to record learning curves")
5551
parser.add_argument("--checkpoint", type=str, help="Checkpoint path.")
5652
parser.add_argument("--num_envs", type=int, default=128, help="Number of parallel environment.")
5753
parser.add_argument("--seed", type=int, default=-1, help="Random seed.")
@@ -61,7 +57,7 @@ def main():
6157
parser.add_argument("--max_iterations", type=int, default=1500, help="Number of iteration for policy learning.")
6258
args = parser.parse_args()
6359

64-
update_task_param(args.cfg_path, args.assembly_id, args.train, args.log_eval, args.wandb)
60+
update_task_param(args.cfg_path, args.assembly_id, args.train, args.log_eval)
6561

6662
bash_command = None
6763
if sys.platform.startswith("win"):

0 commit comments

Comments
 (0)