In [ ]:

    

import sys

import numpy as np
import math

# python3 -m pip install pyopengl
from OpenGL.GL import *
from OpenGL.GLU import *
from OpenGL.GLUT import *

# python3 -m pip install pygame
import pygame
from pygame.locals import *

from draw_object import ConfigureEnv


    

In [ ]:

    

class CartPoleState(object):
    def __init__(self):   
        self.x=0.
        self.x1=0.
        self.th=.2;
        self.th1=0. 

        # init constants
        self.tau = 1/60.;
        Mp = 1
        Mc = 1
        self.l = 1
        self.c1 = 1/(Mp+Mc)
        self.c2 = self.l*Mp/(Mp+Mc)
        self.g = 9.8
        
        # change the dynamicsNoise here
        self.dynamicsNoise = 0.
    
    def step(self, u):
        the2 = self.g*math.sin(self.th) + math.cos(self.th)*(-self.c1*u-self.c2*self.th1*self.th1*math.sin(self.th));
        the2 /= self.l*4/3 - self.c2*math.cos(self.th)*math.cos(self.th);
        x2 = self.c1*u + self.c2*(self.th1*self.th1*math.sin(self.th) - the2*math.cos(self.th));
        
        self.x   += self.tau*self.x1;
        self.x1  += self.tau*x2;
        self.th  += self.tau*self.th1;
        self.th1 += self.tau*the2;
        
        if(self.dynamicsNoise):
            self.x1 += self.dynamicsNoise*np.random.normal(0,1);
            self.th1 += self.dynamicsNoise*np.random.normal(0,1);


    

In [ ]:

    

def draw_objects(env, x, theta):
    # draw guide lines
    env.draw_line()  

    # cart
    glColor3f(1,1,1); 
    glPushMatrix()
    glTranslatef(x, 0., 0.)
    env.draw_cube(1., 0.2, 0.2)  
    env.draw_line2()
    
    # pole
    glColor3f(1,0,0); 
    glPushMatrix()
    glRotatef(theta*180./math.pi, 0., 1., 0.)
    glTranslatef(0, 0., 0.5)
    env.draw_cube(0.1, 0.1, 1.)  
    
    glPopMatrix()
    glPopMatrix()


    

In [ ]:

    

def testMove(s):
    # implement the controller gains here
    K = np.zeros(4)
    u = K[0]*s.x + K[1]*s.x1 + K[2]*s.th + K[3]*s.th1;
    s.step(u)
    
    # return translation x and rotation theta
    return s.x, s.th


    

In [ ]:

    

def main():
    pygame.init()
 
    display = (1000,750)
    pygame.display.set_mode(display, DOUBLEBUF|OPENGL, RESIZABLE)

    glRotatef(-90.,1.,0, 0)
    glScaled(0.8, 0.8, 0.8);

    env = ConfigureEnv()
    s = CartPoleState()

    while True:
        for event in pygame.event.get():
            if event.type == pygame.QUIT:
                pygame.quit()
                quit()
            env.mouseMove(event);

        glClear(GL_COLOR_BUFFER_BIT|GL_DEPTH_BUFFER_BIT)
        
        # calculate and update opengl
        x, th = testMove(s)
        draw_objects(env, x, th)
        
        pygame.display.flip()
        pygame.time.wait(10)

main()


    

In [ ]: