Merge e7b47ab into 2372ac7

CHANGELOG.md

@@ -2,20 +2,17 @@
 
 All notable changes to this project will be documented in this file.
 
-## Next unreleased (hum!)
-
-### Added
-
-- Utils: base class and Upkie-specific exceptions
-
 ## Unreleased
 
 ### Added
 
 - Example: lying genuflections
+- Utils: base class and Upkie-specific exceptions
 
 ### Changed
 
+- Environment: ``UpkieServosEnv-v2`` with frequency regulation
+- Environment: ``UpkieWheelsEnv-v3`` with frequency regulation
 - Rename ``ROBOT_NAME`` to ``ROBOT`` in the main Makefile
 - Rename main repository and project to just "upkie"
 - Update Vulp to v1.2.0

README.md

@@ -28,22 +28,19 @@ Connect a USB controller to move the robot around 🎮
 
 ```python
 import gym
-import loop_rate_limiters
 import upkie.envs
 
 upkie.envs.register()
 
-with gym.make("UpkieWheelsEnv-v2") as env:
+with gym.make("UpkieWheelsEnv-v2", frequency=200.0) as env:
     observation = env.reset()
     action = 0.0 * env.action_space.sample()
-    rate = loop_rate_limiters.RateLimiter(frequency=200.0)
     for step in range(1_000_000):
         observation, reward, done, _ = env.step(action)
         if done:
             observation = env.reset()
         pitch = observation[0]
         action[0] = 10.0 * pitch
-        rate.sleep()
 ```
 
 ## Installation

envs/BUILD

@@ -17,6 +17,7 @@ py_library(
         "@upkie//observers/base_pitch",
         "@vulp//spine:python",
         requirement("gym"),
+        requirement("loop_rate_limiters"),
         requirement("pyyaml"),
         requirement("upkie_description"),
     ],

envs/tests/upkie_base_env_test.py

@@ -22,7 +22,7 @@
 import numpy as np
 import posix_ipc
 
-from upkie.envs import UpkieBaseEnv
+from envs import UpkieBaseEnv
 
 
 class MockSpine:
@@ -76,7 +76,10 @@ def setUp(self):
             shm_name, posix_ipc.O_RDWR | posix_ipc.O_CREAT, size=42
         )
         self.env = UpkieBaseChild(
-            config=None, fall_pitch=1.0, shm_name=shm_name
+            config=None,
+            fall_pitch=1.0,
+            frequency=100.0,
+            shm_name=shm_name,
         )
         shared_memory.close_fd()
         self.env._spine = MockSpine()

envs/tests/upkie_servos_env_test.py

@@ -22,7 +22,7 @@
 import numpy as np
 import posix_ipc
 
-from upkie.envs import UpkieServosEnv
+from envs import UpkieServosEnv
 
 
 class MockSpine:
@@ -68,7 +68,11 @@ def setUp(self):
         shared_memory = posix_ipc.SharedMemory(
             shm_name, posix_ipc.O_RDWR | posix_ipc.O_CREAT, size=42
         )
-        self.env = UpkieServosEnv(shm_name=shm_name)
+        self.env = UpkieServosEnv(
+            fall_pitch=1.0,
+            frequency=100.0,
+            shm_name=shm_name,
+        )
         shared_memory.close_fd()
         self.env._spine = MockSpine()
 

envs/tests/upkie_wheels_env_test.py

@@ -22,7 +22,7 @@
 import numpy as np
 import posix_ipc
 
-from upkie.envs import UpkieWheelsEnv
+from envs import UpkieWheelsEnv
 
 
 class MockSpine:
@@ -64,7 +64,11 @@ def setUp(self):
         shared_memory = posix_ipc.SharedMemory(
             shm_name, posix_ipc.O_RDWR | posix_ipc.O_CREAT, size=42
         )
-        self.env = UpkieWheelsEnv(shm_name=shm_name)
+        self.env = UpkieWheelsEnv(
+            fall_pitch=1.0,
+            frequency=100.0,
+            shm_name=shm_name,
+        )
         shared_memory.close_fd()
         self.env._spine = MockSpine()
 

envs/upkie_base_env.py

@@ -16,10 +16,12 @@
 # limitations under the License.
 
 import abc
+import asyncio
 from typing import Dict, Optional, Tuple, Union
 
 import gym
 import numpy as np
+from loop_rate_limiters import AsyncRateLimiter, RateLimiter
 from vulp.spine import SpineInterface
 
 from upkie.observers.base_pitch import compute_base_pitch_from_imu
