Upload folder using huggingface_hub

.gitattributes

@@ -33,3 +33,4 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+replay.mp4 filter=lfs diff=lfs merge=lfs -text

.summary/0/events.out.tfevents.1688740849.qgallouedec-MS-7C84

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:f6ff0eaa13758cff09bdd058ae223b8b01908d06e46c5955d72f79ec1d5483e6
+size 700655

README.md

@@ -0,0 +1,56 @@
+---
+library_name: sample-factory
+tags:
+- deep-reinforcement-learning
+- reinforcement-learning
+- sample-factory
+model-index:
+- name: APPO
+  results:
+  - task:
+      type: reinforcement-learning
+      name: reinforcement-learning
+    dataset:
+      name: bin-picking-v2
+      type: bin-picking-v2
+    metrics:
+    - type: mean_reward
+      value: 168.59 +/- 207.20
+      name: mean_reward
+      verified: false
+---
+
+A(n) **APPO** model trained on the **bin-picking-v2** environment.
+
+This model was trained using Sample-Factory 2.0: https://github.com/alex-petrenko/sample-factory.
+Documentation for how to use Sample-Factory can be found at https://www.samplefactory.dev/
+
+
+## Downloading the model
+
+After installing Sample-Factory, download the model with:
+```
+python -m sample_factory.huggingface.load_from_hub -r qgallouedec/bin-picking-v2
+```
+
+    
+## Using the model
+
+To run the model after download, use the `enjoy` script corresponding to this environment:
+```
+python -m enjoy --algo=APPO --env=bin-picking-v2 --train_dir=./train_dir --experiment=bin-picking-v2
+```
+
+
+You can also upload models to the Hugging Face Hub using the same script with the `--push_to_hub` flag.
+See https://www.samplefactory.dev/10-huggingface/huggingface/ for more details
+    
+## Training with this model
+
+To continue training with this model, use the `train` script corresponding to this environment:
+```
+python -m train --algo=APPO --env=bin-picking-v2 --train_dir=./train_dir --experiment=bin-picking-v2 --restart_behavior=resume --train_for_env_steps=10000000000
+```
+
+Note, you may have to adjust `--train_for_env_steps` to a suitably high number as the experiment will resume at the number of steps it concluded at.
+    

checkpoint_p0/best_000018160_9297920_reward_202.428.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:83cb11421e93c30a4d5668d31b2d615e4a6526c4f88f6ac2ecd1ed5631765ebd
+size 98239

checkpoint_p0/checkpoint_000019296_9879552.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:cf8931a2ed9a5502496b448578eb87f863bb3703cafccfab6fc9171810e860e4
+size 98567

checkpoint_p0/checkpoint_000019544_10006528.pth

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:5f3aad58d66e16763d6e8c9e7dcdc53df4f7c7c3f9fb5e8bcd847343ff9ab654
+size 98567

config.json

