add style linting

.flake8

@@ -0,0 +1,14 @@
+[flake8]
+exclude = .git
+max-line-length = 88
+select = E,F,W,C
+ignore=W503,
+       E203,
+       E731,
+       E722,
+       F401,
+       F841,
+       E402,
+       E741,
+       E501,
+       C406,

.github/workflows/pre-commit.yaml

@@ -0,0 +1,14 @@
+name: pre-commit
+
+on:
+  pull_request:
+  push:
+    branches: [main]
+
+jobs:
+  pre-commit:
+    runs-on: ubuntu-latest
+    steps:
+    - uses: actions/checkout@v3
+    - uses: actions/setup-python@v3
+    - uses: pre-commit/[email protected]

.gitignore

@@ -0,0 +1,7 @@
+# Ignore Python files
+*.pyc
+__pycache__/
+*.pyo
+
+# MUJOCO_LOG.TXT
+MUJOCO_LOG.TXT

.pre-commit-config.yaml

@@ -0,0 +1,34 @@
+default_language_version:
+    python: python3.10
+repos:
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v2.3.0
+    hooks:
+      - id: check-yaml
+      - id: check-ast
+      - id: check-added-large-files
+      - id: check-case-conflict
+      - id: check-merge-conflict
+      - id: end-of-file-fixer
+      - id: trailing-whitespace
+      - id: detect-private-key
+      - id: debug-statements
+  - repo: https://github.com/psf/black
+    rev: 22.10.0
+    hooks:
+      - id: black
+  - repo: https://github.com/PyCQA/flake8
+    rev: 6.1.0
+    hooks:
+      - id: flake8
+  - repo: https://github.com/pycqa/isort
+    rev: 5.12.0
+    hooks:
+      - id: isort
+  - repo: https://github.com/srstevenson/nb-clean
+    rev: 3.1.0
+    hooks:
+      - id: nb-clean
+        args:
+          - --remove-empty-cells
+          - --preserve-cell-outputs

README.md

@@ -1,5 +1,7 @@
 # SERL: A Software Suite for Sample-Efficient Robotic Reinforcement Learning
 
+![](https://github.com/rail-berkeley/serl/workflows/pre-commit/badge.svg)
+
 ## Installation
 - Conda Environment:
-    - create an environment with `conda create -n serl_dev python=3.10`
+    - create an environment with `conda create -n serl_dev python=3.10`

franka_sim/.gitignore

@@ -159,4 +159,4 @@ cython_debug/
 #  option (not recommended) you can uncomment the following to ignore the entire idea folder.
 #.idea/
 
-*.mp4
+*.mp4

franka_sim/franka_sim/__init__.py

@@ -1,18 +1,20 @@
-from franka_sim.mujoco_gym_env import MujocoGymEnv, GymRenderingSpec
+from franka_sim.mujoco_gym_env import GymRenderingSpec, MujocoGymEnv
+
 __all__ = [
     "MujocoGymEnv",
     "GymRenderingSpec",
 ]
 
 from gymnasium.envs.registration import register
+
 register(
-    id='PandaPickCube-v0',
-    entry_point='franka_sim.envs:PandaPickCubeGymEnv',
+    id="PandaPickCube-v0",
+    entry_point="franka_sim.envs:PandaPickCubeGymEnv",
     max_episode_steps=100,
 )
 register(
-    id='PandaPickCubeVision-v0',
-    entry_point='franka_sim.envs:PandaPickCubeGymEnv',
+    id="PandaPickCubeVision-v0",
+    entry_point="franka_sim.envs:PandaPickCubeGymEnv",
     max_episode_steps=100,
-    kwargs={'image_obs': True},
-)
+    kwargs={"image_obs": True},
+)

franka_sim/franka_sim/controllers/opspace.py

@@ -1,7 +1,8 @@
 from typing import Optional, Tuple, Union
+
+import mujoco
 import numpy as np
 from dm_robotics.transformations import transformations as tr
