| 114 | return actor_state |
| 115 | |
| 116 | class WebLogger: |
| 117 | def __init__(self, dt): |
| 118 | self.state_log = defaultdict(list) |
| 119 | self.rew_log = defaultdict(list) |
| 120 | self.dt = dt |
| 121 | self.num_episodes = 0 |
| 122 | self.app = Flask(__name__) |
| 123 | self.socketio = SocketIO(self.app) |
| 124 | self.thread = None |
| 125 | |
| 126 | def set_robot_limits(self, dof_pos_limits, dof_vel_limits, dof_torque_limits): |
| 127 | self.log_dof_pos_limits = dof_pos_limits |
| 128 | self.log_dof_vel_limits = dof_vel_limits |
| 129 | self.log_dof_torque_limits = dof_torque_limits |
| 130 | |
| 131 | def set_robot_num_dofs(self, num_dofs): |
| 132 | self.robot_num_dofs = num_dofs |
| 133 | |
| 134 | def log_state(self, key, value): |
| 135 | self.state_log[key].append(value) |
| 136 | |
| 137 | def log_states(self, dict): |
| 138 | for key, value in dict.items(): |
| 139 | self.log_state(key, value) |
| 140 | |
| 141 | def log_rewards(self, dict, num_episodes): |
| 142 | for key, value in dict.items(): |
| 143 | if 'rew' in key: |
| 144 | self.rew_log[key].append(value.item() * num_episodes) |
| 145 | self.num_episodes += num_episodes |
| 146 | |
| 147 | def reset(self): |
| 148 | self.state_log.clear() |
| 149 | self.rew_log.clear() |
| 150 | |
| 151 | def plot_states(self): |
| 152 | if self.thread is None: |
| 153 | self.thread = threading.Thread(target=self._run_server) |
| 154 | self.thread.start() |
| 155 | self._update_plot() |
| 156 | |
| 157 | def _run_server(self): |
| 158 | @self.app.route('/') |
| 159 | def index(): |
| 160 | return send_file('analysis_plot_template.html') |
| 161 | |
| 162 | self.socketio.run(self.app, debug=False, use_reloader=False) |
| 163 | |
| 164 | def _update_plot(self): |
| 165 | log = self.state_log.copy() |
| 166 | total_time = len(next(iter(log.values()))) * self.dt |
| 167 | time = np.linspace(0, total_time, len(next(iter(log.values())))) |
| 168 | |
| 169 | for key in log: |
| 170 | if isinstance(log[key], list): |
| 171 | log[key] = np.array(log[key]) |
| 172 | |
| 173 | BLUE = '#005A9D' |