Update plot.py

explorationlib/local_gym.py

@@ -1,5 +1,6 @@
 #! /usr/bin/env python
 import numpy as np
+import random
 
 from copy import deepcopy
 from itertools import cycle
@@ -17,6 +18,8 @@
 import warnings
 warnings.filterwarnings("ignore")
 
+#countSteps = 0
+
 # -------------------------------------------------------------------------
 # Enviroments
 
@@ -587,6 +590,150 @@ def reset(self):
         self.state_obs = (self.state, self.obs)
         self.last()
 
+# ---
+class ScentGridMovingTargets(Grid):
+    """An open-grid, with scent and moving targets"""
+    def __init__(self, mode="discrete"):
+        super().__init__(mode=mode)
+        self.scent = None
+        self.scent_fn = None
+        self.obs = 0.0
+
+    def add_scent(self,
+                  target,
+                  value,
+                  coord,
+                  scent,
+                  detection_radius=1.0,
+                  noise_sigma=0.0,
+                  p_target=1.0):
+        # Offset coords by target location
+        # align them in other words
+        self.scent_x_coord = coord[0] + target[0]
+        self.scent_y_coord = coord[1] + target[1]
+
+        # Set scent and its params
+        self.noise_sigma = noise_sigma
+        self.scent_pdf = scent
+        self.add_targets([target], [value],
+                         detection_radius=detection_radius,
+                         kd_kwargs=None,
+                         p_target=p_target)
+
+        def scent_fn(state):
+            x, y = state
+            i = find_nearest(self.scent_x_coord, x)
+            j = find_nearest(self.scent_y_coord, y)
+
+            # Add noise?
+            noise = np.abs(self.np_random.normal(0, self.noise_sigma))
+            return self.scent_pdf[i, j] + noise
+
+        self.scent_fn = scent_fn
+
+    def add_scents(
+            self,
+            targets,
+            values,
+            coord,  # assume shared
+            scents,
+            detection_radius=1.0,
+            noise_sigma=0.0,
+            p_target=1.0):
+        """Add several scents, and targets"""
+        self.noise_sigma = noise_sigma
+        self.scent_pdfs = scents
+        self.add_targets(targets,
+                         values,
+                         detection_radius=detection_radius,
+                         kd_kwargs=None,
+                         p_target=p_target)
+
+        # Offset coords by target location
+        # align them in other words
+        self.scent_x_coords = []
+        self.scent_y_coords = []
+        for target in self.targets:
+            self.scent_x_coords.append(coord[0] + target[0])
+            self.scent_y_coords.append(coord[1] + target[1])
+
+        def scent_fn(state):
+            # Pos
+            x, y = state
+
+            # Sum scents from all targets @ pos
+            summed = 0.0
+            for ind in range(self.num_targets):
+                i = find_nearest(self.scent_x_coords[ind], x)
+                j = find_nearest(self.scent_y_coords[ind], y)
+                summed += self.scent_pdfs[ind][i, j]
+
+            # Add noise?
+            noise = np.abs(self.np_random.normal(0, self.noise_sigma))
+            return summed + noise
+
+        self.scent_fn = scent_fn
+
+    def step(self, action):
+        # Move
+        super().step(action)  # sets self.ind, etc
+        #global countSteps += 1
+
+        # Scent
+        if self.scent_fn is not None:
+            x, y = int(self.state[0]), int(self.state[1])
+            self.obs = self.scent_fn((x, y))
+        else:
+            self.obs = 0.0
+
+        # New FP Code
+        newTargets = []
+        for target in self.targets:
+            #if countSteps%50 == 0:
+            #print("countSteps = ", countSteps)
+            target[0] += random.randint(-3, 3)
+            if target[0] > 10:
+                target[0] = 10
+            elif target[0] < -10:
+                target[0] = -10
+            target[1] += random.randint(-3, 3)
+            if target[1] > 10:
+                target[1] = 10
+            elif target[1] < -10:
+                target[1] = -10
+            newTargets.append(target)
+        #self.add_targets(newTargets, self.values)
+        values = self.values
+        targets = newTargets
+
+        scentsX = []
+        count = 0
+        for _ in range(len(targets)):
+          if values[count] == 1:
+            newAmp = 1
+            newSig = 2
+          else:
+            newAmp = 1.5
+            newSig = 1.75
+
+          coord, scent = create_grid_scent_patches(
+            (10, 10), p=1.0, amplitude=newAmp, sigma=newSig)
+          scentsX.append(scent)
+          count += 1
+        self.add_scents(targets, values, coord, scentsX, noise_sigma=self.noise_sigma)
+
+        # !
+        self.state_obs = (self.state, self.obs)
+        return self.last()
+
+    def last(self):
+        return (self.state_obs, self.reward, self.done, self.info)
+
+    def reset(self):
+        self.state = np.zeros(2, dtype=int)
+        self.obs = 0.0
+        self.state_obs = (self.state, self.obs)
+        self.last()
 
 # ---
 # Multi

