Gymnasium api

crafts.py

@@ -22,6 +22,16 @@
 
 class Rotor:
     def __init__(self, x=[0., 0., 0.], wmax=4000., Tmax=4., k=0., cm=0.01, tau=0.03, Izz=1e-6, dir='cw'):
+        '''
+        x: position of the rotor in the body frame
+        wmax: maximum angular velocity
+        Tmax: maximum thrust
+        k: thrust coefficient
+        cm: moment arm
+        tau: time constant
+        Izz: moment of inertia
+        dir: direction of rotation
+        '''
         self.x = np.asarray(x)
         self.wmax = wmax
         self.Tmax = Tmax

helpers.py

@@ -0,0 +1,101 @@
+import numpy as np
+
+# tianshou code
+from tianshou.policy import SACPolicy, BasePolicy
+from tianshou.utils.net.continuous import ActorProb, Critic
+from tianshou.utils.net.common import Net
+from tianshou.data import VectorReplayBuffer
+from tianshou.trainer import OffpolicyTrainer
+from tianshou.highlevel.logger import LoggerFactoryDefault
+from tianshou.utils import WandbLogger
+
+
+import torch
+class NumpyDeque(object):
+    def __init__(self, shape:tuple, device='cpu') -> None:
+        self.shape_arr = shape
+
+        self.array = np.zeros((self.shape_arr), dtype=np.float32)
+
+    def __len__(self):
+        return self.shape_arr[1]
+
+    def append(self, els):
+        assert els.shape[0] == self.shape_arr[0] 
+
+        self.array = np.roll(self.array, els.shape[1], axis=1)
+        self.array[:,0:els.shape[1]] = els.astype(np.float32)
+
+    def reset(self, vecs=None):
+        if vecs is None:
+            self.array = np.zeros((self.shape_arr), dtype=np.float32)
+        elif isinstance(vecs,np.ndarray):
+            self.array[vecs==1] = 0.
+
+    def __call__(self):
+        return self.array
+    def __repr__(self):
+        return str(self.array)
+    def __array__(self, dtype=None):
+        if dtype:
+            return self.array.astype(dtype)
+        return self.array
+    @property
+    def shape(self):
+        return self.shape_arr
+
+
+if __name__=='__main__':
+    def create_policy(env):
+        observation_space = env.observation_space.shape or env.observation_space.n
+        action_space = env.action_space.shape or env.action_space.n
+        # create the networks behind actors and critics
+        net_a = Net(state_shape=observation_space,
+                    hidden_sizes=[64,64], device='cpu')
+        net_c1 = Net(state_shape=observation_space,action_shape=action_space,
+                        hidden_sizes=[64,64],
+                        concat=True,)
+        net_c2 = Net(state_shape=observation_space,action_shape=action_space,
+                        hidden_sizes=[64,64],
+                        concat=True,)
+
+        # create actors and critics
+        actor = ActorProb(
+            net_a,
+            action_space,
+            unbounded=True,
+            conditioned_sigma=True,
+        )
+        critic1 = Critic(net_c1, device='cpu')
+        critic2 = Critic(net_c2, device='cpu')
+
+        # create the optimizers
+        actor_optim = torch.optim.Adam(actor.parameters(), lr=1e-3)
+        critic_optim = torch.optim.Adam(critic1.parameters(), lr=1e-3)
+        critic2_optim = torch.optim.Adam(critic2.parameters(), lr=1e-3)
+
+        # create the policy
+        policy = SACPolicy(actor=actor, actor_optim=actor_optim, \
+                            critic=critic1, critic_optim=critic_optim,\
+                            critic2=critic2, critic2_optim=critic2_optim,\
+                            action_space=env.action_space,\
+                            observation_space=env.observation_space, \
+                            action_scaling=True) # make sure actions are scaled properly
+        return policy
+
+    def forward_policy(policy, state):
+        return policy(state)
+
+
+    test_qeue = NumpyDeque((3,5))
+    print(type(test_qeue.array))
+    print(test_qeue)
+    ones_1 = np.ones((3,1))
+    twos_2 = np.ones((3,2))*2
+    test_qeue.append(ones_1)
+    print(test_qeue)
+    test_qeue.append(twos_2)
+    print(test_qeue)
+    test_qeue.reset(np.array([0,1,0]))
+    print(test_qeue)
+

