dpendulum.py

import time
 
import numpy as np
from numpy import pi
from pendulum import Pendulum
 
p = Pendulum(1)
 
NQ = 20  # Discretization steps for position
NV = 19  # Discretization steps for velocity
VMAX = 5  # Max velocity (v in [-vmax,vmax])
NU = 9  # Discretization steps for torque
UMAX = 5  # Max torque (u in [-umax,umax])
DT = 3e-1
 
DQ = 2 * pi / NQ
DV = 2.0 * (VMAX) / NV
DU = 2.0 * (UMAX) / NU
 
 
# Continuous to discrete
def c2dq(q):
    q = (q + pi) % (2 * pi)
    return int(round(q / DQ)) % NQ
 
 
def c2dv(v):
    v = np.clip(v, -VMAX + 1e-3, VMAX - 1e-3)
    return int(np.floor((v + VMAX) / DV))
 
 
def c2du(u):
    u = np.clip(u, -UMAX + 1e-3, UMAX - 1e-3)
    return int(np.floor((u + UMAX) / DU))
 
 
def c2d(qv):
    """From continuous to discrete."""
    return c2dq(qv[0]), c2dv(qv[1])
 
 
# Discrete to continuous
def d2cq(iq):
    iq = np.clip(iq, 0, NQ - 1)
    return iq * DQ - pi
 
 
def d2cv(iv):
    iv = np.clip(iv, 0, NV - 1) - (NV - 1) / 2
    return iv * DV
 
 
def d2cu(iu):
    iu = np.clip(iu, 0, NU - 1) - (NU - 1) / 2
    return iu * DU
 
 
def d2c(iqv):
    """From discrete to continuous"""
    return d2cq(iqv[0]), d2cv(iqv[1])
 
 
def x2i(x):
    return x[0] + x[1] * NQ
 
 
def i2x(i):
    return [i % NQ, i / NQ]
 
 
# --- PENDULUM
 
 
class DPendulum:
    def __init__(self):
        self.pendulum = Pendulum(1)
        self.pendulum.DT = DT
        self.pendulum.NDT = 5
 
    @property
    def nqv(self):
        return [NQ, NV]
 
    @property
    def nx(self):
        return NQ * NV
 
    @property
    def nu(self):
        return NU
 
    @property
    def goal(self):
        return x2i(c2d([0.0, 0.0]))
 
    def reset(self, x=None):
        if x is None:
            x = [np.random.randint(0, NQ), np.random.randint(0, NV)]
        else:
            x = i2x(x)
        assert len(x) == 2
        self.x = x
        return x2i(self.x)
 
    def step(self, iu):
        self.x = self.dynamics(self.x, iu)
        reward = 1 if x2i(self.x) == self.goal else 0
        return x2i(self.x), reward
 
    def render(self):
        q = d2cq(self.x[0])
        self.pendulum.display(
            np.array(
                [
                    q,
                ]
            )
        )
        time.sleep(self.pendulum.DT)
 
    def dynamics(self, ix, iu):
        x = np.array(d2c(ix))
        u = d2cu(iu)
 
        self.xc, _ = self.pendulum.dynamics(x, u)
        return c2d(x.T.tolist()[0])
 
 
"""
env = DPendulum()
 
print env.reset(x2i([14,11]))
hq = []
hv = []
hqc = []
hvc = []
u = 0
for i in range(100):
    ix,r=env.step(u)
    q,v = i2x(ix)
    env.render()
    if d2cv(v)==0.0: u = MAXU-1 if u==0 else 0
    hq.append( d2cq(env.x[0]) )
    hv.append( d2cv(env.x[1]) )
    hqc.append( env.xc[0,0] )
    hvc.append( env.xc[1,0] )
 
"""
 
"""
 
 
EPS = 1e-3
q = 0.0
v = -VMAX
hq = []
hv = []
hiq = []
hiv = []
hqa = []
hva = []
 
while q<2*pi:
    hq.append(q)
    iq = c2dq(q)
    hiq.append(iq)
    hqa.append(d2cq(iq))
    q += EPS
while v<VMAX:
    iv = c2dv(v)
    hv.append(v)
    hiv.append(iv)
    hva.append(d2cv(iv))
    v += EPS
 
 
 
 
"""