1
0
mirror of https://github.com/microsoft/qlib.git synced 2026-07-14 16:26:55 +08:00
This commit is contained in:
Yuchen Fang
2021-01-28 00:41:02 +08:00
parent 98086e4fdc
commit a03b08bb4c
21 changed files with 154 additions and 563 deletions

View File

@@ -55,18 +55,14 @@ class Collector(object):
def _default_rew_metric(x: Union[Number, np.number]) -> Union[Number, np.number]:
# this internal function is designed for single-agent RL
# for multi-agent RL, a reward_metric must be provided
assert np.asanyarray(x).size == 1, (
"Please specify the reward_metric " "since the reward is not a scalar."
)
assert np.asanyarray(x).size == 1, "Please specify the reward_metric " "since the reward is not a scalar."
return x
def reset(self) -> None:
"""Reset all related variables in the collector."""
# use empty Batch for ``state`` so that ``self.data`` supports slicing
# convert empty Batch to None when passing data to policy
self.data = Batch(
state={}, obs={}, act={}, rew={}, done={}, info={}, obs_next={}, policy={}
)
self.data = Batch(state={}, obs={}, act={}, rew={}, done={}, info={}, obs_next={}, policy={})
self.reset_env()
self.reset_buffer()
self.reset_stat()
@@ -96,9 +92,7 @@ class Collector(object):
self.data.obs = obs
for b in self._cached_buf:
b.reset()
self._ready_env_ids = np.array(
[x for x in self._ready_env_ids if x not in stop_id]
)
self._ready_env_ids = np.array([x for x in self._ready_env_ids if x not in stop_id])
def _reset_state(self, id: Union[int, List[int]]) -> None:
"""Reset the hidden state: self.data.state[id]."""
@@ -187,9 +181,7 @@ class Collector(object):
if isinstance(n_episode, list):
assert len(n_episode) == self.get_env_num()
finished_env_ids = [i for i in self._ready_env_ids if n_episode[i] <= 0]
self._ready_env_ids = np.array(
[x for x in self._ready_env_ids if x not in finished_env_ids]
)
self._ready_env_ids = np.array([x for x in self._ready_env_ids if x not in finished_env_ids])
while True:
if step_count >= 100000 and episode_count.sum() == 0:
warnings.warn(
@@ -249,13 +241,9 @@ class Collector(object):
log_fn(info)
else:
# store computed actions, states, etc
_batch_set_item(
whole_data, self._ready_env_ids, self.data, self.env_num
)
_batch_set_item(whole_data, self._ready_env_ids, self.data, self.env_num)
# fetch finished data
obs_next, rew, done, info = self.env.step(
self.data.act, id=self._ready_env_ids
)
obs_next, rew, done, info = self.env.step(self.data.act, id=self._ready_env_ids)
self._ready_env_ids = np.array([i["env_id"] for i in info])
# get the stepped data
self.data = whole_data[self._ready_env_ids]
@@ -264,9 +252,7 @@ class Collector(object):
step_time += time.time() - start
# move data to self.data
self.data.update(
obs_next=obs_next, rew=rew, done=done, info=[{} for i in info]
)
self.data.update(obs_next=obs_next, rew=rew, done=done, info=[{} for i in info])
if render:
self.env.render()
@@ -288,20 +274,13 @@ class Collector(object):
self._cached_buf[i].add(**self.data[j])
if done[j]:
if not (
isinstance(n_episode, list) and episode_count[i] >= n_episode[i]
):
if not (isinstance(n_episode, list) and episode_count[i] >= n_episode[i]):
episode_count[i] += 1
rewards.append(
self._rew_metric(np.sum(self._cached_buf[i].rew, axis=0))
)
rewards.append(self._rew_metric(np.sum(self._cached_buf[i].rew, axis=0)))
step_count += len(self._cached_buf[i])
if self.buffer is not None:
self.buffer.update(self._cached_buf[i])
if (
isinstance(n_episode, list)
and episode_count[i] >= n_episode[i]
):
if isinstance(n_episode, list) and episode_count[i] >= n_episode[i]:
# env i has collected enough data, it has finished
finished_env_ids.append(i)
self._cached_buf[i].reset()
@@ -318,23 +297,17 @@ class Collector(object):
# env_ind_local.remove(_ready_env_ids.index(i))
if len(env_ind_local) > 0:
if self.preprocess_fn:
obs_reset = self.preprocess_fn(obs=obs_reset).get(
"obs", obs_reset
)
obs_reset = self.preprocess_fn(obs=obs_reset).get("obs", obs_reset)
obs_next[env_ind_local] = obs_reset
reset_time += time.time() - start
self.data.obs = obs_next
if is_async:
# set data back
whole_data = deepcopy(whole_data) # avoid reference in ListBuf
_batch_set_item(
whole_data, self._ready_env_ids, self.data, self.env_num
)
_batch_set_item(whole_data, self._ready_env_ids, self.data, self.env_num)
# let self.data be the data in all environments again
self.data = whole_data
self._ready_env_ids = np.array(
[x for x in self._ready_env_ids if x not in finished_env_ids]
)
self._ready_env_ids = np.array([x for x in self._ready_env_ids if x not in finished_env_ids])
if n_step:
if step_count >= n_step:
break