Neural mlp

ball/pybullet_ball.py

@@ -1,4 +1,4 @@
-from typing import List
+from typing import List, Tuple
 
 from ball.abc_ball import ABCBall
 
@@ -36,10 +36,27 @@ def set_ball_angular_velocity(self, angular_velocity):
             self.id, angularVelocity=angular_velocity
         )
 
+    def set_ball_velocity(self, linear_velocity, angular_velocity):
+        self.pybullet_client.resetBaseVelocity(
+            self.id, linearVelocity=linear_velocity, angularVelocity=angular_velocity
+        )
+
+    def stabilize_ball(self):
+        self.set_ball_velocity([0, 0, 0], [0, 0, 0])
+
     def set_position(self, position, orientation):
         self.pybullet_client.resetBasePositionAndOrientation(
             self.id, position, orientation
         )
 
     def get_position(self) -> List[float]:
         return self.pybullet_client.getBasePositionAndOrientation(self.id)[0]
+
+    def get_orientation(self) -> List[float]:
+        return self.pybullet_client.getBasePositionAndOrientation(self.id)[1]
+
+    def get_velocity(self) -> Tuple[List[float], List[float]]:
+        linear_velocity, angular_velocity = self.pybullet_client.getBaseVelocity(
+            self.id
+        )
+        return linear_velocity, angular_velocity

benchmarkmain.py

@@ -11,6 +11,7 @@
 
 parser = argparse.ArgumentParser()
 DEFAULT_FILE_NAME = "polynomial_prediction_benchmark"
+
 DEFAULT_N_PREDICT = 10
 DEFAULT_FETCH_TIME = 1 / 20  # 20 is the maximum number of camera outputs per second.
 
@@ -27,6 +28,7 @@
 parser.add_argument("--delete", action="store_true")
 args = parser.parse_args()
 
+
 N_DELAYED = args.d
 N_PREDICT = args.p
 FETCH_TIME = args.f

data_gatherer.py

@@ -0,0 +1,136 @@
+#!/usr/bin/env python3
+import argparse
+import time
+
+
+import pybullet as p
+import pandas as pd
+import numpy as np
+
+
+from trackers.ball_tracker import BallTracker
+from utils.environment import init_minimalistic_env_and_load_assets
+from utils.pid_performer import PidPerformer
+
+INITIAL_WAIT_TIME = 3.0
+DEFAULT_FETCH_TIME = 0.1
+RANDOM_EVENT_TIME = 0.2
+DEFAULT_TRIAL_NUM = 5
+
+
+parser = argparse.ArgumentParser(description="Mode manager")
+parser.add_argument("--pid", action="store_true")
+parser.add_argument("--filename", action="store", type=str, default="gathered_data.csv")
+parser.add_argument(
+    "--fetch_time", action="store", type=float, default=DEFAULT_FETCH_TIME
+)
+parser.add_argument(
+    "--initial_wait_time", action="store", type=float, default=INITIAL_WAIT_TIME
+)
+parser.add_argument(
+    "--random_event_time", action="store", type=float, default=RANDOM_EVENT_TIME
+)
+parser.add_argument(
+    "--number_of_trials", action="store", type=int, default=DEFAULT_TRIAL_NUM
+)
+args = parser.parse_args()
+filename = args.filename
+pid_flag = args.pid
+fetch_time = args.fetch_time
+initial_wait_time = args.initial_wait_time
+random_event_time = args.random_event_time
+trial_num = args.number_of_trials
+
+(ball, paddle) = init_minimalistic_env_and_load_assets(p)
+
+
+paddle.create_joint_controllers()
+
+if pid_flag:
+    pid_performer = PidPerformer(p, BallTracker(ball, paddle), paddle)
+
+initial_wait_timer = time.time()
+fetch_timer = time.time()
+random_event_timer = time.time()
+
+
+def distance_from_paddle_center(pos):
+    return np.linalg.norm(np.asarray(pos) - np.asarray([0.0, 0.0, 0.5]), ord=2)
+
+
+def vector_norm(vec):
+    return np.linalg.norm(np.asarray(vec), ord=2)
+
+
+trial_id = 0
+df = pd.DataFrame(
+    columns=[
+        "trial_id",
+        "pos_x",
+        "pos_y",
+        "pos_z",
+        "lin_x",
+        "lin_y",
+        "lin_z",
+        "ang_x",
+        "ang_y",
+        "ang_z",
+        "y_roll",
+        "x_roll",
+    ]
+)
+
+while True:
+    paddle.read_and_update_joint_position()
+    if time.time() - initial_wait_timer >= initial_wait_time:
+        if time.time() - fetch_timer >= fetch_time:
+            linear_velocity, angular_velocity = ball.get_velocity()
+            position = ball.get_position()
+            _, y_joint_state, x_joint_state = paddle.get_joint_rolls()
+            # rolls are at index 0 in those tuples - see pybullet documentation for further details
+            y_roll, x_roll = y_joint_state[0], x_joint_state[0]
+
+            df.loc[len(df)] = [
+                trial_id,
+                *position,
+                *linear_velocity,
+                *angular_velocity,
+                y_roll,
+                x_roll,
+            ]
+            fetch_timer = time.time()
+
+        if pid_flag:
+            pid_performer.perform_pid_step()
+
+        if time.time() - random_event_timer >= random_event_time:
+            linear_velocity, angular_velocity = ball.get_velocity()
+            if (
+                vector_norm(linear_velocity) + vector_norm(angular_velocity) < 0.5
+                and distance_from_paddle_center(ball.get_position()) < 0.2
+            ):
+                random_angular_velocity = np.random.uniform(
+                    low=-0.3, high=0.3, size=3
+                ).tolist()
+                random_linear_velocity = np.random.uniform(
+                    low=-0.7, high=0.7, size=3
+                ).tolist()
+                ball.set_ball_velocity(random_linear_velocity, random_angular_velocity)
+
+            random_event_timer = time.time()
+
+        if distance_from_paddle_center(ball.get_position()) > 0.4:
+            ball.stabilize_ball()
+            ball.set_position([0, 0, 0.5 + 0.1], ball.get_orientation())
+            # setting new trial id
+            trial_id += 1
+            if trial_id == trial_num:
+                break
+            # delaying random event occurence
+            random_event_timer = time.time()
+    p.stepSimulation()
+
+    time.sleep(0.01)  # sometimes pybullet crashes, this line helps a lot
+
+df.to_csv(filename, sep=",")
+print("DATA GATHERED TO:", filename)

