[Feature] Add support for loading datasets from local Minari cache

[Ibinarriaga8]

Description

This PR adds support for loading datasets directly from the local Minari cache in the MinariExperienceReplay class. Specifically, it introduces the load_from_local_minari argument, which, when set to True, instructs the class to load the dataset from the user's local Minari cache (typically at ~/.minari/datasets) and skip any fetching from the Minari server (i.e., no remote download or overwrite will occur). After loading from the local cache, all subsequent preprocessing and loading steps continue as usual, ensuring the dataset is processed and made available correctly. This is especially useful for custom Minari datasets or when you want to avoid network access.

Documentation has been updated in the class docstring to clearly state the new behavior of the load_from_local_minari argument, including details about local cache prioritization and unchanged downstream preprocessing.

This PR also includes comprehensive test coverage for the new feature, confirming that datasets created and stored in the local Minari cache can be loaded, sampled, and validated for correctness using the new argument. The provided test (test_local_minari_dataset_loading) creates a custom dataset, loads it from cache, verifies sample integrity, and cleans up afterwards.

Motivation and Context

Previously, MinariExperienceReplay required datasets to be downloaded via its own interface, which was incompatible with custom or preloaded Minari datasets, attempting to load these would result in a FileNotFoundError.
This change allows users to work with their own datasets, datasets created with minari.DataCollector(...).create_dataset(...), or any dataset present in the local Minari cache, without requiring redundant downloads or manual metadata copying.

If the dataset is not found in the local cache, a FileNotFoundError is raised with a clear message.

• I have raised an issue to propose this change (required for new features and bug fixes)

Solves #3067

Types of changes

• New feature (non-breaking change which adds core functionality)
• Documentation (update in the documentation)

Checklist

• I have read the CONTRIBUTION guide (required)
• My change requires a change to the documentation.
• I have updated the tests accordingly (required for a bug fix or a new feature).
• I have updated the documentation accordingly.

[pytorch-bot]

🔗 Helpful Links

🧪 See artifacts and rendered test results at hud.pytorch.org/pr/pytorch/rl/3068

• 📄 Preview Python docs built from this PR

Note: Links to docs will display an error until the docs builds have been completed.

❌ 27 New Failures, 1 Pending, 4 Unrelated Failures

As of commit cc43f9a with merge base 77c00b9 ():

NEW FAILURES - The following jobs have failed:

