1 files changed, 35 insertions, 0 deletions
diff --git a/controllers/final/arm_controller.py b/controllers/final/arm_controller.py
new file mode 100644
index 0000000..bcaf5b7
--- /dev/null
+++ b/controllers/final/arm_controller.py
@@ -0,0 +1,35 @@
+from py_trees.behaviour import Behaviour
+from py_trees.common import Status
+from setup import robot, motors, sensors
+from constants import JOINT_TOLERANCE, TIMEOUT
+
+# Move the provided joints towards the target values.
+# Mainly used to control the arm.
+class MoveArm(Behaviour):
+    def __init__(self, name, targets):
+        super().__init__(name)
+        self.targets = targets
+
+    def initialise(self):
+        self.initial_time = robot.getTime()
+        for joint, target in self.targets.items():
+            motors[joint].setPosition(target)
+        print("%s.initialise()" % (self.name))
+
+    def update(self):
+        if robot.getTime() - self.initial_time > TIMEOUT:
+            print(
+                "%s.update()[%s->%s], Timed out"
+                % (self.name, self.status, Status.FAILURE)
+            )
+            return Status.FAILURE
+
+        # All joints are within JOINT_TOLERANCE from their respective targets.
+        if all([abs(sensors[joint].getValue() - target) < JOINT_TOLERANCE
+                for joint, target in self.targets.items()]):
+            print(
+                "%s.update()[%s->%s]"
+                % (self.name, self.status, Status.SUCCESS)
+            )
+            return Status.SUCCESS
+        return Status.RUNNING