98 lines
2.7 KiB
Python
98 lines
2.7 KiB
Python
import numpy as np
|
|
import scipy
|
|
from scipy import integrate
|
|
from .env import Env
|
|
from ..common.utils import update_state_with_Runge_Kutta
|
|
|
|
class NonlinearSampleEnv(Env):
|
|
""" Nonlinear Sample Env
|
|
"""
|
|
def __init__(self):
|
|
"""
|
|
"""
|
|
self.config = {"state_size" : 2,\
|
|
"input_size" : 1,\
|
|
"dt" : 0.01,\
|
|
"max_step" : 250,\
|
|
"input_lower_bound": [-0.5],\
|
|
"input_upper_bound": [0.5],
|
|
}
|
|
|
|
super(NonlinearSampleEnv, self).__init__(self.config)
|
|
|
|
def reset(self, init_x=np.array([2., 0.])):
|
|
""" reset state
|
|
Returns:
|
|
init_x (numpy.ndarray): initial state, shape(state_size, )
|
|
info (dict): information
|
|
"""
|
|
self.step_count = 0
|
|
|
|
self.curr_x = np.zeros(self.config["state_size"])
|
|
|
|
if init_x is not None:
|
|
self.curr_x = init_x
|
|
|
|
# goal
|
|
self.g_x = np.array([0., 0.])
|
|
|
|
# clear memory
|
|
self.history_x = []
|
|
self.history_g_x = []
|
|
|
|
return self.curr_x, {"goal_state": self.g_x}
|
|
|
|
def step(self, u):
|
|
"""
|
|
Args:
|
|
u (numpy.ndarray) : input, shape(input_size, )
|
|
Returns:
|
|
next_x (numpy.ndarray): next state, shape(state_size, )
|
|
cost (float): costs
|
|
done (bool): end the simulation or not
|
|
info (dict): information
|
|
"""
|
|
# clip action
|
|
u = np.clip(u,
|
|
self.config["input_lower_bound"],
|
|
self.config["input_upper_bound"])
|
|
|
|
funtions = [self._func_x_1, self._func_x_2]
|
|
|
|
next_x = update_state_with_Runge_Kutta(self._curr_x, u,
|
|
functions, self.config["dt"])
|
|
|
|
# cost
|
|
cost = 0
|
|
cost = np.sum(u**2)
|
|
cost += np.sum((self.curr_x - self.g_x)**2)
|
|
|
|
# save history
|
|
self.history_x.append(next_x.flatten())
|
|
self.history_g_x.append(self.g_x.flatten())
|
|
|
|
# update
|
|
self.curr_x = next_x.flatten()
|
|
# update costs
|
|
self.step_count += 1
|
|
|
|
return next_x.flatten(), cost, \
|
|
self.step_count > self.config["max_step"], \
|
|
{"goal_state" : self.g_x}
|
|
|
|
def _func_x_1(self, x_1, x_2, u):
|
|
"""
|
|
"""
|
|
x_dot = x_2
|
|
return x_dot
|
|
|
|
def _func_x_2(self, x_1, x_2, u):
|
|
"""
|
|
"""
|
|
x_dot = (1. - x_1**2 - x_2**2) * x_2 - x_1 + u
|
|
return x_dot
|
|
|
|
def plot_func(self, to_plot, i=None, history_x=None, history_g_x=None):
|
|
"""
|
|
"""
|
|
raise ValueError("NonlinearSampleEnv does not have animation") |