imu.py

@@ -0,0 +1,111 @@
+import numpy as np
+from helpers import NumpyDeque
+import matplotlib.pyplot as plt
+
+def quaternion_rotation_matrix(Q):
+    """
+    Covert a quaternion into a full three-dimensional rotation matrix.
+
+    Input
+    :param Q: A 4 element array representing the quaternion (q0,q1,q2,q3) 
+
+    Output
+    :return: A 3x3 element matrix representing the full 3D rotation matrix. 
+             This rotation matrix converts a point in the local reference 
+             frame to a point in the global reference frame.
+    """
+    # Extract the values from Q
+    q0 = Q[0]
+    q1 = Q[1]
+    q2 = Q[2]
+    q3 = Q[3]
+
+    # First row of the rotation matrix
+    r00 = 2 * (q0 * q0 + q1 * q1) - 1
+    r01 = 2 * (q1 * q2 - q0 * q3)
+    r02 = 2 * (q1 * q3 + q0 * q2)
+
+    # Second row of the rotation matrix
+    r10 = 2 * (q1 * q2 + q0 * q3)
+    r11 = 2 * (q0 * q0 + q2 * q2) - 1
+    r12 = 2 * (q2 * q3 - q0 * q1)
+
+    # Third row of the rotation matrix
+    r20 = 2 * (q1 * q3 - q0 * q2)
+    r21 = 2 * (q2 * q3 + q0 * q1)
+    r22 = 2 * (q0 * q0 + q3 * q3) - 1
+
+    # 3x3 rotation matrix
+    rot_matrix = np.array([[r00, r01, r02],
+                           [r10, r11, r12],
+                           [r20, r21, r22]])
+
+    return rot_matrix
+
+class IMU:
+    def __init__(self, noise=np.array([0, 0, 0, 0, 0, 0]), bias=np.array([0, 0, 0, 0, 0, 0]), dt=0.01, offset=np.array([0, 0, 0])):
+        '''
+        noise: standard deviation of the noise
+        bias: initial value of the bias (modelled as brownian motion)
+        offset: offset of the sensor in body frame'''
+        assert len(noise) == 6
+        assert len(bias) == 6
+        assert len(offset) == 3
+
+        self.dt = dt
+
+        self.noise = noise/np.sqrt(self.dt)
+        self.bias = bias
+
+        self.accel = np.array([0, 0, 0])
+        self.gyro = np.array([0, 0, 0])
+
+        self.offset = offset
+
+        self.vel_history = NumpyDeque((100, 3))
+        self.accel_history = NumpyDeque((100, 3))
+        self.gyro_history = NumpyDeque((100, 3))
+
+    def add_brownian_bias(self):
+        self.bias += np.random.normal(0, 0.01, size=(6,)) * np.sqrt(self.dt)
+    def add_noise(self):
+        self.add_brownian_bias()
+        self.accel = np.random.normal(self.accel, self.noise) - self.bias
+        self.gyro = np.random.normal(self.gyro, self.noise) - self.bias
+
+
+    def simulate(self, state):
+        '''
+        state: [x, y, z, vx, vy, vz, qw, qx, qy, qz, wx, wy, wz]
+        '''
+        # update acceleration
+        vel = state[3:6]
+        self.vel_history.append(vel)
+        q = state[6:10]
+        R = quaternion_rotation_matrix(q)
+
+        # transfrom velocity to body frame
+        vel_body = np.dot(R.T, vel)
+        # transform gravity to body frame
+        accel_grav = np.dot(R, np.array([0, 0, -9.81])) 
+        # compute acceleration from velocity history
+        accel_vel_change = (self.vel_history[0] - self.vel_history[1]) / self.dt
+
+        self.accel = accel_grav + accel_vel_change
+
+        self.gyro = state[10:13]
+
+        # update angular velocity
+        self.gyro = state[10:13]
+
+        self.add_noise()
+        return np.concatenate([self.accel, self.gyro])
+
+    def reset(self):
+        self.accel = np.array([0, 0, 0])
+        self.gyro = np.array([0, 0, 0])
+        self.bias = np.array([0, 0, 0, 0, 0, 0])
+
+
+    def render(self):
+        pass