small improvs

Author: FedeClaudi

proj/environment/environment.py

@@ -101,11 +101,11 @@ def plan(self, curr_x, g_traj, itern):
 
         end = min_idx + n_ahead + pred_len
 
-        if (min_idx + n_ahead + pred_len) > len(g_traj):
-            end = len(g_traj)
+        if start + pred_len > len(g_traj):
+            end = len(g_traj) - 2
 
-        if abs(start - end) != pred_len:
-            return np.tile(g_traj[-1], (pred_len, 1))
+        # if abs(start - end) != pred_len:
+        #     return np.tile(g_traj[-1], (pred_len, 1))
 
         return g_traj[start:end]
 

proj/run/runner.py

@@ -6,13 +6,26 @@
 )
 import numpy as np
 from rich import print
+from rich.text import Text
 import logging
 
-# define progress bar
+
+class SpeedColumn(TextColumn):
+    _renderable_cache = {}
+
+    def __init__(self, *args):
+        pass
+
+    def render(self, task):
+        if task.speed is None:
+            return Text("no speed")
+        else:
+            return Text(f"{task.speed:.3f} steps/s")
+
+
 progress = Progress(
-    TextColumn(
-        "[bold magenta]Step {task.completed}/{task.total} - {task.speed} steps/s"
-    ),
+    TextColumn("[bold magenta]Step {task.completed}/{task.total}"),
+    SpeedColumn(),
     # "[progress.description]{task.description}",
     BarColumn(bar_width=None),
     "•",

run_allocentric.py

@@ -9,9 +9,6 @@
 
 # from proj.plotting.trajectories import plot_trajectory
 
-# TODO get trials from M4 catwalk and M6 catwalk only.
-# TODO trajectory: get moved awaay from start location as start frame
-# TODO find metrics to say if it was good/bad?
 
 # model = ModelPolar()
 model = Model()