updates

README.md

@@ -3,3 +3,6 @@
 ### Get familiar with Hugging Face Hub
 
 Link to the model on the hub : 
+
+# REINFORCE
+Plot showing the total reward accross episodes: ![Alt text](images/reinforce_rewards.png)

a2c_sb3_cartpole.py

@@ -2,8 +2,15 @@ import gymnasium as gym
 import cv2
 from stable_baselines3 import A2C
 from huggingface_sb3 import package_to_hub, push_to_hub
+from gym import envs
+from gymnasium.envs.registration import register
 
-env = gym.make("CartPole-v1", render_mode="rgb_array")
+env_id = "CartPole-v1"
+
+# Register the environment
+register(id=env_id, entry_point='gym.envs.classic_control:CartPoleEnv', max_episode_steps=500)
+
+env = gym.make(env_id)
 
 model = A2C("MlpPolicy", env, verbose=1)
 model.learn(total_timesteps=10_000)
@@ -18,15 +25,16 @@ for i in range(1000):
     # if done:
     #   obs = vec_env.reset()
 
-if __name__ == "__main__":
-    package_to_hub(model=model, 
+# Package and push the model to the Hugging Face Hub
+model_package_id = package_to_hub(model=model, 
                                   model_name="a2c-CartPole-v1",
                                   model_architecture="a2c",
-                    env_id="CartPole-v1",
+                                  env_id=env_id,
                                   eval_env=env,
                                   repo_id="oscarchaufour/a2c-CartPole-v1",
-                    commit_message="Test commit")
+                                  commit_message="Initial commit of A2C CartPole model")
 
+# Push the model package to the Hub
 push_to_hub(repo_id="oscarchaufour/a2c-CartPole-v1",
-                filename="a2c_sb3_cartpole.zip",
+            filename=model_package_id + ".zip",
             commit_message="Added A2C CartPole model")
\ No newline at end of file

reinforce_cartpole.py

@@ -5,6 +5,7 @@ import torch.optim as optim
 import matplotlib.pyplot as plt
 from tqdm import tqdm 
 import numpy as np
+from torch.distributions import Categorical
 
 # Define the neural network model
 class Policy(nn.Module):
@@ -24,77 +25,68 @@ class Policy(nn.Module):
         x = self.softmax(x)
         return x
 
-
-def reinforce():
-    # Create the environment
-    env = gym.make("CartPole-v1")
-
-    # Set up the agent
-    policy = Policy(
-        input_size=env.observation_space.shape[0],
-        output_size=env.action_space.n
-    )
-    optimizer = optim.Adam(policy.parameters(), lr=5e-3)
-
+def reinforce(policy, env, optimizer):
     # Training loop
-    num_episodes = 500
+    num_episodes = 200
     gamma = 0.99
-    episode_rewards = []
+    episodes_rewards = []
 
     for episode in tqdm(range(num_episodes)):
-        action_probabilities = []
-        episode_rewards_weighted = []
-
         # Reset the environment and get the initial observation
-        observation = env.reset()[0]
+        observation = env.reset()
 
+        rewards = []
+        log_probabilities = []
         terminated = False
         episode_reward = 0
-        step = 0
 
         while not terminated:
-            step += 1
             # Compute action probabilities
-            action_probs = policy(torch.FloatTensor(observation).unsqueeze(0))
+            observation_array = observation[0] if isinstance(observation, tuple) else observation
+            observation_tensor = torch.FloatTensor(observation_array).unsqueeze(0)
+            action_probabilities = policy(observation_tensor)
             
+            cat = Categorical(action_probabilities)
             # Sample action based on probabilities and store its probability in the buffer
-            action = torch.multinomial(action_probs, num_samples=1).item()
+            action = cat.sample()
+            log_probability = cat.log_prob(action)
 
             # Step the environment with the action
-            observation, reward, terminated, truncated, info = env.step(action)
+            observation, reward, terminated, _, _ = env.step(action.item())
             env.render()
 
-            # Compute and store the return in the buffer
+            # Compute and store the reward in the buffer
             episode_reward += reward
-            episode_rewards_weighted.append(reward * gamma ** step)
-            #episode_rewards.append(episode_reward)
+            log_probabilities.append(log_probability)
+            rewards.append(reward)
         
-            # Store the action probabilities
-            action_probabilities.append(action_probs[0][action])
+        episodes_rewards.append(episode_reward)
 
-        # Normalize the return
-        # Convert action_probabilities to a tensor with requires_grad=True
-        action_probabilities_tensor = torch.FloatTensor(action_probabilities).requires_grad_(True)
-
-        episode_rewards_tensor = torch.FloatTensor(episode_rewards_weighted)
-        episode_rewards_tensor -= torch.mean(episode_rewards_tensor)
-        episode_rewards_tensor /= torch.std(episode_rewards_tensor)
+        # Compute the returns
+        weighted_rewards = 0
+        returns = []
+        for r in rewards[::-1]:
+            weighted_rewards = r + gamma * weighted_rewards
+            returns.insert(0, weighted_rewards)
+        returns = torch.tensor(returns)
+        returns = (returns - returns.mean()) / (returns.std() + 1e-5)
 
         # Compute policy loss
