Repository navigation
Add a shared multi-task policy across five tasks - #11
Open
Manas-arumalla wants to merge 12 commits into
Open
Manas-arumalla wants to merge 12 commits into
Manas-arumalla wants to merge 12 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.
One policy, one set of weights, trained on reach, valve, door, drawer and puck at once. Lift is excluded: its effort-controlled finger gives it a different action space from the 8-dim joint-position action the other five share. Three pieces, none of which works alone: Disjoint observation slots. Every task's object block previously sat at the same indices, padded to the widest, so one index carried valve_angle for valve, door_angle for door, drawer_pos for drawer and a Cartesian offset for reach and puck. rsl_rl keeps one running mean/std per index pooled over all tasks, so those statistics described a mixture of unrelated quantities, and an index that was a structural zero for most tasks had its std collapse and amplify whichever task used it (index 41 std 0.006, a 4.3 sigma distortion for reach and 7.2 for puck). Reserving a range per task costs 18 dims and makes every index mean one thing. The padding also goes before "actions" so that term lands at a fixed column for all five; appending it last put "actions" at a different raw index per task and ballooned the normalizer std there to 5-15 against 0.1-0.5 in the aligned joint region. Concurrent batches. Training one task at a time and varying the schedule cannot remove catastrophic forgetting, because sequential fine-tuning is the structure that causes it. MultiTaskMixtureVecEnv holds one env per task and presents them as a single vectorised env, so every PPO batch spans all five and there is nothing to forget between them. Allocation can be weighted, since an even split starves the hardest task. Per-task capacity and advantage normalization. A single Linear(128 -> 8) head per task is little private capacity over a trunk serving five behaviours; each head becomes Linear + ELU + output, giving 17,544 private parameters per task instead of 1,032. And stock PPO standardizes advantages once over the whole batch, so with five reward scales present the largest dominates and the rest scale into noise: measured, reach alone learned while the other four sat at exactly 0.000 for 1000 iterations. MultiTaskPPO standardizes within each task, identifying tasks from the observation one-hot so it stays correct under unequal allocation. Measured with the tasks' own termination terms, 256 episodes each, at the even split this ships with: reach 1.000, valve 1.000, door 1.000, puck 0.977, drawer 0.000, mean 0.795, against 0.437 for the best configuration without these three changes. A second seed gives 1.000/1.000/0.996/0.918/0.000, mean 0.783, so the two agree to within 1.2 points. Drawer does not transfer and is a known exclusion rather than an open question. It is not data starvation: five times the environments moved its mean opening from 7.0mm to 7.3mm of the 90mm required. Its success gate wants opening > 90mm AND speed < 0.15 m/s AND contact simultaneously, and the only policy that has ever satisfied it yanks the drawer open in five steps at 1.7 m/s and is quasi-static only because it has hit the end stop. A shared policy does not find that corridor.
rsl_rl's JIT and ONNX wrappers run obs_normalizer and then mlp directly, never calling get_latent -- which is the only place the task one-hot is read. torch.jit.trace therefore froze whichever index the tracing input produced, and an all-zero observation gives head 0, so every exported policy ran reach's head no matter which task it was handed. Nothing raised; the actions were simply wrong. Override as_jit and as_onnx with wrappers that take the argmax over the one-hot on the raw observation, so the selection becomes a node in the traced graph rather than a Python value captured at trace time. The new tests fail on the previous code, where the exported valve policy disagrees with the eager model by 510.0, and pass here.
The previous commit described the wrong mechanism and only half fixed the problem. Checking what mjlab actually calls, rather than assuming, gives a different picture for each export path. export_policy_to_jit runs torch.jit.script, not torch.jit.trace. On the original code it does not silently pick the wrong head -- it raises "Module '_MultiHead' has no attribute '_owner'" and produces nothing, so JIT export of a multi-task checkpoint has never worked. The wrapper added in the previous commit did not fix that either: TorchScript cannot close over the module-level NUM_TASKS, nor comprehend over an nn.ModuleList. The task count is now an instance attribute and the heads run in a plain loop, and torch.jit.script succeeds with output identical to the eager model. export_policy_to_onnx is the path that fails quietly. It traces against an all-zero dummy observation, so argmax over an all-zero one-hot fixes the head index at 0 and every exported policy runs reach's head. That part of the previous commit's description was right. The tests now exercise those two paths instead of torch.jit.trace, which is not what either exporter uses. Against the original code the JIT test fails on the _owner error and the ONNX test on the wrong head.
This branch predates the puck PR, so it carried its own copy of the puck task. With that merged into main the two copies collide, and main's is the canonical one: it is strictly newer, adding fingers_on_handle_obs, terminated_by and an __all__ to the mdp module, fifteen lines of contact-model reasoning to the env config, and a stricter test that push_contact never reaches the actor. So the three add/add conflicts in puck/mdp.py, puck/puck_env_cfg.py and tests/test_puck_env.py all resolve to main's version, and this branch keeps only its multitask additions. tasks/__init__.py merged cleanly and registers both. Verified on the merged tree: OpenArm-Puck and the multitask tasks all build, and the full suite passes, 51 tests.
…multitask branch Conflict in tasks/__init__.py: keep both the lift and multitask imports.
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.
One policy, one set of weights, trained on reach, valve, door, drawer and puck at once. Each task keeps its own rewards, terminations and actions; only observations are unified, into disjoint per-task slots plus a task-identity one-hot. Three changes carry it, none of which works alone: reserved observation slots so a shared normalizer never pools unrelated quantities, a mixture environment putting all five tasks in every PPO batch, and per-task output heads over a shared trunk with advantages standardized within each task. Registers
OpenArm-MultiTask-{Reach,Valve,Door,Drawer,Puck}and anopenarm-mjlab-multitask-trainentry point.Results: 256 held-out episodes per task, scored on each task's own termination term, over three seeds reach 1.000, valve 0.996, puck 0.991, door 0.607, drawer 0.003.The valve is turned by pressing the lever with one jaw rather than grasping it,