2024-03-13 18:04:30 +00:00
|
|
|
{
|
|
|
|
"cells": [
|
|
|
|
{
|
|
|
|
"cell_type": "code",
|
2024-03-13 20:57:23 +00:00
|
|
|
"execution_count": null,
|
2024-03-13 18:04:30 +00:00
|
|
|
"metadata": {},
|
|
|
|
"outputs": [],
|
|
|
|
"source": [
|
|
|
|
"import gym\n",
|
|
|
|
"import gym_wordle\n",
|
|
|
|
"from stable_baselines3 import DQN\n",
|
2024-03-13 20:57:23 +00:00
|
|
|
"import numpy as np\n",
|
|
|
|
"import tqdm"
|
2024-03-13 18:04:30 +00:00
|
|
|
]
|
|
|
|
},
|
|
|
|
{
|
|
|
|
"cell_type": "code",
|
2024-03-13 20:57:23 +00:00
|
|
|
"execution_count": null,
|
2024-03-13 18:04:30 +00:00
|
|
|
"metadata": {},
|
2024-03-13 20:57:23 +00:00
|
|
|
"outputs": [],
|
2024-03-13 18:04:30 +00:00
|
|
|
"source": [
|
|
|
|
"env = gym.make(\"Wordle-v0\")\n",
|
|
|
|
"\n",
|
|
|
|
"print(env)"
|
|
|
|
]
|
|
|
|
},
|
|
|
|
{
|
|
|
|
"cell_type": "code",
|
2024-03-13 20:57:23 +00:00
|
|
|
"execution_count": 35,
|
2024-03-13 18:04:30 +00:00
|
|
|
"metadata": {},
|
2024-03-13 20:57:23 +00:00
|
|
|
"outputs": [],
|
2024-03-13 18:04:30 +00:00
|
|
|
"source": [
|
2024-03-13 20:57:23 +00:00
|
|
|
"total_timesteps = 100000\n",
|
|
|
|
"model = DQN(\"MlpPolicy\", env, verbose=0)\n",
|
|
|
|
"model.learn(total_timesteps=total_timesteps, progress_bar=True)"
|
2024-03-13 18:04:30 +00:00
|
|
|
]
|
|
|
|
},
|
|
|
|
{
|
|
|
|
"cell_type": "code",
|
2024-03-13 20:57:23 +00:00
|
|
|
"execution_count": null,
|
2024-03-13 18:04:30 +00:00
|
|
|
"metadata": {},
|
|
|
|
"outputs": [],
|
|
|
|
"source": [
|
|
|
|
"def test(model):\n",
|
|
|
|
"\n",
|
|
|
|
" end_rewards = []\n",
|
|
|
|
"\n",
|
|
|
|
" for i in range(1000):\n",
|
|
|
|
" \n",
|
|
|
|
" state = env.reset()\n",
|
|
|
|
"\n",
|
|
|
|
" done = False\n",
|
|
|
|
"\n",
|
|
|
|
" while not done:\n",
|
|
|
|
"\n",
|
|
|
|
" action, _states = model.predict(state, deterministic=True)\n",
|
|
|
|
"\n",
|
|
|
|
" state, reward, done, info = env.step(action)\n",
|
2024-03-13 20:57:23 +00:00
|
|
|
" \n",
|
2024-03-13 18:04:30 +00:00
|
|
|
" end_rewards.append(reward == 0)\n",
|
2024-03-13 20:57:23 +00:00
|
|
|
" \n",
|
2024-03-13 18:04:30 +00:00
|
|
|
" return np.sum(end_rewards) / len(end_rewards)"
|
|
|
|
]
|
|
|
|
},
|
|
|
|
{
|
|
|
|
"cell_type": "code",
|
2024-03-13 20:57:23 +00:00
|
|
|
"execution_count": null,
|
2024-03-13 18:04:30 +00:00
|
|
|
"metadata": {},
|
2024-03-13 20:57:23 +00:00
|
|
|
"outputs": [],
|
2024-03-13 18:04:30 +00:00
|
|
|
"source": [
|
|
|
|
"model.save(\"dqn_wordle\")"
|
|
|
|
]
|
|
|
|
},
|
|
|
|
{
|
|
|
|
"cell_type": "code",
|
|
|
|
"execution_count": null,
|
|
|
|
"metadata": {},
|
|
|
|
"outputs": [],
|
2024-03-13 20:57:23 +00:00
|
|
|
"source": [
|
|
|
|
"model = DQN.load(\"dqn_wordle\")"
|
|
|
|
]
|
|
|
|
},
|
|
|
|
{
|
|
|
|
"cell_type": "code",
|
|
|
|
"execution_count": null,
|
|
|
|
"metadata": {},
|
|
|
|
"outputs": [],
|
|
|
|
"source": [
|
|
|
|
"print(test(model))"
|
|
|
|
]
|
2024-03-13 18:04:30 +00:00
|
|
|
}
|
|
|
|
],
|
|
|
|
"metadata": {
|
|
|
|
"kernelspec": {
|
|
|
|
"display_name": "Python 3 (ipykernel)",
|
|
|
|
"language": "python",
|
|
|
|
"name": "python3"
|
|
|
|
},
|
|
|
|
"language_info": {
|
|
|
|
"codemirror_mode": {
|
|
|
|
"name": "ipython",
|
|
|
|
"version": 3
|
|
|
|
},
|
|
|
|
"file_extension": ".py",
|
|
|
|
"mimetype": "text/x-python",
|
|
|
|
"name": "python",
|
|
|
|
"nbconvert_exporter": "python",
|
|
|
|
"pygments_lexer": "ipython3",
|
|
|
|
"version": "3.8.10"
|
|
|
|
}
|
|
|
|
},
|
|
|
|
"nbformat": 4,
|
|
|
|
"nbformat_minor": 2
|
|
|
|
}
|