Merge branch 'main' into dev

.devcontainer/Dockerfile

@@ -25,11 +25,7 @@ ARG USER_UID=1000
 ARG USER_GID=$USER_UID
 
 RUN groupadd --gid $USER_GID $USERNAME \
-    && useradd --uid $USER_UID --gid $USER_GID -m $USERNAME \
-    && apt-get update \
-    && apt-get install -y sudo \
-    && echo $USERNAME ALL=\(root\) NOPASSWD:ALL > /etc/sudoers.d/$USERNAME \
-    && chmod 0440 /etc/sudoers.d/$USERNAME
+    && useradd --uid $USER_UID --gid $USER_GID -m $USERNAME
 
 # Set environment variables for CUDA
 ENV PATH="/usr/local/cuda-12.0/bin:${PATH}"
@@ -39,7 +35,18 @@ ENV NVIDIA_VISIBLE_DEVICES ${NVIDIA_VISIBLE_DEVICES:-all}
 ENV NVIDIA_DRIVER_CAPABILITIES ${NVIDIA_DRIVER_CAPABILITIES:+$NVIDIA_DRIVER_CAPABILITIES,}graphics
 
 # Create and set proper permissions for workspace directory
-RUN mkdir -p /workspace && chown -R $USERNAME:$USERNAME /workspace && sudo chmod -R 777 /workspace
+RUN mkdir -p /workspace && chown -R $USERNAME:$USERNAME /workspace && chmod -R 755 /workspace
+
+# Create and set permissions for Artemis testing directory
+# RUN mkdir -p /var/tmp/testing-dir && \
+#     chmod -R 755 /var/tmp/testing-dir && \
+#     chown -R $USERNAME:$USERNAME /var/tmp/testing-dir
+
+# # Ensure /var/tmp exists with correct permissions
+# RUN chmod 1777 /var/tmp && \
+#     touch /var/tmp/script.sh && \
+#     chmod 755 /var/tmp/script.sh && \
+#     chown $USERNAME:$USERNAME /var/tmp/script.sh
 
 # Create virtual environment with proper permissions
 ENV VIRTUAL_ENV=/home/vscode/venv
@@ -62,5 +69,5 @@ COPY --chown=$USERNAME:$USERNAME pyproject.toml .
 # Install dependencies as non-root user
 RUN pip install --no-cache-dir -e ".[test,gpu]"
 
-# Set final working directory and ensure proper permissions
+# Set final working directory
 WORKDIR /workspace

.devcontainer/devcontainer.linux.json

