Updates the environment using the action and returns a `TimeStep`.
(self, action)
| 410 | ) |
| 411 | |
| 412 | def step(self, action): |
| 413 | """Updates the environment using the action and returns a `TimeStep`.""" |
| 414 | if self._reset_next_step: |
| 415 | self._reset_next_step = False |
| 416 | return self.reset() |
| 417 | |
| 418 | self._hooks.before_step(self._physics_proxy, action, self._random_state) |
| 419 | self._observation_updater.prepare_for_next_control_step() |
| 420 | |
| 421 | try: |
| 422 | for i in range(self._n_sub_steps): |
| 423 | self._substep(action) |
| 424 | # The final observation update must happen after all the hooks in |
| 425 | # `self._hooks.after_step` is called. Otherwise, if any of these hooks |
| 426 | # modify the physics state then we might capture an observation that is |
| 427 | # inconsistent with the final physics state. |
| 428 | if i < self._n_sub_steps - 1: |
| 429 | self._observation_updater.update() |
| 430 | physics_is_divergent = False |
| 431 | except control.PhysicsError as e: |
| 432 | if not self._raise_exception_on_physics_error: |
| 433 | logging.warning(e) |
| 434 | physics_is_divergent = True |
| 435 | else: |
| 436 | raise |
| 437 | |
| 438 | self._hooks.after_step(self._physics_proxy, self._random_state) |
| 439 | self._observation_updater.update() |
| 440 | |
| 441 | if not physics_is_divergent: |
| 442 | reward = self._task.get_reward(self._physics_proxy) |
| 443 | discount = self._task.get_discount(self._physics_proxy) |
| 444 | terminating = ( |
| 445 | self._task.should_terminate_episode(self._physics_proxy) |
| 446 | or self._physics.time() >= self._time_limit |
| 447 | ) |
| 448 | else: |
| 449 | reward = 0.0 |
| 450 | discount = 0.0 |
| 451 | terminating = True |
| 452 | |
| 453 | obs = self._observation_updater.get_observation() |
| 454 | |
| 455 | if not terminating: |
| 456 | return dm_env.TimeStep(dm_env.StepType.MID, reward, discount, obs) |
| 457 | else: |
| 458 | self._reset_next_step = True |
| 459 | return dm_env.TimeStep(dm_env.StepType.LAST, reward, discount, obs) |
| 460 | |
| 461 | def _substep(self, action): |
| 462 | self._hooks.before_substep( |