[NOMERG] PPO-Myo

examples/distributed/collectors/single_machine/generic.py

@@ -32,6 +32,7 @@
 from torchrl.collectors.distributed import DistributedDataCollector
 from torchrl.envs import EnvCreator, ParallelEnv
 from torchrl.envs.libs.gym import GymEnv
+from torchrl.envs.libs.robohive import RoboHiveEnv
 
 parser = ArgumentParser()
 parser.add_argument(
@@ -80,6 +81,16 @@
     default="ALE/Pong-v5",
     help="Gym environment to be run.",
 )
+LIBS = {
+    "gym": GymEnv,
+    "robohive": RoboHiveEnv,
+}
+parser.add_argument(
+    "--lib",
+    default="gym",
+    help="Lib backend",
+    choices=list(LIBS.keys()),
+)
 if __name__ == "__main__":
     args = parser.parse_args()
     num_workers = args.num_workers
@@ -89,7 +100,8 @@
 
     device_count = torch.cuda.device_count()
 
-    make_env = EnvCreator(lambda: GymEnv(args.env))
+    lib = LIBS[args.lib]
+    make_env = EnvCreator(lambda: lib(args.env))
     if args.worker_parallelism == "collector" or num_workers == 1:
         action_spec = make_env().action_spec
     else:

examples/distributed/collectors/single_machine/rpc.py

@@ -27,6 +27,7 @@
 from torchrl.collectors.distributed import RPCDataCollector
 from torchrl.envs import EnvCreator, ParallelEnv
 from torchrl.envs.libs.gym import GymEnv
+from torchrl.envs.libs.robohive import RoboHiveEnv
 
 parser = ArgumentParser()
 parser.add_argument(
@@ -63,6 +64,16 @@
     default="ALE/Pong-v5",
     help="Gym environment to be run.",
 )
+LIBS = {
+    "gym": GymEnv,
+    "robohive": RoboHiveEnv,
+}
+parser.add_argument(
+    "--lib",
+    default="gym",
+    help="Lib backend",
+    choices=list(LIBS.keys()),
+)
 if __name__ == "__main__":
     args = parser.parse_args()
     num_workers = args.num_workers

examples/distributed/collectors/single_machine/sync.py

@@ -32,6 +32,7 @@
 from torchrl.collectors.distributed import DistributedSyncDataCollector
 from torchrl.envs import EnvCreator, ParallelEnv
 from torchrl.envs.libs.gym import GymEnv
+from torchrl.envs.libs.robohive import RoboHiveEnv
 
 parser = ArgumentParser()
 parser.add_argument(
@@ -75,6 +76,16 @@
     default="ALE/Pong-v5",
     help="Gym environment to be run.",
 )
+LIBS = {
+    "gym": GymEnv,
+    "robohive": RoboHiveEnv,
+}
+parser.add_argument(
+    "--lib",
+    default="gym",
+    help="Lib backend",
+    choices=list(LIBS.keys()),
+)
 if __name__ == "__main__":
     args = parser.parse_args()
     num_workers = args.num_workers

examples/ppo/config_atari.yaml

@@ -0,0 +1,35 @@
+# Environment
+env:
+  env_name: PongNoFrameskip-v4
+
+# collector
+collector:
+  frames_per_batch: 4096
+  total_frames: 40_000_000
+
+# logger
+logger:
+  backend: wandb
+  exp_name: Atari_Schulman17
+  test_interval: 40_000_000
+  num_test_episodes: 3
+
+# Optim
+optim:
+  lr: 2.5e-4
+  eps: 1.0e-6
+  weight_decay: 0.0
+  max_grad_norm: 0.5
+  anneal_lr: True
+
+# loss
+loss:
+  gamma: 0.99
+  mini_batch_size: 1024
+  ppo_epochs: 3
+  gae_lambda: 0.95
+  clip_epsilon: 0.1
+  anneal_clip_epsilon: True
+  critic_coef: 1.0
+  entropy_coef: 0.01
+  loss_critic_type: l2

examples/ppo/config_mujoco.yaml

@@ -0,0 +1,32 @@
+# task and env
+env:
+  env_name: HalfCheetah-v3
+
+# collector
+collector:
+  frames_per_batch: 2048
+  total_frames: 1_000_000
+
+# logger
+logger:
+  backend: wandb
+  exp_name: Mujoco_Schulman17
+  test_interval: 1_000_000
+  num_test_episodes: 5
+
+# Optim
+optim:
+  lr: 3e-4
+  weight_decay: 0.0
+  anneal_lr: False
+
+# loss
+loss:
+  gamma: 0.99
+  mini_batch_size: 64
+  ppo_epochs: 10
+  gae_lambda: 0.95
+  clip_epsilon: 0.2
+  critic_coef: 0.25
+  entropy_coef: 0.0
+  loss_critic_type: l2

examples/ppo/config_myo.yaml

@@ -0,0 +1,33 @@
+# task and env
+env:
+  env_name: myoHandReachRandom-v0
+
+# collector
+collector:
+  frames_per_batch: 2048
+  total_frames: 1_000_000
+  num_envs: 1
+
+# logger
+logger:
+  backend: wandb
+  exp_name: myo_hand_reach
+  test_interval: 1_000_000
+  num_test_episodes: 5
+
+# Optim
+optim:
+  lr: 3e-4
+  weight_decay: 0.0
+  anneal_lr: False
+
+# loss
+loss:
+  gamma: 0.99
+  mini_batch_size: 64
+  ppo_epochs: 10
+  gae_lambda: 0.95
+  clip_epsilon: 0.2
+  critic_coef: 0.25
+  entropy_coef: 0.0
+  loss_critic_type: l2