Update README.md

README.md

@@ -30,8 +30,63 @@ TODO: Add your code
 
 
 ```python
-from stable_baselines3 import ...
-from huggingface_sb3 import load_from_hub
+import panda_gym
+import gym
 
+from huggingface_sb3 import package_to_hub
+
+from stable_baselines3 import A2C
+from stable_baselines3.common.env_util import make_vec_env
+from stable_baselines3.common.vec_env import SubprocVecEnv
+from stable_baselines3.common.vec_env import DummyVecEnv, VecNormalize
+from stable_baselines3.common.evaluation import evaluate_policy
+
+env_id = "PandaReachDense-v2"
+model_name = "PandaReachDenseA2C-n8"
+env_name = f"{env_id}_vec_normalize.pkl"
+
+if __name__=="__main__":
+    env = make_vec_env(env_id, n_envs=6, vec_env_cls=SubprocVecEnv)
+    # 3
+    env = VecNormalize(env, norm_obs=True, norm_reward=False, clip_obs=10.)
+    def linear_scheduler(progress_remaining: float):
+        # from https://github.com/DLR-RM/rl-baselines3-zoo/blob/33eba22eb36128412a5b22b57a7a10bfe71e6278/rl_zoo3/utils.py
+        return progress_remaining * 0.0009
+    # 4
+    model = A2C(policy = "MultiInputPolicy",
+            env = env,
+            verbose=1,
+            device='cpu',
+            learning_rate=linear_scheduler,
+            use_rms_prop=True,
+            gae_lambda=0.9,
+            use_sde=True,
+            n_steps=8,
+            )
+    # 5
+    model.learn(1_500_000)
+
+    model.save(model_name)
+    env.save(env_name)
+    del env
+
+    eval_env = DummyVecEnv([lambda: gym.make("PandaReachDense-v2")])
+    eval_env = VecNormalize.load(env_name, eval_env)
+
+    eval_env.training = False
+    eval_env.norm_reward = False
+
+    mean_reward, std_reward = evaluate_policy(model, eval_env)
+    print(f"Mean reward = {mean_reward:.2f} +/- {std_reward:.2f}")
+    
+    package_to_hub(
+        model=model,
+        model_name=model_name,
+        model_architecture="A2C",
+        env_id=env_id,
+        eval_env=eval_env,
+        repo_id=f"bobobert4/a2c-{env_id}",
+        commit_message="Another commit",
+    )
 ...
 ```