diff --git a/.gitignore b/.gitignore
new file mode 100644
index 00000000..0c84676c
--- /dev/null
+++ b/.gitignore
@@ -0,0 +1,91 @@
+__pycache__/
+*.py[cod]
+
+examples/experiments/**/debug_*/
+examples/experiments/**/*.pkl
+examples/experiments/**/checkpoint_*
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# Distribution / packaging
+.eggs/
+dist/
+build/
+sdist/
+*.egg-info/
+*.whl
+
+# Environments
+.venv/
+venv/
+env/
+ENV/
+*.env
+
+# PyInstaller
+# Usually these files are written by a python script; excluding them is not always
+# appropriate.
+# https://docs.pyinstaller.org/en/stable/when-things-go-wrong.html
+# _MEIPASS*
+# _MEI*
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.coverage
+.coverage.*
+nosetests.xml
+coverage.xml
+*.log
+*.rpt
+*.db
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.sqlite3
+*.sqlitedb
+local_settings.py
+/static/
+/media/
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx stuff:
+docs/_build
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# mkdocs
+site/
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.sock
+
+# VS Code
+.vscode/
+
+# PyCharm
+.idea/
+demo_data/
+
+examples/experiments/**/debug_*/
+examples/experiments/**/*.pkl
+examples/experiments/**/checkpoint_*
+**/*.out
\ No newline at end of file
diff --git a/check_server.sh b/check_server.sh
new file mode 100644
index 00000000..dafcfae4
--- /dev/null
+++ b/check_server.sh
@@ -0,0 +1,21 @@
+echo "Getting position..."
+curl -X POST localhost:5000/getpos_euler
+
+echo ""
+echo "Activating gripper..."
+curl -X POST localhost:5000/activate_gripper
+echo ""
+echo ""
+
+echo "Closing gripper in 1s..."
+sleep 1
+
+curl -X POST localhost:5000/close_gripper
+echo ""
+echo ""
+
+echo "Opening gripper in 3s..."
+sleep 3
+
+curl -X POST localhost:5000/reset_gripper
+echo ""
diff --git a/environment.yml b/environment.yml
new file mode 100644
index 00000000..5ac0fddf
--- /dev/null
+++ b/environment.yml
@@ -0,0 +1,94 @@
+name: serl2
+channels:
+ - defaults
+ - https://repo.anaconda.com/pkgs/main
+ - https://repo.anaconda.com/pkgs/r
+dependencies:
+ - _libgcc_mutex=0.1=main
+ - _openmp_mutex=5.1=1_gnu
+ - bzip2=1.0.8=h5eee18b_6
+ - ca-certificates=2025.2.25=h06a4308_0
+ - ld_impl_linux-64=2.40=h12ee557_0
+ - libffi=3.4.4=h6a678d5_1
+ - libgcc-ng=11.2.0=h1234567_1
+ - libgomp=11.2.0=h1234567_1
+ - libstdcxx-ng=11.2.0=h1234567_1
+ - libuuid=1.41.5=h5eee18b_0
+ - ncurses=6.4=h6a678d5_0
+ - openssl=3.0.16=h5eee18b_0
+ - pip=25.0=py310h06a4308_0
+ - python=3.10.16=he870216_1
+ - readline=8.2=h5eee18b_0
+ - setuptools=75.8.0=py310h06a4308_0
+ - sqlite=3.45.3=h5eee18b_0
+ - tk=8.6.14=h39e8969_0
+ - tzdata=2025a=h04d1e81_0
+ - wheel=0.45.1=py310h06a4308_0
+ - xz=5.6.4=h5eee18b_1
+ - zlib=1.2.13=h5eee18b_1
+ - pip:
+ - absl-py==2.2.2
+ - blinker==1.9.0
+ - catkin-pkg==1.0.0
+ - certifi==2025.1.31
+ - cffi==1.17.1
+ - charset-normalizer==3.4.1
+ - click==8.1.8
+ - cloudpickle==3.1.1
+ - cython==3.0.12
+ - defusedxml==0.7.1
+ - distro==1.9.0
+ - docutils==0.21.2
+ - easyhid==0.0.10
+ - empy==4.2
+ - evdev==1.9.1
+ - flask==3.1.0
+ - gym==0.26.2
+ - gym-notices==0.0.8
+ - hidapi==0.14.0.post4
+ - idna==3.10
+ - itsdangerous==2.2.0
+ - jax==0.4.35
+ - jax-cuda12-pjrt==0.4.35
+ - jax-cuda12-plugin==0.4.35
+ - jaxlib==0.4.34
+ - jinja2==3.1.6
+ - lz4==4.4.4
+ - markupsafe==3.0.2
+ - ml-dtypes==0.5.1
+ - numpy==2.2.4
+ - nvidia-cublas-cu12==12.8.4.1
+ - nvidia-cuda-cupti-cu12==12.8.90
+ - nvidia-cuda-nvcc-cu12==12.8.93
+ - nvidia-cuda-runtime-cu12==12.8.90
+ - nvidia-cudnn-cu12==9.8.0.87
+ - nvidia-cufft-cu12==11.3.3.83
+ - nvidia-cusolver-cu12==11.7.3.90
+ - nvidia-cusparse-cu12==12.5.8.93
+ - nvidia-nccl-cu12==2.26.2
+ - nvidia-nvjitlink-cu12==12.8.93
+ - opencv-python==4.11.0.86
+ - opt-einsum==3.4.0
+ - pycparser==2.22
+ - pymodbus==2.5.3
+ - pynput==1.8.1
+ - pyopengl==3.1.9
+ - pyparsing==3.2.3
+ - pyquaternion==0.9.9
+ - pyrealsense2==2.55.1.6486
+ - pyserial==3.5
+ - pyspacemouse==1.1.4
+ - python-dateutil==2.9.0.post0
+ - python-xlib==0.33
+ - pyyaml==6.0.2
+ - pyzmq==26.3.0
+ - requests==2.32.3
+ - rospkg==1.6.0
+ - scipy==1.15.2
+ - six==1.17.0
+ - typing==3.7.4.3
+ - typing-extensions==4.13.1
+ - urllib3==2.3.0
+ - werkzeug==3.1.3
+ - zmq==0.0.0
+prefix: /home/robot/miniconda3/envs/serl
diff --git a/examples/analysis.py b/examples/analysis.py
new file mode 100644
index 00000000..2011aab2
--- /dev/null
+++ b/examples/analysis.py
@@ -0,0 +1,157 @@
+
+from absl import app, flags
+import time
+import numpy as np
+import os
+import pickle
+import imageio
+import cv2
+import queue
+from pynput import keyboard
+import threading
+from flax.training import checkpoints
+import jax
+import jax.numpy as jnp
+
+# from experiments.mappings import CONFIG_MAPPING
+from serl_launcher.agents.continuous.sac import SACAgent
+from serl_launcher.agents.continuous.sac_hybrid_single import SACAgentHybridSingleArm
+from serl_launcher.agents.continuous.sac_hybrid_dual import SACAgentHybridDualArm
+
+from serl_launcher.utils.launcher import (
+ make_sac_pixel_agent,
+ make_sac_pixel_agent_hybrid_single_arm,
+ make_sac_pixel_agent_hybrid_dual_arm,
+ make_trainer_config,
+ make_wandb_logger,
+)
+
+checkpoint_path = "/home/qirico/Desktop/All-Weird/Human-Interventions/jax-hitl-hil-serl/examples/experiments/franka_sim/debug_rlif_2"
+
+from experiments.config import DefaultTrainingConfig
+class TrainConfig(DefaultTrainingConfig):
+ image_keys = ["front", "wrist"]
+ classifier_keys = ["front", "wrist"]
+ proprio_keys = ['panda/tcp_pos', 'panda/tcp_vel', 'panda/gripper_pos']
+ # buffer_period = 1000
+ # checkpoint_period = 5000
+ # steps_per_update = 50
+ pretraining_steps = 0 # How many steps to pre-train the model for using RLPD on offline data only.
+ reward_scale = 1 # How much to scale actual rewards (not RLIF penalties) for RLIF.
+ rlif_minus_one = False
+ checkpoint_period = 2000
+ cta_ratio = 2
+ random_steps = 0
+ discount = 0.98
+ buffer_period = 1000
+ batch_size = 64
+ encoder_type = "resnet-pretrained"
+ setup_mode = "single-arm-learned-gripper"
+
+exp_name = "franka_sim"
+config = TrainConfig()
+# env = config.get_environment(fake_env=True,save_video=False,classifier=True)
+
+intervene_steps = 0 # Default number of steps between pre and post intervention states
+constraint_eps = 0.1 # Default constraint epsilon
+
+# obs_key_shapes = [('front', (1, 128, 128, 3)), ('state', (1, 7)), ('wrist', (1, 128, 128, 3))]
+obs_sample = {
+ 'front': np.zeros((1, 128, 128, 3), dtype=np.uint8),
+ 'state': np.zeros((1, 7), dtype=np.float32),
+ 'wrist': np.zeros((1, 128, 128, 3), dtype=np.uint8),
+}
+action_sample = np.zeros(7, dtype=np.float32)
+
+agent: SACAgentHybridSingleArm = make_sac_pixel_agent_hybrid_single_arm(
+ seed=0,
+ sample_obs=obs_sample,
+ sample_action=action_sample,
+ image_keys=config.image_keys,
+ encoder_type=config.encoder_type,
+ discount=config.discount,
+ enable_cl=False,
+ soft_cl = False,
+ intervene_steps=intervene_steps,
+ constraint_eps=constraint_eps,
+)
+
+ckpt = checkpoints.restore_checkpoint(
+ os.path.abspath(checkpoint_path),
+ agent.state,
+ step='40000',
+ )
+agent = agent.replace(state=ckpt)
+
+
+preference_buffer_base_path="experiments/franka_sim/debug_rlif_2/interventions/transitions"
+preference_buffer_paths = [f"{preference_buffer_base_path}_{i}.pkl" for i in range(1000, 14000, 1000)]
+
+preference_buffer = []
+
+for preference_buffer_path in preference_buffer_paths:
+ if not os.path.exists(preference_buffer_path):
+ print(f"Preference buffer path {preference_buffer_path} does not exist.")
+ continue
+
+ # Load the preference buffer
+ with open(preference_buffer_path, 'rb') as f:
+ preference_buffer_part = pickle.load(f)
+ preference_buffer.extend(preference_buffer_part)
+
+rng = jax.random.PRNGKey(0)
+
+def get_action(obs, rng):
+ rng, key = jax.random.split(rng)
+ actions = agent.sample_actions(
+ observations=jax.device_put(obs),
+ argmax=True,
+ seed=key
+ )
+ return actions, rng
+
+pre_intervention_obs = [p['observations'][0] for p in preference_buffer]
+intervene_actions = [p['actions'][0] for p in preference_buffer]
+policy_actions = [p['policy_actions'][0] for p in preference_buffer]
+post_intervention_obs = [p['observations'][-1] for p in preference_buffer]
+
+pre_intervention_obs = {
+ 'front': np.array([obs['front'] for obs in pre_intervention_obs]),
+ 'state': np.array([obs['state'] for obs in pre_intervention_obs]),
+ 'wrist': np.array([obs['wrist'] for obs in pre_intervention_obs]),
+}
+pre_intervention_expert_action, rng = get_action(pre_intervention_obs, rng)
+
+
+post_intervention_obs = {
+ 'front': np.array([obs['front'] for obs in post_intervention_obs]),
+ 'state': np.array([obs['state'] for obs in post_intervention_obs]),
+ 'wrist': np.array([obs['wrist'] for obs in post_intervention_obs]),
+}
+post_intervention_expert_action, rng = get_action(post_intervention_obs, rng)
+
+policy_actions = np.array(policy_actions)
+intervene_actions = np.array(intervene_actions)
+
+key, rng = jax.random.split(rng)
+q_pre_expert = agent.forward_critic(pre_intervention_obs, pre_intervention_expert_action[:, :6], key)
+key, rng = jax.random.split(rng)
+q_post_expert = agent.forward_critic(post_intervention_obs, post_intervention_expert_action[:, :6], key)
+key, rng = jax.random.split(rng)
+q_pre_policy = agent.forward_critic(pre_intervention_obs, policy_actions[:, :6], key)
+key, rng = jax.random.split(rng)
+q_pre_intervene = agent.forward_critic(pre_intervention_obs, intervene_actions[:, :6], key)
+# q_post_expert = agent.q_network.apply(agent.state.params, post_intervention_obs, post_intervention_expert_action)
+
+# q_pre_policy = agent.q_network.apply(agent.state.params, pre_intervention_obs, policy_actions)
+# q_post_policy = agent.q_network.apply(agent.state.params, pre_intervention_obs, policy_actions)
+constraint1_acc = ((q_pre_expert - q_post_expert) < 0).mean()
+qvalue_based_learning_intervene = ((q_pre_policy - q_pre_intervene) < 0).mean()
+qvalue_based_learning_expert = ((q_pre_policy - q_pre_expert) < 0).mean()
+constraint2_acc = ((q_pre_intervene - q_post_expert) < 0).mean()
+
+print(f"Constraint 1 accuracy: {constraint1_acc}")
+print(f"Q-value based learning intervene accuracy: {qvalue_based_learning_intervene}")
+print(f"Q-value based learning expert accuracy: {qvalue_based_learning_expert}")
+print(f"Constraint 2 accuracy: {constraint2_acc}")
+breakpoint()
\ No newline at end of file
diff --git a/examples/eval_sim_checkpoints.py b/examples/eval_sim_checkpoints.py
new file mode 100644
index 00000000..520838ca
--- /dev/null
+++ b/examples/eval_sim_checkpoints.py
@@ -0,0 +1,176 @@
+from absl import app, flags
+import time
+import numpy as np
+import os
+import pickle
+import imageio
+import cv2
+import queue
+from pynput import keyboard
+import threading
+from flax.training import checkpoints
+import jax
+import jax.numpy as jnp
+
+from experiments.mappings import CONFIG_MAPPING
+from serl_launcher.agents.continuous.sac import SACAgent
+from serl_launcher.agents.continuous.sac_hybrid_single import SACAgentHybridSingleArm
+from serl_launcher.agents.continuous.sac_hybrid_dual import SACAgentHybridDualArm
+from serl_launcher.utils.tools import ImageDisplayer, q_image
+
+from serl_launcher.utils.launcher import (
+ make_sac_pixel_agent,
+ make_sac_pixel_agent_hybrid_single_arm,
+ make_sac_pixel_agent_hybrid_dual_arm,
+ make_trainer_config,
+ make_wandb_logger,
+)
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string("exp_name", None, "Name of experiment corresponding to folder.")
+flags.DEFINE_string("method", "rlif", "valid values: rlif, cl, hil, soft_cl")
+flags.DEFINE_integer("seed", 42, "Random seed.")
+flags.DEFINE_string("checkpoint_path", None, "Path to save checkpoints.")
+flags.DEFINE_integer("eval_checkpoint_step", 0, "Step to evaluate the checkpoint.")
+flags.DEFINE_integer("eval_n_trajs", 0, "Number of trajectories to evaluate.")
+flags.DEFINE_string("save_to_txt", None, "Where to save the results to.")
+flags.DEFINE_boolean("show_q_values", False, "")
+
+
+def print_green(x):
+ return print("\033[92m {}\033[00m".format(x))
+
+def print_yellow(x):
+ return print("\033[93m {}\033[00m".format(x))
+
+def print_cyan(x):
+ return print("\033[96m {}\033[00m".format(x))
+
+
+def main(_):
+ print_green(FLAGS)
+ exp_name = FLAGS.exp_name
+ enable_cl = FLAGS.method in ["cl", "soft_cl"]
+
+ config = CONFIG_MAPPING[exp_name]()
+ env = config.get_environment(
+ fake_env=True,
+ save_video=False,
+ classifier=True,
+ state_based=False,
+ )
+ rng = jax.random.PRNGKey(FLAGS.seed)
+
+ if config.setup_mode == 'single-arm-fixed-gripper' or config.setup_mode == 'dual-arm-fixed-gripper':
+ agent: SACAgent = make_sac_pixel_agent(
+ seed=FLAGS.seed,
+ sample_obs=env.observation_space.sample(),
+ sample_action=env.action_space.sample(),
+ image_keys=config.image_keys,
+ encoder_type=config.encoder_type,
+ discount=config.discount,
+ )
+ include_grasp_penalty = False
+ elif config.setup_mode == 'single-arm-learned-gripper':
+ cl_config = {
+ "enabled": enable_cl,
+ "soft": FLAGS.method == "soft_cl",
+ "enable_margin_constraint": True,
+ "enable_action_constraint": True,
+ "constraint_coeff": 1.0,
+ "constraint_eps": 0.0,
+ "reward_coeff": 1.0,
+ }
+ print_green(f"Using CL Config: {cl_config}")
+
+ agent: SACAgentHybridSingleArm = make_sac_pixel_agent_hybrid_single_arm(
+ seed=FLAGS.seed,
+ sample_obs=env.observation_space.sample(),
+ sample_action=env.action_space.sample(),
+ image_keys=config.image_keys,
+ encoder_type=config.encoder_type,
+ discount=config.discount,
+ enable_cl=enable_cl,
+ cl_config=cl_config,
+ )
+ include_grasp_penalty = True
+ elif config.setup_mode == 'dual-arm-learned-gripper':
+ agent: SACAgentHybridDualArm = make_sac_pixel_agent_hybrid_dual_arm(
+ seed=FLAGS.seed,
+ sample_obs=env.observation_space.sample(),
+ sample_action=env.action_space.sample(),
+ image_keys=config.image_keys,
+ encoder_type=config.encoder_type,
+ discount=config.discount,
+ )
+ include_grasp_penalty = True
+ else:
+ raise NotImplementedError(f"Unknown setup mode: {config.setup_mode}")
+
+ success_counter = 0
+ time_list = []
+
+ ckpt = checkpoints.restore_checkpoint(
+ os.path.abspath(FLAGS.checkpoint_path),
+ agent.state,
+ step=FLAGS.eval_checkpoint_step,
+ )
+ agent = agent.replace(state=ckpt)
+ print_green(f"Overriding agent with checkpoint at {FLAGS.eval_checkpoint_step}.")
+
+ if FLAGS.show_q_values:
+ q_queue = queue.Queue()
+ q_display = ImageDisplayer(q_queue, "q_display")
+ q_display.start()
+ for episode in range(FLAGS.eval_n_trajs):
+ print("reset start")
+ ### receive signal from learner and then reset
+ obs, _ = env.reset()
+ print("reset end")
+ done = False
+ start_time = time.time()
+ while not done:
+ rng, key = jax.random.split(rng)
+ actions = agent.sample_actions(
+ observations=jax.device_put(obs),
+ argmax=True,
+ seed=key
+ )
+ if FLAGS.show_q_values:
+ q_value = np.asarray(jax.device_get(agent.forward_critic(jax.device_put(obs), actions[:-1], rng=key).min()))
+ q_value_grasp = np.asarray(jax.device_get(agent.forward_grasp_critic(jax.device_put(obs), rng=key)))
+ q_value_grasp_index = int(actions[-1] + 1)
+ actions = np.asarray(jax.device_get(actions))
+
+ next_obs, reward, done, truncated, info = env.step(actions)
+ obs = next_obs
+
+ if 'intervene_action' in info:
+ print(info['intervene_action'][-1])
+
+ if FLAGS.show_q_values:
+ q_queue.put({'q_image': q_image(q_value, q_value_grasp, q_value_grasp_index, 'intervene_action' in info)})
+
+ if done or truncated:
+ if reward:
+ dt = time.time() - start_time
+ time_list.append(dt)
+ print(dt)
+
+ success_counter += reward
+ print(reward)
+ print(f"{success_counter}/{episode + 1}")
+
+ print(f"success rate: {success_counter / FLAGS.eval_n_trajs}")
+ print(f"average time: {np.mean(time_list)}")
+
+ if FLAGS.show_q_values:
+ q_queue.put(None)
+ cv2.destroyAllWindows()
+ q_display.join()
+
+ with open(FLAGS.save_to_txt, "a") as f:
+ f.write(f"{FLAGS.eval_checkpoint_step} - {success_counter / FLAGS.eval_n_trajs} ({FLAGS.eval_n_trajs})\n")
+
+if __name__ == "__main__":
+ app.run(main)
diff --git a/examples/experiments/cube_reach/config.py b/examples/experiments/cube_reach/config.py
new file mode 100644
index 00000000..3ae90be2
--- /dev/null
+++ b/examples/experiments/cube_reach/config.py
@@ -0,0 +1,167 @@
+import os
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from franka_env.envs.wrappers import (
+ Quat2EulerWrapper,
+ SpacemouseIntervention,
+ MultiCameraBinaryRewardClassifierWrapper,
+ GripperCloseEnv
+)
+from franka_env.envs.relative_env import RelativeFrame
+from franka_env.envs.franka_env import DefaultEnvConfig
+from serl_launcher.wrappers.serl_obs_wrappers import SERLObsWrapper
+from serl_launcher.wrappers.chunking import ChunkingWrapper
+from serl_launcher.networks.reward_classifier import load_classifier_func
+
+from experiments.config import DefaultTrainingConfig
+from experiments.cube_reach.wrapper import CubeReachEnv
+from experiments.usb_pickup_insertion.wrapper import GripperPenaltyWrapper
+
+class EnvConfig(DefaultEnvConfig):
+ SERVER_URL = "http://localhost:5000/"
+ REALSENSE_CAMERAS = {
+ "wrist_1": {
+ "serial_number": "123622270810",
+ "dim": (1280, 720),
+ "exposure": 13000,
+ },
+ "side_1": {
+ "serial_number": "032522250211",
+ "dim": (1280, 720),
+ "exposure": 13000,
+ }
+ }
+ IMAGE_CROP = {
+ "wrist_1": lambda img: img[50:550, 150:1200],
+ # "wrist_2": lambda img: img[100:500, 400:900],
+ # "side_1": lambda img: img[440:700, 440:900],
+ "side_1": lambda img: img[320:600, 500:880],
+ }
+ # (440, 440) -> (900, 700)
+ # TARGET_POSE = np.array([0.5881241235410154,-0.03578590131997776,0.27843494179085326, np.pi, 0, 0])
+ # GRASP_POSE = np.array([0.5857508505445138,-0.22036261105675414,0.2731021902359492, np.pi, 0, 0])
+
+ # Step 1: add reset pose
+ # RESET_POSE = np.array([0.7, -0.05, 0.225,np.pi, 0, -np.pi/2])
+ RESET_POSE = np.array([
+ 0.6227042846188979,
+ 0.10877613017684733,
+ 0.15120951063136454,
+ 3.092543056432925,
+ 0.03848632982111533,
+ np.pi / 2,
+ ])
+
+ # Step 2: add bounding boxes
+ # upper right corner: {"pose":[0.8245290917409726,-0.27010709908619496,0.18881752346236647]}
+ # upper left corner: "pose":[0.6235723397920869,0.12971844992327827,0.09420440879204567]}
+
+ # ABS_POSE_LIMIT_LOW = np.array([0.6, -0.3, 0.09, np.pi-0.1, -0.1, -np.pi/2-0.3])
+ # ABS_POSE_LIMIT_HIGH = np.array([0.8, 0.13, 0.25, np.pi+0.1, 0.1, -np.pi/2+0.3])
+ # ABS_POSE_LIMIT_LOW = np.array([0.54, -0.285, 0.085, np.pi - 0.1, -0.1, np.pi / 4 ])
+ # ABS_POSE_LIMIT_HIGH = np.array([0.80, 0.085, 0.115, np.pi + 0.1, 0.1, np.pi / 4 + 0.2])
+
+ ABS_POSE_LIMIT_LOW = np.array([0.51, -0.09, 0.09, np.pi - 0.1, -0.1, np.pi / 2 - 0.1])
+ ABS_POSE_LIMIT_HIGH = np.array([0.73, 0.30, 0.16, np.pi + 0.1, 0.1, np.pi / 2 + 0.1])
+
+ # Step 3: set reset randomization
+ RANDOM_RESET = True
+ RANDOM_XY_RANGE = 0.05
+ RANDOM_RZ_RANGE = np.pi / 6
+ ACTION_SCALE = np.array([0.1, 0.3, 1]) # Testing x
+ DISPLAY_IMAGE = True
+ MAX_EPISODE_LENGTH = 400
+
+
+ ## Step 4: Copy the same control parameters -- keep it the same as below : TODO: figure out how to tune these+
+ COMPLIANCE_PARAM = {
+ "translational_stiffness": 2000,
+ "translational_damping": 89,
+ "rotational_stiffness": 150,
+ "rotational_damping": 7,
+ "translational_Ki": 0,
+ "translational_clip_x": 0.005,
+ "translational_clip_y": 0.005,
+ "translational_clip_z": 0.005,
+ "translational_clip_neg_x": 0.005,
+ "translational_clip_neg_y": 0.005,
+ "translational_clip_neg_z": 0.005,
+ "rotational_clip_x": 0.05,
+ "rotational_clip_y": 0.05,
+ "rotational_clip_z": 0.02,
+ "rotational_clip_neg_x": 0.05,
+ "rotational_clip_neg_y": 0.05,
+ "rotational_clip_neg_z": 0.02,
+ "rotational_Ki": 0,
+ }
+ PRECISION_PARAM = {
+ "translational_stiffness": 3000,
+ "translational_damping": 89,
+ "rotational_stiffness": 300,
+ "rotational_damping": 9,
+ "translational_Ki": 0.1,
+ "translational_clip_x": 0.01,
+ "translational_clip_y": 0.01,
+ "translational_clip_z": 0.01,
+ "translational_clip_neg_x": 0.01,
+ "translational_clip_neg_y": 0.01,
+ "translational_clip_neg_z": 0.01,
+ "rotational_clip_x": 0.05,
+ "rotational_clip_y": 0.05,
+ "rotational_clip_z": 0.05,
+ "rotational_clip_neg_x": 0.05,
+ "rotational_clip_neg_y": 0.05,
+ "rotational_clip_neg_z": 0.05,
+ "rotational_Ki": 0.1,
+ }
+
+
+class TrainConfig(DefaultTrainingConfig):
+ image_keys = ["wrist_1", "side_1"]
+ classifier_keys = ["wrist_1", "side_1"]
+ proprio_keys = ["tcp_pose", "tcp_vel", "tcp_force", "tcp_torque", "gripper_pose"]
+ # buffer_period = 1000
+ # checkpoint_period = 5000
+ # steps_per_update = 50
+ checkpoint_period = 2000
+ cta_ratio = 2
+ random_steps = 0
+ discount = 0.98
+ buffer_period = 1000
+ batch_size = 64
+ encoder_type = "resnet-pretrained"
+ setup_mode = "single-arm-learned-gripper"
+
+ def get_environment(self, fake_env=False, save_video=False, classifier=False):
+ env = CubeReachEnv(
+ fake_env=fake_env,
+ save_video=save_video,
+ config=EnvConfig(),
+ )
+ if not fake_env:
+ env = SpacemouseIntervention(env)
+ env = RelativeFrame(env)
+ env = Quat2EulerWrapper(env)
+ env = SERLObsWrapper(env, proprio_keys=self.proprio_keys)
+ env = ChunkingWrapper(env, obs_horizon=1, act_exec_horizon=None)
+ if classifier:
+ classifier = load_classifier_func(
+ key=jax.random.PRNGKey(0),
+ sample=env.observation_space.sample(),
+ image_keys=self.classifier_keys,
+ checkpoint_path=os.path.abspath("classifier_ckpt/"),
+ )
+
+ def reward_func(obs):
+ sigmoid = lambda x: 1 / (1 + jnp.exp(-x))
+ # added check for z position to further robustify classifier, but should work without as well
+ prob = sigmoid(classifier(obs))[0]
+ print(prob)
+ return int(prob > 0.85) # and obs['state'][0, 3] < 0.12
+ # return 0.0
+
+ env = MultiCameraBinaryRewardClassifierWrapper(env, reward_func)
+ env = GripperPenaltyWrapper(env, penalty=-0.02)
+ return env
\ No newline at end of file
diff --git a/examples/experiments/cube_reach/run_actor.sh b/examples/experiments/cube_reach/run_actor.sh
new file mode 100755
index 00000000..6087415b
--- /dev/null
+++ b/examples/experiments/cube_reach/run_actor.sh
@@ -0,0 +1,6 @@
+export XLA_PYTHON_CLIENT_PREALLOCATE=false && \
+export XLA_PYTHON_CLIENT_MEM_FRACTION=.1 && \
+/home/jacob/miniforge3/envs/serl2/bin/python ../../train_rlpd.py "$@" \
+ --exp_name=cube_reach \
+ --checkpoint_path=../../experiments/cube_reach/debug \
+ --actor \
\ No newline at end of file
diff --git a/examples/experiments/cube_reach/run_learner.sh b/examples/experiments/cube_reach/run_learner.sh
new file mode 100755
index 00000000..3f467303
--- /dev/null
+++ b/examples/experiments/cube_reach/run_learner.sh
@@ -0,0 +1,8 @@
+export XLA_PYTHON_CLIENT_PREALLOCATE=false && \
+export XLA_PYTHON_CLIENT_MEM_FRACTION=.3 && \
+/home/jacob/miniforge3/envs/serl2/bin/python ../../train_rlpd.py "$@" \
+ --exp_name=cube_reach \
+ --checkpoint_path=../../experiments/cube_reach/debug \
+ --demo_path=/home/jacob/hil-serl/examples/experiments/cube_reach/demo_data/cube_reach_20_demos_2025-03-30_15-16-49.pkl \
+ --learner \
+ --save_video \
\ No newline at end of file
diff --git a/examples/experiments/cube_reach/wrapper.py b/examples/experiments/cube_reach/wrapper.py
new file mode 100644
index 00000000..e44c40ad
--- /dev/null
+++ b/examples/experiments/cube_reach/wrapper.py
@@ -0,0 +1,73 @@
+import copy
+import time
+from franka_env.utils.rotations import euler_2_quat
+from scipy.spatial.transform import Rotation as R
+import numpy as np
+import requests
+from pynput import keyboard
+
+from franka_env.envs.franka_env import FrankaEnv
+
+class CubeReachEnv(FrankaEnv):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ def go_to_reset(self, joint_reset=False):
+ """
+ Move to the rest position defined in base class.
+ Add a small z offset before going to rest to avoid collision with object.
+ """
+ # use compliance mode for coupled reset
+ self._update_currpos()
+ self._send_pos_command(self.currpos)
+ time.sleep(0.3)
+ requests.post(self.url + "update_param", json=self.config.PRECISION_PARAM)
+
+ # pull up
+ self._update_currpos()
+ reset_pose = copy.deepcopy(self.currpos)
+ reset_pose[2] = self.resetpos[2] + 0.04
+ self.interpolate_move(reset_pose, timeout=1)
+
+ # perform joint reset if needed
+ if joint_reset:
+ print("JOINT RESET")
+ requests.post(self.url + "jointreset")
+ time.sleep(0.5)
+
+ # perform Cartesian reset
+ if self.randomreset: # randomize reset position in xy plane
+ reset_pose = self.resetpos.copy()
+ reset_pose[:2] += np.random.uniform(
+ -self.random_xy_range, self.random_xy_range, (2,)
+ )
+ euler_random = self._RESET_POSE[3:].copy()
+ euler_random[-1] += np.random.uniform(
+ -self.random_rz_range, self.random_rz_range
+ )
+ reset_pose[3:] = euler_2_quat(euler_random)
+ self._send_pos_command(reset_pose)
+ else:
+ reset_pose = self.resetpos.copy()
+ self._send_pos_command(reset_pose)
+ time.sleep(0.5)
+
+ # Change to compliance mode
+ requests.post(self.url + "update_param", json=self.config.COMPLIANCE_PARAM)
+
+
+ def reset(self, joint_reset=False, **kwargs):
+ self.last_gripper_act = time.time()
+ if self.save_video:
+ self.save_video_recording()
+
+ self._recover()
+ self.go_to_reset(joint_reset=False)
+ self._recover()
+ self.curr_path_length = 0
+
+ self._update_currpos()
+ obs = self._get_obs()
+ requests.post(self.url + "update_param", json=self.config.COMPLIANCE_PARAM)
+ self.terminate = False
+ return obs, {}
\ No newline at end of file
diff --git a/examples/experiments/cube_reach2/config.py b/examples/experiments/cube_reach2/config.py
new file mode 100644
index 00000000..ffaf9d1b
--- /dev/null
+++ b/examples/experiments/cube_reach2/config.py
@@ -0,0 +1,159 @@
+import os
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from franka_env.envs.wrappers import (
+ Quat2EulerWrapper,
+ SpacemouseIntervention,
+ MultiCameraBinaryRewardClassifierWrapper,
+ GripperCloseEnv,
+ LastNGripperActionsWrapper,
+)
+from franka_env.envs.relative_env import RelativeFrame
+from franka_env.envs.franka_env import DefaultEnvConfig
+from serl_launcher.wrappers.serl_obs_wrappers import SERLObsWrapper
+from serl_launcher.wrappers.chunking import ChunkingWrapper
+from serl_launcher.networks.reward_classifier import load_classifier_func
+
+from experiments.config import DefaultTrainingConfig
+from experiments.cube_reach2.wrapper import CubeReach2Env
+from experiments.usb_pickup_insertion.wrapper import GripperPenaltyWrapper
+
+class EnvConfig(DefaultEnvConfig):
+ SERVER_URL = "http://localhost:5000/"
+ REALSENSE_CAMERAS = {
+ "wrist_1": {
+ "camera_type": "rs",
+ "serial_number": "123622270810",
+ "dim": (1280, 720),
+ "exposure": 13000,
+ },
+ "side_1": {
+ "camera_type": "zed",
+ "serial_number": "032522250211",
+ "dim": (1280, 720),
+ "exposure": 13000,
+ }
+ }
+ IMAGE_CROP = {
+ "wrist_1": lambda img: img[50:550, 150:1200],
+ # "wrist_2": lambda img: img[100:500, 400:900],
+ # "side_1": lambda img: img[440:700, 440:900],
+ "side_1": lambda img: img[320:700, 600:950],
+ }
+ # (440, 440) -> (900, 700)
+ # TARGET_POSE = np.array([0.5881241235410154,-0.03578590131997776,0.27843494179085326, np.pi, 0, 0])
+ # GRASP_POSE = np.array([0.5857508505445138,-0.22036261105675414,0.2731021902359492, np.pi, 0, 0])
+
+ # Step 1: add reset pose
+ # RESET_POSE = np.array([0.7, -0.05, 0.225,np.pi, 0, -np.pi/2])
+ RESET_POSE = np.array([0.62, 0.099, 0.15, np.pi, 0, np.pi / 2])
+
+ # Step 2: add bounding boxes
+ ABS_POSE_LIMIT_LOW = np.array([0.58, -0.056, 0.09, np.pi - 0.1, -0.1, np.pi / 2 - 0.1])
+ ABS_POSE_LIMIT_HIGH = np.array([0.66, 0.267, 0.18, np.pi + 0.1, 0.1, np.pi / 2 + 0.1])
+
+ # Step 3: set reset randomization
+ RANDOM_RESET = True
+ RANDOM_XY_RANGE = 0.05
+ RANDOM_RZ_RANGE = np.pi / 6
+ ACTION_SCALE = np.array([0.1, 0.3, 1]) # Testing x
+ DISPLAY_IMAGE = True
+ MAX_EPISODE_LENGTH = 300
+
+
+ ## Step 4: Copy the same control parameters -- keep it the same as below : TODO: figure out how to tune these+
+ COMPLIANCE_PARAM = {
+ "translational_stiffness": 2000,
+ "translational_damping": 89,
+ "rotational_stiffness": 150,
+ "rotational_damping": 7,
+ "translational_Ki": 0,
+ "translational_clip_x": 0.005,
+ "translational_clip_y": 0.005,
+ "translational_clip_z": 0.005,
+ "translational_clip_neg_x": 0.005,
+ "translational_clip_neg_y": 0.005,
+ "translational_clip_neg_z": 0.005,
+ "rotational_clip_x": 0.05,
+ "rotational_clip_y": 0.05,
+ "rotational_clip_z": 0.02,
+ "rotational_clip_neg_x": 0.05,
+ "rotational_clip_neg_y": 0.05,
+ "rotational_clip_neg_z": 0.02,
+ "rotational_Ki": 0,
+ }
+ PRECISION_PARAM = {
+ "translational_stiffness": 3000,
+ "translational_damping": 89,
+ "rotational_stiffness": 300,
+ "rotational_damping": 9,
+ "translational_Ki": 0.1,
+ "translational_clip_x": 0.01,
+ "translational_clip_y": 0.01,
+ "translational_clip_z": 0.01,
+ "translational_clip_neg_x": 0.01,
+ "translational_clip_neg_y": 0.01,
+ "translational_clip_neg_z": 0.01,
+ "rotational_clip_x": 0.05,
+ "rotational_clip_y": 0.05,
+ "rotational_clip_z": 0.05,
+ "rotational_clip_neg_x": 0.05,
+ "rotational_clip_neg_y": 0.05,
+ "rotational_clip_neg_z": 0.05,
+ "rotational_Ki": 0.1,
+ }
+
+
+class TrainConfig(DefaultTrainingConfig):
+ image_keys = ["wrist_1", "side_1"]
+ classifier_keys = ["wrist_1", "side_1"]
+ proprio_keys = ["tcp_pose", "tcp_vel", "tcp_force", "tcp_torque", "gripper_pose"]
+ # buffer_period = 1000
+ # checkpoint_period = 5000
+ # steps_per_update = 50
+ pretraining_steps = 30000 # How many steps to pre-train the model for using RLPD on offline data only.
+ reward_scale = 10 # How much to scale actual rewards (not RLIF penalties) for RLIF.
+ checkpoint_period = 4000
+ cta_ratio = 2
+ random_steps = 0
+ discount = 0.98
+ buffer_period = 1000
+ batch_size = 64
+ encoder_type = "resnet-pretrained"
+ setup_mode = "single-arm-learned-gripper"
+
+ def get_environment(self, fake_env=False, save_video=False, classifier=False):
+ env = CubeReach2Env(
+ fake_env=fake_env,
+ save_video=save_video,
+ config=EnvConfig(),
+ )
+ if not fake_env:
+ env = SpacemouseIntervention(env)
+ env = RelativeFrame(env)
+ env = Quat2EulerWrapper(env)
+ env = SERLObsWrapper(env, proprio_keys=self.proprio_keys)
+ env = LastNGripperActionsWrapper(env, 8)
+ env = ChunkingWrapper(env, obs_horizon=1, act_exec_horizon=None)
+ if classifier:
+ classifier = load_classifier_func(
+ key=jax.random.PRNGKey(0),
+ sample=env.observation_space.sample(),
+ image_keys=self.classifier_keys,
+ checkpoint_path=os.path.abspath("classifier_ckpt/"),
+ )
+
+ def reward_func(obs):
+ sigmoid = lambda x: 1 / (1 + jnp.exp(-x))
+ # added check for z position to further robustify classifier, but should work without as well
+ prob = sigmoid(classifier(obs[:len(obs) - 8]))[0]
+ if prob > 0.1:
+ print(f"{prob:.2f}")
+ return int(prob > 0.90) # and obs['state'][0, 3] < 0.12
+ # return 0.0
+
+ env = MultiCameraBinaryRewardClassifierWrapper(env, reward_func)
+ env = GripperPenaltyWrapper(env, penalty=-0.02)
+ return env
\ No newline at end of file
diff --git a/examples/experiments/cube_reach2/eval_checkpoint.sh b/examples/experiments/cube_reach2/eval_checkpoint.sh
new file mode 100644
index 00000000..0bebbce0
--- /dev/null
+++ b/examples/experiments/cube_reach2/eval_checkpoint.sh
@@ -0,0 +1,8 @@
+export XLA_PYTHON_CLIENT_PREALLOCATE=false && \
+export XLA_PYTHON_CLIENT_MEM_FRACTION=.1 && \
+/home/jacob/miniforge3/envs/serl2/bin/python ../../train_rlif.py "$@" \
+ --exp_name=cube_reach2 \
+ --checkpoint_path=../../experiments/cube_reach2/debug_rlif \
+ --actor \
+ --eval_checkpoint_step=100000 \
+ --eval_n_trajs=10 \
diff --git a/examples/experiments/cube_reach2/run_actor.sh b/examples/experiments/cube_reach2/run_actor.sh
new file mode 100755
index 00000000..1d4c6240
--- /dev/null
+++ b/examples/experiments/cube_reach2/run_actor.sh
@@ -0,0 +1,6 @@
+export XLA_PYTHON_CLIENT_PREALLOCATE=false && \
+export XLA_PYTHON_CLIENT_MEM_FRACTION=.1 && \
+/home/jacob/miniforge3/envs/serl2/bin/python ../../train_rlif.py "$@" \
+ --exp_name=cube_reach2 \
+ --checkpoint_path=../../experiments/cube_reach2/debug_rlif \
+ --actor \
\ No newline at end of file
diff --git a/examples/experiments/cube_reach2/run_learner.sh b/examples/experiments/cube_reach2/run_learner.sh
new file mode 100755
index 00000000..cc296374
--- /dev/null
+++ b/examples/experiments/cube_reach2/run_learner.sh
@@ -0,0 +1,8 @@
+export XLA_PYTHON_CLIENT_PREALLOCATE=false && \
+export XLA_PYTHON_CLIENT_MEM_FRACTION=.3 && \
+/home/jacob/miniforge3/envs/serl2/bin/python ../../train_rlif.py "$@" \
+ --exp_name=cube_reach2 \
+ --checkpoint_path=../../experiments/cube_reach2/debug_rlif \
+ --demo_path=/home/jacob/Desktop/rico/jax-hitl-hil-serl/examples/experiments/cube_reach2/demo_data/cube_reach2_20_demos_2025-04-04_10-06-29.pkl \
+ --learner \
+ --save_video \
\ No newline at end of file
diff --git a/examples/experiments/cube_reach2/wrapper.py b/examples/experiments/cube_reach2/wrapper.py
new file mode 100644
index 00000000..7fab5cea
--- /dev/null
+++ b/examples/experiments/cube_reach2/wrapper.py
@@ -0,0 +1,73 @@
+import copy
+import time
+from franka_env.utils.rotations import euler_2_quat
+from scipy.spatial.transform import Rotation as R
+import numpy as np
+import requests
+from pynput import keyboard
+
+from franka_env.envs.franka_env import FrankaEnv
+
+class CubeReach2Env(FrankaEnv):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ def go_to_reset(self, joint_reset=False):
+ """
+ Move to the rest position defined in base class.
+ Add a small z offset before going to rest to avoid collision with object.
+ """
+ # use compliance mode for coupled reset
+ self._update_currpos()
+ self._send_pos_command(self.currpos)
+ time.sleep(0.3)
+ requests.post(self.url + "update_param", json=self.config.PRECISION_PARAM)
+
+ # pull up
+ self._update_currpos()
+ reset_pose = copy.deepcopy(self.currpos)
+ reset_pose[2] = self.resetpos[2] + 0.04
+ self.interpolate_move(reset_pose, timeout=1)
+
+ # perform joint reset if needed
+ if joint_reset:
+ print("JOINT RESET")
+ requests.post(self.url + "jointreset")
+ time.sleep(0.5)
+
+ # perform Cartesian reset
+ if self.randomreset: # randomize reset position in xy plane
+ reset_pose = self.resetpos.copy()
+ reset_pose[:2] += np.random.uniform(
+ -self.random_xy_range, self.random_xy_range, (2,)
+ )
+ euler_random = self._RESET_POSE[3:].copy()
+ euler_random[-1] += np.random.uniform(
+ -self.random_rz_range, self.random_rz_range
+ )
+ reset_pose[3:] = euler_2_quat(euler_random)
+ self._send_pos_command(reset_pose)
+ else:
+ reset_pose = self.resetpos.copy()
+ self._send_pos_command(reset_pose)
+ time.sleep(0.5)
+
+ # Change to compliance mode
+ requests.post(self.url + "update_param", json=self.config.COMPLIANCE_PARAM)
+
+
+ def reset(self, joint_reset=False, **kwargs):
+ self.last_gripper_act = time.time()
+ if self.save_video:
+ self.save_video_recording()
+
+ self._recover()
+ self.go_to_reset(joint_reset=False)
+ self._recover()
+ self.curr_path_length = 0
+
+ self._update_currpos()
+ obs = self._get_obs()
+ requests.post(self.url + "update_param", json=self.config.COMPLIANCE_PARAM)
+ self.terminate = False
+ return obs, {}
\ No newline at end of file
diff --git a/examples/experiments/cube_reach3/config.py b/examples/experiments/cube_reach3/config.py
new file mode 100644
index 00000000..17bb753e
--- /dev/null
+++ b/examples/experiments/cube_reach3/config.py
@@ -0,0 +1,156 @@
+import os
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from franka_env.envs.wrappers import (
+ Quat2EulerWrapper,
+ SpacemouseIntervention,
+ MultiCameraBinaryRewardClassifierWrapper,
+ GripperCloseEnv,
+ LastNGripperActionsWrapper,
+)
+from franka_env.envs.relative_env import RelativeFrame
+from franka_env.envs.franka_env import DefaultEnvConfig
+from serl_launcher.wrappers.serl_obs_wrappers import SERLObsWrapper
+from serl_launcher.wrappers.chunking import ChunkingWrapper
+from serl_launcher.networks.reward_classifier import load_classifier_func
+
+from experiments.config import DefaultTrainingConfig
+from experiments.cube_reach3.wrapper import CubeReach3Env
+from experiments.usb_pickup_insertion.wrapper import GripperPenaltyWrapper
+
+class EnvConfig(DefaultEnvConfig):
+ SERVER_URL = "http://127.0.0.1:5000/"
+ REALSENSE_CAMERAS = {
+ "side_1": {
+ "camera_type": "rs",
+ "serial_number": "215322079254",
+ "dim": (1280, 720),
+ "exposure": 13000,
+ },
+ # "wrist_1": {
+ # "camera_type": "zed",
+ # "serial_number": "16744838",
+ # "dim": (720, 1280), # height, width
+ # "exposure": 12000,
+ # }
+ }
+ IMAGE_CROP = {
+ # "wrist_1": lambda img: img[0:720, 450:1150], # did you know? you can check these values on realsense-viewer even for zed cameras.
+ "side_1": lambda img: img[260:700, 420:1270],
+ }
+
+ # Step 1: add reset pose
+ RESET_POSE = np.array([0.48, 0.04, 0.24, np.pi, 0, np.pi / 2])
+
+ # Step 2: add bounding boxes
+ ABS_POSE_LIMIT_LOW = np.array([0.42, -0.14, 0.20, np.pi - 0.05, -0.05, np.pi / 2 - 0.05])
+ ABS_POSE_LIMIT_HIGH = np.array([0.51, 0.24, 0.29, np.pi + 0.05, 0.05, np.pi / 2 + 0.05])
+
+ # Step 3: set reset randomization
+ RANDOM_RESET = True
+ RANDOM_XY_RANGE = 0.05
+ RANDOM_RZ_RANGE = np.pi / 6
+ ACTION_SCALE = np.array([0.05, 0.3, 1]) # Testing x
+ DISPLAY_IMAGE = True
+ MAX_EPISODE_LENGTH = 400
+
+
+ ## Step 4: Copy the same control parameters -- keep it the same as below : TODO: figure out how to tune these+
+ COMPLIANCE_PARAM = {
+ "translational_stiffness": 2000,
+ "translational_damping": 89,
+ "rotational_stiffness": 150,
+ "rotational_damping": 7,
+ "translational_Ki": 0,
+ "translational_clip_x": 0.005,
+ "translational_clip_y": 0.005,
+ "translational_clip_z": 0.005,
+ "translational_clip_neg_x": 0.005,
+ "translational_clip_neg_y": 0.005,
+ "translational_clip_neg_z": 0.005,
+ "rotational_clip_x": 0.05,
+ "rotational_clip_y": 0.05,
+ "rotational_clip_z": 0.02,
+ "rotational_clip_neg_x": 0.05,
+ "rotational_clip_neg_y": 0.05,
+ "rotational_clip_neg_z": 0.02,
+ "rotational_Ki": 0,
+ }
+ PRECISION_PARAM = {
+ "translational_stiffness": 3000,
+ "translational_damping": 89,
+ "rotational_stiffness": 300,
+ "rotational_damping": 9,
+ "translational_Ki": 0.1,
+ "translational_clip_x": 0.01,
+ "translational_clip_y": 0.01,
+ "translational_clip_z": 0.01,
+ "translational_clip_neg_x": 0.01,
+ "translational_clip_neg_y": 0.01,
+ "translational_clip_neg_z": 0.01,
+ "rotational_clip_x": 0.05,
+ "rotational_clip_y": 0.05,
+ "rotational_clip_z": 0.05,
+ "rotational_clip_neg_x": 0.05,
+ "rotational_clip_neg_y": 0.05,
+ "rotational_clip_neg_z": 0.05,
+ "rotational_Ki": 0.1,
+ }
+
+
+class TrainConfig(DefaultTrainingConfig):
+ image_keys = ["side_1"]
+ classifier_keys = ["side_1"]
+ proprio_keys = ["tcp_pose", "tcp_vel", "tcp_force", "tcp_torque", "gripper_pose"]
+ # buffer_period = 1000
+ # checkpoint_period = 5000
+ # steps_per_update = 50
+ pretraining_steps = 0 # How many steps to pre-train the model for using RLPD on offline data only.
+ reward_scale = 1 # How much to scale actual rewards (not RLIF penalties) for RLIF.
+ rlif_minus_one = False
+ checkpoint_period = 4000
+ cta_ratio = 2
+ random_steps = 0
+ discount = 0.98
+ buffer_period = 1000
+ batch_size = 64
+ encoder_type = "resnet-pretrained"
+ setup_mode = "single-arm-learned-gripper"
+
+ def get_environment(self, fake_env=False, save_video=False, classifier=False):
+ env_config = EnvConfig()
+ assert np.logical_and(env_config.ABS_POSE_LIMIT_LOW <= env_config.RESET_POSE, env_config.RESET_POSE <= env_config.ABS_POSE_LIMIT_HIGH).all()
+ env = CubeReach3Env(
+ fake_env=fake_env,
+ save_video=save_video,
+ config=env_config,
+ )
+ if not fake_env:
+ env = SpacemouseIntervention(env)
+ env = RelativeFrame(env)
+ env = Quat2EulerWrapper(env)
+ env = SERLObsWrapper(env, proprio_keys=self.proprio_keys)
+ # env = LastNGripperActionsWrapper(env, 8)
+ env = ChunkingWrapper(env, obs_horizon=1, act_exec_horizon=None)
+ if classifier:
+ classifier = load_classifier_func(
+ key=jax.random.PRNGKey(0),
+ sample=env.observation_space.sample(),
+ image_keys=self.classifier_keys,
+ checkpoint_path=os.path.abspath("classifier_ckpt/"),
+ )
+
+ def reward_func(obs):
+ sigmoid = lambda x: 1 / (1 + jnp.exp(-x))
+ # added check for z position to further robustify classifier, but should work without as well
+ prob = sigmoid(classifier(obs))[0]
+ if prob > 0.1:
+ print(f"{prob:.2f}")
+ return int(prob > 0.90) # and obs['state'][0, 3] < 0.12
+ # return 0.0
+
+ env = MultiCameraBinaryRewardClassifierWrapper(env, reward_func)
+ env = GripperPenaltyWrapper(env, penalty=-0.02)
+ return env
diff --git a/examples/experiments/cube_reach3/eval_checkpoint.sh b/examples/experiments/cube_reach3/eval_checkpoint.sh
new file mode 100644
index 00000000..414c41cd
--- /dev/null
+++ b/examples/experiments/cube_reach3/eval_checkpoint.sh
@@ -0,0 +1,8 @@
+export XLA_PYTHON_CLIENT_PREALLOCATE=false && \
+export XLA_PYTHON_CLIENT_MEM_FRACTION=.1 && \
+sudo /home/qirico/miniconda3/envs/hilserl3/bin/python ../../../train_rlif.py "$@" \
+ --exp_name=cube_reach3 \
+ --checkpoint_path=debug_cl_2 \
+ --actor \
+ --eval_checkpoint_step=32000 \
+ --eval_n_trajs=20 \
diff --git a/examples/experiments/cube_reach3/run_actor.sh b/examples/experiments/cube_reach3/run_actor.sh
new file mode 100755
index 00000000..bb8659c1
--- /dev/null
+++ b/examples/experiments/cube_reach3/run_actor.sh
@@ -0,0 +1,6 @@
+export XLA_PYTHON_CLIENT_PREALLOCATE=false && \
+export XLA_PYTHON_CLIENT_MEM_FRACTION=.1 && \
+sudo /home/qirico/miniconda3/envs/hilserl/bin/python ../../../train_rlif.py "$@" \
+ --exp_name=cube_reach3 \
+ --checkpoint_path=debug_rlif \
+ --actor \
diff --git a/examples/experiments/cube_reach3/run_learner.sh b/examples/experiments/cube_reach3/run_learner.sh
new file mode 100755
index 00000000..f1228648
--- /dev/null
+++ b/examples/experiments/cube_reach3/run_learner.sh
@@ -0,0 +1,8 @@
+export XLA_PYTHON_CLIENT_PREALLOCATE=false && \
+export XLA_PYTHON_CLIENT_MEM_FRACTION=.3 && \
+sudo /home/qirico/miniconda3/envs/hilserl/bin/python ../../../train_rlif.py "$@" \
+ --exp_name=cube_reach3 \
+ --checkpoint_path=debug_rlif \
+ --demo_path=/home/qirico/Desktop/All-Weird/Human-Interventions/jax-hitl-hil-serl/examples/experiments/cube_reach3/side_only/demo_data/cube_reach3_20_demos_2025-04-16_12-07-35.pkl \
+ --learner \
+ --save_video \
diff --git a/examples/experiments/cube_reach3/side_only/c.sh b/examples/experiments/cube_reach3/side_only/c.sh
new file mode 100644
index 00000000..955391c4
--- /dev/null
+++ b/examples/experiments/cube_reach3/side_only/c.sh
@@ -0,0 +1,14 @@
+# sudo /home/qirico/miniconda3/envs/hilserl3/bin/python ../../../record_success_fail.py --exp_name=cube_reach3
+# sudo /home/qirico/miniconda3/envs/hilserl3/bin/python ../../../train_reward_classifier.py --exp_name=cube_reach3
+sudo /home/qirico/miniconda3/envs/hilserl3/bin/python ../../../record_demos.py --exp_name=cube_reach3 --successes_needed=20
+
+# sudo /home/qirico/miniconda3/envs/hilserl/bin/python ../../../train_bc.py \
+# --exp_name=cube_reach3 \
+# --bc_checkpoint_path=debug_bc \
+# --train_steps=100000
+
+# sudo /home/qirico/miniconda3/envs/hilserl/bin/python ../../../train_bc.py \
+# --exp_name=cube_reach3 \
+# --bc_checkpoint_path=/home/qirico/Desktop/All-Weird/Human-Interventions/jax-hitl-hil-serl/examples/experiments/cube_reach3/side_only/debug_bc/checkpoint_99990 \
+# --eval_n_trajs=20
+
diff --git a/examples/experiments/cube_reach3/side_only/run_cl_actor.sh b/examples/experiments/cube_reach3/side_only/run_cl_actor.sh
new file mode 100644
index 00000000..1e4767dd
--- /dev/null
+++ b/examples/experiments/cube_reach3/side_only/run_cl_actor.sh
@@ -0,0 +1,7 @@
+export XLA_PYTHON_CLIENT_PREALLOCATE=false && \
+export XLA_PYTHON_CLIENT_MEM_FRACTION=.1 && \
+sudo /home/qirico/miniconda3/envs/hilserl3/bin/python ../../../train_rlif.py "$@" \
+ --exp_name=cube_reach3 \
+ --checkpoint_path=debug_rlif_2 \
+ --actor \
+ --method=rlif \
diff --git a/examples/experiments/cube_reach3/side_only/run_cl_learner.sh b/examples/experiments/cube_reach3/side_only/run_cl_learner.sh
new file mode 100644
index 00000000..9c675cdc
--- /dev/null
+++ b/examples/experiments/cube_reach3/side_only/run_cl_learner.sh
@@ -0,0 +1,9 @@
+export XLA_PYTHON_CLIENT_PREALLOCATE=false && \
+export XLA_PYTHON_CLIENT_MEM_FRACTION=.3 && \
+sudo /home/qirico/miniconda3/envs/hilserl3/bin/python ../../../train_rlif.py "$@" \
+ --exp_name=cube_reach3 \
+ --checkpoint_path=debug_rlif_2 \
+ --demo_path=/home/qirico/Desktop/All-Weird/Human-Interventions/jax-hitl-hil-serl/examples/experiments/cube_reach3/side_only/demo_data/cube_reach3_20_demos_2025-04-22_20-19-06.pkl \
+ --learner \
+ --save_video \
+ --method=rlif \
diff --git a/examples/experiments/cube_reach3/wrapper.py b/examples/experiments/cube_reach3/wrapper.py
new file mode 100644
index 00000000..1a80e1e8
--- /dev/null
+++ b/examples/experiments/cube_reach3/wrapper.py
@@ -0,0 +1,73 @@
+import copy
+import time
+from franka_env.utils.rotations import euler_2_quat
+from scipy.spatial.transform import Rotation as R
+import numpy as np
+import requests
+from pynput import keyboard
+
+from franka_env.envs.franka_env import FrankaEnv
+
+class CubeReach3Env(FrankaEnv):
+ def __init__(self, **kwargs):
+ super().__init__(**kwargs)
+
+ def go_to_reset(self, joint_reset=False):
+ """
+ Move to the rest position defined in base class.
+ Add a small z offset before going to rest to avoid collision with object.
+ """
+ # use compliance mode for coupled reset
+ self._update_currpos()
+ self._send_pos_command(self.currpos)
+ time.sleep(0.3)
+ requests.post(self.url + "update_param", json=self.config.PRECISION_PARAM)
+
+ # pull up
+ self._update_currpos()
+ reset_pose = copy.deepcopy(self.currpos)
+ reset_pose[2] = self.resetpos[2] + 0.04
+ self.interpolate_move(reset_pose, timeout=1)
+
+ # perform joint reset if needed
+ if joint_reset:
+ print("JOINT RESET")
+ requests.post(self.url + "jointreset")
+ time.sleep(0.5)
+
+ # perform Cartesian reset
+ if self.randomreset: # randomize reset position in xy plane
+ reset_pose = self.resetpos.copy()
+ reset_pose[:2] += np.random.uniform(
+ -self.random_xy_range, self.random_xy_range, (2,)
+ )
+ euler_random = self._RESET_POSE[3:].copy()
+ euler_random[-1] += np.random.uniform(
+ -self.random_rz_range, self.random_rz_range
+ )
+ reset_pose[3:] = euler_2_quat(euler_random)
+ self._send_pos_command(reset_pose)
+ else:
+ reset_pose = self.resetpos.copy()
+ self._send_pos_command(reset_pose)
+ time.sleep(0.5)
+
+ # Change to compliance mode
+ requests.post(self.url + "update_param", json=self.config.COMPLIANCE_PARAM)
+
+
+ def reset(self, joint_reset=False, **kwargs):
+ self.last_gripper_act = time.time()
+ if self.save_video:
+ self.save_video_recording()
+
+ self._recover()
+ self.go_to_reset(joint_reset=False)
+ self._recover()
+ self.curr_path_length = 0
+
+ self._update_currpos()
+ obs = self._get_obs()
+ requests.post(self.url + "update_param", json=self.config.COMPLIANCE_PARAM)
+ self.terminate = False
+ return obs, {}
\ No newline at end of file
diff --git a/examples/experiments/franka_sim/config.py b/examples/experiments/franka_sim/config.py
new file mode 100644
index 00000000..81654e62
--- /dev/null
+++ b/examples/experiments/franka_sim/config.py
@@ -0,0 +1,65 @@
+import os
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from franka_env.envs.wrappers import (
+ SimSpacemouseIntervention,
+ Quat2EulerWrapper,
+ Useless4To7Wrapper,
+ BruhImageResizingWrapper,
+)
+from serl_launcher.wrappers.serl_obs_wrappers import SERLObsWrapper, AWrapperThatFlattensState
+
+from experiments.config import DefaultTrainingConfig
+from experiments.franka_sim.wrapper import FrankaSimEnv
+from serl_launcher.wrappers.chunking import ChunkingWrapper
+from experiments.usb_pickup_insertion.wrapper import GripperPenaltyWrapper
+from franka_sim.mujoco_gym_env import GymRenderingSpec
+
+
+class TrainConfig(DefaultTrainingConfig):
+ image_keys = ["front", "wrist"]
+ classifier_keys = ["front", "wrist"]
+ proprio_keys = ['panda/tcp_pos', 'panda/tcp_vel', 'panda/gripper_pos']
+ # buffer_period = 1000
+ # checkpoint_period = 5000
+ # steps_per_update = 50
+ pretraining_steps = 0 # How many steps to pre-train the model for using RLPD on offline data only.
+ reward_scale = 10 # How much to scale actual rewards (not RLIF penalties) for RLIF.
+ rlif_minus_one = False
+ checkpoint_period = 1000
+ cta_ratio = 5
+ random_steps = 0
+ discount = 0.98
+ buffer_period = 1000
+ batch_size = 64
+ encoder_type = "resnet-pretrained"
+ setup_mode = "single-arm-learned-gripper"
+
+ def get_environment(self, fake_env: bool, save_video: bool, classifier: bool, state_based: bool = False):
+ if state_based:
+ self.image_keys = [] # hacky but works
+ from franka_sim.envs import PandaPickCubeGymEnv
+ env = PandaPickCubeGymEnv(
+ action_scale=(0.1, 1),
+ render_mode="rgb_array" if fake_env else "human",
+ image_obs=not state_based,
+ control_dt=0.1,
+ time_limit=30.0,
+ show_viewer=not fake_env,
+ render_spec=GymRenderingSpec(height=512, width=512),
+ fixed_block_position=False,
+ )
+ env = FrankaSimEnv(env, action_scale=[0.1, 0.1, 0.1], show_viewer=False, fake_env=fake_env)
+ if not state_based:
+ env = SERLObsWrapper(env, proprio_keys=self.proprio_keys)
+ env = SimSpacemouseIntervention(env, fake_env=fake_env)
+ env = ChunkingWrapper(env, obs_horizon=1, act_exec_horizon=None)
+ env = Useless4To7Wrapper(env)
+ env = GripperPenaltyWrapper(env, penalty=-0.02)
+ env = BruhImageResizingWrapper(env, (128, 128))
+ if state_based:
+ env = AWrapperThatFlattensState(env)
+
+ return env
diff --git a/examples/experiments/franka_sim/eval_checkpoint.sh b/examples/experiments/franka_sim/eval_checkpoint.sh
new file mode 100644
index 00000000..4ce25206
--- /dev/null
+++ b/examples/experiments/franka_sim/eval_checkpoint.sh
@@ -0,0 +1,19 @@
+export EVAL_N_TRAJS=20
+export CHECKPOINT_PATH=debug_soft_cl_cta5
+export ENV_NAME=franka_sim
+
+declare -A status
+
+mkdir "${CHECKPOINT_PATH}_log"
+
+for (( i = 14000; i <= 14000; i += 1000)); do
+ XLA_PYTHON_CLIENT_MEM_FRACTION=.4 /home/qirico/miniconda3/envs/hilserl3/bin/python ../../eval_sim_checkpoints.py \
+ --exp_name=$ENV_NAME \
+ --checkpoint_path=$CHECKPOINT_PATH \
+ --eval_checkpoint_step=$i \
+ --eval_n_trajs=$EVAL_N_TRAJS \
+ --save_to_txt="${CHECKPOINT_PATH}.txt" \
+ >> "${CHECKPOINT_PATH}_log/${i}_2.out" 2>&1
+done
+
+echo "Everything done."
diff --git a/examples/experiments/franka_sim/run_actor.sh b/examples/experiments/franka_sim/run_actor.sh
new file mode 100755
index 00000000..089e2f87
--- /dev/null
+++ b/examples/experiments/franka_sim/run_actor.sh
@@ -0,0 +1,11 @@
+XLA_PYTHON_CLIENT_MEM_FRACTION=0.10 /home/qirico/miniconda3/envs/hilserl3/bin/python ../../train_rlif.py "$@" \
+ --exp_name=franka_sim \
+ --checkpoint_path=debug_soft_cl_cta5_utd1_nostateconstraint \
+ --actor \
+ --method=soft_cl \
+
+# XLA_PYTHON_CLIENT_MEM_FRACTION=0.10 /home/qirico/miniconda3/envs/hilserl3/bin/python ../../train_rlif.py "$@" \
+# --exp_name=franka_sim \
+# --checkpoint_path=debug_soft_cl_cta2_state_fixed_3 \
+# --actor \
+# --method=soft_cl \
diff --git a/examples/experiments/franka_sim/run_learner.sh b/examples/experiments/franka_sim/run_learner.sh
new file mode 100755
index 00000000..04cd15f9
--- /dev/null
+++ b/examples/experiments/franka_sim/run_learner.sh
@@ -0,0 +1,16 @@
+XLA_PYTHON_CLIENT_MEM_FRACTION=0.20 /home/qirico/miniconda3/envs/hilserl3/bin/python ../../train_rlif.py "$@" \
+ --exp_name=franka_sim \
+ --checkpoint_path=debug_soft_cl_cta5_utd1_nostateconstraint \
+ --demo_path=/home/qirico/Desktop/All-Weird/Human-Interventions/jax-hitl-hil-serl/examples/experiments/franka_sim/demo_data/franka_sim_20_demos_2025-04-25_21-16-29.pkl \
+ --demo_path=/home/qirico/Desktop/All-Weird/Human-Interventions/jax-hitl-hil-serl/examples/experiments/franka_sim/demo_data/franka_sim_30_demos_2025-04-27_21-13-05.pkl \
+ --learner \
+ --method=soft_cl \
+
+
+# XLA_PYTHON_CLIENT_MEM_FRACTION=0.30 /home/qirico/miniconda3/envs/hilserl3/bin/python ../../train_rlif.py "$@" \
+# --exp_name=franka_sim \
+# --checkpoint_path=debug_soft_cl_cta2_state_fixed_3 \
+# --demo_path=/home/qirico/Desktop/All-Weird/Human-Interventions/jax-hitl-hil-serl/examples/experiments/franka_sim/demo_data/franka_sim_20_fixed.pkl \
+# --learner \
+# --method=soft_cl \
+# --state_based \
diff --git a/examples/experiments/franka_sim/wrapper.py b/examples/experiments/franka_sim/wrapper.py
new file mode 100644
index 00000000..3c0ed04f
--- /dev/null
+++ b/examples/experiments/franka_sim/wrapper.py
@@ -0,0 +1,49 @@
+import copy
+import time
+from franka_env.utils.rotations import euler_2_quat
+from scipy.spatial.transform import Rotation as R
+import numpy as np
+import requests
+from pynput import keyboard
+import gymnasium as gym
+from franka_sim.mujoco_gym_env import MujocoGymEnv
+
+import mujoco
+import mujoco.viewer
+
+
+class FrankaSimEnv(gym.Wrapper):
+ show_viewer: bool
+
+ def __init__(
+ self,
+ env,
+ action_scale: list,
+ show_viewer: bool = True,
+ fake_env: bool = False,
+ ):
+ assert isinstance(env, MujocoGymEnv)
+ super().__init__(env)
+ self.fake_env = fake_env
+ m = env.model
+ d = env.data
+ self.action_scale = np.array(action_scale)
+ self.show_viewer = show_viewer
+ if self.show_viewer:
+ self.viewer = mujoco.viewer.launch_passive(m, d)
+
+ def reset(self, *args, **kwargs):
+ return super().reset(*args, **kwargs)
+
+
+ def step(self, action):
+ assert action.shape == (4,)
+ action[:3] *= self.action_scale
+ step_start = time.time()
+ vals = super().step(action)
+ if self.show_viewer:
+ self.viewer.sync()
+ time_until_next_step = self.env.control_dt - (time.time() - step_start)
+ if time_until_next_step > 0 and not self.fake_env:
+ time.sleep(time_until_next_step)
+ return vals
\ No newline at end of file
diff --git a/examples/experiments/mappings.py b/examples/experiments/mappings.py
index 8bdabd09..8d44cc67 100644
--- a/examples/experiments/mappings.py
+++ b/examples/experiments/mappings.py
@@ -2,10 +2,18 @@
from experiments.usb_pickup_insertion.config import TrainConfig as USBPickupInsertionTrainConfig
from experiments.object_handover.config import TrainConfig as ObjectHandoverTrainConfig
from experiments.egg_flip.config import TrainConfig as EggFlipTrainConfig
+from experiments.cube_reach.config import TrainConfig as CubeReachTrainConfig
+from experiments.cube_reach2.config import TrainConfig as CubeReach2TrainConfig
+from experiments.cube_reach3.config import TrainConfig as CubeReach3TrainConfig
+from experiments.franka_sim.config import TrainConfig as FrankaSimTrainConfig
CONFIG_MAPPING = {
"ram_insertion": RAMInsertionTrainConfig,
"usb_pickup_insertion": USBPickupInsertionTrainConfig,
"object_handover": ObjectHandoverTrainConfig,
"egg_flip": EggFlipTrainConfig,
+ "cube_reach": CubeReachTrainConfig,
+ "cube_reach2": CubeReach2TrainConfig,
+ "cube_reach3": CubeReach3TrainConfig,
+ "franka_sim": FrankaSimTrainConfig,
}
\ No newline at end of file
diff --git a/examples/experiments/object_handover/__pycache__/config.cpython-310.pyc b/examples/experiments/object_handover/__pycache__/config.cpython-310.pyc
deleted file mode 100644
index 7b38b449..00000000
Binary files a/examples/experiments/object_handover/__pycache__/config.cpython-310.pyc and /dev/null differ
diff --git a/examples/experiments/object_handover/__pycache__/wrapper.cpython-310.pyc b/examples/experiments/object_handover/__pycache__/wrapper.cpython-310.pyc
deleted file mode 100644
index 7064e7cd..00000000
Binary files a/examples/experiments/object_handover/__pycache__/wrapper.cpython-310.pyc and /dev/null differ
diff --git a/examples/experiments/ram_insertion/__pycache__/config.cpython-310.pyc b/examples/experiments/ram_insertion/__pycache__/config.cpython-310.pyc
deleted file mode 100644
index 70c78f29..00000000
Binary files a/examples/experiments/ram_insertion/__pycache__/config.cpython-310.pyc and /dev/null differ
diff --git a/examples/experiments/ram_insertion/__pycache__/wrapper.cpython-310.pyc b/examples/experiments/ram_insertion/__pycache__/wrapper.cpython-310.pyc
deleted file mode 100644
index 2249aaba..00000000
Binary files a/examples/experiments/ram_insertion/__pycache__/wrapper.cpython-310.pyc and /dev/null differ
diff --git a/examples/experiments/usb_pickup_insertion/__pycache__/config.cpython-310.pyc b/examples/experiments/usb_pickup_insertion/__pycache__/config.cpython-310.pyc
deleted file mode 100644
index 897a85ed..00000000
Binary files a/examples/experiments/usb_pickup_insertion/__pycache__/config.cpython-310.pyc and /dev/null differ
diff --git a/examples/experiments/usb_pickup_insertion/__pycache__/wrapper.cpython-310.pyc b/examples/experiments/usb_pickup_insertion/__pycache__/wrapper.cpython-310.pyc
deleted file mode 100644
index c010c027..00000000
Binary files a/examples/experiments/usb_pickup_insertion/__pycache__/wrapper.cpython-310.pyc and /dev/null differ
diff --git a/examples/experiments/usb_pickup_insertion/wrapper.py b/examples/experiments/usb_pickup_insertion/wrapper.py
index b28eb685..b1225357 100644
--- a/examples/experiments/usb_pickup_insertion/wrapper.py
+++ b/examples/experiments/usb_pickup_insertion/wrapper.py
@@ -102,7 +102,7 @@ def go_to_reset(self, joint_reset=False):
class GripperPenaltyWrapper(gym.Wrapper):
def __init__(self, env, penalty=-0.05):
super().__init__(env)
- assert env.action_space.shape == (7,)
+ assert env.action_space.shape == (7,), f"{env.action_space.shape}"
self.penalty = penalty
self.last_gripper_pos = None
diff --git a/examples/inspect_data.py b/examples/inspect_data.py
new file mode 100644
index 00000000..c4818a53
--- /dev/null
+++ b/examples/inspect_data.py
@@ -0,0 +1,215 @@
+from absl import app, flags
+import time
+import numpy as np
+import os
+import pickle
+import imageio
+import cv2
+
+FLAGS = flags.FLAGS
+
+
+def _get_folder():
+ print()
+ print("Please provide a folder that's either a /buffer/, /demo_buffer/, /preference_buffer/, or /interventions/.")
+ folder = input(" ")
+ print(folder)
+ return folder
+
+
+def _save_video(frames: np.ndarray, default_path: str):
+ print()
+ print("Output path: ")
+ output_path = input(f"[{default_path}] ")
+ if output_path == "":
+ output_path = f"{default_path}"
+ output_path = os.path.abspath(output_path)
+ assert os.path.exists(os.path.dirname(output_path))
+ if os.path.exists(output_path):
+ print()
+ yn = input("Will override previous video. Proceed? [y/n] [n] ")
+ if yn != 'y':
+ return
+ print(output_path)
+ print()
+ imageio.mimsave(
+ uri=output_path,
+ ims=frames,
+ fps=10,
+ macro_block_size=None # This can help with sizes not multiples of 16
+ )
+
+
+def inspect_buffer(path: str):
+ if os.path.isfile(path):
+ files = [path]
+ else:
+ files = [os.path.join(path, fp) for fp in os.listdir(path) if os.path.isfile(os.path.join(path, fp))]
+ transitions = []
+ for fp in files:
+ assert os.path.basename(fp).endswith(".pkl")
+ if not (os.path.basename(fp).endswith(".pkl")):
+ continue
+ with open(fp, "rb") as f:
+ obj = pickle.load(f)
+ assert isinstance(obj, list)
+ assert len(obj) > 0
+ assert isinstance(obj[0], dict)
+ assert set(obj[0].keys()).issuperset(["observations", "actions", "next_observations"])
+ transitions += obj
+
+ done_indices = [i for i, o in enumerate(transitions) if o['dones'] == 1]
+ print()
+ print("Pick a trajectory from the following end timesteps:")
+ print(done_indices + ["all"])
+ done_index = input(f"[{done_indices[0]}] ")
+ if done_index == "":
+ done_index = done_indices[0]
+ if done_index == "all":
+ done_index = len(transitions) - 1
+ start_index = 0
+ else:
+ done_index = int(done_index)
+ assert done_index in done_indices
+ start_index = done_indices[done_indices.index(done_index) - 1] + 1 if done_indices.index(done_index) - 1 >= 0 else 0
+ print(done_index)
+
+ camera_keys = list(transitions[start_index]["observations"].keys() - {"state"})
+ assert len(camera_keys) > 0
+ print()
+ print("Pick a camera from the following cameras:")
+ print(camera_keys)
+ camera_key = input(f"['{camera_keys[0]}'] ")
+ if camera_key == "":
+ camera_key = camera_keys[0]
+ assert camera_key in camera_keys
+ print(camera_key)
+
+ frames = []
+ for i, tr in enumerate(transitions[start_index:done_index+1]):
+ frame = tr["observations"][camera_key]
+ assert isinstance(frame, np.ndarray)
+ assert frame.dtype == np.uint8
+ assert frame.shape == (1, 128, 128, 3)
+ frame = frame[0]
+ new_frame = np.zeros((450, 800, 3), dtype=np.uint8)
+ new_frame[50:,:400,:] = cv2.resize(frame, (400, 400))
+ new_frame[50:,400:,:] = cv2.resize(tr["next_observations"][camera_key][0], (400, 400))
+
+ new_frame = cv2.putText(
+ img = new_frame,
+ text = f"t={i}",
+ org = (10, 40),
+ fontFace = cv2.FONT_HERSHEY_SIMPLEX,
+ fontScale = 1,
+ color = (255, 255, 255),
+ thickness = 2,
+ lineType = cv2.LINE_AA,
+ )
+
+ new_frame = cv2.putText(
+ img = new_frame,
+ text = f"r={tr['rewards']}",
+ org = (120, 40),
+ fontFace = cv2.FONT_HERSHEY_SIMPLEX,
+ fontScale = 1,
+ color = (255, 255, 255),
+ thickness = 2,
+ lineType = cv2.LINE_AA,
+ )
+
+ if "info" in tr and "intervene_action" in tr["info"]:
+ new_frame = cv2.putText(
+ img = new_frame,
+ text = f"intervene",
+ org = (220, 40),
+ fontFace = cv2.FONT_HERSHEY_SIMPLEX,
+ fontScale = 1,
+ color = (255, 255, 255),
+ thickness = 2,
+ lineType = cv2.LINE_AA,
+ )
+
+ frames.append(new_frame)
+ frames = np.stack(frames, axis=0)
+
+ _save_video(frames, "./video.mp4")
+
+
+def inspect_classifier_data(path: str):
+ assert os.path.exists(path), path
+ files = [os.path.join(path, fp) for fp in os.listdir(path) if os.path.isfile(os.path.join(path, fp))]
+ successes = []
+ failures = []
+ for fp in files:
+ assert os.path.basename(fp).endswith(".pkl")
+ if not os.path.basename(fp).endswith(".pkl"):
+ continue
+ with open(fp, "rb") as f:
+ obj = pickle.load(f)
+ assert isinstance(obj, list)
+ if len(obj) == 0:
+ continue
+ assert isinstance(obj[0], dict)
+ assert set(obj[0].keys()).issuperset(['observations', 'actions', 'next_observations', 'rewards', 'masks', 'dones'])
+ if "success" in os.path.basename(fp):
+ successes += obj
+ elif "failure" in os.path.basename(fp):
+ failures += obj
+ else:
+ raise Exception()
+
+ camera_keys = list(successes[0]["observations"].keys() - {"state"})
+ assert len(camera_keys) > 0
+ print()
+ print("Pick a camera from the following cameras:")
+ print(camera_keys)
+ camera_key = input(f"['{camera_keys[0]}'] ")
+ if camera_key == "":
+ camera_key = camera_keys[0]
+ assert camera_key in camera_keys
+ print(camera_key)
+
+ success_frames = []
+ for i, tr in enumerate(successes):
+ frame = tr["observations"][camera_key]
+ assert isinstance(frame, np.ndarray)
+ assert frame.dtype == np.uint8
+ assert frame.shape == (1, 128, 128, 3)
+ frame = frame[0]
+ new_frame = np.zeros((450, 400, 3), dtype=np.uint8)
+ new_frame[50:,:,:] = cv2.resize(frame, (400, 400))
+ success_frames.append(new_frame)
+ success_frames = np.stack(success_frames, axis=0)
+
+ failure_frames = []
+ for i, tr in enumerate(failures):
+ frame = tr["observations"][camera_key]
+ assert isinstance(frame, np.ndarray)
+ assert frame.dtype == np.uint8
+ assert frame.shape == (1, 128, 128, 3)
+ frame = frame[0]
+ new_frame = np.zeros((450, 400, 3), dtype=np.uint8)
+ new_frame[50:,:,:] = cv2.resize(frame, (400, 400))
+ failure_frames.append(new_frame)
+ failure_frames = np.stack(failure_frames, axis=0)
+
+ _save_video(success_frames, "./successes.mp4")
+ _save_video(failure_frames, "./failures.mp4")
+
+
+def main(_):
+ folder = _get_folder()
+ path = os.path.normpath(folder)
+ folder_name = os.path.basename(path) if os.path.isdir(path) else os.path.basename(os.path.dirname(path))
+ if folder_name == "buffer" or folder_name == "bc_buffer":
+ inspect_buffer(path)
+ elif folder_name == "classifier_data":
+ inspect_classifier_data(path)
+ elif folder_name == "demo_data":
+ inspect_buffer(path)
+ else:
+ raise Exception()
+
+if __name__ == "__main__":
+ app.run(main)
diff --git a/examples/inspect_demo_environment.py b/examples/inspect_demo_environment.py
new file mode 100644
index 00000000..114306e7
--- /dev/null
+++ b/examples/inspect_demo_environment.py
@@ -0,0 +1,178 @@
+# sudo /home/qirico/miniconda3/envs/hilserl3/bin/python examples/inspect_demo_environment.py
+
+from absl import app, flags
+import time
+import numpy as np
+import os
+import pickle
+import imageio
+import cv2
+import queue
+from pynput import keyboard
+import threading
+
+from experiments.ram_insertion.config import EnvConfig as RAMInsertionTrainConfig
+from experiments.usb_pickup_insertion.config import EnvConfig as USBPickupInsertionTrainConfig
+from experiments.cube_reach.config import EnvConfig as CubeReachTrainConfig
+from experiments.cube_reach2.config import EnvConfig as CubeReach2TrainConfig
+from experiments.cube_reach3.config import EnvConfig as CubeReach3TrainConfig
+
+ENV_CONFIG_MAPPING = {
+ "ram_insertion": RAMInsertionTrainConfig,
+ "usb_pickup_insertion": USBPickupInsertionTrainConfig,
+ "cube_reach": CubeReachTrainConfig,
+ "cube_reach2": CubeReach2TrainConfig,
+ "cube_reach3": CubeReach3TrainConfig,
+}
+
+from franka_env.camera.video_capture import VideoCapture
+from franka_env.camera.rs_capture import RSCapture
+from franka_env.camera.zed_capture import ZedCapture
+
+from run_demos import pull_demos
+
+class ImageDisplayer(threading.Thread):
+ def __init__(self, queue, name: str, keys: list[str], imgsize: int):
+ threading.Thread.__init__(self)
+ self.queue = queue
+ self.daemon = True # make this a daemon thread
+ self.name = name
+ self.keys = keys
+ self.imgsize = imgsize
+
+ def run(self):
+ while True:
+ img_array = self.queue.get() # retrieve an image from the queue
+ if img_array is None: # None is our signal to exit
+ break
+
+ xn = len(self.keys)
+ yn = len(self.keys[0])
+ sqsize = self.imgsize
+ margin = 5
+ cellxl = sqsize + margin * 2
+ cellyl = sqsize + margin * 2
+
+ frame = np.zeros((cellxl * xn, cellyl * yn, 3), dtype=np.uint8)
+ for x in range(xn):
+ for y in range(yn):
+ k = self.keys[x][y]
+ cur_frame = np.full((cellxl, cellyl, 3), 255, dtype=np.uint8)
+ cur_frame[margin : margin + sqsize, margin : margin + sqsize, :] = cv2.resize(img_array[k], (sqsize, sqsize))
+ frame[x * cellxl : (x + 1) * cellxl, y * cellyl : (y + 1) * cellyl, :] = cur_frame
+
+ cv2.imshow(self.name, frame)
+ cv2.waitKey(1)
+
+FLAGS = flags.FLAGS
+
+def print_green(x):
+ return print("\033[92m {}\033[00m".format(x))
+
+def print_yellow(x):
+ return print("\033[93m {}\033[00m".format(x))
+
+def print_cyan(x):
+ return print("\033[96m {}\033[00m".format(x))
+
+quit_key = False
+def on_press(key):
+ global quit_key
+ try:
+ if str(key) == "'q'":
+ print_yellow("quit")
+ quit_key = True
+ except AttributeError:
+ print("error")
+ pass
+
+def _get_exp_name(config_mapping: dict, default_exp_name: str = "cube_reach3"):
+ print()
+ print("Please provide a exp_name from the following:")
+ print(f"{list(config_mapping.keys())}")
+ exp_name = input(f"['{default_exp_name}'] ")
+ if exp_name == "":
+ exp_name = default_exp_name
+ print(exp_name)
+ return exp_name
+
+def _get_demo_file_name(default_demo_file_name: str = "/home/qirico/Desktop/All-Weird/Human-Interventions/jax-hitl-hil-serl/examples/experiments/cube_reach3/side_only/demo_data/cube_reach3_20_demos_2025-04-25_09-41-14.pkl"):
+ print()
+ print("Please provide a demo_file_name that ends with .pkl:")
+ demo_file_name = input(f"['{default_demo_file_name}'] ")
+ if demo_file_name == "":
+ demo_file_name = default_demo_file_name
+ assert demo_file_name.endswith(".pkl")
+ print(demo_file_name)
+
+ return demo_file_name
+
+def main(_):
+ global quit_key
+
+ listener = keyboard.Listener(
+ on_press=on_press)
+ listener.start()
+
+ exp_name = _get_exp_name(ENV_CONFIG_MAPPING)
+ config = ENV_CONFIG_MAPPING[exp_name]()
+
+ demo_file_name = _get_demo_file_name()
+ demos = pull_demos(demo_file_name)
+ assert len(demos) > 0
+
+ demo_images = demos[0]["observations"]
+ del demo_images['state']
+ assert set(demo_images.keys()) == set(config.REALSENSE_CAMERAS.keys())
+ for k in demo_images.keys():
+ if demo_images[k].ndim == 4 and demo_images[k].shape[0] == 1:
+ demo_images[k] = demo_images[k][0]
+
+ keys = [
+ ['side_1_cur', 'side_1_orig', 'side_1_avg', 'side_1_diff']
+ ]
+ displaysize = 400
+
+ img_queue = queue.Queue()
+ displayer = ImageDisplayer(img_queue, config.SERVER_URL, keys=keys, imgsize=displaysize)
+ displayer.start()
+
+ caps = {}
+ for cam_name, kwargs in config.REALSENSE_CAMERAS.items():
+ assert "camera_type" in kwargs
+ camera_type = kwargs['camera_type']
+ del kwargs['camera_type']
+ if camera_type == "rs":
+ cap = VideoCapture(
+ RSCapture(name=cam_name, **kwargs)
+ )
+ elif camera_type == "zed":
+ cap = VideoCapture(
+ ZedCapture(name=cam_name, **kwargs)
+ )
+ else:
+ raise NotImplementedError
+ caps[cam_name] = cap
+
+ while True:
+ if quit_key:
+ break
+ images = {}
+ for key in caps.keys():
+ cur_image = cap.read()
+ cur_image = config.IMAGE_CROP[key](cur_image) if key in config.IMAGE_CROP else cur_image
+ cur_image = cv2.resize(cur_image, demo_images[key].shape[:2])
+ demo_image = demo_images[key][:,:,::-1]
+ images[key + "_cur"] = cur_image
+ images[key + "_orig"] = demo_image
+ images[key + "_avg"] = np.uint8((np.float32(cur_image) + np.float32(demo_image)) / 2)
+ images[key + "_diff"] = np.uint8(np.abs(np.float32(cur_image) - np.float32(demo_image)) / 2)
+ img_queue.put(images)
+ time.sleep(0.1)
+
+ img_queue.put(None)
+ cv2.destroyAllWindows()
+ displayer.join()
+
+if __name__ == "__main__":
+ app.run(main)
diff --git a/examples/record_demos.py b/examples/record_demos.py
index 0467a85d..5a658cf0 100644
--- a/examples/record_demos.py
+++ b/examples/record_demos.py
@@ -1,3 +1,5 @@
+# How to use: sudo /home/jacob/miniforge3/envs/serl2/bin/python ../../record_demos.py --exp_name cube_reach2 --successes_needed 10
+
import os
from tqdm import tqdm
import numpy as np
@@ -16,8 +18,15 @@
def main(_):
assert FLAGS.exp_name in CONFIG_MAPPING, 'Experiment folder not found.'
config = CONFIG_MAPPING[FLAGS.exp_name]()
- env = config.get_environment(fake_env=False, save_video=False, classifier=True)
-
+ env = config.get_environment(fake_env=False, save_video=False, classifier=True, state_based=True)
+ if not os.path.exists("./demo_data"):
+ os.makedirs("./demo_data")
+ uuid = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+ file_name = f"./demo_data/{FLAGS.exp_name}_{FLAGS.successes_needed}_demos_{uuid}.pkl"
+ def dump_data(transitions):
+ with open(file_name, "wb") as f:
+ pkl.dump(transitions, f)
+ print(f"saved {success_count} demos to {file_name}")
obs, info = env.reset()
print("Reset done")
transitions = []
@@ -27,6 +36,8 @@ def main(_):
trajectory = []
returns = 0
+ cs = 0
+ t0 = time.time()
while success_count < success_needed:
actions = np.zeros(env.action_space.sample().shape)
next_obs, rew, done, truncated, info = env.step(actions)
@@ -49,6 +60,7 @@ def main(_):
pbar.set_description(f"Return: {returns}")
obs = next_obs
+ cs += 1
if done:
if info["succeed"]:
for transition in trajectory:
@@ -58,14 +70,14 @@ def main(_):
trajectory = []
returns = 0
obs, info = env.reset()
-
- if not os.path.exists("./demo_data"):
- os.makedirs("./demo_data")
- uuid = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
- file_name = f"./demo_data/{FLAGS.exp_name}_{success_needed}_demos_{uuid}.pkl"
- with open(file_name, "wb") as f:
- pkl.dump(transitions, f)
- print(f"saved {success_needed} demos to {file_name}")
+ dump_data(transitions)
+ print(f"{cs} steps in {time.time() - t0}s")
+ print("reset start")
+ time.sleep(1.0)
+ cs = 0
+ t0 = time.time()
+ print("reset end")
+ print("==")
if __name__ == "__main__":
app.run(main)
\ No newline at end of file
diff --git a/examples/record_success_fail.py b/examples/record_success_fail.py
index a727fbba..d938f80b 100644
--- a/examples/record_success_fail.py
+++ b/examples/record_success_fail.py
@@ -1,3 +1,6 @@
+# How to use:
+# sudo /home/qirico/miniconda3/envs/hilserl/bin/python ../../record_success_fail.py --exp_name=cube_reach2 --successes_needed=200
+
import copy
import os
from tqdm import tqdm
@@ -6,6 +9,7 @@
import datetime
from absl import app, flags
from pynput import keyboard
+import time
from experiments.mappings import CONFIG_MAPPING
@@ -15,16 +19,21 @@
success_key = False
+failure_key = False
def on_press(key):
global success_key
+ global failure_key
try:
if str(key) == 'Key.space':
success_key = True
+ elif str(key) == "f":
+ failure_key = True
except AttributeError:
pass
def main(_):
global success_key
+ global failure_key
listener = keyboard.Listener(
on_press=on_press)
listener.start()
@@ -38,7 +47,25 @@ def main(_):
success_needed = FLAGS.successes_needed
pbar = tqdm(total=success_needed)
- while len(successes) < success_needed:
+ if not os.path.exists("./classifier_data"):
+ os.makedirs("./classifier_data")
+ uuid = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+ success_file_name = f"./classifier_data/{FLAGS.exp_name}_{success_needed}_success_images_{uuid}.pkl"
+ failure_file_name = f"./classifier_data/{FLAGS.exp_name}_failure_images_{uuid}.pkl"
+
+ def dump_data(successes, failures):
+ print("dumping...")
+ with open(success_file_name, "wb") as f:
+ pkl.dump(successes, f)
+ print(f"saved {len(successes)} successful transitions to {success_file_name}")
+
+ with open(failure_file_name, "wb") as f:
+ pkl.dump(failures, f)
+ print(f"saved {len(failures)} failure transitions to {failure_file_name}")
+ print("dumped")
+
+ in_success = False
+ while True:
actions = np.zeros(env.action_space.sample().shape)
next_obs, rew, done, truncated, info = env.step(actions)
if "intervene_action" in info:
@@ -55,28 +82,42 @@ def main(_):
)
)
obs = next_obs
- if success_key:
+ if success_key or in_success:
+ in_success = True
successes.append(transition)
pbar.update(1)
success_key = False
+ # elif failure_key:
+ # failures.append(transition)
+ # failure_key = False
+ # print("failure recorded")
else:
failures.append(transition)
if done or truncated:
+ in_success = False
obs, _ = env.reset()
+ print("reset start")
+ dump_data(successes, failures)
+ time.sleep(7.0)
+ print("reset end")
+ print("==")
- if not os.path.exists("./classifier_data"):
- os.makedirs("./classifier_data")
- uuid = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
- file_name = f"./classifier_data/{FLAGS.exp_name}_{success_needed}_success_images_{uuid}.pkl"
- with open(file_name, "wb") as f:
- pkl.dump(successes, f)
- print(f"saved {success_needed} successful transitions to {file_name}")
+ # except Exception as e:
+ # if not os.path.exists("./classifier_data"):
+ # os.makedirs("./classifier_data")
+ # uuid = datetime.datetime.now().strftime("%Y-%m-%d_%H-%M-%S")
+ # file_name = f"./classifier_data/{FLAGS.exp_name}_{success_needed}_success_images_{uuid}.pkl"
+ # with open(file_name, "wb") as f:
+ # pkl.dump(successes, f)
+ # print(f"saved {success_needed} successful transitions to {file_name}")
- file_name = f"./classifier_data/{FLAGS.exp_name}_failure_images_{uuid}.pkl"
- with open(file_name, "wb") as f:
- pkl.dump(failures, f)
- print(f"saved {len(failures)} failure transitions to {file_name}")
+ # file_name = f"./classifier_data/{FLAGS.exp_name}_failure_images_{uuid}.pkl"
+ # with open(file_name, "wb") as f:
+ # pkl.dump(failures, f)
+ # print(f"saved {len(failures)} failure transitions to {file_name}")
+
+ return
if __name__ == "__main__":
app.run(main)
diff --git a/examples/run_demos.py b/examples/run_demos.py
new file mode 100644
index 00000000..18ada967
--- /dev/null
+++ b/examples/run_demos.py
@@ -0,0 +1,52 @@
+from absl import app, flags
+import time
+import numpy as np
+import pickle
+
+from experiments.mappings import CONFIG_MAPPING
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string("exp_name", None, "Name of experiment corresponding to folder.")
+flags.DEFINE_string("demo_path", None, "Path to the demos.")
+
+def pull_demos(demo_path: str):
+ with open(demo_path, "rb") as f:
+ obj = pickle.load(f)
+ assert isinstance(obj, list)
+ demos = []
+ for transition in obj:
+ assert isinstance(transition, dict)
+ assert set(["observations", "actions", "next_observations", "dones"]).issubset(set(transition.keys())), f"{transition.keys()}"
+ demos.append(transition)
+ return demos
+
+def main(_):
+ assert FLAGS.exp_name in CONFIG_MAPPING, 'Experiment folder not found.'
+ config = CONFIG_MAPPING[FLAGS.exp_name]()
+ env = config.get_environment(fake_env=False, save_video=False, classifier=True)
+
+ assert FLAGS.demo_path is not None
+ demos = pull_demos(FLAGS.demo_path)
+
+ obs, info = env.reset()
+ print("reset start")
+ time.sleep(5.0)
+ print("reset end")
+
+ done = False
+ truncated = False
+ i = 0
+ while not (done or truncated):
+ actions = demos[i]['actions']
+ t1 = time.time()
+ next_obs, rew, done, truncated, info = env.step(actions)
+ t2 = time.time()
+ print(f"Robot operating at {1 / (t2 - t1)} Hz.")
+ obs = next_obs
+ if demos[i]['dones']:
+ break
+ i += 1
+
+
+if __name__ == "__main__":
+ app.run(main)
diff --git a/examples/run_environment.py b/examples/run_environment.py
new file mode 100644
index 00000000..c3c7e4c2
--- /dev/null
+++ b/examples/run_environment.py
@@ -0,0 +1,34 @@
+from absl import app, flags
+import time
+import numpy as np
+
+from experiments.mappings import CONFIG_MAPPING
+
+FLAGS = flags.FLAGS
+flags.DEFINE_string("exp_name", None, "Name of experiment corresponding to folder.")
+
+def main(_):
+ assert FLAGS.exp_name in CONFIG_MAPPING, 'Experiment folder not found.'
+ config = CONFIG_MAPPING[FLAGS.exp_name]()
+ env = config.get_environment(fake_env=False, save_video=False, classifier=True)
+
+ while True:
+ obs, info = env.reset()
+ print("reset start")
+ time.sleep(5.0)
+ obs, info = env.reset()
+ print("reset end")
+
+ done = False
+ truncated = False
+ while not (done or truncated):
+ actions = np.zeros(env.action_space.sample().shape)
+ t1 = time.time()
+ next_obs, rew, done, truncated, info = env.step(actions)
+ t2 = time.time()
+ print(f"Robot operating at {1 / (t2 - t1)} Hz.")
+ obs = next_obs
+
+
+if __name__ == "__main__":
+ app.run(main)
diff --git a/examples/train_cl.py b/examples/train_cl.py
new file mode 100644
index 00000000..e255a46f
--- /dev/null
+++ b/examples/train_cl.py
@@ -0,0 +1,545 @@
+#!/usr/bin/env python3
+
+import glob
+import time
+import jax
+import jax.numpy as jnp
+import numpy as np
+import tqdm
+from absl import app, flags
+from flax.training import checkpoints
+import os
+import copy
+import pickle as pkl
+from gymnasium.wrappers.record_episode_statistics import RecordEpisodeStatistics
+from natsort import natsorted
+
+from serl_launcher.agents.continuous.cl import CLAgent
+from serl_launcher.utils.timer_utils import Timer
+from serl_launcher.utils.train_utils import concat_batches
+
+from agentlace.trainer import TrainerServer, TrainerClient
+from agentlace.data.data_store import QueuedDataStore
+
+from serl_launcher.utils.launcher import (
+ make_cl_pixel_agent,
+ make_trainer_config,
+ make_wandb_logger,
+)
+from serl_launcher.data.data_store import MemoryEfficientReplayBufferDataStore
+from serl_launcher.data.our_thing import ReplayBuffer
+
+from experiments.mappings import CONFIG_MAPPING
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("exp_name", None, "Name of experiment corresponding to folder.")
+flags.DEFINE_integer("seed", 42, "Random seed.")
+flags.DEFINE_boolean("learner", False, "Whether this is a learner.")
+flags.DEFINE_boolean("actor", False, "Whether this is an actor.")
+flags.DEFINE_string("ip", "localhost", "IP address of the learner.")
+flags.DEFINE_multi_string("demo_path", None, "Path to the demo data.")
+flags.DEFINE_string("checkpoint_path", None, "Path to save checkpoints.")
+flags.DEFINE_integer("eval_checkpoint_step", 0, "Step to evaluate the checkpoint.")
+flags.DEFINE_integer("eval_n_trajs", 0, "Number of trajectories to evaluate.")
+flags.DEFINE_boolean("save_video", False, "Save video.")
+
+flags.DEFINE_boolean(
+ "debug", False, "Debug mode."
+) # debug mode will disable wandb logging
+
+
+devices = jax.local_devices()
+num_devices = len(devices)
+sharding = jax.sharding.PositionalSharding(devices)
+
+
+def print_green(x):
+ return print("\033[92m {}\033[00m".format(x))
+
+
+##############################################################################
+
+
+def actor(agent, data_store, intvn_data_store, pref_data_store, env, sampling_rng):
+ """
+ This is the actor loop, which runs when "--actor" is set to True.
+ """
+ if FLAGS.eval_checkpoint_step:
+ success_counter = 0
+ time_list = []
+
+ ckpt = checkpoints.restore_checkpoint(
+ os.path.abspath(FLAGS.checkpoint_path),
+ agent.state,
+ step=FLAGS.eval_checkpoint_step,
+ )
+ agent = agent.replace(state=ckpt)
+
+ for episode in range(FLAGS.eval_n_trajs):
+ obs, _ = env.reset()
+ done = False
+ start_time = time.time()
+ while not done:
+ sampling_rng, key = jax.random.split(sampling_rng)
+ actions = agent.sample_actions(
+ observations=jax.device_put(obs),
+ argmax=False,
+ seed=key
+ )
+ actions = np.asarray(jax.device_get(actions))
+
+ next_obs, reward, done, truncated, info = env.step(actions)
+ obs = next_obs
+
+ if done:
+ if reward:
+ dt = time.time() - start_time
+ time_list.append(dt)
+ print(dt)
+
+ success_counter += reward
+ print(reward)
+ print(f"{success_counter}/{episode + 1}")
+
+ print(f"success rate: {success_counter / FLAGS.eval_n_trajs}")
+ print(f"average time: {np.mean(time_list)}")
+ return # after done eval, return and exit
+
+ start_step = (
+ int(os.path.basename(natsorted(glob.glob(os.path.join(FLAGS.checkpoint_path, "buffer/*.pkl")))[-1])[12:-4]) + 1
+ if FLAGS.checkpoint_path and os.path.exists(FLAGS.checkpoint_path)
+ else 0
+ )
+
+ datastore_dict = {
+ "actor_env": data_store,
+ "actor_env_intvn": intvn_data_store,
+ "actor_env_pref": pref_data_store,
+ }
+
+ client = TrainerClient(
+ "actor_env",
+ FLAGS.ip,
+ make_trainer_config(),
+ data_stores=datastore_dict,
+ wait_for_server=True,
+ timeout_ms=3000,
+ )
+
+ # Function to update the agent with new params
+ def update_params(params):
+ nonlocal agent
+ agent = agent.replace(state=agent.state.replace(params=params))
+
+ client.recv_network_callback(update_params)
+
+ transitions = []
+ demo_transitions = []
+
+ obs, _ = env.reset()
+ done = False
+
+ # training loop
+ timer = Timer()
+ running_return = 0.0
+ already_intervened = False
+ intervention_count = 0
+ intervention_steps = 0
+
+ o1 = None
+ o2 = None
+ a1 = None
+ a2 = None
+
+ pbar = tqdm.tqdm(range(start_step, config.max_steps), dynamic_ncols=True)
+ for step in pbar:
+ timer.tick("total")
+
+ with timer.context("sample_actions"):
+ if step < config.random_steps:
+ actions = env.action_space.sample()
+ else:
+ sampling_rng, key = jax.random.split(sampling_rng)
+ actions = agent.sample_actions(
+ observations=jax.device_put(obs),
+ seed=key,
+ argmax=False,
+ )
+ actions = np.asarray(jax.device_get(actions))
+
+ # Step environment
+ with timer.context("step_env"):
+
+ next_obs, reward, done, truncated, info = env.step(actions)
+ if "left" in info:
+ info.pop("left")
+ if "right" in info:
+ info.pop("right")
+
+ if "intervene_action" in info and not already_intervened:
+ # Start intervention
+ o1 = obs
+ a1 = actions
+ elif "intervene_action" not in info and already_intervened:
+ # End intervention
+ assert o1 is not None and a1 is not None
+ assert o2 is not None and a2 is not None
+ transition = dict(o1=o1, a1=a1, o2=o2, a2=a2)
+ pref_data_store.insert(transition)
+ o1 = a1 = o2 = a2 = None
+
+ # override the action with the intervention action
+ if "intervene_action" in info:
+ o2 = obs
+ a2 = actions
+ actions = info.pop("intervene_action")
+ intervention_steps += 1
+ if not already_intervened:
+ intervention_count += 1
+ already_intervened = True
+ else:
+ already_intervened = False
+
+ running_return += reward
+ transition = dict(
+ observations=obs,
+ actions=actions,
+ next_observations=next_obs,
+ rewards=reward,
+ masks=1.0 - done,
+ dones=done,
+ )
+ if 'grasp_penalty' in info:
+ transition['grasp_penalty']= info['grasp_penalty']
+ data_store.insert(transition)
+ transitions.append(copy.deepcopy(transition))
+ if already_intervened:
+ intvn_data_store.insert(transition)
+ demo_transitions.append(copy.deepcopy(transition))
+
+ obs = next_obs
+ if done or truncated:
+ info["episode"]["intervention_count"] = intervention_count
+ info["episode"]["intervention_steps"] = intervention_steps
+ stats = {"environment": info} # send stats to the learner to log
+ client.request("send-stats", stats)
+ pbar.set_description(f"last return: {running_return}")
+ running_return = 0.0
+ intervention_count = 0
+ intervention_steps = 0
+ already_intervened = False
+ client.update()
+ obs, _ = env.reset()
+
+ if step > 0 and config.buffer_period > 0 and step % config.buffer_period == 0:
+ # dump to pickle file
+ buffer_path = os.path.join(FLAGS.checkpoint_path, "buffer")
+ demo_buffer_path = os.path.join(FLAGS.checkpoint_path, "demo_buffer")
+ if not os.path.exists(buffer_path):
+ os.makedirs(buffer_path)
+ if not os.path.exists(demo_buffer_path):
+ os.makedirs(demo_buffer_path)
+ with open(os.path.join(buffer_path, f"transitions_{step}.pkl"), "wb") as f:
+ pkl.dump(transitions, f)
+ transitions = []
+ with open(
+ os.path.join(demo_buffer_path, f"transitions_{step}.pkl"), "wb"
+ ) as f:
+ pkl.dump(demo_transitions, f)
+ demo_transitions = []
+
+ timer.tock("total")
+
+ if step % config.log_period == 0:
+ stats = {"timer": timer.get_average_times()}
+ client.request("send-stats", stats)
+
+
+##############################################################################
+
+
+def learner(rng, agent, replay_buffer, demo_buffer, preference_buffer: ReplayBuffer, wandb_logger=None):
+ """
+ The learner loop, which runs when "--learner" is set to True.
+ """
+ start_step = (
+ int(os.path.basename(checkpoints.latest_checkpoint(os.path.abspath(FLAGS.checkpoint_path)))[11:])
+ + 1
+ if FLAGS.checkpoint_path and os.path.exists(FLAGS.checkpoint_path)
+ else 0
+ )
+ step = start_step
+
+ def stats_callback(type: str, payload: dict) -> dict:
+ """Callback for when server receives stats request."""
+ assert type == "send-stats", f"Invalid request type: {type}"
+ if wandb_logger is not None:
+ wandb_logger.log(payload, step=step)
+ return {} # not expecting a response
+
+ # Create server
+ server = TrainerServer(make_trainer_config(), request_callback=stats_callback)
+ server.register_data_store("actor_env", replay_buffer)
+ server.register_data_store("actor_env_intvn", demo_buffer)
+ server.register_data_store("actor_env_pref", preference_buffer)
+ server.start(threaded=True)
+
+ # Loop to wait until replay_buffer is filled
+ pbar = tqdm.tqdm(
+ total=config.training_starts,
+ initial=len(replay_buffer),
+ desc="Filling up replay buffer",
+ position=0,
+ leave=True,
+ )
+ while len(replay_buffer) < config.training_starts:
+ pbar.update(len(replay_buffer) - pbar.n) # Update progress bar
+ time.sleep(1)
+ pbar.update(len(replay_buffer) - pbar.n) # Update progress bar
+ pbar.close()
+
+ # send the initial network to the actor
+ server.publish_network(agent.state.params)
+ print_green("sent initial network to actor")
+
+ # 50/50 sampling from RLPD, half from demo and half from online experience
+ replay_iterator = replay_buffer.get_iterator(
+ sample_args={
+ "batch_size": config.batch_size // 2,
+ "pack_obs_and_next_obs": True,
+ },
+ device=sharding.replicate(),
+ )
+ demo_iterator = demo_buffer.get_iterator(
+ sample_args={
+ "batch_size": config.batch_size // 2,
+ "pack_obs_and_next_obs": True,
+ },
+ device=sharding.replicate(),
+ )
+ # preference_iterator = preference_buffer.get_iterator(
+ # sample_args={
+ # "batch_size": config.batch_size,
+ # "pack_obs_and_next_obs": True,
+ # },
+ # device=sharding.replicate(),
+ # )
+
+ # wait till the replay buffer is filled with enough data
+ timer = Timer()
+
+ if isinstance(agent, CLAgent):
+ train_critic_networks_to_update = frozenset({"critic"})
+ train_networks_to_update = frozenset({"critic", "actor", "temperature"})
+ else:
+ raise Exception("what")
+
+ for step in tqdm.tqdm(
+ range(start_step, config.max_steps), dynamic_ncols=True, desc="learner"
+ ):
+ # run n-1 critic updates and 1 critic + actor update.
+ # This makes training on GPU faster by reducing the large batch transfer time from CPU to GPU
+ for critic_step in range(config.cta_ratio - 1):
+ with timer.context("sample_replay_buffer"):
+ batch = next(replay_iterator)
+ demo_batch = next(demo_iterator)
+ # preference_batch = next(preference_iterator)
+ preference_batch = preference_buffer.sample(config.batch_size)
+ batch = concat_batches(batch, demo_batch, axis=0)
+ assert not (preference_batch.keys() & batch.keys()), breakpoint()
+ batch = {**batch, **preference_batch}
+
+ with timer.context("train_critics"):
+ agent, critics_info = agent.update(
+ batch,
+ networks_to_update=train_critic_networks_to_update,
+ )
+
+ with timer.context("train"):
+ batch = next(replay_iterator)
+ demo_batch = next(demo_iterator)
+ batch = concat_batches(batch, demo_batch, axis=0)
+ agent, update_info = agent.update(
+ batch,
+ networks_to_update=train_networks_to_update,
+ )
+ # publish the updated network
+ if step > 0 and step % (config.steps_per_update) == 0:
+ agent = jax.block_until_ready(agent)
+ server.publish_network(agent.state.params)
+
+ if step % config.log_period == 0 and wandb_logger:
+ wandb_logger.log(update_info, step=step)
+ wandb_logger.log({"timer": timer.get_average_times()}, step=step)
+
+ if (
+ step > 0
+ and config.checkpoint_period
+ and step % config.checkpoint_period == 0
+ ):
+ checkpoints.save_checkpoint(
+ os.path.abspath(FLAGS.checkpoint_path), agent.state, step=step, keep=100
+ )
+
+
+##############################################################################
+
+
+def main(_):
+ global config
+ config = CONFIG_MAPPING[FLAGS.exp_name]()
+
+ assert config.batch_size % num_devices == 0
+ # seed
+ rng = jax.random.PRNGKey(FLAGS.seed)
+ rng, sampling_rng = jax.random.split(rng)
+
+ assert FLAGS.exp_name in CONFIG_MAPPING, "Experiment folder not found."
+ env = config.get_environment(
+ fake_env=FLAGS.learner,
+ save_video=FLAGS.save_video,
+ classifier=True,
+ )
+ env = RecordEpisodeStatistics(env)
+
+ rng, sampling_rng = jax.random.split(rng)
+
+ if config.setup_mode == 'single-arm-fixed-gripper' or config.setup_mode == 'dual-arm-fixed-gripper':
+ agent: CLAgent = make_cl_pixel_agent(
+ seed=FLAGS.seed,
+ sample_obs=env.observation_space.sample(),
+ sample_action=env.action_space.sample(),
+ image_keys=config.image_keys,
+ encoder_type=config.encoder_type,
+ discount=config.discount,
+ )
+ include_grasp_penalty = False
+ else:
+ raise NotImplementedError(f"Unknown setup mode: {config.setup_mode}")
+
+ # replicate agent across devices
+ # need the jnp.array to avoid a bug where device_put doesn't recognize primitives
+ agent = jax.device_put(
+ jax.tree_util.tree_map(jnp.array, agent), sharding.replicate()
+ )
+
+ if FLAGS.checkpoint_path is not None and os.path.exists(FLAGS.checkpoint_path):
+ input("Checkpoint path already exists. Press Enter to resume training.")
+ ckpt = checkpoints.restore_checkpoint(
+ os.path.abspath(FLAGS.checkpoint_path),
+ agent.state,
+ )
+ agent = agent.replace(state=ckpt)
+ ckpt_number = os.path.basename(
+ checkpoints.latest_checkpoint(os.path.abspath(FLAGS.checkpoint_path))
+ )[11:]
+ print_green(f"Loaded previous checkpoint at step {ckpt_number}.")
+
+ def create_replay_buffer_and_wandb_logger():
+ replay_buffer = MemoryEfficientReplayBufferDataStore(
+ env.observation_space,
+ env.action_space,
+ capacity=config.replay_buffer_capacity,
+ image_keys=config.image_keys,
+ include_grasp_penalty=include_grasp_penalty,
+ )
+ # set up wandb and logging
+ wandb_logger = make_wandb_logger(
+ project="hil-serl",
+ description=FLAGS.exp_name,
+ debug=FLAGS.debug,
+ )
+ return replay_buffer, wandb_logger
+
+ if FLAGS.learner:
+ sampling_rng = jax.device_put(sampling_rng, device=sharding.replicate())
+ replay_buffer, wandb_logger = create_replay_buffer_and_wandb_logger()
+ demo_buffer = MemoryEfficientReplayBufferDataStore(
+ env.observation_space,
+ env.action_space,
+ capacity=config.replay_buffer_capacity,
+ image_keys=config.image_keys,
+ include_grasp_penalty=include_grasp_penalty,
+ )
+ preference_buffer = ReplayBuffer.create(
+ dict(
+ o1=env.observation_space.sample(),
+ o2=env.observation_space.sample(),
+ a1=env.action_space.sample(),
+ a2=env.action_space.sample(),
+ ),
+ size=int(1e6),
+ )
+
+ assert FLAGS.demo_path is not None
+ for path in FLAGS.demo_path:
+ with open(path, "rb") as f:
+ transitions = pkl.load(f)
+ for transition in transitions:
+ if 'infos' in transition and 'grasp_penalty' in transition['infos']:
+ transition['grasp_penalty'] = transition['infos']['grasp_penalty']
+ demo_buffer.insert(transition)
+ print_green(f"demo buffer size: {len(demo_buffer)}")
+ print_green(f"online buffer size: {len(replay_buffer)}")
+
+ if FLAGS.checkpoint_path is not None and os.path.exists(
+ os.path.join(FLAGS.checkpoint_path, "buffer")
+ ):
+ for file in glob.glob(os.path.join(FLAGS.checkpoint_path, "buffer/*.pkl")):
+ with open(file, "rb") as f:
+ transitions = pkl.load(f)
+ for transition in transitions:
+ replay_buffer.insert(transition)
+ print_green(
+ f"Loaded previous buffer data. Replay buffer size: {len(replay_buffer)}"
+ )
+
+ if FLAGS.checkpoint_path is not None and os.path.exists(
+ os.path.join(FLAGS.checkpoint_path, "demo_buffer")
+ ):
+ for file in glob.glob(
+ os.path.join(FLAGS.checkpoint_path, "demo_buffer/*.pkl")
+ ):
+ with open(file, "rb") as f:
+ transitions = pkl.load(f)
+ for transition in transitions:
+ demo_buffer.insert(transition)
+ print_green(
+ f"Loaded previous demo buffer data. Demo buffer size: {len(demo_buffer)}"
+ )
+
+ # learner loop
+ print_green("starting learner loop")
+ learner(
+ sampling_rng,
+ agent,
+ replay_buffer,
+ demo_buffer=demo_buffer,
+ preference_buffer=preference_buffer,
+ wandb_logger=wandb_logger,
+ )
+
+ elif FLAGS.actor:
+ sampling_rng = jax.device_put(sampling_rng, sharding.replicate())
+ data_store = QueuedDataStore(50000) # the queue size on the actor
+ intvn_data_store = QueuedDataStore(50000)
+ pref_data_store = QueuedDataStore(50000)
+
+ # actor loop
+ print_green("starting actor loop")
+ actor(
+ agent,
+ data_store,
+ intvn_data_store,
+ pref_data_store,
+ env,
+ sampling_rng,
+ )
+
+ else:
+ raise NotImplementedError("Must be either a learner or an actor")
+
+
+if __name__ == "__main__":
+ app.run(main)
diff --git a/examples/train_hgdagger.py b/examples/train_hgdagger.py
index ced81c7e..69f46e52 100644
--- a/examples/train_hgdagger.py
+++ b/examples/train_hgdagger.py
@@ -362,7 +362,7 @@ def main(_):
# replicate agent across devices
# need the jnp.array to avoid a bug where device_put doesn't recognize primitives
agent: BCAgent = jax.device_put(
- jax.tree_map(jnp.array, agent), sharding.replicate()
+ jax.tree_util.tree_map(jnp.array, agent), sharding.replicate()
)
if FLAGS.learner:
diff --git a/examples/train_reward_classifier.py b/examples/train_reward_classifier.py
index 4939e6bf..e5cbd241 100644
--- a/examples/train_reward_classifier.py
+++ b/examples/train_reward_classifier.py
@@ -9,6 +9,8 @@
import optax
from tqdm import tqdm
from absl import app, flags
+from sklearn.metrics import roc_curve, auc
+import matplotlib.pyplot as plt
from serl_launcher.data.data_store import ReplayBuffer
from serl_launcher.utils.train_utils import concat_batches
@@ -17,7 +19,6 @@
from experiments.mappings import CONFIG_MAPPING
-
FLAGS = flags.FLAGS
flags.DEFINE_string("exp_name", None, "Name of experiment corresponding to folder.")
flags.DEFINE_integer("num_epochs", 150, "Number of training epochs.")
@@ -40,15 +41,27 @@ def main(_):
include_label=True,
)
+ eval_pos_buffer = ReplayBuffer(
+ env.observation_space,
+ env.action_space,
+ capacity=20000,
+ include_label=True,
+ )
+
success_paths = glob.glob(os.path.join(os.getcwd(), "classifier_data", "*success*.pkl"))
for path in success_paths:
success_data = pkl.load(open(path, "rb"))
- for trans in success_data:
+ train_success_data_len = int(.9*len(success_data))
+ for i, trans in enumerate(success_data):
if "images" in trans['observations'].keys():
continue
trans["labels"] = 1
trans['actions'] = env.action_space.sample()
- pos_buffer.insert(trans)
+ if i < train_success_data_len:
+ pos_buffer.insert(trans)
+ else:
+ pos_buffer.insert(trans)
+ # eval_pos_buffer.insert(trans)
pos_iterator = pos_buffer.get_iterator(
sample_args={
@@ -64,17 +77,28 @@ def main(_):
capacity=50000,
include_label=True,
)
+
+ eval_neg_buffer = ReplayBuffer(
+ env.observation_space,
+ env.action_space,
+ capacity=50000,
+ include_label=True,
+ )
failure_paths = glob.glob(os.path.join(os.getcwd(), "classifier_data", "*failure*.pkl"))
for path in failure_paths:
failure_data = pkl.load(
open(path, "rb")
)
- for trans in failure_data:
+ train_failure_data_len = int(.9*len(failure_data))
+ for i, trans in enumerate(failure_data):
if "images" in trans['observations'].keys():
continue
trans["labels"] = 0
trans['actions'] = env.action_space.sample()
- neg_buffer.insert(trans)
+ if i < train_failure_data_len:
+ neg_buffer.insert(trans)
+ else:
+ eval_neg_buffer.insert(trans)
neg_iterator = neg_buffer.get_iterator(
sample_args={
@@ -156,6 +180,69 @@ def loss_fn(params):
step=FLAGS.num_epochs,
overwrite=True,
)
+
+ # Calculate logits and ROC curve
+ def calculate_logits_and_plot_roc(pos_buffer, neg_buffer, classifier, config):
+ all_logits = []
+ all_labels = []
+ pos_logits = []
+ neg_logits = []
+ batch_size = 256
+ # Collect logits for all positive samples
+ pos_data = pos_buffer.dataset_dict
+ for start in tqdm(range(0, len(pos_buffer), batch_size)):
+ end = start + batch_size
+ end = min(len(pos_buffer), end)
+ obs = {k: v[start:end] for k,v in pos_data["observations"].items()}
+ labels = pos_data["labels"][start:end]
+ logits = classifier.apply_fn({"params": classifier.params}, obs, train=False)
+ all_logits.append(logits)
+ pos_logits.append(logits)
+ all_labels.append(labels)
+
+ neg_data = neg_buffer.dataset_dict
+ for start in tqdm(range(0, len(neg_buffer), batch_size)):
+ end = start + batch_size
+ end = min(len(neg_buffer), end)
+ obs = {k: v[start:end] for k,v in neg_data["observations"].items()}
+ labels = neg_data["labels"][start:end]
+ logits = classifier.apply_fn({"params": classifier.params}, obs, train=False)
+ all_logits.append(logits)
+ neg_logits.append(logits)
+ all_labels.append(labels)
+
+ # Flatten the logits and labels
+ all_logits = jnp.concatenate(all_logits, axis=0)
+ all_labels = jnp.concatenate(all_labels, axis=0)
+
+ pos_logits = nn.sigmoid(jnp.concatenate(pos_logits, axis=0))
+ neg_logits = nn.sigmoid(jnp.concatenate(neg_logits, axis=0))
+ # Calculate ROC curve
+ fpr, tpr, _ = roc_curve(all_labels, nn.sigmoid(all_logits))
+ roc_auc = auc(fpr, tpr)
+
+ # Plot ROC curve
+ plt.figure()
+ plt.plot(fpr, tpr, color='blue', lw=2, label=f'ROC curve (area = {roc_auc:.2f})')
+ plt.plot([0, 1], [0, 1], color='gray', linestyle='--')
+ plt.xlim([0.0, 1.0])
+ plt.ylim([0.0, 1.05])
+ plt.xlabel('False Positive Rate')
+ plt.ylabel('True Positive Rate')
+ plt.title('Receiver Operating Characteristic')
+ plt.legend(loc='lower right')
+ plt.show()
+
+ acc1 = ((all_logits>0.8)==(all_labels)).mean()
+ plt.hist(pos_logits, bins='auto', alpha=0.5, label="Positive Samples", color="green")
+ plt.hist(neg_logits, bins='auto', alpha=0.5, label="Negative Samples", color="red")
+ plt.title(f"Training Accuracy: {acc1}")
+ plt.show()
+ get_acc = lambda x: ((all_logits>x)==(all_labels)).mean()
+ breakpoint()
+
+ # Calculate and plot ROC after training
+ # calculate_logits_and_plot_roc(eval_pos_buffer, eval_neg_buffer, classifier, config)
if __name__ == "__main__":
diff --git a/examples/train_rlif.py b/examples/train_rlif.py
new file mode 100644
index 00000000..fb57a98f
--- /dev/null
+++ b/examples/train_rlif.py
@@ -0,0 +1,966 @@
+#!/usr/bin/env python3
+
+import glob
+import time
+import jax
+import jax.numpy as jnp
+import numpy as np
+import tqdm
+from absl import app, flags
+from flax.training import checkpoints
+import os
+import copy
+import pickle as pkl
+from gymnasium.wrappers.record_episode_statistics import RecordEpisodeStatistics
+from natsort import natsorted
+from pynput import keyboard
+import requests
+import cv2
+import threading
+import queue
+
+from serl_launcher.agents.continuous.sac import SACAgent
+from serl_launcher.agents.continuous.sac_hybrid_single import SACAgentHybridSingleArm
+from serl_launcher.agents.continuous.sac_hybrid_dual import SACAgentHybridDualArm
+from serl_launcher.utils.timer_utils import Timer
+from serl_launcher.utils.train_utils import concat_batches
+from serl_launcher.utils.tools import ImageDisplayer, q_image
+
+from agentlace.trainer import TrainerServer, TrainerClient
+from agentlace.data.data_store import QueuedDataStore
+
+from serl_launcher.utils.launcher import (
+ make_sac_pixel_agent,
+ make_sac_pixel_agent_hybrid_single_arm,
+ make_sac_pixel_agent_hybrid_dual_arm,
+ make_trainer_config,
+ make_wandb_logger,
+)
+from serl_launcher.data.data_store import MemoryEfficientReplayBufferDataStore, PreferenceBufferDataStore
+
+from experiments.mappings import CONFIG_MAPPING
+
+FLAGS = flags.FLAGS
+
+flags.DEFINE_string("exp_name", None, "Name of experiment corresponding to folder.")
+flags.DEFINE_string("method", "rlif", "Valid values: rlif, cl, soft_cl, hgdagger")
+flags.DEFINE_integer("seed", 42, "Random seed.")
+flags.DEFINE_boolean("learner", False, "Whether this is a learner.")
+flags.DEFINE_boolean("actor", False, "Whether this is an actor.")
+flags.DEFINE_string("ip", "localhost", "IP address of the learner.")
+flags.DEFINE_multi_string("demo_path", None, "Path to the demo data.")
+flags.DEFINE_string("checkpoint_path", None, "Path to save checkpoints.")
+flags.DEFINE_integer("eval_checkpoint_step", 0, "Step to evaluate the checkpoint.")
+flags.DEFINE_integer("eval_n_trajs", 0, "Number of trajectories to evaluate.")
+flags.DEFINE_boolean("save_video", False, "Save video.")
+flags.DEFINE_boolean("use_bc_loss", False, "Whether to add a crude bc max likelihood loss on the policy.")
+
+flags.DEFINE_float("optimism", 0.0, "Whether or not to add a small amount of bonus to the post-state of interventions.")
+flags.DEFINE_boolean("optimism_done_mask", False, "The done to be set for the optimism transition.")
+flags.DEFINE_boolean("show_q_values", False, "Whether or not to open another window that shows the live Q values.")
+flags.DEFINE_boolean("state_based", False, "Whether or not to use states instead of image + encoder.")
+flags.DEFINE_boolean(
+ "debug", False, "Debug mode."
+) # debug mode will disable wandb logging
+
+
+devices = jax.local_devices()
+num_devices = len(devices)
+sharding = jax.sharding.PositionalSharding(devices)
+
+
+def print_green(x):
+ return print("\033[92m {}\033[00m".format(x))
+
+def print_yellow(x):
+ return print("\033[93m {}\033[00m".format(x))
+
+def print_cyan(x):
+ return print("\033[96m {}\033[00m".format(x))
+
+def print_cyan(x):
+ return print("\033[96m {}\033[00m".format(x))
+
+
+##############################################################################
+
+
+failure_key = False
+checkpoint_key = False
+pause_key = False
+def on_press(key):
+ global failure_key
+ global checkpoint_key
+ global pause_key
+ try:
+ if str(key) == "'f'":
+ failure_key = True
+ elif str(key) == "'c'":
+ checkpoint_key = True
+ elif str(key) == "'r'":
+ print_yellow("reset")
+ requests.post("http://localhost:5000/reset_gripper")
+ elif str(key) == "'t'":
+ print_yellow("close gripper for reset")
+ requests.post("http://localhost:5000/close_gripper")
+ elif str(key) == "'p'":
+ print_yellow("pause")
+ pause_key = not pause_key
+ except AttributeError:
+ print("error")
+ pass
+
+def actor(agent, data_store, intvn_data_store, env, sampling_rng, pref_data_store = None, bc_data_store = None):
+ """
+ This is the actor loop, which runs when "--actor" is set to True.
+ """
+ global failure_key
+ global checkpoint_key
+ global pause_key
+
+ listener = keyboard.Listener(
+ on_press=on_press)
+ listener.start()
+
+ if FLAGS.eval_checkpoint_step:
+ success_counter = 0
+ time_list = []
+
+ ckpt = checkpoints.restore_checkpoint(
+ os.path.abspath(FLAGS.checkpoint_path),
+ agent.state,
+ step=FLAGS.eval_checkpoint_step,
+ )
+ agent = agent.replace(state=ckpt)
+ print_green(f"Overriding agent with checkpoint at {FLAGS.eval_checkpoint_step}.")
+
+ for episode in range(FLAGS.eval_n_trajs):
+ print("reset start")
+ ### receive signal from learner and then reset
+ obs, _ = env.reset()
+ time.sleep(7.0)
+ obs, _ = env.reset()
+ print("reset end")
+ done = False
+ start_time = time.time()
+ while not done:
+ sampling_rng, key = jax.random.split(sampling_rng)
+ actions = agent.sample_actions(
+ observations=jax.device_put(obs),
+ argmax=True,
+ seed=key
+ )
+ actions = np.asarray(jax.device_get(actions))
+
+ next_obs, reward, done, truncated, info = env.step(actions)
+ if failure_key:
+ print("failure detected")
+ failure_key = False
+ done = True
+ obs = next_obs
+
+ if done or truncated:
+ if reward:
+ dt = time.time() - start_time
+ time_list.append(dt)
+ print(dt)
+
+ success_counter += reward
+ print(reward)
+ print(f"{success_counter}/{episode + 1}")
+
+ print(f"success rate: {success_counter / FLAGS.eval_n_trajs}")
+ print(f"average time: {np.mean(time_list)}")
+ return # after done eval, return and exit
+
+ buffers = (
+ natsorted(glob.glob(os.path.join(FLAGS.checkpoint_path, "buffer/*.pkl")))
+ if FLAGS.checkpoint_path and os.path.exists(FLAGS.checkpoint_path)
+ else []
+ )
+ start_step = 1 if len(buffers) == 0 else int(os.path.basename(buffers[-1])[12:-4]) + 1
+
+ datastore_dict = {
+ "actor_env": data_store,
+ "actor_env_intvn": intvn_data_store,
+ }
+
+ if pref_data_store is not None:
+ datastore_dict["actor_env_pref"] = pref_data_store
+ if FLAGS.use_bc_loss:
+ assert bc_data_store is not None
+ datastore_dict["actor_env_bc"] = bc_data_store
+
+ client = TrainerClient(
+ "actor_env",
+ FLAGS.ip,
+ make_trainer_config(),
+ data_stores=datastore_dict,
+ wait_for_server=True,
+ timeout_ms=3000,
+ )
+
+ # Function to update the agent with new params
+ learner_step = start_step
+ def update_params(params):
+ if isinstance(params, dict) and set(params.keys()) == set(["step"]):
+ nonlocal learner_step
+ learner_step = params['step']
+ else:
+ nonlocal agent
+ agent = agent.replace(state=agent.state.replace(params=params))
+
+ client.recv_network_callback(update_params)
+
+ transitions = []
+ transitions_full_trajs = []
+ demo_transitions = []
+ demo_transitions_full_trajs = []
+ interventions = []
+ this_intervention = None
+ preference_datas = []
+ bc_transitions = []
+
+ print("reset start")
+ obs, _ = env.reset()
+ time.sleep(5.0)
+ obs, _ = env.reset()
+ print("reset end")
+ done = False
+
+ # training loop
+ timer = Timer()
+ running_return = 0.0
+ already_intervened = False
+ intervention_count = 0
+ intervention_steps = 0
+ total_interventions = 0
+
+ pre_int_obs = None
+ post_int_obs = None
+ a_int_pi = None
+ a_int_exp = None
+
+ if FLAGS.show_q_values:
+ q_queue = queue.Queue()
+ q_display = ImageDisplayer(q_queue, "q_display")
+ q_display.start()
+
+ pbar = tqdm.tqdm(range(start_step, config.max_steps), dynamic_ncols=True)
+ print(config.buffer_period)
+ from_time = time.time()
+ cur_steps = 0
+ for step in pbar:
+ while pause_key:
+ time.sleep(0.5)
+ timer.tick("total")
+
+ with timer.context("sample_actions"):
+ if step < config.random_steps:
+ actions = env.action_space.sample()
+ else:
+ sampling_rng, key = jax.random.split(sampling_rng)
+ actions = agent.sample_actions(
+ observations=jax.device_put(obs),
+ seed=key,
+ argmax=False,
+ )
+ actions = np.asarray(jax.device_get(actions))
+
+ # Step environment
+ if FLAGS.show_q_values:
+ q_value = float(np.asarray(jax.device_get(agent.forward_critic(jax.device_put(obs), actions[:-1], rng=key).min())))
+ q_value_grasp = float(np.asarray(jax.device_get(agent.forward_grasp_critic(jax.device_put(obs), rng=key)[int(actions[-1] + 1)])))
+ with timer.context("step_env"):
+
+ next_obs, reward, done, truncated, info = env.step(actions)
+ reward *= config.reward_scale
+ cur_steps += 1
+ if "left" in info:
+ info.pop("left")
+ if "right" in info:
+ info.pop("right")
+ if failure_key:
+ print("failure detected")
+ failure_key = False
+ truncated = True
+ if FLAGS.show_q_values:
+ info['q_value_grasp'] = q_value_grasp
+ info['q_value'] = q_value
+ q_queue.put({'q_image': q_image(q_value, q_value_grasp, 'intervene_action' in info)})
+
+ # override the action with the intervention action
+ if "intervene_action" in info:
+ policy_actions = actions
+ info["policy_action"] = actions
+ actions = info["intervene_action"]
+ info["policy_action"] = actions
+ actions = info["intervene_action"]
+ intervention_steps += 1
+
+ post_int_obs = next_obs
+
+ if not already_intervened:
+ print_cyan("Started intervention.")
+ intervention_count += 1
+
+ pre_int_obs = obs
+ a_int_exp = actions
+ a_int_pi = policy_actions
+
+ this_intervention = dict(
+ t0=cur_steps,
+ t1=cur_steps+1,
+ observations=[obs, next_obs],
+ actions=[actions],
+ policy_actions=[policy_actions],
+ dones=[done],
+ truncateds=[truncated],
+ )
+ if config.rlif_minus_one:
+ print(f"Detected intervention; RLIF replaced reward {reward} with {-1}.")
+ reward = -1
+ else:
+ this_intervention['observations'].append(next_obs)
+ this_intervention['actions'].append(actions)
+ this_intervention['policy_actions'].append(policy_actions)
+ this_intervention['dones'].append(done)
+ this_intervention['truncateds'].append(truncated)
+ this_intervention['t1'] = cur_steps+1
+ already_intervened = True
+ else:
+ if already_intervened:
+ print_cyan(f"Ended intervention of {this_intervention['t1'] - this_intervention['t0']} steps.")
+ if this_intervention is None:
+ print("Error: Should not be None")
+ interventions.append(this_intervention)
+ # add to preference buffer
+ if FLAGS.method in ["cl", "soft_cl"]:
+ pref_datapoint = dict(
+ pre_obs=pre_int_obs,
+ post_obs=post_int_obs,
+ a_pi=a_int_pi,
+ a_exp=a_int_exp,
+ t=np.array([this_intervention['t1'] - this_intervention['t0']]),
+ )
+ pref_data_store.insert(pref_datapoint)
+ preference_datas.append(pref_datapoint)
+ this_intervention = None
+
+ if abs(FLAGS.optimism) > 1e-9:
+ print_cyan(f"Adding optimism transition with reward={FLAGS.optimism} and done={FLAGS.optimism_done_mask}.")
+ transition = dict(
+ observations=obs,
+ actions=actions,
+ next_observations=next_obs,
+ rewards=FLAGS.optimism,
+ masks=1.0 - FLAGS.optimism_done_mask, # Used in training, denoting whether or not we're at the end of a trajectory.
+ dones=FLAGS.optimism_done_mask, # Not actually used in training.
+ )
+ if 'grasp_penalty' in info:
+ transition['grasp_penalty'] = 0
+ data_store.insert(transition)
+ transitions.append(copy.deepcopy(transition) | {'info': info | {'optimism': True}})
+ already_intervened = False
+
+ if (done or truncated) and this_intervention is not None:
+ print_cyan(f"Ended intervention of {this_intervention['t1'] - this_intervention['t0']} steps.")
+ interventions.append(this_intervention)
+ # add to preference buffer
+ if FLAGS.method in ["cl", "soft_cl"]:
+ pref_datapoint = dict(
+ pre_obs=pre_int_obs,
+ post_obs=post_int_obs,
+ a_pi=a_int_pi,
+ a_exp=a_int_exp,
+ t=np.array([this_intervention['t1'] - this_intervention['t0']]),
+ )
+ pref_data_store.insert(pref_datapoint)
+ preference_datas.append(pref_datapoint)
+ this_intervention = None
+
+ running_return += reward
+ transition = dict(
+ observations=obs,
+ actions=actions,
+ next_observations=next_obs,
+ rewards=reward,
+ masks=1.0 - done,
+ dones=done,
+ )
+ # if checkpoint_key:
+ # breakpoint()
+ if 'grasp_penalty' in info:
+ transition['grasp_penalty']= info['grasp_penalty']
+ data_store.insert(transition)
+ transitions.append(copy.deepcopy(transition) | {'info': info})
+ if already_intervened:
+ intvn_data_store.insert(transition)
+ demo_transitions.append(copy.deepcopy(transition) | {'info': info})
+ if FLAGS.use_bc_loss and "intervene_action" in info:
+ bc_data_store.insert(transition)
+ bc_transitions.append(transition)
+
+ obs = next_obs
+ if done or truncated:
+ failure_key = False
+ if "episode" not in info:
+ info["episode"] = {}
+ info["episode"]["intervention_count"] = intervention_count
+ info["episode"]["intervention_steps"] = intervention_steps
+
+ total_interventions += intervention_count
+
+ info["episode"]["total_interventions"] = total_interventions
+ info["episode"]["intervention_rate"] = total_interventions / (step + 1)
+ info["episode"]["current_intervention_rate"] = intervention_steps / cur_steps
+ info["episode"]["episode_duration"] = time.time() - from_time
+ info["episode"]["success_rate"] = running_return
+ info["episode"]["episode_steps"] = cur_steps
+ info["episode"]["environment_step"] = step
+ stats = {"environment": info} # send stats to the learner to log
+ client.request("send-stats", stats)
+ pbar.set_description(f"last return: {running_return}")
+ running_return = 0.0
+ intervention_count = 0
+ intervention_steps = 0
+ cur_steps = 0
+
+ already_intervened = False
+ client.update()
+ print("reset start")
+ obs, _ = env.reset()
+ transitions_full_trajs = transitions
+ demo_transitions_full_trajs = demo_transitions
+ # input("Waiting for input to proceed...")
+ time.sleep(7.0)
+ obs, _ = env.reset()
+ # For synchronizing learner and actor...
+ if step > learner_step * 1 + 200:
+ print("Stopped actor")
+ while step > learner_step * 1 + 200:
+ time.sleep(0.5)
+ print("Released actor")
+ print("reset end")
+ from_time = time.time()
+
+ if step > 0 and config.buffer_period > 0 and step % config.buffer_period == 0:
+ # dump to pickle file
+ buffer_path = os.path.join(FLAGS.checkpoint_path, "buffer")
+ demo_buffer_path = os.path.join(FLAGS.checkpoint_path, "demo_buffer")
+ interventions_buffer_path = os.path.join(FLAGS.checkpoint_path, "interventions")
+ preference_buffer_path = os.path.join(FLAGS.checkpoint_path, "preference_buffer")
+ bc_buffer_path = os.path.join(FLAGS.checkpoint_path, "bc_buffer")
+
+ if not os.path.exists(buffer_path):
+ os.makedirs(buffer_path)
+ if not os.path.exists(demo_buffer_path):
+ os.makedirs(demo_buffer_path)
+ if not os.path.exists(interventions_buffer_path):
+ os.makedirs(interventions_buffer_path)
+ if not os.path.exists(preference_buffer_path):
+ os.makedirs(preference_buffer_path)
+ if not os.path.exists(bc_buffer_path):
+ os.makedirs(bc_buffer_path)
+
+ with open(os.path.join(buffer_path, f"transitions_{step}.pkl"), "wb") as f:
+ fp = os.path.join(buffer_path, f"transitions_{step}.pkl")
+ print(f"Dumping {len(transitions_full_trajs)} transitions out of {len(transitions)} to {fp} !!!")
+ pkl.dump(transitions_full_trajs, f)
+ transitions = transitions[len(transitions_full_trajs):]
+ transitions_full_trajs = []
+ with open(os.path.join(demo_buffer_path, f"transitions_{step}.pkl"), "wb") as f:
+ fp = os.path.join(demo_buffer_path, f"transitions_{step}.pkl")
+ print(f"Dumping {len(demo_transitions_full_trajs)} expert transitions out of {len(demo_transitions)} to {fp} !!!")
+ pkl.dump(demo_transitions_full_trajs, f)
+ demo_transitions = demo_transitions[len(demo_transitions_full_trajs):]
+ demo_transitions_full_trajs = []
+ with open(os.path.join(interventions_buffer_path, f"transitions_{step}.pkl"), "wb") as f:
+ fp = os.path.join(interventions_buffer_path, f"transitions_{step}.pkl")
+ print(f"Dumping {len(interventions)} interventions to {fp}")
+ pkl.dump(interventions, f)
+ interventions = []
+ with open(os.path.join(bc_buffer_path, f"transitions_{step}.pkl"), "wb") as f:
+ fp = os.path.join(bc_buffer_path, f"transitions_{step}.pkl")
+ print(f"Dumping {len(bc_transitions)} interventions to {fp}")
+ pkl.dump(bc_transitions, f)
+ bc_transitions = []
+
+ if (
+ step > 0
+ and config.checkpoint_period
+ and step % config.checkpoint_period == 0
+ ):
+ checkpoints.save_checkpoint(
+ os.path.abspath(FLAGS.checkpoint_path), agent.state, step=step, keep=100
+ )
+
+ timer.tock("total")
+
+ if step % config.log_period == 0:
+ stats = {"timer": timer.get_average_times()}
+ client.request("send-stats", stats)
+
+ if FLAGS.show_q_values:
+ q_queue.put(None)
+ cv2.destroyAllWindows()
+ q_display.join()
+
+
+##############################################################################
+
+
+def learner(rng, agent: SACAgentHybridSingleArm, replay_buffer, demo_buffer, preference_buffer = None, bc_buffer = None, wandb_logger=None):
+ """
+ The learner loop, which runs when "--learner" is set to True.
+ """
+ start_step = (
+ int(os.path.basename(checkpoints.latest_checkpoint(os.path.abspath(FLAGS.checkpoint_path)))[11:])
+ + 1
+ if FLAGS.checkpoint_path and os.path.exists(FLAGS.checkpoint_path)
+ else 0
+ )
+ step = start_step
+
+ if isinstance(agent, SACAgent):
+ train_critic_networks_to_update = frozenset({"critic"})
+ train_networks_to_update = frozenset({"critic", "actor", "temperature"})
+ else:
+ train_critic_networks_to_update = frozenset({"critic", "grasp_critic"})
+ train_networks_to_update = frozenset({"critic", "grasp_critic", "actor", "temperature"})
+
+ if FLAGS.method == "cl" and FLAGS.method != "soft_cl":
+ assert "modules_log_alpha_state" in agent.state.params
+ assert "modules_log_alpha_gripper_state" in agent.state.params
+ train_critic_networks_to_update = frozenset(train_critic_networks_to_update | {"log_alpha_state", "log_alpha_grasp_state"})
+ train_networks_to_update = frozenset(train_networks_to_update | {"log_alpha_state", "log_alpha_grasp_state"})
+
+
+ def stats_callback(type: str, payload: dict) -> dict:
+ """Callback for when server receives stats request."""
+ assert type == "send-stats", f"Invalid request type: {type}"
+ if wandb_logger is not None:
+ wandb_logger.log(payload, step=step) # + config.pretraining_steps)
+ return {} # not expecting a response
+
+ # Create server
+ server = TrainerServer(make_trainer_config(), request_callback=stats_callback)
+ server.register_data_store("actor_env", replay_buffer)
+ server.register_data_store("actor_env_intvn", demo_buffer)
+ if preference_buffer is not None:
+ server.register_data_store("actor_env_pref", preference_buffer)
+ server.start(threaded=True)
+
+ if FLAGS.use_bc_loss and config.pretraining_steps > 0:
+ pbar = tqdm.tqdm(range(config.pretraining_steps))
+ for step in pbar:
+ bc_batch = bc_buffer.sample(config.batch_size)
+ agent, update_info = agent.update_bc(bc_batch)
+ pbar.set_description(f"bc_loss = {round(update_info['actor']['bc_loss'], 3)}")
+ wandb_logger.log(update_info, step=step)
+ step = config.pretraining_steps
+ agent = jax.block_until_ready(agent)
+ server.publish_network(agent.state.params)
+ checkpoints.save_checkpoint(
+ os.path.abspath(FLAGS.checkpoint_path), agent.state, step=1, keep=100
+ )
+
+ # Loop to wait until replay_buffer is filled
+ pbar = tqdm.tqdm(
+ total=config.training_starts,
+ initial=len(replay_buffer),
+ desc="Filling up replay buffer",
+ position=0,
+ leave=True,
+ )
+ while len(replay_buffer) < config.training_starts:
+ pbar.update(len(replay_buffer) - pbar.n) # Update progress bar
+ time.sleep(1)
+ pbar.update(len(replay_buffer) - pbar.n) # Update progress bar
+ pbar.close()
+
+ # send the initial network to the actor
+ server.publish_network(agent.state.params)
+ print_green("sent initial network to actor")
+
+ # 50/50 sampling from RLPD, half from demo and half from online experience
+ replay_iterator = replay_buffer.get_iterator(
+ sample_args={
+ "batch_size": config.batch_size // 2,
+ "pack_obs_and_next_obs": True,
+ },
+ device=sharding.replicate(),
+ )
+ demo_iterator = demo_buffer.get_iterator(
+ sample_args={
+ "batch_size": config.batch_size // 2,
+ "pack_obs_and_next_obs": True,
+ },
+ device=sharding.replicate(),
+ )
+ if FLAGS.method in ["cl", "soft_cl"]:
+ assert preference_buffer is not None
+ preference_iterator = preference_buffer.get_iterator(
+ sample_args={
+ "batch_size": config.batch_size,
+ },
+ device=sharding.replicate(),
+ )
+ if FLAGS.use_bc_loss:
+ assert bc_buffer is not None
+ bc_iterator = bc_buffer.get_iterator(
+ sample_args={
+ "batch_size": config.batch_size,
+ },
+ device=sharding.replicate(),
+ )
+
+ # wait till the replay buffer is filled with enough data
+ timer = Timer()
+
+
+ '''
+ update_steps = 0
+ current_envs_steps = 0
+ while update_steps < max_steps:
+ new_steps = getfromactor - current_env_steps
+ num_updates = new_steps * utd
+ for _ in range(num)
+ update###
+ send signal to actor to get next trajectory
+
+
+
+ '''
+ for step in tqdm.tqdm(
+ range(start_step + config.pretraining_steps, config.max_steps + config.pretraining_steps), dynamic_ncols=True, desc="learner"
+ ):
+ if step - config.pretraining_steps > len(replay_buffer) * 1 + 1 + 300:
+ while step - config.pretraining_steps > len(replay_buffer) * 1 + 1 + 300:
+ time.sleep(0.5)
+ print(f"Training for another {(len(replay_buffer) * 1 + 1 + 300) - (step - config.pretraining_steps) + 1} steps...")
+ # run n-1 critic updates and 1 critic + actor update.
+ # This makes training on GPU faster by reducing the large batch transfer time from CPU to GPU
+ if FLAGS.method != "hgdagger":
+ for critic_step in range(config.cta_ratio - 1):
+ with timer.context("sample_replay_buffer"):
+ batch = next(replay_iterator)
+ demo_batch = next(demo_iterator)
+ batch = concat_batches(batch, demo_batch, axis=0)
+ pref_batch = next(preference_iterator) if preference_buffer is not None else None
+ bc_batch = next(bc_iterator) if bc_buffer is not None else None
+
+ with timer.context("train_critics"):
+ agent, critics_info = agent.update(
+ batch,
+ networks_to_update=train_critic_networks_to_update,
+ pref_batch=pref_batch,
+ bc_batch=bc_batch,
+ )
+
+ with timer.context("train"):
+ batch = next(replay_iterator)
+ demo_batch = next(demo_iterator)
+ batch = concat_batches(batch, demo_batch, axis=0)
+ pref_batch = next(preference_iterator) if preference_buffer is not None else None
+ bc_batch = next(bc_iterator) if bc_buffer is not None else None
+
+ if FLAGS.method == "hgdagger":
+ agent, update_info = agent.update_bc(bc_batch)
+ else:
+ agent, update_info = agent.update(
+ batch,
+ networks_to_update=train_networks_to_update,
+ pref_batch=pref_batch,
+ bc_batch=bc_batch,
+ )
+
+ # publish the updated network
+ if step > 0 and step % (config.steps_per_update) == 0:
+ agent = jax.block_until_ready(agent)
+ server.publish_network(agent.state.params)
+ server.publish_network({'step': step - config.pretraining_steps})
+
+ if step % config.log_period == 0 and wandb_logger:
+ wandb_logger.log(update_info, step=step) # + config.pretraining_steps)
+ wandb_logger.log({"timer": timer.get_average_times()}, step=step) # + config.pretraining_steps)
+
+
+##############################################################################
+
+
+def main(_):
+ global config
+ config = CONFIG_MAPPING[FLAGS.exp_name]()
+ enable_cl = FLAGS.method in ["cl", "soft_cl"]
+ enable_cl = FLAGS.method in ["cl", "soft_cl"]
+
+ if config.rlif_minus_one:
+ print_green("Using RLIF.")
+ if enable_cl:
+ print_green("Using CL.")
+ if FLAGS.method == "hgdagger":
+ print_green("Using HG-DAgger.")
+ assert FLAGS.use_bc_loss
+
+ assert config.batch_size % num_devices == 0
+ # seed
+ rng = jax.random.PRNGKey(FLAGS.seed)
+ rng, sampling_rng = jax.random.split(rng)
+
+ assert FLAGS.exp_name in CONFIG_MAPPING, "Experiment folder not found."
+ env = config.get_environment(
+ fake_env=FLAGS.learner,
+ save_video=FLAGS.save_video,
+ classifier=True,
+ )
+ env = RecordEpisodeStatistics(env)
+
+ rng, sampling_rng = jax.random.split(rng)
+
+ if config.setup_mode == 'single-arm-fixed-gripper' or config.setup_mode == 'dual-arm-fixed-gripper':
+ agent: SACAgent = make_sac_pixel_agent(
+ seed=FLAGS.seed,
+ sample_obs=env.observation_space.sample(),
+ sample_action=env.action_space.sample(),
+ image_keys=config.image_keys,
+ encoder_type=config.encoder_type,
+ discount=config.discount,
+ )
+ include_grasp_penalty = False
+ elif config.setup_mode == 'single-arm-learned-gripper':
+ cl_config = {
+ "enabled": enable_cl,
+ "soft": FLAGS.method == "soft_cl",
+ "enable_margin_constraint": True,
+ "enable_action_constraint": True,
+ "constraint_coeff": 1.0,
+ "constraint_eps": 0.0,
+ "reward_coeff": 1.0,
+ }
+ print_green(f"Using CL Config: {cl_config}")
+
+ agent: SACAgentHybridSingleArm = make_sac_pixel_agent_hybrid_single_arm(
+ seed=FLAGS.seed,
+ sample_obs=env.observation_space.sample(),
+ sample_action=env.action_space.sample(),
+ image_keys=config.image_keys,
+ encoder_type=config.encoder_type,
+ discount=config.discount,
+ enable_cl=enable_cl,
+ cl_config=cl_config,
+ has_image=not FLAGS.state_based,
+ use_bc_loss=FLAGS.use_bc_loss
+ )
+ include_grasp_penalty = True
+ elif config.setup_mode == 'dual-arm-learned-gripper':
+ agent: SACAgentHybridDualArm = make_sac_pixel_agent_hybrid_dual_arm(
+ seed=FLAGS.seed,
+ sample_obs=env.observation_space.sample(),
+ sample_action=env.action_space.sample(),
+ image_keys=config.image_keys,
+ encoder_type=config.encoder_type,
+ discount=config.discount,
+ )
+ include_grasp_penalty = True
+ else:
+ raise NotImplementedError(f"Unknown setup mode: {config.setup_mode}")
+
+ # replicate agent across devices
+ # need the jnp.array to avoid a bug where device_put doesn't recognize primitives
+ agent = jax.device_put(
+ jax.tree_util.tree_map(jnp.array, agent), sharding.replicate()
+ )
+
+ if FLAGS.checkpoint_path is not None and os.path.exists(FLAGS.checkpoint_path):
+ input("Checkpoint path already exists. Press Enter to resume training.")
+ ckpt = checkpoints.restore_checkpoint(
+ os.path.abspath(FLAGS.checkpoint_path),
+ agent.state,
+ )
+ agent = agent.replace(state=ckpt)
+ ckpt_number = os.path.basename(
+ checkpoints.latest_checkpoint(os.path.abspath(FLAGS.checkpoint_path))
+ )[11:]
+ print_green(f"Loaded previous checkpoint at step {ckpt_number}.")
+
+ def create_replay_buffer_and_wandb_logger():
+ replay_buffer = MemoryEfficientReplayBufferDataStore(
+ env.observation_space,
+ env.action_space,
+ capacity=config.replay_buffer_capacity,
+ image_keys=config.image_keys,
+ include_grasp_penalty=include_grasp_penalty,
+ )
+ # set up wandb and logging
+ wandb_logger = make_wandb_logger(
+ project="hil-serl-cubereach2",
+ description=FLAGS.exp_name,
+ debug=FLAGS.debug,
+ )
+ return replay_buffer, wandb_logger
+
+ def create_demo_buffer():
+ demo_buffer = MemoryEfficientReplayBufferDataStore(
+ env.observation_space,
+ env.action_space,
+ capacity=config.replay_buffer_capacity,
+ image_keys=config.image_keys,
+ include_grasp_penalty=include_grasp_penalty,
+ )
+ return demo_buffer
+
+ def create_preference_buffer():
+ if FLAGS.method in ["rlif", "hgdagger", "hil"]:
+ return None
+ preference_buffer = PreferenceBufferDataStore(
+ env.observation_space,
+ env.observation_space,
+ env.action_space,
+ env.action_space,
+ config.replay_buffer_capacity,
+ )
+ return preference_buffer
+
+ def create_bc_buffer():
+ if not FLAGS.use_bc_loss:
+ return None
+ bc_buffer = MemoryEfficientReplayBufferDataStore(
+ env.observation_space,
+ env.action_space,
+ capacity=config.replay_buffer_capacity,
+ image_keys=config.image_keys,
+ include_grasp_penalty=include_grasp_penalty,
+ )
+ return bc_buffer
+
+ if FLAGS.learner:
+ sampling_rng = jax.device_put(sampling_rng, device=sharding.replicate())
+ replay_buffer, wandb_logger = create_replay_buffer_and_wandb_logger()
+ demo_buffer = create_demo_buffer()
+ preference_buffer = create_preference_buffer()
+ bc_buffer = create_bc_buffer()
+
+ prev_checkpoint_exist: bool = FLAGS.checkpoint_path is not None and os.path.exists(FLAGS.checkpoint_path)
+
+ if prev_checkpoint_exist:
+ assert os.path.exists(os.path.join(FLAGS.checkpoint_path, "buffer"))
+ for file in glob.glob(os.path.join(FLAGS.checkpoint_path, "buffer/*.pkl")):
+ with open(file, "rb") as f:
+ transitions = pkl.load(f)
+ for transition in transitions:
+ replay_buffer.insert(transition)
+ print_green(
+ f"Loaded previous buffer data. Replay buffer size: {len(replay_buffer)}"
+ )
+
+ if prev_checkpoint_exist and preference_buffer is not None:
+ assert os.path.exists(os.path.join(FLAGS.checkpoint_path, "preference_buffer"))
+ for file in glob.glob(os.path.join(FLAGS.checkpoint_path, "preference_buffer/*.pkl")):
+ with open(file, "rb") as f:
+ preferences = pkl.load(f)
+ for preference in preferences:
+ preference_buffer.insert(preference)
+ print_green(
+ f"Loaded previous preference buffer data. Preference buffer size: {len(preference_buffer)}"
+ )
+
+ if prev_checkpoint_exist:
+ assert os.path.exists(os.path.join(FLAGS.checkpoint_path, "demo_buffer"))
+ for file in glob.glob(os.path.join(FLAGS.checkpoint_path, "demo_buffer/*.pkl")):
+ with open(file, "rb") as f:
+ transitions = pkl.load(f)
+ for transition in transitions:
+ demo_buffer.insert(transition)
+ print_green(
+ f"Loaded previous demo buffer data. Demo buffer size: {len(demo_buffer)}"
+ )
+
+ if prev_checkpoint_exist and bc_buffer is not None:
+ assert os.path.exists(os.path.join(FLAGS.checkpoint_path, "bc_buffer"))
+ for file in glob.glob(os.path.join(FLAGS.checkpoint_path, "bc_buffer/*.pkl")):
+ with open(file, "rb") as f:
+ transitions = pkl.load(f)
+ for transition in transitions:
+ bc_buffer.insert(transition)
+ print_green(
+ f"Loaded previous bc buffer data. Bc buffer size: {len(bc_buffer)}"
+ )
+
+ assert FLAGS.demo_path is not None
+ if len(demo_buffer) == 0:
+ num_demos = 0
+ for path in FLAGS.demo_path:
+ with open(path, "rb") as f:
+ transitions = pkl.load(f)
+ for transition in transitions:
+ # Transforming the data
+ for k in set(transition['observations'].keys()) - set(config.image_keys + ['state']):
+ del transition['observations'][k]
+ del transition['next_observations'][k]
+ for k in config.image_keys:
+ img = transition['observations'][k]
+ if img.ndim == 4 and img.shape[0] == 1:
+ img = img[0]
+ transition['observations'][k] = cv2.resize(img, (128, 128))
+ img = transition['next_observations'][k]
+ if img.ndim == 4 and img.shape[0] == 1:
+ img = img[0]
+ transition['next_observations'][k] = cv2.resize(img, (128, 128))
+ if transition['actions'].shape == (4,):
+ transition['actions'] = np.concatenate([transition['actions'][:3], np.zeros((3,)), transition['actions'][3:]], axis=0)
+ transition['grasp_penalty'] = 0
+ assert transition['rewards'] < 1 + 1e-6 and transition['rewards'] > -1e-6, f"{transition['rewards']}"
+
+ num_demos += transition['rewards']
+ transition['rewards'] *= config.reward_scale
+ demo_buffer.insert(transition)
+ if bc_buffer is not None:
+ bc_buffer.insert(transition)
+
+ print_green(f"demo buffer size: {len(demo_buffer)}")
+ print_green(f"demo count: {num_demos}")
+ if bc_buffer is not None:
+ print_green(f"bc buffer size: {len(bc_buffer)}")
+
+ # learner loop
+ print_green("starting learner loop")
+ learner(
+ sampling_rng,
+ agent,
+ replay_buffer,
+ demo_buffer=demo_buffer,
+ wandb_logger=wandb_logger,
+ preference_buffer=preference_buffer,
+ bc_buffer=bc_buffer,
+ )
+
+ elif FLAGS.actor:
+ sampling_rng = jax.device_put(sampling_rng, sharding.replicate())
+ data_store = QueuedDataStore(10000) # the queue size on the actor
+ intvn_data_store = QueuedDataStore(10000)
+ pref_data_store = QueuedDataStore(10000) if FLAGS.method in ["cl", "soft_cl"] else None
+ bc_data_store = QueuedDataStore(10000) if FLAGS.use_bc_loss else None
+
+ if FLAGS.method in ["cl", "soft_cl"]:
+ pref_data_store = QueuedDataStore(10000)
+ else:
+ pref_data_store = None
+ # actor loop
+ print_green("starting actor loop")
+ actor(
+ agent,
+ data_store,
+ intvn_data_store,
+ env,
+ sampling_rng,
+ pref_data_store=pref_data_store,
+ bc_data_store=bc_data_store
+ )
+
+ else:
+ raise NotImplementedError("Must be either a learner or an actor")
+
+
+if __name__ == "__main__":
+ app.run(main)
diff --git a/examples/train_rlpd.py b/examples/train_rlpd.py
index 6b334920..a2a3256a 100644
--- a/examples/train_rlpd.py
+++ b/examples/train_rlpd.py
@@ -13,6 +13,8 @@
import pickle as pkl
from gymnasium.wrappers.record_episode_statistics import RecordEpisodeStatistics
from natsort import natsorted
+from pynput import keyboard
+import requests
from serl_launcher.agents.continuous.sac import SACAgent
from serl_launcher.agents.continuous.sac_hybrid_single import SACAgentHybridSingleArm
@@ -60,14 +62,50 @@
def print_green(x):
return print("\033[92m {}\033[00m".format(x))
+def print_yellow(x):
+ return print("\033[93m {}\033[00m".format(x))
+
##############################################################################
+failure_key = False
+checkpoint_key = False
+pause_key = False
+def on_press(key):
+ global failure_key
+ global checkpoint_key
+ global pause_key
+ try:
+ if str(key) == "'f'":
+ failure_key = True
+ elif str(key) == "'c'":
+ checkpoint_key = True
+ elif str(key) == "'r'":
+ print_yellow("reset")
+ requests.post("http://localhost:5000/reset_gripper")
+ elif str(key) == "'t'":
+ print_yellow("close gripper for reset")
+ requests.post("http://localhost:5000/close_gripper")
+ elif str(key) == "'p'":
+ print_yellow("pause")
+ pause_key = not pause_key
+ except AttributeError:
+ print("error")
+ pass
+
def actor(agent, data_store, intvn_data_store, env, sampling_rng):
"""
This is the actor loop, which runs when "--actor" is set to True.
"""
+ global failure_key
+ global checkpoint_key
+ global pause_key
+
+ listener = keyboard.Listener(
+ on_press=on_press)
+ listener.start()
+
if FLAGS.eval_checkpoint_step:
success_counter = 0
time_list = []
@@ -78,9 +116,13 @@ def actor(agent, data_store, intvn_data_store, env, sampling_rng):
step=FLAGS.eval_checkpoint_step,
)
agent = agent.replace(state=ckpt)
+ print_green(f"Overriding agent with checkpoint at {FLAGS.eval_checkpoint_step}.")
for episode in range(FLAGS.eval_n_trajs):
+ print("reset start")
obs, _ = env.reset()
+ time.sleep(7.0)
+ print("reset end")
done = False
start_time = time.time()
while not done:
@@ -92,10 +134,25 @@ def actor(agent, data_store, intvn_data_store, env, sampling_rng):
)
actions = np.asarray(jax.device_get(actions))
+ t1 = time.time()
+ # breakpoint()
next_obs, reward, done, truncated, info = env.step(actions)
+ t2 = time.time()
+ diff = (t2-t1)
+ # print(f"Control frequency: {1/diff}")
+ # if actions[-1] > 0.5 or actions[-1] < -0.5:
+ # print(f"prev_obs: {obs['state'][0,0]}")
+ # print(f"actions: {actions[-1]}")
+ # print(f"next_obs: {next_obs['state'][0,0]}")
+ if failure_key:
+ # print(next_obs['state'][0,0])
+ # failure_key = False
+ print("failure detected")
+ failure_key = False
+ done = True
obs = next_obs
- if done:
+ if done or truncated:
if reward:
dt = time.time() - start_time
time_list.append(dt)
@@ -108,12 +165,13 @@ def actor(agent, data_store, intvn_data_store, env, sampling_rng):
print(f"success rate: {success_counter / FLAGS.eval_n_trajs}")
print(f"average time: {np.mean(time_list)}")
return # after done eval, return and exit
-
- start_step = (
- int(os.path.basename(natsorted(glob.glob(os.path.join(FLAGS.checkpoint_path, "buffer/*.pkl")))[-1])[12:-4]) + 1
+
+ buffers = (
+ natsorted(glob.glob(os.path.join(FLAGS.checkpoint_path, "buffer/*.pkl")))
if FLAGS.checkpoint_path and os.path.exists(FLAGS.checkpoint_path)
- else 0
+ else []
)
+ start_step = 1 if len(buffers) == 0 else int(os.path.basename(buffers[-1])[12:-4]) + 1
datastore_dict = {
"actor_env": data_store,
@@ -137,9 +195,16 @@ def update_params(params):
client.recv_network_callback(update_params)
transitions = []
+ transitions_full_trajs = []
demo_transitions = []
+ demo_transitions_full_trajs = []
+ interventions = []
+ this_intervention = None
+ print("reset start")
obs, _ = env.reset()
+ time.sleep(5.0)
+ print("reset end")
done = False
# training loop
@@ -148,9 +213,15 @@ def update_params(params):
already_intervened = False
intervention_count = 0
intervention_steps = 0
+ total_interventions = 0
pbar = tqdm.tqdm(range(start_step, config.max_steps), dynamic_ncols=True)
+ print(config.buffer_period)
+ from_time = time.time()
+ cur_steps = 0
for step in pbar:
+ while pause_key:
+ time.sleep(0.5)
timer.tick("total")
with timer.context("sample_actions"):
@@ -167,22 +238,61 @@ def update_params(params):
# Step environment
with timer.context("step_env"):
-
+ t1 = time.time()
+ # breakpoint()
next_obs, reward, done, truncated, info = env.step(actions)
+ t2 = time.time()
+ diff = (t2-t1)
+ # print(f"Control frequency: {1/diff}")
+ # print(f"Generated action: {actions}, got: {reward, done, truncated},")
+ # print(info)
+ # print("---")
+ reward *= config.reward_scale
+ cur_steps += 1
if "left" in info:
info.pop("left")
if "right" in info:
info.pop("right")
+ if failure_key:
+ print("failure detected")
+ failure_key = False
+ truncated = True
# override the action with the intervention action
if "intervene_action" in info:
+ policy_actions = actions
actions = info.pop("intervene_action")
intervention_steps += 1
if not already_intervened:
intervention_count += 1
+ this_intervention = dict(
+ t0=cur_steps,
+ t1=cur_steps+1,
+ observations=[obs, next_obs],
+ actions=[actions],
+ policy_actions=[policy_actions],
+ dones=[done],
+ truncateds=[truncated],
+ )
+ else:
+ this_intervention['observations'].append(next_obs)
+ this_intervention['actions'].append(actions)
+ this_intervention['policy_actions'].append(policy_actions)
+ this_intervention['dones'].append(done)
+ this_intervention['truncateds'].append(truncated)
+ this_intervention['t1'] = cur_steps+1
already_intervened = True
else:
+ if already_intervened:
+ if this_intervention is None:
+ print("Error: Should not be None")
+ interventions.append(this_intervention)
+ this_intervention = None
already_intervened = False
+
+ if (done or truncated) and this_intervention is not None:
+ interventions.append(this_intervention)
+ this_intervention = None
running_return += reward
transition = dict(
@@ -193,6 +303,8 @@ def update_params(params):
masks=1.0 - done,
dones=done,
)
+ # if checkpoint_key:
+ # breakpoint()
if 'grasp_penalty' in info:
transition['grasp_penalty']= info['grasp_penalty']
data_store.insert(transition)
@@ -203,34 +315,70 @@ def update_params(params):
obs = next_obs
if done or truncated:
+ failure_key = False
+ if "episode" not in info:
+ info["episode"] = {}
info["episode"]["intervention_count"] = intervention_count
info["episode"]["intervention_steps"] = intervention_steps
+
+ total_interventions += intervention_count
+
+ info["episode"]["total_interventions"] = total_interventions
+ info["episode"]["intervention_rate"] = total_interventions / (step + 1)
+ info["episode"]["current_intervention_rate"] = intervention_steps / cur_steps
+ info["episode"]["episode_duration"] = time.time() - from_time
+ info["episode"]["success_rate"] = running_return
+ info["episode"]["episode_steps"] = cur_steps
+ info["episode"]["environment_step"] = step
stats = {"environment": info} # send stats to the learner to log
client.request("send-stats", stats)
pbar.set_description(f"last return: {running_return}")
running_return = 0.0
intervention_count = 0
intervention_steps = 0
+ cur_steps = 0
+
already_intervened = False
client.update()
+ print("reset start")
obs, _ = env.reset()
+ transitions_full_trajs = transitions
+ demo_transitions_full_trajs = demo_transitions
+ # input("Waiting for input to proceed...")
+ time.sleep(7.0)
+ print("reset end")
+ from_time = time.time()
if step > 0 and config.buffer_period > 0 and step % config.buffer_period == 0:
# dump to pickle file
buffer_path = os.path.join(FLAGS.checkpoint_path, "buffer")
demo_buffer_path = os.path.join(FLAGS.checkpoint_path, "demo_buffer")
+ interventions_buffer_path = os.path.join(FLAGS.checkpoint_path, "interventions")
if not os.path.exists(buffer_path):
os.makedirs(buffer_path)
if not os.path.exists(demo_buffer_path):
os.makedirs(demo_buffer_path)
+ if not os.path.exists(interventions_buffer_path):
+ os.makedirs(interventions_buffer_path)
with open(os.path.join(buffer_path, f"transitions_{step}.pkl"), "wb") as f:
- pkl.dump(transitions, f)
- transitions = []
+ fp = os.path.join(buffer_path, f"transitions_{step}.pkl")
+ print(f"Dumping {len(transitions_full_trajs)} expert transitions out of {len(transitions)} to {fp} !!!")
+ pkl.dump(transitions_full_trajs, f)
+ transitions = transitions[len(transitions_full_trajs):]
+ transitions_full_trajs = []
with open(
os.path.join(demo_buffer_path, f"transitions_{step}.pkl"), "wb"
) as f:
- pkl.dump(demo_transitions, f)
- demo_transitions = []
+ fp = os.path.join(demo_buffer_path, f"transitions_{step}.pkl")
+ print(f"Dumping {len(demo_transitions_full_trajs)} expert transitions out of {len(demo_transitions)} to {fp} !!!")
+ pkl.dump(demo_transitions_full_trajs, f)
+ demo_transitions = demo_transitions[len(demo_transitions_full_trajs):]
+ demo_transitions_full_trajs = []
+ with open(os.path.join(interventions_buffer_path, f"transitions_{step}.pkl"), "wb") as f:
+ fp = os.path.join(interventions_buffer_path, f"transitions_{step}.pkl")
+ print(f"Dumping {len(interventions)} interventions to {fp}")
+ pkl.dump(interventions, f)
+ interventions = []
timer.tock("total")
@@ -254,11 +402,19 @@ def learner(rng, agent, replay_buffer, demo_buffer, wandb_logger=None):
)
step = start_step
+ if isinstance(agent, SACAgent):
+ train_critic_networks_to_update = frozenset({"critic"})
+ train_networks_to_update = frozenset({"critic", "actor", "temperature"})
+ else:
+ train_critic_networks_to_update = frozenset({"critic", "grasp_critic"})
+ train_networks_to_update = frozenset({"critic", "grasp_critic", "actor", "temperature"})
+
+
def stats_callback(type: str, payload: dict) -> dict:
"""Callback for when server receives stats request."""
assert type == "send-stats", f"Invalid request type: {type}"
if wandb_logger is not None:
- wandb_logger.log(payload, step=step)
+ wandb_logger.log(payload, step=step + config.pretraining_steps)
return {} # not expecting a response
# Create server
@@ -267,6 +423,20 @@ def stats_callback(type: str, payload: dict) -> dict:
server.register_data_store("actor_env_intvn", demo_buffer)
server.start(threaded=True)
+ if step == 0 and config.pretraining_steps > 0:
+ epochs = config.pretraining_steps // config.batch_size
+ print(f"Pretraining on {len(demo_buffer)} demo steps for {config.pretraining_steps} steps ({epochs} epochs)...")
+ for epoch in tqdm.tqdm(range(epochs)):
+ batch = demo_buffer.sample(config.batch_size)
+ agent, update_info = agent.update(batch, networks_to_update=frozenset({"critic", "grasp_critic", "actor"}))
+ # agent, update_info = agent.update_bc(batch)
+ wandb_logger.log({'pretraining': update_info}, step=(epoch + 1) * config.batch_size)
+ agent = jax.block_until_ready(agent)
+ server.publish_network(agent.state.params)
+ checkpoints.save_checkpoint(
+ os.path.abspath(FLAGS.checkpoint_path), agent.state, step=1, keep=100
+ )
+
# Loop to wait until replay_buffer is filled
pbar = tqdm.tqdm(
total=config.training_starts,
@@ -303,13 +473,6 @@ def stats_callback(type: str, payload: dict) -> dict:
# wait till the replay buffer is filled with enough data
timer = Timer()
-
- if isinstance(agent, SACAgent):
- train_critic_networks_to_update = frozenset({"critic"})
- train_networks_to_update = frozenset({"critic", "actor", "temperature"})
- else:
- train_critic_networks_to_update = frozenset({"critic", "grasp_critic"})
- train_networks_to_update = frozenset({"critic", "grasp_critic", "actor", "temperature"})
for step in tqdm.tqdm(
range(start_step, config.max_steps), dynamic_ncols=True, desc="learner"
@@ -342,8 +505,8 @@ def stats_callback(type: str, payload: dict) -> dict:
server.publish_network(agent.state.params)
if step % config.log_period == 0 and wandb_logger:
- wandb_logger.log(update_info, step=step)
- wandb_logger.log({"timer": timer.get_average_times()}, step=step)
+ wandb_logger.log(update_info, step=step + config.pretraining_steps)
+ wandb_logger.log({"timer": timer.get_average_times()}, step=step + config.pretraining_steps)
if (
step > 0
@@ -376,8 +539,8 @@ def main(_):
env = RecordEpisodeStatistics(env)
rng, sampling_rng = jax.random.split(rng)
-
- if config.setup_mode == 'single-arm-fixed-gripper' or config.setup_mode == 'dual-arm-fixed-gripper':
+
+ if config.setup_mode == 'single-arm-fixed-gripper' or config.setup_mode == 'dual-arm-fixed-gripper':
agent: SACAgent = make_sac_pixel_agent(
seed=FLAGS.seed,
sample_obs=env.observation_space.sample(),
@@ -413,7 +576,7 @@ def main(_):
# replicate agent across devices
# need the jnp.array to avoid a bug where device_put doesn't recognize primitives
agent = jax.device_put(
- jax.tree_map(jnp.array, agent), sharding.replicate()
+ jax.tree_util.tree_map(jnp.array, agent), sharding.replicate()
)
if FLAGS.checkpoint_path is not None and os.path.exists(FLAGS.checkpoint_path):
@@ -438,7 +601,7 @@ def create_replay_buffer_and_wandb_logger():
)
# set up wandb and logging
wandb_logger = make_wandb_logger(
- project="hil-serl",
+ project="hil-serl-cubereach2",
description=FLAGS.exp_name,
debug=FLAGS.debug,
)
@@ -456,14 +619,19 @@ def create_replay_buffer_and_wandb_logger():
)
assert FLAGS.demo_path is not None
+ num_demos = 0
for path in FLAGS.demo_path:
with open(path, "rb") as f:
transitions = pkl.load(f)
for transition in transitions:
if 'infos' in transition and 'grasp_penalty' in transition['infos']:
transition['grasp_penalty'] = transition['infos']['grasp_penalty']
+ assert transition['rewards'] < 1 + 1e-6 and transition['rewards'] > -1e-6, f"{transition['rewards']}"
+ num_demos += transition['rewards']
+ transition['rewards'] *= config.reward_scale
demo_buffer.insert(transition)
print_green(f"demo buffer size: {len(demo_buffer)}")
+ print_green(f"demo count: {num_demos}")
print_green(f"online buffer size: {len(replay_buffer)}")
if FLAGS.checkpoint_path is not None and os.path.exists(
@@ -504,8 +672,8 @@ def create_replay_buffer_and_wandb_logger():
elif FLAGS.actor:
sampling_rng = jax.device_put(sampling_rng, sharding.replicate())
- data_store = QueuedDataStore(50000) # the queue size on the actor
- intvn_data_store = QueuedDataStore(50000)
+ data_store = QueuedDataStore(10000) # the queue size on the actor
+ intvn_data_store = QueuedDataStore(10000)
# actor loop
print_green("starting actor loop")
diff --git a/franka_sim/.gitignore b/franka_sim/.gitignore
new file mode 100644
index 00000000..61bfd845
--- /dev/null
+++ b/franka_sim/.gitignore
@@ -0,0 +1,162 @@
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+# Usually these files are written by a python script from a template
+# before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+# For a library or package, you might want to ignore these files since the code is
+# intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+# However, in case of collaboration, if having platform-specific dependencies or dependencies
+# having no cross-platform support, pipenv may install dependencies that don't work, or not
+# install all needed dependencies.
+#Pipfile.lock
+
+# poetry
+# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
+# This is especially recommended for binary packages to ensure reproducibility, and is more
+# commonly ignored for libraries.
+# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
+#poetry.lock
+
+# pdm
+# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
+#pdm.lock
+# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
+# in version control.
+# https://pdm.fming.dev/#use-with-ide
+.pdm.toml
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
+
+# PyCharm
+# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
+# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
+# and can be added to the global gitignore or merged into this file. For a more nuclear
+# option (not recommended) you can uncomment the following to ignore the entire idea folder.
+#.idea/
+
+*.mp4
diff --git a/franka_sim/README.md b/franka_sim/README.md
new file mode 100644
index 00000000..486f0e05
--- /dev/null
+++ b/franka_sim/README.md
@@ -0,0 +1,22 @@
+# Intro:
+This package provide a simple Franka arm and Robotiq Gripper simulator written in Mujoco.
+It includes a state-based and a vision-based Franka lift cube task environment.
+
+# Installation:
+- From `serl` folder, cd into `franka_sim`.
+- In your `serl` conda environment, run `pip install -e .` to install this package.
+- run `pip install -r requirements.txt` to install sim dependencies.
+
+# Explore the Environments
+- Run `python franka_sim/test/test_gym_env_human.py` to launch a display window and visualize the task.
+
+# Credits:
+- This simulation is initially built by [Kevin Zakka](https://kzakka.com/).
+- Under Kevin's permission, we adopted a Gymnasium environment based on it.
+
+# Notes:
+- Error due to `egl` when running on a CPU machine:
+```bash
+export MUJOCO_GL=egl
+conda install -c conda-forge libstdcxx-ng
+```
diff --git a/franka_sim/franka_sim/__init__.py b/franka_sim/franka_sim/__init__.py
new file mode 100644
index 00000000..967e9e5d
--- /dev/null
+++ b/franka_sim/franka_sim/__init__.py
@@ -0,0 +1,20 @@
+from franka_sim.mujoco_gym_env import GymRenderingSpec, MujocoGymEnv
+
+__all__ = [
+ "MujocoGymEnv",
+ "GymRenderingSpec",
+]
+
+from gym.envs.registration import register
+
+register(
+ id="PandaPickCube-v0",
+ entry_point="franka_sim.envs:PandaPickCubeGymEnv",
+ max_episode_steps=100,
+)
+register(
+ id="PandaPickCubeVision-v0",
+ entry_point="franka_sim.envs:PandaPickCubeGymEnv",
+ max_episode_steps=100,
+ kwargs={"image_obs": True},
+)
diff --git a/franka_sim/franka_sim/controllers/__init__.py b/franka_sim/franka_sim/controllers/__init__.py
new file mode 100644
index 00000000..f309a66d
--- /dev/null
+++ b/franka_sim/franka_sim/controllers/__init__.py
@@ -0,0 +1,5 @@
+from franka_sim.controllers.opspace import opspace
+
+__all__ = [
+ "opspace",
+]
diff --git a/franka_sim/franka_sim/controllers/opspace.py b/franka_sim/franka_sim/controllers/opspace.py
new file mode 100644
index 00000000..4546d7d2
--- /dev/null
+++ b/franka_sim/franka_sim/controllers/opspace.py
@@ -0,0 +1,180 @@
+from typing import Optional, Tuple, Union
+
+import mujoco
+import numpy as np
+from dm_robotics.transformations import transformations as tr
+
+
+def pd_control(
+ x: np.ndarray,
+ x_des: np.ndarray,
+ dx: np.ndarray,
+ kp_kv: np.ndarray,
+ ddx_max: float = 0.0,
+) -> np.ndarray:
+ # Compute error.
+ x_err = x - x_des
+ dx_err = dx
+
+ # Apply gains.
+ x_err *= -kp_kv[:, 0]
+ dx_err *= -kp_kv[:, 1]
+
+ # Limit maximum error.
+ if ddx_max > 0.0:
+ x_err_sq_norm = np.sum(x_err**2)
+ ddx_max_sq = ddx_max**2
+ if x_err_sq_norm > ddx_max_sq:
+ x_err *= ddx_max / np.sqrt(x_err_sq_norm)
+
+ return x_err + dx_err
+
+
+def pd_control_orientation(
+ quat: np.ndarray,
+ quat_des: np.ndarray,
+ w: np.ndarray,
+ kp_kv: np.ndarray,
+ dw_max: float = 0.0,
+) -> np.ndarray:
+ # Compute error.
+ quat_err = tr.quat_diff_active(source_quat=quat_des, target_quat=quat)
+ ori_err = tr.quat_to_axisangle(quat_err)
+ w_err = w
+
+ # Apply gains.
+ ori_err *= -kp_kv[:, 0]
+ w_err *= -kp_kv[:, 1]
+
+ # Limit maximum error.
+ if dw_max > 0.0:
+ ori_err_sq_norm = np.sum(ori_err**2)
+ dw_max_sq = dw_max**2
+ if ori_err_sq_norm > dw_max_sq:
+ ori_err *= dw_max / np.sqrt(ori_err_sq_norm)
+
+ return ori_err + w_err
+
+
+def opspace(
+ model,
+ data,
+ site_id,
+ dof_ids: np.ndarray,
+ pos: Optional[np.ndarray] = None,
+ ori: Optional[np.ndarray] = None,
+ joint: Optional[np.ndarray] = None,
+ pos_gains: Union[Tuple[float, float, float], np.ndarray] = (200.0, 200.0, 200.0),
+ ori_gains: Union[Tuple[float, float, float], np.ndarray] = (200.0, 200.0, 200.0),
+ damping_ratio: float = 1.0,
+ nullspace_stiffness: float = 0.5,
+ max_pos_acceleration: Optional[float] = None,
+ max_ori_acceleration: Optional[float] = None,
+ gravity_comp: bool = True,
+) -> np.ndarray:
+ if pos is None:
+ x_des = data.site_xpos[site_id]
+ else:
+ x_des = np.asarray(pos)
+ if ori is None:
+ xmat = data.site_xmat[site_id].reshape((3, 3))
+ quat_des = tr.mat_to_quat(xmat.reshape((3, 3)))
+ else:
+ ori = np.asarray(ori)
+ if ori.shape == (3, 3):
+ quat_des = tr.mat_to_quat(ori)
+ else:
+ quat_des = ori
+ if joint is None:
+ q_des = data.qpos[dof_ids]
+ else:
+ q_des = np.asarray(joint)
+
+ kp = np.asarray(pos_gains)
+ kd = damping_ratio * 2 * np.sqrt(kp)
+ kp_kv_pos = np.stack([kp, kd], axis=-1)
+
+ kp = np.asarray(ori_gains)
+ kd = damping_ratio * 2 * np.sqrt(kp)
+ kp_kv_ori = np.stack([kp, kd], axis=-1)
+
+ kp_joint = np.full((len(dof_ids),), nullspace_stiffness)
+ kd_joint = damping_ratio * 2 * np.sqrt(kp_joint)
+ kp_kv_joint = np.stack([kp_joint, kd_joint], axis=-1)
+
+ ddx_max = max_pos_acceleration if max_pos_acceleration is not None else 0.0
+ dw_max = max_ori_acceleration if max_ori_acceleration is not None else 0.0
+
+ # Get current state.
+ q = data.qpos[dof_ids]
+ dq = data.qvel[dof_ids]
+
+ # Compute Jacobian of the eef site in world frame.
+ J_v = np.zeros((3, model.nv), dtype=np.float64)
+ J_w = np.zeros((3, model.nv), dtype=np.float64)
+ mujoco.mj_jacSite(
+ model,
+ data,
+ J_v,
+ J_w,
+ site_id,
+ )
+ J_v = J_v[:, dof_ids]
+ J_w = J_w[:, dof_ids]
+ J = np.concatenate([J_v, J_w], axis=0)
+
+ # Compute position PD control.
+ x = data.site_xpos[site_id]
+ dx = J_v @ dq
+ ddx = pd_control(
+ x=x,
+ x_des=x_des,
+ dx=dx,
+ kp_kv=kp_kv_pos,
+ ddx_max=ddx_max,
+ )
+
+ # Compute orientation PD control.
+ quat = tr.mat_to_quat(data.site_xmat[site_id].reshape((3, 3)))
+ if quat @ quat_des < 0.0:
+ quat *= -1.0
+ w = J_w @ dq
+ dw = pd_control_orientation(
+ quat=quat,
+ quat_des=quat_des,
+ w=w,
+ kp_kv=kp_kv_ori,
+ dw_max=dw_max,
+ )
+
+ # Compute inertia matrix in joint space.
+ M = np.zeros((model.nv, model.nv), dtype=np.float64)
+ mujoco.mj_fullM(model, M, data.qM)
+ M = M[dof_ids, :][:, dof_ids]
+
+ # Compute inertia matrix in task space.
+ M_inv = np.linalg.inv(M)
+ Mx_inv = J @ M_inv @ J.T
+ if abs(np.linalg.det(Mx_inv)) >= 1e-2:
+ Mx = np.linalg.inv(Mx_inv)
+ else:
+ Mx = np.linalg.pinv(Mx_inv, rcond=1e-2)
+
+ # Compute generalized forces.
+ ddx_dw = np.concatenate([ddx, dw], axis=0)
+ tau = J.T @ Mx @ ddx_dw
+
+ # Add joint task in nullspace.
+ ddq = pd_control(
+ x=q,
+ x_des=q_des,
+ dx=dq,
+ kp_kv=kp_kv_joint,
+ ddx_max=0.0,
+ )
+ Jnull = M_inv @ J.T @ Mx
+ tau += (np.eye(len(q)) - J.T @ Jnull.T) @ ddq
+
+ if gravity_comp:
+ tau += data.qfrc_bias[dof_ids]
+ return tau
diff --git a/franka_sim/franka_sim/envs/__init__.py b/franka_sim/franka_sim/envs/__init__.py
new file mode 100644
index 00000000..50c68828
--- /dev/null
+++ b/franka_sim/franka_sim/envs/__init__.py
@@ -0,0 +1,5 @@
+from franka_sim.envs.panda_pick_gym_env import PandaPickCubeGymEnv
+
+__all__ = [
+ "PandaPickCubeGymEnv",
+]
diff --git a/franka_sim/franka_sim/envs/panda_pick_gym_env.py b/franka_sim/franka_sim/envs/panda_pick_gym_env.py
new file mode 100644
index 00000000..ed452684
--- /dev/null
+++ b/franka_sim/franka_sim/envs/panda_pick_gym_env.py
@@ -0,0 +1,358 @@
+from pathlib import Path
+from typing import Any, Literal, Tuple, Dict
+
+import gymnasium as gym
+import mujoco
+import numpy as np
+from gymnasium import spaces
+import threading
+import cv2
+import queue
+
+try:
+ import mujoco_py
+except ImportError as e:
+ MUJOCO_PY_IMPORT_ERROR = e
+else:
+ MUJOCO_PY_IMPORT_ERROR = None
+
+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"
+_PANDA_HOME = np.asarray((0, -0.785, 0, -2.35, 0, 1.57, np.pi / 4))
+_CARTESIAN_BOUNDS = np.asarray([[0.2, -0.3, 0], [0.6, 0.3, 0.5]])
+_SAMPLING_BOUNDS = np.asarray([[0.25, -0.25], [0.55, 0.25]])
+
+
+class ImageDisplayer(threading.Thread):
+ def __init__(self, queue, name):
+ threading.Thread.__init__(self)
+ self.queue = queue
+ self.daemon = True # make this a daemon thread
+ self.name = name
+
+ def run(self):
+ while True:
+ img_array = self.queue.get() # retrieve an image from the queue
+ if img_array is None: # None is our signal to exit
+ break
+
+ assert isinstance(img_array, dict)
+ frame = np.concatenate(
+ [cv2.resize(v, (512, 512)) for k, v in img_array.items() if "full" not in k], axis=1
+ )
+
+ cv2.imshow(self.name, frame)
+ cv2.waitKey(1)
+
+
+class PandaPickCubeGymEnv(MujocoGymEnv):
+ metadata = {"render_modes": ["rgb_array", "human"]}
+
+ fixed_block_position: bool
+
+ def __init__(
+ self,
+ action_scale: np.ndarray = np.asarray([0.1, 1]),
+ seed: int = 0,
+ control_dt: float = 0.02,
+ physics_dt: float = 0.002,
+ time_limit: float = 10.0,
+ render_spec: GymRenderingSpec = GymRenderingSpec(),
+ render_mode: Literal["rgb_array", "human"] = "rgb_array",
+ image_obs: bool = False,
+ show_viewer: bool = False,
+ fixed_block_position: bool = False,
+ ):
+ self._action_scale = action_scale
+ self.render_spec = render_spec
+ self.show_viewer = show_viewer
+ self.fixed_block_position = fixed_block_position
+
+ super().__init__(
+ xml_path=_XML_PATH,
+ seed=seed,
+ control_dt=control_dt,
+ physics_dt=physics_dt,
+ time_limit=time_limit,
+ render_spec=render_spec,
+ )
+
+ self.metadata = {
+ "render_modes": [
+ "human",
+ "rgb_array",
+ ],
+ "render_fps": int(np.round(1.0 / self.control_dt)),
+ }
+
+ self.render_mode = render_mode
+ self.camera_id = (0, 1)
+ self.image_obs = image_obs
+
+ # Caching.
+ self._panda_dof_ids = np.asarray(
+ [self._model.joint(f"joint{i}").id for i in range(1, 8)]
+ )
+ self._panda_ctrl_ids = np.asarray(
+ [self._model.actuator(f"actuator{i}").id for i in range(1, 8)]
+ )
+ self._gripper_ctrl_id = self._model.actuator("fingers_actuator").id
+ 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
+ ),
+ }
+ ),
+ }
+ )
+
+ 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.action_space = gym.spaces.Box(
+ low=np.asarray([-1.0, -1.0, -1.0, -1.0]),
+ high=np.asarray([1.0, 1.0, 1.0, 1.0]),
+ dtype=np.float32,
+ )
+
+ # NOTE: gymnasium is used here since MujocoRenderer is not available in gym. It
+ # is possible to add a similar viewer feature with gym, but that can be a future TODO
+ from gymnasium.envs.mujoco.mujoco_rendering import MujocoRenderer
+
+ self.model.vis.global_.offwidth = self.render_spec.width
+ self.model.vis.global_.offheight = self.render_spec.height
+ self._viewer = MujocoRenderer(
+ self.model,
+ self.data,
+ {
+ 'elevation': -30,
+ 'azimuth': 180
+ }
+ )
+ self._viewer.render(self.render_mode)
+
+ if self.show_viewer:
+ self.img_queue = queue.Queue()
+ self.displayer = ImageDisplayer(self.img_queue, "http://127.0.0.1:5000/")
+ self.displayer.start()
+
+ def reset(
+ self, seed=None, **kwargs
+ ) -> Tuple[Dict[str, np.ndarray], Dict[str, Any]]:
+ """Reset the environment."""
+ mujoco.mj_resetData(self._model, self._data)
+
+ # Reset arm to home position.
+ self._data.qpos[self._panda_dof_ids] = _PANDA_HOME
+ mujoco.mj_forward(self._model, self._data)
+
+ # Reset mocap body to home position.
+ tcp_pos = self._data.sensor("2f85/pinch_pos").data
+ self._data.mocap_pos[0] = tcp_pos
+
+ # Sample a new block position.
+ if self.fixed_block_position:
+ block_xy = np.array([0.45905119, 0 ])
+ else:
+ block_xy = np.random.uniform(*_SAMPLING_BOUNDS)
+ self._data.jnt("block").qpos[:3] = (*block_xy, self._block_z)
+ mujoco.mj_forward(self._model, self._data)
+
+ # Cache the initial block height.
+ self._z_init = self._data.sensor("block_pos").data[2]
+ self._z_success = self._z_init + 0.2
+
+ obs = self._compute_observation()
+ return obs, {}
+
+ def step(
+ self, action: np.ndarray
+ ) -> Tuple[Dict[str, np.ndarray], float, bool, bool, Dict[str, Any]]:
+ """
+ take a step in the environment.
+ Params:
+ action: np.ndarray
+
+ Returns:
+ observation: dict[str, np.ndarray],
+ reward: float,
+ done: bool,
+ truncated: bool,
+ info: dict[str, Any]
+ """
+ x, y, z, grasp = action
+
+ # Set the mocap position.
+ pos = self._data.mocap_pos[0].copy()
+ dpos = np.asarray([x, y, z]) * self._action_scale[0]
+ npos = np.clip(pos + dpos, *_CARTESIAN_BOUNDS)
+ self._data.mocap_pos[0] = npos
+
+ # Set gripper grasp.
+ g = self._data.ctrl[self._gripper_ctrl_id] / 255
+ dg = grasp * self._action_scale[1]
+ ng = np.clip(g + dg, 0.0, 1.0)
+ self._data.ctrl[self._gripper_ctrl_id] = ng * 255
+
+ for _ in range(self._n_substeps):
+ tau = opspace(
+ model=self._model,
+ data=self._data,
+ site_id=self._pinch_site_id,
+ dof_ids=self._panda_dof_ids,
+ pos=self._data.mocap_pos[0],
+ ori=self._data.mocap_quat[0],
+ joint=_PANDA_HOME,
+ gravity_comp=True,
+ )
+ self._data.ctrl[self._panda_ctrl_ids] = tau
+ mujoco.mj_step(self._model, self._data)
+
+ obs = self._compute_observation()
+ rew = self._compute_reward()
+ success = self._compute_success()
+ terminated = self.time_limit_exceeded()
+
+ return obs, bool(success), success or terminated, False, {'success': success}
+
+ def render(self):
+ rendered_frames = []
+ rendered_frames_dict = {}
+ for cam_id in self.camera_id:
+ img = self._viewer.render(render_mode="rgb_array", camera_id=cam_id)
+ rendered_frames.append(img)
+ rendered_frames_dict[str(cam_id)] = img
+ if self.show_viewer:
+ self.img_queue.put(rendered_frames_dict)
+ return rendered_frames
+
+ # Helper methods.
+
+ def _compute_observation(self) -> dict:
+ obs = {}
+ obs["state"] = {}
+
+ tcp_pos = self._data.sensor("2f85/pinch_pos").data
+ 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
+
+ # joint_pos = np.stack(
+ # [self._data.sensor(f"panda/joint{i}_pos").data for i in range(1, 8)],
+ # ).ravel()
+ # obs["panda/joint_pos"] = joint_pos.astype(np.float32)
+
+ # joint_vel = np.stack(
+ # [self._data.sensor(f"panda/joint{i}_vel").data for i in range(1, 8)],
+ # ).ravel()
+ # obs["panda/joint_vel"] = joint_vel.astype(np.float32)
+
+ # joint_torque = np.stack(
+ # [self._data.sensor(f"panda/joint{i}_torque").data for i in range(1, 8)],
+ # ).ravel()
+ # obs["panda/joint_torque"] = symlog(joint_torque.astype(np.float32))
+
+ # wrist_force = self._data.sensor("panda/wrist_force").data.astype(np.float32)
+ # obs["panda/wrist_force"] = symlog(wrist_force.astype(np.float32))
+
+ if self.image_obs:
+ 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
+
+ if self.render_mode == "human":
+ self._viewer.render(self.render_mode)
+
+ return obs
+
+ def _compute_reward(self) -> float:
+ block_pos = self._data.sensor("block_pos").data
+ tcp_pos = self._data.sensor("2f85/pinch_pos").data
+ dist = np.linalg.norm(block_pos - tcp_pos)
+ r_close = np.exp(-20 * dist)
+ r_lift = (block_pos[2] - self._z_init) / (self._z_success - self._z_init)
+ r_lift = np.clip(r_lift, 0.0, 1.0)
+ rew = 0.3 * r_close + 0.7 * r_lift
+ return rew
+
+ def _compute_success(self) -> bool:
+ block_pos = self._data.sensor("block_pos").data
+ return block_pos[2] > self._z_success
+
+ def close(self):
+ super().close()
+ if self.show_viewer:
+ self.img_queue.put(None)
+ cv2.destroyAllWindows()
+ self.displayer.join()
+
+
+if __name__ == "__main__":
+ env = PandaPickCubeGymEnv(render_mode="human")
+ env.reset()
+ for i in range(100):
+ env.step(np.random.uniform(-1, 1, 4))
+ env.render()
+ env.close()
diff --git a/franka_sim/franka_sim/envs/utils.py b/franka_sim/franka_sim/envs/utils.py
new file mode 100644
index 00000000..dff1aeec
--- /dev/null
+++ b/franka_sim/franka_sim/envs/utils.py
@@ -0,0 +1,9 @@
+import numpy as np
+
+
+def symlog(x: np.ndarray) -> np.ndarray:
+ return np.sign(x) * np.log1p(abs(x))
+
+
+def symexp(x: np.ndarray) -> np.ndarray:
+ return np.sign(x) * (np.exp(np.abs(x)) - 1)
diff --git a/franka_sim/franka_sim/envs/xmls/arena.xml b/franka_sim/franka_sim/envs/xmls/arena.xml
new file mode 100644
index 00000000..e8b69cd9
--- /dev/null
+++ b/franka_sim/franka_sim/envs/xmls/arena.xml
@@ -0,0 +1,39 @@
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
+
diff --git a/franka_sim/franka_sim/envs/xmls/franka_emika_panda/LICENSE b/franka_sim/franka_sim/envs/xmls/franka_emika_panda/LICENSE
new file mode 100644
index 00000000..d9a10c0d
--- /dev/null
+++ b/franka_sim/franka_sim/envs/xmls/franka_emika_panda/LICENSE
@@ -0,0 +1,176 @@
+ Apache License
+ Version 2.0, January 2004
+ http://www.apache.org/licenses/
+
+ TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+ 1. Definitions.
+
+ "License" shall mean the terms and conditions for use, reproduction,
+ and distribution as defined by Sections 1 through 9 of this document.
+
+ "Licensor" shall mean the copyright owner or entity authorized by
+ the copyright owner that is granting the License.
+
+ "Legal Entity" shall mean the union of the acting entity and all
+ other entities that control, are controlled by, or are under common
+ control with that entity. For the purposes of this definition,
+ "control" means (i) the power, direct or indirect, to cause the
+ direction or management of such entity, whether by contract or
+ otherwise, or (ii) ownership of fifty percent (50%) or more of the
+ outstanding shares, or (iii) beneficial ownership of such entity.
+
+ "You" (or "Your") shall mean an individual or Legal Entity
+ exercising permissions granted by this License.
+
+ "Source" form shall mean the preferred form for making modifications,
+ including but not limited to software source code, documentation
+ source, and configuration files.
+
+ "Object" form shall mean any form resulting from mechanical
+ transformation or translation of a Source form, including but
+ not limited to compiled object code, generated documentation,
+ and conversions to other media types.
+
+ "Work" shall mean the work of authorship, whether in Source or
+ Object form, made available under the License, as indicated by a
+ copyright notice that is included in or attached to the work
+ (an example is provided in the Appendix below).
+
+ "Derivative Works" shall mean any work, whether in Source or Object
+ form, that is based on (or derived from) the Work and for which the
+ editorial revisions, annotations, elaborations, or other modifications
+ represent, as a whole, an original work of authorship. For the purposes
+ of this License, Derivative Works shall not include works that remain
+ separable from, or merely link (or bind by name) to the interfaces of,
+ the Work and Derivative Works thereof.
+
+ "Contribution" shall mean any work of authorship, including
+ the original version of the Work and any modifications or additions
+ to that Work or Derivative Works thereof, that is intentionally
+ submitted to Licensor for inclusion in the Work by the copyright owner
+ or by an individual or Legal Entity authorized to submit on behalf of
+ the copyright owner. For the purposes of this definition, "submitted"
+ means any form of electronic, verbal, or written communication sent
+ to the Licensor or its representatives, including but not limited to
+ communication on electronic mailing lists, source code control systems,
+ and issue tracking systems that are managed by, or on behalf of, the
+ Licensor for the purpose of discussing and improving the Work, but
+ excluding communication that is conspicuously marked or otherwise
+ designated in writing by the copyright owner as "Not a Contribution."
+
+ "Contributor" shall mean Licensor and any individual or Legal Entity
+ on behalf of whom a Contribution has been received by Licensor and
+ subsequently incorporated within the Work.
+
+ 2. Grant of Copyright License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ copyright license to reproduce, prepare Derivative Works of,
+ publicly display, publicly perform, sublicense, and distribute the
+ Work and such Derivative Works in Source or Object form.
+
+ 3. Grant of Patent License. Subject to the terms and conditions of
+ this License, each Contributor hereby grants to You a perpetual,
+ worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+ (except as stated in this section) patent license to make, have made,
+ use, offer to sell, sell, import, and otherwise transfer the Work,
+ where such license applies only to those patent claims licensable
+ by such Contributor that are necessarily infringed by their
+ Contribution(s) alone or by combination of their Contribution(s)
+ with the Work to which such Contribution(s) was submitted. If You
+ institute patent litigation against any entity (including a
+ cross-claim or counterclaim in a lawsuit) alleging that the Work
+ or a Contribution incorporated within the Work constitutes direct
+ or contributory patent infringement, then any patent licenses
+ granted to You under this License for that Work shall terminate
+ as of the date such litigation is filed.
+
+ 4. Redistribution. You may reproduce and distribute copies of the
+ Work or Derivative Works thereof in any medium, with or without
+ modifications, and in Source or Object form, provided that You
+ meet the following conditions:
+
+ (a) You must give any other recipients of the Work or
+ Derivative Works a copy of this License; and
+
+ (b) You must cause any modified files to carry prominent notices
+ stating that You changed the files; and
+
+ (c) You must retain, in the Source form of any Derivative Works
+ that You distribute, all copyright, patent, trademark, and
+ attribution notices from the Source form of the Work,
+ excluding those notices that do not pertain to any part of
+ the Derivative Works; and
+
+ (d) If the Work includes a "NOTICE" text file as part of its
+ distribution, then any Derivative Works that You distribute must
+ include a readable copy of the attribution notices contained
+ within such NOTICE file, excluding those notices that do not
+ pertain to any part of the Derivative Works, in at least one
+ of the following places: within a NOTICE text file distributed
+ as part of the Derivative Works; within the Source form or
+ documentation, if provided along with the Derivative Works; or,
+ within a display generated by the Derivative Works, if and
+ wherever such third-party notices normally appear. The contents
+ of the NOTICE file are for informational purposes only and
+ do not modify the License. You may add Your own attribution
+ notices within Derivative Works that You distribute, alongside
+ or as an addendum to the NOTICE text from the Work, provided
+ that such additional attribution notices cannot be construed
+ as modifying the License.
+
+ You may add Your own copyright statement to Your modifications and
+ may provide additional or different license terms and conditions
+ for use, reproduction, or distribution of Your modifications, or
+ for any such Derivative Works as a whole, provided Your use,
+ reproduction, and distribution of the Work otherwise complies with
+ the conditions stated in this License.
+
+ 5. Submission of Contributions. Unless You explicitly state otherwise,
+ any Contribution intentionally submitted for inclusion in the Work
+ by You to the Licensor shall be under the terms and conditions of
+ this License, without any additional terms or conditions.
+ Notwithstanding the above, nothing herein shall supersede or modify
+ the terms of any separate license agreement you may have executed
+ with Licensor regarding such Contributions.
+
+ 6. Trademarks. This License does not grant permission to use the trade
+ names, trademarks, service marks, or product names of the Licensor,
+ except as required for reasonable and customary use in describing the
+ origin of the Work and reproducing the content of the NOTICE file.
+
+ 7. Disclaimer of Warranty. Unless required by applicable law or
+ agreed to in writing, Licensor provides the Work (and each
+ Contributor provides its Contributions) on an "AS IS" BASIS,
+ WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+ implied, including, without limitation, any warranties or conditions
+ of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+ PARTICULAR PURPOSE. You are solely responsible for determining the
+ appropriateness of using or redistributing the Work and assume any
+ risks associated with Your exercise of permissions under this License.
+
+ 8. Limitation of Liability. In no event and under no legal theory,
+ whether in tort (including negligence), contract, or otherwise,
+ unless required by applicable law (such as deliberate and grossly
+ negligent acts) or agreed to in writing, shall any Contributor be
+ liable to You for damages, including any direct, indirect, special,
+ incidental, or consequential damages of any character arising as a
+ result of this License or out of the use or inability to use the
+ Work (including but not limited to damages for loss of goodwill,
+ work stoppage, computer failure or malfunction, or any and all
+ other commercial damages or losses), even if such Contributor
+ has been advised of the possibility of such damages.
+
+ 9. Accepting Warranty or Additional Liability. While redistributing
+ the Work or Derivative Works thereof, You may choose to offer,
+ and charge a fee for, acceptance of support, warranty, indemnity,
+ or other liability obligations and/or rights consistent with this
+ License. However, in accepting such obligations, You may act only
+ on Your own behalf and on Your sole responsibility, not on behalf
+ of any other Contributor, and only if You agree to indemnify,
+ defend, and hold each Contributor harmless for any liability
+ incurred by, or claims asserted against, such Contributor by reason
+ of your accepting any such warranty or additional liability.
+
+ END OF TERMS AND CONDITIONS
diff --git a/franka_sim/franka_sim/envs/xmls/franka_emika_panda/README.md b/franka_sim/franka_sim/envs/xmls/franka_emika_panda/README.md
new file mode 100644
index 00000000..0eab07b7
--- /dev/null
+++ b/franka_sim/franka_sim/envs/xmls/franka_emika_panda/README.md
@@ -0,0 +1,45 @@
+# Franka Emika Panda Description (MJCF)
+
+Requires MuJoCo 2.3.3 or later.
+
+## Overview
+
+This package contains a simplified robot description (MJCF) of the [Franka Emika
+Panda](https://www.franka.de/) developed by [Franka
+Emika](https://www.franka.de/company). It is derived from the [publicly
+available URDF
+description](https://github.com/frankaemika/franka_ros/tree/develop/franka_description).
+
+
+
+
+
+## URDF → MJCF derivation steps
+
+1. Converted the DAE [mesh
+ files](https://github.com/frankaemika/franka_ros/tree/develop/franka_description/meshes/visual)
+ to OBJ format using [Blender](https://www.blender.org/).
+2. Processed `.obj` files with [`obj2mjcf`](https://github.com/kevinzakka/obj2mjcf).
+3. Eliminated the perfectly flat `link0_6` from the resulting submeshes created for `link0`.
+4. Created a convex decomposition of the STL collision [mesh
+ file](https://github.com/frankaemika/franka_ros/tree/develop/franka_description/meshes/collision)
+ for `link5` using [V-HACD](https://github.com/kmammou/v-hacd).
+5. Added `` to the
+ [URDF](https://github.com/frankaemika/franka_ros/tree/develop/franka_description/robots)'s
+ `` clause in order to preserve visual geometries.
+6. Loaded the URDF into MuJoCo and saved a corresponding MJCF.
+7. Matched inertial parameters with [inertial.yaml](
+ https://github.com/frankaemika/franka_ros/blob/develop/franka_description/robots/common/inertial.yaml).
+8. Added a tracking light to the base.
+9. Manually edited the MJCF to extract common properties into the `` section.
+10. Added `
+