plane_data2.py

"""
new interface to Plane Data.

simpler implementation: read in the obstacles

simulate a single long trajectory to allow robot to cover space.

This time, robot will be larger and square shaped.


"""

import matplotlib.pyplot as plt

import numpy as np
from numpy.random import randint
import os
from dataset import DataSet
import ipdb as pdb

"""
- each dataset comes up with it's own method of storing trajectories.
- this one can only save x_env and keep track of U,ps then reconstruct X
to save a lot of space

"""

num_t=128 # number of trajectories (i.e. number of initial states)
T=1000 # length of each trajectory sequence
u_dim=2 # control (action) dimension
w,h=40,40
x_dim=w*h
rw=1 # robot half-width

def get_params():
  return x_dim,u_dim,T

class PlaneData(DataSet):
  """docstring for PlaneData"""
  def __init__(self, fname, env_file):
    super(PlaneData, self).__init__()
    self.cache=fname
    self.initialized=False
    self.im=plt.imread(env_file) # grayscale
    self.params=(x_dim,u_dim,T)
    
  def is_colliding(self,p):
    if np.any([p-rw<0, p+rw>=w]):
      return True
    # check robot body overlap with obstacle field
    return np.mean(self.im[p[0]-rw:p[0]+rw+1, p[1]-rw:p[1]+rw+1]) > 0.05

  def compute_traj(self, max_dist=1):
    # computes P,U data for single trajectory
    # all P,U share the same environment obstacles.png
    P=np.zeros((T,2),dtype=np.int) # r,c position
    U=np.zeros((T,u_dim),dtype=np.int)
    P[0,:]=[rw,randint(rw,w-rw)] # initial location
    for t in range(1,T):
      p=np.copy(P[t-1,:])
      # dr direction
      d=randint(-1,2) # direction
      nsteps=randint(max_dist+1)
      dr=d*nsteps # applied control
      for i in range(nsteps):
        p[0]+=d
        if self.is_colliding(p):
          p[0]-=d
          break
      # dc direction
      d=randint(-1,2) # direction
      nsteps=randint(max_dist+1)
      dc=d*nsteps # applied control
      for i in range(nsteps):
        p[1]+=d
        if self.is_colliding(p):
          p[1]-=d # step back
          break
      P[t,:]=p
      U[t,:]=[dr,dc]
    return P,U

  def initialize(self):
    if os.path.exists(self.cache):
      self.load()
    else:
      self.precompute()
    self.initialized=True

  def compute_data(self):
    # compute multiple trajectories
    P=np.zeros((num_t,T,2),dtype=np.int)
    U=np.zeros((num_t,T,u_dim),dtype=np.int)
    for i in range(num_t):
      P[i,:,:], U[i,:,:] = self.compute_traj(max_dist=1)
    return P,U

  def precompute(self):
    print("Precomputing P,U...")
    self.P, self.U = self.compute_data()

  def save(self):
    print("Saving P,U...")
    np.savez(self.cache, P=self.P, U=self.U)

  def load(self):
    print("Loading P,U from %s..." % (self.cache))
    D=np.load(self.cache)
    self.P, self.U = D['P'], D['U']

  def getXp(self,p):
    # return image X given true state p (position) of robot
    x=np.copy(self.im)
    x[p[0]-rw:p[0]+rw+1, p[1]-rw:p[1]+rw+1]=1. # robot is white on black background
    return x.flat

  def getX(self,i,t):
    # i=trajectory index, t=time step
    return self.getXp(self.P[i,t,:])
  
  def getXTraj(self,i):
    # i=traj index
    X=np.zeros((T,x_dim),dtype=np.float)
    for t in range(T):
      X[t,:]=self.getX(i,t)
    return X

  def sample(self, batch_size, replace=True):
    """
    computes (x_t,u_t,x_{t+1}) pair
    returns tuple of 3 ndarrays with shape
    (batch,x_dim), (batch, u_dim), (batch, x_dim)
    """
    if not self.initialized:
      raise ValueError("Dataset not loaded - call PlaneData.initialize() first.")
    #traj=randint(0,num_t,size=batch_size) # which trajectory
    #tt=randint(0,T-1,size=batch_size) # time step t for each batch
    traj=np.random.choice(num_t,size=batch_size,replace=replace)
    tt=np.random.choice(T-1,size=batch_size,replace=replace)
    X0=np.zeros((batch_size,x_dim))
    U0=np.zeros((batch_size,u_dim),dtype=np.int)
    X1=np.zeros((batch_size,x_dim))
    for i in range(batch_size):
      t=tt[i]
      p=self.P[traj[i], t, :]
      X0[i,:]=self.getX(traj[i],t)
      X1[i,:]=self.getX(traj[i],t+1)
      U0[i,:]=self.U[traj[i], t, :]
    return (X0,U0,X1)

  def sample_seq(self, batch_size, seq_len, replace=True):
    """
    like sample, but returns a sequence X,U where len(X)=seq_len, len(U)=seq_len-1
    """
    # traj=randint(0,num_t,size=batch_size) # which trajectory
    # tt=randint(0,T-seq_len,size=batch_size) # start time step t for each batch
    traj=np.random.choice(num_t,size=batch_size,replace=replace)
    tt=np.random.choice(T-seq_len,size=batch_size,replace=replace)
    Xs=[0]*seq_len
    Us=[0]*seq_len
    for t in range(seq_len):
      Xs[t]=np.zeros((batch_size,x_dim))
      Us[t]=np.zeros((batch_size,u_dim),dtype=np.int)
    for i in range(batch_size):
      start_t=tt[i]
      for t in range(seq_len):
        Xs[t][i,:]=self.getX(traj[i], start_t+t)    
        Us[t][i,:]=self.U[traj[i], start_t+t, :]       
    return Xs, Us[:-1]

  def getPSpace(self):
    """
    Returns all possible positions of agent
    """
    #ww=h-2*rw
    P=np.zeros((w*h,2)) # max possible positions
    NP=np.zeros((w*h,2))
    i,j=(0,0)
    p=np.array([rw,rw]) # initial location
    for dr in range(h):
      for dc in range(w):
        pp=p+np.array([dr,dc])
        if not self.is_colliding(pp):
          P[i,:]=pp
          i+=1
        else:
          NP[j,:]=pp
          j+=1
    return P[:i,:], NP[:j,:]

  def getXPs(self, Ps):
    X=np.zeros((Ps.shape[0],x_dim))
    for i in range(Ps.shape[0]):
      X[i,:]=self.getXp(Ps[i,:])
    return X