@@ -54,6 +56,9 @@ class UpkieBaseEnv(abc.ABC, gym.Env):
     real robot, as it relies on the same spine interface that runs on Upkie.
     """
 
+    __frequency: float
+    __async_rate: Optional[AsyncRateLimiter]
+    __rate: Optional[RateLimiter]
     _spine: SpineInterface
     config: dict
     fall_pitch: float
@@ -62,17 +67,20 @@ def __init__(
         self,
         config: Optional[dict],
         fall_pitch: float,
+        frequency: float,
         shm_name: str,
     ) -> None:
         """!
         Initialize environment.
 
         @param config Configuration dictionary, also sent to the spine.
         @param fall_pitch Fall pitch angle, in radians.
+        @param frequency Regulated frequency of the control loop, in Hz.
         @param shm_name Name of shared-memory file.
         """
         if config is None:
             config = DEFAULT_CONFIG
+        self.__frequency = frequency
         self._spine = SpineInterface(shm_name)
         self.config = config
         self.fall_pitch = fall_pitch
@@ -103,6 +111,7 @@ def reset(
                 It is only returned if ``return_info`` is set to true.
         """
         # super().reset(seed=seed)  # we are pinned at gym==0.21.0
+        self.__reset_rates()
         self._spine.stop()
         self._spine.start(self.config)
         self._spine.get_observation()  # might be a pre-reset observation
@@ -114,11 +123,42 @@ def reset(
         else:  # return_info
             return observation, observation_dict
 
+    def __reset_rates(self):
+        self.__async_rate = None
+        self.__rate = None
+        try:
+            asyncio.get_running_loop()
+            self.__async_rate = AsyncRateLimiter(self.__frequency)
+        except RuntimeError:  # not asyncio
+            self.__rate = RateLimiter(self.__frequency)
+
     def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, dict]:
         """!
-        Run one timestep of the environment's dynamics. When end of episode is
-        reached, you are responsible for calling `reset()` to reset the
-        environment's state.
+        Run one timestep of the environment's dynamics. When the end of the
+        episode is reached, you are responsible for calling `reset()` to reset
+        the environment's state.
+
+        @param action Action from the agent.
+        @returns
+            - ``observation``: Agent's observation of the environment.
+            - ``reward``: Amount of reward returned after previous action.
+            - ``done``: Whether the agent reaches the terminal state, which can
+              be a good or a bad thing. If true, the user needs to call
+              :func:`reset()`.
+            - ``info``: Contains auxiliary diagnostic information (helpful for
+              debugging, logging, and sometimes learning).
+        """
+        action_dict = self.dictionarize_action(action)
+        self.__rate.sleep()  # wait until clock tick to send the action
+        return self.__step(action_dict)
+
+    async def async_step(
+        self, action: np.ndarray
+    ) -> Tuple[np.ndarray, float, bool, dict]:
+        """!
+        Run one timestep of the environment's dynamics using asynchronous I/O.
+        When the end of the episode is reached, you are responsible for calling
+        `reset()` to reset the environment's state.
 
         @param action Action from the agent.
         @returns
@@ -131,27 +171,32 @@ def step(self, action: np.ndarray) -> Tuple[np.ndarray, float, bool, dict]:
               debugging, logging, and sometimes learning).
         """
         action_dict = self.dictionarize_action(action)
+        await self.__async_rate.sleep()  # send action at next clock tick
+        return self.__step(action_dict)
+
+    def __step(
+        self, action_dict: dict
+    ) -> Tuple[np.ndarray, float, bool, dict]:
         self._spine.set_action(action_dict)
         observation_dict = self._spine.get_observation()
-        imu = observation_dict["imu"]
-        # TODO(scaron): use tilt (angle to the vertical) rather than pitch
-        pitch = compute_base_pitch_from_imu(imu["orientation"])
         observation = self.vectorize_observation(observation_dict)
         reward = self.reward.get(observation)
-        done = self.detect_fall(pitch)
+        done = self.detect_fall(observation_dict)
         info = {
             "action": action_dict,
             "observation": observation_dict,
         }
         return observation, reward, done, info
 
-    def detect_fall(self, pitch: float) -> bool:
+    def detect_fall(self, observation_dict: dict) -> bool:
         """!
         Detect a fall based on the body-to-world pitch angle.
 
         @param pitch Current pitch angle in [rad].
         @returns True if and only if a fall is detected.
         """
+        imu = observation_dict["imu"]
+        pitch = compute_base_pitch_from_imu(imu["orientation"])
         return abs(pitch) > self.fall_pitch
 
     @abc.abstractmethod

envs/upkie_servos_env.py

