Fix experiment example to support new changes (#224)

* fix experiment example * Update experiment.py * Update experiment.py * add more detail to example/experiment.py Co-authored-by: Chris Nota <cpnota@gmail.com>
cpnota · Jan 27, 2021 · 8f687f2 · 8f687f2
1 parent 90e38e5
commit 8f687f2
Show file tree

Hide file tree

Showing 2 changed files with 13 additions and 5 deletions.
diff --git a/examples/experiment.py b/examples/experiment.py
@@ -7,11 +7,19 @@
 
 
 def main():
-    device = 'cpu'
+    DEVICE = 'cpu'
+    # DEVICE = 'cuda' # uncomment for gpu support
     timesteps = 40000
     run_experiment(
-        [dqn(), a2c()],
-        [GymEnvironment('CartPole-v0', device), GymEnvironment('Acrobot-v1', device)],
+        [
+            # DQN with default hyperparameters
+            dqn.device(DEVICE),
+            # DQN with a custom hyperparameters and a custom name.
+            dqn.device(DEVICE).hyperparameters(replay_buffer_size=100).name('dqn-small-buffer'),
+            # A2C with a custom name
+            a2c.device(DEVICE).name('not-dqn')
+        ],
+        [GymEnvironment('CartPole-v0', DEVICE), GymEnvironment('Acrobot-v1', DEVICE)],
         timesteps,
     )
     plot_returns_100('runs', timesteps=timesteps)

diff --git a/examples/slurm_experiment.py b/examples/slurm_experiment.py
@@ -4,14 +4,14 @@
 For real experiments, you will surely need a modified version of this script.
 '''
 from all.experiments import SlurmExperiment
-from all.presets.atari import a2c
+from all.presets.atari import a2c, dqn
 from all.environments import AtariEnvironment
 
 
 def main():
     device = 'cuda'
     envs = [AtariEnvironment(env, device) for env in ['Pong', 'Breakout', 'SpaceInvaders']]
-    SlurmExperiment(a2c(device=device), envs, 1e6, sbatch_args={
+    SlurmExperiment([a2c, dqn], envs, 1e6, sbatch_args={
         'partition': '1080ti-short'
     })