explorationlib/plot.py

@@ -24,14 +24,14 @@ def show_gif(name):
     return display.HTML(f'<img src="data:image/gif;base64,{b64}" />')
 
 
-# def render_2d(name,
+# def render_(name,
 #               env,
 #               exp_data,
 #               num_experiment=0,
 #               figsize=(4, 4),
 #               boundary=(50, 50),
 #               interval=200):
-#     """Replay an experiment, as a movie.
+#     """Replay an experiment, as a movie. FUCK
 
 #     NOTE: can be very slow to run for experiments
 #     with more than a couple thousand steps.
@@ -261,17 +261,18 @@ def plot_scent_grid(env,
 def plot_targets2d(env,
                    figsize=(3, 3),
                    boundary=(1, 1),
-                   color="black",
+                   color=None,
                    alpha=1.0,
                    label=None,
                    title=None,
+                   differ=None,
                    ax=None):
 
     # No targets no plot
     if env.targets is None:
         return None
 
-    # Fmt
+        # Fmt
     try:
         vec = np.vstack(env.initial_targets)
     except AttributeError:
@@ -281,9 +282,10 @@ def plot_targets2d(env,
     if ax is None:
         fig = plt.figure(figsize=figsize)
         ax = fig.add_subplot(111)
+
 
-    # !
-    ax.scatter(
+
+    ax.scatter(    
         vec[:, 0],
         vec[:, 1],
         env.values,  # value is size, literal
@@ -294,16 +296,67 @@ def plot_targets2d(env,
     ax.set_ylim(-boundary[1], boundary[1])
     ax.set_xlabel("x")
     ax.set_ylabel("y")
-
+
+
     # Labels, legends, titles?
     if title is not None:
         ax.set_title(title)
     if label is not None:
         ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
 
     return ax
+'''
+def plot_targets2d(env,
+                   figsize=(3, 3),
+                   boundary=(1, 1),
+                   color=None,
+                   alpha=1.0,
+                   label=None,
+                   title=None,
+                   differ=None,
+                   ax=None):
 
+    # No targets no plot
+    if env.targets is None:
+        return None
+
+        # Fmt
+    try:
+        vec = np.vstack(env.initial_targets)
+    except AttributeError:
+        vec = np.vstack(env.targets)
+
+    # Create a fig obj?
+    if ax is None:
+        fig = plt.figure(figsize=figsize)
+        ax = fig.add_subplot(111)
+
+    if differ == None:
+        color = "black"
+
+    else:
+        color = "red"
+
+    ax.scatter(    
+        vec[:, 0],
+        vec[:, 1],
+        env.values,  # value is size, literal
+        color=color,
+        label=label,
+        alpha=alpha)
+    ax.set_xlim(-boundary[0], boundary[0])
+    ax.set_ylim(-boundary[1], boundary[1])     
+    ax.set_xlabel("x")
+    ax.set_ylabel("y")     
+
+    # Labels, legends, titles?
+    if title is not None:
+        ax.set_title(title)
+    if label is not None:
+        ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
 
+    return ax
+'''
 def plot_position2d(exp_data,
                     boundary=(1, 1),
                     figsize=(3, 3),
@@ -495,4 +548,7 @@ def plot_length_hist(exp_data,
     if label is not None:
         ax.legend(loc='center left', bbox_to_anchor=(1, 0.5))
 
-    return ax
+    return ax
+
+
+