Skip to content

Commit 517eae9

Browse files
aaronxie0000budzianowskijingxiangmoCopilot
authored
Fixed projected gravity computation + Added deployment visualization (#62)
* add fields "model" for model used, timestamp, and quat to the log deployments. add 3d visualizer for the quat over the course of rollout * linting * measure imu freq * remove aritificial fixing, fix proj gravity computation * add xax requirement * Update ksim_kbot/deploy/deployment_logs/plot_logs.py Co-authored-by: Copilot <[email protected]> * Update ksim_kbot/deploy/deployment_logs/plot_3d_imu.py Co-authored-by: Copilot <[email protected]> --------- Co-authored-by: budzianowski <[email protected]> Co-authored-by: JX <[email protected]> Co-authored-by: Copilot <[email protected]>
1 parent 7e9713d commit 517eae9

File tree

8 files changed

+737
-123
lines changed

8 files changed

+737
-123
lines changed

‎.gitignore‎

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -14,7 +14,6 @@ __pycache__/
1414

1515
# Databases
1616
*.db
17-
*.pdf
1817

1918
# Build artifacts
2019
build/
@@ -28,8 +27,9 @@ experimental/
2827

2928
# Deployment Assets
3029
# */deploy/assets/*/*
31-
3230
*.pkl
31+
*.pdf
32+
*.mp4
3333

3434
# Development
3535
scratch/

‎ksim_kbot/deploy/__init__.py‎

Lines changed: 0 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1 +0,0 @@
1-
from ksim_kbot.deploy.deploy import Deploy, FixedArmDeploy

‎ksim_kbot/deploy/deploy.py‎

Lines changed: 10 additions & 51 deletions
Original file line numberDiff line numberDiff line change
@@ -157,7 +157,14 @@ def save_rollout(self) -> None:
157157
if self.rollout_dict is not None:
158158
file_dir = os.path.dirname(os.path.abspath(__file__))
159159
timestamp = time.strftime("%Y%m%d-%H%M%S")
160-
with open(f"{file_dir}/deployment_checks/{self.mode}_{timestamp}.pkl", "wb") as f:
160+
date = time.strftime("%Y%m%d")
161+
date_dir = f"{file_dir}/deployment_logs/{date}"
162+
163+
# Check if date directory exists, create if it doesn't
164+
if not os.path.exists(date_dir):
165+
os.makedirs(date_dir)
166+
167+
with open(f"{date_dir}/{self.mode}_{timestamp}.pkl", "wb") as f:
161168
pickle.dump(self.rollout_dict, f)
162169

163170
@abstractmethod
@@ -272,8 +279,10 @@ async def run(self, episode_length: int) -> None:
272279
if time.time() < target_time:
273280
logger.debug(f"Sleeping for {max(0, target_time - time.time())} seconds")
274281
await asyncio.sleep(max(0, target_time - time.time()))
282+
self.rollout_dict["loop_overrun_time"].append(0.0)
275283
else:
276284
logger.info(f"Loop overran by {time.time() - target_time} seconds")
285+
self.rollout_dict["loop_overrun_time"].append(time.time() - target_time)
277286

278287
target_time += self.DT
279288

@@ -283,53 +292,3 @@ async def run(self, episode_length: int) -> None:
283292
raise KeyboardInterrupt
284293

285294
await self.postflight()
286-
287-
288-
class FixedArmDeploy(Deploy):
289-
"""Deploy class for fixed-arm policies."""
290-
291-
def __init__(self, model_path: str, mode: str, ip: str) -> None:
292-
super().__init__(model_path, mode, ip)
293-
# Override the arm positions.
294-
self.arm_positions = {
295-
11: 0.0,
296-
12: 0.0,
297-
13: 0.0,
298-
14: -80.0,
299-
21: 0.0,
300-
22: 0.0,
301-
23: 0.0,
302-
24: 80.0,
303-
25: 0.0,
304-
}
305-
306-
async def send_actions(self, position: np.ndarray, velocity: np.ndarray) -> None:
307-
"""Send actions to the robot's actuators with additional functionality.
308-
309-
Args:
310-
position: Position commands in radians
311-
velocity: Velocity commands in radians/s
312-
"""
313-
position = np.rad2deg(position)
314-
velocity = np.rad2deg(velocity)
315-
316-
actuator_commands: list[pykos.services.actuator.ActuatorCommand] = [
317-
{
318-
"actuator_id": ac.actuator_id,
319-
"position": (
320-
self.arm_positions[ac.actuator_id] if ac.actuator_id in self.arm_positions else position[ac.nn_id]
321-
),
322-
"velocity": (0.0 if ac.actuator_id in self.arm_positions else velocity[ac.nn_id]),
323-
}
324-
for ac in self.actuator_list
325-
]
326-
327-
if self.mode == "real-deploy":
328-
await self.kos.actuator.command_actuators(actuator_commands)
329-
self.rollout_dict["command"].append(actuator_commands)
330-
elif self.mode == "real-check":
331-
self.rollout_dict["command"].append(actuator_commands)
332-
elif self.mode == "sim":
333-
# For all other modes, log and send commands
334-
await self.kos.actuator.command_actuators(actuator_commands)
335-
self.rollout_dict["command"].append(actuator_commands)

‎ksim_kbot/deploy/deploy_joystick.py‎

Lines changed: 21 additions & 10 deletions
Original file line numberDiff line numberDiff line change
@@ -4,6 +4,7 @@
44
import asyncio
55
import os
66
import sys
7+
import time
78

89
import numpy as np
910
from loguru import logger # to be removed
@@ -48,6 +49,9 @@ def __init__(self, enable_joystick: bool, model_path: str, mode: str, ip: str) -
4849
self.phase = np.array([0, np.pi])
4950

5051
self.rollout_dict = {
52+
"model_name": "/".join(model_path.split("/")[-2:]),
53+
"timestamp": [],
54+
"loop_overrun_time": [],
5155
"command": [],
5256
"pos_diff": [],
5357
"vel_obs": [],
@@ -98,14 +102,14 @@ async def get_observation(self) -> np.ndarray:
98102

99103
cmd = self.get_command()
100104

101-
if self.mode in ["sim", "real-check", "real-deploy"]:
102-
self.rollout_dict["pos_diff"].append(pos_diff)
103-
self.rollout_dict["vel_obs"].append(vel_obs)
104-
self.rollout_dict["imu_accel"].append(imu_accel)
105-
self.rollout_dict["imu_gyro"].append(imu_gyro)
106-
self.rollout_dict["controller_cmd"].append(cmd)
107-
self.rollout_dict["prev_action"].append(self.prev_action)
108-
self.rollout_dict["phase"].append(phase_vec)
105+
self.rollout_dict["timestamp"].append(time.time())
106+
self.rollout_dict["pos_diff"].append(pos_diff)
107+
self.rollout_dict["vel_obs"].append(vel_obs)
108+
self.rollout_dict["imu_accel"].append(imu_accel)
109+
self.rollout_dict["imu_gyro"].append(imu_gyro)
110+
self.rollout_dict["controller_cmd"].append(cmd)
111+
self.rollout_dict["prev_action"].append(self.prev_action)
112+
self.rollout_dict["phase"].append(phase_vec)
109113

110114
observation = np.concatenate(
111115
[phase_vec, pos_diff, vel_obs, imu_accel, imu_gyro, cmd, self.gait, self.prev_action]
@@ -124,7 +128,7 @@ def main() -> None:
124128
parser.add_argument("--enable_joystick", action="store_true", help="Enable joystick")
125129
parser.add_argument("--scale_action", type=float, default=0.1, help="Action Scale, default 0.1")
126130
parser.add_argument("--ip", type=str, default="localhost", help="IP address of KOS")
127-
parser.add_argument("--episode_length", type=int, default=5, help="Length of episode in seconds")
131+
parser.add_argument("--episode_length", type=int, default=30, help="Length of episode in seconds")
128132
parser.add_argument("--debug", action="store_true", help="Enable debug logging")
129133

130134
args = parser.parse_args()
@@ -138,7 +142,6 @@ def main() -> None:
138142
# Set global log level
139143
logger.remove()
140144
logger.add(sys.stderr, level=log_level) # This will keep the default colorized format
141-
logger.add(f"{file_dir}/deployment_checks/last_deployment.log", level=log_level)
142145

143146
deploy = JoystickDeploy(args.enable_joystick, model_path, args.mode, args.ip)
144147
deploy.ACTION_SCALE = args.scale_action
@@ -151,5 +154,13 @@ def main() -> None:
151154
raise e
152155

153156

157+
"""
158+
python -m ksim_kbot.deploy.deploy_joystick \
159+
--model_path noisy_joystick_example/tf_model_1576 \
160+
--mode sim \
161+
--scale_action 1.0 \
162+
--debug
163+
"""
164+
154165
if __name__ == "__main__":
155166
main()

‎ksim_kbot/deploy/deploy_joystick_rnn.py‎

Lines changed: 36 additions & 23 deletions
Original file line numberDiff line numberDiff line change
@@ -7,11 +7,11 @@
77
import time
88

99
import numpy as np
10+
import xax
1011
from askin import KeyboardController
1112
from loguru import logger # to be removed
12-
from scipy.spatial.transform import Rotation
1313

14-
from ksim_kbot.deploy.deploy import FixedArmDeploy
14+
from ksim_kbot.deploy.deploy import Deploy
1515

1616

1717
class JoystickCommand:
@@ -30,7 +30,7 @@ async def update(self, key: str) -> None:
3030
self.command[0] += -1 * self.step_size
3131

3232

33-
class JoystickRNNDeploy(FixedArmDeploy):
33+
class JoystickRNNDeploy(Deploy):
3434
"""Deploy class for joystick-controlled policies."""
3535

3636
def __init__(
@@ -76,14 +76,20 @@ def __init__(
7676
self.phase = np.array([0, np.pi])
7777

7878
self.rollout_dict = {
79+
"model_name": "/".join(model_path.split("/")[-2:]),
80+
"timestamp": [],
81+
"loop_overrun_time": [],
7982
"command": [],
8083
"pos_diff": [],
8184
"vel_obs": [],
8285
"projected_gravity": [],
86+
"quat": [],
8387
"imu_gyro": [],
8488
"controller_cmd": [],
8589
"prev_action": [],
8690
"phase": [],
91+
"imu_accel": [],
92+
"euler_angles": [],
8793
}
8894

8995
def get_command(self) -> np.ndarray:
@@ -108,13 +114,7 @@ async def get_observation(self) -> np.ndarray:
108114
)
109115

110116
imu_gyro = np.array([imu.gyro_x, imu.gyro_y, imu.gyro_z])
111-
euler_angles = np.deg2rad(np.array([euler_angles.roll, euler_angles.pitch, euler_angles.yaw]))
112-
113-
# r = Rotation.from_euler("xyz", euler_angles)
114-
r = Rotation.from_quat(np.array([quat.w, quat.x, quat.y, quat.z]), scalar_first=True)
115-
proj_grav_world = r.apply(np.array([0.0, 0.0, 1.0]), inverse=True)
116-
projected_gravity = proj_grav_world
117-
# print(projected_gravity)
117+
projected_gravity = xax.get_projected_gravity_vector_from_quat(np.array([quat.w, quat.x, quat.y, quat.z]))
118118

119119
# * Pos Diff. Difference of current position from default position
120120
state_dict_pos = {state.actuator_id: state.position for state in actuator_states.states}
@@ -135,18 +135,22 @@ async def get_observation(self) -> np.ndarray:
135135

136136
cmd = self.get_command()
137137

138-
if self.mode in ["sim", "real-check", "real-deploy"]:
139-
self.rollout_dict["pos_diff"].append(pos_diff)
140-
self.rollout_dict["vel_obs"].append(vel_obs)
141-
self.rollout_dict["projected_gravity"].append(projected_gravity)
142-
self.rollout_dict["imu_gyro"].append(imu_gyro)
143-
self.rollout_dict["controller_cmd"].append(cmd)
144-
self.rollout_dict["prev_action"].append(self.prev_action)
145-
self.rollout_dict["phase"].append(phase_vec)
138+
imu_accel = np.array([imu.accel_x, imu.accel_y, imu.accel_z])
139+
euler_angles = np.array([euler_angles.roll, euler_angles.pitch, euler_angles.yaw])
140+
141+
self.rollout_dict["timestamp"].append(time.time())
142+
self.rollout_dict["pos_diff"].append(pos_diff)
143+
self.rollout_dict["vel_obs"].append(vel_obs)
144+
self.rollout_dict["projected_gravity"].append(projected_gravity)
145+
self.rollout_dict["quat"].append(np.array([quat.w, quat.x, quat.y, quat.z]))
146+
self.rollout_dict["imu_gyro"].append(imu_gyro)
147+
self.rollout_dict["controller_cmd"].append(cmd)
148+
self.rollout_dict["prev_action"].append(self.prev_action)
149+
self.rollout_dict["phase"].append(phase_vec)
150+
self.rollout_dict["imu_accel"].append(imu_accel)
151+
self.rollout_dict["euler_angles"].append(euler_angles)
146152

147-
observation = np.concatenate(
148-
[phase_vec, pos_diff, vel_obs, projected_gravity, imu_gyro, cmd, self.gait]
149-
).reshape(1, -1)
153+
observation = np.concatenate([phase_vec, pos_diff, vel_obs, projected_gravity, cmd, self.gait]).reshape(1, -1)
150154

151155
return observation
152156

@@ -198,8 +202,10 @@ async def run(self, episode_length: int) -> None:
198202
if time.time() < target_time:
199203
logger.debug(f"Sleeping for {max(0, target_time - time.time())} seconds")
200204
await asyncio.sleep(max(0, target_time - time.time()))
205+
self.rollout_dict["loop_overrun_time"].append(0.0)
201206
else:
202207
logger.info(f"Loop overran by {time.time() - target_time} seconds")
208+
self.rollout_dict["loop_overrun_time"].append(time.time() - target_time)
203209

204210
target_time += self.DT
205211

@@ -221,7 +227,7 @@ def main() -> None:
221227
parser.add_argument("--enable_joystick", action="store_true", help="Enable joystick")
222228
parser.add_argument("--scale_action", type=float, default=0.1, help="Action Scale, default 0.1")
223229
parser.add_argument("--ip", type=str, default="localhost", help="IP address of KOS")
224-
parser.add_argument("--episode_length", type=int, default=5, help="Length of episode in seconds")
230+
parser.add_argument("--episode_length", type=int, default=30, help="Length of episode in seconds")
225231
parser.add_argument("--debug", action="store_true", help="Enable debug logging")
226232

227233
args = parser.parse_args()
@@ -235,7 +241,6 @@ def main() -> None:
235241
# Set global log level
236242
logger.remove()
237243
logger.add(sys.stderr, level=log_level) # This will keep the default colorized format
238-
logger.add(f"{file_dir}/deployment_checks/last_deployment.log", level=log_level)
239244

240245
deploy = JoystickRNNDeploy(args.enable_joystick, model_path, args.mode, args.ip, carry_shape=(5, 256))
241246
deploy.ACTION_SCALE = args.scale_action
@@ -248,5 +253,13 @@ def main() -> None:
248253
raise e
249254

250255

256+
"""
257+
python -m ksim_kbot.deploy.deploy_joystick_rnn \
258+
--model_path noisy_joystick_example/tf_model_1407 \
259+
--mode sim \
260+
--scale_action 1.0 \
261+
--debug
262+
"""
263+
251264
if __name__ == "__main__":
252265
main()

0 commit comments

Comments
 (0)