Add success function to env trait (#28)

* Add permutations as twists in training

* Fix lint

* Add success to env trait. Check it in training

* Revert notebooks reformat

* Revert notebooks reformat

* Add succes in test structures

Author: victor-villar

examples/grid_world/src/lib.rs

@@ -154,6 +154,10 @@ impl Env for GridWorld {
         }
     }
 
+    fn success(&self) -> bool {
+        self.at_goal()
+    }
+
     fn observe(&self) -> Vec<usize> {
         self.get_state().iter().enumerate().map(|(i, v)| i * self.height * self.width + v).collect()  
     }

examples/puzzle.ipynb

@@ -35,17 +35,18 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 3,
+   "execution_count": null,
    "metadata": {},
    "outputs": [],
    "source": [
     "import numpy as np\n",
     "\n",
+    "\n",
     "def display(state):\n",
     "    padding = len(str(max(state)))\n",
     "    txt = \"\"\n",
     "    for row in np.array(state).reshape([config[\"env\"][\"width\"], config[\"env\"][\"height\"]]):\n",
-    "        txt += \"|\" + \"|\".join((\"{:\"+str(padding)+\"d}\").format(num) if num > 0 else \" \"*padding for num in row) + \"|\\n\"\n",
+    "        txt += \"|\" + \"|\".join((\"{:\" + str(padding) + \"d}\").format(num) if num > 0 else \" \" * padding for num in row) + \"|\\n\"\n",
     "    print(txt)"
    ]
   },
@@ -121,7 +122,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 6,
+   "execution_count": null,
    "metadata": {},
    "outputs": [
     {

rust/src/collector/az.rs

@@ -156,6 +156,7 @@ mod tests {
         fn is_final(&self) -> bool { self.step >= 1 }
         fn reward(&self) -> f32 { 1.0 }
         fn observe(&self) -> Vec<usize> { vec![0] }
+        fn success(&self) -> bool { true }
     }
 
     fn dummy_policy() -> Policy {

rust/src/collector/ppo.rs

@@ -152,6 +152,7 @@ mod tests {
         fn is_final(&self) -> bool { self.step >= 1 }
         fn reward(&self) -> f32 { 1.0 }
         fn observe(&self) -> Vec<usize> { vec![0] }
+        fn success(&self) -> bool { true }
     }
 
     fn dummy_policy() -> Policy {

rust/src/envs/puzzle.rs

@@ -176,6 +176,10 @@ impl Env for Puzzle {
         }
     }
 
+    fn success(&self) -> bool {
+        self.solved()
+    }
+
 fn observe(&self,) -> Vec<usize> {
         self.state.iter().enumerate().map(|(i, v)| i * self.height * self.width + v).collect()
     }

rust/src/python_interface/pyenv.rs

@@ -138,6 +138,16 @@ impl Env for PyEnvImpl {
         })
     }
 
+    fn success(&self) -> bool {
+        Python::with_gil(|py| {
+            let py_env = self.py_env.borrow();
+            py_env
+                .call_method0(py, "success")
+                .and_then(|val| val.extract::<bool>(py))
+                .expect("Python `success` method must return a bool.")
+        })
+    }
+
     fn observe(&self) -> Vec<usize> {
         Python::with_gil(|py| {
             let py_env = self.py_env.borrow();

rust/src/rl/env.rs

@@ -46,6 +46,9 @@ pub trait Env : DynClone + Send + Sync {
     // Returns True if the given state is a terminal state
     fn is_final(&self) -> bool;
 
+    // Returns True if the goal is achieved
+    fn success(&self) -> bool;
+
     // Returns the value of current state
     fn reward(&self) -> f32;
 
@@ -56,4 +59,4 @@ pub trait Env : DynClone + Send + Sync {
     fn twists(&self) -> (Vec<Vec<usize>>, Vec<Vec<usize>>) {(vec![], vec![])}
 }
 
-dyn_clone::clone_trait_object!(Env);
+dyn_clone::clone_trait_object!(Env);

rust/src/rl/solve.rs

@@ -59,7 +59,9 @@ pub fn single_solve(
     let val = env.reward();
     total_val += val;
 
-    (((val == 1.0) as usize as f32, total_val), solution)
+    let success = if env.success() { 1.0 } else { 0.0 };
+
+    ((success, total_val), solution)
 }
 
 pub fn solve(