example_csv.csv

@@ -0,0 +1,6 @@
+1,2,3
+4,5,6
+1,1,1
+2,2,2
+3,3,3
+10,10,11

neural_net_example.py

@@ -0,0 +1,27 @@
+import neural_networks.models as models
+import neural_networks.net_utils as utils
+import torch
+
+net = models.MLP("example_MLP_network", [2, 10, 1])
+
+dataset = utils.CSVDataset("example_csv.csv", [0, 1], [2])
+dataloader = utils.create_dataloader(dataset, 2)
+
+print(dataset.data, dataloader)
+for input, labels in dataloader:
+    print("input: ", input)
+    print("labels: ", labels)
+    print("predicted:", net(input))
+
+optimizer = torch.optim.SGD(net.parameters(), lr=1, momentum=0.9)
+
+# using same dataloader on train and eval just for example
+utils.train(net, dataloader, dataloader, optimizer, torch.nn.L1Loss(), 5, 2)
+
+for input, labels in dataloader:
+    print("input: ", input)
+    print("labels: ", labels)
+    print("predicted:", net(input))
+
+# saving net
+utils.pickle_net(net, ".", suffix="example_suffix")

neural_networks/models.py

@@ -0,0 +1,29 @@
+import torch.nn as nn
+from typing import List
+from abc import ABC
+
+
+class AbstractNet(ABC):
+    def __init__(self, *args, **kwargs):
+        if "name" not in kwargs:
+            raise "Net's constructor needs 'name' argument"
+        self.name = kwargs["name"]
+
+    def get_name(self) -> str:
+        return self.name
+
+
+class MLP(nn.Module, AbstractNet):
+    def __init__(self, name: str, sizes: List[int]):
+        AbstractNet.__init__(self, name=name)
+        nn.Module.__init__(self)
+        self.mod_list = nn.ModuleList()
+        for i in range(len(sizes) - 1):
+            self.mod_list.append(nn.Linear(sizes[i], sizes[i + 1]))
+            if i + 1 != len(sizes):
+                self.mod_list.append(nn.ReLU())
+
+    def forward(self, x):
+        for module in self.mod_list:
+            x = module(x)
+        return x

neural_networks/net_utils.py

