diff --git a/ball/pybullet_ball.py b/ball/pybullet_ball.py index 4bf586d..08bb211 100644 --- a/ball/pybullet_ball.py +++ b/ball/pybullet_ball.py @@ -1,4 +1,4 @@ -from typing import List +from typing import List, Tuple from ball.abc_ball import ABCBall @@ -36,6 +36,14 @@ 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 @@ -43,3 +51,12 @@ def set_position(self, 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 diff --git a/benchmarkmain.py b/benchmarkmain.py index 5be8051..2700fc7 100644 --- a/benchmarkmain.py +++ b/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 diff --git a/data_gatherer.py b/data_gatherer.py new file mode 100644 index 0000000..27023ab --- /dev/null +++ b/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) diff --git a/example_csv.csv b/example_csv.csv new file mode 100644 index 0000000..8dc346a --- /dev/null +++ b/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 diff --git a/neural_net_example.py b/neural_net_example.py new file mode 100644 index 0000000..164d67e --- /dev/null +++ b/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") diff --git a/neural_networks/__init__.py b/neural_networks/__init__.py new file mode 100644 index 0000000..e69de29 diff --git a/neural_networks/models.py b/neural_networks/models.py new file mode 100644 index 0000000..cab6659 --- /dev/null +++ b/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 diff --git a/neural_networks/net_utils.py b/neural_networks/net_utils.py new file mode 100644 index 0000000..6175107 --- /dev/null +++ b/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() diff --git a/paddle/paddle.py b/paddle/paddle.py index ccf7a34..e0e080e 100644 --- a/paddle/paddle.py +++ b/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( diff --git a/position_prediction/linear_regression.py b/position_prediction/linear_regression.py new file mode 100644 index 0000000..739f471 --- /dev/null +++ b/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 diff --git a/utils/environment.py b/utils/environment.py index e4aa46f..d2f493b 100644 --- a/utils/environment.py +++ b/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