-import mujoco
 
 
 def pd_control(

franka_sim/franka_sim/envs/panda_pick_gym_env.py

@@ -1,12 +1,13 @@
-import mujoco
 from pathlib import Path
-import numpy as np
 from typing import Any, Literal
+
 import gymnasium as gym
+import mujoco
+import numpy as np
 from gymnasium import spaces
 
-from franka_sim.mujoco_gym_env import MujocoGymEnv, GymRenderingSpec
 from franka_sim.controllers import opspace
+from franka_sim.mujoco_gym_env import GymRenderingSpec, MujocoGymEnv
 
 _HERE = Path(__file__).parent
 _XML_PATH = _HERE / "xmls" / "arena.xml"
@@ -63,41 +64,65 @@ def __init__(
         self._pinch_site_id = self._model.site("pinch").id
         self._block_z = self._model.geom("block").size[2]
 
-        self.observation_space = gym.spaces.Dict({
-            'state': gym.spaces.Dict(
-                {
-                    "panda/tcp_pos": spaces.Box(-np.inf, np.inf, shape=(3,), dtype=np.float32),
-                    "panda/tcp_vel": spaces.Box(-np.inf, np.inf, shape=(3,), dtype=np.float32),
-                    "panda/gripper_pos": spaces.Box(-np.inf, np.inf, shape=(1,), dtype=np.float32),
-                    # "panda/joint_pos": spaces.Box(-np.inf, np.inf, shape=(7,), dtype=np.float32),
-                    # "panda/joint_vel": spaces.Box(-np.inf, np.inf, shape=(7,), dtype=np.float32),
-                    # "panda/joint_torque": specs.Array(shape=(21,), dtype=np.float32),
-                    # "panda/wrist_force": specs.Array(shape=(3,), dtype=np.float32),
-                    "block_pos": spaces.Box(-np.inf, np.inf, shape=(3,), dtype=np.float32),
-                }
-            ),
-        })
+        self.observation_space = gym.spaces.Dict(
+            {
+                "state": gym.spaces.Dict(
+                    {
+                        "panda/tcp_pos": spaces.Box(
+                            -np.inf, np.inf, shape=(3,), dtype=np.float32
+                        ),
+                        "panda/tcp_vel": spaces.Box(
+                            -np.inf, np.inf, shape=(3,), dtype=np.float32
+                        ),
+                        "panda/gripper_pos": spaces.Box(
+                            -np.inf, np.inf, shape=(1,), dtype=np.float32
+                        ),
+                        # "panda/joint_pos": spaces.Box(-np.inf, np.inf, shape=(7,), dtype=np.float32),
+                        # "panda/joint_vel": spaces.Box(-np.inf, np.inf, shape=(7,), dtype=np.float32),
+                        # "panda/joint_torque": specs.Array(shape=(21,), dtype=np.float32),
+                        # "panda/wrist_force": specs.Array(shape=(3,), dtype=np.float32),
+                        "block_pos": spaces.Box(
+                            -np.inf, np.inf, shape=(3,), dtype=np.float32
+                        ),
+                    }
+                ),
+            }
+        )
 
         if self.image_obs:
-            self.observation_space = gym.spaces.Dict({
-                'state': gym.spaces.Dict({
-                        "panda/tcp_pos": spaces.Box(-np.inf, np.inf, shape=(3,), dtype=np.float32),
-                        "panda/tcp_vel": spaces.Box(-np.inf, np.inf, shape=(3,), dtype=np.float32),
-                        "panda/gripper_pos": spaces.Box(-np.inf, np.inf, shape=(1,), dtype=np.float32),
-                    }),
-                'images': gym.spaces.Dict({
-                        "front": gym.spaces.Box(
-                            low=0,
-                            high=255,
-                            shape=(render_spec.height, render_spec.width, 3),
-                            dtype=np.uint8,),
-                        "wrist": gym.spaces.Box(
-                            low=0,
-                            high=255,
-                            shape=(render_spec.height, render_spec.width, 3),
-                            dtype=np.uint8,),
-                }),
-            })
+            self.observation_space = gym.spaces.Dict(
+                {
+                    "state": gym.spaces.Dict(
+                        {
+                            "panda/tcp_pos": spaces.Box(
+                                -np.inf, np.inf, shape=(3,), dtype=np.float32
+                            ),
+                            "panda/tcp_vel": spaces.Box(
+                                -np.inf, np.inf, shape=(3,), dtype=np.float32
+                            ),
+                            "panda/gripper_pos": spaces.Box(
+                                -np.inf, np.inf, shape=(1,), dtype=np.float32
+                            ),
+                        }
+                    ),
+                    "images": gym.spaces.Dict(
+                        {
+                            "front": gym.spaces.Box(
+                                low=0,
+                                high=255,
+                                shape=(render_spec.height, render_spec.width, 3),
+                                dtype=np.uint8,
+                            ),
+                            "wrist": gym.spaces.Box(
+                                low=0,
+                                high=255,
+                                shape=(render_spec.height, render_spec.width, 3),
+                                dtype=np.uint8,
+                            ),
+                        }
+                    ),
+                }
+            )
 
         self.action_space = gym.spaces.Box(
             low=np.asarray([-1.0, -1.0, -1.0, -1.0]),
@@ -106,13 +131,16 @@ def __init__(
         )
 
         from gymnasium.envs.mujoco.mujoco_rendering import MujocoRenderer
+
         self._viewer = MujocoRenderer(
             self.model,
             self.data,
         )
         self._viewer.render(self.render_mode)
 
-    def reset(self, seed=None, **kwargs) -> tuple[dict[str, np.ndarray], dict[str, Any]]:
+    def reset(
+        self, seed=None, **kwargs
+    ) -> tuple[dict[str, np.ndarray], dict[str, Any]]:
         mujoco.mj_resetData(self._model, self._data)
 
         # Reset arm to home position.
@@ -135,7 +163,7 @@ def reset(self, seed=None, **kwargs) -> tuple[dict[str, np.ndarray], dict[str, A
         obs = self._compute_observation()
         return obs, {}
 
-    '''
+    """
     take a step in the environment.
     Params:
         action: np.ndarray
@@ -146,8 +174,11 @@ def reset(self, seed=None, **kwargs) -> tuple[dict[str, np.ndarray], dict[str, A
         done: bool,
         truncated: bool,
         info: dict[str, Any]
-    '''
-    def step(self, action: np.ndarray) -> tuple[dict[str, np.ndarray], float, bool, bool, dict[str, Any]]:
+    """
+
+    def step(
+        self, action: np.ndarray
+    ) -> tuple[dict[str, np.ndarray], float, bool, bool, dict[str, Any]]:
         x, y, z, grasp = action
 
         # Set the mocap position.
@@ -185,14 +216,16 @@ def step(self, action: np.ndarray) -> tuple[dict[str, np.ndarray], float, bool,
     def render(self):
         rendered_frames = []
         for cam_id in self.camera_id:
-            rendered_frames.append(self._viewer.render(render_mode='rgb_array', camera_id=cam_id))
+            rendered_frames.append(
+                self._viewer.render(render_mode="rgb_array", camera_id=cam_id)
+            )
         return rendered_frames
 
     # Helper methods.
 
     def _compute_observation(self) -> dict:
         obs = {}
-        obs['state'] = {}
+        obs["state"] = {}
 
         # joint_pos = np.stack(
         #     [self._data.sensor(f"panda/joint{i}_pos").data for i in range(1, 8)],
@@ -205,11 +238,13 @@ def _compute_observation(self) -> dict:
         # obs["panda/joint_vel"] = joint_vel.astype(np.float32)
 
         tcp_pos = self._data.sensor("2f85/pinch_pos").data
-        obs['state']["panda/tcp_pos"] = tcp_pos.astype(np.float32)
+        obs["state"]["panda/tcp_pos"] = tcp_pos.astype(np.float32)
         tcp_vel = self._data.sensor("2f85/pinch_vel").data
-        obs['state']["panda/tcp_vel"] = tcp_vel.astype(np.float32)
-        gripper_pos = np.array(self._data.ctrl[self._gripper_ctrl_id] / 255, dtype=np.float32)
-        obs['state']["panda/gripper_pos"] = gripper_pos
+        obs["state"]["panda/tcp_vel"] = tcp_vel.astype(np.float32)
+        gripper_pos = np.array(
+            self._data.ctrl[self._gripper_ctrl_id] / 255, dtype=np.float32
+        )
+        obs["state"]["panda/gripper_pos"] = gripper_pos
         # joint_torque = np.stack(
         # [self._data.sensor(f"panda/joint{i}_torque").data for i in range(1, 8)],
         # ).ravel()
@@ -219,13 +254,13 @@ def _compute_observation(self) -> dict:
         # obs["panda/wrist_force"] = symlog(wrist_force.astype(np.float32))
 
         if self.image_obs:
-            obs['images'] = {}
-            obs['images']['front'], obs['images']['wrist'] = self.render()
+            obs["images"] = {}
+            obs["images"]["front"], obs["images"]["wrist"] = self.render()
         else:
             block_pos = self._data.sensor("block_pos").data.astype(np.float32)
-            obs['state']["block_pos"] = block_pos
+            obs["state"]["block_pos"] = block_pos
 
-        if self.render_mode == 'human':
+        if self.render_mode == "human":
             self._viewer.render(self.render_mode)
 
         return obs
@@ -240,8 +275,9 @@ def _compute_reward(self) -> float:
         rew = 0.3 * r_close + 0.7 * r_lift
         return rew
 
-if __name__ == '__main__':
-    env = PandaPickCubeGymEnv(render_mode='human')
+
+if __name__ == "__main__":
+    env = PandaPickCubeGymEnv(render_mode="human")
     env.reset()
     for i in range(100):
         env.step(np.random.uniform(-1, 1, 4))

franka_sim/franka_sim/mujoco_gym_env.py

@@ -1,9 +1,11 @@
-import gymnasium as gym
-from typing import Optional, Literal
+from dataclasses import dataclass
 from pathlib import Path
+from typing import Literal, Optional
+
+import gymnasium as gym
 import mujoco
 import numpy as np
-from dataclasses import dataclass
+
 
 @dataclass(frozen=True)
 class GymRenderingSpec:
@@ -12,6 +14,7 @@ class GymRenderingSpec:
     camera_id: str | int = -1
     mode: Literal["rgb_array", "human"] = "rgb_array"
 
+
 class MujocoGymEnv(gym.Env):
     """MujocoEnv with gym interface."""
 

franka_sim/franka_sim/test/test_gym_env_human.py

@@ -1,7 +1,9 @@
 import time
+
 import mujoco
 import mujoco.viewer
 import numpy as np
+
 from franka_sim import envs
 
 env = envs.PandaPickCubeGymEnv(action_scale=(0.1, 1))
@@ -19,6 +21,7 @@ def sample():
 reset = False
 KEY_SPACE = 32
 
+
 def key_callback(keycode):
     if keycode == KEY_SPACE:
         global reset