diff --git a/js/agent.js b/js/agent.js
index ada4556..d8e98ab 100644
--- a/js/agent.js
+++ b/js/agent.js
@@ -5,10 +5,11 @@ function Agent(opt, world) {
 
     this.world = world
     this.frequency = 20
-    this.reward = 0
     this.loaded = false
 
-    this.loss = 0
+    this.infos = []
+    this.maxInfos = 2000
+
     this.timer = 0
     this.timerFrequency = 60 / this.frequency
     this.resetFrequency = 30 / this.frequency
@@ -39,7 +40,7 @@ Agent.prototype.init = function (actor, critic) {
         temporalWindow: temporal, 
 
         discount: 0.97, 
-        rate: 0.004,
+        rate: 0.002,
         theta: 0.05, // progressive copy
         alpha: 0.1, // advantage learning
 
@@ -68,21 +69,21 @@ Agent.prototype.step = function () {
     this.timer++
 
     if (this.timer % this.timerFrequency === 0) {
-        
-        // reward from last step
-        this.reward = this.walker.reward
-        // this.reward = this.walker.score * 10
-        // // punish for using energy
-        // this.reward -= this.walker.joints.map(j=>j.GetJointSpeed()).reduce((sum,speed)=>sum+speed**2)
-        // // console.log(this.reward)
-        // this.walker.score = 0
+        var [state, reward, done, info] = this.walker.simulationStep()
+
+        if (done) {
+            // TODO reset?
+        }
+
+        this.infos.push(info)
+        if (this.infos.length>this.maxInfos) this.infos = this.infos.slice(1)
         
         // train
-        this.loss = this.brain.learn(this.reward) 
-        this.action = this.brain.policy(this.walker.getState())
+        this.loss = this.brain.learn(reward) 
+        this.action = this.brain.policy(state)
     }
     if (this.action) {
-        this.walker.simulationStep(this.action)
+        this.walker.simulationPreStep(this.action)
     }
 
     return this.timer % this.timerFrequency === 0
diff --git a/js/walker.js b/js/walker.js
index 8054dcc..23049a8 100644
--- a/js/walker.js
+++ b/js/walker.js
@@ -329,19 +329,23 @@ Walker.prototype.getState = function () {
     }, [])
 }
 
-Walker.prototype.simulationStep = function (motorSpeeds) {
-  
+Walker.prototype.simulationPreStep = function (motorSpeeds) {  
   // act
-  for(var k = 0; k < this.joints.length; k++) {
-    this.joints[k].SetMotorSpeed(motorSpeeds[k]*3); // action can range from -3 to 3, radians per second
+  for (var k = 0; k < this.joints.length; k++) {
+    this.joints[k].SetMotorSpeed(motorSpeeds[k] * 3); // action can range from -3 to 3, radians per second
   }
+}
+
+Walker.prototype.simulationStep = function (motorSpeeds) {
 
   /* score/reward */
   var head_height_reward = this.head.head.GetPosition().y * 2;  // it's head should be above it's feet 2*(-0.25-2)
 
   // TODO reward for moving one leg beyond the other?
   var left_leg_forward = this.right_leg.foot.GetPosition().x > this.left_leg.foot.GetPosition().x;
-  var leg_switch_reward = (left_leg_forward!=this.last_left_left_forward)? 5:0
+  var leg_switch_reward = (left_leg_forward!=this.last_left_left_forward)? 1:0
+
+  var joint_angle_cost = - 0.02 * this.joints.map(j=>j.GetJointAngle()-j.GetReferenceAngle()).reduce((o,v)=>o+v*v,0) // - 0.02 * (0 - 20)
 
   // reward copied from OpenAI Gym Humanoid Walker https://github.com/openai/gym/blob/master/gym/envs/mujoco/humanoid.py
   // also see https://github.com/AdamStelmaszczyk/learning2run/blob/master/osim-rl/osim/env/run.py#L67
@@ -355,7 +359,7 @@ Walker.prototype.simulationStep = function (motorSpeeds) {
 
   // punish for using energy, squared
   var quad_ctrl_cost = -0.01 * this.joints.map(j => j.GetJointSpeed()).reduce((sum, speed) => sum + speed ** 2)
-  var alive_bonus = 1
+  var alive_bonus = 5 // May they find happiness for all their days
 
   // we don't have data on external forces, so I will just punish for contact
   var contacts = this.bodies.map(b => b.GetContactList()).filter(b => b).length
@@ -365,6 +369,7 @@ Walker.prototype.simulationStep = function (motorSpeeds) {
     lin_vel_reward,
     quad_ctrl_cost,
     quad_impact_cost,
+    joint_angle_cost,
     alive_bonus,
     head_height_reward,
     leg_switch_reward
@@ -374,5 +379,7 @@ Walker.prototype.simulationStep = function (motorSpeeds) {
 
   this.last_left_left_forward = left_leg_forward
 
-  return
+  var info = {}
+  var done = 0
+  return [this.getState(), this.reward, done, info]
 }