Clean up

Author: root

pufferlib/ocean/bitflip/bitflip.h

@@ -47,16 +47,9 @@ void add_log(BitFlip *env) {
 }
 
 void c_reset(BitFlip *env) {
-  // Clear observations
   memset(env->observations, OFF, env->size * 3 * sizeof(char));
-
-  // Clear n_correct
   env->n_correct = 0;
-
-  // Always make the first bit 1 to avoid "free" rounds just by chance
   env->observations[0] = ON;
-
-  // Initialize target pattern
   for (int i = 1; i < env->size; i++) {
     env->observations[i] = (rand() % 2 == 1) ? ON : OFF;
 
@@ -65,32 +58,30 @@ void c_reset(BitFlip *env) {
       env->n_correct++;
     }
   }
-
-  // Initialize starting position in the middle
   env->pos = 2 * env->size + (env->size - 1) / 2;
   env->observations[env->pos] = CURSOR;
-
-  // Clear number of steps
   env->tick = 0;
 }
 
 void c_step(BitFlip *env) {
-  env->tick++;
+  env->tick += 1;
 
-  env->observations[env->pos] = EMPTY;
+  int action = env->actions[0];
+  env->terminals[0] = 0;
+  env->rewards[0] = 0.0;
 
-  if (env->actions[0] == LEFT) {
-    env->pos--;
-  }
+  env->observations[env->pos] = EMPTY;
 
-  if (env->actions[0] == RIGHT) {
-    env->pos++;
+  if (action == LEFT) {
+    env->pos -= 1;
+  } else if (action == RIGHT) {
+    env->pos += 1;
   }
 
   if (env->tick == 12 * env->size || env->pos < 2 * env->size ||
       env->pos >= env->size * 3) {
-    env->rewards[0] = -1.0;
     env->terminals[0] = 1;
+    env->rewards[0] = -1.0;
     add_log(env);
     c_reset(env);
     return;
@@ -101,14 +92,13 @@ void c_step(BitFlip *env) {
   int state_idx = env->pos - env->size;
   int target_idx = env->pos - 2 * env->size;
 
-  // Flip bit
-  if (env->actions[0] == FLIP) {
+  if (action == FLIP) {
     env->observations[state_idx] ^= 1;
 
     if (env->observations[state_idx] == env->observations[target_idx]) {
-      env->n_correct++;
+      env->n_correct += 1;
     } else {
-      env->n_correct--;
+      env->n_correct -= 1;
     }
   }
 
@@ -119,9 +109,6 @@ void c_step(BitFlip *env) {
     c_reset(env);
     return;
   }
-
-  env->rewards[0] = 0.0;
-  env->terminals[0] = 0;
 }
 
 void c_render(BitFlip *env) {