@@ -0,0 +1,134 @@
+{
+  "help": false,
+  "algo": "APPO",
+  "env": "bin-picking-v2",
+  "experiment": "bin-picking-v2",
+  "train_dir": "/home/qgallouedec/data/gia/data/envs/metaworld/train_dir",
+  "restart_behavior": "restart",
+  "device": "cpu",
+  "seed": null,
+  "num_policies": 1,
+  "async_rl": false,
+  "serial_mode": false,
+  "batched_sampling": false,
+  "num_batches_to_accumulate": 2,
+  "worker_num_splits": 2,
+  "policy_workers_per_policy": 1,
+  "max_policy_lag": 1000,
+  "num_workers": 8,
+  "num_envs_per_worker": 8,
+  "batch_size": 1024,
+  "num_batches_per_epoch": 4,
+  "num_epochs": 2,
+  "rollout": 64,
+  "recurrence": 1,
+  "shuffle_minibatches": false,
+  "gamma": 0.99,
+  "reward_scale": 0.1,
+  "reward_clip": 1000.0,
+  "value_bootstrap": true,
+  "normalize_returns": true,
+  "exploration_loss_coeff": 0.0,
+  "value_loss_coeff": 1.3,
+  "kl_loss_coeff": 0.1,
+  "exploration_loss": "entropy",
+  "gae_lambda": 0.95,
+  "ppo_clip_ratio": 0.2,
+  "ppo_clip_value": 1.0,
+  "with_vtrace": false,
+  "vtrace_rho": 1.0,
+  "vtrace_c": 1.0,
+  "optimizer": "adam",
+  "adam_eps": 1e-06,
+  "adam_beta1": 0.9,
+  "adam_beta2": 0.999,
+  "max_grad_norm": 3.5,
+  "learning_rate": 0.00295,
+  "lr_schedule": "linear_decay",
+  "lr_schedule_kl_threshold": 0.008,
+  "lr_adaptive_min": 1e-06,
+  "lr_adaptive_max": 0.01,
+  "obs_subtract_mean": 0.0,
+  "obs_scale": 1.0,
+  "normalize_input": true,
+  "normalize_input_keys": null,
+  "decorrelate_experience_max_seconds": 0,
+  "decorrelate_envs_on_one_worker": true,
+  "actor_worker_gpus": [],
+  "set_workers_cpu_affinity": true,
+  "force_envs_single_thread": false,
+  "default_niceness": 0,
+  "log_to_file": true,
+  "experiment_summaries_interval": 3,
+  "flush_summaries_interval": 30,
+  "stats_avg": 100,
+  "summaries_use_frameskip": true,
+  "heartbeat_interval": 20,
+  "heartbeat_reporting_interval": 180,
+  "train_for_env_steps": 10000000,
+  "train_for_seconds": 10000000000,
+  "save_every_sec": 15,
+  "keep_checkpoints": 2,
+  "load_checkpoint_kind": "latest",
+  "save_milestones_sec": -1,
+  "save_best_every_sec": 5,
+  "save_best_metric": "reward",
+  "save_best_after": 100000,
+  "benchmark": false,
+  "encoder_mlp_layers": [
+    64,
+    64
+  ],
+  "encoder_conv_architecture": "convnet_simple",
+  "encoder_conv_mlp_layers": [
+    512
+  ],
+  "use_rnn": false,
+  "rnn_size": 512,
+  "rnn_type": "gru",
+  "rnn_num_layers": 1,
+  "decoder_mlp_layers": [],
+  "nonlinearity": "tanh",
+  "policy_initialization": "torch_default",
+  "policy_init_gain": 1.0,
+  "actor_critic_share_weights": true,
+  "adaptive_stddev": false,
+  "continuous_tanh_scale": 0.0,
+  "initial_stddev": 1.0,
+  "use_env_info_cache": false,
+  "env_gpu_actions": false,
+  "env_gpu_observations": true,
+  "env_frameskip": 1,
+  "env_framestack": 1,
+  "pixel_format": "CHW",
+  "use_record_episode_statistics": false,
+  "with_wandb": true,
+  "wandb_user": "qgallouedec",
+  "wandb_project": "sample_facotry_metaworld",
+  "wandb_group": null,
+  "wandb_job_type": "SF",
+  "wandb_tags": [],
+  "with_pbt": false,
+  "pbt_mix_policies_in_one_env": true,
+  "pbt_period_env_steps": 5000000,
+  "pbt_start_mutation": 20000000,
+  "pbt_replace_fraction": 0.3,
+  "pbt_mutation_rate": 0.15,
+  "pbt_replace_reward_gap": 0.1,
+  "pbt_replace_reward_gap_absolute": 1e-06,
+  "pbt_optimize_gamma": false,
+  "pbt_target_objective": "true_objective",
+  "pbt_perturb_min": 1.1,
+  "pbt_perturb_max": 1.5,
+  "command_line": "--env bin-picking-v2 --experiment bin-picking-v2 --with_wandb True --wandb_user qgallouedec --wandb_project sample_facotry_metaworld",
+  "cli_args": {
+    "env": "bin-picking-v2",
+    "experiment": "bin-picking-v2",
+    "with_wandb": true,
+    "wandb_user": "qgallouedec",
+    "wandb_project": "sample_facotry_metaworld"
+  },
+  "git_hash": "aed90d9e164e44f91bab1d70c09fac4dee064031",
+  "git_repo_name": "https://github.com/huggingface/gia",
+  "wandb_unique_id": "bin-picking-v2_20230707_164044_957676"
+}

git.diff