@@ -97,22 +97,29 @@ class UpkieServosEnv(UpkieBaseEnv):
 
     reward: StandingReward
     robot: pin.RobotWrapper
-    version: int = 1
+    version: int = 2
 
     def __init__(
         self,
         config: Optional[dict] = None,
         fall_pitch: float = 1.0,
+        frequency: float = 200.0,
         shm_name: str = "/vulp",
     ):
         """!
         Initialize environment.
 
         @param config Configuration dictionary, also sent to the spine.
         @param fall_pitch Fall pitch angle, in radians.
+        @param frequency Regulated frequency of the control loop, in Hz.
         @param shm_name Name of shared-memory file.
         """
-        super().__init__(config, fall_pitch, shm_name)
+        super().__init__(
+            config=config,
+            fall_pitch=fall_pitch,
+            frequency=frequency,
+            shm_name=shm_name,
+        )
 
         robot = upkie_description.load_in_pinocchio(root_joint=None)
         model = robot.model

envs/upkie_wheels_env.py

@@ -20,8 +20,8 @@
 from typing import Optional
 
 import numpy as np
-
 from gym import spaces
+
 from upkie.observers.base_pitch import compute_base_pitch_from_imu
 
 from .standing_reward import StandingReward
@@ -79,7 +79,7 @@ class UpkieWheelsEnv(UpkieBaseEnv):
 
     fall_pitch: float
     max_ground_velocity: float
-    version: int = 2
+    version: int = 3
     wheel_radius: float
 
     LEG_JOINTS = [
@@ -92,6 +92,7 @@ def __init__(
         self,
         config: Optional[dict] = None,
         fall_pitch: float = 1.0,
+        frequency: float = 200.0,
         max_ground_velocity: float = 1.0,
         shm_name: str = "/vulp",
         wheel_radius: float = 0.06,
@@ -101,11 +102,17 @@ def __init__(
 
         @param config Configuration dictionary, also sent to the spine.
         @param fall_pitch Fall pitch angle, in radians.
+        @param frequency Regulated frequency of the control loop, in Hz.
         @param max_ground_velocity Maximum commanded ground velocity in m/s.
         @param shm_name Name of shared-memory file.
         @param wheel_radius Wheel radius in [m].
         """
-        super().__init__(config, fall_pitch, shm_name)
+        super().__init__(
+            config=config,
+            fall_pitch=fall_pitch,
+            frequency=frequency,
+            shm_name=shm_name,
+        )
 
         observation_limit = np.array(
             [

examples/BUILD

@@ -11,11 +11,30 @@ load("//tools/lint:lint.bzl", "add_lint_tests")
 load("@pip_upkie//:requirements.bzl", "requirement")
 
 py_binary(
-    name = "wheeled_balancing",
-    srcs = ["wheeled_balancing.py"],
+    name = "cpu_isolation",
+    srcs = ["cpu_isolation.py"],
+    deps = ["//envs"],
+)
+
+py_binary(
+    name = "lying_genuflection",
+    srcs = ["lying_genuflection.py"],
     deps = [
         "//envs",
+        "//observers/base_pitch",
     ],
 )
 
+py_binary(
+    name = "wheeled_balancing",
+    srcs = ["wheeled_balancing.py"],
+    deps = ["//envs"],
+)
+
+py_binary(
+    name = "wheeled_balancing_async",
+    srcs = ["wheeled_balancing_async.py"],
+    deps = ["//envs"],
+)
+
 add_lint_tests()

examples/README.md

@@ -1,7 +1,5 @@
 # Examples
 
-Wheeled balancing is our entry-level
+Wheeled balancing is our simplest example: [`wheeled_balancing.py`](wheeled_balancing.py).
 
-- [`wheeled_balancing.py`](wheeled_balancing.py): smallest example, it balances Upkie with a proportional wheel controller.
-- [`wheeled_balancing_with_logging.py`](wheeled_balancing_with_logging.py): adds logging of actions and observations using [`asynchronous I/O`](https://docs.python.org/3/library/asyncio.html).
-- [`wheeled_balancing_cpu_isolation.py`](wheeled_balancing_cpu_isolation.py): on the real robot, the previous example sometimes warns that the "rate limiter is late", as its process runs on the same CPU core as all others. This examples fixes this by isolating the process to its own core.
+When running on Upkie, we often add logging (check out [`wheeled_balancing_async.py`](wheeled_balancing_async.py)) as well as CPU isolation for real-time performance (see [`wheeled_balancing_cpu_isolation.py`](wheeled_balancing_cpu_isolation.py)).