Replace a2c_sb3_panda_reach.py

a2c_sb3_panda_reach.py

@@ -2,20 +2,22 @@ import wandb
 wandb.init(project="PandaReachJointsDense-v2", sync_tensorboard=True, monitor_gym=True)
 
 import gym
+import panda_gym
 import torch
 from stable_baselines3 import A2C
 from stable_baselines3.common.monitor import Monitor
+from huggingface_sb3 import push_to_hub
 
 # Set up the CartPole environement
 # Create the environment
 env = gym.make("PandaReachJointsDense-v2")
-#env = Monitor(env)
+env = Monitor(env)
 
 # Reset the environment and get the initial observation
 observation = env.reset()
 
 
-model = A2C("MlpPolicy", env, verbose=1)
+model = A2C("MultiInputPolicy", env, verbose=1)
 print(model)
 model.learn(total_timesteps=500000)
 
@@ -36,3 +38,14 @@ for _ in range(500):
     env.render()
     
 wandb.finish()
+
+# Save the trained model
+model.save("a2c_sb3_panda_reach.zip")
+
+# Load the trained model
+model = A2C.load("a2c_sb3_panda_reach.zip")
+
+push_to_hub(
+    repo_id="Younes-hands-on-rl/a2c_sb3_panda_reach",
+    filename="a2c_sb3_panda_reach.zip",
+    commit_message="Commit model")
\ No newline at end of file