Conversation
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.
🔗 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 FailureAs of commit 487b112 with merge base 2d258fe ( NEW FAILURE - The following job has failed:
This comment was automatically generated by Dr. CI and updates every 15 minutes. |
Collaborator
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.
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.
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
Description
Follow-up to #4386, stacked on it: the base branch is
menagerie-envin this repository (a mirror of #4386's head), so the diff is this PR's own delta. Retarget tomainonce #4386 merges. It implements the suggestion from that review: a model-resolution layer that is not tied to Menagerie.MujocoModelEnvis the bare-simulator env of [Feature] MenagerieEnv: any MuJoCo Menagerie robot by name #4386 (qpos,qvel,sensordataandsite_positionsobservations, keyframe reset,MujocoModelTaskconfig with the hold-pose terms) taking a model source: the path of an XML, or any object withresolve(download=...)(theModelSourceprotocol).download=Trueremains 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)namesowner/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_TOKENis used when set) on the firstdownload=True, and the mapping is cached underrefs/, 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, thenTORCHRL_MUJOCO_MENAGERIE_PATH, then themujoco-menageriepackage cache), extracted fromMenagerieEnv.resolve_model, which now delegates to it.MenagerieEnvbecomes a thin subclass ofMujocoModelEnvover that source. Its constructor, attributes and behaviour are unchanged;MenagerieTaskstays as an alias ofMujocoModelTask.The pre-resolution step lives in
_MujocoMetaitself: after validating the backend and batching arguments and before building any env or worker, it calls the class'_resolve_before_batchinghook (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 aliasModelSource = str | Pathin_backends.pyis renamedXmlSourceso the public protocol keeps the name.examples/menagerie/ppo.pytakes--repo,--revisionand--entry, so the same PPO recipe andrlrenderfactories work on any MuJoCo model on GitHub throughMujocoModelEnvandGitHubModelSource.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, resolvesmainto that commit and caches it, passescheck_env_specs, resets to thehomekeyframe and rolls out. A wrong entry reports the tree's XML files.Tests
test_github_model_sourcemocks the archive download and the API call, and asserts the cache layout by owner, repository and commit, that nothing is fetched withoutdownload=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.Docs
envs_api.rst: paragraph and autosummary entries forMujocoModelEnv,MujocoModelTask,GitHubModelSource,MenagerieModelSourceandModelSource.