Repository navigation
Add OpenArm-Puck-Language task - #13
Open
Manas-arumalla wants to merge 14 commits into
Open
Manas-arumalla wants to merge 14 commits into
Manas-arumalla wants to merge 14 commits into
Conversation
Right-arm contact-only pushing task: push a free-body puck into a fixed goal disc and settle it there. Registers two task_ids sharing the same rewards/terminations/events: OpenArm-Puck (privileged puck position observations) and OpenArm-Puck-Vision (a fixed overhead depth camera replaces the privileged terms for the actor, asymmetric actor-critic). Self-contained.
Tracking cameras follow the body's subtree COM, so rendered videos swayed with every arm motion. Same change as on the reach task.
Same change requested on the valve PR, applied here for consistency.
The puck is declared with friction 0.4 to make it slide, and that value reached neither contact it matters on. MuJoCo resolves a pair by taking the higher-priority geom's parameters, mixing only when priorities are equal, and that mix is an elementwise max. The right-arm finger geoms are priority=1, so they governed the push contact at mu=1.0; the table is priority=0 as the puck was, so that pair mixed to max(0.4, 1.0) = 1.0. The puck has therefore behaved exactly like a mu=1.0 object throughout, and dr_puck_friction randomized a number read on neither pair. priority=2 puts the puck above both. Measured on a two-geom model carrying the real finger and table parameters, effective friction goes from 1.0/1.0 to 0.4/0.4 at finger/table, and the randomization spans the declared range instead of collapsing to a point. condim=4 and the fingers' own solref are set so the finger contact model is unchanged and only the friction moves. The table pair does gain torsional friction (condim 3 to 4), which is the more faithful model for a box sliding on a surface. This changes how the task behaves, so it is a retrain, not just a config edit.
The puck task has one hardcoded target. Give it a small table of named targets and a per-env buffer selecting between them, so a caller can point individual envs at a different one. "right" is the coordinate the task already used and stays the default, so an existing run is unaffected -- there is an assert to keep it that way. "left" mirrors it across the puck's start line: same push distance, opposite lateral direction. Two things this touches: - reset_puck_uniform seeded the progress-shaping baseline from the module constant. With one fixed target that was the per-env goal; with two it is not, and the puck's spawn jitter then pays out a few centimetres of free progress on first contact, for whichever goal is not the constant. It now seeds from the per-env goal. - puck_to_midpoint_obs is an alternative observation for callers that let something else choose the target: it points at the midpoint of the two, which is the same kind of vector at the same scale but says nothing about which target is active.
The puck task has a fixed target. This adds a variant where a natural- language instruction chooses between two of them, so the policy has to read the sentence to know where to push. How it works. Instructions are embedded once, offline, by a frozen sentence encoder, and the result is saved as a lookup table. Training and evaluation read the table, so the encoder is a build-time dependency only -- it is in an optional extra, and nothing at runtime imports it. Each reset draws a goal and a phrasing for that goal, and the policy sees the embedding alongside its usual observation, projected down so it cannot swamp the proprioception. Two things the task has to get right or the result means nothing: - The observation must not already reveal the goal. The stock observation includes a vector from the puck to the active target, which answers the question the instruction is asking. The variant substitutes the vector to the midpoint of the two targets: same kind of quantity, same scale, identical for both goals. - Both goals must be reachable and worth reaching, or a policy that ignores the instruction loses nothing by doing so. Evaluation measures whether the sentence is doing the work, by scoring goal A while showing the instruction for goal B. The diagonal of that 2x2 is execution, the off-diagonal is suppression, and separation is the gap. A single success rate cannot tell the two apart, and a policy with a favourite goal scores well on one of them by accident. Held-out phrasings are disjoint from the training pool, so the number reported is generalisation to unseen wordings rather than recall of ten sentences. Also included: offline diagnostics that predict, before any GPU time is spent, whether the encoder separates the instructions and whether the observation leaks the goal; two console scripts wrapping them; and tests covering the samplers, the table format and its validation, the model's guard rails, and the evaluation's statistics.
…tioning # Conflicts: # src/openarm_mjlab/tasks/__init__.py
The result this task reports is a table of numbers. This makes it watchable: the viewer opens, the instruction alternates between the two goals, and each sentence is printed as it changes, so a reviewer can check the claim by eye instead of taking the table's word for it. Phrasings come from the held-out pool by default, so what is on screen is wording the policy never trained on. The goal can also be driven by hand -- typed into the terminal, not pressed in the viewer window. That looks like the wrong choice until you try the other one: MuJoCo binds every letter A-Z to a visualization toggle and every digit to a geom group, and mjlab binds several more on top, so there is no free key. A viewer keybinding silently does two things at once -- "h" also draws convex hulls, "n" also recolours the limbs by constraint island, "2" also hides the robot's 56 visual meshes. Each of those reads as a bug in the scene. stdin has no such conflict. Switching the goal touches the env, so it cannot happen on the reader thread while the viewer's thread is stepping. A typed command only records what was asked for; the change is applied inside the policy, which the viewer calls on its own stepping thread. Also adds a test that every declared console script resolves to a callable. A typo in pyproject.toml otherwise surfaces only when a user runs the command, and the entry points are the documented way in.
The check asks whether the observation gives the goal away. At the old default of 64 envs x 6 resets it answered that from 384 samples, where the standard error on a proportion near 0.5 is ~0.026. A measured run returned 0.598 -- clearing the 0.10 margin by 0.002 and printing 'at chance'. The same environment measures 0.501 at 3072 samples. Raises the default to 256 x 12 and prints the sample count and standard error next to the result, so the reader can see what the number is worth. Adds a middle verdict: when the margin exceeds three standard errors the check no longer says 'at chance', even though it is under the 0.10 bar. It reports the margin and asks for a re-run with more samples instead. Takes 38s on CPU.
…sk has landed This branch was built on the puck PR's commits as they stood at the time, plus language-specific work, so with the puck task now merged into main the two copies of the puck files collide. Main's copy is the canonical, reviewed one, but this branch also changes the puck task on purpose: 2b79bf7 adds named goals, a per-env goal buffer, set_puck_goal and puck_to_midpoint_obs, which the language task depends on. Taking either side wholesale would be wrong -- main's alone drops that feature and breaks the language task, and this branch's alone reverts the review fixes in main. So the puck files are main's version with 2b79bf7 re-applied on top. The patch applies cleanly, and the names its new code uses are all imported in main's file. puck_env_cfg.py and the puck test are main's, and tasks/__init__.py registers both the puck and the language tasks. Main's review changed the puck's contact model from condim 4 to 3 and refactored two reward and observation helpers into shared ones. The helpers were checked to be equivalent. For the contact model change, the language evaluation was re-run on the merged code with the same command, splits and seed: train separation is +1.000 in all three runs, suppression is unchanged, and held-out separation moves by at most 0.042 per run (mean +0.722 before, +0.736 after). The language observation layout is identical before and after, so existing checkpoints load. The full test suite passes, 68 tests.
…language branch Conflicts: tasks/__init__.py keeps both the language and lift imports; puck reset keeps the per-episode goal for the progress baseline and uses the new env_buffer helper. puck_to_midpoint_obs now uses object_pos_w and to_base_frame from common_mdp, which replaced the helpers it called.
…tioning # Conflicts: # pyproject.toml
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.
OpenArm-Puck-Language: the puck task, but with two targets, and a written instruction tells the arm which one to push to. Nothing else in the observation says which so the policy has to read the sentence.
Sentences are turned into vectors once, up front, by a frozen text encoder (bge-base-en-v1.5), and saved to a lookup file. Training and eval just read that file, so the encoder is only needed at build time. It's an optional extra and nothing at runtime imports it.
Result
"New wording" means phrasings the policy never saw in training.
Three separate training runs, evaluated on this branch:
The new-wording gap varies by about one episode between training runs and between evaluations of the same checkpoint, so it's a range, not a single number.
A plain success rate can't tell apart a policy that follows instructions and one that always goes the same way and is right half the time. So the policy is scored on one target while being shown the other target's instruction.
Two separate things fall out: does it go where told, and does it avoid the other target. On new wording, only the first one dips (0.708 to 0.875) while the second stays high (0.958 in every run) so it isn't mixing the instructions up, it just sometimes fails to finish the job. One number would have hidden that.
The observation must not give away the answer. The normal puck observation has a vector pointing at the active target, which is exactly what the instruction is meant to tell you. This swaps in a vector to the midpoint of the two targets: same kind of value, same size, but the same for both goals.
Both targets must be reachable and worth reaching, or ignoring the instruction costs nothing.
"right" is the target the task already used and stays the default (there's an assert to keep it that way), so existing runs don't change. "left" mirrors it same distance, other direction.
Running it
Watching it
The result above is a table. openarm-mjlab-lang-play makes it watchable: the viewer opens, the instruction alternates between the two goals on its own, and each sentence is printed as it changes so you can check the claim by eye rather than take the table's word for it. It uses held-out phrasings by default, so what's on screen is wording the policy never trained on.
You can also drive it by hand: type l, r, n, t, h into the terminal and press Enter.