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