if __name__ == "__main__":
  import matplotlib.animation as animation
  p=PlaneData("plane1.npz","env1.png")
  p.initialize()
  p.save()
  im=p.im
  A,B=im.shape

  # show sample tuples
  if True:
    fig, aa = plt.subplots(1,2)
    x0,u0,x1=p.sample(2)
    m1=aa[0].matshow(x0[0,:].reshape(w,w), cmap=plt.cm.gray, vmin = 0., vmax = 1.)
    aa[0].set_title('x(t)')
    m2=aa[1].matshow(x1[0,:].reshape(w,w), cmap=plt.cm.gray, vmin = 0., vmax = 1.)
    aa[1].set_title('x(t+1), u=(%d,%d)' % (u0[0,0],u0[0,1]))
    fig.tight_layout()
    def updatemat2(t):
      x0,u0,x1=p.sample(2)
      m1.set_data(x0[0,:].reshape(w,w))
      m2.set_data(x1[0,:].reshape(w,w))
      return m1,m2
    
    anim=animation.FuncAnimation(fig, updatemat2, frames=100, interval=1000, blit=True, repeat=True)

    Writer = animation.writers['imagemagick'] # animation.writers.avail
    writer = Writer(fps=1, metadata=dict(artist='Me'), bitrate=1800)
    anim.save('sample_obs.gif', writer=writer)

  #show trajectory 
  if True:
    fig, ax = plt.subplots()
    X=p.getXTraj(0)
    mat=ax.matshow(X[0,:].reshape((A,B)), cmap=plt.cm.gray, vmin = 0., vmax = 1.)
    def updatemat(t):
      mat.set_data(X[t,:].reshape((A,B)))
      return mat,
    anim = animation.FuncAnimation(fig, updatemat, frames=T-1, interval=30, blit=True, repeat=True)
    plt.show()