Merge branch 'trebuchet' of github.com:MikeInnes/differentiable-rl in…

…to trebuchet
kianaf · Mar 3, 2019 · 237c09f · 237c09f
2 parents 87794b3 + 6708baa
commit 237c09f
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 6 deletions.
diff --git a/games/pendulum/DDPG.jl b/games/pendulum/DDPG.jl
@@ -186,11 +186,11 @@ function episode!(env, train=true)
     a = action(s, train)
     s′, r, done, _ = step!(env, a)
     total_reward += data(r)[1]
-    s = s′ 
     if train    
       remember(s, a, r, s′, done)
       replay()
     end
+    s = s′ 
   end
   total_reward
 end

diff --git a/games/pendulum/DiffRL.jl b/games/pendulum/DiffRL.jl
@@ -14,6 +14,7 @@ env = PendulumEnv()
 STATE_SIZE = length(reset!(env)) # returns state from obs space
 ACTION_SIZE = 1#length(env.actions)
 ACTION_BOUND = 2#env.action_space.hi
+MAX_REWARD = 0f0 # Max reward in a timestep
 MAX_EP = 15_000
 MAX_EP_LENGTH = 1000
 SEQ_LEN = 4
@@ -28,11 +29,7 @@ model = Chain(Dense(STATE_SIZE, 24, relu),
 
 opt = ADAM(η)
 
-function loss(r)
-  seq_len = size(r, 1)
-  z = zeros(Float32, seq_len) |> gpu
-  Flux.mse(r, z)
-end
+loss(r) = Flux.mse(r, MAX_REWARD)
 
 # ----------------------------- Helper Functions -------------------------------