@@ -0,0 +1,79 @@
+from neural_networks.models import AbstractNet
+import torch.nn as nn
+import torch
+import pickle
+import time
+from pathlib import Path
+from torch.utils.data import Dataset, DataLoader
+import numpy as np
+from typing import List
+
+
+class CSVDataset(Dataset):
+    # TODO: ustalic delimiter
+    def __init__(
+        self,
+        csv_path: Path,
+        x_indicies: List[int],
+        y_indicies: List[int],
+        delimiter: str = ",",
+    ) -> None:
+        super().__init__()
+        self.data = np.genfromtxt(csv_path, delimiter=delimiter)
+        self.x = torch.from_numpy(self.data[:, x_indicies]).float()
+        self.y = torch.from_numpy(self.data[:, y_indicies]).float()
+        self.size = self.data.shape[0]
+
+    def __getitem__(self, idx):
+        return self.x[idx], self.y[idx]
+
+    def __len__(self):
+        return self.size
+
+
+def create_dataloader(dataset: Dataset, batch_size: int) -> DataLoader:
+    # TODO: shuffle=True? num_workes=>1? other args
+    return DataLoader(dataset, batch_size=batch_size)
+
+
+def train(
+    net: nn.Module,
+    dataloader_train,
+    dataloader_test,
+    optim,
+    loss_function,
+    epochs,
+    eval_gap,
+):
+    for epoch in range(1, epochs + 1):
+        if epoch % eval_gap == 0:
+            net.eval()
+            total_test_loss = 0
+            for input, labels in dataloader_test:
+                out = net(input)
+                loss = loss_function(out, labels)
+                total_test_loss += torch.sum(loss).item()
+            print(
+                "AVG LOSS ON TEST SET:", total_test_loss / len(dataloader_test.dataset)
+            )
+
+        net.train()
+        total_train_loss = 0
+        for input, labels in dataloader_train:
+            optim.zero_grad()
+
+            out = net(input)
+            loss = loss_function(out, labels)
+            total_train_loss += torch.sum(loss).item()
+            loss.backward()
+            optim.step()
+        print(
+            "AVG LOSS ON TRAIN SET:", total_train_loss / len(dataloader_train.dataset)
+        )
+
+
+def pickle_net(net: AbstractNet, path: Path, suffix: str = ""):
+    path = Path(path)
+    filehandler = open(path / (net.get_name() + suffix + str(time.time())), "wb")
+    pickle.dump(net, filehandler)
+    filehandler.close()

paddle/paddle.py

@@ -52,6 +52,12 @@ def create_joint_controllers(self):
             self.pybullet_client.addUserDebugParameter("x_roll", -3.14, 3.14, 0)
         )
 
+    def get_joint_rolls(self):
+        # (jointPosition, jointVelocity, jointReactionForces, appliedJointMotorTorque)
+        return tuple(
+            self.pybullet_client.getJointState(self.robot_id, i) for i in [3, 4, 5]
+        )
+
     def read_and_update_joint_position(self):
         for i in range(len(self.joint_controllers)):
             self.pybullet_client.setJointMotorControl2(

position_prediction/linear_regression.py

@@ -0,0 +1,35 @@
+from collections import deque
+from turtle import pos
+from typing import Tuple, List
+
+from scipy import stats
+import numpy as np
+
+from position_prediction.abc_predicter import ABCPredicter
+
+
+class LinearRegressionPredicter(ABCPredicter):
+    def __init__(self, n_predict):
+        self.confirmed_positions = deque(maxlen=n_predict)
+
+    def add_position(self, position: List[float]):
+        self.confirmed_positions.append(position)
+
+    def next_position(self) -> List[float]:
+        a = list(zip(*self.confirmed_positions))
+        x_positions, y_positions, z_positions = a[0], a[1], a[2]
+        # print("ZWRACAM", self.predict(list(x_positions)))
+        return [
+            self.predict(list(x_positions)),
+            self.predict(list(y_positions)),
+            self.predict(list(z_positions)),
+        ]
+
+    def predict(self, positions: List[float]) -> float:
+        n = len(positions)
+        if n == 1:
+            return positions[0]
+
+        time_series = np.arange(n)
+        res = stats.linregress(time_series, positions)
+        return res.slope * n + res.intercept

utils/environment.py

@@ -133,3 +133,14 @@ def init_env_and_load_assets(
     ball_controller = PyBulletBallController(ball)
     paddle = load_paddle(p)
     return ball_controller, ball, paddle, wind_controllers, force_controllers
+
+
+def init_minimalistic_env_and_load_assets(
+    p,
+) -> Tuple[PyBulletBall, Paddle,]:
+    init_environment(p)
+    load_plane(p)
+    ball = PyBulletBall(p)
+    paddle = load_paddle(p)
+
+    return ball, paddle