@@ -0,0 +1,714 @@
+diff --git a/data/envs/download_expert_scores.py b/data/envs/download_expert_scores.py
+index 4c3f06b..88b6c45 100644
+--- a/data/envs/download_expert_scores.py
++++ b/data/envs/download_expert_scores.py
+@@ -12,162 +12,162 @@ from tqdm import tqdm
+ 
+ 
+ ENV_NAMES = [
+-    "atari-alien",
+-    "atari-amidar",
+-    "atari-assault",
+-    "atari-asterix",
+-    "atari-asteroids",
+-    "atari-atlantis",
+-    "atari-bankheist",
+-    "atari-battlezone",
+-    "atari-beamrider",
+-    "atari-berzerk",
+-    "atari-bowling",
+-    "atari-boxing",
+-    "atari-breakout",
+-    "atari-centipede",
+-    "atari-choppercommand",
+-    "atari-crazyclimber",
+-    "atari-defender",
+-    "atari-demonattack",
+-    "atari-doubledunk",
+-    "atari-enduro",
+-    "atari-fishingderby",
+-    "atari-freeway",
+-    "atari-frostbite",
+-    "atari-gopher",
+-    "atari-gravitar",
+-    "atari-hero",
+-    "atari-icehockey",
+-    "atari-jamesbond",
+-    "atari-kangaroo",
+-    "atari-krull",
+-    "atari-kungfumaster",
+-    "atari-montezumarevenge",
+-    "atari-mspacman",
+-    "atari-namethisgame",
+-    "atari-phoenix",
+-    "atari-pitfall",
+-    "atari-pong",
+-    "atari-privateeye",
+-    "atari-qbert",
+-    "atari-riverraid",
+-    "atari-roadrunner",
+-    "atari-robotank",
+-    "atari-seaquest",
+-    "atari-skiing",
+-    "atari-solaris",
+-    "atari-spaceinvaders",
+-    "atari-stargunner",
+-    # "atari-surround", # Not in the dataset
+-    "atari-tennis",
+-    "atari-timepilot",
+-    "atari-tutankham",
+-    "atari-upndown",
+-    "atari-venture",
+-    "atari-videopinball",
+-    "atari-wizardofwor",
+-    "atari-yarsrevenge",
+-    "atari-zaxxon",
+-    "babyai-action-obj-door",
+-    "babyai-blocked-unlock-pickup",
+-    "babyai-boss-level-no-unlock",
+-    "babyai-boss-level",
+-    "babyai-find-obj-s5",
+-    "babyai-go-to-door",
+-    # "babyai-go-to-imp-unlock",  # Not in the dataset
+-    "babyai-go-to-local",
+-    "babyai-go-to-obj-door",
+-    "babyai-go-to-obj",
+-    "babyai-go-to-red-ball-grey",
+-    "babyai-go-to-red-ball-no-dists",
+-    "babyai-go-to-red-ball",
+-    "babyai-go-to-red-blue-ball",
+-    "babyai-go-to-seq",
+-    "babyai-go-to",
+-    "babyai-key-corridor",
+-    "babyai-key-in-box",
+-    "babyai-mini-boss-level",
+-    "babyai-move-two-across",
+-    "babyai-one-room-s8",
+-    "babyai-open-door",
+-    "babyai-open-doors-order",
+-    "babyai-open-red-door",
+-    "babyai-open-two-doors",
+-    "babyai-open",
+-    "babyai-pickup-above",
+-    "babyai-pickup-dist",
+-    "babyai-pickup-loc",
+-    "babyai-pickup",
+-    "babyai-synth-loc",
+-    "babyai-synth-seq",
+-    "babyai-synth",
+-    "babyai-unblock-pickup",
+-    "babyai-unlock-local",
+-    "babyai-unlock-pickup",
+-    # "babyai-unlock-to-unlock",  # Not in the dataset
+-    # "babyai-unlock",  # Not in the dataset
++    # "atari-alien",
++    # "atari-amidar",
++    # "atari-assault",
++    # "atari-asterix",
++    # "atari-asteroids",
++    # "atari-atlantis",
++    # "atari-bankheist",
++    # "atari-battlezone",
++    # "atari-beamrider",
++    # "atari-berzerk",
++    # "atari-bowling",
++    # "atari-boxing",
++    # "atari-breakout",
++    # "atari-centipede",
++    # "atari-choppercommand",
++    # "atari-crazyclimber",
++    # "atari-defender",
++    # "atari-demonattack",
++    # "atari-doubledunk",
++    # "atari-enduro",
++    # "atari-fishingderby",
++    # "atari-freeway",
++    # "atari-frostbite",
++    # "atari-gopher",
++    # "atari-gravitar",
++    # "atari-hero",
++    # "atari-icehockey",
++    # "atari-jamesbond",
++    # "atari-kangaroo",
++    # "atari-krull",
++    # "atari-kungfumaster",
++    # "atari-montezumarevenge",
++    # "atari-mspacman",
++    # "atari-namethisgame",
++    # "atari-phoenix",
++    # "atari-pitfall",
++    # "atari-pong",
++    # "atari-privateeye",
++    # "atari-qbert",
++    # "atari-riverraid",
++    # "atari-roadrunner",
++    # "atari-robotank",
++    # "atari-seaquest",
++    # "atari-skiing",
++    # "atari-solaris",
++    # "atari-spaceinvaders",
++    # "atari-stargunner",
++    # # "atari-surround", # Not in the dataset
++    # "atari-tennis",
++    # "atari-timepilot",
++    # "atari-tutankham",
++    # "atari-upndown",
++    # "atari-venture",
++    # "atari-videopinball",
++    # "atari-wizardofwor",
++    # "atari-yarsrevenge",
++    # "atari-zaxxon",
++    # "babyai-action-obj-door",
++    # "babyai-blocked-unlock-pickup",
++    # "babyai-boss-level-no-unlock",
++    # "babyai-boss-level",
++    # "babyai-find-obj-s5",
++    # "babyai-go-to-door",
++    # # "babyai-go-to-imp-unlock",  # Not in the dataset
++    # "babyai-go-to-local",
++    # "babyai-go-to-obj-door",
++    # "babyai-go-to-obj",
++    # "babyai-go-to-red-ball-grey",
++    # "babyai-go-to-red-ball-no-dists",
++    # "babyai-go-to-red-ball",
++    # "babyai-go-to-red-blue-ball",
++    # "babyai-go-to-seq",
++    # "babyai-go-to",
++    # "babyai-key-corridor",
++    # "babyai-key-in-box",
++    # "babyai-mini-boss-level",
++    # "babyai-move-two-across",
++    # "babyai-one-room-s8",
++    # "babyai-open-door",
++    # "babyai-open-doors-order",
++    # "babyai-open-red-door",
++    # "babyai-open-two-doors",
++    # "babyai-open",
++    # "babyai-pickup-above",
++    # "babyai-pickup-dist",
++    # "babyai-pickup-loc",
++    # "babyai-pickup",
++    # "babyai-synth-loc",
++    # "babyai-synth-seq",
++    # "babyai-synth",
++    # "babyai-unblock-pickup",
++    # "babyai-unlock-local",
++    # "babyai-unlock-pickup",
++    # # "babyai-unlock-to-unlock",  # Not in the dataset
++    # # "babyai-unlock",  # Not in the dataset
+     "metaworld-assembly",
+-    "metaworld-basketball",
+-    "metaworld-bin-picking",
+-    "metaworld-box-close",
+-    "metaworld-button-press-topdown-wall",
+-    "metaworld-button-press-topdown",
+-    "metaworld-button-press-wall",
+-    "metaworld-button-press",
+-    "metaworld-coffee-button",
+-    "metaworld-coffee-pull",
+-    "metaworld-coffee-push",
+-    "metaworld-dial-turn",
+-    "metaworld-disassemble",
+-    "metaworld-door-close",
+-    "metaworld-door-lock",
+-    "metaworld-door-open",
+-    "metaworld-door-unlock",
+-    "metaworld-drawer-close",
+-    "metaworld-drawer-open",
+-    "metaworld-faucet-close",
+-    "metaworld-faucet-open",
+-    "metaworld-hammer",
+-    "metaworld-hand-insert",
+-    "metaworld-handle-press-side",
+-    "metaworld-handle-press",
+-    "metaworld-handle-pull-side",
+-    "metaworld-handle-pull",
+-    "metaworld-lever-pull",
+-    "metaworld-peg-insert-side",
+-    "metaworld-peg-unplug-side",
+-    "metaworld-pick-out-of-hole",
+-    "metaworld-pick-place-wall",
+-    "metaworld-pick-place",
+-    "metaworld-plate-slide-back-side",
+-    "metaworld-plate-slide-back",
+-    "metaworld-plate-slide-side",
+-    "metaworld-plate-slide",
+-    "metaworld-push-back",
+-    "metaworld-push-wall",
+-    "metaworld-push",
+-    "metaworld-reach-wall",
+-    "metaworld-reach",
+-    "metaworld-shelf-place",
+-    "metaworld-soccer",
+-    "metaworld-stick-pull",
+-    "metaworld-stick-push",
+-    "metaworld-sweep-into",
+-    "metaworld-sweep",
+-    "metaworld-window-close",
+-    "metaworld-window-open",
+-    "mujoco-ant",
+-    "mujoco-doublependulum",
+-    "mujoco-halfcheetah",
+-    "mujoco-hopper",
++    # "metaworld-basketball",
++    # "metaworld-bin-picking",
++    # "metaworld-box-close",
++    # "metaworld-button-press-topdown-wall",
++    # "metaworld-button-press-topdown",
++    # "metaworld-button-press-wall",
++    # "metaworld-button-press",
++    # "metaworld-coffee-button",
++    # "metaworld-coffee-pull",
++    # "metaworld-coffee-push",
++    # "metaworld-dial-turn",
++    # "metaworld-disassemble",
++    # "metaworld-door-close",
++    # "metaworld-door-lock",
++    # "metaworld-door-open",
++    # "metaworld-door-unlock",
++    # "metaworld-drawer-close",
++    # "metaworld-drawer-open",
++    # "metaworld-faucet-close",
++    # "metaworld-faucet-open",
++    # "metaworld-hammer",
++    # "metaworld-hand-insert",
++    # "metaworld-handle-press-side",
++    # "metaworld-handle-press",
++    # "metaworld-handle-pull-side",
++    # "metaworld-handle-pull",
++    # "metaworld-lever-pull",
++    # "metaworld-peg-insert-side",
++    # "metaworld-peg-unplug-side",
++    # "metaworld-pick-out-of-hole",
++    # "metaworld-pick-place-wall",
++    # "metaworld-pick-place",
++    # "metaworld-plate-slide-back-side",
++    # "metaworld-plate-slide-back",
++    # "metaworld-plate-slide-side",
++    # "metaworld-plate-slide",
++    # "metaworld-push-back",
++    # "metaworld-push-wall",
++    # "metaworld-push",
++    # "metaworld-reach-wall",
++    # "metaworld-reach",
++    # "metaworld-shelf-place",
++    # "metaworld-soccer",
++    # "metaworld-stick-pull",
++    # "metaworld-stick-push",
++    # "metaworld-sweep-into",
++    # "metaworld-sweep",
++    # "metaworld-window-close",
++    # "metaworld-window-open",
++    # "mujoco-ant",
++    # "mujoco-doublependulum",
++    # "mujoco-halfcheetah",
++    # "mujoco-hopper",
+     # "mujoco-humanoid",  # Not in the dataset
+-    "mujoco-pendulum",
+-    # "mujoco-pusher",  # Not in the dataset
+-    "mujoco-reacher",
++    # "mujoco-pendulum",
++    # # "mujoco-pusher",  # Not in the dataset
++    # "mujoco-reacher",
+     # "mujoco-standup",  # Not in the dataset
+-    "mujoco-swimmer",
+-    "mujoco-walker",
++    # "mujoco-swimmer",
++    # "mujoco-walker",
+ ]
+ 
+ 
+diff --git a/data/envs/metaworld/generate_dataset.py b/data/envs/metaworld/generate_dataset.py
+index e21b237..c2b1907 100644
+--- a/data/envs/metaworld/generate_dataset.py
++++ b/data/envs/metaworld/generate_dataset.py
+@@ -142,7 +142,8 @@ def create_dataset(cfg: Config, dataset_size: int = 100_000, split: str = "train
+ 
+             # Actions shape should be [num_agents, num_actions] even if it's [1, 1]
+             actions = preprocess_actions(env_info, actions)
+-
++            # Clamp actions to be in the range of the action space
++            actions = np.clip(actions, env.action_space.low, env.action_space.high)
+             rnn_states = policy_outputs["new_rnn_states"]
+             dataset["continuous_observations"][-1].append(observations["obs"].cpu().numpy()[0])
+             dataset["continuous_actions"][-1].append(actions[0])
+diff --git a/data/envs/metaworld/generate_dataset_all.sh b/data/envs/metaworld/generate_dataset_all.sh
+index cfdae2f..5db8c4b 100755
+--- a/data/envs/metaworld/generate_dataset_all.sh
++++ b/data/envs/metaworld/generate_dataset_all.sh
+@@ -2,58 +2,58 @@
+ 
+ ENVS=(
+     assembly
+-    basketball
+-    bin-picking
+-    box-close
+-    button-press-topdown
+-    button-press-topdown-wall
+-    button-press
+-    button-press-wall
+-    coffee-button
+-    coffee-pull
+-    coffee-push
+-    dial-turn
+-    disassemble
+-    door-close
+-    door-lock
+-    door-open
+-    door-unlock
+-    drawer-close
+-    drawer-open
+-    faucet-close
+-    faucet-open
+-    hammer
+-    hand-insert
+-    handle-press-side
+-    handle-press
+-    handle-pull-side
+-    handle-pull
+-    lever-pull
+-    peg-insert-side
+-    peg-unplug-side
+-    pick-out-of-hole
+-    pick-place
+-    pick-place-wall
+-    plate-slide-back-side
+-    plate-slide-back
+-    plate-slide-side
+-    plate-slide
+-    push-back
+-    push
+-    push-wall
+-    reach
+-    reach-wall
+-    shelf-place
+-    soccer
+-    stick-pull
+-    stick-push
+-    sweep-into
+-    sweep
+-    window-close
+-    window-open
++    # basketball
++    # bin-picking
++    # box-close
++    # button-press-topdown
++    # button-press-topdown-wall
++    # button-press
++    # button-press-wall
++    # coffee-button
++    # coffee-pull
++    # coffee-push
++    # dial-turn
++    # disassemble
++    # door-close
++    # door-lock
++    # door-open
++    # door-unlock
++    # drawer-close
++    # drawer-open
++    # faucet-close
++    # faucet-open
++    # hammer
++    # hand-insert
++    # handle-press-side
++    # handle-press
++    # handle-pull-side
++    # handle-pull
++    # lever-pull
++    # peg-insert-side
++    # peg-unplug-side
++    # pick-out-of-hole
++    # pick-place
++    # pick-place-wall
++    # plate-slide-back-side
++    # plate-slide-back
++    # plate-slide-side
++    # plate-slide
++    # push-back
++    # push
++    # push-wall
++    # reach
++    # reach-wall
++    # shelf-place
++    # soccer
++    # stick-pull
++    # stick-push
++    # sweep-into
++    # sweep
++    # window-close
++    # window-open
+ )
+ 
+ for ENV in "${ENVS[@]}"; do
+-    python -m sample_factory.huggingface.load_from_hub -r qgallouedec/sample-factory-$ENV-v2
+-    python generate_dataset.py --env $ENV-v2 --experiment sample-factory-$ENV-v2 --train_dir=./train_dir
++    python -m sample_factory.huggingface.load_from_hub -r qgallouedec/$ENV-v2
++    python generate_dataset.py --env $ENV-v2 --experiment $ENV-v2 --train_dir=./train_dir
+ done
+diff --git a/data/envs/metaworld/push_all.sh b/data/envs/metaworld/push_all.sh
+index 9d71467..5b05c6d 100755
+--- a/data/envs/metaworld/push_all.sh
++++ b/data/envs/metaworld/push_all.sh
+@@ -2,57 +2,57 @@
+ 
+ ENVS=(
+     assembly
+-    basketball
+-    bin-picking
+-    box-close
+-    button-press-topdown
+-    button-press-topdown-wall
+-    button-press
+-    button-press-wall
+-    coffee-button
+-    coffee-pull
+-    coffee-push
+-    dial-turn
+-    disassemble
+-    door-close
+-    door-lock
+-    door-open
+-    door-unlock
+-    drawer-close
+-    drawer-open
+-    faucet-close
+-    faucet-open
+-    hammer
+-    hand-insert
+-    handle-press-side
+-    handle-press
+-    handle-pull-side
+-    handle-pull
+-    lever-pull
+-    peg-insert-side
+-    peg-unplug-side
+-    pick-out-of-hole
+-    pick-place
+-    pick-place-wall
+-    plate-slide-back-side
+-    plate-slide-back
+-    plate-slide-side
+-    plate-slide
+-    push-back
+-    push
+-    push-wall
+-    reach
+-    reach-wall
+-    shelf-place
+-    soccer
+-    stick-pull
+-    stick-push
+-    sweep-into
+-    sweep
+-    window-close
+-    window-open
++    # basketball
++    # bin-picking
++    # box-close
++    # button-press-topdown
++    # button-press-topdown-wall
++    # button-press
++    # button-press-wall
++    # coffee-button
++    # coffee-pull
++    # coffee-push
++    # dial-turn
++    # disassemble
++    # door-close
++    # door-lock
++    # door-open
++    # door-unlock
++    # drawer-close
++    # drawer-open
++    # faucet-close
++    # faucet-open
++    # hammer
++    # hand-insert
++    # handle-press-side
++    # handle-press
++    # handle-pull-side
++    # handle-pull
++    # lever-pull
++    # peg-insert-side
++    # peg-unplug-side
++    # pick-out-of-hole
++    # pick-place
++    # pick-place-wall
++    # plate-slide-back-side
++    # plate-slide-back
++    # plate-slide-side
++    # plate-slide
++    # push-back
++    # push
++    # push-wall
++    # reach
++    # reach-wall
++    # shelf-place
++    # soccer
++    # stick-pull
++    # stick-push
++    # sweep-into
++    # sweep
++    # window-close
++    # window-open
+ )
+ 
+ for ENV in "${ENVS[@]}"; do
+-    python enjoy.py --algo=APPO --env $ENV-v2 --experiment $ENV-v2 --train_dir=./train_dir --max_num_episodes=10 --push_to_hub --hf_repository=qgallouedec/sample-factory-$ENV-v2 --save_video --no_render --enjoy_script=enjoy --train_script=train --load_checkpoint_kind best
++    python enjoy.py --algo=APPO --env $ENV-v2 --experiment $ENV-v2 --train_dir=./train_dir --max_num_episodes=10 --push_to_hub --hf_repository=qgallouedec/$ENV-v2 --save_video --no_render --enjoy_script=enjoy --train_script=train --load_checkpoint_kind best
+ done
+diff --git a/data/envs/metaworld/train.py b/data/envs/metaworld/train.py
+index 46dc581..c72f289 100644
+--- a/data/envs/metaworld/train.py
++++ b/data/envs/metaworld/train.py
+@@ -79,7 +79,7 @@ def override_defaults(parser: argparse.ArgumentParser) -> argparse.ArgumentParse
+         num_workers=8,
+         num_envs_per_worker=8,
+         worker_num_splits=2,
+-        train_for_env_steps=100_000_000,
++        train_for_env_steps=10_000_000,
+         encoder_mlp_layers=[64, 64],
+         env_frameskip=1,
+         nonlinearity="tanh",
+diff --git a/data/envs/metaworld/train_all.sh b/data/envs/metaworld/train_all.sh
+index dbf328a..1b3c4c8 100755
+--- a/data/envs/metaworld/train_all.sh
++++ b/data/envs/metaworld/train_all.sh
+@@ -1,7 +1,7 @@
+ #!/bin/bash
+ 
+ ENVS=(
+-    assembly
++    # assembly
+     basketball
+     bin-picking
+     box-close
+diff --git a/gia/eval/callback.py b/gia/eval/callback.py
+index 5c3a080..4b6198f 100644
+--- a/gia/eval/callback.py
++++ b/gia/eval/callback.py
+@@ -2,10 +2,10 @@ import glob
+ import json
+ import subprocess
+ 
+-import wandb
+ from accelerate import Accelerator
+ from transformers import TrainerCallback, TrainerControl, TrainerState, TrainingArguments
+ 
++import wandb
+ from gia.config import Arguments
+ from gia.eval.utils import is_slurm_available
+ 
+diff --git a/gia/eval/evaluator.py b/gia/eval/evaluator.py
+index 91b645c..3e2cae7 100644
+--- a/gia/eval/evaluator.py
++++ b/gia/eval/evaluator.py
+@@ -1,3 +1,5 @@
++from typing import Optional
++
+ import torch
+ 
+ from gia.config.arguments import Arguments
+@@ -5,11 +7,12 @@ from gia.model import GiaModel
+ 
+ 
+ class Evaluator:
+-    def __init__(self, args: Arguments, task: str) -> None:
++    def __init__(self, args: Arguments, task: str, mean_random: Optional[float] = None) -> None:
+         self.args = args
+         self.task = task
++        self.mean_random = mean_random
+ 
+-    @torch.no_grad()
++    @torch.inference_mode()
+     def evaluate(self, model: GiaModel) -> float:
+         return self._evaluate(model)
+ 
+diff --git a/gia/eval/mappings.py b/gia/eval/mappings.py
+deleted file mode 100644
+index e7ba9d3..0000000
+--- a/gia/eval/mappings.py
++++ /dev/null
+@@ -1,11 +0,0 @@
+-TASK_TO_ENV_MAPPING = {
+-    "mujoco-ant": "Ant-v4",
+-    "mujoco-halfcheetah": "HalfCheetah-v4",
+-    "mujoco-hopper": "Hopper-v4",
+-    "mujoco-doublependulum": "InvertedDoublePendulum-v4",
+-    "mujoco-pendulum": "InvertedPendulum-v4",
+-    "mujoco-reacher": "Reacher-v4",
+-    "mujoco-swimmer": "Swimmer-v4",
+-    "mujoco-walker": "Walker2d-v4",
+-    # Atari etc...
+-}
+diff --git a/gia/eval/rl/__init__.py b/gia/eval/rl/__init__.py
+index 36d890b..da5e0c7 100644
+--- a/gia/eval/rl/__init__.py
++++ b/gia/eval/rl/__init__.py
+@@ -1,4 +1,5 @@
++from .envs.core import make
+ from .gym_evaluator import GymEvaluator
+ 
+ 
+-__all__ = ["GymEvaluator"]
++__all__ = ["GymEvaluator", "make"]
+diff --git a/gia/eval/rl/gia_agent.py b/gia/eval/rl/gia_agent.py
+index f0d0b9b..04b9637 100644
+--- a/gia/eval/rl/gia_agent.py
++++ b/gia/eval/rl/gia_agent.py
+@@ -75,6 +75,11 @@ class GiaAgent:
+     ) -> Tuple[Tuple[Tensor, Tensor], ...]:
+         return tuple((k[:, :, -self._max_length :], v[:, :, -self._max_length :]) for (k, v) in past_key_values)
+ 
++    def set_model(self, model: GiaModel) -> None:
++        self.model = model
++        self.device = next(model.parameters()).device
++        self._max_length = self.model.config.max_position_embeddings
++
+     def reset(self, num_envs: int = 1) -> None:
+         if self.prompter is not None:
+             prompts = self.prompter.generate_prompts(num_envs)
+diff --git a/gia/eval/rl/gym_evaluator.py b/gia/eval/rl/gym_evaluator.py
+index f8531ee..754c05d 100644
+--- a/gia/eval/rl/gym_evaluator.py
++++ b/gia/eval/rl/gym_evaluator.py
+@@ -1,7 +1,7 @@
+ import gym
+ from gym.vector.vector_env import VectorEnv
+ 
+-from gia.eval.mappings import TASK_TO_ENV_MAPPING
++# from gia.eval.rl.envs.mappings import TASK_TO_ENV_MAPPING
+ from gia.eval.rl.rl_evaluator import RLEvaluator
+ 
+ 
+diff --git a/gia/eval/rl/rl_evaluator.py b/gia/eval/rl/rl_evaluator.py
+index c5cc423..91189f3 100644
+--- a/gia/eval/rl/rl_evaluator.py
++++ b/gia/eval/rl/rl_evaluator.py
+@@ -8,6 +8,10 @@ from gia.eval.rl.gia_agent import GiaAgent
+ 
+ 
+ class RLEvaluator(Evaluator):
++    def __init__(self, args, task):
++        super().__init__(args, task)
++        self.agent = GiaAgent()
++
+     def _build_env(self) -> VectorEnv:  # TODO: maybe just a gym.Env ?
+         raise NotImplementedError
+ 
+diff --git a/gia/eval/rl/scores_dict.json b/gia/eval/rl/scores_dict.json
+index 1b8ebee..ff7d030 100644
+--- a/gia/eval/rl/scores_dict.json
++++ b/gia/eval/rl/scores_dict.json
+@@ -929,8 +929,8 @@
+     },
+     "metaworld-assembly": {
+         "expert": {
+-            "mean": 311.29314618777823,
+-            "std": 75.04282151450695
++            "mean": 3523.81468486244,
++            "std": 63.22745220327798
+         },
+         "random": {
+             "mean": 220.65601680730813,

replay.mp4

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:d7d292cb6c8b454365d8a689263ddb710b0149da82dafb6442f0524ee8a0c486
+size 2343193