def __init__(self, learningArgs):
lr = learningArgs[0]
lrDecay_stepSize = learningArgs[1]
lrDecay_gamma = learningArgs[2]
weight_decay = learningArgs[3]
self.device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
self.MotionModel = simpleMotionModel().to(self.device)
self.optimizer = Adam(self.MotionModel.parameters(),
lr=lr,
weight_decay=weight_decay)
self.lrScheduler = torch.optim.lr_scheduler.StepLR(
self.optimizer, step_size=lrDecay_stepSize, gamma=lrDecay_gamma)
self.criterion = torch.nn.MSELoss()
from simController import simController
from motionModel import simpleMotionModel
import matplotlib.pyplot as plt
import torch
import numpy as np
physicsClientId = p.connect(p.GUI)
sim = simController(physicsClientId=physicsClientId)
sim.terrain.generate()
#sim.terrain.generate(cellHeightScale=0,perlinHeightScale=0)
sim.resetClifford()
startState = sim.controlLoopStep([0,0])
predictedPose = [startState[3][0][0],startState[3][0][1],startState[4]]
device = 'cuda:0' if torch.cuda.is_available() else 'cpu'
motionModel = simpleMotionModel().to(device)
#motionModel.load_state_dict(torch.load('motionModels/simple.pt'))
motionModel.eval()
actualX = [startState[3][0][0]]
actualY = [startState[3][0][1]]
predX = [startState[3][0][0]]
predY = [startState[3][0][1]]
for i in range(50):
print(i)
data = sim.controlLoopStep(sim.randomDriveAction())
if data[2]:
break
inAction = torch.FloatTensor(data[0][2]).unsqueeze(0).to(device)
prediction = motionModel([inAction])[0]
print(prediction)