Skip to content

[Feature] MujocoModelEnv with GitHub and Menagerie model sources - #4401

Draft
theap06 wants to merge 5 commits into
pytorch:menagerie-envfrom
theap06:mujoco-model-source
Draft

theap06 wants to merge 5 commits into
pytorch:menagerie-envfrom
theap06:mujoco-model-source

Conversation

@theap06

@theap06 theap06 commented Sep 16, 2026

Copy link
Copy Markdown
Collaborator

Description

Follow-up to #4386, stacked on it: the base branch is menagerie-env in this repository (a mirror of #4386's head), so the diff is this PR's own delta. Retarget to main once #4386 merges. It implements the suggestion from that review: a model-resolution layer that is not tied to Menagerie.

  • MujocoModelEnv is the bare-simulator env of [Feature] MenagerieEnv: any MuJoCo Menagerie robot by name #4386 (qpos, qvel, sensordata and site_positions observations, keyframe reset, MujocoModelTask config with the hold-pose terms) taking a model source: the path of an XML, or any object with resolve(download=...) (the ModelSource protocol). download=True remains the explicit permission for network access, and the metaclass resolves the source once before the native backend spawns workers.

  • GitHubModelSource(repo, revision, entry, root=None) names owner/name, a pinned revision (commit SHA, tag or branch) and an explicit repository-relative entry, so nothing is guessed. The whole repository tree at that revision is fetched as a GitHub zip archive and cached under ~/.cache/torchrl/github_models/owner/name/<commit> with an atomic move, so includes and sibling asset directories resolve as in a checkout. A tag or branch is resolved to its commit through the GitHub API (GITHUB_TOKEN is used when set) on the first download=True, and the mapping is cached under refs/, so later resolutions are offline and stay pinned until the cache entry is removed. A missing entry lists the XML files the tree contains instead of picking one.

  • MenagerieModelSource(robot, entry=None, menagerie_path=None) is the curated source (checkout, then TORCHRL_MUJOCO_MENAGERIE_PATH, then the mujoco-menagerie package cache), extracted from MenagerieEnv.resolve_model, which now delegates to it.

  • MenagerieEnv becomes a thin subclass of MujocoModelEnv over that source. Its constructor, attributes and behaviour are unchanged; MenagerieTask stays as an alias of MujocoModelTask.

  • The pre-resolution step lives in _MujocoMeta itself: after validating the backend and batching arguments and before building any env or worker, it calls the class' _resolve_before_batching hook (identity by default). That removes the per-subclass metaclasses (_MujocoModelMeta, and the existing _MicroDuckMeta), makes argument errors precede any download, and lets MicroDuck's downloader reuse the shared GitHub tree fetcher. The backend alias ModelSource = str | Path in _backends.py is renamed XmlSource so the public protocol keeps the name.

  • examples/menagerie/ppo.py takes --repo, --revision and --entry, so the same PPO recipe and rlrender factories work on any MuJoCo model on GitHub through MujocoModelEnv and GitHubModelSource.

Left for later, as suggested: a manifest format and conservative discovery of a unique compile-valid root.

Verified on the repository cited in the review: MujocoModelEnv(GitHubModelSource("SouthColumn76/universal_robots_ur3e", revision="5f042ffca6b5885fd18f5448e17b71ab46274fa3", entry="ur3e.xml"), download=True) downloads the tree in about two seconds, resolves main to that commit and caches it, passes check_env_specs, resets to the home keyframe and rolls out. A wrong entry reports the tree's XML files.

Tests

  • test_github_model_source mocks the archive download and the API call, and asserts the cache layout by owner, repository and commit, that nothing is fetched without download=True, pinned offline resolution afterwards, no API call for a full SHA, the entry error, argument validation, and one download for batched native workers.
  • test_mujoco_model_env_from_path; the existing Menagerie tests pass unchanged apart from one error message.
  • Doctests of the three modules.

Docs

  • envs_api.rst: paragraph and autosummary entries for MujocoModelEnv, MujocoModelTask, GitHubModelSource, MenagerieModelSource and ModelSource.

MujocoModelEnv is the bare simulator of any MuJoCo model, from the path
of an XML or a model source. GitHubModelSource pins a repository to a
revision with an explicit repository-relative entry and caches the whole
tree by owner, repository and commit; MenagerieModelSource locates a
Menagerie robot in a checkout or the mujoco-menagerie package cache.
MenagerieEnv becomes a thin subclass over that source, and MenagerieTask
an alias of MujocoModelTask.
@pytorch-bot

pytorch-bot Bot commented Sep 16, 2026

Copy link
Copy Markdown

🔗 Helpful Links

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

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

❌ 1 New Failure

As of commit 487b112 with merge base 2d258fe (image):

NEW FAILURE - The following job has failed:

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

@meta-cla meta-cla Bot added the CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. label Sep 16, 2026
@github-actions github-actions Bot added Documentation Improvements or additions to documentation CI Has to do with CI setup (e.g. wheels & builds, tests...) Environments Adds or modifies an environment wrapper Examples Feature New feature labels Sep 16, 2026
@vmoens

vmoens commented Sep 16, 2026

Copy link
Copy Markdown
Collaborator

the diff includes its commits until it merges

Now that you're maintainer you can push to PyTorch (and not your fork) to make "real" gh PR stacks!

Move the pre-resolution hook into _MujocoMeta, which now validates the
backend and batching arguments before calling MujocoEnv._resolve_before_batching
and building anything, so argument errors precede any download; this
removes _MujocoModelMeta and _MicroDuckMeta, and MicroDuck's downloader
becomes a call to the shared GitHub tree fetcher. GitHubModelSource
validates owner and name segments, ignores torn ref files and writes
them atomically, fetches archives through urlopen with a timeout and
the GITHUB_TOKEN when set, and its cache root reuses the datasets
helper. The backend alias ModelSource becomes XmlSource so the public
protocol keeps the name, _resolve_model_source is private and lets URLs
through to MujocoEnv, MenagerieEnv keeps its typed task, backend and
max_episode_steps keywords, and error messages name the robot again.
Carries the review fixes of pytorch#4386 into the model-source layer: the
entry check for a file menagerie_path and the FileNotFoundError for a
robot the package cannot fetch move to MenagerieModelSource, and the
substep-lag, floating-base and hold_pose_task docstrings to
MujocoModelEnv.
@theap06
theap06 changed the base branch from main to menagerie-env September 16, 2026 22:41
Deriving the GitHub cache root from torchrl.data.datasets.utils made
torchrl.envs import torchrl.data while it was initializing, a circular
import whenever torchrl.collectors was imported first; the root is the
same path spelled locally, and a subprocess test pins the import order.
The PPO example takes --repo, --revision and --entry to train on any
MuJoCo model on GitHub through MujocoModelEnv and GitHubModelSource.
Brings the walk-task fixes and the examples README; the README gains
the GitHub-source recipe on the UR3e.
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

CI Has to do with CI setup (e.g. wheels & builds, tests...) CLA Signed This label is managed by the Facebook bot. Authors need to sign the CLA before a PR can be reviewed. Documentation Improvements or additions to documentation Environments Adds or modifies an environment wrapper Examples Feature New feature

Projects

None yet

Development

Successfully merging this pull request may close these issues.

2 participants