@@ -20,6 +20,12 @@
             ],
             "settings": {
                 "python.defaultInterpreterPath": "/home/vscode/venv/bin/python",
+                "jupyter.kernels.excludePythonEnvironments": [
+                    "/usr/bin/python",
+                    "/usr/bin/python3.10",
+                    "/usr/python3.10",
+                    "/usr/python3.10"
+                ], // hide unrequired python interpreters, so that the right one gets selected for jupyter https://github.com/microsoft/vscode/issues/130946#issuecomment-1899389049
                 "python.testing.pytestEnabled": true, // test framework
                 "python.testing.pytestArgs": [
                     "-m",

.devcontainer/devcontainer.wsl2.json

@@ -20,6 +20,12 @@
             ],
             "settings": {
                 "python.defaultInterpreterPath": "/home/vscode/venv/bin/python",
+                "jupyter.kernels.excludePythonEnvironments": [
+                    "/usr/bin/python",
+                    "/usr/bin/python3.10",
+                    "/usr/python3.10",
+                    "/usr/python3.10"
+                ], // hide unrequired python interpreters, so that the right one gets selected for jupyter https://github.com/microsoft/vscode/issues/130946#issuecomment-1899389049
                 "python.testing.pytestEnabled": true, // test framework
                 "python.testing.pytestArgs": [
                     "-m",
@@ -47,7 +53,8 @@
         "source=/tmp/.X11-unix,target=/tmp/.X11-unix,type=bind",
         "source=/mnt/wslg,target=/mnt/wslg,type=bind",
         "source=/usr/lib/wsl,target=/usr/lib/wsl,type=bind"
-    ],"containerEnv": {
+    ],
+    "containerEnv": {
         "DISPLAY": "${localEnv:DISPLAY}",
         "XAUTHORITY": "${localEnv:XAUTHORITY}",
         "WAYLAND_DISPLAY": "${localEnv:WAYLAND_DISPLAY}", // WSL2

README.md

@@ -25,7 +25,9 @@ For windows, we require WSL2 to run the devcontainer. (So its actually Linux wit
 
 **MacOS**
 
-Should work like Linux. Untested.
+Unfortunately, we did not get the devcontainer to work with MacOS yet, even after following [those](https://gist.github.com/sorny/969fe55d85c9b0035b0109a31cbcb088) steps. We expect that the issue is related to Mujoco rendering from inside the Docker container and display forwarding with X11. There is also an [unresolved Issue](https://github.com/google-deepmind/mujoco/issues/1047) on GitHub. If you manage to make it work, please let us know.
+
+Until then, MacOS users are required to install this project using an python environment manager such as [conda](https://docs.anaconda.com/anaconda/install/) or [mamba](https://mamba.readthedocs.io/en/latest/). If you use conda, these are the required commands: ```conda create --name crazyflow -c conda-forge python=3.11```, ```conda activate crazyflow```, ```conda install pip```, ```pip install -e .```.
 
 ____________
 

crazyflow/gymnasium_envs/crazyflow.py

@@ -243,8 +243,9 @@ def _obs(self) -> dict[str, Array]:
 class CrazyflowEnvReachGoal(CrazyflowBaseEnv):
     """JAX Gymnasium environment for Crazyflie simulation."""
 
-    def __init__(self, **kwargs: dict):
+    def __init__(self, render_goal_marker: bool = False, **kwargs: dict):
         super().__init__(**kwargs)
+        self.render_goal_marker = render_goal_marker
         spec = {k: v for k, v in self.single_observation_space.items()}
         spec["difference_to_goal"] = spaces.Box(-np.inf, np.inf, shape=(3,))
         self.single_observation_space = spaces.Dict(spec)
@@ -277,6 +278,18 @@ def reset_masked(self, mask: Array) -> None:
         )
         self.goal = self.goal.at[mask].set(new_goals[mask])
 
+    def step(self, action: Array) -> tuple[Array, Array, Array, Array, dict]:
+        if self.render_goal_marker:
+            for i in range(self.sim.n_worlds):
+                if hasattr(self.sim, "viewer") and self.sim.viewer is not None:
+                    self.sim.viewer.viewer.add_marker(
+                        type=mujoco.mjtGeom.mjGEOM_SPHERE,
+                        size=np.array([0.02, 0.02, 0.02]),
+                        pos=np.array(self.goal[i]),
+                        rgba=np.array([1, 0, 0, 0.5]),
+                    )
+        return super().step(action)
+
     def _obs(self) -> dict[str, Array]:
         obs = super()._obs()
         obs["difference_to_goal"] = [self.goal - self.sim.data.states.pos]
@@ -330,8 +343,9 @@ def _obs(self) -> dict[str, Array]:
 class CrazyflowEnvLanding(CrazyflowBaseEnv):
     """JAX Gymnasium environment for Crazyflie simulation."""
 
-    def __init__(self, **kwargs: dict):
+    def __init__(self, render_landing_marker: bool = False, **kwargs: dict):
         super().__init__(**kwargs)
+        self.render_landing_target = render_landing_marker
         spec = {k: v for k, v in self.single_observation_space.items()}
         spec["difference_to_goal"] = spaces.Box(-np.inf, np.inf, shape=(3,))
         self.single_observation_space = spaces.Dict(spec)
@@ -356,6 +370,18 @@ def _reward(prev_done: Array, terminated: Array, states: SimState, goal: Array)
     def reset_masked(self, mask: Array) -> None:
         super().reset_masked(mask)
 
+    def step(self, action: Array) -> tuple[Array, Array, Array, Array, dict]:
+        if self.render_landing_target:
+            for i in range(self.sim.n_worlds):
+                if hasattr(self.sim, "viewer") and self.sim.viewer is not None:
+                    self.sim.viewer.viewer.add_marker(
+                        type=mujoco.mjtGeom.mjGEOM_SPHERE,
+                        size=np.array([0.02, 0.02, 0.02]),
+                        pos=np.array(self.goal[i]),
+                        rgba=np.array([1, 0, 0, 0.5]),
+                    )
+        return super().step(action)
+
     def _obs(self) -> dict[str, Array]:
         obs = super()._obs()
         obs["difference_to_goal"] = [self.goal - self.sim.data.states.pos]
@@ -379,11 +405,9 @@ def render_trajectory(viewer: MujocoRenderer | None, pos: Array) -> None:
 class CrazyflowEnvFigureEightTrajectory(CrazyflowBaseEnv):
     """JAX Gymnasium environment for Crazyfly simulation.
 
-    This environment is used to follow a figure-eight trajectory. The trajectory is defined as a
-    parametrized scipy spline using `splprep` in 3D space. Sampling of the trajectory for the
-    observations can be configured in the `__init__` method. The observations contain the relative
-    position errors to the next `n_trajectory_sample_points` points that are distanced by
-    `dt_trajectory_sample_points`. The reward is based on the distance to the next trajectory point.
+    This environment is used to follow a figure-eight trajectory. The observations contain the
+    relative position errors to the next `n_samples` points that are distanced by `samples_dt`. The
+    reward is based on the distance to the next trajectory point.
     """
 
     def __init__(
@@ -398,7 +422,7 @@ def __init__(
 
         Args:
             n_samples: Number of next trajectory points to sample for observations.
-            dt_trajectory_sample_points: Time between trajectory sample points in seconds.
+            samples_dt: Time between trajectory sample points in seconds.
             trajectory_time: Total time for completing the figure-eight trajectory in seconds.
             render_samples: Flag to enable/disable rendering of the trajectory sample.
             **kwargs: Arguments passed to the Crazyfly simulation.
@@ -506,3 +530,21 @@ def actions(self, actions: Array) -> Array:
         # Ensure actions are within the valid range of the simulation action space
         rescaled_actions = np.clip(rescaled_actions, self.action_sim_low, self.action_sim_high)
         return rescaled_actions
+
+
+def render_trajectory(viewer: MujocoRenderer | None, pos: Array) -> None:
+    """Render trajectory."""
+    if viewer is None:
+        return
+
+    pos = np.array(pos[0]).transpose(1, 0, 2)
+    n_trace, n_drones = len(pos) - 1, len(pos[0])
+
+    for i in range(n_trace):
+        for j in range(n_drones):
+            viewer.viewer.add_marker(
+                type=mujoco.mjtGeom.mjGEOM_SPHERE,
+                size=np.array([0.02, 0.02, 0.02]),
+                pos=pos[i][j],
+                rgba=np.array([1, 0, 0, 0.8]),
+            )

pyproject.toml

@@ -1,6 +1,7 @@
 [build-system]
 requires = ["setuptools>=61.0.0", "wheel", "numpy"]
 build-backend = "setuptools.build_meta"
+requires-python = "3.11" # tested in python 3.11
 
 [project]
 name = "crazyflow"