| 313 | return actor_state |
| 314 | |
| 315 | class WebLogger: |
| 316 | def __init__(self, dt): |
| 317 | self.state_log = defaultdict(list) |
| 318 | self.rew_log = defaultdict(list) |
| 319 | self.dt = dt |
| 320 | self.num_episodes = 0 |
| 321 | self.app = Flask(__name__) |
| 322 | self.socketio = SocketIO(self.app) |
| 323 | self.thread = None |
| 324 | |
| 325 | def set_robot_limits(self, dof_pos_limits, dof_vel_limits, dof_torque_limits): |
| 326 | self.log_dof_pos_limits = dof_pos_limits |
| 327 | self.log_dof_vel_limits = dof_vel_limits |
| 328 | self.log_dof_torque_limits = dof_torque_limits |
| 329 | |
| 330 | def set_robot_num_dofs(self, num_dofs): |
| 331 | self.robot_num_dofs = num_dofs |
| 332 | |
| 333 | def log_state(self, key, value): |
| 334 | self.state_log[key].append(value) |
| 335 | |
| 336 | def log_states(self, dict): |
| 337 | for key, value in dict.items(): |
| 338 | self.log_state(key, value) |
| 339 | |
| 340 | def log_rewards(self, dict, num_episodes): |
| 341 | for key, value in dict.items(): |
| 342 | if 'rew' in key: |
| 343 | self.rew_log[key].append(value.item() * num_episodes) |
| 344 | self.num_episodes += num_episodes |
| 345 | |
| 346 | def reset(self): |
| 347 | self.state_log.clear() |
| 348 | self.rew_log.clear() |
| 349 | |
| 350 | def plot_states(self): |
| 351 | if self.thread is None: |
| 352 | self.thread = threading.Thread(target=self._run_server) |
| 353 | self.thread.start() |
| 354 | self._update_plot() |
| 355 | |
| 356 | def _run_server(self): |
| 357 | @self.app.route('/') |
| 358 | def index(): |
| 359 | return send_file('analysis_plot_template.html') |
| 360 | |
| 361 | self.socketio.run(self.app, debug=False, use_reloader=False) |
| 362 | |
| 363 | def _update_plot(self): |
| 364 | log = self.state_log.copy() |
| 365 | |
| 366 | total_time = len(next(iter(log.values()))) * self.dt |
| 367 | time = np.linspace(0, total_time, len(next(iter(log.values())))) |
| 368 | |
| 369 | for key in log: |
| 370 | if isinstance(log[key], list): |
| 371 | log[key] = np.array(log[key]) |
| 372 | |