feat: Add RLlib BC and MARWIL by noamonti748 · Pull Request #38 · GPUOpen-LibrariesAndSDKs/Schola

noamonti748 · 2026-08-21T16:06:05Z

Summary

The PR enables the use of RLlib's BC and MARWIL with Schola.

Can connect to a running Imitation Environment, automatically collect data to parquet episodes, and perform training.

Related issues

None.

Type of change

New feature or enhancement
Documentation

Changes

Add BC and MARWIL through schola rllib bc and schola rllib marwil
Store demonstrations as RLlib Parquet episodes
Add a collector that records gameplay from Unreal though the existing simulator subcommands
Run collection and training in one command. Can also train from an existing dataset with --input
Extend the RLlib train CLI, checkpouint export, and command template for offline workflows
Add an [offline] pip extra for RLlib offline dependencies
Add an imitation learning guide and tests for offline data, CLI, and training

Testing

Python (`Resources/python`, `Test`)

Installed test dependencies: pip install --group test -e "./Resources/python[all]"
Ran: python -m pytest Test --import-mode=importlib -n 0

Other Tests

Tested end-to-end pipeline on a version of the Basic Environment adapted for IL.
Tested on a Python process running a cartpole example. One over gRPC, and one that writes to a Parquet set directly.

Checklist

Code follows project style (Unreal coding standard for C++; Black for Python)
Comments / docstrings added or updated where behavior is non-obvious
README or Sphinx docs updated if user-facing behavior changed
No unrelated changes included in this PR
Appropriate license headers added to new files

Keep BC/MARWIL offline training on top of main's typing cleanup.

…ecutable for data collection

…in in one command

amd-alexcann · 2026-08-21T16:39:02Z


 """
 Cyclopts template for generating Schola subcommands (multi-simulator and algorithm dispatch).
 """


Why are you modifying command_template.py in this PR? This code has proven difficult for AI to interpret, so there is a good chance these changes were unnecessary, or extremely complicated for what should be a small change.

What issue does the changes to command_template solve?

It's for sure overkill and will cut back, but the one thing I think it solves is that it allows us to run without binding to a simulator, because the default template binds to external by default whenever a simulator command is not included. Would it be ok to just isolate this part and kill the rest (should be < 50 lines I think)

The command template should support this without changes (0-N algorithms, 0-N simulators). This case has tests as well so should be working. If you want to change the default sim, just override the function in the child class (it should be ignored when there is no supplied simulators).

If you really want to change it, making it str or None to indicate No Default should be fine. In general, the actual structure is rather delicate and bolting on features rather than building them in properly is error prone.

amd-alexcann · 2026-08-21T16:48:36Z

@@ -1,29 +1,43 @@
-# Copyright (c) 2024-2025 Advanced Micro Devices, Inc. All Rights Reserved.
+# Copyright (c) 2024-2026 Advanced Micro Devices, Inc. All Rights Reserved.


This is a bit of a mess, seems like a separate offline-train command is the way to go, alongside a collect command with functionality matching minari.

Split commands into these.

amd-alexcann · 2026-08-21T16:52:34Z

    )
+
+
+def load_rl_module_from_algorithm_checkpoint(


Does another version of this code exist elsewhere already? I believe this functionality is required by the export command (to load the policy before converting it to a Schola Model API).

Yep, that's true, but I think the issue is that RLModule load, where the rllib export goes through load_rl_module_from_algorithm_checkpoint or MultiRLModule.from_checkpoint is not shared between them. And of course, the load_rl... functionality is different.

What about for eval? The comment here says it shares layout handling with eval.

It may also be possible to make eval depend on this, ideally, we keep the number of distinct handlers for loading RLLib objects down lol. Between eval and rllib we already have a lot.

Feel free to make more drastic changes here if you see a way to make everything fit nicely. (e.g a generic Load function that everything else relies on in different places)

amd-alexcann · 2026-08-21T17:14:20Z

@@ -0,0 +1,480 @@
+# Copyright (c) 2026 Advanced Micro Devices, Inc. All Rights Reserved.


How much of this can be outsourced to RLLIB?
As an example, see the JSON Writer for writing the dataset to a file.

JSON Writer
Instructions Recommending Ray Data tools

JSONWriter/DatasetWriter seems to use the old-stack data format.
We can use the Parquet + msgpack over write_parquet though.
Also the SingleAgentEpisode

Sounds good to me.

My philosophy here is that it will be easier in the long run to use the rllib tools, rather than reimplement, where possible since any re-implementations will be coupled to the output of the functions. If the output format changes, we will need to adjust our tools to match.

- Update imitation_learning guide for generic BC/IL with Minari - `command_template` kept to original standard - Shared checkpoint loading path - Use RLlib functions for dataset writing

Noah Monti added 5 commits August 18, 2026 15:02

Add initial implementation of BC and MARWIL for RLlib

a195037

Merge origin/main into rllib-imitation-learning

34a0f6e

Keep BC/MARWIL offline training on top of main's typing cleanup.

update to use parquet dataset instead of Minari, ability to launch ex…

79fe0c9

…ecutable for data collection

adjust so that the IL process is pipelined, users can collect and tra…

61e84fe

…in in one command

improve path handling

c4e6103

noamonti748 changed the title ~~Noamonti748/rllib imitation learning~~ [feat] Add RLlib BC and MARWIL Aug 21, 2026

noamonti748 self-assigned this Aug 21, 2026

noamonti748 added the enhancement New feature or request label Aug 21, 2026

noamonti748 changed the title ~~[feat] Add RLlib BC and MARWIL~~ feat: Add RLlib BC and MARWIL Aug 21, 2026

Noah Monti added 2 commits August 21, 2026 12:10

update doc to increase coherency

3e812d6

black format and update copyright headers

7e846a4

noamonti748 marked this pull request as ready for review August 21, 2026 16:14

noamonti748 requested review from amd-alexcann and kainotkai August 21, 2026 16:16

amd-alexcann requested changes Aug 21, 2026

View reviewed changes

Noah Monti added 3 commits August 21, 2026 17:01

Apply first round of changes from PR comments.

89428ad

Second round of changes from PR comments:

50ba68a

- Update imitation_learning guide for generic BC/IL with Minari - `command_template` kept to original standard - Shared checkpoint loading path - Use RLlib functions for dataset writing

use pyarrow for rllib writing

3e8757e

noamonti748 requested a review from amd-alexcann August 24, 2026 19:43

		@@ -1,29 +1,43 @@
		# Copyright (c) 2024-2025 Advanced Micro Devices, Inc. All Rights Reserved.
		# Copyright (c) 2024-2026 Advanced Micro Devices, Inc. All Rights Reserved.

		@@ -0,0 +1,480 @@
		# Copyright (c) 2026 Advanced Micro Devices, Inc. All Rights Reserved.

		)


		def load_rl_module_from_algorithm_checkpoint(

Conversation

noamonti748 commented Aug 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Summary

Related issues

Type of change

Changes

Testing

Python (Resources/python, Test)

Other Tests

Checklist

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

noamonti748 Aug 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

noamonti748 Aug 21, 2026 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

noamonti748 commented Aug 21, 2026 •

edited

Loading

Python (`Resources/python`, `Test`)

noamonti748 Aug 21, 2026 •

edited

Loading

noamonti748 Aug 21, 2026 •

edited

Loading