-        log_probs = torch.log(action_probabilities_tensor.squeeze(0))
-        policy_loss = -torch.sum(log_probs * torch.tensor(episode_rewards_weighted))
+        loss = []
+        for log_prob, weighted_rewards in zip(log_probabilities, returns):
+            loss.append(log_prob * weighted_rewards)
+        policy_loss = - torch.cat(loss).sum()
 
         # Update the policy
         optimizer.zero_grad()
         policy_loss.backward()
         optimizer.step()
-        episode_rewards.append(episode_reward)
-    return episode_rewards
+    return episodes_rewards
 
 
-def plot_rewards(episode_rewards):
+def plot_rewards(episodes_rewards):
     # Plot the total reward across episodes
-    plt.plot(episode_rewards)
+    plt.plot(episodes_rewards)
     plt.xlabel('Episode')
     plt.ylabel('Total Reward')
     plt.title('REINFORCE: Total Reward across Episodes')
@@ -102,5 +94,33 @@ def plot_rewards(episode_rewards):
 
 
 if __name__ == "__main__":
-    episode_rewards = reinforce()
-    plot_rewards(episode_rewards)
\ No newline at end of file
+    # Create the environment
+    env = gym.make("CartPole-v1")
+
+    # Set up the agent
+    policy = Policy(
+        input_size=env.observation_space.shape[0],
+        output_size=env.action_space.n
+    )
+    optimizer = optim.Adam(policy.parameters(), lr=5e-3)
+    episodes_rewards = reinforce(policy, env, optimizer)
+    env.close()
+
+    # plot the rewards
+    plot_rewards(episodes_rewards)
+
+
+    # Serialize the model and save it to a .zip file
+    # import pickle
+    # import zipfile
+
+    # Assuming `policy` is your neural network model
+
+    # # Step 1: Serialize the model
+    # model_bytes = pickle.dumps(policy)
+
+    # # Step 2: Create a .zip file containing the serialized model
+    # zip_filename = ".zip"
+    # with zipfile.ZipFile(zip_filename, 'w') as zipf:
+    #     zipf.writestr("model.pkl", model_bytes)
+

test.ipynb

@@ -708,13 +708,41 @@
    "cell_type": "markdown",
    "metadata": {},
    "source": [
-    "## Reinforce"
+    "## Upload model to Hugging Face hub"
    ]
   },
   {
-   "cell_type": "markdown",
+   "cell_type": "code",
+   "execution_count": 8,
    "metadata": {},
-   "source": []
+   "outputs": [
+    {
+     "ename": "NameError",
+     "evalue": "name 'model' is not defined",
+     "output_type": "error",
+     "traceback": [
+      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
+      "\u001b[1;31mNameError\u001b[0m                                 Traceback (most recent call last)",
+      "Cell \u001b[1;32mIn[8], line 3\u001b[0m\n\u001b[0;32m      1\u001b[0m \u001b[38;5;28;01mfrom\u001b[39;00m \u001b[38;5;21;01mhuggingface_sb3\u001b[39;00m \u001b[38;5;28;01mimport\u001b[39;00m package_to_hub, push_to_hub\n\u001b[1;32m----> 3\u001b[0m package_to_hub(model\u001b[38;5;241m=\u001b[39m\u001b[43mmodel\u001b[49m, \n\u001b[0;32m      4\u001b[0m                 model_name\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124ma2c-CartPole-v1\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m      5\u001b[0m                 model_architecture\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124ma2c\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m      6\u001b[0m                 env_id\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mCartPole-v1\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m      7\u001b[0m                 eval_env\u001b[38;5;241m=\u001b[39m\u001b[38;5;28;01mNone\u001b[39;00m,\n\u001b[0;32m      8\u001b[0m                 repo_id\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124moscarchaufour/a2c-CartPole-v1\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m      9\u001b[0m                 commit_message\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mTest commit\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n\u001b[0;32m     11\u001b[0m push_to_hub(repo_id\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124moscarchaufour/a2c-CartPole-v1\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m     12\u001b[0m             filename\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124ma2c_sb3_cartpole.zip\u001b[39m\u001b[38;5;124m\"\u001b[39m,\n\u001b[0;32m     13\u001b[0m             commit_message\u001b[38;5;241m=\u001b[39m\u001b[38;5;124m\"\u001b[39m\u001b[38;5;124mAdded A2C CartPole model\u001b[39m\u001b[38;5;124m\"\u001b[39m)\n",
+      "\u001b[1;31mNameError\u001b[0m: name 'model' is not defined"
+     ]
+    }
+   ],
+   "source": [
+    "from huggingface_sb3 import package_to_hub, push_to_hub\n",
+    "\n",
+    "package_to_hub(model=model, \n",
+    "                model_name=\"a2c-CartPole-v1\",\n",
+    "                model_architecture=\"a2c\",\n",
+    "                env_id=\"CartPole-v1\",\n",
+    "                eval_env=None,\n",
+    "                repo_id=\"oscarchaufour/a2c-CartPole-v1\",\n",
+    "                commit_message=\"Test commit\")\n",
+    "    \n",
+    "push_to_hub(repo_id=\"oscarchaufour/a2c-CartPole-v1\",\n",
+    "            filename=\"a2c_sb3_cartpole.zip\",\n",
+    "            commit_message=\"Added A2C CartPole model\")\n"
+   ]
   }
  ],
  "metadata": {