[Proposal] Improve Current Ray Training Scripts

[binw666]

Proposal

This proposal aims to enhance the current Ray-based training scripts in the isaaclab project, with a focus on improving usability and flexibility for distributed training workflows. The goal is to address several limitations related to custom module support, file management, resource specification, and configuration handling.

Motivation

The existing Ray scripts located in the scripts/reinforcement_learning/ray directory exhibit several usability issues that hinder effective deployment and integration in practical scenarios:

• Lack of support for py_modules prevents users from distributing custom Python packages or third-party libraries to Ray workers, limiting extensibility and customization.
• Model files are written to local or relative directories, such as logs/, which makes it difficult to access trained models—especially in remote or cloud environments where local filesystems are not directly accessible.
• Resource allocation is not intuitive—users often encounter confusion or make mistakes when specifying CPU, GPU, or memory requirements.
• Heavy reliance on command-line arguments only results in long, complex, and error-prone commands, reducing readability and maintainability.

I'm often frustrated when I need to manually patch the Ray scripts just to load a custom library or retrieve trained models without having to SSH into a remote machine.

Alternatives

Some potential workarounds include:

• Modifying the Ray scripts manually to add support for py_modules.
• Using external tools (e.g., scp, or cloud CLI utilities) to move files after training.
• Writing custom Ray submission scripts to better manage job configuration, resource allocation, and post-training file handling.

While these alternatives can provide temporary relief, they increase complexity and maintenance burden, and do not offer a sustainable or user-friendly solution in the long run.

Additional context

Currently, all training outputs (e.g., model checkpoints and logs) are saved to local working directories used by Ray. This setup works reasonably well in local environments but becomes problematic in remote or cluster-based executions. To improve accessibility and streamline deployment, there should be a built-in way to move or copy these files to mounted storage locations (e.g., network drives, cloud buckets like OSS) after training completes.

Additionally, while the current scripts do allow some level of resource specification, the interface is not intuitive.

Checklist

• I have checked that there is no similar issue in the repo (required)

Acceptance Criteria

• Support for py_modules to allow distribution of custom Python packages
• Ability to move or copy training output files (e.g., models, logs) to specified directories or mounted storage (e.g., OSS)
• Use of a structured configuration format (e.g., YAML) instead of relying solely on CLI arguments to specify tasks and hardware resources (CPU, GPU, memory)

[RandomOakForest]

Thank you for posting this proposal. The team will review it.

[binw666]

@glvov-bdai Thank you very much for introducing Ray into IsaacLab. I would like to make some improvements to the Ray-related sections based on the above content. Do you have any suggestions?

[kellyguo11]

@garylvov for viz

[garylvov]

Hi I'm currently on vacation this summer before starting my PhD so I apologize for the short/messy response as I don't have a computer with me. Just fyi I no longer monitor the @glvov-bdai account as my internship ended so please feel free to ping @garylvov instead.

The suggested changes mostly sound good. It should be pretty straight forward to add pymodules, but the use of the special isaaclab.sh may make this difficult: see how I had to patch the ray shebang.

As for retrieving the logs,there should be a way to do this without scp. I think scp may be a brittle solution. Currently, the configuration supports automatically fetching the logs with MLFlow without SCP from remote clusters: check out the video tutorial linked in the tutorial. Ray is able to integrate with services like Wandb to save logs AND checkpoints(I don't believe I currently save checkpoints), which I think having optional support for would be better than manual scp stuff. With buckets, it may be hard to integrate with auth keys, and I think this is reinventing the wheel for stuff that Ray should have functionality for. So, basically,yes I support this, but we should do this the Ray integrated way that makes this functionality optional.

As for files specifying resources instead of command line arguments, this sounds fine, but we should retain support for the command line arguments. A potential way to do this is use a config system like hydra to be consistent with the rest of the project.

Feel free to ping me with any questions or feedback or on the PR I will be happy to help to the extent that I can :)

[garylvov]

As for resource specification being unintuitive, I think it may serve well to simply make the documentation more explicit as opposed to changing the code. I'm rather averse to breaking changes.

For example, although the resource allocation is not intuitive for you, it is currently intuitive for me. If we change it, it may become intuitive for you, but not intuitive to someone else. I think resource allocation intuition will always be a source of contention. That being said, feel free to change my mind here maybe I'm missing an simple obvious way to make it easy for everyone.

[garylvov]

Also, for the PR, @ozhanozen may be a good final reviewer in addition to myself (if he has time) as he is familiar with the Ray workflow having recently made improvements as well.

[binw666]

@garylvov Thanks a lot for your response and for taking the time to provide feedback while on vacation — really appreciate it!

I wanted to ask a quick follow-up question about one part of your message:

"The use of the special isaaclab.sh may make this difficult."

Could you elaborate a bit more on what makes isaaclab.sh special or tricky in this context? I'm trying to understand the underlying issue with integrating Ray and pymodules when using that script.

Again, thanks so much for your help and guidance!

[garylvov]

Basically, Isaac lab doesn't use the system python, but instead uses the isaaclab.sh bash script that does some special wrapping stuff and then defers to an another bash script packaged with Isaac sim that then calls some python executable.

I'm not the best person to explain this as I'm unsure why we do this maybe @hhansen-bdai has a better idea.

The consequence of this, is that if you just "pip install ...", it may install it on the system python, which may not be visible to the isaac lab python. Instead, you have to "./isaaclab.sh -p -m pip install ...", so that it shows up to Isaac lab. In order to get the "ray" commands to work smoothly in the docker container, I had to modify the Ray commands shebang to not be the system python, but instead be the Isaac lab python; check out the Ray Dockerfile to see this.

In the context of py modules for Ray, I suspect that it may not work out of the box as it maybe just tries to pip install everything. That being said, maybe the modified Ray shebang that is already included helps address this issue; I'm not sure. Anyways, there's just some nuance here be aware of; I once burned a lot of time figuring it out myself ;) Maybe there is some way to specify which pip command to use, again I'm not sure you'd have to check the Ray docs.

Felix Yu @feyu-bdai showed me the Ray shebang trick, he may know how to do the pymodules and if the python thing is an issue

[ozhanozen]

Also, for the PR, @ozhanozen may be a good final reviewer in addition to myself (if he has time) as he is familiar with the Ray workflow having recently made improvements as well.

Hi @garylvov, sure, I would be glad to help and review if you assign me. Just keep in mind that I currently only have a single-GPU remote system, so I can test only non-cluster configs.

Congrats and good luck with the PhD! 🎓

[feyu-bdai]

You don't need to worry about the pip install issue when specifying py_modules.

What Ray does under the hood for any specified py_modules is zip up the local files of the module and uploads them directly to the Ray Worker under /tmp/ray/session_latest/runtime_resources/py_modules_files/_ray_pkg_<hash>. This module path is then appended to the beginning of the PYTHONPATH of any script launched by the Ray Job so any imports of the module will resolve to the uploaded module.

The good news is that means no pip install occurs: the module's files are directly referenced. The bad news is that there is no dependency resolution: any package specified through py_modules should already have their dependencies within the native environment.

The alternatives that Ray provides to py_modules is the pip flag, which actually tries the pip install. This runs into the issue Gary pointed out within this thread, and is not something I have tried to get working.

[feyu-bdai]

There are additional potential complications with how py_modules interacts with custom Omniverse Extensions as well, but those are more case-specific so I won't elaborate on it here. If users are struggling with getting their Extensions working with py_modules, I can then go into further depth.

[garylvov]

Thank you for the insight @feyu-bdai it's much appreciated!!!!!