• Habitat Tests on Linux / tests (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t de3a7f21e58a8af9782d474033edc1a1466988cd8519cd24247bc11fb3f9c9a5 /exec failed with exit code 1
• Libs Tests on Linux / unittests-brax (3.11, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t b9bf9223f8bf4f5d8f57c6f9a718ba2137f684d62d775dc6a428bb71b423d304 /exec failed with exit code 2
• Libs Tests on Linux / unittests-d4rl (3.10, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 9d900238c51f398f88c04a2bb79a0634d00277f762ee7cfaa5abb5cd7ff214ae /exec failed with exit code 2
• Libs Tests on Linux / unittests-envpool (3.10, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 37007adc7d0a1d4b189efec36ffe7efb14ddae9d93215dcd9ce4720f12fb7fc7 /exec failed with exit code 2
• Libs Tests on Linux / unittests-gendgrl (3.10, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 721c8ce36d38a21b8ac546e0b88e1493885742141af3448447866d815b216ed0 /exec failed with exit code 2
• Libs Tests on Linux / unittests-gym (3.9, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t fa5035e966b972d91d9d4c2cff7fa5094fd46c1816a198e28a1067b4d94358f6 /exec failed with exit code 2
• Libs Tests on Linux / unittests-jumanji (3.10, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 3154a033f1d5180aad17365610f4ae3f0ef896acaca175250b269fe8cf1b4693 /exec failed with exit code 2
• Libs Tests on Linux / unittests-meltingpot / linux-job (gh)
  RuntimeError: Command docker exec -t 14d584b5a34652fcdf2884bc9451494689dd55440da200841ed1fa9e6799ed67 /exec failed with exit code 2
• Libs Tests on Linux / unittests-minari (3.10, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 1d3ae4797830e6281e7bc140cdb5d07880fb4ec441f265ab6346064906e19db0 /exec failed with exit code 2
• Libs Tests on Linux / unittests-open_spiel (3.10, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t f7b624b8405809af3d667a3bee8f6d468b2612b98a2b5345b5133a31163dfb84 /exec failed with exit code 2
• Libs Tests on Linux / unittests-openx (3.10, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t f528abea536420b920c21e4b87ed08e2420820135813956ae886ec45187be215 /exec failed with exit code 2
• Libs Tests on Linux / unittests-pettingzoo / linux-job (gh)
  RuntimeError: Command docker exec -t a5063ed7c534872e739d4386ffa592f1ed6cf2697ed9015ea5e21fefcb18a618 /exec failed with exit code 2
• Libs Tests on Linux / unittests-roboset (3.10, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t a007486b65b770d60bf85eeaac96eca42bccd829aa12f7f46fb219e4291c453e /exec failed with exit code 2
• Libs Tests on Linux / unittests-sklearn (3.10, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t c14552bb0d36d7cacf56bf637fce434138ce44bbc3fd3a737cb18b16567337e3 /exec failed with exit code 2
• Libs Tests on Linux / unittests-smacv2 (3.10, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 9ea3576922b9d08fd2fcd56dc8193b9df9920ad4f28265b06227cc72f775dbda /exec failed with exit code 2
• Libs Tests on Linux / unittests-unity_mlagents (3.10.12, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t fc9febe1ec8b0de222fce33d8095dd3bf635bf2958c43accf00cf5d0d41ffa04 /exec failed with exit code 2
• Libs Tests on Linux / unittests-vd4rl (3.10, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 633cce66d6c8774e481f1fb531ab9187ce80ce2f8d117f6b72104826025d8e3d /exec failed with exit code 2
• Libs Tests on Linux / unittests-vmas (3.10, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t bfbede12d1a5459bbb9f371673859ae91eceb65b3506e0285c8882ce19bd0c99 /exec failed with exit code 2
• Unit-tests on Linux / tests-cpu (3.10) / linux-job (gh)
  RuntimeError: Command docker exec -t 74c50f455ce42b44bd283db218e295c1f1e257540c11bef14222bdd34584474a /exec failed with exit code 2
• Unit-tests on Linux / tests-cpu (3.11) / linux-job (gh)
  RuntimeError: Command docker exec -t f0803bc3acc7eecf1cbc5f6a3ff91c5a23d9028a3858543b6597158eaf7804f1 /exec failed with exit code 2
• Unit-tests on Linux / tests-cpu (3.12) / linux-job (gh)
  RuntimeError: Command docker exec -t e10903c9290e8ce6083042309d132173504bde767b4883184de8e5d5a2f7ef6d /exec failed with exit code 2
• Unit-tests on Linux / tests-cpu (3.9) / linux-job (gh)
  RuntimeError: Command docker exec -t 84d20cd621efe4b7c3ad79ceec2a5b3211c8160292f38ba0abe255d2741e0b20 /exec failed with exit code 2
• Unit-tests on Linux / tests-gpu (3.11, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t e0ff8fd316a3f3caf50b871622c5e0197ded3fce6114190d17930c435d771a8f /exec failed with exit code 2
• Unit-tests on Linux / tests-olddeps (3.9, 11.6) / linux-job (gh)
  RuntimeError: Command docker exec -t bd63dfa8caf7fc530d7e56038aa6c2426f94527913f1a670111b4d1a431eec50 /exec failed with exit code 2
• Unit-tests on Linux / tests-optdeps (3.11, 12.8) / linux-job (gh)
  RuntimeError: Command docker exec -t ef37b89dabeef9ac6b05a7ba256744e6325ecc537725145219f4bd0a86406718 /exec failed with exit code 2
• Unit-tests on Linux / tests-stable-gpu (3.10, 11.8) / linux-job (gh)
  RuntimeError: Command docker exec -t 0444b45ee452d96b7fcdf21bd9218fda0021720e000b46fdbf0a6d8acb2fd3db /exec failed with exit code 2
• Unit-tests on Windows / unittests-cpu (3.10, windows.4xlarge, cpu) / windows-job (gh)
  Process completed with exit code 2.

FLAKY - The following job failed but was likely due to flakiness present on trunk:

• Continuous Benchmark (PR) / GPU Pytest benchmark (gh) (similar failure)
  RuntimeError: PassManager::run failed

BROKEN TRUNK - The following jobs failed but was present on the merge base:

👉 Rebase onto the `viable/strict` branch to avoid these failures

• Continuous Benchmark / GPU Pytest benchmark (gh) (trunk failure)
  RuntimeError: PassManager::run failed
• Libs Tests on Linux / unittests-robohive (3.9, 12.8) / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.
• LLM Tests on Linux / unittests (3.9, 12.8) / linux-job (gh) (trunk failure)
  ##[error]The operation was canceled.

This comment was automatically generated by Dr. CI and updates every 15 minutes.

[Ibinarriaga8]

Custom minari dataset creation from a gymnasium environment

[vmoens]

That's really handy I love it!

[vmoens]

@Ibinarriaga8

Do you think we have an opportunity here to reduce the number of datasets we download to test minari, and use customly built datasets instead like you do in your test?

[Ibinarriaga8]

@Ibinarriaga8

Do you think we have an opportunity here to reduce the number of datasets we download to test minari, and use customly built datasets instead like you do in your test?

Absolutely, instead of downloading 20 datasets from the Minari server, we could generate smaller, custom datasets from any gymnasium environment as part of our tests. This approach is especially valuable for D4RL datasets, which tend to be very large and can significantly slow down testing.

[vmoens]

This is a global import - we must avoid these at all cost.
Can you make it local?

[Ibinarriaga8]

Ok, I can fix it

[vmoens]

Absolutely, instead of downloading 20 datasets from the Minari server, we could generate smaller, custom datasets from any gymnasium environment as part of our tests. This approach is especially valuable for D4RL datasets, which tend to be very large and can significantly slow down testing.

Do you want to give it a go? Otherwise I can do it no worry

[Ibinarriaga8]

Absolutely, instead of downloading 20 datasets from the Minari server, we could generate smaller, custom datasets from any gymnasium environment as part of our tests. This approach is especially valuable for D4RL datasets, which tend to be very large and can significantly slow down testing.

Do you want to give it a go? Otherwise I can do it no worry

Yes, no problem

[Ibinarriaga8]

@vmoens Do you think we should replace testing with all Minari datasets by using smaller, custom datasets generated from gymnasium environments? Or should we still download and test with some datasets from the Minari server to ensure that the downloading functionality in MinariExperienceReplay works correctly?

[vmoens]

@vmoens Do you think we should replace testing with all Minari datasets by using smaller, custom datasets generated from gymnasium environments? Or should we still download and test with some datasets from the Minari server to ensure that the downloading functionality in MinariExperienceReplay works correctly?

Maybe just one defined (not random) dataset - to make sure?

[Ibinarriaga8]

@vmoens Do you think we should replace testing with all Minari datasets by using smaller, custom datasets generated from gymnasium environments? Or should we still download and test with some datasets from the Minari server to ensure that the downloading functionality in MinariExperienceReplay works correctly?

Maybe just one defined (not random) dataset - to make sure?

Ok, I agree that testing with a single, well-defined dataset makes sense to ensure the download and loading functionality are covered.

However, with most gym environments, creating custom datasets isn’t always straightforward. For example, I ran into this error:

ValueError: Dict key mismatch; expected keys: ['reward_ctrl', 'reward_forward', 'reward_survive', 'x_position', 'x_velocity', 'z_distance_from_origin']; dict: {'x_position': np.float64(0.001823518632481435), 'z_distance_from_origin': np.float64(-0.004461789811977868)}

I’ll keep investigating to see if there’s a workaround or a more general approach that works across environments.

[vmoens]

Ok maybe we could proceed with this in its current state and refactor the tests later!

[Ibinarriaga8]

Ok maybe we could proceed with this in its current state and refactor the tests later!

Ok, I have refactored the tests to use this custom datasets approach. I am still fixing some minor errors and making sure it all works.
I can either commit these changes to this PR or create a separate PR for the test refactor. Whichever you prefer

[vmoens]

If you're ready feel free to commit here directly

[vmoens]

LGTM!
This actually makes me think that we should have a similar pipeline where one creates a datacollector with some replay buffer and serializes it as a dataset later on.
It's mostly a matter of documenting it rather than creating the feature, everything already exists I think!

[Ibinarriaga8]

LGTM! This actually makes me think that we should have a similar pipeline where one creates a datacollector with some replay buffer and serializes it as a dataset later on. It's mostly a matter of documenting it rather than creating the feature, everything already exists I think!

Ok, should I provide a tutorial documenting the pipeline from DataCollector with a replay buffer to dataset serialization?